Upload 3 files

agent.py

@@ -0,0 +1,150 @@
+"""LangGraph Agent"""
+import os
+from dotenv import load_dotenv
+from langgraph.graph import START, StateGraph, MessagesState
+from langgraph.prebuilt import tools_condition
+from langgraph.prebuilt import ToolNode
+from langchain_community.tools.tavily_search import TavilySearchResults
+from langchain_community.document_loaders import WikipediaLoader
+from langchain_core.messages import HumanMessage
+from langchain_core.tools import tool
+from langchain_mistralai import ChatMistralAI
+
+load_dotenv()
+
+
+def web_search(query: str) -> str:
+    """Search Tavily for a query and return maximum 3 results.
+    
+    Args:
+        query: The search query."""
+    search_docs = TavilySearchResults(max_results=3).invoke(input=query)
+    formatted_search_docs = "\n\n---\n\n".join(
+    [
+        f'<Document source="{doc["title"]}" page="{doc.get("url", "")}"/>\n{doc['content']}\n</Document>'
+        for doc in search_docs
+    ])
+    return {"web_results": formatted_search_docs}
+
+
+def wiki_search(query: str) -> str:
+    """Search Wikipedia for a query and return maximum 2 results.
+    
+    Args:
+        query: The search query."""
+    search_docs = WikipediaLoader(query=query, load_max_docs=2).load()
+    formatted_search_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 search_docs
+        ])
+    return {"wiki_results": formatted_search_docs}
+
+
+@tool
+def multiply(a: int, b: int) -> int:
+    """Multiply two numbers.
+
+    Args:
+        a: first int
+        b: second int
+    """
+    return a * b
+
+@tool
+def add(a: int, b: int) -> int:
+    """Add two numbers.
+    
+    Args:
+        a: first int
+        b: second int
+    """
+    return a + b
+
+@tool
+def subtract(a: int, b: int) -> int:
+    """Subtract two numbers.
+    
+    Args:
+        a: first int
+        b: second int
+    """
+    return a - b
+
+@tool
+def divide(a: int, b: int) -> int:
+    """Divide two numbers.
+    
+    Args:
+        a: first int
+        b: second int
+    """
+    if b == 0:
+        raise ValueError("Cannot divide by zero.")
+    return a / b
+
+@tool
+def modulus(a: int, b: int) -> int:
+    """Get the modulus of two numbers.
+    
+    Args:
+        a: first int
+        b: second int
+    """
+    return a % b
+
+tools = [
+    multiply,
+    add,
+    subtract,
+    divide,
+    modulus,
+    wiki_search,
+    web_search,
+    
+]
+
+def build_graph():
+
+    
+    llm = ChatMistralAI(
+            model="mistral-small-2503",
+            temperature=0,
+            max_retries=2,
+        )
+
+    llm_with_tools = llm.bind_tools(tools)
+
+
+    # Node
+    def assistant(state: MessagesState):
+        """Assistant node"""
+        return {"messages": [llm_with_tools.invoke(state["messages"])]}
+
+
+    builder = StateGraph(MessagesState)
+    builder.add_node("assistant", assistant)
+    builder.add_node("tools", ToolNode(tools))
+    builder.add_edge(START, "assistant")
+    builder.add_conditional_edges(
+        "assistant",
+        tools_condition,
+    )
+    builder.add_edge("tools", "assistant")
+
+    # Compile graph
+    return builder.compile()
+
+if __name__=="__main__":
+    system_prompt = "You are a general AI assistant. I will ask you a question. Report your thoughts, and finish your answer with the following template: FINAL ANSWER: [YOUR FINAL ANSWER]. YOUR FINAL ANSWER should be a number OR as few words as possible OR a comma separated list of numbers and/or strings. If you are asked for a number, don't use comma to write your number neither use units such as $ or percent sign unless specified otherwise. If you are asked for a string, don't use articles, neither abbreviations (e.g. for cities), and write the digits in plain text unless specified otherwise. If you are asked for a comma separated list, apply the above rules depending of whether the element to be put in the list is a number or a string."
+    question = "whats 2 + 15?"
+    question = f"{system_prompt} \
+    The question: \
+    {question} \
+    "
+    graph = build_graph()
+    messages = [HumanMessage(content=question)]
+    messages = graph.invoke({"messages": messages})
+    for m in messages["messages"]:
+        m.pretty_print()
+        

app.py

@@ -3,8 +3,9 @@ import gradio as gr
 import requests
 import inspect
 import pandas as pd
+from langchain_core.messages import HumanMessage
+from agent import build_graph
 
-from mistralai import Mistral
 
 # (Keep Constants as is)
 # --- Constants ---
@@ -18,45 +19,22 @@ DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
 class BasicAgent:
     def __init__(self):
         print("BasicAgent initialized.")
+        self.graph = build_graph()
 
     def __call__(self, question: str) -> str:
         print(f"Agent received question (first 50 chars): {question[:50]}...")
 
-        
-        api_key = os.environ["MISTRAL_API_KEY"]
-        mistral_client = Mistral(api_key=api_key)
-        model = "mistral-small-2503"
-
-        system_prompt = "You are a general AI assistant. I will ask you a question. Report your thoughts, and finish your answer with the following template: FINAL ANSWER: [YOUR FINAL ANSWER]. YOUR FINAL ANSWER should be a number OR as few words as possible OR a comma separated list of numbers and/or strings. If you are asked for a number, don't use comma to write your number neither use units such as $ or percent sign unless specified otherwise. If you are asked for a string, don't use articles, neither abbreviations (e.g. for cities), and write the digits in plain text unless specified otherwise. If you are asked for a comma separated list, apply the above rules depending of whether the element to be put in the list is a number or a string."
-        question = f"{system_prompt} \
-        The question: \
-        {question} \
-        "
-
-
-
-        response = mistral_client.chat.complete(
-        model= model,
-        messages = [
-            {
-                "role": "user",
-                    "content": question,
-                },
-            ]
-        )
-
-        answer = response.choices[0].message.content
-        
         #answer = 'static answer'
-
         #answer = "This is a default answer."
         #print(f"Agent returning fixed answer: {answer}")
-        return answer
-
-
-
 
 
+        # Wrap the question in a HumanMessage from langchain_core
+        messages = [HumanMessage(content=question)]
+        messages = self.graph.invoke({"messages": messages})
+        answer = messages['messages'][-1].content
+        return answer[14:]
+        
 
 
 

requirements.txt

@@ -1,3 +1,12 @@
 gradio
 requests
-mistralai
+mistralai
+langchain_community
+duckduckgo-search
+langchain
+langchain-core
+langchain-tavily
+langgraph
+pymupdf
+wikipedia
+python-dotenv