Update agent.py

agent.py

@@ -1,93 +1,111 @@
 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.duckduckgo_search import DuckDuckGoSearchResults
+from langchain_community.document_loaders import WikipediaLoader, ArxivLoader
+from langchain_core.messages import SystemMessage, HumanMessage
+from langchain_core.tools import tool
 from langchain_google_genai import ChatGoogleGenerativeAI
-from langchain.agents import initialize_agent, Tool, AgentType
-from langchain_community.tools import DuckDuckGoSearchResults, WikipediaQueryRun
-from langchain_experimental.tools import PythonREPLTool
-from langchain.tools import tool
-from langchain.memory import ConversationBufferMemory
-from langchain_core.messages import SystemMessage
-import pandas as pd
-
-# API Key automatisch aus Environment ziehen
-google_api_key = os.getenv("GOOGLE_API_KEY")
-
-# LLM: Gemini 2.0 Flash
-llm = ChatGoogleGenerativeAI(
-    model="gemini-2.0-flash",
-    google_api_key=google_api_key,
-    temperature=0,
-    max_output_tokens=2048,
-    system_message=SystemMessage(content=(
-        "You are a highly accurate AI assistant. "
-        "You must answer precisely, concisely, and only if you are confident. "
-        "Use the available tools like Web Search, Wikipedia, Python REPL, or Table Analysis if needed. "
-        "Always prefer exact information over assumptions."
-    ))
-)
-
-# Tool 1: Web Search
-web_search = DuckDuckGoSearchResults()
-
-# Tool 2: Wikipedia Search
-wiki_search = WikipediaQueryRun()
-
-# Tool 3: Python REPL
-python_repl = PythonREPLTool()
-
-# Tool 4: Analyze CSV files (sehr einfaches Tool)
+
+load_dotenv()
+GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY")
+
+# --- Define Tools ---
+
+@tool 
+def multiply(a: int, b: int) -> int:
+    """Multiplies two numbers."""
+    return a * b
+
+@tool 
+def add(a: int, b: int) -> int:
+    """Adds two numbers."""
+    return a + b
+
+@tool 
+def subtract(a: int, b: int) -> int:
+    """Subtracts two numbers."""
+    return a - b
+
+@tool 
+def divide(a: int, b: int) -> float:
+    """Divides two numbers."""
+    if b == 0:
+        raise ValueError("Cannot divide by zero.")
+    return a / b
+
 @tool
-def analyze_csv(content: str) -> str:
-    """Analyzes CSV data and provides basic statistics and insights."""
-    try:
-        from io import StringIO
-        df = pd.read_csv(StringIO(content))
-        return str(df.describe())
-    except Exception as e:
-        return f"Failed to analyze CSV: {str(e)}"
-
-# Tool 5: Analyze Excel files
+def modulo(a: int, b: int) -> int:
+    """Returns the remainder of dividing two numbers."""
+    return a % b
+
 @tool
-def analyze_excel(content: bytes) -> str:
-    """Analyzes Excel data and provides basic statistics and insights."""
-    try:
-        from io import BytesIO
-        df = pd.read_excel(BytesIO(content))
-        return str(df.describe())
-    except Exception as e:
-        return f"Failed to analyze Excel: {str(e)}"
-
-# Alle Tools zusammen
+def wiki_search(query: str) -> str:
+    """Search Wikipedia for a query and return up to 2 results."""
+    search_docs = WikipediaLoader(query=query, load_max_docs=2).load()
+    formatted = "\n\n---\n\n".join(
+        [f'<Document source="{doc.metadata["source"]}">\n{doc.page_content}\n</Document>'
+         for doc in search_docs]
+    )
+    return formatted
+
+@tool
+def arxiv_search(query: str) -> str:
+    """Search Arxiv for scientific papers matching the query."""
+    search_docs = ArxivLoader(query=query, load_max_docs=3).load()
+    formatted = "\n\n---\n\n".join(
+        [f'<Document source="{doc.metadata["source"]}">\n{doc.page_content[:1000]}\n</Document>'
+         for doc in search_docs]
+    )
+    return formatted
+
+@tool
+def web_search(query: str) -> str:
+    """Search the web using DuckDuckGo."""
+    search = DuckDuckGoSearchResults()
+    return search.run(query)
+
+# --- Load System Prompt ---
+
+with open("system_prompt.txt", "r", encoding="utf-8") as f:
+    system_prompt = f.read()
+
+sys_msg = SystemMessage(content=system_prompt)
+
+# --- Define Tools List ---
 tools = [
-    Tool(
-        name="WebSearch",
-        func=web_search.run,
-        description="Use this to search the internet for up-to-date or unknown information."
-    ),
-    Tool(
-        name="WikipediaSearch",
-        func=wiki_search.run,
-        description="Use this to search Wikipedia articles when a direct lookup of factual information is needed."
-    ),
-    Tool(
-        name="Python_REPL",
-        func=python_repl.run,
-        description="Use this for math problems, small code executions, or calculations."
-    ),
-    analyze_csv,
-    analyze_excel,
+    multiply,
+    add,
+    subtract,
+    divide,
+    modulo,
+    wiki_search,
+    arxiv_search,
+    web_search,
 ]
 
-# Memory (optional, für Chat-History)
-memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)
-
-# Agent
-agent_executor = initialize_agent(
-    tools=tools,
-    llm=llm,
-    agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
-    verbose=True,
-    memory=memory,
-    handle_parsing_errors=True,
-)
+# --- Build Graph Function ---
 
+def build_graph():
+    llm = ChatGoogleGenerativeAI(
+        model="gemini-2.0-flash",
+        google_api_key=GOOGLE_API_KEY,
+        temperature=0,
+        max_output_tokens=2048,
+        system_message=sys_msg
+    )
+    llm_with_tools = llm.bind_tools(tools)
 
+    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")
+    return builder.compile()