Upload agent.py

agent.py

@@ -0,0 +1,157 @@
+import os
+from dotenv import load_dotenv
+from typing import TypedDict, Annotated
+from langgraph.graph import START, StateGraph, MessagesState
+from langgraph.graph.message import add_messages
+from langchain_core.messages import AnyMessage, HumanMessage, AIMessage, SystemMessage
+from langgraph.prebuilt import ToolNode
+from langgraph.graph import START, StateGraph
+from langgraph.prebuilt import tools_condition
+from langchain_huggingface import HuggingFaceEndpoint, ChatHuggingFace
+from langchain_core.tools import tool
+from langchain_community.document_loaders import WikipediaLoader
+from langchain_google_genai import ChatGoogleGenerativeAI
+from langchain_community.tools import DuckDuckGoSearchRun
+
+load_dotenv()
+
+# ReAct System Prompt
+REACT_SYSTEM_PROMPT = """You are a helpful assistant tasked with answering questions using a set of tools.
+
+You use the ReAct (Reasoning and Acting) methodology to solve problems.
+
+For each user query, you should:
+1. **Thought**: Think step by step about what you need to do to answer the question
+2. **Action**: Use the appropriate tool(s) to gather information or perform calculations
+3. **Observation**: Analyze the results from your actions
+4. Repeat the Thought-Action-Observation cycle as needed until you have enough information to provide a complete answer
+
+When you need to use tools:
+- Think about which tool is most appropriate for the task
+- Use tools to gather information or perform calculations
+- Analyze the results and determine if you need additional information
+- Continue until you can provide a comprehensive answer
+
+Available tools:
+- multiply, add, subtract, divide: For mathematical calculations
+- wikidata_search: Search Wikipedia for factual information
+- duckduckgo_search: Search the web for current information
+
+Now, 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.
+Your answer should only start with "FINAL ANSWER: ", then follows with the answer.
+
+Always explain your reasoning process and show your work step by step."""
+
+@tool
+def multiply(a:int, b:int) -> int:
+    """
+    Multiply two numbers
+    """
+    return a * b
+
+@tool
+def add(a:int, b:int) -> int:
+    """
+    Add two numbers
+    """
+    return a + b
+
+@tool
+def subtract(a:int, b:int) -> int:
+    """
+    Subtract two numbers
+    """
+    return a - b
+
+@tool
+def divide(a:int, b:int) -> int:
+    """
+    Divide two numbers
+    """
+    return a / b
+
+@tool
+def wikidata_search(query: str) -> str:
+    """
+    Search for information on Wikipedia and return maximum 2 results.
+    
+    Args:
+        query: The search query.
+    """
+    loader = WikipediaLoader(query=query, load_max_docs=2)
+    docs = loader.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 docs
+        ])
+    return {"wiki_results": formatted_search_docs}
+
+@tool
+def duckduckgo_search(query: str) -> str:
+    """
+    Search for information on DuckDuckGo and return maximum 3 results.
+    
+    Args:
+        query: The search query.
+    """
+    search = DuckDuckGoSearchRun()
+    # DuckDuckGo returns a string, so we need to format it differently
+    results = search.run(query)
+    formatted_search_docs = f'<Document source="DuckDuckGo Search" page=""/>\n{results}\n</Document>'
+    return {"web_results": formatted_search_docs}
+
+tools = [multiply, add, subtract, divide, wikidata_search, duckduckgo_search]
+
+def build_graph():
+    llm = ChatGoogleGenerativeAI(model="gemini-2.0-flash", api_key=os.getenv("GOOGLE_API_KEY"))
+    llm_with_tools = llm.bind_tools(tools)
+
+    def agent_node(state: MessagesState) -> MessagesState:
+        """This is the agent node with ReAct methodology"""
+        messages = state["messages"]
+        
+        # Add system prompt if not already present
+        if not messages or not isinstance(messages[0], SystemMessage):
+            messages = [SystemMessage(content=REACT_SYSTEM_PROMPT)] + messages
+            
+        return {"messages": [llm_with_tools.invoke(messages)]}
+    
+    
+
+    builder = StateGraph(MessagesState)
+    builder.add_node("agent", agent_node)
+    builder.add_node("tools", ToolNode(tools))
+    
+    
+    builder.add_edge(START, "agent")
+    builder.add_conditional_edges("agent", tools_condition)
+    builder.add_edge("tools", "agent")
+
+    return builder.compile()
+
+class LangGraphAgent:
+    def __init__(self):
+        self.graph = build_graph()
+        print("LangGraphAgent initialized with tools.")
+        
+    def __call__(self, question: str) -> str:
+        """Run the agent on a question and return the answer"""
+        try:
+            messages = [HumanMessage(content=question)]
+            result = self.graph.invoke({"messages": messages})
+            for m in result["messages"]:
+                m.pretty_print()
+            return result["messages"][-1].content
+        except Exception as e:
+            return f"Error: {str(e)}"
+
+if __name__ == "__main__":
+    agent = LangGraphAgent()
+    question = ".rewsna eht sa \"tfel\" drow eht fo etisoppo eht etirw ,ecnetnes siht dnatsrednu uoy fI"
+    answer = agent(question)
+
+
+