sable-sphynx 0.7.0.dev0__py3-none-any.whl

sphynx/integrated_rag.py

@@ -0,0 +1,352 @@
+"""Versioned router/rewrite/retrieval/context/answer execution.
+
+This boundary executes retrieval. It does not reinterpret the native supplied-
+context v1 runner. Component calls are trusted local execution, not OS isolation.
+Lifecycle journals retain stage artifacts, including partial work on failure.
+"""
+from __future__ import annotations
+
+from copy import deepcopy
+from dataclasses import dataclass
+import hashlib
+import json
+import time
+
+from .assessment_contracts import CaseView, _thaw
+from .components import (
+    AdapterDescriptor, EvaluationTarget, OPERATIONS_BY_KIND, PIPELINE_COMPONENT_ORDER,
+    adapter_component_identity_sha256,
+)
+from .integration_contracts import value_digest
+from .lifecycle import (
+    AttemptOutput, DriverDescriptor, DriverFailure, RunPolicy, execute_lifecycle,
+)
+from .pipeline import validate_pipeline_response
+from .pipeline_workflow import (
+    PipelineBackend, PipelineCandidate, _stage_definition, validate_stage_input,
+    validate_stage_output,
+)
+from ._composed_pipeline import compose_pipeline, validate_chunks
+from .workflow_contracts import PayloadSchema, PlanBinding, WorkflowContractError, _text, contract_digest
+
+INTEGRATED_SCHEMA = "sphynx.rag.integrated.response.v1"
+INTEGRATED_STAGE_SCHEMA = "sphynx.rag.integrated.stage.v1"
+INTEGRATED_STAGES = ("router", "rewriter", *PIPELINE_COMPONENT_ORDER, "context", "generator")
+_FULL_TARGET = EvaluationTarget("level3", "full_rag")
+
+
+def _require(value, message):
+    if not value:
+        raise WorkflowContractError(message)
+
+
+def json_bytes(value):
+    return json.dumps(_thaw(value), ensure_ascii=False, allow_nan=False,
+                      sort_keys=True, separators=(",", ":")).encode("utf-8")
+
+
+def text_digest(value):
+    return "sha256:" + hashlib.sha256(value.encode("utf-8")).hexdigest()
+
+
+def validate_extra_output(kind, value):
+    _require(type(value) is dict, "integrated stage output must be an object")
+    if kind == "router":
+        _require(set(value) <= {"should_call_rag", "latency_ms"} and type(value.get("should_call_rag")) is bool,
+                 "router must return a boolean should_call_rag")
+    elif kind == "rewriter":
+        _require("rewritten_query" in value and set(value) <= {"query", "rewritten_query", "latency_ms"},
+                 "rewriter output fields differ")
+        _text(value["rewritten_query"], "rewritten query")
+    elif kind == "generator":
+        _require({"answer", "abstained"} <= set(value)
+                 and set(value) <= {"answer", "abstained", "citations", "token_counts", "latency_ms"},
+                 "generator output fields differ")
+        _require(type(value["answer"]) is str and type(value["abstained"]) is bool,
+                 "generator answer/action differs")
+        _require(type(value.get("citations", [])) is list, "generator citations must be an array")
+        if "token_counts" in value:
+            counts = value["token_counts"]
+            _require(type(counts) is dict and set(counts) == {"prompt", "completion"}
+                     and all(type(n) is int and n >= 0 for n in counts.values()), "token count fields differ")
+    elif kind == "context":
+        _require(type(value.get("rows")) is list and type(value.get("text")) is str
+                 and type(value.get("bytes")) is int, "context output fields differ")
+    else:
+        raise WorkflowContractError("unknown integrated stage")
+
+
+def build_context(response, documents, max_context_bytes):
+    """Select whole source-bound chunks in final-rank order under a UTF-8 budget."""
+    _require(type(max_context_bytes) is int and max_context_bytes > 0, "context byte limit must be positive")
+    chunks = validate_chunks(response["construction"]["chunks"], documents)
+    source = {row["document_id"]: row for row in documents}
+    selected = sorted((row for row in response["candidates"] if row["final_rank"] is not None),
+                      key=lambda row: row["final_rank"])
+    rows, texts, count = [], [], 0
+    for selected_row in selected:
+        chunk = chunks[selected_row["candidate_id"]]
+        text = chunk["input"]
+        extra = len(text.encode("utf-8")) + (2 if texts else 0)
+        if count + extra > max_context_bytes:
+            break
+        document = source[chunk["document_id"]]
+        _require(document.get("content_sha256") == text_digest(document["text"]).removeprefix("sha256:"),
+                 "context source document digest differs")
+        rows.append({"candidate_id": chunk["candidate_id"], "document_id": chunk["document_id"],
+                     "start": chunk["start"], "end": chunk["end"], "input": text,
+                     "text_sha256": text_digest(text), "source_document_sha256": text_digest(document["text"])})
+        texts.append(text)
+        count += extra
+    text = "\n\n".join(texts)
+    return {"rows": rows, "text": text, "bytes": count, "text_sha256": text_digest(text),
+            "selection_policy": "whole_chunks_in_final_rank_order_utf8_bytes.v1"}
+
+
+def validate_answer(output, context):
+    validate_extra_output("generator", output)
+    allowed = {row["candidate_id"] for row in context["rows"]}
+    ids = set()
+    for citation in output.get("citations", []):
+        _require(type(citation) is dict and set(citation) == {"citation_id", "evidence_id", "start", "end"},
+                 "integrated citation fields differ")
+        _text(citation["citation_id"], "citation_id")
+        _require(citation["citation_id"] not in ids and citation["evidence_id"] in allowed,
+                 "citation leaves the exact generator context")
+        ids.add(citation["citation_id"])
+        _require(type(citation["start"]) is int and type(citation["end"]) is int
+                 and 0 <= citation["start"] < citation["end"] <= len(output["answer"]),
+                 "citation offsets leave the answer")
+
+
+@dataclass(frozen=True)
+class IntegratedCandidate:
+    pipeline: PipelineCandidate
+    router_id: str
+    rewriter_id: str
+    generator_id: str
+    revision: str
+
+    def __post_init__(self):
+        _require(isinstance(self.pipeline, PipelineCandidate), "integrated candidate needs PipelineCandidate")
+        for name in ("router_id", "rewriter_id", "generator_id", "revision"):
+            _text(getattr(self, name), name)
+        _require(self.revision == self.pipeline.revision, "component/subsystem/system candidate revisions differ")
+
+
+class IntegratedDriver:
+    """Cooperative lifecycle driver over explicitly registered native components."""
+
+    def __init__(self, pipeline: PipelineBackend, *, router_id, rewriter_id, generator_id,
+                 max_context_bytes=16000):
+        _require(isinstance(pipeline, PipelineBackend), "integrated driver requires PipelineBackend")
+        _require(type(max_context_bytes) is int and max_context_bytes > 0, "context byte limit must be positive")
+        self.pipeline = pipeline
+        self.router_id, self.rewriter_id, self.generator_id = router_id, rewriter_id, generator_id
+        self.max_context_bytes = max_context_bytes
+        self._configuration = {"pipeline_binding": pipeline.describe(pipeline.suite_id).definition_digest,
+                               "router_id": router_id, "rewriter_id": rewriter_id,
+                               "generator_id": generator_id, "max_context_bytes": max_context_bytes}
+        self.descriptor = DriverDescriptor(
+            "sphynx.rag.integrated.driver", "sphynx.rag.integrated.driver.v1",
+            ("sphynx.lifecycle.deadline.cooperative.v1",))
+        self._descriptor_digest = contract_digest(self.descriptor.to_dict())
+        for kind, identifier in (("router", router_id), ("rewriter", rewriter_id), ("generator", generator_id)):
+            self._component(kind, identifier, resolve=False)
+
+    def candidate(self, pipeline_candidate):
+        return IntegratedCandidate(pipeline_candidate, self.router_id, self.rewriter_id,
+                                   self.generator_id, pipeline_candidate.revision)
+
+    def _component(self, kind, identifier, *, resolve):
+        component = self.pipeline.components.by_id.get(identifier)
+        _require(component is not None and component.kind == kind, "integrated component kind/identity differs")
+        _require(_FULL_TARGET in component.capabilities, "integrated component lacks its native capability")
+        _require(component.model is None or not component.model.training_required, "integrated model requires forbidden training")
+        if resolve:
+            _require(component.implementation.kind in self.pipeline.allowed_implementation_kinds,
+                     "integrated component implementation is not authorized")
+            adapter = self.pipeline.adapters.resolve(identifier, _FULL_TARGET, OPERATIONS_BY_KIND[kind][0])
+            _require(adapter.descriptor == AdapterDescriptor.from_component(component), "integrated adapter identity drift")
+        return component
+
+    def candidate_identity(self, candidate):
+        _require(isinstance(candidate, IntegratedCandidate), "candidate must be IntegratedCandidate")
+        self.pipeline.descriptor
+        _require(contract_digest(self.descriptor.to_dict()) == self._descriptor_digest,
+                 "integrated driver descriptor drift")
+        _require(self._configuration == {
+            "pipeline_binding": self.pipeline.describe(self.pipeline.suite_id).definition_digest,
+            "router_id": self.router_id, "rewriter_id": self.rewriter_id,
+            "generator_id": self.generator_id, "max_context_bytes": self.max_context_bytes},
+            "integrated driver configuration drift")
+        native = self.pipeline._candidate(candidate.pipeline)
+        extra = {}
+        for kind in ("router", "rewriter", "generator"):
+            component = self._component(kind, getattr(candidate, kind + "_id"), resolve=False)
+            extra[kind] = {"component_id": component.component_id, "component_version": component.version,
+                           "descriptor_digest": "sha256:" + adapter_component_identity_sha256(
+                               AdapterDescriptor.from_component(component))}
+        return {"schema_version": "sphynx.rag.integrated.candidate.v1",
+                "revision": candidate.revision, "pipeline": native.to_dict(), "components": extra,
+                "context_policy": {"max_bytes": self.max_context_bytes,
+                                   "selection": "whole_chunks_in_final_rank_order_utf8_bytes.v1"},
+                "candidate_source_bytes_verified": False}
+
+    def admit(self, candidate):
+        self.candidate_identity(candidate)
+        plan = PlanBinding(self.pipeline.describe(self.pipeline.suite_id), candidate.pipeline, "subsystem", "local")
+        admitted = self.pipeline.admit(plan, self.pipeline.resolve(plan))
+        for kind in ("router", "rewriter", "generator"):
+            self._component(kind, getattr(candidate, kind + "_id"), resolve=True)
+        return admitted
+
+    def cases(self):
+        return self.pipeline.cases(self.pipeline.describe(self.pipeline.suite_id))
+
+    def prepare(self, candidate, context):
+        admitted = self.admit(candidate)
+        return {"source_digest": self.pipeline.source.sha256,
+                "pipeline_admission_digest": contract_digest(admitted.to_dict()),
+                "case_digests": {case.case_id: contract_digest(case.to_dict()) for case in self.cases()},
+                "context_policy": self.candidate_identity(candidate)["context_policy"]}
+
+    def execute(self, candidate, case: CaseView, context):
+        _require(case.case_id == context.case_id
+                 and context.prepared.get("case_digests", {}).get(case.case_id) == contract_digest(case.to_dict()),
+                 "integrated case differs from the prepared source")
+        identity = self.candidate_identity(candidate)
+        stages, events = {}, []
+        result = {"schema_version": INTEGRATED_SCHEMA, "case_id": case.case_id,
+                  "trial": context.trial, "attempt": context.attempt, "candidate_revision": candidate.revision,
+                  "input_digest": contract_digest(case.to_dict()), "candidate_identity_digest": value_digest(identity),
+                  "stages": stages, "events": events, "retrieval": None, "context": None, "answer": None,
+                  "citations": [], "abstained": None, "actual_rag_call_count": None,
+                  "token_counts": None, "failure_stage": None, "skipped_stages": []}
+
+        def partial(failure, *, acknowledged=False):
+            return AttemptOutput(outputs=result, failure=failure,
+                metadata={"boundary": "integrated_retrieval_context_answer",
+                          "timing_authority": "parent_observed_local_stage_wall",
+                          "cancellation_acknowledged": acknowledged})
+
+        def checkpoint():
+            if context.cancellation.requested:
+                raise DriverFailure("cancelled", partial=partial("cancelled", acknowledged=True))
+            if context.deadline is not None and time.monotonic() >= context.deadline:
+                raise DriverFailure("timeout", partial=partial("timeout"))
+
+        def retain(kind, phase, record):
+            raw = json_bytes(record)
+            binding = context.emit("integrated." + kind + "." + phase, raw)
+            events.append({"sequence": len(events), "stage_id": kind, "occurrence": 1, "phase": phase,
+                           "artifact": dict(binding)})
+            if binding["size_bytes"] != len(raw):
+                raise DriverFailure("output_limit", partial=partial("output_limit"))
+            return dict(binding)
+
+        def call(kind, operation, payload, *, internal=None):
+            checkpoint()
+            record = {"schema_version": INTEGRATED_STAGE_SCHEMA, "run_id": context.run_id,
+                      "case_id": case.case_id, "trial": context.trial, "attempt": context.attempt,
+                      "stage_id": kind, "occurrence": 1, "operation": operation,
+                      "input_digest": value_digest(payload), "input": deepcopy(payload),
+                      "output": None, "output_digest": None, "status": "started"}
+            if kind in PIPELINE_COMPONENT_ORDER:
+                ref = candidate.pipeline.configuration.by_kind[kind]
+                component = self.pipeline.components.by_id[ref.component_id]
+                validate_stage_input(kind, payload)
+                definition = _stage_definition(component).to_dict()
+                target = EvaluationTarget("level2", "pipeline")
+            elif kind != "context":
+                component = self._component(kind, getattr(candidate, kind + "_id"), resolve=True)
+                definition = {"component_id": component.component_id,
+                              "descriptor_digest": "sha256:" + adapter_component_identity_sha256(
+                                  AdapterDescriptor.from_component(component))}
+                target = _FULL_TARGET
+            else:
+                definition = {"implementation_id": "sphynx.integrated.context-builder.v1",
+                              "configuration": identity["context_policy"],
+                              "authority": "package_declared_not_code_attestation"}
+            record["implementation"] = definition
+            retain(kind, "started", record)
+            begin = time.perf_counter()
+            try:
+                output = internal() if internal is not None else self.pipeline.adapters.resolve(
+                    component.component_id, target, operation).invoke(operation, deepcopy(payload))
+                def validator(value):
+                    if kind in PIPELINE_COMPONENT_ORDER:
+                        validate_stage_output(kind, value)
+                    else:
+                        validate_extra_output(kind, value)
+                    if kind == "generator":
+                        validate_answer(value, result["context"])
+                schema = PayloadSchema("sphynx.integrated." + kind + ".output.v1", "v1",
+                                       "sphynx.integrated.stage-validator.v1", validator)
+                output = _thaw(schema.validate(output))
+                record.update(status="completed", output=output, output_digest=value_digest(output),
+                              parent_elapsed_seconds=time.perf_counter() - begin)
+                binding = retain(kind, "completed", record)
+                stages[kind] = {**deepcopy(record), "artifact": binding}
+                checkpoint()
+                return output
+            except DriverFailure:
+                raise
+            except Exception as exc:
+                record.update(status="failed", error_type=type(exc).__name__,
+                              parent_elapsed_seconds=time.perf_counter() - begin)
+                binding = retain(kind, "failed", record)
+                stages[kind] = {**deepcopy(record), "artifact": binding}
+                result["failure_stage"] = kind
+                raise DriverFailure("candidate_error", partial=partial("candidate_error")) from None
+
+        query = case.inputs["query"]
+        routing = call("router", "route", {"task_id": case.case_id, "query": query,
+                                           "conversation_context": _thaw(case.inputs["conversation_context"])})
+        result["routing_decision"] = routing["should_call_rag"]
+        result["actual_rag_call_count"] = 0
+        if not routing["should_call_rag"]:
+            result.update(abstained=True, skipped_stages=list(INTEGRATED_STAGES[1:]))
+            return AttemptOutput(outputs=result)
+        rewriting = call("rewriter", "rewrite", {"task_id": case.case_id, "query": query})
+        rewritten = rewriting["rewritten_query"]
+        result["rewritten_query"] = rewritten
+        native = self.pipeline._candidate(candidate.pipeline)
+        request = {"schema_version": "sphynx.rag.integrated.retrieval-input.v1", "case_id": case.case_id,
+                   "query": rewritten, "documents": _thaw(case.inputs["documents"]),
+                   "policy": {"boundary": "composed_source_documents",
+                              "controlled_config": _thaw(native.parameters["controlled_config"])}}
+        request["input_digest"] = value_digest(request)
+        result["actual_rag_call_count"] = 1
+        try:
+            response = compose_pipeline(request, call)
+            validate_pipeline_response(response, request)
+        except DriverFailure:
+            raise
+        except Exception:
+            result["failure_stage"] = "retrieval_verification"
+            raise DriverFailure("candidate_error", partial=partial("candidate_error")) from None
+        result["retrieval"] = {"request": request, "response": response}
+        context_output = call("context", "construct_context",
+            {"retrieval_input_digest": request["input_digest"], "max_context_bytes": self.max_context_bytes},
+            internal=lambda: build_context(response, request["documents"], self.max_context_bytes))
+        result["context"] = context_output
+        generated = call("generator", "generate", {"task_id": case.case_id, "query": rewritten,
+                                                  "contexts": deepcopy(context_output["rows"])})
+        result.update(answer=generated["answer"], abstained=generated["abstained"],
+                      citations=generated.get("citations", []), token_counts=generated.get("token_counts"))
+        usage = {} if result["token_counts"] is None else {"tokens": sum(result["token_counts"].values())}
+        return AttemptOutput(outputs=result, usage=usage,
+                             metadata={"boundary": "integrated_retrieval_context_answer",
+                                       "timing_authority": "parent_observed_local_stage_wall"})
+
+
+def run_integrated(driver: IntegratedDriver, candidate: IntegratedCandidate, output_dir, *,
+                   policy: RunPolicy | None = None, cancellation=None, resume=None):
+    """Execute the integrated boundary through the versioned lifecycle driver."""
+    driver.admit(candidate)
+    cases = driver.cases()
+    return execute_lifecycle(driver, candidate, cases, output_dir,
+        identity={"suite_binding": driver.pipeline.describe(driver.pipeline.suite_id).definition_digest,
+                  "boundary": "sphynx.rag.integrated.v1", "source": driver.pipeline.source.to_dict()},
+        policy=policy, cancellation=cancellation, resume=resume)