memplex 3.2.0__py3-none-any.whl

memnex/storage/vector.py

@@ -0,0 +1,265 @@
+"""Vector store abstraction -- migrated from the legacy ``storage/`` package.
+
+Provides a ``VectorStore`` Protocol plus two implementations:
+
+* ``InMemoryVectorStore`` -- TF-IDF bag-of-words cosine similarity, zero
+  external dependencies.
+* ``ChromaVectorStore`` -- ChromaDB + sentence-transformers for production
+  quality embeddings.
+
+Usage::
+
+    from memnex.storage.vector import create_vector_store
+
+    vs = create_vector_store("auto")   # ChromaDB if available, else InMemory
+"""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass
+from typing import Dict, List, Optional, Protocol, runtime_checkable
+
+logger = logging.getLogger(__name__)
+
+# Re-export the list[float] alias for convenience
+Vector = List[float]
+
+
+@dataclass
+class VectorSearchResult:
+    """A single vector search hit."""
+
+    id: str
+    score: float
+    text: str
+
+
+# ── Protocol ────────────────────────────────────────────────────────
+
+
+@runtime_checkable
+class VectorStore(Protocol):
+    """Minimal vector store interface."""
+
+    def add(self, id: str, text: str, metadata: Optional[dict] = None) -> None: ...
+
+    def upsert(self, id: str, vector: Vector, text: str = "") -> None: ...
+
+    def upsert_batch(self, items: Dict[str, Vector]) -> None: ...
+
+    def search(
+        self, query: str, top_k: int = 5, query_vector: Optional[Vector] = None
+    ) -> List[VectorSearchResult]: ...
+
+    def delete(self, id: str) -> None: ...
+
+    def clear(self) -> None: ...
+
+
+# ── InMemory implementation ─────────────────────────────────────────
+
+
+class InMemoryVectorStore:
+    """In-memory TF-IDF cosine similarity store.
+
+    Zero external dependencies.  Suitable for Lite backend and testing.
+    """
+
+    def __init__(self) -> None:
+        self._vectors: Dict[str, tuple] = {}   # id -> (text, embedding)
+        self._stored_vectors: Dict[str, Vector] = {}  # id -> pre-stored vector
+        self._all_words: set = set()
+
+    def add(self, id: str, text: str, metadata: Optional[dict] = None) -> None:
+        words = self._get_words(text)
+        self._all_words.update(words)
+        embedding = self._encode_with_vocab(text, words)
+        self._vectors[id] = (text, embedding)
+        self._stored_vectors.pop(id, None)
+
+    def upsert(self, id: str, vector: Vector, text: str = "") -> None:
+        self._stored_vectors[id] = vector
+        self._vectors[id] = (text, [0])  # placeholder text embedding
+
+    def upsert_batch(self, items: Dict[str, Vector]) -> None:
+        for id, vector in items.items():
+            self.upsert(id, vector)
+
+    def search(
+        self,
+        query: str,
+        top_k: int = 5,
+        query_vector: Optional[Vector] = None,
+    ) -> List[VectorSearchResult]:
+        if query_vector is not None:
+            return self._search_by_vector(query_vector, top_k)
+        query_emb = self._encode_with_vocab(query, self._get_words(query))
+        return self._search_by_embedding(query_emb, top_k)
+
+    def _search_by_vector(
+        self, query_vec: Vector, top_k: int
+    ) -> List[VectorSearchResult]:
+        scores: list = []
+        for vid, vec in self._stored_vectors.items():
+            score = self._cosine(query_vec, vec)
+            text = self._vectors.get(vid, ("", None))[0]
+            scores.append((vid, score, text))
+        scores.sort(key=lambda x: x[1], reverse=True)
+        return [VectorSearchResult(id=s[0], score=s[1], text=s[2]) for s in scores[:top_k]]
+
+    def _search_by_embedding(
+        self, query_emb: list, top_k: int
+    ) -> List[VectorSearchResult]:
+        scores: list = []
+        for vid, (text, emb) in self._vectors.items():
+            score = self._cosine(query_emb, emb)
+            scores.append((vid, score, text))
+        scores.sort(key=lambda x: x[1], reverse=True)
+        return [VectorSearchResult(id=s[0], score=s[1], text=s[2]) for s in scores[:top_k]]
+
+    def delete(self, id: str) -> None:
+        self._vectors.pop(id, None)
+        self._stored_vectors.pop(id, None)
+
+    def clear(self) -> None:
+        self._vectors.clear()
+        self._stored_vectors.clear()
+        self._all_words.clear()
+
+    # ── Helpers ──────────────────────────────────────────────────
+
+    @staticmethod
+    def _get_words(text: str) -> set:
+        text_lower = text.lower()
+        if any("一" <= c <= "鿿" for c in text):
+            return set(list(text_lower))
+        return set(text_lower.split())
+
+    def _encode_with_vocab(self, text: str, words: set) -> list:
+        if not self._all_words:
+            return [0]
+        return [1 if w in words else 0 for w in sorted(self._all_words)]
+
+    @staticmethod
+    def _cosine(a: list, b: list) -> float:
+        dot = sum(x * y for x, y in zip(a, b))
+        norm_a = sum(x * x for x in a) ** 0.5
+        norm_b = sum(x * x for x in b) ** 0.5
+        return dot / (norm_a * norm_b + 1e-10)
+
+
+# ── ChromaDB implementation ─────────────────────────────────────────
+
+try:
+    import chromadb  # type: ignore
+    from chromadb.config import Settings  # type: ignore
+
+    _CHROMA_AVAILABLE = True
+except ImportError:
+    _CHROMA_AVAILABLE = False
+    chromadb = None  # type: ignore[assignment]
+
+
+class ChromaVectorStore:
+    """ChromaDB-backed vector store with sentence-transformers embeddings."""
+
+    def __init__(
+        self,
+        collection_name: str = "functions",
+        embedding_model: str = "all-MiniLM-L6-v2",
+    ) -> None:
+        if not _CHROMA_AVAILABLE:
+            raise ImportError(
+                "chromadb not installed: pip install chromadb sentence-transformers"
+            )
+        self.client = chromadb.Client(Settings(anonymized_telemetry=False))
+        self.collection = self.client.get_or_create_collection(
+            name=collection_name
+        )
+        self._embedding_model = embedding_model
+        self._model = None
+
+    def _get_model(self):
+        if self._model is None:
+            from sentence_transformers import SentenceTransformer  # type: ignore
+
+            self._model = SentenceTransformer(self._embedding_model)
+        return self._model
+
+    def add(self, id: str, text: str, metadata: Optional[dict] = None) -> None:
+        embedding = self._get_model().encode([text])[0]
+        self.collection.upsert(
+            ids=[id],
+            embeddings=[embedding.tolist()],
+            documents=[text],
+            metadatas=[metadata or {}],
+        )
+
+    def upsert(self, id: str, vector: Vector, text: str = "") -> None:
+        self.collection.upsert(
+            ids=[id],
+            embeddings=[vector if isinstance(vector, list) else list(vector)],
+            documents=[text],
+            metadatas=[{}],
+        )
+
+    def upsert_batch(self, items: Dict[str, Vector]) -> None:
+        ids = list(items.keys())
+        vectors = [v if isinstance(v, list) else list(v) for v in items.values()]
+        self.collection.upsert(
+            ids=ids,
+            embeddings=vectors,
+            documents=[""] * len(ids),
+            metadatas=[{}] * len(ids),
+        )
+
+    def search(
+        self,
+        query: str,
+        top_k: int = 5,
+        query_vector: Optional[Vector] = None,
+    ) -> List[VectorSearchResult]:
+        if query_vector is not None:
+            q_emb = query_vector if isinstance(query_vector, list) else list(query_vector)
+        else:
+            q_emb = self._get_model().encode([query])[0].tolist()
+        results = self.collection.query(
+            query_embeddings=[q_emb], n_results=top_k
+        )
+        ids = results.get("ids", [[]])[0]
+        distances = results.get("distances", [[]])[0]
+        documents = results.get("documents", [[]])[0]
+        return [
+            VectorSearchResult(id=ids[i], score=float(distances[i]), text=documents[i])
+            for i in range(len(ids))
+        ]
+
+    def delete(self, id: str) -> None:
+        self.collection.delete(ids=[id])
+
+    def clear(self) -> None:
+        self.collection.delete(where={})
+
+
+# ── Factory ──────────────────────────────────────────────────────────
+
+
+def create_vector_store(backend: str = "auto") -> VectorStore:
+    """Create a vector store by backend name.
+
+    Parameters
+    ----------
+    backend:
+        ``"inmemory"`` | ``"chroma"`` | ``"auto"`` (default).
+        ``"auto"`` prefers ChromaDB and falls back to InMemory.
+    """
+    if backend == "chroma" and _CHROMA_AVAILABLE:
+        return ChromaVectorStore()
+    if backend == "inmemory":
+        return InMemoryVectorStore()
+    if backend == "auto":
+        if _CHROMA_AVAILABLE:
+            return ChromaVectorStore()
+        return InMemoryVectorStore()
+    raise ValueError(f"Unknown vector store backend: {backend!r}")

memnex/wiki/__init__.py

@@ -0,0 +1,11 @@
+"""MemNex Wiki Layer -- compile, generate, search, and lint wiki pages."""
+
+from memnex.wiki.compiler import WikiCompiler
+from memnex.wiki.generator import LLMWikiGenerator
+from memnex.wiki.search import DualIndexSearch
+
+__all__ = [
+    "WikiCompiler",
+    "LLMWikiGenerator",
+    "DualIndexSearch",
+]

memnex/wiki/community.py

@@ -0,0 +1,221 @@
+"""GraphRAG community detection using the Louvain algorithm.
+
+Detects communities (clusters) in the knowledge graph and generates
+concept pages for each community.  Falls back to domain-based grouping
+when python-louvain / networkx are not available.
+"""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass, field
+from typing import Dict, List, Optional, TYPE_CHECKING
+
+from memnex.models import (
+    Function,
+    GraphData,
+    WikiPage,
+)
+
+if TYPE_CHECKING:
+    from memnex.storage.base import MemoryStore
+
+logger = logging.getLogger(__name__)
+
+# Only use strong-relationship edge types for community detection
+STRONG_EDGE_TYPES = frozenset({
+    "REFERENCES",
+    "DEPENDS_ON",
+    "IMPLEMENTS",
+    "ASSOCIATED_WITH",
+})
+
+DEFAULT_MIN_COMMUNITY_SIZE = 3
+
+
+@dataclass
+class Community:
+    """A detected community of function nodes."""
+
+    community_id: int
+    node_ids: List[str] = field(default_factory=list)
+    theme: str = ""
+
+    @property
+    def size(self) -> int:
+        return len(self.node_ids)
+
+
+class GraphCommunityDetector:
+    """Louvain-based graph community detector.
+
+    Detects communities (clusters) in the knowledge graph using the
+    Louvain algorithm.  Falls back to domain-based grouping when
+    ``python-louvain`` or ``networkx`` are not installed.
+
+    Parameters
+    ----------
+    min_community_size:
+        Minimum number of nodes for a community to be reported.
+        Communities below this threshold are discarded as noise.
+    """
+
+    def __init__(
+        self,
+        min_community_size: int = DEFAULT_MIN_COMMUNITY_SIZE,
+    ) -> None:
+        self.min_community_size = min_community_size
+
+    def detect_communities(
+        self,
+        graph: GraphData,
+        min_size: Optional[int] = None,
+    ) -> List[Community]:
+        """Detect communities in the graph.
+
+        Parameters
+        ----------
+        graph:
+            The graph data containing nodes and edges.
+        min_size:
+            Override minimum community size for this call.
+
+        Returns
+        -------
+        List of Community objects, each with a list of node IDs.
+        """
+        threshold = min_size if min_size is not None else self.min_community_size
+
+        try:
+            import networkx as nx  # type: ignore
+            import community as community_louvain  # type: ignore
+        except ImportError:
+            logger.info(
+                "python-louvain/networkx not available, "
+                "falling back to domain grouping"
+            )
+            return self._fallback_domain_grouping(graph, threshold)
+
+        # Build networkx graph with strong edges only
+        G = nx.Graph()
+        for node in graph.nodes:
+            node_id = node.id if hasattr(node, "id") else str(node)
+            G.add_node(node_id)
+        for edge in graph.edges:
+            if edge.edge_type in STRONG_EDGE_TYPES:
+                G.add_edge(edge.source, edge.target, weight=edge.weight)
+
+        if G.number_of_edges() == 0:
+            logger.info("No strong edges found, falling back to domain grouping")
+            return self._fallback_domain_grouping(graph, threshold)
+
+        partition = community_louvain.best_partition(G)
+        groups: Dict[int, List[str]] = {}
+        for node_id, comm_id in partition.items():
+            groups.setdefault(comm_id, []).append(node_id)
+
+        communities: List[Community] = []
+        for comm_id, node_ids in groups.items():
+            if len(node_ids) >= threshold:
+                communities.append(
+                    Community(
+                        community_id=comm_id,
+                        node_ids=node_ids,
+                    )
+                )
+
+        return communities
+
+    def generate_concept_pages(
+        self,
+        communities: List[Community],
+        store: MemoryStore,
+    ) -> List[WikiPage]:
+        """Generate concept pages for detected communities.
+
+        Each community gets a concept page listing its member functions.
+
+        Parameters
+        ----------
+        communities:
+            Detected communities.
+        store:
+            MemoryStore for looking up Function details.
+
+        Returns
+        -------
+        List of WikiPage objects for the communities.
+        """
+        pages: List[WikiPage] = []
+
+        for community in communities:
+            lines: list[str] = [
+                f"# Community {community.community_id}",
+                "",
+            ]
+            if community.theme:
+                lines.append(f"**Theme:** {community.theme}")
+                lines.append("")
+
+            lines.append(f"**Members:** {len(community.node_ids)} functions")
+            lines.append("")
+
+            lines.append("## Functions")
+            lines.append("")
+
+            for node_id in community.node_ids:
+                func = store.get(node_id)
+                if func:
+                    lines.append(
+                        f"- [[{func.id}]] -- "
+                        f"{func.domain or 'uncategorized'} "
+                        f"(confidence: {func.confidence:.2f})"
+                    )
+                else:
+                    lines.append(f"- [[{node_id}]]")
+
+            lines.append("")
+
+            page_id = f"community_{community.community_id}"
+            content = "\n".join(lines)
+
+            pages.append(WikiPage(
+                page_id=page_id,
+                content=content,
+                metadata={
+                    "type": "community",
+                    "community_id": community.community_id,
+                    "member_count": len(community.node_ids),
+                    "theme": community.theme,
+                },
+            ))
+
+        return pages
+
+    # ── Private helpers ───────────────────────────────────────────────
+
+    def _fallback_domain_grouping(
+        self,
+        graph: GraphData,
+        min_size: int,
+    ) -> List[Community]:
+        """Degraded community detection: group nodes by domain field."""
+        groups: Dict[str, List[str]] = {}
+        for node in graph.nodes:
+            node_id = node.id if hasattr(node, "id") else str(node)
+            domain = (
+                getattr(node, "domain", None) or "uncategorized"
+            )
+            groups.setdefault(domain, []).append(node_id)
+
+        communities: List[Community] = []
+        for idx, (domain, node_ids) in enumerate(sorted(groups.items())):
+            if len(node_ids) >= min_size:
+                communities.append(
+                    Community(
+                        community_id=idx,
+                        node_ids=node_ids,
+                        theme=domain,
+                    )
+                )
+        return communities