final-agent-course / agent.py
jjvelezo's picture
Update agent.py
427cae4 verified
raw
history blame
2.33 kB
import time
from langchain.memory import ConversationBufferMemory
from langchain.agents import AgentExecutor, Tool
from langchain.chat_models import ChatGoogleGenerativeAI
import genai
from gemini import GeminiAgent
from smolagent_tools import WikipediaSearchTool, WebSearchTool, SmolagentToolWrapper
# Función para configuración de llm
def setup_llm(api_key: str, model_name: str = "gemini-2.0-flash"):
"""Configuración del modelo de lenguaje con Gemini."""
genai.configure(api_key=api_key)
return ChatGoogleGenerativeAI(model_name=model_name)
# Herramientas adicionales
class EnhancedAgent:
def __init__(self, api_key: str, model_name: str = "gemini-2.0-flash"):
self.api_key = api_key
self.model_name = model_name
self.llm = setup_llm(api_key, model_name)
self.memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)
self.tools = [
SmolagentToolWrapper(WikipediaSearchTool()),
Tool(name="web_search", func=self._web_search, description="Realizar búsqueda web avanzada")
]
self.agent = GeminiAgent(
tools=self.tools,
llm=self.llm,
memory=self.memory,
)
def _web_search(self, query: str, domain: Optional[str] = None) -> str:
"""Realiza una búsqueda en la web usando DuckDuckGo con limitación de tasa y reintentos."""
try:
search_tool = WebSearchTool()
results = search_tool.search(query, domain)
return results
except Exception as e:
return f"Error al realizar la búsqueda: {str(e)}"
def run(self, query: str) -> str:
"""Procesa las consultas del usuario con reintentos y manejo de errores."""
max_retries = 3
base_sleep = 1
for attempt in range(max_retries):
try:
response = self.agent.run(query)
return response
except Exception as e:
sleep_time = base_sleep * (attempt + 1)
if attempt < max_retries - 1:
print(f"Intento {attempt + 1} fallido. Reintentando en {sleep_time} segundos...")
time.sleep(sleep_time)
continue
return f"Error procesando la consulta: {str(e)}"