memplex 3.2.0__py3-none-any.whl

memnex/wiki/generator.py

@@ -0,0 +1,270 @@
+"""LLMWikiGenerator -- LLM-augmented wiki page generation.
+
+Uses the LLM to produce enriched Entity Pages, summaries, concept pages,
+and cross-reference suggestions.  All LLM calls go through
+``LLMPromptSanitizer.build_structured_prompt`` for safe input handling.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+from typing import List, Optional, TYPE_CHECKING
+
+from memnex.models import (
+    FieldValue,
+    Function,
+    WikiPage,
+)
+from memnex.llm.sanitizer import LLMPromptSanitizer
+
+if TYPE_CHECKING:
+    from memnex.llm.enhancer import LLMEnhancer
+
+logger = logging.getLogger(__name__)
+
+# Default maximum input length for prompt sanitization
+DEFAULT_MAX_INPUT_LENGTH = 10000
+
+
+class LLMWikiGenerator:
+    """LLM-enhanced wiki page generator.
+
+    Uses LLM understanding to produce higher-quality Wiki content than
+    rule-based compilation alone.
+
+    Parameters
+    ----------
+    llm_enhancer:
+        The :class:`LLMEnhancer` instance used for LLM calls.
+    sep:
+        Separator string used to delimit sections in generated pages.
+    max_input_length:
+        Maximum character length for sanitized inputs sent to the LLM.
+    """
+
+    def __init__(
+        self,
+        llm_enhancer: LLMEnhancer,
+        sep: str = "---",
+        max_input_length: int = DEFAULT_MAX_INPUT_LENGTH,
+    ) -> None:
+        self._llm = llm_enhancer
+        self.sep = sep
+        self.max_input_length = max_input_length
+
+    # ── Public API ────────────────────────────────────────────────────
+
+    async def generate_entity_page(self, func: Function) -> str:
+        """Use LLM to generate an enhanced Entity Page for a Function.
+
+        The LLM receives structured Function data and returns a natural
+        language wiki page covering: one-sentence summary, trigger
+        conditions, execution flow, and related functions.
+        """
+        func_data = {
+            "name": func.name,
+            "domain": func.domain or "uncategorized",
+            "trigger": [fv.desc for fv in func.trigger],
+            "condition": [fv.desc for fv in func.condition],
+            "action": [fv.desc for fv in func.action],
+            "benefit": [fv.desc for fv in func.benefit],
+        }
+
+        prompt = LLMPromptSanitizer.build_structured_prompt(
+            instruction=(
+                "为以下函数生成一份清晰的 Wiki 页面，包含："
+                "1.一句话概括 2.触发条件与前置条件 "
+                "3.执行流程自然语言描述 4.关联函数（如有）"
+            ),
+            user_input=json.dumps(func_data, ensure_ascii=False),
+            max_length=self.max_input_length,
+        )
+        return await self._llm.llm.complete(prompt)
+
+    async def generate_summary(self, functions: List[Function]) -> str:
+        """Use LLM to generate a domain summary from multiple Functions.
+
+        Produces a concept-level overview including core responsibilities,
+        functional components, collaboration patterns, and key workflows.
+        """
+        funcs_data = [
+            {"name": f.name, "action": [fv.desc for fv in f.action]}
+            for f in functions
+        ]
+
+        prompt = LLMPromptSanitizer.build_structured_prompt(
+            instruction=(
+                "分析以下函数列表，生成简洁领域摘要，包含："
+                "1.核心职责 2.主要功能组件 3.协作关系 4.关键业务流程"
+            ),
+            user_input=json.dumps(funcs_data, ensure_ascii=False),
+            max_length=self.max_input_length,
+        )
+        return await self._llm.llm.complete(prompt)
+
+    async def generate_concept_page(
+        self,
+        domain: str,
+        functions: List[Function],
+    ) -> str:
+        """Generate a concept page that aggregates Functions by domain.
+
+        Parameters
+        ----------
+        domain:
+            The domain / topic label for this concept page.
+        functions:
+            Functions belonging to this domain.
+        """
+        funcs_summary = [
+            {
+                "name": f.name,
+                "trigger": [fv.desc for fv in f.trigger[:2]],
+                "action": [fv.desc for fv in f.action[:2]],
+            }
+            for f in functions[:20]
+        ]
+
+        prompt = LLMPromptSanitizer.build_structured_prompt(
+            instruction=(
+                f"为领域 \"{domain}\" 生成一份概念聚合页面，包含："
+                "1.领域概述 2.核心功能列表 3.功能间协作关系 4.业务流程图描述"
+            ),
+            user_input=json.dumps(funcs_summary, ensure_ascii=False),
+            max_length=self.max_input_length,
+        )
+        return await self._llm.llm.complete(prompt)
+
+    async def update_cross_references(
+        self,
+        pages: List[WikiPage],
+    ) -> List[WikiPage]:
+        """Use LLM to discover and update cross-references across pages.
+
+        For each page, the LLM analyses the content and suggests relevant
+        ``[[wikilinks]]`` to other pages in the corpus.
+
+        Parameters
+        ----------
+        pages:
+            All wiki pages to analyse for cross-references.
+
+        Returns
+        -------
+        Updated list of WikiPage objects with enriched cross-references.
+        """
+        if not pages:
+            return pages
+
+        # Build a lightweight summary index for the LLM
+        page_summaries: list[dict] = []
+        for p in pages:
+            # Extract first non-heading, non-blank line as summary
+            summary = ""
+            for line in p.content.splitlines():
+                stripped = line.strip()
+                if stripped and not stripped.startswith("#") and not stripped.startswith("---"):
+                    summary = stripped[:200]
+                    break
+            page_summaries.append({"id": p.page_id, "summary": summary})
+
+        updated: List[WikiPage] = []
+        for page in pages:
+            try:
+                candidates = [
+                    s for s in page_summaries if s["id"] != page.page_id
+                ][:20]
+
+                prompt = LLMPromptSanitizer.build_structured_prompt(
+                    instruction=(
+                        "分析当前 Wiki 页面，从候选页面中选择相关页面，"
+                        "返回相关页面的 ID 列表和关联理由。"
+                        "格式：{\"related\": [{\"id\": \"str\", \"reason\": \"str\"}]}"
+                    ),
+                    user_input=json.dumps(
+                        {
+                            "current_page": {"id": page.page_id, "summary": page.content[:500]},
+                            "candidates": candidates,
+                        },
+                        ensure_ascii=False,
+                    ),
+                    output_schema={
+                        "related": [{"id": "str", "reason": "str"}],
+                    },
+                    max_length=self.max_input_length,
+                )
+                result = await self._llm.llm.complete_json(prompt)
+
+                # Inject cross-references into page content
+                related = result.get("related", [])
+                if related:
+                    link_lines = [
+                        f"- [[{r.get('id', '')}]] -- {r.get('reason', '')}"
+                        for r in related
+                        if r.get("id")
+                    ]
+                    if link_lines:
+                        cross_ref_block = (
+                            "\n## Cross-References (LLM)\n"
+                            + "\n".join(link_lines)
+                            + "\n"
+                        )
+                        new_content = page.content.rstrip() + cross_ref_block
+                        updated.append(
+                            WikiPage(
+                                page_id=page.page_id,
+                                content=new_content,
+                                metadata=page.metadata,
+                            )
+                        )
+                        continue
+            except Exception:
+                logger.warning(
+                    "Cross-reference generation failed for %s, keeping original",
+                    page.page_id,
+                )
+
+            updated.append(page)
+
+        return updated
+
+    async def generate_community_page(
+        self,
+        community_funcs: List[Function],
+        community_id: int,
+    ) -> dict:
+        """Generate a Concept Page for a GraphRAG-detected community.
+
+        At most 20 functions are sent to the LLM to avoid token overflow.
+
+        Parameters
+        ----------
+        community_funcs:
+            Functions belonging to this community.
+        community_id:
+            Numeric identifier for the community.
+
+        Returns
+        -------
+        Parsed JSON dict from the LLM response.
+        """
+        func_summaries = [
+            f"{f.name}: {', '.join(fv.desc for fv in f.action[:1])}"
+            for f in community_funcs[:20]
+        ]
+        safe_text = LLMPromptSanitizer.sanitize(
+            "\n".join(func_summaries), self.max_input_length,
+        )
+
+        prompt = LLMPromptSanitizer.build_structured_prompt(
+            instruction="分析以下功能节点的聚类，识别共同主题并生成简洁社区摘要",
+            user_input=safe_text,
+            output_schema={
+                "community_theme": "str",
+                "core_functions": ["str"],
+                "summary": "str",
+                "key_relationships": ["str"],
+            },
+        )
+        return await self._llm.llm.complete_json(prompt)

memnex/wiki/search.py

@@ -0,0 +1,282 @@
+"""DualIndexSearch -- hybrid FTS + vector search for wiki pages.
+
+Combines full-text keyword search with vector semantic search using
+Reciprocal Rank Fusion (RRF) to merge results.
+
+ID mapping convention: Wiki page filename = Function.id + ".md"
+FTS and vector results are unified by Function.id, so RRF merging is
+unambiguous.
+"""
+
+from __future__ import annotations
+
+import logging
+import re
+from pathlib import Path
+from typing import Any, Dict, List, Optional, TYPE_CHECKING
+
+from memnex.models import (
+    SearchResult,
+    SourceType,
+    WikiPage,
+)
+
+if TYPE_CHECKING:
+    from memnex.retrieval.embedding import EmbeddingService
+
+logger = logging.getLogger(__name__)
+
+# Wikilink pattern
+_WIKILINK_RE = re.compile(r"\[\[([^\]]+)\]\]")
+
+
+class DualIndexSearch:
+    """Hybrid FTS + vector search over wiki pages.
+
+    Parameters
+    ----------
+    wiki_dir:
+        Root directory for wiki files.
+    embedding_service:
+        EmbeddingService for generating and comparing vector embeddings.
+    """
+
+    def __init__(
+        self,
+        wiki_dir: Path,
+        embedding_service: EmbeddingService,
+    ) -> None:
+        self.wiki_dir = wiki_dir
+        self._embedding = embedding_service
+
+        # In-memory FTS index: page_id -> content (lowercased)
+        self._fts_index: Dict[str, str] = {}
+        # In-memory vector index: page_id -> embedding vector
+        self._vector_index: Dict[str, List[float]] = {}
+
+    # ── Public API ────────────────────────────────────────────────────
+
+    def add_page(self, page: WikiPage) -> None:
+        """Add a wiki page to both FTS and vector indices.
+
+        Parameters
+        ----------
+        page:
+            The WikiPage to index.
+        """
+        # FTS index: store lowered content for keyword matching
+        self._fts_index[page.page_id] = page.content.lower()
+
+        # Vector index: embed the page content
+        try:
+            vector = self._embedding.embed(page.content)
+            self._vector_index[page.page_id] = vector
+        except Exception:
+            logger.warning(
+                "Failed to embed page %s, skipping vector index",
+                page.page_id,
+            )
+
+    def search(self, query: str, top_k: int = 10) -> List[SearchResult]:
+        """Execute hybrid FTS + vector search with RRF merging.
+
+        Steps:
+        1. FTS keyword search over indexed page content
+        2. Vector semantic search using embedding similarity
+        3. Merge results with Reciprocal Rank Fusion
+
+        Parameters
+        ----------
+        query:
+            Search query string.
+        top_k:
+            Maximum number of results to return.
+        """
+        # 1. FTS search
+        fts_results = self._fts_search(query)
+
+        # 2. Vector search
+        vector_results = self._vector_search(query)
+
+        # 3. RRF merge
+        merged = self._reciprocal_rank_fusion(fts_results, vector_results)
+
+        return merged[:top_k]
+
+    def rebuild_index(self) -> None:
+        """Rebuild both FTS and vector indices from wiki files on disk.
+
+        Reads all ``*.md`` files from ``wiki_dir/entities/`` and
+        re-indexes them.
+        """
+        self._fts_index.clear()
+        self._vector_index.clear()
+
+        entities_dir = self.wiki_dir / "entities"
+        if not entities_dir.exists():
+            logger.info("No entities directory at %s, index empty", entities_dir)
+            return
+
+        for md_file in entities_dir.glob("*.md"):
+            page_id = md_file.stem
+            content = md_file.read_text(encoding="utf-8")
+            page = WikiPage(page_id=page_id, content=content)
+            self.add_page(page)
+
+        logger.info(
+            "Rebuilt index: %d pages (fts=%d, vector=%d)",
+            len(self._fts_index),
+            len(self._fts_index),
+            len(self._vector_index),
+        )
+
+    # ── RRF merge ─────────────────────────────────────────────────────
+
+    @staticmethod
+    def _reciprocal_rank_fusion(
+        fts_results: List[SearchResult],
+        vector_results: List[SearchResult],
+        k: int = 60,
+    ) -> List[SearchResult]:
+        """Merge FTS and vector results using Reciprocal Rank Fusion.
+
+        RRF score = sum of 1/(k + rank_i) across both result sets.
+
+        Parameters
+        ----------
+        fts_results:
+            Results from full-text search.
+        vector_results:
+            Results from vector similarity search.
+        k:
+            RRF constant (default 60).  Higher k dampens the effect
+            of individual rank positions.
+        """
+        item_map: Dict[str, SearchResult] = {}
+        for item in fts_results + vector_results:
+            item_map[item.func_id] = item
+
+        scores: Dict[str, float] = {}
+        for rank, item in enumerate(fts_results):
+            scores[item.func_id] = scores.get(item.func_id, 0.0) + 1.0 / (k + rank + 1)
+        for rank, item in enumerate(vector_results):
+            scores[item.func_id] = scores.get(item.func_id, 0.0) + 1.0 / (k + rank + 1)
+
+        sorted_ids = sorted(scores.items(), key=lambda x: x[1], reverse=True)
+
+        results: List[SearchResult] = []
+        for page_id, rrf_score in sorted_ids:
+            item = item_map.get(page_id)
+            if item is None:
+                continue
+            # Create a new result with the RRF score
+            results.append(SearchResult(
+                func_id=item.func_id,
+                name=item.name,
+                domain=item.domain,
+                relevance_score=rrf_score,
+                summary=item.summary,
+                source_type=item.source_type,
+                created_at=item.created_at,
+                updated_at=item.updated_at,
+            ))
+        return results
+
+    # ── Private: FTS search ───────────────────────────────────────────
+
+    def _fts_search(self, query: str) -> List[SearchResult]:
+        """Keyword search over indexed page content.
+
+        Simple TF-based scoring: count of query word occurrences.
+        """
+        query_lower = query.lower()
+        query_terms = set(query_lower.split())
+
+        results: List[SearchResult] = []
+        for page_id, content in self._fts_index.items():
+            score = 0.0
+            for term in query_terms:
+                count = content.count(term)
+                if count > 0:
+                    score += min(count / 10.0, 1.0)
+            if score > 0:
+                # Extract a brief summary from content
+                summary = self._extract_summary(content)
+                results.append(SearchResult(
+                    func_id=page_id,
+                    name=page_id,
+                    domain="",
+                    relevance_score=score,
+                    summary=summary,
+                    source_type=SourceType.WIKI,
+                ))
+
+        results.sort(key=lambda r: r.relevance_score, reverse=True)
+        return results
+
+    # ── Private: vector search ────────────────────────────────────────
+
+    def _vector_search(self, query: str) -> List[SearchResult]:
+        """Semantic search using embedding cosine similarity."""
+        if not self._vector_index:
+            return []
+
+        try:
+            query_vector = self._embedding.embed(query)
+        except Exception:
+            logger.warning("Failed to embed query, skipping vector search")
+            return []
+
+        results: List[SearchResult] = []
+        for page_id, doc_vector in self._vector_index.items():
+            similarity = self._cosine_similarity(query_vector, doc_vector)
+            if similarity > 0.1:  # Minimum relevance threshold
+                content = self._fts_index.get(page_id, "")
+                summary = self._extract_summary(content)
+                results.append(SearchResult(
+                    func_id=page_id,
+                    name=page_id,
+                    domain="",
+                    relevance_score=float(similarity),
+                    summary=summary,
+                    source_type=SourceType.WIKI,
+                ))
+
+        results.sort(key=lambda r: r.relevance_score, reverse=True)
+        return results
+
+    # ── Utility helpers ───────────────────────────────────────────────
+
+    @staticmethod
+    def _cosine_similarity(a: List[float], b: List[float]) -> float:
+        """Compute cosine similarity between two vectors."""
+        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
+        if norm_a == 0 or norm_b == 0:
+            return 0.0
+        return dot / (norm_a * norm_b)
+
+    @staticmethod
+    def _extract_summary(content: str, max_len: int = 200) -> str:
+        """Extract a brief summary from page content.
+
+        Skips frontmatter and headings, returns the first substantive line.
+        """
+        in_frontmatter = False
+        frontmatter_dashes = 0
+        for line in content.splitlines():
+            stripped = line.strip()
+            if stripped == "---":
+                frontmatter_dashes += 1
+                if frontmatter_dashes == 1:
+                    in_frontmatter = True
+                    continue
+                elif frontmatter_dashes == 2:
+                    in_frontmatter = False
+                    continue
+            if in_frontmatter:
+                continue
+            if stripped and not stripped.startswith("#") and not stripped.startswith("|"):
+                return stripped[:max_len]
+        return ""