94 lines
2.7 KiB
Python
94 lines
2.7 KiB
Python
import asyncio
|
|
from typing import Dict, Optional
|
|
from agent.memory import MemoryManager
|
|
|
|
|
|
class SessionManager:
|
|
"""Manages agent sessions and their associated memory.
|
|
|
|
Coordinates between the gateway's UserSession (for conversation state)
|
|
and the agent's MemoryManager (for persistent memory with Zep).
|
|
"""
|
|
|
|
def __init__(self, memory_manager: MemoryManager):
|
|
"""Initialize session manager.
|
|
|
|
Args:
|
|
memory_manager: MemoryManager instance for persistent storage
|
|
"""
|
|
self.memory = memory_manager
|
|
self._locks: Dict[str, asyncio.Lock] = {}
|
|
|
|
async def get_session_lock(self, session_id: str) -> asyncio.Lock:
|
|
"""Get or create a lock for a session.
|
|
|
|
Prevents concurrent execution for the same session.
|
|
|
|
Args:
|
|
session_id: Session identifier
|
|
|
|
Returns:
|
|
asyncio.Lock for this session
|
|
"""
|
|
if session_id not in self._locks:
|
|
self._locks[session_id] = asyncio.Lock()
|
|
return self._locks[session_id]
|
|
|
|
async def save_message(
|
|
self,
|
|
session_id: str,
|
|
role: str,
|
|
content: str,
|
|
metadata: Optional[Dict] = None
|
|
) -> None:
|
|
"""Save a message to persistent memory.
|
|
|
|
Args:
|
|
session_id: Session identifier
|
|
role: Message role ("user" or "assistant")
|
|
content: Message content
|
|
metadata: Optional metadata
|
|
"""
|
|
await self.memory.add_memory(session_id, role, content, metadata)
|
|
|
|
async def get_relevant_context(
|
|
self,
|
|
session_id: str,
|
|
query: str,
|
|
limit: int = 5
|
|
) -> str:
|
|
"""Get relevant historical context for a query.
|
|
|
|
Uses semantic search over past conversations.
|
|
|
|
Args:
|
|
session_id: Session identifier
|
|
query: Search query
|
|
limit: Maximum results
|
|
|
|
Returns:
|
|
Formatted string of relevant past messages
|
|
"""
|
|
results = await self.memory.search_memory(session_id, query, limit)
|
|
|
|
if not results:
|
|
return ""
|
|
|
|
context_parts = ["## Relevant Past Context\n"]
|
|
for i, result in enumerate(results, 1):
|
|
role = result["role"].capitalize()
|
|
content = result["content"]
|
|
score = result.get("score", 0.0)
|
|
context_parts.append(f"{i}. [{role}, relevance: {score:.2f}] {content}\n")
|
|
|
|
return "\n".join(context_parts)
|
|
|
|
async def cleanup_session(self, session_id: str) -> None:
|
|
"""Cleanup session resources.
|
|
|
|
Args:
|
|
session_id: Session to cleanup
|
|
"""
|
|
if session_id in self._locks:
|
|
del self._locks[session_id]
|