Update agent.py

Modified agent to use langgraph

agent.py

@@ -1,253 +1,130 @@
-from smolagents import HfApiModel, CodeAgent, DuckDuckGoSearchTool, WikipediaSearchTool, Tool, LiteLLMModel
 from langchain_tavily import TavilySearch
-from langchain_community.document_loaders import WikipediaLoader, ArxivLoader
-import os, time, math, pandas, numpy, PIL
-
-
-class add(Tool):
-    name = "add"
-    description = """
-        This tool adds two integers together and returns an integer.
-
-        Args:
-            a: first integer
-            b: second integer
-        """
-    inputs = {
-        "a":{
-            "type":"integer",
-            "description":"first integer"
-        },
-        "b":{
-            "type":"integer",
-            "description":"second integer"
-        }
-    }
-    output_type = "integer"
-
-    def forward(self, a: int, b: int):
-        return a + b
-
-
-class subtract(Tool):
-    name = "subtract"
-    description = """
-        This tool subtracts two integers and returns an integer.
+from langchain_community.document_loaders import WikipediaLoader
+from langchain_ollama import ChatOllama
+from langgraph.prebuilt import create_react_agent
+from langgraph_supervisor import create_supervisor
+
+llm_model=ChatOllama(model="qwen2:7b", temperature=0.0, max_tokens=1000, base_url="http://localhost:11434")
+
+# Tool for searching wikipedia and returning the content
+def search_wikipedia(query: str) -> str:
+    """Searches Wikipedia for the given query and returns the first result's content.
+    Args:
+        query (str): The search query to use for Wikipedia. Should be a simple keyword or phrase. Do not enter the URL or any other complex query.
+    """
 
-        Args:
+    loader = WikipediaLoader(query=query, load_max_docs=1)
+    docs = loader.load()
+    if docs:
+        return "From: " + docs[0].metadata['source'] + "\n\nContent: " + docs[0].page_content
+    return "No content found."
+
+# Tool for searching the web using Tavily
+def search_web(query: str) -> str:
+    """Searches the web for the given query and returns the first result's content.
+    Args:
+        query (str): The search query to use for web search. Should be a simple keyword or phrase. Do not enter the URL or any other complex query.
+    """
+    results = TavilySearch(max_results=1).invoke(input=query)
+    if results:
+        return "From: " + results['results'][0]['url'] + "\n\nContent: " + results['results'][0]['content']
+    return "No content found."
+
+research_agent = create_react_agent(
+    model=llm_model,
+    tools=[search_wikipedia, search_web],
+    prompt=(
+        "You are a research agent.\n\n"
+        "INSTRUCTIONS:\n"
+        "- Assist ONLY with research-related tasks, DO NOT do any math\n"
+        "- After you're done with your tasks, respond to the supervisor directly\n"
+        "- Respond ONLY with the results of your work, do NOT include ANY other text."
+    ),
+    name="research_agent",
+)
+
+def add(a: int, b: int) -> int:
+    """
+        Add two integers and return the result.
+        Inputs:
             a: first integer
             b: second integer
-        """
-    inputs = {
-        "a":{
-            "type":"integer",
-            "description":"first integer"
-        },
-        "b":{
-            "type":"integer",
-            "description":"second integer"
-        }
-    }
-    output_type = "integer"
-
-    def forward(self, a: int, b: int):
-        return a - b
-
-class multiply(Tool):
-    name = "multiply"
-    description = """
-        This tool multiplies two integers and returns an integer.
+    """
+    return a + b
 
-        Args:
+def subtract(a: int, b: int) -> int:
+    """
+        Subtract two integers and return the result.
+        Inputs:
             a: first integer
             b: second integer
-        """
-    inputs = {
-        "a":{
-            "type":"integer",
-            "description":"first integer"
-        },
-        "b":{
-            "type":"integer",
-            "description":"second integer"
-        }
-    }
-    output_type = "integer"
-
-    def forward(self, a: int, b: int):
-        return a * b
-
-class divide(Tool):
-    name = "divide"
-    description = """
-        This tool divides two integers and returns an integer.
+    """
+    return a - b
 
-        Args:
+def multiply(a: int, b: int) -> int:
+    """
+        Multiply two integers and return the result.
+        Inputs:
             a: first integer
             b: second integer
-        """
-    inputs = {
-        "a":{
-            "type":"integer",
-            "description":"first integer"
-        },
-        "b":{
-            "type":"integer",
-            "description":"second integer"
-        }
-    }
-    output_type = "integer"
-
-    def forward(self, a: int, b: int):
-        if b == 0:
-            raise ValueError("Cannot divide by zero.")
-        return a / b
-
-class modulo(Tool):
-    name = "modulo"
-    description = """
-        This tool returns the modulus of two integers
+    """
+    return a * b
 
-        Args:
+def divide(a: int, b: int) -> float:
+    """
+        Divide two integers and return the result.
+        Inputs:
             a: first integer
             b: second integer
-        """
-    inputs = {
-        "a":{
-            "type":"integer",
-            "description":"first integer"
-        },
-        "b":{
-            "type":"integer",
-            "description":"second integer"
-        }
-    }
-    output_type = "integer"
-
-    def forward(self, a: int, b: int):
-        return a % b
-
-class WikipediaSearchTool(Tool):
-    name = "wikipedia_search_tool"
-    
-    description = """
-        Search Wikipedia for a query and return top 2 results.
-    
-        Args:
-            query: the search query.
-    """
-    
-    inputs = {
-        "query":{
-            "type":"string",
-            "description":"the search query"
-        }
-    }
-    output_type = "string"
-
-    def forward(self, query: str) -> str:
-        documents = WikipediaLoader(query=query, load_max_docs=2).load()
-        condensed_docs = "\n\n---\n\n".join(
-        [
-            f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content}\n</Document>'
-            for doc in documents
-        ])
-        return {"wikipedia_results": condensed_docs}
-
-class TavilySearchTool(Tool):
-    name = "tavily_search_tool"
-    
-    description = """
-        Search the internet using the browser Tavily and return the top 3 results.
-    
-        Args:
-            query: the search query.
+        Note: If b is 0, return "Error: Division by zero is not allowed."
     """
-    
-    inputs = {
-        "query":{
-            "type":"string",
-            "description":"the search query"
-        }
-    }
-    output_type = "string"
-
-    def forward(self, query: str) -> str:
-        documents = TavilySearch(max_results=3).invoke(query=query)
-        condensed_docs = "\n\n---\n\n".join(
-        [
-            f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content}\n</Document>'
-            for doc in documents
-        ])
-        return {"web_search_results": condensed_docs}
-
-
-class ArvixSearchTool(Tool):
-    name = "arvix_search_tool"
-    
-    description = """
-        Search arxiv for a query and return maximum 3 result.
-    
-        Args:
-            query: The search query.
-        """
-    
-    inputs = {
-        "query":{
-            "type":"string",
-            "description":"the search query"
-        }
-    }
-    output_type = "string"
+    if b == 0:
+        return "Error: Division by zero is not allowed."
+    return a / b
 
-    def forward(self, query: str) -> str:
-        documents = ArxivLoader(query=query, load_max_docs=3).load()
-        condensed_docs = "\n\n---\n\n".join(
-        [
-            f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.content[:1000]}\n</Document>'
-            for doc in documents
-        ])
-        return {"arvix_search_results": condensed_docs}
-
-
-tools=[
-    add(),
-    subtract(),
-    multiply(),
-    divide(),
-    modulo(),
-    WikipediaSearchTool(),
-    ArvixSearchTool(),
-    TavilySearchTool()
-]
-
-# model = LiteLLMModel(model_id="gemini/gemini-2.0-flash-exp",
-#                      api_key=os.getenv("GEMINI_API_KEY"))
-
-model = LiteLLMModel(
-        model_id="ollama/gemma3:27b",  # Or try other Ollama-supported models
-        api_base="http://localhost:11434",  # Default Ollama local server
-        num_ctx=8192,
-    )
-
-def delay_execution_10(pagent, **kwargs) -> bool:
+def modulo(a: int, b: int) -> int:
     """
-    Delays the execution for 10 seconds.
+        Calculate the modulo of two integers and return the result.
+        Inputs:
+            a: first integer
+            b: second integer
+        Note: If b is 0, return "Error: Division by zero is not allowed."
     """
-    time.sleep(10)
-    return True
-
+    if b == 0:
+        return "Error: Division by zero is not allowed."
+    return a % b
+
+math_agent = create_react_agent(
+    model=llm_model,
+    tools=[add, subtract, multiply, divide, modulo],
+    prompt=(
+        "You are a math agent.\n\n"
+        "INSTRUCTIONS:\n"
+        "- Assist ONLY with math-related tasks\n"
+        "- After you're done with your tasks, respond to the supervisor directly\n"
+        "- Respond ONLY with the results of your work, do NOT include ANY other text."
+    ),
+    name="math_agent",
+)
 def create_agent():
-    agent = CodeAgent(
-        model = model,
-        tools = tools,
-        max_steps=10,
-        verbosity_level=2,
-        additional_authorized_imports=['*'],
-        planning_interval=5,  
-        step_callbacks=[delay_execution_10]
-    )
-
-    return agent
-    
-
+    supervisor = create_supervisor(
+        model=llm_model,
+        agents=[research_agent, math_agent],
+        prompt=(
+            "You are a supervisor managing two agents:\n"
+            "- a research agent. Assign research-related tasks to this agent\n"
+            "- a math agent. Assign math-related tasks to this agent\n"
+            "Assign work to one agent at a time, do not call agents in parallel.\n"
+            "Do not do any work yourself. When you are ready to give the final answer, do so as simply as possible.If the answer is a number, write only the number do not use commas or any special characters or any other text. \n"
+            "If the answer is a string, write only the string without any additional text.\n"
+            "If the answer is a list, write the items in the list separated by commas while following the above rules for numbers and string where appropriate.\n"
+            "Preface your answer with FINAL ANSWER: {answer}\n"
+            "Example: How many planets are in the solar system?\n Answer: 8\n"
+            "Example: What is the capital of France?\n Answer: Paris\n"
+            "Example: What is 2 + 2?\n Answer: 4\n"
+        ),
+        add_handoff_back_messages=True,
+        output_mode="full_history",
+    ).compile()
+
+    return supervisor