pratikshahp/agent-persistence
1
1# agent.py2import os3from typing import List, Dict, Optional4from openai import OpenAI5from strategy import StrategyFactory, ExecutionStrategy6from persistence import AgentPersistence7from datetime import datetime8 9class Agent:10 def __init__(self, name: str, persistence: Optional[AgentPersistence] = None):11 """12 Initialize an agent with a name and optional persistence manager.13 If no persistence manager is provided, a default one will be created.14 """15 self._name = name16 self._persona = ""17 self._instruction = ""18 self._task = ""19 self._api_key = os.getenv('OPENAI_API_KEY', '')20 self._model = "gpt-4o-mini"21 self._history: List[Dict[str, str]] = []22 self._strategy: Optional[ExecutionStrategy] = None23 self._persistence = persistence or AgentPersistence()24 25 # Try to load existing state26 self._persistence.load_agent_state(self)27 28 @property29 def name(self) -> str:30 """Get the agent's name."""31 return self._name32 33 @property34 def persona(self) -> str:35 """Get the agent's persona."""36 return self._persona37 38 @persona.setter39 def persona(self, value: str):40 """Set the agent's persona."""41 self._persona = value42 43 @property44 def instruction(self) -> str:45 """Get the agent's global instruction."""46 return self._instruction47 48 @instruction.setter49 def instruction(self, value: str):50 """Set the agent's global instruction."""51 self._instruction = value52 53 @property54 def task(self) -> str:55 """Get the current task."""56 return self._task57 58 @task.setter59 def task(self, value: str):60 """Set the current task."""61 self._task = value62 63 @property64 def strategy(self) -> Optional[ExecutionStrategy]:65 """Get the current execution strategy."""66 return self._strategy67 68 @strategy.setter69 def strategy(self, strategy_name: str):70 """Set the execution strategy by name."""71 self._strategy = StrategyFactory.create_strategy(strategy_name)72 73 @property74 def history(self) -> List[Dict[str, str]]:75 """Get the conversation history."""76 return self._history77 78 def get_history_states(self, limit: int = 10) -> List[Dict]:79 """80 Retrieve the last N states with their timestamps.81 """82 return self._persistence.get_agent_history(self.name, limit)83 84 def _build_messages(self, task: Optional[str] = None) -> List[Dict[str, str]]:85 """Build the messages list including persona, instruction, and history."""86 messages = [{"role": "system", "content": self.persona}]87 88 if self.instruction:89 messages.append({90 "role": "user", 91 "content": f"Global Instruction: {self.instruction}"92 })93 94 # Add conversation history95 messages.extend(self._history)96 97 # Use provided task or stored task98 current_task = task if task is not None else self._task99 100 # Apply strategy if set101 if self._strategy and current_task:102 current_task = self._strategy.build_prompt(current_task, self.instruction)103 104 # Add the current task if it exists105 if current_task:106 messages.append({"role": "user", "content": current_task})107 108 return messages109 110 def execute(self, task: Optional[str] = None) -> str:111 """Execute a task using the configured LLM."""112 if task is not None:113 self._task = task114 115 if not self._api_key:116 return "API key not found. Please set the OPENAI_API_KEY environment variable."117 118 if not self._task:119 return "No task specified. Please provide a task to execute."120 121 client = OpenAI(api_key=self._api_key)122 messages = self._build_messages()123 124 try:125 response = client.chat.completions.create(126 model=self._model,127 messages=messages128 )129 130 response_content = response.choices[0].message.content131 132 # Process response through strategy if set133 if self._strategy:134 response_content = self._strategy.process_response(response_content)135 136 # Store the interaction in history137 self._history.append({"role": "user", "content": self._task})138 self._history.append({139 "role": "assistant",140 "content": response_content141 })142 143 # Save state after successful execution144 self.save_state()145 146 # Clear the task after execution147 self._task = ""148 149 return response_content150 except Exception as e:151 return f"An error occurred: {str(e)}"152 153 def save_state(self) -> bool:154 """Save the current state of the agent."""155 return self._persistence.save_agent_state(self)156 157 def load_state(self, agent_name: Optional[str] = None) -> bool:158 """Load a saved state into the agent."""159 return self._persistence.load_agent_state(self, agent_name)160 161 def clear_history(self, keep_last: int = 0):162 """163 Clear the conversation history, optionally keeping the last N states.164 If keep_last > 0, it will clean up old states but retain the specified number.165 If keep_last = 0, it clears all history.166 """167 if keep_last > 0:168 self._persistence.cleanup_old_states(self.name, keep_last)169 # Reload the state to get the kept history170 self.load_state()171 else:172 self._history = []173 self.save_state()174 175 def pause(self) -> bool:176 """Pause the agent by saving its current state."""177 return self.save_state()178 179 def resume(self, agent_name: Optional[str] = None) -> bool:180 """Resume the agent by loading its saved state."""181 return self.load_state(agent_name)182 183 def available_strategies(self) -> List[str]:184 """Return a list of available strategy names."""185 return StrategyFactory.available_strategies()186 187 def delete_agent(self) -> bool:188 """Delete all data for this agent from the database."""189 return self._persistence.delete_agent_state(self.name)190 191 @staticmethod192 def list_saved_agents() -> Dict[str, datetime]:193 """194 List all saved agents and their last update times.195 Returns a dictionary of agent names mapped to their last update timestamps.196 """197 persistence = AgentPersistence()198 return persistence.list_saved_agents()