Spaces:
Sleeping
Sleeping
| from __future__ import annotations | |
| from collections import defaultdict | |
| from datetime import UTC, datetime | |
| from typing import Any | |
| from uuid import uuid4 | |
| from langchain_core.documents import Document | |
| from src.config import settings | |
| class LongTermMemoryStore: | |
| def __init__(self) -> None: | |
| self._fallback: dict[str, list[Document]] = defaultdict(list) | |
| self._init_warning: str | None = None | |
| self._vectorstore = None | |
| try: | |
| from langchain_chroma import Chroma | |
| self._vectorstore = Chroma( | |
| collection_name=settings.agent_memory_collection, | |
| persist_directory=settings.chroma_persist_dir, | |
| ) | |
| except Exception as exc: | |
| self._init_warning = ( | |
| f"Long-term vector memory fallback mode: {exc}" | |
| ) | |
| def init_warning(self) -> str | None: | |
| return self._init_warning | |
| def remember_fact( | |
| self, | |
| *, | |
| user_id: str, | |
| fact: str, | |
| memory_type: str = "general", | |
| source: str = "agent_memory", | |
| extra_metadata: dict[str, Any] | None = None, | |
| ) -> str: | |
| cleaned = fact.strip() | |
| if not cleaned: | |
| return "" | |
| memory_id = f"mem-{uuid4().hex}" | |
| metadata: dict[str, Any] = { | |
| "memory_id": memory_id, | |
| "user_id": user_id, | |
| "memory_type": memory_type, | |
| "created_at": datetime.now(UTC).isoformat(), | |
| "source": source, | |
| } | |
| if extra_metadata: | |
| metadata.update(extra_metadata) | |
| doc = Document(page_content=cleaned, metadata=metadata) | |
| self._fallback[user_id].append(doc) | |
| if self._vectorstore is not None: | |
| try: | |
| self._vectorstore.add_documents([doc], ids=[memory_id]) | |
| except Exception: | |
| # Keep service available even if vector persistence fails. | |
| pass | |
| return memory_id | |
| def recall_user_facts( | |
| self, | |
| *, | |
| user_id: str, | |
| query: str, | |
| limit: int = 5, | |
| ) -> list[Document]: | |
| safe_limit = max(1, min(limit, 20)) | |
| where = {"user_id": user_id} | |
| if self._vectorstore is not None and query.strip(): | |
| try: | |
| docs = self._vectorstore.similarity_search( | |
| query.strip(), | |
| k=safe_limit, | |
| filter=where, | |
| ) | |
| if docs: | |
| return docs | |
| except Exception: | |
| pass | |
| if self._vectorstore is not None: | |
| try: | |
| payload = self._vectorstore.get( | |
| where=where, | |
| limit=safe_limit, | |
| include=["documents", "metadatas"], | |
| ) | |
| documents = payload.get("documents", []) | |
| metadatas = payload.get("metadatas", []) | |
| docs: list[Document] = [] | |
| for text, metadata in zip(documents, metadatas): | |
| if text: | |
| docs.append( | |
| Document( | |
| page_content=text, | |
| metadata=dict(metadata or {}), | |
| ) | |
| ) | |
| if docs: | |
| return docs | |
| except Exception: | |
| pass | |
| return self._fallback.get(user_id, [])[-safe_limit:] | |