Tribh/devops-copilot

1import chromadb2from chromadb.config import Settings3from typing import List, Dict, Any, Optional4import os5from devops_copilot.utils.logger import logger6 7class MemorySystem:8    """Persistent memory using ChromaDB."""9    10    def __init__(self, persist_directory: str = "./chroma_db"):11        self.client = chromadb.PersistentClient(path=persist_directory)12        self.collection = self.client.get_or_create_collection(name="agent_memory")13        logger.info(f"Memory system initialized at {persist_directory}")14 15    def add_memories(self, documents: List[str], metadatas: List[Dict[str, Any]], ids: List[str]):16        """Adds documents to the vector store."""17        self.collection.add(18            documents=documents,19            metadatas=metadatas,20            ids=ids21        )22        logger.info(f"Added {len(documents)} memories to vector store.")23 24    def search_memories(self, query: str, n_results: int = 5) -> Dict[str, Any]:25        """Searches for relevant memories."""26        results = self.collection.query(27            query_texts=[query],28            n_results=n_results29        )30        return results31 32class SessionManager:33    """Handles ephemeral session state."""34    35    def __init__(self):36        self._sessions: Dict[str, Dict[str, Any]] = {}37 38    def get_session(self, session_id: str) -> Dict[str, Any]:39        if session_id not in self._sessions:40            self._sessions[session_id] = {"history": [], "context": {}}41        return self._sessions[session_id]42 43    def update_session(self, session_id: str, key: str, value: Any):44        session = self.get_session(session_id)45        session[key] = value46 47    def clear_session(self, session_id: str):48        if session_id in self._sessions:49            del self._sessions[session_id]50