deeptutor-knowledge 1.7.0__py3-none-any.whl

deeptutor/services/rag/pipelines/graphrag/pipeline.py

@@ -0,0 +1,284 @@
+"""GraphRAG-backed RAG pipeline orchestration.
+
+Implements the same contract as :class:`LlamaIndexPipeline` (see
+``..base.RAGPipeline``) but delegates indexing and retrieval to a local
+microsoft/graphrag project. Each KB owns a self-contained GraphRAG project under
+its ``version-N`` directory (see ``storage``); documents are parsed to text by
+DeepTutor first (see ``ingestion``) so GraphRAG only ever sees ``.txt`` input.
+
+GraphRAG is an optional dependency: every method fails with a clear, actionable
+message when it is not installed instead of an opaque ``ImportError``.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+from pathlib import Path
+import shutil
+import tempfile
+import traceback
+from typing import Any, Dict, List, Optional
+
+from deeptutor.logging import PROCESS_LOG_PRIVATE_ATTR
+from deeptutor.runtime.home import get_runtime_data_root
+from deeptutor.services.rag.index_versioning import (
+    resolve_storage_dir_for_read,
+    resolve_storage_dir_for_rebuild,
+)
+from deeptutor.services.rag.kb_paths import resolve_kb_dir
+
+from . import config as gr_config
+from . import ingestion, storage
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_KB_BASE_DIR = str(get_runtime_data_root() / "knowledge_bases")
+
+
+class GraphRagPipeline:
+    """Index/retrieve KB content via a local microsoft/graphrag project."""
+
+    def __init__(self, kb_base_dir: Optional[str] = None, **_: Any) -> None:
+        self.logger = logging.getLogger(__name__)
+        self.kb_base_dir = kb_base_dir or DEFAULT_KB_BASE_DIR
+
+    # ----- helpers --------------------------------------------------------
+
+    def _ensure_available(self) -> None:
+        if not gr_config.is_graphrag_available():
+            raise gr_config.GraphRagNotAvailableError(
+                "GraphRAG is not installed. Install it with "
+                "`pip install 'deeptutor[graphrag]'` to use GraphRAG knowledge bases."
+            )
+
+    def _resolve_mode(self, kb_name: str, kwargs: dict[str, Any]) -> str:
+        from ..modes import resolve_kb_mode
+
+        return resolve_kb_mode(
+            self.kb_base_dir,
+            kb_name,
+            storage.PROVIDER,
+            explicit=kwargs.get("mode"),
+            supported=gr_config.SUPPORTED_MODES,
+            default=gr_config.DEFAULT_MODE,
+        )
+
+    def _cleanup_failed_version_dir(self, root_dir: Path) -> None:
+        try:
+            if root_dir.is_dir() and not (root_dir / storage.META_FILENAME).exists():
+                shutil.rmtree(root_dir)
+        except Exception as exc:  # pragma: no cover - best-effort
+            self.logger.warning("Could not clean up failed version dir %s: %s", root_dir, exc)
+
+    async def _preflight_settings(self, settings: dict[str, Any]) -> None:
+        """Probe a settings snapshot without touching an existing KB version."""
+        from . import engine
+
+        with tempfile.TemporaryDirectory(prefix="deeptutor-graphrag-preflight-") as temp_dir:
+            probe_root = Path(temp_dir)
+            gr_config.write_settings_payload(probe_root, settings)
+            # Check the cheaper embedding request first so an invalid endpoint
+            # does not trigger a structured completion call unnecessarily.
+            await engine.preflight_embedding(probe_root)
+            await engine.preflight_completion(probe_root)
+
+    # ----- indexing -------------------------------------------------------
+
+    async def initialize(self, kb_name: str, file_paths: List[str], **kwargs) -> bool:
+        self._ensure_available()
+        kb_dir = resolve_kb_dir(self.kb_base_dir, kb_name)
+        root_dir = resolve_storage_dir_for_rebuild(kb_dir, None)
+        self.logger.info(
+            "Initializing KB '%s' with %d file(s) using GraphRAG", kb_name, len(file_paths)
+        )
+        try:
+            gr_config.write_settings(root_dir)
+            source_for_input: dict[str, str] = {}
+            count = await ingestion.prepare_input(
+                file_paths, root_dir, source_for_input=source_for_input
+            )
+            if count == 0:
+                self.logger.error("GraphRAG: no extractable documents for '%s'", kb_name)
+                self._cleanup_failed_version_dir(root_dir)
+                return False
+            await self._build(root_dir, is_update=False)
+            storage.write_meta(root_dir)
+            if indexed_file_callback := kwargs.get("indexed_file_callback"):
+                names = await asyncio.to_thread(storage.indexed_input_names, root_dir)
+                indexed_file_callback(
+                    sorted(source_for_input[name] for name in names if name in source_for_input)
+                )
+            self.logger.info("KB '%s' initialized with GraphRAG (%d docs)", kb_name, count)
+            return True
+        except Exception as exc:
+            self.logger.error("Failed to initialize GraphRAG KB: %s", exc)
+            self.logger.error(
+                traceback.format_exc(),
+                extra={PROCESS_LOG_PRIVATE_ATTR: True},
+            )
+            self._cleanup_failed_version_dir(root_dir)
+            raise
+
+    async def add_documents(self, kb_name: str, file_paths: List[str], **kwargs) -> bool:
+        self._ensure_available()
+        kb_dir = resolve_kb_dir(self.kb_base_dir, kb_name)
+        existing = resolve_storage_dir_for_read(kb_dir, None)
+        from deeptutor.services.rag.embedding_binding import bound_graph_storage_root
+
+        existing = bound_graph_storage_root(kb_dir, storage.PROVIDER, existing)
+        is_update = existing is not None and storage.has_output(existing)
+        root_dir = (
+            existing if existing is not None else resolve_storage_dir_for_rebuild(kb_dir, None)
+        )
+
+        self.logger.info(
+            "Adding %d document(s) to GraphRAG KB '%s' (update=%s)",
+            len(file_paths),
+            kb_name,
+            is_update,
+        )
+        try:
+            # Refresh settings so a changed model/endpoint is picked up.
+            settings = gr_config.build_settings()
+            if is_update:
+                # An update reuses the active version directory. Validate the
+                # exact settings snapshot first so a bad provider configuration
+                # cannot overwrite a working settings.yaml or add input files.
+                await self._preflight_settings(settings)
+            gr_config.write_settings_payload(root_dir, settings)
+            count = await ingestion.prepare_input(file_paths, root_dir)
+            if count == 0:
+                self.logger.warning("GraphRAG: no extractable documents to add for '%s'", kb_name)
+                return False
+            await self._build(
+                root_dir,
+                is_update=is_update,
+                preflight_embedding_model=not is_update,
+            )
+            storage.write_meta(root_dir)
+            self.logger.info("Added %d doc(s) to GraphRAG KB '%s'", count, kb_name)
+            return True
+        except Exception as exc:
+            self.logger.error("Failed to add documents to GraphRAG KB: %s", exc)
+            self.logger.error(
+                traceback.format_exc(),
+                extra={PROCESS_LOG_PRIVATE_ATTR: True},
+            )
+            if not is_update:
+                self._cleanup_failed_version_dir(root_dir)
+            raise
+
+    async def _build(
+        self,
+        root_dir: Path,
+        *,
+        is_update: bool,
+        preflight_embedding_model: bool = True,
+    ) -> None:
+        from . import engine
+
+        await engine.build(
+            root_dir,
+            is_update=is_update,
+            preflight_embedding_model=preflight_embedding_model,
+        )
+
+    # ----- retrieval ------------------------------------------------------
+
+    async def search(self, query: str, kb_name: str, **kwargs) -> Dict[str, Any]:
+        kb_dir = resolve_kb_dir(self.kb_base_dir, kb_name)
+        root_dir = resolve_storage_dir_for_read(kb_dir, None)
+        from deeptutor.services.rag.embedding_binding import bound_graph_storage_root
+
+        root_dir = bound_graph_storage_root(kb_dir, storage.PROVIDER, root_dir)
+
+        if root_dir is None or not storage.has_output(root_dir):
+            return {
+                "query": query,
+                "answer": (
+                    "This GraphRAG knowledge base has no index yet. Add documents before querying."
+                ),
+                "content": "",
+                "sources": [],
+                "provider": storage.PROVIDER,
+                "needs_reindex": True,
+            }
+
+        mode = self._resolve_mode(kb_name, kwargs)
+        try:
+            self._ensure_available()
+            from . import engine
+
+            response, context_data = await engine.search(root_dir, query, mode)
+        except gr_config.GraphRagNotAvailableError as exc:
+            return self._error_result(query, exc, error_type="not_configured")
+        except Exception as exc:
+            self.logger.error("GraphRAG search failed: %s", exc)
+            self.logger.error(
+                traceback.format_exc(),
+                extra={PROCESS_LOG_PRIVATE_ATTR: True},
+            )
+            return self._error_result(query, exc, error_type="retrieval_error")
+
+        return {
+            "query": query,
+            "answer": response,
+            "content": response,
+            "sources": _context_to_sources(context_data),
+            "provider": storage.PROVIDER,
+            "mode": mode,
+        }
+
+    def _error_result(self, query: str, exc: Exception, *, error_type: str) -> Dict[str, Any]:
+        return {
+            "query": query,
+            "answer": str(exc),
+            "content": "",
+            "sources": [],
+            "provider": storage.PROVIDER,
+            "error_type": error_type,
+        }
+
+    # ----- lifecycle ------------------------------------------------------
+
+    async def delete(self, kb_name: str, **kwargs) -> bool:
+        kb_dir = resolve_kb_dir(self.kb_base_dir, kb_name)
+        if kb_dir.exists():
+            shutil.rmtree(kb_dir)
+            self.logger.info("Deleted GraphRAG KB '%s'", kb_name)
+            return True
+        return False
+
+
+def _context_to_sources(context_data: dict[str, Any]) -> list[dict[str, Any]]:
+    """Map GraphRAG context records into DeepTutor's source-citation shape."""
+    sources: list[dict[str, Any]] = []
+    if not isinstance(context_data, dict):
+        return sources
+    # ``sources`` are the text units; ``reports`` are community summaries. Prefer
+    # the most concrete provenance available.
+    for key in ("sources", "reports", "entities"):
+        records = context_data.get(key)
+        if not isinstance(records, list):
+            continue
+        for rec in records:
+            if not isinstance(rec, dict):
+                continue
+            text = str(rec.get("text") or rec.get("content") or rec.get("description") or "")
+            sources.append(
+                {
+                    "title": str(rec.get("title") or rec.get("name") or f"GraphRAG {key}"),
+                    "content": text[:200],
+                    "source": str(rec.get("source") or ""),
+                    "page": "",
+                    "chunk_id": str(rec.get("id") or ""),
+                    "score": rec.get("rank") or rec.get("score") or "",
+                }
+            )
+        if sources:
+            break
+    return sources
+
+
+__all__ = ["GraphRagPipeline"]

deeptutor/services/rag/pipelines/graphrag/provider.py

@@ -0,0 +1,134 @@
+"""Resolve DeepTutor model bindings at the GraphRAG/LiteLLM boundary."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from deeptutor.services.llm.reasoning_params import build_openai_compatible_reasoning_kwargs
+from deeptutor.services.provider_registry import (
+    effective_backend,
+    find_by_name,
+    strip_provider_prefix,
+)
+
+from .errors import GraphRagUnsupportedProviderError
+
+COMPLETION_TYPE = "deeptutor_litellm"
+
+
+def _runtime_binding(llm_cfg: Any) -> str:
+    return str(
+        getattr(llm_cfg, "binding", None) or getattr(llm_cfg, "provider_name", None) or "openai"
+    ).strip()
+
+
+def resolve_completion_provider(llm_cfg: Any) -> str:
+    """Map a resolved DeepTutor provider to the narrow LiteLLM transport GraphRAG needs."""
+    binding = _runtime_binding(llm_cfg)
+    spec = find_by_name(binding)
+    if spec is None:
+        return "openai"
+    backend = effective_backend(spec, getattr(llm_cfg, "api_format", "auto"))
+    if backend == "anthropic":
+        return "anthropic"
+    if backend == "azure_openai":
+        return "azure"
+    if backend in {"openai_codex", "github_copilot"} or spec.is_oauth:
+        raise GraphRagUnsupportedProviderError(
+            "GraphRAG cannot use this OAuth-only model provider. Choose an API-key profile."
+        )
+    if backend == "openai_compat":
+        # DeepSeek's LiteLLM provider owns its parameter compatibility logic.
+        # Other DeepTutor OpenAI-compatible profiles already expose an OpenAI
+        # chat-completions endpoint and are safest on the generic transport.
+        return "deepseek" if spec.name == "deepseek" else "openai"
+    raise GraphRagUnsupportedProviderError(
+        "GraphRAG does not support the selected model provider transport."
+    )
+
+
+def resolve_completion_model(llm_cfg: Any) -> str:
+    """Return the model identifier expected by the selected LiteLLM transport."""
+    model = str(getattr(llm_cfg, "model", "") or "")
+    spec = find_by_name(_runtime_binding(llm_cfg))
+    return strip_provider_prefix(model, spec)
+
+
+def resolve_completion_call_args(llm_cfg: Any) -> dict[str, Any]:
+    """Return the request options shared by GraphRAG probes and indexing."""
+    call_args: dict[str, Any] = {}
+    extra_headers = getattr(llm_cfg, "extra_headers", None)
+    if isinstance(extra_headers, dict) and extra_headers:
+        call_args["extra_headers"] = dict(extra_headers)
+
+    binding = _runtime_binding(llm_cfg)
+    spec = find_by_name(binding)
+    reasoning_effort = getattr(llm_cfg, "reasoning_effort", None)
+    if spec is not None and spec.backend == "openai_compat":
+        call_args.update(
+            build_openai_compatible_reasoning_kwargs(
+                spec=spec,
+                binding=binding,
+                model=resolve_completion_model(llm_cfg),
+                reasoning_effort=reasoning_effort,
+            )
+        )
+    elif reasoning_effort:
+        # Preserve the existing GraphRAG behavior for Anthropic, Azure, and
+        # unknown transports; the OpenAI-compatible normalizer is not valid for
+        # those backends.
+        call_args["reasoning_effort"] = reasoning_effort
+    return call_args
+
+
+def resolve_persisted_completion_provider(model_config: Any) -> str:
+    """Recover a provider for old DeepTutor GraphRAG settings without rewriting them."""
+    current = str(getattr(model_config, "model_provider", "") or "openai")
+    if current != "openai":
+        return current
+
+    model = str(getattr(model_config, "model", "") or "")
+    api_base = str(getattr(model_config, "api_base", "") or "").rstrip("/")
+    try:
+        from deeptutor.services.config import (
+            get_model_catalog_service,
+            resolve_llm_runtime_config,
+        )
+
+        service = get_model_catalog_service()
+        catalog = service.load()
+        profiles = catalog.get("services", {}).get("llm", {}).get("profiles", [])
+        matches: list[tuple[str, str]] = []
+        for profile in profiles:
+            profile_base = str(profile.get("base_url") or "").rstrip("/")
+            if profile_base != api_base:
+                continue
+            for candidate in profile.get("models", []):
+                if str(candidate.get("model") or "") == model:
+                    matches.append((str(profile.get("id") or ""), str(candidate.get("id") or "")))
+        if len(matches) == 1:
+            profile_id, model_id = matches[0]
+            resolved = resolve_llm_runtime_config(
+                catalog=catalog,
+                service=service,
+                llm_selection={"profile_id": profile_id, "model_id": model_id},
+            )
+            return resolve_completion_provider(resolved)
+    except GraphRagUnsupportedProviderError:
+        raise
+    except Exception:
+        # Old settings remain usable even when the catalog is unavailable.
+        pass
+
+    if "api.deepseek.com" in api_base.lower():
+        return "deepseek"
+    return current
+
+
+__all__ = [
+    "COMPLETION_TYPE",
+    "resolve_completion_call_args",
+    "resolve_completion_model",
+    "resolve_completion_provider",
+    "resolve_persisted_completion_provider",
+]

deeptutor/services/rag/pipelines/graphrag/storage.py

@@ -0,0 +1,127 @@
+"""On-disk layout for a GraphRAG-backed knowledge base.
+
+A GraphRAG KB keeps a self-contained project inside the KB's flat ``version-N``
+directory (reused from ``index_versioning`` with a ``None`` signature, exactly
+like the PageIndex pipeline). The version dir doubles as GraphRAG's project
+root::
+
+    <kb_dir>/version-N/
+        settings.yaml        # generated from DeepTutor config (see config.py)
+        input/               # parsed *.txt fed to the indexer
+        output/              # GraphRAG's parquet artefacts + lancedb
+        cache/  logs/
+        meta.json            # synthetic "ready" marker (see write_meta)
+
+The synthetic ``meta.json`` is what makes the existing "is this KB initialised?"
+and Index-versions UI checks treat a GraphRAG KB as ready, without teaching the
+manager about GraphRAG internals.
+"""
+
+from __future__ import annotations
+
+from datetime import datetime, timezone
+import logging
+from pathlib import Path
+
+from deeptutor.services.file_io import atomic_write_json
+
+logger = logging.getLogger(__name__)
+
+META_FILENAME = "meta.json"
+PROVIDER = "graphrag"
+
+INPUT_DIRNAME = "input"
+OUTPUT_DIRNAME = "output"
+
+# Parquet artefacts GraphRAG writes on a successful index; their presence is our
+# "the index actually built" signal (independent of the synthetic meta marker).
+OUTPUT_TABLES = (
+    "entities",
+    "communities",
+    "community_reports",
+    "text_units",
+    "relationships",
+)
+
+
+def input_dir(root_dir: Path) -> Path:
+    return Path(root_dir) / INPUT_DIRNAME
+
+
+def output_dir(root_dir: Path) -> Path:
+    return Path(root_dir) / OUTPUT_DIRNAME
+
+
+def has_output(root_dir: Path | None) -> bool:
+    """True when GraphRAG has produced at least its core parquet tables."""
+    if root_dir is None:
+        return False
+    out = output_dir(root_dir)
+    if not out.is_dir():
+        return False
+    return any((out / f"{name}.parquet").exists() for name in OUTPUT_TABLES)
+
+
+def indexed_input_names(root_dir: Path) -> set[str]:
+    """Input filenames with persisted GraphRAG text units; empty on unknown output."""
+    documents = output_dir(root_dir) / "documents.parquet"
+    text_units = output_dir(root_dir) / "text_units.parquet"
+    if not documents.is_file() or not text_units.is_file():
+        return set()
+    try:
+        import pandas as pd
+
+        rows = pd.read_parquet(documents, columns=["title", "text_unit_ids"])
+        unit_ids = set(pd.read_parquet(text_units, columns=["id"])["id"])
+    except Exception as exc:
+        logger.warning("Could not inspect GraphRAG document receipts: %s", exc)
+        return set()
+
+    names: set[str] = set()
+    for title, text_unit_ids in rows.itertuples(index=False, name=None):
+        if not isinstance(title, str) or isinstance(text_unit_ids, (str, bytes)):
+            continue
+        try:
+            has_persisted_unit = any(unit_id in unit_ids for unit_id in text_unit_ids)
+        except TypeError:
+            has_persisted_unit = False
+        if has_persisted_unit:
+            names.add(Path(title).name)
+    return names
+
+
+def write_meta(root_dir: Path) -> None:
+    """Write a flat-layout ``meta.json`` so the version is listed as ready.
+
+    Mirrors ``index_versioning.write_version_meta`` but carries a synthetic
+    ``graphrag`` signature instead of an embedding hash. The embedding identity
+    is stamped alongside so an externally-linked index can be checked for
+    embedding compatibility at connect time (GraphRAG otherwise fails retrieval
+    silently on a dimension mismatch).
+    """
+    from deeptutor.services.rag.embedding_signature import embedding_meta_fields
+
+    target = Path(root_dir)
+    payload = {
+        "version": target.name,
+        "signature": PROVIDER,
+        "provider": PROVIDER,
+        "layout": "flat",
+        "created_at": datetime.now(timezone.utc).replace(tzinfo=None).isoformat() + "Z",
+        **embedding_meta_fields(),
+    }
+    atomic_write_json(target / META_FILENAME, payload)
+
+
+__all__ = [
+    "META_FILENAME",
+    "PROVIDER",
+    "INPUT_DIRNAME",
+    "OUTPUT_DIRNAME",
+    "OUTPUT_TABLES",
+    "input_dir",
+    "output_dir",
+    "has_output",
+    "indexed_input_names",
+    "write_meta",
+]

deeptutor/services/rag/pipelines/ima/__init__.py

@@ -0,0 +1,28 @@
+"""Tencent IMA integration — retrieval, inventory and authoring over a live library.
+
+A KB bound to the ``ima`` provider is a connection pointer to a library the user
+keeps in IMA (https://ima.qq.com) and curates there. DeepTutor never indexes or
+stores anything locally; everything is an OpenAPI call:
+
+* :mod:`.pipeline` — the ``rag`` engine: ``search_knowledge`` plus a bounded
+  full-text top-up for matches whose snippet is too thin to reason from.
+* :mod:`.inventory` — what the library *contains* (``get_knowledge_list``), which
+  is what makes "list the documents in this knowledge base" answerable; read from
+  the synchronous manifest layer, so it is blocking and cached.
+* :mod:`.notes` — the notes module, including the sort orders that answer "my
+  most recent items" (knowledge items carry no timestamps; notes do).
+* :mod:`.client` — the knowledge-base method table, with :mod:`.transport`,
+  :mod:`.envelope`, :mod:`.models` and :mod:`.media` underneath it.
+* :mod:`.probe` — the connect-time health check.
+
+The IMA credentials are configured once for the account on the engine page (a KB
+may override them to reach another account); the knowledge base id is always
+per-KB. The agentic surface over these calls is
+:mod:`deeptutor.capabilities.ima`.
+"""
+
+from __future__ import annotations
+
+from .pipeline import ImaPipeline
+
+__all__ = ["ImaPipeline"]