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