File size: 837 Bytes
9023ac5 | 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 |
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
]
})
|