"""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, )