agent.py

import os
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)
@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
@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
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,
]

# 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,
)