agent-cli 0.70.5__py3-none-any.whl

agent_cli/memory/_store.py

@@ -0,0 +1,169 @@
+"""ChromaDB helpers for memory storage."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from agent_cli.constants import DEFAULT_OPENAI_EMBEDDING_MODEL
+from agent_cli.core.chroma import delete as delete_docs
+from agent_cli.core.chroma import init_collection, upsert
+from agent_cli.memory._filters import to_chroma_where
+from agent_cli.memory.models import MemoryMetadata, StoredMemory
+
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+    from pathlib import Path
+
+    from chromadb import Collection
+
+
+def init_memory_collection(
+    persistence_path: Path,
+    *,
+    collection_name: str = "memory",
+    embedding_model: str = DEFAULT_OPENAI_EMBEDDING_MODEL,
+    openai_base_url: str | None = None,
+    openai_api_key: str | None = None,
+) -> Collection:
+    """Initialize or create the memory collection."""
+    return init_collection(
+        persistence_path,
+        name=collection_name,
+        embedding_model=embedding_model,
+        openai_base_url=openai_base_url,
+        openai_api_key=openai_api_key,
+        subdir="chroma",
+    )
+
+
+def upsert_memories(
+    collection: Collection,
+    ids: list[str],
+    contents: list[str],
+    metadatas: Sequence[MemoryMetadata],
+) -> None:
+    """Persist memory entries."""
+    upsert(collection, ids=ids, documents=contents, metadatas=metadatas)
+
+
+def query_memories(
+    collection: Collection,
+    *,
+    conversation_id: str,
+    text: str,
+    n_results: int,
+    filters: dict[str, Any] | None = None,
+) -> list[StoredMemory]:
+    """Query for relevant memory entries and return structured results."""
+    base_filters: list[dict[str, Any]] = [
+        {"conversation_id": conversation_id},
+        {"role": {"$ne": "summary"}},
+    ]
+    if filters:
+        chroma_filters = to_chroma_where(filters)
+        if chroma_filters:
+            base_filters.append(chroma_filters)
+    raw = collection.query(
+        query_texts=[text],
+        n_results=n_results,
+        where={"$and": base_filters},
+        include=["documents", "metadatas", "distances", "embeddings"],
+    )
+    docs_list = raw.get("documents")
+    docs = docs_list[0] if docs_list else []
+
+    metas_list = raw.get("metadatas")
+    metas = metas_list[0] if metas_list else []
+
+    ids_list = raw.get("ids")
+    ids = ids_list[0] if ids_list else []
+
+    dists_list = raw.get("distances")
+    distances = dists_list[0] if dists_list else []
+
+    raw_embeddings = raw.get("embeddings")
+    embeddings: list[Any] = []
+    if raw_embeddings and len(raw_embeddings) > 0 and raw_embeddings[0] is not None:
+        embeddings = raw_embeddings[0]
+
+    if len(embeddings) != len(docs):
+        msg = f"Chroma returned embeddings of unexpected length: {len(embeddings)} vs {len(docs)}"
+        raise ValueError(msg)
+    records: list[StoredMemory] = []
+    for doc, meta, doc_id, dist, emb in zip(
+        docs,
+        metas,
+        ids,
+        distances,
+        embeddings,
+        strict=False,
+    ):
+        assert doc_id is not None
+        records.append(
+            StoredMemory(
+                id=doc_id,
+                content=doc,
+                metadata=MemoryMetadata(**dict(meta)),
+                distance=float(dist) if dist is not None else None,
+                embedding=[float(x) for x in emb] if emb is not None else None,
+            ),
+        )
+    return records
+
+
+def get_summary_entry(
+    collection: Collection,
+    conversation_id: str,
+    *,
+    role: str = "summary",
+) -> StoredMemory | None:
+    """Return the latest summary entry for a conversation, if present."""
+    result = collection.get(
+        where={"$and": [{"conversation_id": conversation_id}, {"role": role}]},
+    )
+    docs = result.get("documents") or []
+    metas = result.get("metadatas") or []
+    ids = result.get("ids") or []
+
+    if not docs or not metas or not ids:
+        return None
+
+    return StoredMemory(
+        id=ids[0],
+        content=docs[0],
+        metadata=MemoryMetadata(**dict(metas[0])),
+        distance=None,
+    )
+
+
+def list_conversation_entries(
+    collection: Collection,
+    conversation_id: str,
+    *,
+    include_summary: bool = False,
+) -> list[StoredMemory]:
+    """List all entries for a conversation (optionally excluding summary)."""
+    filters: list[dict[str, Any]] = [{"conversation_id": conversation_id}]
+    if not include_summary:
+        filters.append({"role": {"$ne": "summary"}})
+    result = collection.get(where={"$and": filters} if len(filters) > 1 else filters[0])
+    docs = result.get("documents") or []
+    metas = result.get("metadatas") or []
+    ids = result.get("ids") or []
+
+    records: list[StoredMemory] = []
+    for doc, meta, entry_id in zip(docs, metas, ids, strict=False):
+        records.append(
+            StoredMemory(
+                id=entry_id,
+                content=doc,
+                metadata=MemoryMetadata(**dict(meta)),
+                distance=None,
+            ),
+        )
+    return records
+
+
+def delete_entries(collection: Collection, ids: list[str]) -> None:
+    """Delete entries by ID."""
+    delete_docs(collection, ids)

agent_cli/memory/_streaming.py

@@ -0,0 +1,44 @@
+"""Streaming helpers for chat completions."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from agent_cli.core.sse import extract_content_from_chunk, parse_chunk
+
+if TYPE_CHECKING:
+    from collections.abc import AsyncGenerator
+
+
+async def stream_chat_sse(
+    *,
+    openai_base_url: str,
+    payload: dict[str, Any],
+    headers: dict[str, str] | None = None,
+    request_timeout: float = 120.0,
+) -> AsyncGenerator[str, None]:
+    """Stream Server-Sent Events from an OpenAI-compatible chat completion endpoint."""
+    import httpx  # noqa: PLC0415
+
+    url = f"{openai_base_url.rstrip('/')}/chat/completions"
+    async with (
+        httpx.AsyncClient(timeout=request_timeout) as client,
+        client.stream("POST", url, json=payload, headers=headers) as response,
+    ):
+        if response.status_code != 200:  # noqa: PLR2004
+            error_text = await response.aread()
+            yield f"data: {error_text.decode(errors='ignore')}\n\n"
+            return
+        async for line in response.aiter_lines():
+            if line:
+                yield line
+
+
+def accumulate_assistant_text(line: str, buffer: list[str]) -> None:
+    """Parse SSE line and append any assistant text delta into buffer."""
+    chunk = parse_chunk(line)
+    if chunk is None:
+        return
+    piece = extract_content_from_chunk(chunk)
+    if piece:
+        buffer.append(piece)

agent_cli/memory/_tasks.py

@@ -0,0 +1,48 @@
+"""Utilities for tracking background tasks in the memory proxy."""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from collections.abc import Coroutine
+
+LOGGER = logging.getLogger(__name__)
+
+_BACKGROUND_TASKS: set[asyncio.Task[Any]] = set()
+
+
+def _track_background(task: asyncio.Task[Any], label: str) -> asyncio.Task[Any]:
+    """Track background tasks and surface failures."""
+    _BACKGROUND_TASKS.add(task)
+
+    def _done_callback(done: asyncio.Task[Any]) -> None:
+        _BACKGROUND_TASKS.discard(done)
+        if done.cancelled():
+            LOGGER.debug("Background task %s cancelled", label)
+            return
+        exc = done.exception()
+        if exc:
+            LOGGER.exception("Background task %s failed", label, exc_info=exc)
+
+    task.add_done_callback(_done_callback)
+    return task
+
+
+def run_in_background(
+    coro: asyncio.Task[Any] | Coroutine[Any, Any, Any],
+    label: str,
+) -> asyncio.Task[Any]:
+    """Create and track a background asyncio task."""
+    task = coro if isinstance(coro, asyncio.Task) else asyncio.create_task(coro)
+    task.set_name(f"memory-{label}")
+    return _track_background(task, label)
+
+
+async def wait_for_background_tasks() -> None:
+    """Await any in-flight background tasks (useful in tests)."""
+    while _BACKGROUND_TASKS:
+        tasks = list(_BACKGROUND_TASKS)
+        await asyncio.gather(*tasks, return_exceptions=False)

agent_cli/memory/api.py

@@ -0,0 +1,113 @@
+"""FastAPI application factory for memory proxy."""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, Any
+
+from fastapi import FastAPI, Request
+from fastapi.middleware.cors import CORSMiddleware
+
+from agent_cli.constants import DEFAULT_OPENAI_EMBEDDING_MODEL
+from agent_cli.core.openai_proxy import proxy_request_to_upstream
+from agent_cli.memory.client import MemoryClient
+from agent_cli.memory.models import ChatRequest  # noqa: TC001
+
+if TYPE_CHECKING:
+    from pathlib import Path
+
+LOGGER = logging.getLogger(__name__)
+
+
+def create_app(
+    memory_path: Path,
+    openai_base_url: str,
+    embedding_model: str = DEFAULT_OPENAI_EMBEDDING_MODEL,
+    embedding_api_key: str | None = None,
+    chat_api_key: str | None = None,
+    default_top_k: int = 5,
+    enable_summarization: bool = True,
+    max_entries: int = 500,
+    mmr_lambda: float = 0.7,
+    recency_weight: float = 0.2,
+    score_threshold: float | None = None,
+    enable_git_versioning: bool = True,
+) -> FastAPI:
+    """Create the FastAPI app for memory-backed chat."""
+    LOGGER.info("Initializing memory client...")
+
+    client = MemoryClient(
+        memory_path=memory_path,
+        openai_base_url=openai_base_url,
+        embedding_model=embedding_model,
+        embedding_api_key=embedding_api_key,
+        chat_api_key=chat_api_key,
+        default_top_k=default_top_k,
+        enable_summarization=enable_summarization,
+        max_entries=max_entries,
+        mmr_lambda=mmr_lambda,
+        recency_weight=recency_weight,
+        score_threshold=score_threshold,
+        start_watcher=False,  # We control start/stop via app events
+        enable_git_versioning=enable_git_versioning,
+    )
+
+    app = FastAPI(title="Memory Proxy")
+
+    app.add_middleware(
+        CORSMiddleware,
+        allow_origins=["*"],
+        allow_credentials=True,
+        allow_methods=["*"],
+        allow_headers=["*"],
+    )
+
+    @app.post("/v1/chat/completions")
+    async def chat_completions(request: Request, chat_request: ChatRequest) -> Any:
+        auth_header = request.headers.get("Authorization")
+        api_key = None
+        if auth_header and auth_header.startswith("Bearer "):
+            api_key = auth_header.split(" ")[1]
+
+        return await client.chat(
+            messages=chat_request.messages,
+            conversation_id=chat_request.memory_id or "default",
+            model=chat_request.model,
+            stream=chat_request.stream or False,
+            api_key=api_key,
+            memory_top_k=chat_request.memory_top_k,
+            recency_weight=chat_request.memory_recency_weight,
+            score_threshold=chat_request.memory_score_threshold,
+        )
+
+    @app.on_event("startup")
+    async def start_watch() -> None:
+        client.start()
+
+    @app.on_event("shutdown")
+    async def stop_watch() -> None:
+        await client.stop()
+
+    @app.get("/health")
+    def health() -> dict[str, str]:
+        return {
+            "status": "ok",
+            "memory_store": str(client.memory_path),
+            "openai_base_url": client.openai_base_url,
+            "default_top_k": str(client.default_top_k),
+        }
+
+    @app.api_route(
+        "/{path:path}",
+        methods=["GET", "POST", "PUT", "DELETE", "OPTIONS", "HEAD", "PATCH"],
+    )
+    async def proxy_catch_all(request: Request, path: str) -> Any:
+        """Forward any other request to the upstream provider."""
+        return await proxy_request_to_upstream(
+            request,
+            path,
+            client.openai_base_url,
+            client.chat_api_key,
+        )
+
+    return app

agent_cli/memory/client.py

@@ -0,0 +1,272 @@
+"""High-level client for interacting with the memory system."""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+from contextlib import suppress
+from typing import TYPE_CHECKING, Any, Self
+
+from agent_cli.constants import DEFAULT_OPENAI_EMBEDDING_MODEL, DEFAULT_OPENAI_MODEL
+from agent_cli.core.reranker import get_reranker_model
+from agent_cli.memory._files import ensure_store_dirs
+from agent_cli.memory._git import init_repo
+from agent_cli.memory._indexer import MemoryIndex, initial_index, watch_memory_store
+from agent_cli.memory._ingest import extract_and_store_facts_and_summaries
+from agent_cli.memory._persistence import evict_if_needed
+from agent_cli.memory._retrieval import augment_chat_request
+from agent_cli.memory._store import init_memory_collection, list_conversation_entries
+from agent_cli.memory.engine import process_chat_request
+from agent_cli.memory.models import ChatRequest, MemoryRetrieval, Message
+
+if TYPE_CHECKING:
+    from pathlib import Path
+
+    from chromadb import Collection
+
+    from agent_cli.core.reranker import OnnxCrossEncoder
+
+
+logger = logging.getLogger("agent_cli.memory.client")
+
+
+class MemoryClient:
+    """A client for interacting with the memory system (add, search, chat).
+
+    This class decouples the memory logic from the HTTP server, allowing
+    direct usage in other applications or scripts.
+    """
+
+    def __init__(
+        self,
+        memory_path: Path,
+        openai_base_url: str,
+        embedding_model: str = DEFAULT_OPENAI_EMBEDDING_MODEL,
+        embedding_api_key: str | None = None,
+        chat_api_key: str | None = None,
+        enable_summarization: bool = True,
+        default_top_k: int = 5,
+        max_entries: int = 500,
+        mmr_lambda: float = 0.7,
+        recency_weight: float = 0.2,
+        score_threshold: float | None = None,
+        start_watcher: bool = False,
+        enable_git_versioning: bool = True,
+    ) -> None:
+        """Initialize the memory client."""
+        self.memory_path = memory_path.resolve()
+        self.openai_base_url = openai_base_url.rstrip("/")
+        self.chat_api_key = chat_api_key
+        self.enable_summarization = enable_summarization
+        self.default_top_k = default_top_k
+        self.max_entries = max_entries
+        self.mmr_lambda = mmr_lambda
+        self.recency_weight = recency_weight
+        self.score_threshold = score_threshold
+        self.enable_git_versioning = enable_git_versioning
+
+        _, snapshot_path = ensure_store_dirs(self.memory_path)
+
+        if self.enable_git_versioning:
+            init_repo(self.memory_path)
+
+        logger.info("Initializing memory collection...")
+        self.collection: Collection = init_memory_collection(
+            self.memory_path,
+            embedding_model=embedding_model,
+            openai_base_url=self.openai_base_url,
+            openai_api_key=embedding_api_key,
+        )
+
+        self.index = MemoryIndex.from_snapshot(snapshot_path)
+        initial_index(self.collection, self.memory_path, index=self.index)
+
+        logger.info("Loading reranker model...")
+        self.reranker_model: OnnxCrossEncoder = get_reranker_model()
+
+        self._watch_task: asyncio.Task | None = None
+        if start_watcher:
+            self.start()
+
+    def start(self) -> None:
+        """Start the background file watcher."""
+        if self._watch_task is None:
+            self._watch_task = asyncio.create_task(
+                watch_memory_store(self.collection, self.memory_path, index=self.index),
+            )
+
+    async def stop(self) -> None:
+        """Stop the background file watcher."""
+        if self._watch_task:
+            self._watch_task.cancel()
+            with suppress(asyncio.CancelledError):
+                await self._watch_task
+            self._watch_task = None
+
+    async def __aenter__(self) -> Self:
+        """Start the client context."""
+        self.start()
+        return self
+
+    async def __aexit__(self, *args: object) -> None:
+        """Stop the client context."""
+        await self.stop()
+
+    async def add(
+        self,
+        text: str,
+        conversation_id: str = "default",
+        model: str = DEFAULT_OPENAI_MODEL,
+    ) -> None:
+        """Add a memory by extracting facts from text and reconciling them.
+
+        This mimics the 'mem0.add' behavior but uses our advanced reconciliation
+        pipeline (Add/Update/Delete) and updates the conversation summary.
+        """
+        await extract_and_store_facts_and_summaries(
+            collection=self.collection,
+            memory_root=self.memory_path,
+            conversation_id=conversation_id,
+            user_message=text,
+            assistant_message=None,
+            openai_base_url=self.openai_base_url,
+            api_key=self.chat_api_key,
+            model=model,
+            enable_git_versioning=self.enable_git_versioning,
+            enable_summarization=self.enable_summarization,
+        )
+        evict_if_needed(self.collection, self.memory_path, conversation_id, self.max_entries)
+
+    async def search(
+        self,
+        query: str,
+        conversation_id: str = "default",
+        top_k: int | None = None,
+        model: str = DEFAULT_OPENAI_MODEL,
+        recency_weight: float | None = None,
+        score_threshold: float | None = None,
+        filters: dict[str, Any] | None = None,
+    ) -> MemoryRetrieval:
+        """Search for memories relevant to a query.
+
+        Args:
+            query: The search query text.
+            conversation_id: Conversation scope for the search.
+            top_k: Number of results to return.
+            model: Model for any LLM operations.
+            recency_weight: Weight for recency scoring (0-1).
+            score_threshold: Minimum relevance score threshold.
+            filters: Optional metadata filters. Examples:
+                - {"role": "user"} - exact match
+                - {"created_at": {"gte": "2024-01-01"}} - comparison
+                - {"$or": [{"role": "user"}, {"role": "assistant"}]} - logical OR
+                Operators: eq, ne, gt, gte, lt, lte, in, nin
+
+        """
+        dummy_request = ChatRequest(
+            messages=[Message(role="user", content=query)],
+            model=model,
+            memory_id=conversation_id,
+            memory_top_k=top_k or self.default_top_k,
+        )
+
+        _, retrieval, _, _ = await augment_chat_request(
+            dummy_request,
+            self.collection,
+            reranker_model=self.reranker_model,
+            default_top_k=top_k or self.default_top_k,
+            include_global=True,
+            mmr_lambda=self.mmr_lambda,
+            recency_weight=recency_weight if recency_weight is not None else self.recency_weight,
+            score_threshold=score_threshold
+            if score_threshold is not None
+            else self.score_threshold,
+            filters=filters,
+        )
+        return retrieval or MemoryRetrieval(entries=[])
+
+    def list_all(
+        self,
+        conversation_id: str = "default",
+        include_summary: bool = False,
+    ) -> list[dict[str, Any]]:
+        """List all stored memories for a conversation.
+
+        Args:
+            conversation_id: Conversation scope.
+            include_summary: Whether to include summary entries.
+
+        Returns:
+            List of memory entries with id, content, and metadata.
+
+        """
+        entries = list_conversation_entries(
+            self.collection,
+            conversation_id,
+            include_summary=include_summary,
+        )
+        return [
+            {
+                "id": e.id,
+                "content": e.content,
+                "role": e.metadata.role,
+                "created_at": e.metadata.created_at,
+            }
+            for e in entries
+        ]
+
+    async def chat(
+        self,
+        messages: list[dict[str, str]] | list[Any],
+        conversation_id: str = "default",
+        model: str = DEFAULT_OPENAI_MODEL,
+        stream: bool = False,
+        api_key: str | None = None,
+        memory_top_k: int | None = None,
+        recency_weight: float | None = None,
+        score_threshold: float | None = None,
+        filters: dict[str, Any] | None = None,
+    ) -> Any:
+        """Process a chat request (Augment -> LLM -> Update Memory).
+
+        Args:
+            messages: Chat messages.
+            conversation_id: Conversation scope.
+            model: LLM model to use.
+            stream: Whether to stream the response.
+            api_key: Optional API key override.
+            memory_top_k: Number of memories to retrieve.
+            recency_weight: Weight for recency scoring (0-1).
+            score_threshold: Minimum relevance score threshold.
+            filters: Optional metadata filters for memory retrieval.
+
+        """
+        req = ChatRequest(
+            messages=messages,  # type: ignore[arg-type]
+            model=model,
+            memory_id=conversation_id,
+            stream=stream,
+            memory_top_k=memory_top_k if memory_top_k is not None else self.default_top_k,
+            memory_recency_weight=recency_weight,
+            memory_score_threshold=score_threshold,
+        )
+
+        return await process_chat_request(
+            req,
+            collection=self.collection,
+            memory_root=self.memory_path,
+            openai_base_url=self.openai_base_url,
+            reranker_model=self.reranker_model,
+            default_top_k=self.default_top_k,
+            api_key=api_key or self.chat_api_key,
+            enable_summarization=self.enable_summarization,
+            max_entries=self.max_entries,
+            mmr_lambda=self.mmr_lambda,
+            recency_weight=recency_weight if recency_weight is not None else self.recency_weight,
+            score_threshold=score_threshold
+            if score_threshold is not None
+            else self.score_threshold,
+            postprocess_in_background=True,
+            enable_git_versioning=self.enable_git_versioning,
+            filters=filters,
+        )