susmitsil commited on
Commit
533baaa
·
verified ·
1 Parent(s): c68b100

adding agent

Browse files
Files changed (1) hide show
  1. gemini_agent.py +69 -0
gemini_agent.py ADDED
@@ -0,0 +1,69 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from typing import List
2
+ from langchain_community.tools.ddg_search import DuckDuckGoSearchRun
3
+ from langchain.agents import AgentType, initialize_agent
4
+ from langchain.tools.base import BaseTool
5
+ from langchain_google_genai import ChatGoogleGenerativeAI
6
+ from langchain_core.messages import SystemMessage
7
+ from langchain.memory import ConversationBufferMemory
8
+
9
+ class GeminiAgent:
10
+ def __init__(self, api_key: str, model_name: str = "gemini-2.0-flash"):
11
+ # Suppress warnings
12
+ import warnings
13
+ warnings.filterwarnings("ignore", category=UserWarning)
14
+ warnings.filterwarnings("ignore", category=DeprecationWarning)
15
+ warnings.filterwarnings("ignore", message=".*will be deprecated.*")
16
+ warnings.filterwarnings("ignore", "LangChain.*")
17
+
18
+ self.api_key = api_key
19
+ self.model_name = model_name
20
+ self.agent = self._setup_agent()
21
+
22
+ def _setup_agent(self):
23
+ # Initialize model with system message
24
+ model = ChatGoogleGenerativeAI(
25
+ model=self.model_name,
26
+ google_api_key=self.api_key,
27
+ temperature=0, # Lower temperature for faster, more focused responses
28
+ max_output_tokens=200, # Limit response length
29
+ convert_system_message_to_human=True, # Faster processing of system message
30
+ stream=True, # Enable streaming for faster initial response
31
+ system_message=SystemMessage(content="You are a concise AI assistant. Provide a short and accurate answer. Preferable answer should be in one word or line. Unless if query asked expects an elaborate answer.")
32
+ )
33
+
34
+ # Setup tools
35
+ tools: List[BaseTool] = [DuckDuckGoSearchRun()]
36
+
37
+ # Setup memory
38
+ memory = ConversationBufferMemory(
39
+ memory_key="chat_history",
40
+ return_messages=True
41
+ )
42
+
43
+ # Create and return agent
44
+ return initialize_agent(
45
+ tools,
46
+ model,
47
+ agent=AgentType.CHAT_CONVERSATIONAL_REACT_DESCRIPTION,
48
+ memory=memory,
49
+ verbose=False,
50
+ handle_parsing_errors=True
51
+ )
52
+
53
+ def run(self, query: str) -> str:
54
+ try:
55
+ result = self.agent.invoke({"input": query})
56
+ return result["output"]
57
+ except Exception as e:
58
+ return f"Error: {e}"
59
+
60
+ def run_interactive(self):
61
+ print("AI Assistant Ready! (Type 'exit' to quit)")
62
+
63
+ while True:
64
+ query = input("You: ").strip()
65
+ if query.lower() == 'exit':
66
+ print("Goodbye!")
67
+ break
68
+
69
+ print("Assistant:", self.run(query))