Update agent.py

agent.py

@@ -1,118 +1,95 @@
 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 duckduckgo_search import DDGS
-from langchain_community.document_loaders import WikipediaLoader
-from langchain_community.document_loaders import ArxivLoader
-from langchain_core.messages import SystemMessage, HumanMessage
-from langchain_core.tools import tool
+from langgraph.graph import StateGraph, START, MessagesState
+from langgraph.prebuilt import tools_condition, ToolNode
 from langchain_google_genai import ChatGoogleGenerativeAI
+from langchain_core.tools import tool
+from langchain_community.document_loaders import WikipediaLoader, ArxivLoader
+from langchain_community.utilities.duckduckgo_search import DuckDuckGoSearchAPIWrapper
+from langchain_core.messages import SystemMessage, HumanMessage
 
-load_dotenv()
-
+# Lade Umgebungsvariablen (Google API Key)
 GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY")
 
-# --- Tools ---
+# === Tools definieren ===
 @tool
 def multiply(a: int, b: int) -> int:
-    """Multiplies two integers and returns the result."""
+    """Multiplies two numbers."""
     return a * b
 
 @tool
 def add(a: int, b: int) -> int:
-    """Adds two integers and returns the result."""
+    """Adds two numbers."""
     return a + b
 
 @tool
 def subtract(a: int, b: int) -> int:
-    """Subtracts the second integer from the first."""
+    """Subtracts two numbers."""
     return a - b
 
 @tool
 def divide(a: int, b: int) -> float:
-    """Divides the first integer by the second, returns float."""
+    """Divides two numbers."""
     if b == 0:
         raise ValueError("Cannot divide by zero.")
     return a / b
 
 @tool
 def modulo(a: int, b: int) -> int:
-    """Returns the remainder of the division of two integers."""
+    """Returns the remainder of dividing two numbers."""
     return a % b
 
 @tool
 def wiki_search(query: str) -> str:
-    """Search Wikipedia for a given query and return up to 2 results formatted."""
+    """Search Wikipedia for a query and return the result."""
     search_docs = WikipediaLoader(query=query, load_max_docs=2).load()
-    formatted = "\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}
+    return "\n\n".join(doc.page_content for doc in search_docs)
 
 @tool
 def arxiv_search(query: str) -> str:
-    """Search Arxiv for a given query and return up to 3 results formatted."""
+    """Search Arxiv for academic papers about a query."""
     search_docs = ArxivLoader(query=query, load_max_docs=3).load()
-    formatted = "\n\n---\n\n".join(
-        [f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}">\n{doc.page_content[:1000]}\n</Document>' for doc in search_docs]
-    )
-    return {"arxiv_results": formatted}
+    return "\n\n".join(doc.page_content[:1000] for doc in search_docs)
 
 @tool
 def web_search(query: str) -> str:
-    """Search DuckDuckGo (for websearch) for a query and return up to 5 links."""
-    with DDGS() as ddgs:
-        results = ddgs.text(query, max_results=5)
-        if not results:
-            return "No results found."
-        return "\n\n".join(f"{r['title']}: {r['href']}" for r in results)
-
-# --- Setup LLM und Tools ---
-tools = [
-    multiply,
-    add,
-    subtract,
-    divide,
-    modulo,
-    wiki_search,
-    arxiv_search,
-    web_search,
-]
-
-system_prompt = (
-    "You are a highly accurate AI assistant. "
-    "Use tools when needed. Be very concise and precise. "
-    "Do not hallucinate information."
+    """Perform a DuckDuckGo web search."""
+    wrapper = DuckDuckGoSearchAPIWrapper(max_results=5)
+    results = wrapper.run(query)
+    return results
+
+# === System Prompt definieren ===
+system_prompt = SystemMessage(content=(
+    "You are an expert assistant. You MUST answer precisely, factually, and accurately. "
+    "If you do not know the answer, use the available tools such as Wikipedia Search, Arxiv Search, "
+    "or Web Search to find the correct information. "
+    "If a math operation is needed, use the calculation tools. "
+    "Do NOT invent answers. Only return answers you are confident in."
+))
+
+# === LLM definieren ===
+llm = ChatGoogleGenerativeAI(
+    model="gemini-2.0-flash",
+    google_api_key=GOOGLE_API_KEY,
+    temperature=0,
+    max_output_tokens=2048,
+    system_message=system_prompt,
 )
-sys_msg = SystemMessage(content=system_prompt)
-
-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):
-        return {"messages": [llm_with_tools.invoke(state["messages"])]}
+# === Tools in LLM einbinden ===
+tools = [multiply, add, subtract, divide, modulo, wiki_search, arxiv_search, web_search]
+llm_with_tools = llm.bind_tools(tools)
 
-    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")
+# === Nodes für LangGraph ===
+def assistant(state: MessagesState):
+    return {"messages": [llm_with_tools.invoke(state["messages"])]}
 
-    return builder.compile()
+# === LangGraph bauen ===
+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")
 
-# Agent Executor für app.py
-def agent_executor(question: str) -> str:
-    graph = build_graph()
-    messages = [HumanMessage(content=question)]
-    result = graph.invoke({"messages": messages})
-    return result["messages"][-1].content
+# === Agent Executor ===
+agent_executor = builder.compile()