Blame for code/graph/api/app.py

"""Graph-service extension exposing GRM as a typed agent operation."""

from __future__ import annotations

from collections.abc import Callable, Mapping
from dataclasses import replace

from graph.api.models import (
    CorrectionRequest,
    GraphPathResponse,
    QueryRequest,
    ReasoningRequest,
    ReasoningResponse,
)
from graph.graph import GraphStore
from graph.inference import ReasoningRuntime
from graph.model import GraphModel
from graph.training.calibration import TemperatureCalibrator
from graph.training.corrections import (
    CorrectionRecord,
    CorrectionStore,
    HumanAction,
)
from graph.types import GraphPath, GraphQuery, PathStep

GraphProvider = Callable[[str, str | None], GraphStore]

#: ``(tenant_id, credential) -> bool``.  Returning False denies the request.
TenantAuthorizer = Callable[[str, str | None], bool]


class TenantAccessError(PermissionError):
    """Raised when a caller may not act on the requested tenant scope."""


class ReasoningService:
    def __init__(
        self,
        model: GraphModel,
        graphs: Mapping[str, GraphStore] | None = None,
        calibrator: TemperatureCalibrator | None = None,
        graph_provider: GraphProvider | None = None,
        authorizer: TenantAuthorizer | None = None,
        max_predictions: int = 1000,
    ) -> None:
        self.model = model
        self.graphs = dict(graphs or {})
        self.calibrator = calibrator
        self.graph_provider = graph_provider
        #: When set, every request must be authorised for its tenant scope.  When
        #: unset the service is single-tenant-trusted and must not be exposed to
        #: untrusted callers; ``/health`` reports which mode is active so a
        #: misconfiguration is visible rather than silent.
        self.authorizer = authorizer
        self.runtime = ReasoningRuntime(model, self.graphs, calibrator)
        self.corrections: list[CorrectionRequest] = []
        self.correction_store = CorrectionStore()
        #: Predictions are retained per tenant: a single shared list let any
        #: caller read every other tenant's queries and answers.
        self.predictions: dict[str, list[ReasoningResponse]] = {}
        self.max_predictions = max_predictions

    def authorize(self, tenant_id: str, credential: str | None) -> None:
        """Raise unless ``credential`` may act on ``tenant_id``."""
        if self.authorizer is None:
            return
        if not self.authorizer(tenant_id, credential):
            raise TenantAccessError("not authorised for this graph scope")

    def prediction_log(self, tenant_id: str) -> list[ReasoningResponse]:
        return self.predictions.get(tenant_id, [])

    def register_graph(self, tenant_id: str, graph: GraphStore) -> None:
        if not tenant_id:
            raise ValueError("tenant_id is required")
        self.graphs[tenant_id] = graph
        self.runtime.register_graph(tenant_id, graph)

    def reason(
        self,
        request: ReasoningRequest,
        *,
        credential: str | None = None,
        log: bool = True,
    ) -> ReasoningResponse:
        tenant_id = request.graph_scope.tenant_id
        self.authorize(tenant_id, credential)
        # Use the graph this call resolved rather than re-reading the shared dict,
        # which could have been replaced between the two reads.
        graph = self._resolve_graph(tenant_id, request.graph_scope.graph_version)
        query = _to_query(request.query)
        if not query.start_entity_ids:
            resolved = graph.resolve_entities(query.start_entities, query.text)
            query = replace(query, start_entity_ids=resolved)
        result = self.runtime.reason(tenant_id, query)
        response = ReasoningResponse.model_validate(result.to_dict())
        if log:
            entries = self.predictions.setdefault(tenant_id, [])
            entries.append(response)
            # Bound the log so a long-lived service cannot grow without limit.
            del entries[: max(0, len(entries) - self.max_predictions)]
        return response

    def verify(
        self, request: ReasoningRequest, *, credential: str | None = None
    ) -> ReasoningResponse:
        return self.reason(request, credential=credential)

    def correction(
        self, request: CorrectionRequest, *, credential: str | None = None
    ) -> dict[str, object]:
        tenant_id = request.request.graph_scope.tenant_id
        self.authorize(tenant_id, credential)
        graph = self._resolve_graph(
            tenant_id, request.request.graph_scope.graph_version
        )
        # ``log=False``: recording a correction must not also append a synthetic
        # entry to the tenant's prediction log.
        prediction = self.reason(request.request, credential=credential, log=False)
        if any(
            answer_id not in graph.nodes for answer_id in request.correct_answer_ids
        ):
            raise ValueError(
                "correction contains an answer ID absent from the authoritative graph"
            )
        predicted_paths = tuple(
            _to_graph_path(path)
            for path in prediction.supporting_paths + prediction.conflicting_paths
        )
        if any(
            index < 0 or index >= len(predicted_paths)
            for index in request.correct_path_indices
        ):
            raise ValueError(
                "correction path index is outside the returned prediction paths"
            )
        record = CorrectionRecord(
            correction_id=f"{request.request.graph_scope.tenant_id}:{len(self.corrections) + 1}",
            graph_version=graph.graph_version,
            query=_to_query(request.request.query),
            predicted_answer_ids=tuple(
                answer.entity_id for answer in prediction.answers
            ),
            predicted_paths=predicted_paths,
            human_action=HumanAction(request.human_action),
            correct_answer_ids=request.correct_answer_ids,
            correct_paths=tuple(
                predicted_paths[index] for index in request.correct_path_indices
            ),
            reason=request.reason,
        )
        self.correction_store.append(record, graph)
        self.corrections.append(request)
        return {
            "accepted": True,
            "correction_count": len(self.corrections),
            "correction_id": record.correction_id,
        }

    def _resolve_graph(self, tenant_id: str, graph_version: str | None) -> GraphStore:
        """Resolve a tenant graph, registering a provider result only if valid.

        The version check runs *before* registration.  Registering first meant a
        request for the wrong version still mutated the service's cached graph for
        that tenant and only then failed, so a rejected request changed state.
        """
        if not tenant_id:
            raise KeyError("tenant_id is required")
        graph = self.graphs.get(tenant_id)
        if graph is None and self.graph_provider is not None:
            candidate = self.graph_provider(tenant_id, graph_version)
            if graph_version and candidate.graph_version != graph_version:
                raise ValueError("requested graph version is not available")
            self.register_graph(tenant_id, candidate)
            return candidate
        if graph is None:
            raise KeyError("unknown graph scope")
        if graph_version and graph.graph_version != graph_version:
            raise ValueError("requested graph version is not available")
        return graph


def create_app(service: ReasoningService):
    """Create the FastAPI app; imports stay local so model code remains usable without HTTP extras."""
    try:
        from fastapi import APIRouter, FastAPI, Header, HTTPException
    except ImportError as exc:  # pragma: no cover - exercised in dependency checks
        raise RuntimeError(
            "FastAPI is required to create the graph service; install the API extras"
        ) from exc

    app = FastAPI(title="Graph Service", version=service.model.config.model_version)

    def _credential(authorization: str | None) -> str | None:
        """Extract a bearer credential without logging or echoing it."""
        if not authorization:
            return None
        scheme, _, token = authorization.partition(" ")
        return token.strip() if scheme.lower() == "bearer" and token.strip() else None

    def _fail(exc: Exception) -> HTTPException:
        """Map an internal error to a client status without leaking internals.

        Messages are deliberately generic: the previous handler returned
        ``str(exc)``, which disclosed tenant identifiers and loaded graph versions
        to unauthenticated callers.  The original exception is still chained, so
        server-side logs keep the detail.
        """
        if isinstance(exc, TenantAccessError):
            return HTTPException(status_code=403, detail="forbidden")
        if isinstance(exc, KeyError):
            return HTTPException(status_code=404, detail="graph scope not found")
        return HTTPException(
            status_code=409, detail="request conflicts with graph state"
        )

    router = APIRouter(prefix="/v1/graph")

    @app.get("/health")
    async def health() -> dict[str, object]:
        # Deliberately no tenant list: enumerating tenant identifiers to any
        # caller is a disclosure, and the count is enough for a liveness probe.
        return {
            "status": "ok",
            "model_version": service.model.config.model_version,
            "tenant_count": len(service.graphs),
            "authorization": "enforced"
            if service.authorizer
            else "trusted-single-tenant",
        }

    @router.post("/reason", response_model=ReasoningResponse, operation_id="graph_reason")
    async def reason(
        request: ReasoningRequest, authorization: str | None = Header(default=None)
    ) -> ReasoningResponse:
        try:
            return service.reason(request, credential=_credential(authorization))
        except (TenantAccessError, KeyError, ValueError) as exc:
            raise _fail(exc) from exc

    @router.post("/verify", response_model=ReasoningResponse, operation_id="graph_verify")
    async def verify(
        request: ReasoningRequest, authorization: str | None = Header(default=None)
    ) -> ReasoningResponse:
        try:
            return service.verify(request, credential=_credential(authorization))
        except (TenantAccessError, KeyError, ValueError) as exc:
            raise _fail(exc) from exc

    @router.post(
        "/find_entities",
        response_model=ReasoningResponse,
        operation_id="graph_find_entities",
    )
    async def find_entities(
        request: ReasoningRequest, authorization: str | None = Header(default=None)
    ) -> ReasoningResponse:
        return await reason(
            request.model_copy(
                update={
                    "query": request.query.model_copy(
                        update={"operation": "find_entities", "answer_mode": "entity"}
                    )
                }
            ),
            authorization,
        )

    @router.post(
        "/find_neighbors",
        response_model=ReasoningResponse,
        operation_id="graph_find_neighbors",
    )
    async def find_neighbors(
        request: ReasoningRequest, authorization: str | None = Header(default=None)
    ) -> ReasoningResponse:
        return await reason(
            request.model_copy(
                update={
                    "query": request.query.model_copy(
                        update={
                            "operation": "find_neighbors",
                            "max_hops": min(1, request.query.max_hops),
                        }
                    )
                }
            ),
            authorization,
        )

    @router.post(
        "/find_paths", response_model=ReasoningResponse, operation_id="graph_find_paths"
    )
    async def find_paths(
        request: ReasoningRequest, authorization: str | None = Header(default=None)
    ) -> ReasoningResponse:
        return await reason(
            request.model_copy(
                update={
                    "query": request.query.model_copy(
                        update={"operation": "find_paths", "return_paths": True}
                    )
                }
            ),
            authorization,
        )

    @router.post(
        "/explain", response_model=ReasoningResponse, operation_id="graph_explain"
    )
    async def explain(
        request: ReasoningRequest, authorization: str | None = Header(default=None)
    ) -> ReasoningResponse:
        return await reason(
            request.model_copy(
                update={
                    "query": request.query.model_copy(
                        update={"operation": "explain", "return_paths": True}
                    )
                }
            ),
            authorization,
        )

    @router.post(
        "/predict_links",
        response_model=ReasoningResponse,
        operation_id="graph_predict_links",
    )
    async def predict_links(
        request: ReasoningRequest, authorization: str | None = Header(default=None)
    ) -> ReasoningResponse:
        return await reason(
            request.model_copy(
                update={
                    "query": request.query.model_copy(
                        update={"operation": "predict_links"}
                    )
                }
            ),
            authorization,
        )

    @router.post("/corrections", operation_id="graph_correction")
    async def correction(
        request: CorrectionRequest, authorization: str | None = Header(default=None)
    ) -> dict[str, object]:
        try:
            return service.correction(request, credential=_credential(authorization))
        except (TenantAccessError, KeyError, ValueError) as exc:
            raise _fail(exc) from exc

    @router.get("/models", operation_id="graph_models")
    async def models() -> dict[str, object]:
        return {
            "active_model": service.model.config.model_version,
            "parameter_count": sum(
                parameter.numel() for parameter in service.model.parameters()
            ),
            "calibrated": service.calibrator is not None,
            "graph_scope_count": len(service.graphs),
        }

    @router.get("/predictions", operation_id="graph_predictions")
    async def predictions(
        tenant_id: str, authorization: str | None = Header(default=None)
    ) -> dict[str, object]:
        """Prediction log for one tenant.

        Previously this returned every prediction the service had ever made,
        across all tenants, to any caller.
        """
        try:
            service.authorize(tenant_id, _credential(authorization))
        except TenantAccessError as exc:
            raise _fail(exc) from exc
        entries = service.prediction_log(tenant_id)
        return {
            "tenant_id": tenant_id,
            "count": len(entries),
            "items": [item.model_dump() for item in entries[-100:]],
        }

    @router.get("/paths", operation_id="graph_paths_list")
    async def paths() -> dict[str, str]:
        return {
            "description": "Use POST /paths with return_paths=true; paths are verified graph evidence."
        }

    @router.post("/paths", response_model=ReasoningResponse, operation_id="graph_paths")
    async def paths_query(
        request: ReasoningRequest, authorization: str | None = Header(default=None)
    ) -> ReasoningResponse:
        return await find_paths(request, authorization)

    app.include_router(router)
    return app


def _to_query(query: QueryRequest) -> GraphQuery:
    return GraphQuery(
        text=query.text,
        start_entities=query.start_entities,
        start_entity_ids=query.start_entity_ids,
        allowed_relations=query.allowed_relations,
        required_relations=query.required_relations,
        required_relation_sequence=query.required_relation_sequence,
        forbidden_relations=query.forbidden_relations,
        required_types=query.required_types,
        forbidden_types=query.forbidden_types,
        required_intermediate_types=query.required_intermediate_types,
        max_hops=query.max_hops,
        min_hops=query.min_hops,
        exact_hops=query.exact_hops,
        return_paths=query.return_paths,
        verify=query.verify,
        operation=query.operation,
        answer_mode=query.answer_mode,
        target_entity_ids=query.target_entity_ids,
    )


def _to_graph_path(path: GraphPathResponse) -> GraphPath:
    """Convert the validated wire response back into the domain path type."""
    steps = tuple(
        PathStep(
            source_id=step.source_id,
            relation_id=step.relation_id,
            target_id=step.target_id,
            source_type=step.source_type,
            target_type=step.target_type,
            edge_attributes=dict(step.edge_attributes),
        )
        for step in path.steps
    )
    return GraphPath(
        steps=steps,
        score=float(path.score),
        verified=bool(path.verified),
        query_compatibility=path.query_compatibility,
    )