aid-runtime-python 0.1.39__py3-none-any.whl

aid_runtime/providers/http.py

@@ -0,0 +1,344 @@
+"""A single OpenAI-compatible HTTP provider adapter.
+
+One adapter, deliberately. Anthropic, Azure, Bedrock and friends all speak
+slightly different dialects, and the framework's value is the IR seam and the
+eval gate, not the breadth of vendor plumbing. Anything OpenAI-compatible --
+vLLM, Ollama's compatibility endpoint, LiteLLM, OpenAI itself -- works through
+this one class via ``base_url``.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+from collections.abc import AsyncIterator
+from typing import Any
+
+import httpx
+
+from aid_runtime.errors import ProviderError, ProviderResponseError
+from aid_runtime.types import (
+    Capability,
+    CompletionRequest,
+    CompletionResponse,
+    FinishReason,
+    Message,
+    ModelSpec,
+    StreamChunk,
+    TokenUsage,
+    ToolCall,
+)
+
+__all__ = ["OpenAICompatibleProvider"]
+
+DEFAULT_BASE_URL = "https://api.openai.com/v1"
+
+# 429 and 5xx are the retryable classes: the request was well-formed but the
+# endpoint declined for a reason that may not repeat.
+_RETRYABLE_STATUS = frozenset({408, 409, 425, 429, 500, 502, 503, 504})
+
+
+class OpenAICompatibleProvider:
+    """Calls ``POST {base_url}/chat/completions`` and maps the response back."""
+
+    def __init__(
+        self,
+        *,
+        base_url: str = DEFAULT_BASE_URL,
+        api_key: str | None = None,
+        name: str = "openai-compatible",
+        timeout_seconds: float = 60.0,
+        max_attempts: int = 3,
+        backoff_seconds: float = 0.5,
+        client: httpx.AsyncClient | None = None,
+        extra_headers: dict[str, str] | None = None,
+    ) -> None:
+        if not base_url.startswith(("http://", "https://")):
+            raise ValueError(f"base_url must be an absolute http(s) URL, got {base_url!r}")
+        if max_attempts < 1:
+            raise ValueError("max_attempts must be at least 1")
+        self._name = name
+        self._base_url = base_url.rstrip("/")
+        self._api_key = api_key
+        self._timeout = timeout_seconds
+        self._max_attempts = max_attempts
+        self._backoff = backoff_seconds
+        self._extra_headers = dict(extra_headers or {})
+        self._owns_client = client is None
+        self._client = client or httpx.AsyncClient(timeout=timeout_seconds)
+        self._closed = False
+
+    @property
+    def name(self) -> str:
+        return self._name
+
+    @property
+    def capabilities(self) -> frozenset[Capability]:
+        return frozenset(
+            {
+                Capability.TOOL_CALLING,
+                Capability.STRUCTURED_OUTPUT,
+                Capability.STREAMING,
+            }
+        )
+
+    async def complete(self, request: CompletionRequest, *, model: ModelSpec) -> CompletionResponse:
+        payload = self._build_payload(request, model)
+        body = await self._post("/chat/completions", payload)
+        return self._parse(body, model)
+
+    async def stream(
+        self, request: CompletionRequest, *, model: ModelSpec
+    ) -> AsyncIterator[StreamChunk]:
+        """Stream one call as OpenAI-style SSE, yielding one chunk per delta.
+
+        Streaming is a single attempt, not the retry loop ``complete`` uses: a
+        stream that has already emitted a token cannot be safely replayed, so a
+        transport or status error is raised once rather than retried mid-stream.
+        """
+        payload = self._build_payload(request, model)
+        payload["stream"] = True
+        payload["stream_options"] = {"include_usage": True}
+        url = f"{self._base_url}/chat/completions"
+        try:
+            async with self._client.stream(
+                "POST", url, json=payload, headers=self._headers()
+            ) as response:
+                if response.status_code >= 400:
+                    await response.aread()
+                    raise ProviderError(
+                        f"{self._name} returned {response.status_code}: {response.text[:500]}"
+                    )
+                async for line in response.aiter_lines():
+                    if not line.startswith("data:"):
+                        continue
+                    data = line[len("data:") :].strip()
+                    if data == "[DONE]":
+                        continue
+                    try:
+                        decoded = json.loads(data)
+                    except json.JSONDecodeError as exc:
+                        raise ProviderResponseError(
+                            f"{self._name} streamed non-JSON", data
+                        ) from exc
+                    if not isinstance(decoded, dict):
+                        raise ProviderResponseError(
+                            f"{self._name} streamed a non-object body", data
+                        )
+                    yield self._parse_chunk(decoded, model)
+        except httpx.TransportError as exc:
+            raise ProviderError(f"{self._name} stream failed: {exc}") from exc
+
+    async def aclose(self) -> None:
+        if self._closed:
+            return
+        self._closed = True
+        if self._owns_client:
+            await self._client.aclose()
+
+    def _headers(self) -> dict[str, str]:
+        headers = {"content-type": "application/json", **self._extra_headers}
+        if self._api_key:
+            headers["authorization"] = f"Bearer {self._api_key}"
+        return headers
+
+    def _build_payload(self, request: CompletionRequest, model: ModelSpec) -> dict[str, Any]:
+        payload: dict[str, Any] = {
+            "model": model.name,
+            "messages": [_message_payload(message) for message in request.messages],
+        }
+        if request.temperature is not None:
+            payload["temperature"] = request.temperature
+        if request.max_output_tokens is not None:
+            payload["max_tokens"] = min(request.max_output_tokens, model.max_output_tokens)
+        if request.tools:
+            payload["tools"] = [tool.as_openai_tool() for tool in request.tools]
+        if request.stop:
+            payload["stop"] = list(request.stop)
+        return payload
+
+    async def _post(self, path: str, payload: dict[str, Any]) -> dict[str, Any]:
+        url = f"{self._base_url}{path}"
+        last_error: Exception | None = None
+        for attempt in range(1, self._max_attempts + 1):
+            try:
+                response = await self._client.post(url, json=payload, headers=self._headers())
+            except httpx.TransportError as exc:
+                last_error = exc
+                if attempt == self._max_attempts:
+                    break
+                await asyncio.sleep(self._backoff * 2 ** (attempt - 1))
+                continue
+
+            if response.status_code in _RETRYABLE_STATUS and attempt < self._max_attempts:
+                last_error = ProviderError(f"{self._name} returned {response.status_code}")
+                await asyncio.sleep(self._backoff * 2 ** (attempt - 1))
+                continue
+
+            if response.status_code >= 400:
+                raise ProviderError(
+                    f"{self._name} returned {response.status_code}: {response.text[:500]}"
+                )
+
+            try:
+                decoded = response.json()
+            except json.JSONDecodeError as exc:
+                raise ProviderResponseError(
+                    f"{self._name} returned non-JSON", response.text
+                ) from exc
+            if not isinstance(decoded, dict):
+                raise ProviderResponseError(
+                    f"{self._name} returned a non-object body", response.text
+                )
+            return decoded
+
+        raise ProviderError(
+            f"{self._name} failed after {self._max_attempts} attempt(s): {last_error}"
+        ) from last_error
+
+    def _parse(self, body: dict[str, Any], model: ModelSpec) -> CompletionResponse:
+        choices = body.get("choices")
+        if not isinstance(choices, list) or not choices:
+            raise ProviderResponseError(
+                f"{self._name} returned no choices", json.dumps(body, sort_keys=True)
+            )
+        first = choices[0]
+        if not isinstance(first, dict):
+            raise ProviderResponseError(f"{self._name} returned a malformed choice")
+        message = first.get("message")
+        if not isinstance(message, dict):
+            raise ProviderResponseError(
+                f"{self._name} returned a choice without a message",
+                json.dumps(first, sort_keys=True),
+            )
+
+        text = message.get("content")
+        if text is None:
+            text = ""
+        if not isinstance(text, str):
+            raise ProviderResponseError(f"{self._name} returned non-string content")
+
+        tool_calls = self._parse_tool_calls(message.get("tool_calls"))
+
+        raw_usage = body.get("usage")
+        usage = TokenUsage(input_tokens=0, output_tokens=0)
+        if isinstance(raw_usage, dict):
+            usage = TokenUsage(
+                input_tokens=int(raw_usage.get("prompt_tokens") or 0),
+                output_tokens=int(raw_usage.get("completion_tokens") or 0),
+            )
+
+        return CompletionResponse(
+            text=text,
+            model=str(body.get("model") or model.name),
+            provider=self._name,
+            usage=usage,
+            tool_calls=tool_calls,
+            finish_reason=_finish_reason(first.get("finish_reason")),
+        )
+
+    def _parse_chunk(self, body: dict[str, Any], model: ModelSpec) -> StreamChunk:
+        """Map one streamed SSE ``data`` payload to a :class:`StreamChunk`."""
+        choices = body.get("choices")
+        delta = ""
+        finish_reason: FinishReason | None = None
+        if isinstance(choices, list) and choices:
+            first = choices[0]
+            if isinstance(first, dict):
+                delta_obj = first.get("delta")
+                if isinstance(delta_obj, dict):
+                    content = delta_obj.get("content")
+                    if isinstance(content, str):
+                        delta = content
+                raw_reason = first.get("finish_reason")
+                if isinstance(raw_reason, str) and raw_reason:
+                    finish_reason = _finish_reason(raw_reason)
+
+        usage: TokenUsage | None = None
+        raw_usage = body.get("usage")
+        if isinstance(raw_usage, dict):
+            usage = TokenUsage(
+                input_tokens=int(raw_usage.get("prompt_tokens") or 0),
+                output_tokens=int(raw_usage.get("completion_tokens") or 0),
+            )
+        return StreamChunk(delta=delta, finish_reason=finish_reason, usage=usage)
+
+    def _parse_tool_calls(self, raw: Any) -> tuple[ToolCall, ...]:
+        if not isinstance(raw, list):
+            return ()
+        calls: list[ToolCall] = []
+        for index, item in enumerate(raw):
+            if not isinstance(item, dict):
+                continue
+            function = item.get("function")
+            if not isinstance(function, dict):
+                continue
+            name = function.get("name")
+            if not isinstance(name, str) or not name:
+                continue
+            calls.append(
+                ToolCall(
+                    id=str(item.get("id") or f"call_{index}"),
+                    name=name,
+                    arguments=_decode_arguments(function.get("arguments")),
+                )
+            )
+        return tuple(calls)
+
+
+_FINISH_REASONS: dict[str, FinishReason] = {
+    "stop": "stop",
+    "length": "length",
+    "tool_calls": "tool-calls",
+    "content_filter": "content-filter",
+    "error": "error",
+}
+
+
+def _finish_reason(raw: Any) -> FinishReason:
+    return _FINISH_REASONS.get(raw, "stop") if isinstance(raw, str) else "stop"
+
+
+def _decode_arguments(raw: Any) -> dict[str, Any]:
+    """Tool arguments arrive as a JSON *string*; a malformed one is an empty dict.
+
+    Returning ``{}`` rather than raising keeps the failure in the tool executor,
+    where it is reported to the model as a retryable tool error instead of
+    aborting the whole turn.
+    """
+    if isinstance(raw, dict):
+        return raw
+    if not isinstance(raw, str) or not raw.strip():
+        return {}
+    try:
+        decoded = json.loads(raw)
+    except json.JSONDecodeError:
+        return {}
+    return decoded if isinstance(decoded, dict) else {}
+
+
+def _message_payload(message: Message) -> dict[str, Any]:
+    """One OpenAI chat message, omitting fields the transcript did not set.
+
+    Single-turn requests must keep the exact ``{"role", "content"}`` shape they had
+    before tool bookkeeping existed, so absent fields are omitted rather than sent
+    as ``null``/``[]``.
+    """
+    payload: dict[str, Any] = {"role": message.role, "content": message.content}
+    if message.tool_calls:
+        payload["tool_calls"] = [
+            {
+                "id": call.id,
+                "type": "function",
+                "function": {
+                    "name": call.name,
+                    "arguments": json.dumps(call.arguments, sort_keys=True),
+                },
+            }
+            for call in message.tool_calls
+        ]
+    if message.tool_call_id is not None:
+        payload["tool_call_id"] = message.tool_call_id
+    if message.name is not None:
+        payload["name"] = message.name
+    return payload

aid_runtime/retrieval/__init__.py

@@ -0,0 +1,178 @@
+"""The retrieval runtime: connectors, chunkers, embedders, stores, fusion, citations.
+
+The package is one pipeline in source order. A connector yields documents, a chunker
+splits them, an embedder vectorizes them, a store holds them, and retrieval goes back
+out through the store, optionally fusing, optionally reranking, and always attaching
+citations.
+
+Every seam is a protocol with a deterministic implementation and, where the plan
+allows one, a hosted implementation. Phase 1 ships exactly one hosted adapter --
+:class:`OpenAICompatibleEmbedder` -- alongside :class:`FakeDeterministicEmbedder`,
+which is what lets a gate run offline. There is no hosted reranker: the only
+reranker is the deterministic :class:`LexicalReranker`.
+
+Import safety is deliberate. :class:`PgVectorStore` is exported from here, but the
+``psycopg`` driver is imported inside its connect path, never at module scope, so
+``import aid_runtime`` works without the ``pgvector`` extra installed -- which is the
+only state CI is ever in, since CI installs the ``dev`` extra and ``pgvector`` is not
+part of it. Reaching for the driver without the extra raises ``ConfigurationError``
+naming the extra, rather than ``ImportError`` from three frames down.
+"""
+
+from __future__ import annotations
+
+from aid_runtime.retrieval.chunking import (
+    Chunker,
+    CodeChunker,
+    FixedChunker,
+    RecursiveChunker,
+    chunk_document,
+    chunker_for,
+)
+from aid_runtime.retrieval.citations import (
+    Citation,
+    RetrievalResult,
+    RetrievedItem,
+    assemble_context,
+    cite,
+)
+from aid_runtime.retrieval.connectors import (
+    Connector,
+    FileConnector,
+    SourceDocument,
+    open_connector,
+)
+from aid_runtime.retrieval.embedding import (
+    Embedder,
+    EmbeddingBatch,
+    FakeDeterministicEmbedder,
+    OpenAICompatibleEmbedder,
+    embedder_for,
+    estimate_tokens,
+)
+from aid_runtime.retrieval.gate import RetrievalGate
+from aid_runtime.retrieval.hybrid import (
+    RankedId,
+    bm25_scores,
+    fuse,
+    lexical_ranking,
+    rrf_score,
+    tokenize,
+)
+from aid_runtime.retrieval.metrics import (
+    CitationGrounding,
+    MeanReciprocalRank,
+    RecallAtK,
+    RetrievalCase,
+    RetrievalMetric,
+    retrieval_metrics,
+)
+from aid_runtime.retrieval.pgvector import PgVectorStore, resolve_uri
+from aid_runtime.retrieval.rerank import LexicalReranker, Reranker, lexical_score, reranker_for
+from aid_runtime.retrieval.retriever import DEFAULT_CANDIDATE_FACTOR, IngestReport, Retriever
+from aid_runtime.retrieval.stores import (
+    InMemoryVectorStore,
+    SearchHit,
+    StoreRecord,
+    VectorStore,
+    create_store,
+    similarity,
+)
+from aid_runtime.retrieval.types import (
+    Chunk,
+    ChunkingSpec,
+    ChunkingStrategy,
+    DistanceMetric,
+    EmbeddingCostProfile,
+    EmbeddingSpec,
+    IndexSpec,
+    PinnedEmbedding,
+    RefreshMode,
+    RefreshSpec,
+    RerankSpec,
+    RetrievalMetricSpec,
+    RetrieverBinding,
+    RetrieverSpec,
+    ScoredChunk,
+    SourceKind,
+    SourceSpec,
+    StoreKind,
+    VectorStoreSpec,
+    binding_from_spec,
+    chunk_id_for,
+    pinned_embedding,
+)
+
+__all__ = [
+    "DEFAULT_CANDIDATE_FACTOR",
+    "Chunk",
+    "Chunker",
+    "ChunkingSpec",
+    "ChunkingStrategy",
+    "Citation",
+    "CitationGrounding",
+    "CodeChunker",
+    "Connector",
+    "DistanceMetric",
+    "Embedder",
+    "EmbeddingBatch",
+    "EmbeddingCostProfile",
+    "EmbeddingSpec",
+    "FakeDeterministicEmbedder",
+    "FileConnector",
+    "FixedChunker",
+    "InMemoryVectorStore",
+    "IndexSpec",
+    "IngestReport",
+    "LexicalReranker",
+    "MeanReciprocalRank",
+    "OpenAICompatibleEmbedder",
+    "PgVectorStore",
+    "PinnedEmbedding",
+    "RankedId",
+    "RecallAtK",
+    "RecursiveChunker",
+    "RefreshMode",
+    "RefreshSpec",
+    "RerankSpec",
+    "Reranker",
+    "RetrievalCase",
+    "RetrievalGate",
+    "RetrievalMetric",
+    "RetrievalMetricSpec",
+    "RetrievalResult",
+    "RetrievedItem",
+    "Retriever",
+    "RetrieverBinding",
+    "RetrieverSpec",
+    "ScoredChunk",
+    "SearchHit",
+    "SourceDocument",
+    "SourceKind",
+    "SourceSpec",
+    "StoreKind",
+    "StoreRecord",
+    "VectorStore",
+    "VectorStoreSpec",
+    "assemble_context",
+    "binding_from_spec",
+    "bm25_scores",
+    "chunk_document",
+    "chunk_id_for",
+    "chunker_for",
+    "cite",
+    "create_store",
+    "embedder_for",
+    "estimate_tokens",
+    "fuse",
+    "lexical_ranking",
+    "lexical_score",
+    "open_connector",
+    "pinned_embedding",
+    "reranker_for",
+    "resolve_uri",
+    "retrieval_metrics",
+    "rrf_score",
+    "similarity",
+    "tokenize",
+]