Upload agent.py

agent.py

@@ -0,0 +1,109 @@
+import os
+from typing import TypedDict, Annotated, Sequence
+import operator
+from langgraph.graph import StateGraph, END
+from langgraph.prebuilt import ToolNode, tools_condition
+from langchain_google_genai import ChatGoogleGenerativeAI
+from langchain_community.tools import DuckDuckGoSearchRun
+from langchain_community.document_loaders import WikipediaLoader, ArxivLoader
+from langchain_core.messages import BaseMessage, HumanMessage, SystemMessage
+from langchain.agents import Tool
+from langchain_core.tools import tool
+
+# Load environment variables (compatible with Hugging Face Space)
+google_api_key = os.getenv("GOOGLE_API_KEY") or os.environ.get("GOOGLE_API_KEY")
+if not google_api_key:
+    raise ValueError("Missing GOOGLE_API_KEY environment variable")
+
+# --- System Prompt ---
+with open("System_Prompt.txt", "r", encoding="utf-8") as f:
+    system_prompt = f.read()
+sys_msg = SystemMessage(content=system_prompt)
+
+# --- Tool Definitions ---
+@tool
+def multiply(a: int, b: int) -> int:
+    """Multiply two integers together."""
+    return a * b
+
+@tool
+def add(a: int, b: int) -> int:
+    """Add two integers together."""
+    return a + b
+
+@tool
+def subtract(a: int, b: int) -> int:
+    """Subtract b from a."""
+    return a - b
+
+@tool
+def divide(a: int, b: int) -> float:
+    """Divide a by b. Returns float. Raises error if b is zero."""
+    if b == 0:
+        raise ValueError("Cannot divide by zero.")
+    return a / b
+
+@tool
+def wiki_search(query: str) -> str:
+    """Search Wikipedia and return up to 2 relevant documents."""
+    docs = WikipediaLoader(query=query, load_max_docs=2).load()
+    if not docs:
+        return "No Wikipedia results found."
+    return "\n\n".join([d.page_content[:1000] for d in docs])
+
+# Tool inventory with proper categorization
+tools = [
+    Tool(name="Math/Multiply", func=multiply, description="Multiplies two integers"),
+    Tool(name="Math/Add", func=add, description="Adds two integers"),
+    Tool(name="Math/Subtract", func=subtract, description="Subtracts two integers"),
+    Tool(name="Math/Divide", func=divide, description="Divides two numbers"),
+    Tool(name="Search/Wikipedia", func=wiki_search, description="Searches Wikipedia"),
+    Tool(
+        name="Search/Web", 
+        func=DuckDuckGoSearchRun().run,
+        description="Searches the web using DuckDuckGo"
+    )
+]
+
+# --- Graph Definition ---
+class AgentState(TypedDict):
+    """State definition for the agent workflow"""
+    messages: Annotated[Sequence[BaseMessage], operator.add]
+
+def build_graph():
+    """Constructs and compiles the LangGraph workflow"""
+    
+    # Initialize LLM with Gemini 2.0 Flash
+    llm = ChatGoogleGenerativeAI(
+        model="gemini-2.0-flash-exp",
+        temperature=0.3,
+        google_api_key=google_api_key
+    )
+    llm_with_tools = llm.bind_tools(tools)
+
+    # Node definitions
+    def agent_node(state: AgentState):
+        """Main agent node that processes messages"""
+        response = llm_with_tools.invoke(state["messages"])
+        return {"messages": [response]}
+
+    # Graph construction
+    workflow = StateGraph(AgentState)
+    
+    # Add nodes to the workflow
+    workflow.add_node("agent", agent_node)
+    workflow.add_node("tools", ToolNode(tools))
+    
+    # Configure graph flow
+    workflow.set_entry_point("agent")
+    workflow.add_conditional_edges(
+        "agent",
+        tools_condition,  # LangGraph's built-in tool detection
+        {"tools": "tools", "end": END}  # Fixed END reference
+    )
+    workflow.add_edge("tools", "agent")
+    
+    return workflow.compile()
+
+# Initialize the agent graph
+agent_graph = build_graph()