devmem-agents 0.1.0__py3-none-any.whl

devmem/service.py

@@ -0,0 +1,280 @@
+"""Core service logic for devmem.
+
+Wires the API layer to a `RecordStore` backend (in-memory or SQL) and
+embedding provider. Session task descriptions are embedded at session start
+to power `/v1/tasks/similar`.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from devmem.config import settings
+from devmem.embeddings import embed_one, is_fallback_active
+from devmem.models import SessionStartRequest, SessionStartResponse
+from devmem.storage.record_store import RecordStore, build_record_store
+
+
+class NamespaceError(ValueError):
+    """Raised when request namespace is invalid."""
+
+
+class DevMemService:
+    """Contract-first service layer with strict namespace guard."""
+
+    def __init__(self, store: RecordStore | None = None) -> None:
+        self.store: RecordStore = store or build_record_store(settings.record_store_dsn)
+
+    def _require_namespace(self, namespace: str) -> None:
+        if namespace != settings.namespace:
+            raise NamespaceError(
+                f"invalid namespace '{namespace}'. expected '{settings.namespace}'"
+            )
+
+    def health(self) -> dict[str, Any]:
+        return {
+            "service": settings.service_name,
+            "version": settings.service_version,
+            "namespace": settings.namespace,
+            "checks": {
+                "api": "healthy",
+                "milvus": "not_configured" if not settings.milvus_uri else "configured",
+                "neptune": "not_configured" if not settings.neptune_endpoint else "configured",
+                "aurora": "not_configured" if not settings.aurora_dsn else "configured",
+                "record_store": "sql" if settings.record_store_dsn else "memory",
+                "embeddings": "fallback_sha256" if is_fallback_active() else settings.embedding_model,
+            },
+        }
+
+    # ---- sessions --------------------------------------------------------
+
+    def start_session(self, req: SessionStartRequest) -> SessionStartResponse:
+        self._require_namespace(req.namespace)
+        session = SessionStartResponse()
+        task_embedding = embed_one(req.task) if req.task else None
+        self.store.create_session(
+            session_id=session.session_id,
+            namespace=req.namespace,
+            project=req.project,
+            repo=req.repo,
+            branch=req.branch,
+            agent=req.agent,
+            task=req.task,
+            started_at=session.started_at,
+            task_embedding=task_embedding,
+        )
+        return session
+
+    # ---- retrieval -------------------------------------------------------
+
+    def pull_context(
+        self,
+        *,
+        namespace: str,
+        session_id: str,
+        project: str,
+        repo: str,
+        task: str,
+        top_k: int,
+    ) -> dict[str, Any]:
+        self._require_namespace(namespace)
+        artifacts = self.store.query_artifacts(
+            namespace=namespace, project=project, repo=repo, limit=top_k
+        )
+        decisions = self.store.query_decisions(
+            namespace=namespace, project=project, repo=repo, limit=top_k
+        )
+        return {
+            "session_id": session_id,
+            "task": task,
+            "artifacts": artifacts,
+            "decisions": decisions,
+        }
+
+    def hybrid_search(
+        self,
+        *,
+        namespace: str,
+        q: str,
+        project: str,
+        repo: str,
+        top_k: int,
+    ) -> list[dict[str, Any]]:
+        """Semantic search over artifact+decision content, ranked by cosine similarity."""
+        self._require_namespace(namespace)
+        if not (q or "").strip():
+            return []
+        query_vec = embed_one(q)
+
+        # Pool all artifacts and decisions for this project/repo, rank by
+        # cosine similarity between query embedding and a per-record embedding
+        # computed on the fly from title + content.
+        artifacts = self.store.query_artifacts(
+            namespace=namespace, project=project, repo=repo, limit=500
+        )
+        decisions = self.store.query_decisions(
+            namespace=namespace, project=project, repo=repo, limit=500
+        )
+        candidates: list[tuple[str, dict[str, Any]]] = []
+        for a in artifacts:
+            text = f"{a.get('title', '')}\n{a.get('content', '')}"
+            candidates.append((text, {"kind": "artifact", **a}))
+        for d in decisions:
+            text = f"{d.get('title', '')}\n{d.get('decision', '')}\n{d.get('rationale', '')}"
+            candidates.append((text, {"kind": "decision", **d}))
+        if not candidates:
+            return []
+
+        from devmem.embeddings import embed_many, cosine_similarity
+
+        texts = [c[0] for c in candidates]
+        vectors = embed_many(texts)
+        scored: list[tuple[float, dict[str, Any]]] = []
+        for (_, record), vec in zip(candidates, vectors):
+            score = cosine_similarity(query_vec, vec)
+            if score > 0:
+                scored.append((float(score), record))
+        scored.sort(key=lambda x: x[0], reverse=True)
+        return [{"score": score, **rec} for score, rec in scored[:top_k]]
+
+    def similar_tasks(
+        self,
+        *,
+        namespace: str,
+        q: str,
+        project: str | None,
+        repo: str | None,
+        top_k: int,
+    ) -> list[dict[str, Any]]:
+        self._require_namespace(namespace)
+        if not (q or "").strip():
+            return []
+        query_vec = embed_one(q)
+        return self.store.similar_tasks(
+            namespace=namespace,
+            task_embedding=query_vec,
+            project=project,
+            repo=repo,
+            top_k=top_k,
+        )
+
+    # ---- writes ----------------------------------------------------------
+
+    def upsert_artifact(self, namespace: str, payload: dict[str, Any]) -> dict[str, Any]:
+        self._require_namespace(namespace)
+        return self.store.insert_artifact(namespace=namespace, payload=payload)
+
+    def upsert_fact(self, namespace: str, payload: dict[str, Any]) -> dict[str, Any]:
+        self._require_namespace(namespace)
+        return self.store.insert_fact(namespace=namespace, payload=payload)
+
+    def upsert_decision(self, namespace: str, payload: dict[str, Any]) -> dict[str, Any]:
+        self._require_namespace(namespace)
+        return self.store.insert_decision(namespace=namespace, payload=payload)
+
+    def create_handoff(self, namespace: str, payload: dict[str, Any]) -> dict[str, Any]:
+        self._require_namespace(namespace)
+        return self.store.insert_handoff(namespace=namespace, payload=payload)
+
+    def update_task(self, namespace: str, payload: dict[str, Any]) -> dict[str, Any]:
+        self._require_namespace(namespace)
+        return self.store.insert_task_update(namespace=namespace, payload=payload)
+
+    def record_feedback(self, namespace: str, payload: dict[str, Any]) -> dict[str, Any]:
+        self._require_namespace(namespace)
+        return self.store.insert_feedback(namespace=namespace, payload=payload)
+
+    def commit_session(
+        self,
+        *,
+        namespace: str,
+        session_id: str,
+        artifacts: list[dict[str, Any]],
+        decisions: list[dict[str, Any]],
+        handoff: dict[str, Any] | None,
+        task_update: dict[str, Any] | None,
+        client_commit_id: str | None,
+    ) -> dict[str, Any]:
+        self._require_namespace(namespace)
+        return self.store.commit_session(
+            namespace=namespace,
+            session_id=session_id,
+            artifacts=artifacts,
+            decisions=decisions,
+            handoff=handoff,
+            task_update=task_update,
+            client_commit_id=client_commit_id,
+        )
+
+    # ---- listings --------------------------------------------------------
+
+    def list_artifacts(
+        self,
+        *,
+        namespace: str,
+        project: str | None,
+        repo: str | None,
+        session_id: str | None,
+        artifact_type: str | None,
+        since: str | None,
+        limit: int,
+        offset: int,
+    ) -> list[dict[str, Any]]:
+        self._require_namespace(namespace)
+        return self.store.query_artifacts(
+            namespace=namespace,
+            project=project,
+            repo=repo,
+            session_id=session_id,
+            artifact_type=artifact_type,
+            since=since,
+            limit=limit,
+            offset=offset,
+        )
+
+    def list_decisions(
+        self,
+        *,
+        namespace: str,
+        project: str | None,
+        repo: str | None,
+        session_id: str | None,
+        since: str | None,
+        limit: int,
+        offset: int,
+    ) -> list[dict[str, Any]]:
+        self._require_namespace(namespace)
+        return self.store.query_decisions(
+            namespace=namespace,
+            project=project,
+            repo=repo,
+            session_id=session_id,
+            since=since,
+            limit=limit,
+            offset=offset,
+        )
+
+    def list_handoffs(
+        self,
+        *,
+        namespace: str,
+        project: str | None,
+        repo: str | None,
+        session_id: str | None,
+        since: str | None,
+        limit: int,
+        offset: int,
+    ) -> list[dict[str, Any]]:
+        self._require_namespace(namespace)
+        return self.store.query_handoffs(
+            namespace=namespace,
+            project=project,
+            repo=repo,
+            session_id=session_id,
+            since=since,
+            limit=limit,
+            offset=offset,
+        )
+
+
+service = DevMemService()

devmem/storage/__init__.py

@@ -0,0 +1,4 @@
+"""Storage adapters package.
+
+Real adapters for Neptune, Milvus, and Aurora will be implemented here.
+"""

devmem/storage/milvus_store.py

@@ -0,0 +1,321 @@
+"""Milvus adapter for devmem repository ingestion."""
+
+from __future__ import annotations
+
+import hashlib
+import logging
+import re
+from datetime import datetime, timezone
+from typing import Any, Iterable
+
+import numpy as np
+from pymilvus import (
+    Collection,
+    CollectionSchema,
+    DataType,
+    FieldSchema,
+    connections,
+    utility,
+)
+
+from devmem.embeddings import embed_many, embed_one, embedding_dim
+from devmem.live_backend import LiveBackendConfig, parse_host_port_from_url
+
+logger = logging.getLogger(__name__)
+_TOKEN_RE = re.compile(r"[a-zA-Z0-9_]{2,}")
+
+
+class MilvusStore:
+    """Write repository text chunks to Milvus with strict namespace metadata."""
+
+    def __init__(
+        self,
+        cfg: LiveBackendConfig,
+        *,
+        namespace: str,
+        project_id: str,
+        repo_id: str,
+        dim: int | None = None,
+    ) -> None:
+        self.cfg = cfg
+        self.namespace = namespace
+        self.project_id = project_id
+        self.repo_id = repo_id
+        self.dim = dim if dim is not None else embedding_dim()
+        self.alias = f"devmem_{namespace}_{project_id}"
+        self.collection_name = f"{namespace}_{project_id}_{repo_id}_chunks".replace("-", "_")[:200]
+        self.collection: Collection | None = None
+
+    def connect(self) -> None:
+        """Connect to Milvus and ensure target collection exists."""
+        args: dict[str, object] = {"timeout": self.cfg.milvus_timeout}
+        if self.cfg.milvus_uri:
+            args["uri"] = self.cfg.milvus_uri
+        elif self.cfg.milvus_host:
+            args["host"] = self.cfg.milvus_host
+            args["port"] = self.cfg.milvus_port
+            args["secure"] = self.cfg.milvus_secure
+        else:
+            raise ValueError("Milvus host/URI is not configured")
+
+        if self.cfg.milvus_token:
+            args["token"] = self.cfg.milvus_token
+        elif self.cfg.milvus_user:
+            args["user"] = self.cfg.milvus_user
+            args["password"] = self.cfg.milvus_password or ""
+
+        connections.connect(alias=self.alias, **args)
+        self.collection = self._ensure_collection()
+        try:
+            self.collection.load()
+        except Exception:
+            logger.debug(
+                "Milvus collection load skipped collection=%s",
+                self.collection_name,
+                exc_info=True,
+            )
+        logger.info(
+            "Milvus connected alias=%s collection=%s",
+            self.alias,
+            self.collection_name,
+        )
+
+    def close(self) -> None:
+        """Disconnect from Milvus."""
+        try:
+            connections.disconnect(self.alias)
+        except Exception:
+            logger.debug("Milvus disconnect failed alias=%s", self.alias, exc_info=True)
+
+    def _ensure_collection(self) -> Collection:
+        if not utility.has_collection(self.collection_name, using=self.alias):
+            fields = [
+                FieldSchema(name="id", dtype=DataType.VARCHAR, is_primary=True, max_length=64, auto_id=False),
+                FieldSchema(name="vector", dtype=DataType.FLOAT_VECTOR, dim=self.dim),
+                FieldSchema(name="namespace", dtype=DataType.VARCHAR, max_length=128),
+                FieldSchema(name="project_id", dtype=DataType.VARCHAR, max_length=256),
+                FieldSchema(name="repo_id", dtype=DataType.VARCHAR, max_length=256),
+                FieldSchema(name="path", dtype=DataType.VARCHAR, max_length=1024),
+                FieldSchema(name="chunk_index", dtype=DataType.INT64),
+                FieldSchema(name="sha", dtype=DataType.VARCHAR, max_length=128),
+                FieldSchema(name="content", dtype=DataType.VARCHAR, max_length=65535),
+                FieldSchema(name="created_at", dtype=DataType.VARCHAR, max_length=64),
+            ]
+            schema = CollectionSchema(fields=fields, description="devmem repository chunks", enable_dynamic_field=False)
+            coll = Collection(self.collection_name, schema=schema, using=self.alias)
+            coll.create_index(
+                field_name="vector",
+                index_params={"index_type": "IVF_FLAT", "metric_type": "L2", "params": {"nlist": 2048}},
+            )
+        else:
+            coll = Collection(self.collection_name, using=self.alias)
+        return coll
+
+    def _embed(self, text: str) -> list[float]:
+        return embed_one(text)
+
+    def _embed_batch(self, texts: list[str]) -> list[list[float]]:
+        return embed_many(texts)
+
+    @staticmethod
+    def _doc_id(namespace: str, project_id: str, repo_id: str, path: str, chunk_index: int, sha: str) -> str:
+        h = hashlib.sha256()
+        h.update(namespace.encode("utf-8"))
+        h.update(project_id.encode("utf-8"))
+        h.update(repo_id.encode("utf-8"))
+        h.update(path.encode("utf-8"))
+        h.update(str(chunk_index).encode("utf-8"))
+        h.update(sha.encode("utf-8"))
+        return h.hexdigest()
+
+    def clear_project_data(self) -> None:
+        """Delete existing rows for this namespace/project/repo."""
+        if self.collection is None:
+            raise RuntimeError("Milvus collection is not initialized")
+
+        escaped_ns = self.namespace.replace('"', '\\"')
+        escaped_project = self.project_id.replace('"', '\\"')
+        escaped_repo = self.repo_id.replace('"', '\\"')
+        delete_expr = (
+            f'namespace == "{escaped_ns}" and project_id == "{escaped_project}" '
+            f'and repo_id == "{escaped_repo}"'
+        )
+        try:
+            self.collection.delete(expr=delete_expr)
+        except Exception:
+            logger.debug(
+                "Milvus project cleanup skipped namespace=%s project=%s repo=%s",
+                self.namespace,
+                self.project_id,
+                self.repo_id,
+                exc_info=True,
+            )
+
+    def replace_file_chunks(self, *, path: str, sha: str, chunks: Iterable[str]) -> int:
+        """Insert/update chunks for a single file path in the project collection."""
+        if self.collection is None:
+            raise RuntimeError("Milvus collection is not initialized")
+
+        rows = []
+        created_at = datetime.now(timezone.utc).isoformat()
+        chunk_texts = [chunk[:65500] for chunk in chunks]
+        vectors = self._embed_batch(chunk_texts) if chunk_texts else []
+        for idx, (chunk_text, vec) in enumerate(zip(chunk_texts, vectors)):
+            rows.append(
+                {
+                    "id": self._doc_id(self.namespace, self.project_id, self.repo_id, path, idx, sha),
+                    "vector": vec,
+                    "namespace": self.namespace,
+                    "project_id": self.project_id,
+                    "repo_id": self.repo_id,
+                    "path": path,
+                    "chunk_index": idx,
+                    "sha": sha,
+                    "content": chunk_text,
+                    "created_at": created_at,
+                }
+            )
+
+        if not rows:
+            return 0
+
+        if hasattr(self.collection, "upsert"):
+            self.collection.upsert(rows)
+        else:
+            self.collection.insert(rows)
+        return len(rows)
+
+    def flush(self) -> None:
+        if self.collection is not None:
+            self.collection.flush()
+
+    def fetch_project_chunks(self, *, max_rows: int = 20000, batch_size: int = 1000) -> list[dict[str, Any]]:
+        """Read chunk rows for this namespace/project/repo from Milvus."""
+        if self.collection is None:
+            raise RuntimeError("Milvus collection is not initialized")
+
+        escaped_ns = self.namespace.replace('"', '\\"')
+        escaped_project = self.project_id.replace('"', '\\"')
+        escaped_repo = self.repo_id.replace('"', '\\"')
+        expr = (
+            f'namespace == "{escaped_ns}" and project_id == "{escaped_project}" '
+            f'and repo_id == "{escaped_repo}"'
+        )
+        output_fields = ["path", "chunk_index", "content", "sha"]
+        rows: list[dict[str, Any]] = []
+        offset = 0
+
+        while len(rows) < max_rows:
+            limit = min(batch_size, max_rows - len(rows))
+            try:
+                batch = self.collection.query(
+                    expr=expr,
+                    output_fields=output_fields,
+                    limit=limit,
+                    offset=offset,
+                )
+            except TypeError:
+                # Older Milvus clients may not support offset.
+                batch = self.collection.query(
+                    expr=expr,
+                    output_fields=output_fields,
+                    limit=max_rows,
+                )
+                rows.extend(batch)
+                break
+
+            if not batch:
+                break
+            rows.extend(batch)
+            offset += len(batch)
+            if len(batch) < limit:
+                break
+
+        return rows[:max_rows]
+
+    def vector_search(self, *, query_text: str, top_k: int) -> list[dict[str, Any]]:
+        """Semantic ANN search over repo chunks using the configured embedding model."""
+        if self.collection is None:
+            raise RuntimeError("Milvus collection is not initialized")
+        if not (query_text or "").strip():
+            return []
+
+        query_vec = self._embed(query_text)
+        escaped_ns = self.namespace.replace('"', '\\"')
+        escaped_project = self.project_id.replace('"', '\\"')
+        escaped_repo = self.repo_id.replace('"', '\\"')
+        expr = (
+            f'namespace == "{escaped_ns}" and project_id == "{escaped_project}" '
+            f'and repo_id == "{escaped_repo}"'
+        )
+        try:
+            hits = self.collection.search(
+                data=[query_vec],
+                anns_field="vector",
+                param={"metric_type": "L2", "params": {"nprobe": 16}},
+                limit=top_k,
+                expr=expr,
+                output_fields=["path", "chunk_index", "content", "sha"],
+            )
+        except Exception:
+            logger.warning("Milvus vector search failed; falling back to lexical search", exc_info=True)
+            return self.lexical_search(query_text=query_text, top_k=top_k)
+
+        results: list[dict[str, Any]] = []
+        for hit_list in hits:
+            for hit in hit_list:
+                entity = getattr(hit, "entity", None)
+                getter = entity.get if entity is not None else (lambda key, default=None: default)
+                results.append(
+                    {
+                        # Lower L2 distance = closer match; convert to similarity score.
+                        "score": 1.0 / (1.0 + float(getattr(hit, "distance", 0.0))),
+                        "path": getter("path"),
+                        "chunk_index": getter("chunk_index"),
+                        "sha": getter("sha"),
+                        "content": getter("content"),
+                    }
+                )
+        return results
+
+    def lexical_search(self, *, query_text: str, top_k: int) -> list[dict[str, Any]]:
+        """Search ingested chunks with token-overlap scoring."""
+        query_lc = (query_text or "").strip().lower()
+        if not query_lc:
+            return []
+
+        tokens = list(dict.fromkeys(_TOKEN_RE.findall(query_lc)))
+        rows = self.fetch_project_chunks()
+        scored: list[tuple[float, dict[str, Any]]] = []
+
+        for row in rows:
+            content = str(row.get("content") or "")
+            content_lc = content.lower()
+            score = 0.0
+            if query_lc in content_lc:
+                score += 10.0
+            for token in tokens:
+                hits = content_lc.count(token)
+                if hits:
+                    score += min(3.0, float(hits))
+            if score <= 0:
+                continue
+            scored.append((score, row))
+
+        scored.sort(key=lambda item: item[0], reverse=True)
+        return [
+            {
+                "score": score,
+                "path": row.get("path"),
+                "chunk_index": row.get("chunk_index"),
+                "sha": row.get("sha"),
+                "content": row.get("content"),
+            }
+            for score, row in scored[:top_k]
+        ]
+
+    def endpoint_summary(self) -> str:
+        if self.cfg.milvus_uri:
+            host, port = parse_host_port_from_url(self.cfg.milvus_uri, default_port=self.cfg.milvus_port)
+            return f"{host}:{port}"
+        return f"{self.cfg.milvus_host}:{self.cfg.milvus_port}"