Update agent.py

agent.py

@@ -1,98 +1,35 @@
 # agent.py
-
 import os
-from typing import Optional, List, Mapping, Any
-
-from google import genai
-from langchain.tools import tool
-from langchain.agents import initialize_agent, AgentType
-from langchain_community.document_loaders import WikipediaLoader
-from langchain.llms.base import LLM
-
-# ——— 1) Gemini SDK client ———
-GENAI_CLIENT = genai.Client(api_key=os.environ.get("GEMINI_API_KEY", ""))
-GEMINI_MODEL = "gemini-1.5-pro"  # or your trial model
-
-# ——— 2) LLM wrapper with class‐level defaults ———
-class GeminiLLM(LLM):
-    """
-    A LangChain-compatible LLM wrapper around Google Gemini.
-    """
-    client: genai.Client = GENAI_CLIENT
-    model: str = GEMINI_MODEL
-
-    @property
-    def _llm_type(self) -> str:
-        return "gemini"
-
-    @property
-    def _identifying_params(self) -> Mapping[str, Any]:
-        return {"model": self.model}
-
-    def _call(
-        self,
-        prompt: str,
-        stop: Optional[List[str]] = None,
-    ) -> str:
-        response = self.client.generate_content(
-            model=self.model,
-            contents=[prompt]
-        )
-        return response.text
+from langgraph import Graph, LLM, tool #or other graph library
 
-# ——— 3) Tools ———
+# 1) Define tools
 
 @tool
 def calculator(expr: str) -> str:
-    """
-    Safely evaluates a math expression and returns the result.
-    """
+    """Simple math via Python eval"""
     try:
-        return str(eval(expr, {"__builtins__": {}}))
-    except Exception as e:
-        return f"Error: {e}"
-
-@tool
-def wiki_search(query: str) -> str:
-    """
-    Loads up to 2 Wikipedia pages for the query and concatenates their text.
-    """
-    docs = WikipediaLoader(query=query, load_max_docs=2).load()
-    return "\n\n".join(d.page_content for d in docs)
-
-# ——— 4) Agent ———
-
-class BasicAgent:
-    """
-    Zero-Shot React agent powered by Gemini + two tools.
-    """
-    def __init__(self):
-        # Ensure your key is present
-        assert os.environ.get("GEMINI_API_KEY"), "❌ GEMINI_API_KEY not set in Secrets"
-
-        # Instantiate our LLM (defaults are on the class)
-        gemini_llm = GeminiLLM()
-
-        # Build the agent
-        self.agent = initialize_agent(
-            tools=[calculator, wiki_search],
-            llm=gemini_llm,
-            agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
-            verbose=True,
-            max_iterations=5,
-            early_stopping_method="generate",
-        )
-
-    def __call__(self, question: str) -> str:
-        # Build the prompt
-        prompt = (
-            "You have access to two tools:\n"
-            "  • calculator(expr)\n"
-            "  • wiki_search(query)\n"
-            "Think internally; output ONLY the final answer.\n\n"
-            f"Question: {question}"
-        )
-        # Run & extract
-        raw = self.agent.run(prompt)
-        lines = [l.strip() for l in raw.splitlines() if l.strip()]
-        return lines[-1] if lines else raw.strip()
+        return str(eval(expr))
+    except Exception:
+        return "Error"
+# e.g. search, vector_retrieval, etc.
+# @tool
+# def web_search(query:str) -> str:
+# ...
+
+# 2) Build your graph
+def build_graph(provider: str = "huggingface") -> Graph:
+    # 2a) Instantiate your LLM endpoint
+    api_token = os.environ["HF_API_TOKEN"]
+    llm = LLM(provider=provider, token=api_token)
+
+    # 2b) Attach tools
+    llm = llm.with_tools([calculator]) # add more tools here
+
+    # 2c) Compose your graph
+    graph = Graph()
+    graph.add_node("ask", llm) # prompt node
+    graph.add_node("calc", calculator) # tool node
+    graph.add_edge("ask", "calc") # allow ask -> calc
+    graph.set.start("ask")
+    return graph
+