|
|
| from src.setup.cache import RedisDataSource |
| from src.agent.custom_agent import LearningAgent |
|
|
|
|
| class AgentManager: |
| def __init__(self): |
| self.cache = RedisDataSource() |
|
|
| def get_agent(self, session_id, llm_type): |
| cached = self.cache.read(f"agent:{session_id}") |
| agent = LearningAgent(llm_type) |
| if cached and "chat_history" in cached: |
| agent.chat_history = [ |
| msg for msg in cached["chat_history"] |
| if isinstance(msg, dict) and "role" in msg and "content" in msg |
| ] |
| return agent |
|
|
| def save_agent(self, session_id, agent): |
| self.cache.write(f"agent:{session_id}", { |
| "chat_history": [ |
| {"role": msg["role"], "content": msg["content"]} |
| for msg in agent.chat_history |
| ] |
| }) |
|
|