Create agent.py

agent.py

@@ -0,0 +1,61 @@
+import os
+from typing import Dict
+from langchain.agents import initialize_agent, AgentType
+from langchain.agents.tools import Tool
+from langchain_community.tools import DuckDuckGoSearchRun
+from langchain.tools import WikipediaQueryRun
+from langchain.utilities import WikipediaAPIWrapper
+from langchain_experimental.tools.python.tool import PythonREPLTool
+from langchain_google_genai import ChatGoogleGenerativeAI
+
+class Agent:
+    def __init__(self):
+        api_key = os.getenv("GEMINI_API_KEY")
+        if not api_key:
+            raise ValueError("GEMINI_API_KEY not found in environment variables.")
+
+        llm = ChatGoogleGenerativeAI(
+            model="gemini-1.5-pro",
+            google_api_key=api_key,
+            convert_system_message_to_human=True
+        )
+
+        search = DuckDuckGoSearchRun()
+        wiki = WikipediaQueryRun(api_wrapper=WikipediaAPIWrapper())
+        python_tool = PythonREPLTool()
+
+        tools = [
+            Tool(
+                name="DuckDuckGo Search",
+                func=search.run,
+                description="Useful for answering factual, current event, or open-domain questions."
+            ),
+            Tool(
+                name="Wikipedia",
+                func=wiki.run,
+                description="Useful for general knowledge and encyclopedic questions."
+            ),
+            Tool(
+                name="Calculator",
+                func=python_tool.run,
+                description="Useful for solving math and logical problems through Python."
+            )
+        ]
+
+        self.agent = initialize_agent(
+            tools=tools,
+            llm=llm,
+            agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
+            verbose=True,
+            handle_parsing_errors=True
+        )
+
+    def __call__(self, input_data: Dict) -> str:
+        question = input_data.get("question", "")
+        if not question:
+            return "No question provided."
+        try:
+            result = self.agent.run(question)
+            return result.strip()
+        except Exception as e:
+            return f"AGENT ERROR: {str(e)}"