memuron 0.1.1__py3-none-any.whl

memuron/ingest/guardian.py

@@ -0,0 +1,213 @@
+"""OpenRouter Guardian for memory ingest writes."""
+
+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.ingest.prompts import MEMORY_WRITE_GUARDIAN_PROMPT, GuardianWritePlan
+
+logger = logging.getLogger(__name__)
+
+OPENROUTER_CHAT_URL = "https://openrouter.ai/api/v1/chat/completions"
+
+
+class GuardianError(RuntimeError):
+    """Raised when the Guardian LLM call fails after retries."""
+
+
+class AgnoGuardian:
+    """Mandatory LLM Guardian: one atomic write plan per ingest."""
+
+    def __init__(
+        self,
+        *,
+        api_key: str | None = None,
+        model_id: str | None = None,
+        max_retries: int | None = None,
+        max_tokens: int | None = None,
+        timeout_seconds: int | None = None,
+    ) -> None:
+        self.api_key = api_key or settings.openrouter_api_key
+        self.model_id = model_id or settings.guardian_model
+        self.max_retries = max_retries if max_retries is not None else settings.guardian_retries
+        self.max_tokens = max_tokens if max_tokens is not None else settings.guardian_max_tokens
+        self.timeout_seconds = (
+            timeout_seconds
+            if timeout_seconds is not None
+            else settings.guardian_timeout_seconds
+        )
+
+    @staticmethod
+    def _parse_write_plan_content(content: Any) -> GuardianWritePlan:
+        if isinstance(content, GuardianWritePlan):
+            return content
+        if isinstance(content, dict):
+            return GuardianWritePlan.model_validate(content)
+        if isinstance(content, str):
+            text = content.strip()
+            if not text:
+                raise GuardianError("Guardian returned empty content")
+            if text.startswith("```"):
+                text = re.sub(r"^```(?:json)?\s*", "", text)
+                text = re.sub(r"\s*```$", "", text)
+            payload = json.loads(text)
+            if not isinstance(payload, dict):
+                raise GuardianError(
+                    f"Guardian JSON payload must be an object, got {type(payload)!r}"
+                )
+            return GuardianWritePlan.model_validate(payload)
+        raise GuardianError(f"Unexpected Guardian content type: {type(content)!r}")
+
+    @staticmethod
+    def _format_candidates(candidates: list[tuple[str, float, dict[str, Any]]]) -> str:
+        if not candidates:
+            return "(no candidates)"
+        lines: list[str] = []
+        for index, (memory_id, score, memory) in enumerate(candidates, start=1):
+            content = str(memory.get("content", ""))[:400]
+            scope = memory.get("scope") or []
+            node_type = memory.get("node_type", "text")
+            lines.append(
+                f"\n--- Candidate {index} ---\n"
+                f"ID: {memory_id}\n"
+                f"Node type: {node_type}\n"
+                f"Similarity Score: {score:.3f}\n"
+                f"Content: {content}\n"
+                f"Scope: {scope}\n"
+            )
+        return "".join(lines)
+
+    def _build_prompt(
+        self,
+        new_content: str,
+        candidates: list[tuple[str, float, dict[str, Any]]],
+        *,
+        space_context: dict[str, str] | None = None,
+    ) -> str:
+        from memuron.ingest.prompts import SPACE_GUARDIAN_SECTION
+
+        prompt = MEMORY_WRITE_GUARDIAN_PROMPT.format(
+            new_content=new_content,
+            candidates_text=self._format_candidates(candidates),
+        )
+        if space_context:
+            prompt = (
+                prompt
+                + "\n"
+                + SPACE_GUARDIAN_SECTION.format(
+                    active_space_name=space_context["active_space_name"],
+                    active_space_token=space_context["active_space_token"],
+                    active_space_prompt=space_context["active_space_prompt"],
+                    other_spaces_summary=space_context["other_spaces_summary"],
+                    registered_tokens=space_context["registered_tokens"],
+                    all_spaces_detail=space_context.get("all_spaces_detail", "(none)"),
+                    space_mode=space_context.get("space_mode", "assist"),
+                )
+            )
+        return prompt
+
+    def _call_guardian(self, prompt: str) -> GuardianWritePlan:
+        if not self.api_key:
+            raise GuardianError("OPENROUTER_API_KEY is required for Guardian writes")
+        payload = {
+            "model": self.model_id,
+            "messages": [{"role": "user", "content": prompt}],
+            "response_format": {"type": "json_object"},
+            "max_tokens": self.max_tokens,
+        }
+        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 GuardianError(f"Guardian request failed: {exc}") from exc
+        if response.status_code >= 400:
+            raise GuardianError(
+                f"Guardian HTTP {response.status_code}: {response.text[:500]}"
+            )
+        try:
+            body = response.json()
+            choice = body["choices"][0]
+            message = choice["message"]
+            raw_content = message.get("content")
+            finish_reason = choice.get("finish_reason")
+        except (KeyError, IndexError, TypeError, json.JSONDecodeError) as exc:
+            raise GuardianError(f"Guardian response parse error: {exc}") from exc
+        if finish_reason == "length":
+            raise GuardianError(
+                "Guardian response truncated at max_tokens="
+                f"{self.max_tokens}; increase MEMURON_GUARDIAN_MAX_TOKENS"
+            )
+        if not isinstance(raw_content, str) or not raw_content.strip():
+            raise GuardianError("Guardian returned empty content")
+        try:
+            return self._parse_write_plan_content(raw_content)
+        except json.JSONDecodeError as exc:
+            raise GuardianError(f"Guardian response JSON parse error: {exc}") from exc
+
+    async def _call_guardian_async(
+        self,
+        new_content: str,
+        candidates: list[tuple[str, float, dict[str, Any]]],
+        *,
+        space_context: dict[str, str] | None = None,
+    ) -> GuardianWritePlan:
+        prompt = self._build_prompt(
+            new_content,
+            candidates,
+            space_context=space_context,
+        )
+        return await asyncio.to_thread(self._call_guardian, prompt)
+
+    async def plan_write(
+        self,
+        new_content: str,
+        candidates: list[tuple[str, float, dict[str, Any]]],
+        *,
+        space_context: dict[str, str] | None = None,
+    ) -> GuardianWritePlan:
+        last_error: Exception | None = None
+        for attempt in range(1, self.max_retries + 1):
+            try:
+                plan = await self._call_guardian_async(
+                    new_content,
+                    candidates,
+                    space_context=space_context,
+                )
+                action = plan.action
+                if action == "update":
+                    if not plan.target_memory_id:
+                        raise GuardianError("Guardian update plan missing target_memory_id")
+                else:
+                    plan = plan.model_copy(update={"target_memory_id": None})
+                return plan
+            except Exception as exc:
+                last_error = exc
+                logger.warning(
+                    "Guardian attempt %s/%s failed: %s",
+                    attempt,
+                    self.max_retries,
+                    exc,
+                )
+                if attempt < self.max_retries:
+                    await asyncio.sleep(2 ** (attempt - 1))
+        raise GuardianError(
+            f"Guardian failed after {self.max_retries} attempts: {last_error}"
+        ) from last_error
+
+
+def create_guardian() -> AgnoGuardian:
+    return AgnoGuardian()

memuron/ingest/jobs.py

@@ -0,0 +1,424 @@
+"""Durable async ingest job queue backed by PostgreSQL."""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import logging
+import os
+from contextlib import contextmanager
+from datetime import UTC, datetime, timedelta
+from typing import Any, Iterator
+from uuid import uuid4
+
+import psycopg
+from psycopg.rows import dict_row
+
+from memuron.application.config import settings
+
+logger = logging.getLogger(__name__)
+
+JOB_TABLE = "memuron_ingest_jobs"
+
+
+def _utcnow() -> datetime:
+    return datetime.now(UTC).replace(tzinfo=None)
+
+
+def _serialize_job(row: dict[str, Any]) -> dict[str, Any]:
+    return {
+        "id": str(row["id"]),
+        "status": row["status"],
+        "payload_json": row["payload_json"],
+        "result_json": row.get("result_json"),
+        "error_json": row.get("error_json"),
+        "attempt_count": int(row["attempt_count"]),
+        "max_attempts": int(row["max_attempts"]),
+        "retryable": bool(row["retryable"]),
+        "available_at": row["available_at"].isoformat() if row["available_at"] else None,
+        "locked_at": row["locked_at"].isoformat() if row.get("locked_at") else None,
+        "lease_expires_at": row["lease_expires_at"].isoformat()
+        if row.get("lease_expires_at")
+        else None,
+        "started_at": row["started_at"].isoformat() if row.get("started_at") else None,
+        "completed_at": row["completed_at"].isoformat() if row.get("completed_at") else None,
+        "created_at": row["created_at"].isoformat() if row["created_at"] else None,
+        "updated_at": row["updated_at"].isoformat() if row["updated_at"] else None,
+    }
+
+
+class IngestJobStore:
+    def __init__(self, database_url: str) -> None:
+        if not database_url:
+            raise ValueError("ARTHA_DATABASE_URL is required for ingest jobs")
+        self.database_url = database_url
+        self._init_schema()
+
+    @contextmanager
+    def _connect(self) -> Iterator[psycopg.Connection[Any]]:
+        conn = psycopg.connect(self.database_url, row_factory=dict_row)
+        try:
+            yield conn
+        finally:
+            conn.close()
+
+    def _init_schema(self) -> None:
+        with self._connect() as conn:
+            conn.execute(
+                f"""
+                CREATE TABLE IF NOT EXISTS {JOB_TABLE} (
+                  id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
+                  status TEXT NOT NULL DEFAULT 'queued',
+                  payload_json JSONB NOT NULL,
+                  result_json JSONB,
+                  error_json JSONB,
+                  attempt_count INT NOT NULL DEFAULT 0,
+                  max_attempts INT NOT NULL DEFAULT 3,
+                  retryable BOOLEAN NOT NULL DEFAULT TRUE,
+                  available_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
+                  locked_at TIMESTAMPTZ,
+                  lease_expires_at TIMESTAMPTZ,
+                  started_at TIMESTAMPTZ,
+                  completed_at TIMESTAMPTZ,
+                  created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
+                  updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
+                )
+                """
+            )
+            conn.execute(
+                f"""
+                CREATE INDEX IF NOT EXISTS idx_memuron_ingest_jobs_claim
+                ON {JOB_TABLE} (status, available_at, lease_expires_at)
+                """
+            )
+            conn.commit()
+
+    def create_job(self, payload: dict[str, Any], *, max_attempts: int | None = None) -> dict[str, Any]:
+        job_id = str(uuid4())
+        now = _utcnow()
+        with self._connect() as conn:
+            row = conn.execute(
+                f"""
+                INSERT INTO {JOB_TABLE} (
+                  id, status, payload_json, max_attempts, available_at, created_at, updated_at
+                ) VALUES (%s, 'queued', %s::jsonb, %s, %s, %s, %s)
+                RETURNING *
+                """,
+                (
+                    job_id,
+                    json.dumps(payload),
+                    max_attempts or settings.ingest_job_max_attempts,
+                    now,
+                    now,
+                    now,
+                ),
+            ).fetchone()
+            conn.commit()
+        return _serialize_job(row)
+
+    def get_job(self, job_id: str) -> dict[str, Any] | None:
+        with self._connect() as conn:
+            row = conn.execute(
+                f"SELECT * FROM {JOB_TABLE} WHERE id = %s",
+                (job_id,),
+            ).fetchone()
+        return _serialize_job(row) if row else None
+
+    def claim_next_job(self, lease_seconds: int) -> dict[str, Any] | None:
+        now = _utcnow()
+        with self._connect() as conn:
+            row = conn.execute(
+                f"""
+                SELECT * FROM {JOB_TABLE}
+                WHERE (
+                  status = 'queued' AND available_at <= %s
+                ) OR (
+                  status = 'processing'
+                  AND lease_expires_at IS NOT NULL
+                  AND lease_expires_at <= %s
+                )
+                ORDER BY created_at ASC
+                FOR UPDATE SKIP LOCKED
+                LIMIT 1
+                """,
+                (now, now),
+            ).fetchone()
+            if not row:
+                conn.rollback()
+                return None
+            lease_expires = now + timedelta(seconds=lease_seconds)
+            updated = conn.execute(
+                f"""
+                UPDATE {JOB_TABLE}
+                SET status = 'processing',
+                    attempt_count = attempt_count + 1,
+                    locked_at = %s,
+                    lease_expires_at = %s,
+                    started_at = COALESCE(started_at, %s),
+                    updated_at = %s
+                WHERE id = %s
+                RETURNING *
+                """,
+                (now, lease_expires, now, now, row["id"]),
+            ).fetchone()
+            conn.commit()
+        return _serialize_job(updated) if updated else None
+
+    def complete_job(self, job_id: str, result_json: dict[str, Any]) -> dict[str, Any]:
+        now = _utcnow()
+        with self._connect() as conn:
+            row = conn.execute(
+                f"""
+                UPDATE {JOB_TABLE}
+                SET status = 'completed',
+                    result_json = %s::jsonb,
+                    error_json = NULL,
+                    completed_at = %s,
+                    locked_at = NULL,
+                    lease_expires_at = NULL,
+                    updated_at = %s
+                WHERE id = %s
+                RETURNING *
+                """,
+                (json.dumps(result_json), now, now, job_id),
+            ).fetchone()
+            conn.commit()
+        if row is None:
+            raise KeyError(f"Ingest job not found: {job_id}")
+        return _serialize_job(row)
+
+    def fail_job(
+        self,
+        job_id: str,
+        error_json: dict[str, Any],
+        *,
+        retryable: bool = True,
+        backoff_seconds: int | None = None,
+    ) -> dict[str, Any]:
+        now = _utcnow()
+        with self._connect() as conn:
+            row = conn.execute(
+                f"SELECT * FROM {JOB_TABLE} WHERE id = %s",
+                (job_id,),
+            ).fetchone()
+            if row is None:
+                raise KeyError(f"Ingest job not found: {job_id}")
+            should_retry = retryable and int(row["attempt_count"]) < int(row["max_attempts"])
+            if should_retry:
+                delay = backoff_seconds if backoff_seconds is not None else 2 ** max(
+                    int(row["attempt_count"]) - 1, 0
+                )
+                status = "queued"
+                available_at = now + timedelta(seconds=delay)
+                completed_at = None
+            else:
+                status = "failed"
+                available_at = row["available_at"]
+                completed_at = now
+            updated = conn.execute(
+                f"""
+                UPDATE {JOB_TABLE}
+                SET status = %s,
+                    error_json = %s::jsonb,
+                    retryable = %s,
+                    available_at = %s,
+                    locked_at = NULL,
+                    lease_expires_at = NULL,
+                    completed_at = %s,
+                    updated_at = %s
+                WHERE id = %s
+                RETURNING *
+                """,
+                (
+                    status,
+                    json.dumps(error_json),
+                    retryable,
+                    available_at,
+                    completed_at,
+                    now,
+                    job_id,
+                ),
+            ).fetchone()
+            conn.commit()
+        return _serialize_job(updated)
+
+
+def build_job_payload(
+    content: str,
+    scope: list[str] | None,
+    *,
+    metadata: dict[str, Any] | None = None,
+    event_metadata: dict[str, Any] | None = None,
+    space_context: dict[str, str] | None = None,
+    candidate_scope: list[str] | None = None,
+) -> dict[str, Any]:
+    payload: dict[str, Any] = {
+        "content": content,
+        "scope": scope or [],
+        "metadata": metadata or {},
+        "event_metadata": event_metadata or {},
+        "candidate_scope": candidate_scope or [],
+    }
+    if space_context:
+        payload["space_context"] = space_context
+    return payload
+
+
+def _payload_space_context(payload: dict[str, Any]) -> dict[str, str] | None:
+    space_context = payload.get("space_context")
+    if not isinstance(space_context, dict):
+        return None
+    if not space_context.get("active_space_token"):
+        return None
+    return space_context
+
+
+async def _run_ingest_job(job: dict[str, Any], engine: Any, guardian: Any) -> None:
+    from memuron.ingest.guardian import GuardianError
+    from memuron.memory.recipes import ingest_memory, update_memory
+    from memuron.spaces.service import apply_guardian_space_scope
+    from memuron.security.tenant import normalize_tenant_scope
+
+    store = job["_store"]  # type: ignore[index]
+    payload = job["payload_json"] or {}
+    event_metadata = payload.get("event_metadata") or {}
+    space_context = _payload_space_context(payload)
+    if space_context is None:
+        backup = event_metadata.get("space_context")
+        if isinstance(backup, dict):
+            space_context = _payload_space_context({"space_context": backup})
+    tenant_id = event_metadata.get("tenant_id")
+    if not tenant_id:
+        for token in payload.get("scope") or []:
+            if str(token).startswith("org:"):
+                tenant_id = str(token)[4:]
+                break
+    try:
+        result = await ingest_memory(
+            engine,
+            guardian,
+            content=payload["content"],
+            scope=payload.get("scope"),
+            metadata=payload.get("metadata") or None,
+            event_metadata=event_metadata,
+            space_context=space_context,
+            candidate_scope=payload.get("candidate_scope") or None,
+        )
+        if result.get("memory_id"):
+            memory = result.get("memory") or {}
+            current_scope = list(memory.get("scope") or [])
+            fixed_scope = normalize_tenant_scope(
+                current_scope,
+                str(tenant_id) if tenant_id else None,
+            )
+            if space_context:
+                fixed_scope = apply_guardian_space_scope(
+                    fixed_scope,
+                    space_context=space_context,
+                )
+            if fixed_scope != current_scope:
+                patched = update_memory(
+                    engine,
+                    str(result["memory_id"]),
+                    scope=fixed_scope,
+                    event_metadata=event_metadata,
+                )
+                result["memory"] = patched
+        await asyncio.to_thread(store.complete_job, job["id"], result)
+    except GuardianError as exc:
+        await asyncio.to_thread(
+            store.fail_job,
+            job["id"],
+            {"code": "GUARDIAN_FAILED", "message": str(exc), "retryable": True},
+            retryable=True,
+        )
+    except Exception as exc:
+        logger.exception("Ingest job %s failed", job["id"])
+        await asyncio.to_thread(
+            store.fail_job,
+            job["id"],
+            {"code": "INGEST_JOB_FAILED", "message": str(exc), "retryable": True},
+            retryable=True,
+        )
+
+
+async def _ingest_worker_loop(
+    worker_name: str,
+    stop_event: asyncio.Event,
+    *,
+    engine: Any,
+    guardian: Any,
+    job_store: IngestJobStore,
+) -> None:
+    while not stop_event.is_set():
+        try:
+            job = await asyncio.to_thread(
+                job_store.claim_next_job,
+                settings.ingest_job_lease_seconds,
+            )
+            if not job:
+                try:
+                    await asyncio.wait_for(
+                        stop_event.wait(),
+                        timeout=settings.ingest_job_poll_interval_seconds,
+                    )
+                except TimeoutError:
+                    pass
+                continue
+            logger.info("Worker %s processing ingest job %s", worker_name, job["id"])
+            job["_store"] = job_store
+            await _run_ingest_job(job, engine, guardian)
+        except asyncio.CancelledError:
+            raise
+        except Exception:
+            logger.exception("Worker %s hit an unexpected ingest loop error", worker_name)
+            await asyncio.sleep(settings.ingest_job_poll_interval_seconds)
+
+
+async def start_ingest_workers(app: Any, engine: Any, guardian: Any) -> None:
+    worker_count = int(
+        os.environ.get("MEMURON_INGEST_WORKER_COUNT", str(settings.ingest_worker_count))
+    )
+    if worker_count <= 0:
+        app.state.ingest_worker_tasks = []
+        app.state.ingest_worker_stop_event = None
+        return
+    if not settings.database_url.startswith("postgresql"):
+        logger.info("Skipping ingest workers because database is not PostgreSQL")
+        app.state.ingest_worker_tasks = []
+        app.state.ingest_worker_stop_event = None
+        return
+
+    job_store = getattr(app.state, "ingest_job_store", None)
+    if job_store is None:
+        job_store = IngestJobStore(settings.database_url)
+        app.state.ingest_job_store = job_store
+    stop_event = asyncio.Event()
+    tasks = [
+        asyncio.create_task(
+            _ingest_worker_loop(
+                f"ingest-{index + 1}",
+                stop_event,
+                engine=engine,
+                guardian=guardian,
+                job_store=job_store,
+            )
+        )
+        for index in range(worker_count)
+    ]
+    app.state.ingest_job_store = job_store
+    app.state.ingest_worker_stop_event = stop_event
+    app.state.ingest_worker_tasks = tasks
+    logger.info("Started %s ingest worker(s)", len(tasks))
+
+
+async def stop_ingest_workers(app: Any) -> None:
+    stop_event = getattr(app.state, "ingest_worker_stop_event", None)
+    tasks = getattr(app.state, "ingest_worker_tasks", [])
+    if stop_event:
+        stop_event.set()
+    if tasks:
+        for task in tasks:
+            task.cancel()
+        await asyncio.gather(*tasks, return_exceptions=True)
+        logger.info("Stopped %s ingest worker(s)", len(tasks))