76 lines
2.5 KiB
Python
76 lines
2.5 KiB
Python
from typing import Any, Dict, List, Optional
|
|
from langchain_core.messages import BaseMessage, HumanMessage
|
|
from langchain_openai import ChatOpenAI
|
|
from langgraph.graph import StateGraph, END
|
|
from config import Config
|
|
|
|
|
|
class AgentState:
|
|
"""State definition for the agent workflow"""
|
|
messages: List[BaseMessage]
|
|
current_step: str
|
|
context: Dict[str, Any]
|
|
|
|
def __init__(self, messages: List[BaseMessage] = None, current_step: str = "start", context: Dict[str, Any] = None):
|
|
self.messages = messages or []
|
|
self.current_step = current_step
|
|
self.context = context or {}
|
|
|
|
|
|
class BaseAgent:
|
|
"""Base agent class with common functionality"""
|
|
|
|
def __init__(self, model_name: str = Config.DEFAULT_MODEL):
|
|
self.model = ChatOpenAI(
|
|
model=model_name,
|
|
api_key=Config.OPENAI_API_KEY,
|
|
temperature=0.1,
|
|
max_retries=Config.MAX_RETRIES,
|
|
timeout=Config.TIMEOUT
|
|
)
|
|
self.graph = self._build_graph()
|
|
|
|
def _build_graph(self) -> StateGraph:
|
|
"""Build the state graph for the agent"""
|
|
workflow = StateGraph(AgentState)
|
|
|
|
# Add nodes and edges
|
|
workflow.add_node("process_input", self._process_input)
|
|
workflow.add_node("generate_response", self._generate_response)
|
|
|
|
# Define edges
|
|
workflow.add_edge("process_input", "generate_response")
|
|
workflow.add_edge("generate_response", END)
|
|
|
|
# Set entry point
|
|
workflow.set_entry_point("process_input")
|
|
|
|
return workflow.compile()
|
|
|
|
def _process_input(self, state: AgentState) -> AgentState:
|
|
"""Process user input"""
|
|
# This is a base implementation - subclasses should override
|
|
state.current_step = "processed"
|
|
return state
|
|
|
|
def _generate_response(self, state: AgentState) -> AgentState:
|
|
"""Generate response using the LLM"""
|
|
if state.messages:
|
|
response = self.model.invoke(state.messages)
|
|
state.messages.append(response)
|
|
return state
|
|
|
|
def run(self, user_input: str, **kwargs) -> Dict[str, Any]:
|
|
"""Run the agent with user input"""
|
|
initial_state = AgentState(
|
|
messages=[HumanMessage(content=user_input)],
|
|
context=kwargs
|
|
)
|
|
|
|
result = self.graph.invoke(initial_state)
|
|
|
|
return {
|
|
"messages": result.messages,
|
|
"context": result.context,
|
|
"final_step": result.current_step
|
|
} |