deeptutor-knowledge 1.7.0__py3-none-any.whl

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

@@ -0,0 +1,209 @@
+"""What a connected IMA library actually contains — its document inventory.
+
+Retrieval answers "what does the material say"; it cannot answer "what is in
+here" or "did I add X". For an indexed KB the answer is a directory walk. For an
+IMA library the equivalent is ``get_knowledge_list``, which browses the library's
+folder tree, and this module turns that into the same flat list of relative paths
+:mod:`deeptutor.knowledge.manifest` produces for a local KB — so "list what's in
+this knowledge base" works the same whether the documents sit on disk or in IMA.
+
+Two properties make it usable from the manifest layer:
+
+* **Blocking.** The manifest is deliberately synchronous (it is called from a
+  worker thread, and from the chat pipeline's prompt assembly), so the traversal
+  uses the transport's blocking flavour instead of forcing that whole path async.
+* **Cached.** A manifest is rebuilt on every turn's system prompt; without a
+  cache that would be a fresh IMA round-trip per turn. Entries expire after
+  :data:`CACHE_TTL_SECONDS`, which is short enough that a document added in IMA
+  shows up almost immediately and long enough that a burst of turns costs one
+  traversal.
+
+The traversal is bounded on purpose (:data:`MAX_REQUESTS`): a huge library must
+not stall prompt assembly. When the budget runs out the result is marked
+incomplete and the count it reports is a lower bound, never a confident total.
+"""
+
+from __future__ import annotations
+
+from collections import deque
+from dataclasses import dataclass
+import logging
+import threading
+import time
+from typing import Any, Mapping, Optional
+
+from .config import ImaNotConfiguredError, resolve_kb_config
+
+logger = logging.getLogger(__name__)
+
+# Requests one traversal may spend. With IMA's 50-item pages this covers the
+# first few hundred documents, including a level or two of folders.
+MAX_REQUESTS = 8
+
+PAGE_SIZE = 50
+
+CACHE_TTL_SECONDS = 60.0
+
+# Failures are cached too, briefly. Without this, an unreachable library would be
+# retried on every single turn's prompt assembly — paying the timeout each time.
+FAILURE_TTL_SECONDS = 30.0
+
+# This runs while a turn's system prompt is being assembled, so it must never be
+# what makes a turn feel slow: a stalled IMA gives up quickly rather than holding
+# the default 30 s.
+TIMEOUT_SECONDS = 6.0
+
+# Folder nesting to descend. Deep trees are truncated rather than walked
+# exhaustively — the inventory is an overview, not a mirror.
+MAX_DEPTH = 3
+
+
+@dataclass(frozen=True, slots=True)
+class ImaInventory:
+    """A connected IMA library's documents, as folder-relative paths."""
+
+    documents: tuple[str, ...] = ()
+    complete: bool = True
+    """False when the request budget ran out — ``documents`` is then a prefix."""
+
+
+# key -> (expires_at, inventory or None). ``None`` is a cached *failure*.
+_CACHE: dict[str, tuple[float, Optional[ImaInventory]]] = {}
+_CACHE_LOCK = threading.Lock()
+
+# Distinguishes "nothing cached" from "a cached failure", which is itself a
+# ``None`` result worth honouring.
+_MISS = object()
+
+
+def read_inventory(
+    entry: Mapping[str, Any],
+    *,
+    client_factory=None,
+    use_cache: bool = True,
+) -> Optional[ImaInventory]:
+    """Return the inventory of the IMA library *entry* points at.
+
+    ``None`` means "cannot be determined" (missing credentials, or IMA
+    unreachable) — distinct from an empty library, which is an empty inventory.
+    ``client_factory`` (config → client) is the test seam.
+    """
+    try:
+        config = resolve_kb_config(dict(entry))
+    except ImaNotConfiguredError:
+        return None
+
+    cache_key = f"{config.client_id}:{config.knowledge_base_id}"
+    if use_cache:
+        cached = _cached(cache_key)
+        if cached is not _MISS:
+            return cached  # type: ignore[return-value]
+
+    if client_factory is not None:
+        client = client_factory(config)
+    else:
+        from .client import ImaClient
+
+        client = ImaClient(config, timeout=TIMEOUT_SECONDS)
+
+    try:
+        inventory = _traverse(client, config.knowledge_base_id)
+    except Exception as exc:
+        logger.warning(
+            "Could not read the IMA document inventory for '%s' (%s)",
+            config.knowledge_base_id,
+            type(exc).__name__,
+        )
+        if use_cache:
+            _store(cache_key, None, ttl=FAILURE_TTL_SECONDS)
+        return None
+
+    if use_cache:
+        _store(cache_key, inventory, ttl=CACHE_TTL_SECONDS)
+    return inventory
+
+
+def clear_cache() -> None:
+    """Drop every cached inventory (used by tests and after a KB is re-bound)."""
+    with _CACHE_LOCK:
+        _CACHE.clear()
+
+
+def _traverse(client, root_id: str) -> ImaInventory:
+    """Breadth-first walk of the library, folder paths flattened like a local KB.
+
+    ``truncated`` records whether anything was left unvisited — a folder still
+    queued, a page still to fetch, or a subtree below :data:`MAX_DEPTH`. It is
+    tracked explicitly rather than inferred from the request count so a traversal
+    that happens to finish on its last allowed request is still reported as
+    complete.
+    """
+    documents: list[str] = []
+    seen_folders: set[str] = set()
+    # (folder_id, path prefix, depth); the root folder's id is the library id.
+    queue: deque[tuple[str, str, int]] = deque([(root_id, "", 0)])
+    requests = 0
+    truncated = False
+
+    while queue:
+        folder_id, prefix, depth = queue.popleft()
+        cursor = ""
+        while True:
+            if requests >= MAX_REQUESTS:
+                truncated = True
+                break
+            page = client.get_knowledge_list_sync(
+                folder_id=folder_id,
+                cursor=cursor,
+                limit=PAGE_SIZE,
+            )
+            requests += 1
+            for document in page.documents:
+                documents.append(f"{prefix}{document.title}" if prefix else document.title)
+            for folder in page.folders:
+                if folder.folder_id in seen_folders or folder.folder_id == folder_id:
+                    continue
+                seen_folders.add(folder.folder_id)
+                if depth >= MAX_DEPTH:
+                    truncated = True
+                    continue
+                queue.append((folder.folder_id, f"{prefix}{folder.name}/", depth + 1))
+            cursor = page.next_cursor
+            if page.is_end or not cursor:
+                break
+        if truncated and requests >= MAX_REQUESTS:
+            break
+
+    return ImaInventory(documents=tuple(documents), complete=not truncated and not queue)
+
+
+def _cached(key: str) -> Any:
+    """The cached value, or :data:`_MISS` when absent or expired."""
+    now = time.monotonic()
+    with _CACHE_LOCK:
+        hit = _CACHE.get(key)
+        if hit is None:
+            return _MISS
+        expires_at, inventory = hit
+        if now >= expires_at:
+            _CACHE.pop(key, None)
+            return _MISS
+        return inventory
+
+
+def _store(key: str, inventory: Optional[ImaInventory], *, ttl: float) -> None:
+    with _CACHE_LOCK:
+        _CACHE[key] = (time.monotonic() + ttl, inventory)
+
+
+__all__ = [
+    "CACHE_TTL_SECONDS",
+    "FAILURE_TTL_SECONDS",
+    "MAX_DEPTH",
+    "MAX_REQUESTS",
+    "PAGE_SIZE",
+    "TIMEOUT_SECONDS",
+    "ImaInventory",
+    "clear_cache",
+    "read_inventory",
+]

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

@@ -0,0 +1,190 @@
+"""Fetching and reading one IMA item's actual content.
+
+An IMA knowledge item is either a *note* (whose text the notes API returns
+directly) or a *file* (which IMA hands over as a short-lived Tencent COS
+download URL). This module owns the file half — the download and the text
+extraction — so both retrieval (``pipeline``) and the capability's ``ima_read``
+tool read documents through exactly one implementation.
+
+Security boundary
+-----------------
+The download URL arrives inside an API response, so it is treated as untrusted
+input: it must be HTTPS and served from a trusted Tencent domain (Tencent COS
+or IMA's own resource host, :data:`_ALLOWED_MEDIA_ROOT_DOMAINS`), which
+prevents a tampered or buggy response from turning retrieval into an SSRF
+primitive. Redirects are not followed, IMA's own credentials are never attached
+to this separate client, and hop-by-hop / identity headers are stripped from
+whatever header set the response asked us to send. Size is capped
+(:data:`MAX_MEDIA_BYTES`) both by the advertised ``content-length`` and while
+streaming, so a mis-sized response cannot exhaust memory.
+"""
+
+from __future__ import annotations
+
+import asyncio
+from dataclasses import dataclass
+from pathlib import PurePath, PurePosixPath
+from typing import Optional
+from urllib.parse import unquote, urlparse
+
+import httpx
+
+from deeptutor.utils.document_extractor import (
+    SUPPORTED_DOC_EXTENSIONS,
+    extract_text_from_bytes,
+)
+
+from .envelope import ImaAPIError
+
+# Official IMA media links are short-lived Tencent COS download URLs, and
+# note/file content is served from IMA's exact resource host
+# (res-pkb.ima.qq.com). Do not trust every sibling under ima.qq.com: a broad
+# suffix allowlist would let an unrelated or later-compromised service become
+# an SSRF relay.
+_ALLOWED_MEDIA_ROOT_DOMAINS = ("myqcloud.com",)
+_ALLOWED_MEDIA_HOSTS = frozenset({"res-pkb.ima.qq.com"})
+
+# Headers we never forward, whatever the API response asks for: hop-by-hop
+# fields and anything that would carry identity to COS.
+_FORBIDDEN_MEDIA_HEADERS = frozenset(
+    {
+        "connection",
+        "content-length",
+        "cookie",
+        "host",
+        "proxy-authorization",
+        "transfer-encoding",
+    }
+)
+
+MAX_MEDIA_BYTES = 20 * 1024 * 1024
+
+_CONTENT_TYPE_EXTENSIONS: dict[str, str] = {
+    "application/pdf": ".pdf",
+    "application/vnd.openxmlformats-officedocument.presentationml.presentation": ".pptx",
+    "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet": ".xlsx",
+    "application/vnd.openxmlformats-officedocument.wordprocessingml.document": ".docx",
+    "application/json": ".json",
+    "text/csv": ".csv",
+    "text/html": ".html",
+    "text/markdown": ".md",
+    "text/plain": ".txt",
+}
+
+
+@dataclass(frozen=True)
+class ImaMediaContent:
+    """One IMA item's content, as either note text or downloaded file bytes."""
+
+    text: str = ""
+    data: bytes = b""
+    filename: str = ""
+
+
+async def download_media(
+    url: str,
+    *,
+    headers: dict[str, str],
+    timeout: float,
+    transport: Optional[httpx.AsyncBaseTransport] = None,
+) -> ImaMediaContent:
+    """Stream a COS media URL into memory, bounded by :data:`MAX_MEDIA_BYTES`."""
+    validate_media_url(url)
+    async with httpx.AsyncClient(
+        timeout=timeout,
+        transport=transport,
+        follow_redirects=False,
+    ) as client:
+        async with client.stream("GET", url, headers=headers) as response:
+            response.raise_for_status()
+            length = response.headers.get("content-length")
+            if length and length.isdigit() and int(length) > MAX_MEDIA_BYTES:
+                raise ImaAPIError("IMA media exceeds the 20 MB retrieval limit.")
+            body = bytearray()
+            async for chunk in response.aiter_bytes():
+                body.extend(chunk)
+                if len(body) > MAX_MEDIA_BYTES:
+                    raise ImaAPIError("IMA media exceeds the 20 MB retrieval limit.")
+            filename = media_filename(url, response.headers.get("content-type"))
+            return ImaMediaContent(data=bytes(body), filename=filename)
+
+
+async def extract_text(
+    media: ImaMediaContent | None,
+    *title_candidates: str,
+    max_chars: int,
+) -> str:
+    """Plain text for *media*, truncated to *max_chars* (``""`` when unreadable).
+
+    Note text is already plain. File bytes are decoded by
+    :func:`~deeptutor.utils.document_extractor.extract_text_from_bytes`, which
+    needs a filename to pick a decoder — the download's own filename or the
+    item title, whichever carries a supported extension. Extraction is CPU-bound
+    (PDF / Office parsing), so it runs in a worker thread.
+    """
+    if media is None:
+        return ""
+    if media.text:
+        return media.text[:max_chars].strip()
+    if not media.data:
+        return ""
+    filename = _extractable_filename(media.filename, *title_candidates)
+    if not filename:
+        return ""
+    return await asyncio.to_thread(
+        extract_text_from_bytes,
+        filename,
+        media.data,
+        max_bytes=MAX_MEDIA_BYTES,
+        max_chars=max_chars,
+    )
+
+
+def validate_media_url(url: str) -> None:
+    parsed = urlparse(url)
+    hostname = (parsed.hostname or "").rstrip(".").lower()
+    if parsed.scheme != "https" or not hostname:
+        raise ImaAPIError("IMA media URL must use HTTPS.")
+    trusted_root = any(
+        hostname == root or hostname.endswith(f".{root}") for root in _ALLOWED_MEDIA_ROOT_DOMAINS
+    )
+    if hostname not in _ALLOWED_MEDIA_HOSTS and not trusted_root:
+        raise ImaAPIError("IMA media URL is outside the trusted Tencent media hosts.")
+
+
+def media_headers(raw: object) -> dict[str, str]:
+    """The subset of a response's suggested download headers we will forward."""
+    if not isinstance(raw, dict):
+        return {}
+    return {
+        str(key): str(value)
+        for key, value in raw.items()
+        if str(key).lower() not in _FORBIDDEN_MEDIA_HEADERS and isinstance(value, (str, int, float))
+    }
+
+
+def media_filename(url: str, content_type: str | None) -> str:
+    """A filename for the download, inferring an extension when the path lacks one."""
+    name = unquote(PurePosixPath(urlparse(url).path).name).strip()
+    if "." in name:
+        return name
+    media_type = str(content_type or "").partition(";")[0].strip().lower()
+    return f"{name or 'ima-document'}{_CONTENT_TYPE_EXTENSIONS.get(media_type, '')}"
+
+
+def _extractable_filename(*candidates: str) -> str:
+    for candidate in candidates:
+        if candidate and PurePath(candidate).suffix.lower() in SUPPORTED_DOC_EXTENSIONS:
+            return candidate
+    return ""
+
+
+__all__ = [
+    "MAX_MEDIA_BYTES",
+    "ImaMediaContent",
+    "download_media",
+    "extract_text",
+    "media_filename",
+    "media_headers",
+    "validate_media_url",
+]