mnemo-dev 0.1.0__py3-none-any.whl

mnemo/vector_index/__init__.py

@@ -0,0 +1,180 @@
+"""Local vector index with optional ChromaDB backend."""
+
+from __future__ import annotations
+
+import os
+import subprocess
+import sys
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Protocol
+
+from ..chunking import Chunk
+from ..embeddings import KeywordEmbeddingProvider
+
+
+_CHROMA_INSTALL_ATTEMPTED = False
+
+
+def _auto_install_chromadb() -> bool:
+    """Attempt to install chromadb automatically. Only tries once per process."""
+    global _CHROMA_INSTALL_ATTEMPTED
+    if _CHROMA_INSTALL_ATTEMPTED:
+        return False
+    _CHROMA_INSTALL_ATTEMPTED = True
+
+    if getattr(sys, "frozen", False):
+        return False
+
+    # Only auto-install if triggered by `mnemo init`, not on every recall
+    if os.environ.get("MNEMO_AUTO_INSTALL") != "1":
+        return False
+
+    try:
+        subprocess.check_call(
+            [sys.executable, "-m", "pip", "install", "chromadb>=0.5", "--quiet"],
+            stdout=subprocess.DEVNULL,
+            stderr=subprocess.DEVNULL,
+        )
+        return True
+    except Exception:
+        return False
+
+
+class VectorIndex(Protocol):
+    def available(self) -> bool:
+        ...
+
+    def upsert(self, namespace: str, chunks: list[Chunk]) -> None:
+        ...
+
+    def query(self, namespace: str, query: str, limit: int = 10, filters: dict[str, Any] | None = None) -> list[dict[str, Any]]:
+        ...
+
+    def clear(self, namespace: str | None = None) -> None:
+        ...
+
+
+@dataclass
+class _MemoryRecord:
+    id: str
+    text: str
+    metadata: dict[str, Any]
+
+
+class LocalVectorIndex:
+    """Chroma-first index with in-memory keyword fallback."""
+
+    def __init__(self, repo_root: Path):
+        self.repo_root = repo_root
+        self.index_dir = repo_root / ".mnemo" / "index"
+        self.index_dir.mkdir(parents=True, exist_ok=True)
+        self._fallback = KeywordEmbeddingProvider()
+        self._memory_store: dict[str, list[_MemoryRecord]] = {}
+        self._chroma_collection_cache: dict[str, Any] = {}
+        self._chroma_client: Any = None
+        self._chroma_ready = False
+        self._init_chroma()
+
+    def _init_chroma(self) -> None:
+        try:
+            import chromadb
+        except ImportError:
+            if not _auto_install_chromadb():
+                self._chroma_ready = False
+                return
+            try:
+                import chromadb
+            except ImportError:
+                self._chroma_ready = False
+                return
+        except Exception:
+            self._chroma_ready = False
+            return
+        try:
+            self._chroma_client = chromadb.PersistentClient(path=str(self.index_dir / "chroma"))
+            self._chroma_ready = True
+        except Exception:
+            self._chroma_ready = False
+
+    def _collection(self, namespace: str):
+        if not self._chroma_ready or self._chroma_client is None:
+            return None
+        if namespace not in self._chroma_collection_cache:
+            self._chroma_collection_cache[namespace] = self._chroma_client.get_or_create_collection(
+                name=namespace
+            )
+        return self._chroma_collection_cache[namespace]
+
+    def available(self) -> bool:
+        return self._chroma_ready
+
+    def upsert(self, namespace: str, chunks: list[Chunk]) -> None:
+        self._memory_store[namespace] = [
+            _MemoryRecord(
+                id=chunk.id,
+                text=f"{chunk.path} {chunk.symbol}\n{chunk.content}",
+                metadata={**chunk.metadata, "path": chunk.path, "symbol": chunk.symbol, "chunk_type": chunk.chunk_type},
+            )
+            for chunk in chunks
+        ]
+        collection = self._collection(namespace)
+        if not collection:
+            return
+        if chunks:
+            collection.upsert(
+                ids=[chunk.id for chunk in chunks],
+                documents=[f"{chunk.path} {chunk.symbol}\n{chunk.content}" for chunk in chunks],
+                metadatas=[
+                    {
+                        "path": chunk.path,
+                        "language": chunk.language,
+                        "symbol": chunk.symbol,
+                        "chunk_type": chunk.chunk_type,
+                        **chunk.metadata,
+                    }
+                    for chunk in chunks
+                ],
+            )
+
+    def _query_fallback(self, namespace: str, query: str, limit: int, filters: dict[str, Any] | None) -> list[dict[str, Any]]:
+        query_emb = self._fallback.embed(query)
+        records = self._memory_store.get(namespace, [])
+        scored: list[tuple[float, _MemoryRecord]] = []
+        for record in records:
+            if filters and any(record.metadata.get(k) != v for k, v in filters.items()):
+                continue
+            score = query_emb.score(self._fallback.embed(record.text))
+            if score > 0:
+                scored.append((score, record))
+        scored.sort(key=lambda item: item[0], reverse=True)
+        return [
+            {"id": record.id, "score": score, "content": record.text, "metadata": record.metadata}
+            for score, record in scored[:limit]
+        ]
+
+    def query(self, namespace: str, query: str, limit: int = 10, filters: dict[str, Any] | None = None) -> list[dict[str, Any]]:
+        collection = self._collection(namespace)
+        if not collection:
+            return self._query_fallback(namespace, query, limit, filters)
+        try:
+            where = filters or None
+            result = collection.query(query_texts=[query], n_results=limit, where=where)
+            ids = result.get("ids", [[]])[0]
+            docs = result.get("documents", [[]])[0]
+            distances = result.get("distances", [[]])[0] if result.get("distances") else [0.0] * len(ids)
+            metas = result.get("metadatas", [[]])[0] if result.get("metadatas") else [{} for _ in ids]
+            return [
+                {"id": doc_id, "score": 1.0 - float(dist), "content": doc, "metadata": meta}
+                for doc_id, doc, dist, meta in zip(ids, docs, distances, metas)
+            ]
+        except Exception:
+            return self._query_fallback(namespace, query, limit, filters)
+
+    def clear(self, namespace: str | None = None) -> None:
+        if namespace is None:
+            self._memory_store.clear()
+            self._chroma_collection_cache.clear()
+            return
+        self._memory_store.pop(namespace, None)
+        self._chroma_collection_cache.pop(namespace, None)

mnemo/workspace/__init__.py

@@ -0,0 +1,224 @@
+"""Multi-repo workspace — link, discover, and query across repositories."""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Any
+
+from ..config import mnemo_path
+
+LINKS_FILE = "links.json"
+
+
+def _links_path(repo_root: Path) -> Path:
+    return mnemo_path(repo_root) / LINKS_FILE
+
+
+def get_linked_repos(repo_root: Path) -> list[Path]:
+    """Return resolved paths of all linked repos that exist and are initialized."""
+    path = _links_path(repo_root)
+    if not path.exists():
+        return []
+    try:
+        data = json.loads(path.read_text(encoding="utf-8"))
+    except (json.JSONDecodeError, OSError):
+        return []
+    linked: list[Path] = []
+    for entry in data:
+        repo_path = Path(entry["path"]).expanduser().resolve()
+        if repo_path != repo_root.resolve() and (repo_path / ".mnemo").exists():
+            linked.append(repo_path)
+    return linked
+
+
+def _save_links(repo_root: Path, links: list[dict[str, str]]) -> None:
+    path = _links_path(repo_root)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    path.write_text(json.dumps(links, indent=2), encoding="utf-8")
+
+
+def link_repo(repo_root: Path, target: Path) -> str:
+    """Link a sibling repo for cross-repo queries."""
+    target = target.expanduser().resolve()
+    if not target.exists():
+        return f"Path does not exist: {target}"
+    if not (target / ".git").exists() and not (target / ".mnemo").exists():
+        return f"Not a repo: {target} (no .git or .mnemo found)"
+
+    existing = get_linked_repos(repo_root)
+    if target in [r.resolve() for r in existing]:
+        return f"Already linked: {target.name}"
+
+    # Load raw data to append
+    path = _links_path(repo_root)
+    data: list[dict[str, str]] = []
+    if path.exists():
+        try:
+            data = json.loads(path.read_text(encoding="utf-8"))
+        except (json.JSONDecodeError, OSError):
+            data = []
+
+    data.append({"name": target.name, "path": str(target)})
+    _save_links(repo_root, data)
+
+    initialized = (target / ".mnemo").exists()
+    status = "✓ indexed" if initialized else "⚠ needs `mnemo init`"
+    return f"Linked: {target.name} ({status})"
+
+
+def unlink_repo(repo_root: Path, name: str) -> str:
+    """Remove a linked repo by name or path."""
+    path = _links_path(repo_root)
+    if not path.exists():
+        return "No linked repos."
+    try:
+        data = json.loads(path.read_text(encoding="utf-8"))
+    except (json.JSONDecodeError, OSError):
+        return "No linked repos."
+
+    name_lower = name.lower()
+    filtered = [e for e in data if e.get("name", "").lower() != name_lower and name_lower not in e.get("path", "").lower()]
+    if len(filtered) == len(data):
+        return f"No linked repo matching '{name}'."
+    _save_links(repo_root, filtered)
+    return f"Unlinked: {name}"
+
+
+def discover_repos(repo_root: Path, search_dir: Path, auto_init: bool = False) -> str:
+    """Auto-discover and link all repos under a directory. Optionally init uninitialized ones."""
+    search_dir = search_dir.expanduser().resolve()
+    if not search_dir.exists():
+        return f"Directory not found: {search_dir}"
+
+    found: list[Path] = []
+    for child in search_dir.iterdir():
+        if not child.is_dir():
+            continue
+        if child.resolve() == repo_root.resolve():
+            continue
+        if (child / ".git").exists():
+            found.append(child)
+
+    if not found:
+        return f"No git repos found under {search_dir}"
+
+    results = []
+    for repo in sorted(found):
+        result = link_repo(repo_root, repo)
+        results.append(result)
+
+        # Auto-init if requested and not yet initialized
+        if auto_init and not (repo / ".mnemo").exists():
+            try:
+                from ..init import init
+                import sys
+                sys.stderr.write(f"  Initializing {repo.name}...\n")
+                init(repo)
+                results.append(f"  ✓ Initialized {repo.name}")
+            except Exception as e:
+                results.append(f"  ✗ Failed to init {repo.name}: {e}")
+
+    return "\n".join(results)
+
+
+def format_links(repo_root: Path) -> str:
+    """Show all linked repos with status."""
+    path = _links_path(repo_root)
+    if not path.exists():
+        return "No linked repos. Use `mnemo link <path>` or `mnemo link --discover <dir>`."
+
+    try:
+        data = json.loads(path.read_text(encoding="utf-8"))
+    except (json.JSONDecodeError, OSError):
+        return "No linked repos."
+
+    if not data:
+        return "No linked repos."
+
+    lines = [f"# Linked Repos ({len(data)})\n"]
+    for entry in data:
+        name = entry.get("name", "unknown")
+        repo_path = Path(entry.get("path", ""))
+        exists = repo_path.exists()
+        initialized = (repo_path / ".mnemo").exists() if exists else False
+
+        if not exists:
+            status = "✗ path not found"
+        elif initialized:
+            status = "✓ indexed"
+        else:
+            status = "⚠ needs `mnemo init`"
+
+        lines.append(f"- **{name}**  {repo_path}  {status}")
+
+    return "\n".join(lines)
+
+
+def cross_repo_semantic_query(
+    repo_root: Path, namespace: str, query: str, limit: int = 10
+) -> list[dict[str, Any]]:
+    """Query this repo + all linked repos, merge and rank results."""
+    from ..retrieval import semantic_query
+
+    # Query local repo first
+    results = semantic_query(repo_root, namespace, query, limit=limit)
+    for r in results:
+        r["repo"] = repo_root.name
+
+    # Query linked repos
+    for linked in get_linked_repos(repo_root):
+        try:
+            linked_results = semantic_query(linked, namespace, query, limit=limit // 2)
+            for r in linked_results:
+                r["repo"] = linked.name
+            results.extend(linked_results)
+        except Exception:
+            continue
+
+    # Sort by score descending, deduplicate by id
+    seen: set[str] = set()
+    unique: list[dict[str, Any]] = []
+    for r in sorted(results, key=lambda x: x.get("score", 0), reverse=True):
+        rid = r.get("id", "")
+        if rid not in seen:
+            seen.add(rid)
+            unique.append(r)
+
+    return unique[:limit]
+
+
+def cross_repo_impact(repo_root: Path, query: str) -> str:
+    """Check all linked repos for code that depends on the queried service/file."""
+    from ..retrieval import semantic_query
+
+    lines = [f"# Cross-Repo Impact: '{query}'\n"]
+    lines.append("## This Repo")
+
+    local_hits = semantic_query(repo_root, "code", query, limit=5)
+    if local_hits:
+        for hit in local_hits:
+            meta = hit.get("metadata", {})
+            lines.append(f"- `{meta.get('path', '')}` :: `{meta.get('symbol', '')}`")
+    else:
+        lines.append("- No matches in this repo")
+
+    linked = get_linked_repos(repo_root)
+    if not linked:
+        lines.append("\nNo linked repos. Use `mnemo link` to enable cross-repo impact analysis.")
+        return "\n".join(lines)
+
+    for linked_repo in linked:
+        lines.append(f"\n## {linked_repo.name}")
+        try:
+            hits = semantic_query(linked_repo, "code", query, limit=5)
+            if hits:
+                for hit in hits:
+                    meta = hit.get("metadata", {})
+                    lines.append(f"- `{meta.get('path', '')}` :: `{meta.get('symbol', '')}`")
+            else:
+                lines.append("- No matches")
+        except Exception:
+            lines.append("- ⚠ Could not query (run `mnemo init` in that repo)")
+
+    return "\n".join(lines)