memuron 0.1.1__py3-none-any.whl

memuron/documents/link_guardian.py

@@ -0,0 +1,192 @@
+"""Guardian for batch semantic linking after document ingest."""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import logging
+import re
+from typing import Any
+
+import requests
+
+from memuron.application.config import settings
+from memuron.documents.parser import _image_data_url
+from memuron.ingest.prompts import (
+    DOCUMENT_LINK_GUARDIAN_PROMPT,
+    DocumentLinkGuardianPlan,
+)
+
+logger = logging.getLogger(__name__)
+
+OPENROUTER_CHAT_URL = "https://openrouter.ai/api/v1/chat/completions"
+_MAX_VISION_ATTACHMENTS = 12
+
+
+class DocumentLinkGuardianError(RuntimeError):
+    """Raised when the document-link Guardian call fails after retries."""
+
+
+class DocumentLinkGuardian:
+    """One multimodal Guardian call to link all chunks from a document ingest."""
+
+    def __init__(
+        self,
+        *,
+        api_key: str | None = None,
+        model_id: str | None = None,
+        max_retries: int | None = None,
+        timeout_seconds: int = 120,
+    ) -> None:
+        self.api_key = api_key or settings.openrouter_api_key
+        self.model_id = model_id or settings.document_link_model
+        self.max_retries = (
+            max_retries
+            if max_retries is not None
+            else settings.document_link_retries
+        )
+        self.timeout_seconds = timeout_seconds
+
+    @staticmethod
+    def _format_memories(
+        items: list[tuple[str, dict[str, Any], float | None]],
+        *,
+        title: str,
+    ) -> str:
+        if not items:
+            return f"(no {title})"
+        lines: list[str] = []
+        for index, (memory_id, memory, score) in enumerate(items, start=1):
+            if not isinstance(memory, dict):
+                continue
+            content = str(memory.get("content", ""))[:500]
+            node_type = memory.get("node_type", "text")
+            score_text = (
+                f"Similarity: {score:.3f}\n"
+                if isinstance(score, (int, float))
+                else ""
+            )
+            lines.append(
+                f"\n--- {title} {index} ---\n"
+                f"ID: {memory_id}\n"
+                f"Node type: {node_type}\n"
+                f"{score_text}"
+                f"Content: {content}\n"
+                f"Scope: {memory.get('scope') or []}\n"
+            )
+        return "".join(lines)
+
+    def _build_prompt(
+        self,
+        *,
+        external: list[tuple[str, dict[str, Any], float | None]],
+        in_doc_chunks: list[tuple[str, dict[str, Any], float | None]],
+        in_doc_images: list[tuple[str, dict[str, Any], float | None]],
+    ) -> str:
+        return DOCUMENT_LINK_GUARDIAN_PROMPT.format(
+            chunks_text=self._format_memories(in_doc_chunks, title="Chunk"),
+            images_text=self._format_memories(in_doc_images, title="Image"),
+            external_text=self._format_memories(external, title="External"),
+        )
+
+    @staticmethod
+    def _parse_json_content(raw: str) -> DocumentLinkGuardianPlan:
+        text = raw.strip()
+        if text.startswith("```"):
+            text = re.sub(r"^```(?:json)?\s*", "", text)
+            text = re.sub(r"\s*```$", "", text)
+        payload = json.loads(text)
+        return DocumentLinkGuardianPlan.model_validate(payload)
+
+    def _call_guardian(
+        self,
+        prompt: str,
+        *,
+        vision_attachments: list[tuple[str, bytes]],
+    ) -> DocumentLinkGuardianPlan:
+        if not self.api_key:
+            raise DocumentLinkGuardianError(
+                "OPENROUTER_API_KEY is required for document link Guardian"
+            )
+        content: list[dict[str, Any]] = [{"type": "text", "text": prompt}]
+        for media_type, image_bytes in vision_attachments[:_MAX_VISION_ATTACHMENTS]:
+            content.append(
+                {
+                    "type": "image_url",
+                    "image_url": {"url": _image_data_url(image_bytes, media_type)},
+                }
+            )
+        payload = {
+            "model": self.model_id,
+            "messages": [{"role": "user", "content": content}],
+            "response_format": {"type": "json_object"},
+        }
+        try:
+            response = requests.post(
+                OPENROUTER_CHAT_URL,
+                headers={
+                    "Authorization": f"Bearer {self.api_key}",
+                    "Content-Type": "application/json",
+                },
+                json=payload,
+                timeout=self.timeout_seconds,
+            )
+        except requests.RequestException as exc:
+            raise DocumentLinkGuardianError(
+                f"Document link Guardian request failed: {exc}"
+            ) from exc
+        if response.status_code >= 400:
+            raise DocumentLinkGuardianError(
+                f"Document link Guardian HTTP {response.status_code}: "
+                f"{response.text[:500]}"
+            )
+        try:
+            body = response.json()
+            message = body["choices"][0]["message"]["content"]
+        except (KeyError, IndexError, TypeError, json.JSONDecodeError) as exc:
+            raise DocumentLinkGuardianError(
+                f"Document link Guardian response parse error: {exc}"
+            ) from exc
+        if not isinstance(message, str):
+            raise DocumentLinkGuardianError("Document link Guardian returned non-text content")
+        return self._parse_json_content(message)
+
+    async def plan_document_links(
+        self,
+        *,
+        external: list[tuple[str, dict[str, Any], float | None]],
+        in_doc_chunks: list[tuple[str, dict[str, Any], float | None]],
+        in_doc_images: list[tuple[str, dict[str, Any], float | None]],
+        vision_attachments: list[tuple[str, bytes]] | None = None,
+    ) -> DocumentLinkGuardianPlan:
+        prompt = self._build_prompt(
+            external=external,
+            in_doc_chunks=in_doc_chunks,
+            in_doc_images=in_doc_images,
+        )
+        attachments = list(vision_attachments or [])
+        last_error: Exception | None = None
+        for attempt in range(1, self.max_retries + 1):
+            try:
+                return await asyncio.to_thread(
+                    self._call_guardian,
+                    prompt,
+                    vision_attachments=attachments,
+                )
+            except Exception as exc:
+                last_error = exc
+                logger.warning(
+                    "Document link Guardian attempt %s/%s failed: %s",
+                    attempt,
+                    self.max_retries,
+                    exc,
+                )
+                if attempt < self.max_retries:
+                    await asyncio.sleep(2 ** (attempt - 1))
+        raise DocumentLinkGuardianError(
+            f"Document link Guardian failed after {self.max_retries} attempts: {last_error}"
+        ) from last_error
+
+
+def create_document_link_guardian() -> DocumentLinkGuardian:
+    return DocumentLinkGuardian()

memuron/documents/linking.py

@@ -0,0 +1,292 @@
+"""Post-ingest semantic linking for document chunks."""
+
+from __future__ import annotations
+
+import logging
+from collections import Counter
+from typing import Any
+
+from artha_engine import ArthaEngine
+
+from memuron.application.config import settings
+from memuron.documents.link_guardian import (
+    DocumentLinkGuardian,
+    DocumentLinkGuardianError,
+    create_document_link_guardian,
+)
+from memuron.ingest.prompts import ConnectionSpec, DocumentLinkGuardianPlan, GuardianWritePlan
+from memuron.memory.recipes import (
+    MAX_GUARDIAN_LINKS,
+    _apply_guardian_links,
+    _find_link_candidates,
+    _ingest_scope_context,
+    _require_memory_projections,
+    get_memory,
+    refresh_memory_projections,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def _document_ingest_memory_ids(payload: dict[str, Any]) -> set[str]:
+    ids = {
+        str(payload["collection"]["id"]),
+        str(payload["document"]["id"]),
+    }
+    ids.update(str(item) for item in payload.get("chunk_ids") or [])
+    ids.update(str(item) for item in payload.get("image_ids") or [])
+    return ids
+
+
+def collect_document_external_candidates(
+    engine: ArthaEngine,
+    *,
+    chunk_ids: list[str],
+    exclude_ids: set[str],
+    tenant_id: str | None,
+    top_k_per_chunk: int | None = None,
+    budget: int | None = None,
+) -> list[tuple[str, dict[str, Any], float]]:
+    """Rank external memories by how often they appear in per-chunk top-k retrieval."""
+
+    top_k = top_k_per_chunk or settings.document_chunk_neighbor_top_k
+    max_external = budget or settings.document_external_candidate_budget
+    scores_by_id: dict[str, float] = {}
+    counts: Counter[str] = Counter()
+
+    for chunk_id in chunk_ids:
+        try:
+            chunk = get_memory(engine, chunk_id)
+        except KeyError:
+            continue
+        content = str(chunk.get("content", "")).strip()
+        if not content:
+            continue
+        neighbors = _find_link_candidates(
+            engine,
+            content,
+            tenant_id=tenant_id,
+            exclude_memory_id=chunk_id,
+            top_k=top_k,
+        )
+        for memory_id, score, memory in neighbors:
+            if memory_id in exclude_ids:
+                continue
+            counts[memory_id] += 1
+            scores_by_id[memory_id] = max(scores_by_id.get(memory_id, 0.0), score)
+
+    ranked = sorted(
+        counts.keys(),
+        key=lambda memory_id: (-counts[memory_id], -scores_by_id.get(memory_id, 0.0)),
+    )[:max_external]
+    output: list[tuple[str, dict[str, Any], float]] = []
+    for memory_id in ranked:
+        memory = get_memory(engine, memory_id)
+        output.append((memory_id, memory, scores_by_id[memory_id]))
+    return output
+
+
+def _resolve_chunk_id(raw_id: str, chunk_ids: list[str]) -> str | None:
+    """Map Guardian chunk_id to an ingest chunk (exact or unique prefix match)."""
+    token = raw_id.strip()
+    if not token:
+        return None
+    if token in chunk_ids:
+        return token
+    prefix_matches = [chunk_id for chunk_id in chunk_ids if chunk_id.startswith(token)]
+    if len(prefix_matches) == 1:
+        return prefix_matches[0]
+    suffix_matches = [chunk_id for chunk_id in chunk_ids if token in chunk_id]
+    if len(suffix_matches) == 1:
+        return suffix_matches[0]
+    return None
+
+
+def _normalize_document_connections(
+    connections: list[ConnectionSpec],
+) -> list[ConnectionSpec]:
+    normalized: list[ConnectionSpec] = []
+    for conn in connections:
+        description = conn.description.strip()
+        if len(description) < 20:
+            continue
+        if "?" not in description:
+            description = f"{description.rstrip('.')}? What else connects these memories?"
+        normalized.append(
+            ConnectionSpec(target_id=conn.target_id.strip(), description=description)
+        )
+    return normalized
+
+
+def _validate_document_connections(
+    memory_id: str,
+    scope: list[str],
+    raw_connections: list[ConnectionSpec],
+    allowed_target_ids: set[str],
+    memories_by_id: dict[str, dict[str, Any]],
+) -> list[ConnectionSpec]:
+    validated: list[ConnectionSpec] = []
+    for conn in _normalize_document_connections(raw_connections):
+        target_id = conn.target_id
+        if not target_id or target_id == memory_id or target_id not in allowed_target_ids:
+            continue
+        if not memories_by_id.get(target_id):
+            continue
+        if len(validated) >= MAX_GUARDIAN_LINKS:
+            break
+        validated.append(conn)
+    return validated
+
+
+def _apply_document_link_plan(
+    engine: ArthaEngine,
+    *,
+    plan: DocumentLinkGuardianPlan,
+    chunk_ids: list[str],
+    scope: list[str],
+    allowed_target_ids: set[str],
+    memories_by_id: dict[str, dict[str, Any]],
+    event_metadata: dict[str, object] | None,
+) -> int:
+    linked = 0
+    for assignment in plan.chunk_links:
+        chunk_id = _resolve_chunk_id(assignment.chunk_id, chunk_ids)
+        if not chunk_id or chunk_id not in memories_by_id:
+            continue
+        connections = _validate_document_connections(
+            chunk_id,
+            scope,
+            assignment.connections,
+            allowed_target_ids,
+            memories_by_id,
+        )
+        if not connections:
+            continue
+        write_plan = GuardianWritePlan(
+            action="create",
+            target_memory_id=None,
+            reasoning=plan.reasoning,
+            final_content=str(memories_by_id[chunk_id].get("content", "")),
+            scope=scope,
+            connections=connections,
+            links_to_remove=[],
+        )
+        linked += _apply_guardian_links(
+            engine,
+            chunk_id,
+            scope,
+            write_plan,
+            list(allowed_target_ids),
+            memories_by_id,
+            event_metadata={
+                **(event_metadata or {}),
+                "document_link_guardian": True,
+            },
+        )
+    return linked
+
+
+async def link_document_after_ingest(
+    engine: ArthaEngine,
+    *,
+    ingest_payload: dict[str, Any],
+    scope: list[str] | None = None,
+    event_metadata: dict[str, object] | None = None,
+    space_context: dict[str, str] | None = None,  # noqa: ARG001 — reserved
+    image_attachments: list[dict[str, Any]] | None = None,
+    document_link_guardian: DocumentLinkGuardian | None = None,
+) -> int:
+    """Run batch Guardian linking for document chunks; no-op for standalone images."""
+
+    if ingest_payload.get("source_type") == "image":
+        return 0
+
+    chunk_ids = [str(item) for item in ingest_payload.get("chunk_ids") or []]
+    if not chunk_ids:
+        return 0
+
+    if not settings.openrouter_api_key:
+        logger.info("Skipping document link Guardian: OPENROUTER_API_KEY unset")
+        return 0
+
+    _require_memory_projections(engine)
+    request_scope, tenant_id, _search_scope = _ingest_scope_context(
+        scope=scope,
+        event_metadata=event_metadata,
+        space_context=None,
+        candidate_scope=None,
+    )
+
+    exclude_ids = _document_ingest_memory_ids(ingest_payload)
+    external = collect_document_external_candidates(
+        engine,
+        chunk_ids=chunk_ids,
+        exclude_ids=exclude_ids,
+        tenant_id=tenant_id,
+    )
+    if not external and len(chunk_ids) < 2:
+        return 0
+
+    in_doc_chunks: list[tuple[str, dict[str, Any], float | None]] = []
+    in_doc_images: list[tuple[str, dict[str, Any], float | None]] = []
+    memories_by_id: dict[str, dict[str, Any]] = {}
+
+    for chunk_id in chunk_ids:
+        memory = get_memory(engine, chunk_id)
+        in_doc_chunks.append((chunk_id, memory, None))
+        memories_by_id[chunk_id] = memory
+
+    for image_id in ingest_payload.get("image_ids") or []:
+        image_id = str(image_id)
+        memory = get_memory(engine, image_id)
+        in_doc_images.append((image_id, memory, None))
+        memories_by_id[image_id] = memory
+
+    for memory_id, memory, _score in external:
+        memories_by_id[memory_id] = memory
+
+    allowed_target_ids = set(memories_by_id.keys()) - {str(ingest_payload["collection"]["id"]), str(ingest_payload["document"]["id"])}
+
+    guardian = document_link_guardian or create_document_link_guardian()
+    vision: list[tuple[str, bytes]] = []
+    for item in image_attachments or []:
+        raw = item.get("bytes")
+        media_type = str(item.get("media_type") or "image/png")
+        if isinstance(raw, (bytes, bytearray)) and raw:
+            vision.append((media_type, bytes(raw)))
+
+    try:
+        plan = await guardian.plan_document_links(
+            external=external,
+            in_doc_chunks=in_doc_chunks,
+            in_doc_images=in_doc_images,
+            vision_attachments=vision,
+        )
+    except DocumentLinkGuardianError:
+        logger.warning(
+            "Document link Guardian failed for document_key=%s; keeping placement-only graph",
+            ingest_payload.get("document_key"),
+            exc_info=True,
+        )
+        return 0
+
+    linked = _apply_document_link_plan(
+        engine,
+        plan=plan,
+        chunk_ids=chunk_ids,
+        scope=request_scope,
+        allowed_target_ids=allowed_target_ids,
+        memories_by_id=memories_by_id,
+        event_metadata=event_metadata,
+    )
+    logger.info(
+        "Document link Guardian document_key=%s chunks=%s external=%s proposed=%s linked=%s",
+        ingest_payload.get("document_key"),
+        len(chunk_ids),
+        len(external),
+        sum(len(item.connections) for item in plan.chunk_links),
+        linked,
+    )
+    if linked:
+        refresh_memory_projections(engine)
+    return linked