schemarouter 0.2.0a1__py3-none-any.whl

schemarouter/runtime.py

@@ -0,0 +1,686 @@
+from __future__ import annotations
+
+import asyncio
+from collections.abc import AsyncIterator, Awaitable, Callable, Iterator, Sequence
+from typing import Any, TypeVar
+from urllib.parse import urlparse
+from uuid import uuid4
+
+import httpx
+from pydantic import TypeAdapter
+
+from .adapters.base import AdapterRegistry, SourceAdapter
+from .adapters.openapi import OpenAPIRemoteInvoker
+from .adapters.python import PythonCallableInvoker, callable_options, tool_from_callable
+from .errors import ProposalApprovalError, RegistrationError
+from .executor import RegistryExecutor
+from .ingestion import SourceKind, URLSchemaLoader
+from .models import ExecutionPlan, PlanRequest, ToolResult, ToolSpec
+from .planner import QueryAnalyzer, SchemaPlanner
+from .policy import ExecutionPolicy
+from .proposals import DocumentationModelCallable, SchemaProposal, inspect_documentation_url
+from .registry import InMemoryRegistry, ToolRegistry
+from .runs import RunConfig, RunEvent
+
+_T = TypeVar("_T")
+
+
+def _coerce_config(config: RunConfig | dict[str, Any] | None) -> RunConfig:
+    if config is None:
+        return RunConfig()
+    if isinstance(config, RunConfig):
+        return config
+    return RunConfig.model_validate(config)
+
+
+def _run_sync(factory: Callable[[], Awaitable[_T]]) -> _T:
+    try:
+        asyncio.get_running_loop()
+    except RuntimeError:
+        return asyncio.run(factory())
+    raise RuntimeError(
+        "synchronous SchemaRouter API cannot run inside an active event loop; "
+        "use the async API instead"
+    )
+
+
+def _stream_sync(factory: Callable[[], AsyncIterator[_T]]) -> Iterator[_T]:
+    try:
+        asyncio.get_running_loop()
+    except RuntimeError:
+        pass
+    else:
+        raise RuntimeError(
+            "synchronous SchemaRouter streaming cannot run inside an active event loop; "
+            "use the async streaming API instead"
+        )
+
+    loop = asyncio.new_event_loop()
+    iterator = factory()
+    try:
+        while True:
+            try:
+                yield loop.run_until_complete(iterator.__anext__())
+            except StopAsyncIteration:
+                break
+    finally:
+        loop.run_until_complete(iterator.aclose())
+        loop.close()
+
+
+class SchemaRouter:
+    """High-level facade for schema-aware planning and execution."""
+
+    def __init__(
+        self,
+        *,
+        analyzer: QueryAnalyzer | None = None,
+        http_client: httpx.AsyncClient | None = None,
+        policy: ExecutionPolicy | None = None,
+        registry: ToolRegistry | None = None,
+        adapter_registry: AdapterRegistry | None = None,
+    ) -> None:
+        self.registry = registry if registry is not None else InMemoryRegistry()
+        self.planner = SchemaPlanner(self.registry, analyzer=analyzer)
+        self.executor = RegistryExecutor(self.registry, policy=policy)
+        self.loader = URLSchemaLoader(
+            self.registry,
+            self.executor,
+            http_client=http_client,
+            adapters=adapter_registry,
+        )
+
+    @property
+    def input_schema(self) -> dict[str, Any]:
+        return PlanRequest.model_json_schema()
+
+    @property
+    def output_schema(self) -> dict[str, Any]:
+        return TypeAdapter(list[ToolResult]).json_schema()
+
+    @property
+    def config_schema(self) -> dict[str, Any]:
+        return RunConfig.model_json_schema()
+
+    def with_config(
+        self,
+        config: RunConfig | dict[str, Any],
+    ) -> ConfiguredSchemaRouter:
+        return ConfiguredSchemaRouter(self, _coerce_config(config))
+
+    @classmethod
+    async def from_url(
+        cls,
+        url: str,
+        *,
+        kind: SourceKind = "auto",
+        name: str | None = None,
+        namespace: str | None = None,
+        analyzer: QueryAnalyzer | None = None,
+        http_client: httpx.AsyncClient | None = None,
+        policy: ExecutionPolicy | None = None,
+        registry: ToolRegistry | None = None,
+        adapter_registry: AdapterRegistry | None = None,
+        base_url: str | None = None,
+        schema_headers: dict[str, str] | None = None,
+        trusted_headers: dict[str, str] | None = None,
+    ) -> SchemaRouter:
+        router = cls(
+            analyzer=analyzer,
+            http_client=http_client,
+            policy=policy,
+            registry=registry,
+            adapter_registry=adapter_registry,
+        )
+        await router.add_url(
+            url,
+            kind=kind,
+            name=name,
+            namespace=namespace,
+            base_url=base_url,
+            schema_headers=schema_headers,
+            trusted_headers=trusted_headers,
+        )
+        return router
+
+    def add_tool(self, tool: ToolSpec, *, replace: bool = False) -> str:
+        return self.registry.register(tool, replace=replace)
+
+    def register_adapter(
+        self,
+        adapter: SourceAdapter,
+        *,
+        replace: bool = False,
+    ) -> None:
+        self.loader.register_adapter(adapter, replace=replace)
+
+    @property
+    def adapter_registry(self) -> AdapterRegistry:
+        return self.loader.adapters
+
+    def add_callable(
+        self,
+        function: Callable[..., Any],
+        *,
+        name: str | None = None,
+        namespace: str | None = None,
+        description: str | None = None,
+        read_only: bool | None = None,
+        destructive: bool | None = None,
+        replace: bool = False,
+    ) -> str:
+        decorated = callable_options(function)
+        tool = tool_from_callable(
+            function,
+            name=name if name is not None else decorated.get("name"),
+            namespace=(
+                namespace if namespace is not None else decorated.get("namespace")
+            ),
+            description=(
+                description if description is not None else decorated.get("description")
+            ),
+            read_only=(
+                read_only if read_only is not None else decorated.get("read_only")
+            ),
+            destructive=(
+                destructive if destructive is not None else decorated.get("destructive")
+            ),
+        )
+        invoker = PythonCallableInvoker(function)
+        key = self.registry.register(tool, replace=replace)
+        self.executor.bind(key, invoker)
+        return key
+
+    async def inspect_url(
+        self,
+        url: str,
+        *,
+        model: DocumentationModelCallable,
+        timeout: float = 20.0,
+        max_document_chars: int = 60_000,
+    ) -> SchemaProposal:
+        return await inspect_documentation_url(
+            url,
+            model=model,
+            http_client=self.loader.http_client,
+            timeout=timeout,
+            max_document_chars=max_document_chars,
+        )
+
+    def bind_openapi(
+        self,
+        tool_key: str,
+        *,
+        base_url: str,
+        trusted_headers: dict[str, str] | None = None,
+        timeout: float = 20.0,
+    ) -> None:
+        tool = self.registry.get(tool_key)
+        if tool.metadata.get("adapter") != "openapi":
+            raise RegistrationError(
+                f"tool {tool_key!r} was not imported from OpenAPI"
+            )
+        try:
+            invoker = OpenAPIRemoteInvoker(
+                tool,
+                base_url,
+                trusted_headers=trusted_headers,
+                timeout=timeout,
+            )
+        except ValueError as exc:
+            raise RegistrationError("invalid OpenAPI execution binding") from exc
+        updated = tool.model_copy(deep=True)
+        updated.metadata.update(
+            {
+                "execution_bound": True,
+                "approved_base_url": base_url,
+                "requires_explicit_base_url": False,
+            }
+        )
+        self.registry.register(updated, replace=True)
+        self.executor.bind(tool_key, invoker)
+
+    def approve_proposal(
+        self,
+        proposal: SchemaProposal,
+        *,
+        base_url: str,
+        min_grounding_score: float = 0.8,
+        allow_mutations: bool = False,
+        replace: bool = False,
+        trusted_headers: dict[str, str] | None = None,
+        timeout: float = 20.0,
+    ) -> str:
+        if proposal.status != "grounded" or proposal.tool is None:
+            raise ProposalApprovalError("proposal has no grounded tool to approve")
+        if proposal.grounding_score < min_grounding_score:
+            raise ProposalApprovalError(
+                "proposal grounding score is below the required threshold"
+            )
+
+        parsed = urlparse(base_url)
+        if parsed.scheme not in {"http", "https"} or not parsed.netloc:
+            raise ProposalApprovalError("base_url must be an absolute http(s) URL")
+        if parsed.username or parsed.password:
+            raise ProposalApprovalError(
+                "credentials must not be embedded in base_url; use trusted runtime auth"
+            )
+
+        mutating = [
+            endpoint.name
+            for endpoint in proposal.tool.endpoints
+            if endpoint.method not in {"GET", "HEAD", "OPTIONS"}
+        ]
+        if mutating and not allow_mutations:
+            raise ProposalApprovalError(
+                "proposal contains mutating endpoints; set allow_mutations=True "
+                "after explicit review"
+            )
+
+        tool = proposal.tool.model_copy(deep=True)
+        tool.metadata.update(
+            {
+                "approved_from_proposal": True,
+                "executable": True,
+                "grounding_score": proposal.grounding_score,
+                "approved_base_url": base_url,
+            }
+        )
+        try:
+            invoker = OpenAPIRemoteInvoker(
+                tool,
+                base_url,
+                trusted_headers=trusted_headers,
+                timeout=timeout,
+            )
+        except ValueError as exc:
+            raise ProposalApprovalError("invalid proposal execution binding") from exc
+        key = self.registry.register(tool, replace=replace)
+        self.executor.bind(key, invoker)
+        return key
+
+    async def add_url(
+        self,
+        url: str,
+        *,
+        kind: SourceKind = "auto",
+        name: str | None = None,
+        namespace: str | None = None,
+        replace: bool = False,
+        base_url: str | None = None,
+        schema_headers: dict[str, str] | None = None,
+        trusted_headers: dict[str, str] | None = None,
+        timeout: float = 20.0,
+    ) -> ToolSpec:
+        return await self.loader.load(
+            url,
+            kind=kind,
+            name=name,
+            namespace=namespace,
+            replace=replace,
+            base_url=base_url,
+            schema_headers=schema_headers,
+            trusted_headers=trusted_headers,
+            timeout=timeout,
+        )
+
+    def plan(self, request: PlanRequest | str) -> ExecutionPlan:
+        return self.planner.plan(request)
+
+    async def aplan(self, request: PlanRequest | str) -> ExecutionPlan:
+        return await self.planner.aplan(request)
+
+    async def execute(
+        self,
+        plan: ExecutionPlan,
+        *,
+        config: RunConfig | dict[str, Any] | None = None,
+    ) -> list[ToolResult]:
+        run_config = _coerce_config(config)
+        return await self.executor.execute(plan, retry=run_config.retry)
+
+    async def ainvoke(
+        self,
+        request: PlanRequest | str,
+        *,
+        config: RunConfig | dict[str, Any] | None = None,
+    ) -> list[ToolResult]:
+        run_config = _coerce_config(config)
+        plan = await self.aplan(request)
+        return await self.executor.execute(plan, retry=run_config.retry)
+
+    def invoke(
+        self,
+        request: PlanRequest | str,
+        *,
+        config: RunConfig | dict[str, Any] | None = None,
+    ) -> list[ToolResult]:
+        return _run_sync(lambda: self.ainvoke(request, config=config))
+
+    async def abatch(
+        self,
+        requests: Sequence[PlanRequest | str],
+        *,
+        config: RunConfig | dict[str, Any] | None = None,
+        return_exceptions: bool = False,
+    ) -> list[list[ToolResult] | Exception]:
+        run_config = _coerce_config(config)
+        semaphore = asyncio.Semaphore(run_config.max_concurrency)
+
+        async def invoke_one(request: PlanRequest | str) -> list[ToolResult]:
+            async with semaphore:
+                return await self.ainvoke(request, config=run_config)
+
+        return await asyncio.gather(
+            *(invoke_one(request) for request in requests),
+            return_exceptions=return_exceptions,
+        )
+
+    def batch(
+        self,
+        requests: Sequence[PlanRequest | str],
+        *,
+        config: RunConfig | dict[str, Any] | None = None,
+        return_exceptions: bool = False,
+    ) -> list[list[ToolResult] | Exception]:
+        return _run_sync(
+            lambda: self.abatch(
+                requests,
+                config=config,
+                return_exceptions=return_exceptions,
+            )
+        )
+
+    async def abatch_as_completed(
+        self,
+        requests: Sequence[PlanRequest | str],
+        *,
+        config: RunConfig | dict[str, Any] | None = None,
+        return_exceptions: bool = False,
+    ) -> AsyncIterator[tuple[int, list[ToolResult] | Exception]]:
+        run_config = _coerce_config(config)
+        semaphore = asyncio.Semaphore(run_config.max_concurrency)
+
+        async def invoke_indexed(
+            index: int,
+            request: PlanRequest | str,
+        ) -> tuple[int, list[ToolResult] | Exception]:
+            try:
+                async with semaphore:
+                    result = await self.ainvoke(request, config=run_config)
+                return index, result
+            except Exception as exc:
+                if return_exceptions:
+                    return index, exc
+                raise
+
+        tasks = [
+            asyncio.create_task(invoke_indexed(index, request))
+            for index, request in enumerate(requests)
+        ]
+        try:
+            for completed in asyncio.as_completed(tasks):
+                yield await completed
+        finally:
+            for task in tasks:
+                if not task.done():
+                    task.cancel()
+            await asyncio.gather(*tasks, return_exceptions=True)
+
+    def batch_as_completed(
+        self,
+        requests: Sequence[PlanRequest | str],
+        *,
+        config: RunConfig | dict[str, Any] | None = None,
+        return_exceptions: bool = False,
+    ) -> Iterator[tuple[int, list[ToolResult] | Exception]]:
+        return _stream_sync(
+            lambda: self.abatch_as_completed(
+                requests,
+                config=config,
+                return_exceptions=return_exceptions,
+            )
+        )
+
+    async def astream(
+        self,
+        request: PlanRequest | str,
+        *,
+        config: RunConfig | dict[str, Any] | None = None,
+    ) -> AsyncIterator[ToolResult]:
+        run_config = _coerce_config(config)
+        plan = await self.aplan(request)
+        async for result in self.executor.execute_iter(
+            plan,
+            retry=run_config.retry,
+        ):
+            yield result
+
+    def stream(
+        self,
+        request: PlanRequest | str,
+        *,
+        config: RunConfig | dict[str, Any] | None = None,
+    ) -> Iterator[ToolResult]:
+        return _stream_sync(lambda: self.astream(request, config=config))
+
+    async def astream_events(
+        self,
+        request: PlanRequest | str,
+        *,
+        config: RunConfig | dict[str, Any] | None = None,
+    ) -> AsyncIterator[RunEvent]:
+        run_config = _coerce_config(config)
+        run_id = uuid4().hex
+        sequence = 0
+
+        request_payload: dict[str, Any] = {
+            "input_type": type(request).__name__,
+        }
+        if run_config.include_payloads:
+            request_payload["input"] = (
+                request.model_dump(mode="json")
+                if isinstance(request, PlanRequest)
+                else request
+            )
+
+        yield RunEvent.create(
+            event="run.start",
+            run_id=run_id,
+            sequence=sequence,
+            config=run_config,
+            data=request_payload,
+        )
+        sequence += 1
+
+        try:
+            plan = await self.aplan(request)
+        except Exception as exc:
+            data = {"error_type": type(exc).__name__, "stage": "planning"}
+            if run_config.include_payloads:
+                data["message"] = str(exc)
+            yield RunEvent.create(
+                event="run.error",
+                run_id=run_id,
+                sequence=sequence,
+                config=run_config,
+                data=data,
+            )
+            raise
+
+        plan_data: dict[str, Any] = {
+            "call_count": len(plan.calls),
+            "warnings": list(plan.warnings),
+        }
+        if run_config.include_payloads:
+            plan_data["plan"] = plan.model_dump(mode="json")
+        yield RunEvent.create(
+            event="plan.end",
+            run_id=run_id,
+            sequence=sequence,
+            config=run_config,
+            data=plan_data,
+        )
+        sequence += 1
+
+        result_count = 0
+        for call in plan.calls:
+            start_data: dict[str, Any] = {
+                "argument_names": sorted(call.arguments),
+                "fields": list(call.fields),
+            }
+            if run_config.include_payloads:
+                start_data["arguments"] = dict(call.arguments)
+            yield RunEvent.create(
+                event="tool.start",
+                run_id=run_id,
+                sequence=sequence,
+                config=run_config,
+                tool=call.tool,
+                endpoint=call.endpoint,
+                data=start_data,
+            )
+            sequence += 1
+
+            try:
+                result = await self.executor.execute_call(
+                    call,
+                    retry=run_config.retry,
+                )
+            except Exception as exc:
+                error_data = {"error_type": type(exc).__name__}
+                if run_config.include_payloads:
+                    error_data["message"] = str(exc)
+                yield RunEvent.create(
+                    event="tool.error",
+                    run_id=run_id,
+                    sequence=sequence,
+                    config=run_config,
+                    tool=call.tool,
+                    endpoint=call.endpoint,
+                    data=error_data,
+                )
+                sequence += 1
+                yield RunEvent.create(
+                    event="run.error",
+                    run_id=run_id,
+                    sequence=sequence,
+                    config=run_config,
+                    data={
+                        "error_type": type(exc).__name__,
+                        "stage": "execution",
+                    },
+                )
+                raise
+
+            end_data: dict[str, Any] = {
+                "projected_fields": list(result.projected_fields),
+            }
+            if run_config.include_payloads:
+                end_data["result"] = result.model_dump(mode="json")
+            yield RunEvent.create(
+                event="tool.end",
+                run_id=run_id,
+                sequence=sequence,
+                config=run_config,
+                tool=result.tool,
+                endpoint=result.endpoint,
+                data=end_data,
+            )
+            sequence += 1
+            result_count += 1
+
+        yield RunEvent.create(
+            event="run.end",
+            run_id=run_id,
+            sequence=sequence,
+            config=run_config,
+            data={"result_count": result_count},
+        )
+
+    async def run(
+        self,
+        request: PlanRequest | str,
+        *,
+        config: RunConfig | dict[str, Any] | None = None,
+    ) -> list[ToolResult]:
+        return await self.ainvoke(request, config=config)
+
+    async def arun(
+        self,
+        request: PlanRequest | str,
+        *,
+        config: RunConfig | dict[str, Any] | None = None,
+    ) -> list[ToolResult]:
+        return await self.ainvoke(request, config=config)
+
+
+class ConfiguredSchemaRouter:
+    """SchemaRouter with an immutable default RunConfig."""
+
+    def __init__(self, router: SchemaRouter, config: RunConfig) -> None:
+        self.router = router
+        self.config = config
+
+    def invoke(self, request: PlanRequest | str) -> list[ToolResult]:
+        return self.router.invoke(request, config=self.config)
+
+    async def ainvoke(self, request: PlanRequest | str) -> list[ToolResult]:
+        return await self.router.ainvoke(request, config=self.config)
+
+    def batch(
+        self,
+        requests: Sequence[PlanRequest | str],
+        *,
+        return_exceptions: bool = False,
+    ) -> list[list[ToolResult] | Exception]:
+        return self.router.batch(
+            requests,
+            config=self.config,
+            return_exceptions=return_exceptions,
+        )
+
+    async def abatch(
+        self,
+        requests: Sequence[PlanRequest | str],
+        *,
+        return_exceptions: bool = False,
+    ) -> list[list[ToolResult] | Exception]:
+        return await self.router.abatch(
+            requests,
+            config=self.config,
+            return_exceptions=return_exceptions,
+        )
+
+    def batch_as_completed(
+        self,
+        requests: Sequence[PlanRequest | str],
+        *,
+        return_exceptions: bool = False,
+    ) -> Iterator[tuple[int, list[ToolResult] | Exception]]:
+        return self.router.batch_as_completed(
+            requests,
+            config=self.config,
+            return_exceptions=return_exceptions,
+        )
+
+    def abatch_as_completed(
+        self,
+        requests: Sequence[PlanRequest | str],
+        *,
+        return_exceptions: bool = False,
+    ) -> AsyncIterator[tuple[int, list[ToolResult] | Exception]]:
+        return self.router.abatch_as_completed(
+            requests,
+            config=self.config,
+            return_exceptions=return_exceptions,
+        )
+
+    def stream(self, request: PlanRequest | str) -> Iterator[ToolResult]:
+        return self.router.stream(request, config=self.config)
+
+    def astream(self, request: PlanRequest | str) -> AsyncIterator[ToolResult]:
+        return self.router.astream(request, config=self.config)
+
+    def astream_events(self, request: PlanRequest | str) -> AsyncIterator[RunEvent]:
+        return self.router.astream_events(request, config=self.config)