infinity-context 0.1.0__py3-none-any.whl

infinity_context_server/memory_comparison_chat_completions.py

@@ -0,0 +1,402 @@
+"""OpenAI-compatible Chat Completions transport for evaluation LLMs.
+
+This module is intentionally confined to the memory-comparison evaluation
+boundary.  Core memory behavior remains unaware of provider HTTP contracts.
+"""
+
+from __future__ import annotations
+
+import time
+from collections.abc import Callable, Mapping, Sequence
+from dataclasses import dataclass
+from math import isfinite
+from typing import Protocol
+from urllib.parse import SplitResult, urlsplit, urlunsplit
+
+import httpx
+
+from infinity_context_server.memory_comparison_llm import (
+    _judge_prompt,
+    _parse_judge_payload,
+    approximate_token_count,
+    render_answer_prompt,
+)
+from infinity_context_server.memory_comparison_models import (
+    AnswerResult,
+    JudgeResult,
+    RetrievedMemory,
+    TokenUsage,
+)
+from infinity_context_server.memory_comparison_prompt_policy import (
+    resolve_memory_comparison_prompt_policy,
+)
+from infinity_context_server.public_benchmark_models import PublicBenchmarkCase
+
+_RETRYABLE_STATUS_CODES = frozenset({408, 409, 429})
+
+
+class ChatCompletionsError(RuntimeError):
+    """Base error for the evaluation Chat Completions adapter."""
+
+
+class ChatCompletionsRequestError(ChatCompletionsError):
+    """Raised when the compatible endpoint cannot be reached."""
+
+
+class ChatCompletionsHTTPError(ChatCompletionsError):
+    """Raised for a non-success response without retaining its body."""
+
+    def __init__(self, status_code: int) -> None:
+        self.status_code = int(status_code)
+        super().__init__(f"chat completions request failed with HTTP {self.status_code}")
+
+
+class ChatCompletionsMalformedResponseError(ChatCompletionsError):
+    """Raised when a compatible response has no usable assistant text."""
+
+
+@dataclass(frozen=True)
+class ChatCompletion:
+    """Provider-neutral result returned through the evaluation transport port."""
+
+    text: str
+    prompt_tokens: int = 0
+    completion_tokens: int = 0
+
+
+class EvaluationChatCompletionsPort(Protocol):
+    """Narrow transport port used by evaluation answerer and judge adapters."""
+
+    def complete(
+        self,
+        *,
+        model: str,
+        system_prompt: str,
+        user_prompt: str,
+        max_output_tokens: int,
+    ) -> ChatCompletion:
+        """Return one assistant completion for the supplied evaluator messages."""
+
+    def close(self) -> None:
+        """Release transport resources."""
+
+
+class OpenAICompatibleChatCompletionsTransport:
+    """HTTP adapter for an explicitly configured OpenAI-compatible endpoint."""
+
+    def __init__(
+        self,
+        *,
+        base_url: str,
+        api_key: str | None = None,
+        timeout_seconds: float = 600.0,
+        max_retries: int = 2,
+        transport: httpx.BaseTransport | None = None,
+        retry_backoff_seconds: float = 0.25,
+        sleep: Callable[[float], None] = time.sleep,
+    ) -> None:
+        timeout = float(timeout_seconds)
+        if not isfinite(timeout) or timeout <= 0:
+            raise ValueError("chat completions timeout must be positive")
+        if isinstance(max_retries, bool) or not isinstance(max_retries, int):
+            raise ValueError("chat completions max_retries must be an integer")
+        if not 0 <= max_retries <= 10:
+            raise ValueError("chat completions max_retries must be between 0 and 10")
+        backoff = float(retry_backoff_seconds)
+        if not isfinite(backoff) or backoff < 0:
+            raise ValueError("chat completions retry backoff must be non-negative")
+
+        self._endpoint = normalize_chat_completions_endpoint(base_url)
+        normalized_api_key = str(api_key or "").strip()
+        self._authorization = f"Bearer {normalized_api_key}" if normalized_api_key else None
+        self._max_retries = max_retries
+        self._retry_backoff_seconds = backoff
+        self._sleep = sleep
+        self._client = httpx.Client(timeout=timeout, transport=transport)
+
+    def complete(
+        self,
+        *,
+        model: str,
+        system_prompt: str,
+        user_prompt: str,
+        max_output_tokens: int,
+    ) -> ChatCompletion:
+        normalized_model = str(model or "").strip()
+        if not normalized_model:
+            raise ValueError("chat completions model must not be blank")
+        if max_output_tokens <= 0:
+            raise ValueError("chat completions max_output_tokens must be positive")
+        payload: dict[str, object] = {
+            "model": normalized_model,
+            "messages": [
+                {"role": "system", "content": str(system_prompt)},
+                {"role": "user", "content": str(user_prompt)},
+            ],
+            "max_tokens": int(max_output_tokens),
+        }
+        response = self._post(payload)
+        if not 200 <= response.status_code < 300:
+            raise ChatCompletionsHTTPError(response.status_code)
+        try:
+            decoded = response.json()
+        except ValueError:
+            raise ChatCompletionsMalformedResponseError(
+                "chat completions response was not valid JSON"
+            ) from None
+        return _parse_chat_completion(decoded)
+
+    def close(self) -> None:
+        self._client.close()
+
+    def _post(self, payload: Mapping[str, object]) -> httpx.Response:
+        headers = {"Content-Type": "application/json"}
+        if self._authorization is not None:
+            headers["Authorization"] = self._authorization
+        for attempt in range(self._max_retries + 1):
+            try:
+                response = self._client.post(
+                    self._endpoint,
+                    headers=headers,
+                    json=dict(payload),
+                )
+            except httpx.TransportError:
+                if attempt >= self._max_retries:
+                    raise ChatCompletionsRequestError("chat completions request failed") from None
+                self._backoff(attempt)
+                continue
+            if not _retryable_status(response.status_code) or attempt >= self._max_retries:
+                return response
+            self._backoff(attempt)
+        raise AssertionError("chat completions retry loop exhausted unexpectedly")
+
+    def _backoff(self, attempt: int) -> None:
+        delay = self._retry_backoff_seconds * (2**attempt)
+        if delay:
+            self._sleep(delay)
+
+
+class ChatCompletionsAnswerer:
+    """Memory-comparison answerer backed by the narrow chat transport port."""
+
+    def __init__(
+        self,
+        *,
+        transport: EvaluationChatCompletionsPort,
+        model: str,
+        max_output_tokens: int = 400,
+    ) -> None:
+        self.model = _nonblank_model(model)
+        self._transport = transport
+        self._max_output_tokens = int(max_output_tokens)
+
+    def answer(
+        self,
+        case: PublicBenchmarkCase,
+        memories: Sequence[RetrievedMemory],
+        *,
+        backend_name: str,
+        cutoff: int,
+    ) -> AnswerResult:
+        started = time.perf_counter()
+        policy = resolve_memory_comparison_prompt_policy(case)
+        prompt = render_answer_prompt(case, memories, cutoff=cutoff)
+        completion = self._transport.complete(
+            model=self.model,
+            system_prompt=policy.answer_system_instruction,
+            user_prompt=prompt,
+            max_output_tokens=self._max_output_tokens,
+        )
+        return AnswerResult(
+            answer=completion.text,
+            model=self.model,
+            latency_ms=round((time.perf_counter() - started) * 1000, 2),
+            token_usage=_token_usage(completion, prompt=prompt),
+            metadata={
+                "backend_name": backend_name,
+                "cutoff": cutoff,
+                "provider": "openai",
+                "transport": "chat-completions",
+                "prompt_policy_id": policy.prompt_policy_id,
+            },
+        )
+
+    def close(self) -> None:
+        self._transport.close()
+
+
+class ChatCompletionsJudge:
+    """Memory-comparison judge backed by the narrow chat transport port."""
+
+    def __init__(
+        self,
+        *,
+        transport: EvaluationChatCompletionsPort,
+        model: str,
+        max_output_tokens: int = 300,
+    ) -> None:
+        self.model = _nonblank_model(model)
+        self._transport = transport
+        self._max_output_tokens = int(max_output_tokens)
+
+    def judge(
+        self,
+        case: PublicBenchmarkCase,
+        answer: AnswerResult,
+        memories: Sequence[RetrievedMemory],
+        *,
+        backend_name: str,
+        cutoff: int,
+    ) -> JudgeResult:
+        started = time.perf_counter()
+        policy = resolve_memory_comparison_prompt_policy(case)
+        prompt = _judge_prompt(case, answer, memories)
+        completion = self._transport.complete(
+            model=self.model,
+            system_prompt=policy.judge_system_instruction,
+            user_prompt=prompt,
+            max_output_tokens=self._max_output_tokens,
+        )
+        try:
+            payload = _parse_judge_payload(completion.text)
+        except ValueError:
+            raise ChatCompletionsMalformedResponseError(
+                "chat completions judge content was malformed"
+            ) from None
+        return JudgeResult(
+            verdict=payload["verdict"],
+            score=payload["score"],
+            reason=payload["reason"],
+            model=self.model,
+            latency_ms=round((time.perf_counter() - started) * 1000, 2),
+            token_usage=_token_usage(completion, prompt=prompt),
+            metadata={
+                "backend_name": backend_name,
+                "cutoff": cutoff,
+                "provider": "openai",
+                "transport": "chat-completions",
+                "prompt_policy_id": policy.prompt_policy_id,
+            },
+        )
+
+    def close(self) -> None:
+        self._transport.close()
+
+
+def normalize_chat_completions_endpoint(base_url: str) -> str:
+    """Normalize an API base URL or full endpoint without deployment assumptions."""
+
+    raw_url = str(base_url or "").strip()
+    try:
+        parsed = urlsplit(raw_url)
+        port = parsed.port
+    except ValueError:
+        raise ValueError("chat completions base URL is invalid") from None
+    if (
+        parsed.scheme not in {"http", "https"}
+        or not parsed.hostname
+        or parsed.username is not None
+        or parsed.password is not None
+        or parsed.query
+        or parsed.fragment
+    ):
+        raise ValueError(
+            "chat completions base URL must be an http(s) URL without credentials, "
+            "query, or fragment"
+        )
+    netloc = parsed.hostname
+    if ":" in netloc and not netloc.startswith("["):
+        netloc = f"[{netloc}]"
+    if port is not None:
+        netloc = f"{netloc}:{port}"
+    path = parsed.path.rstrip("/")
+    if path.endswith("/chat/completions"):
+        endpoint_path = path
+    elif path.endswith("/v1"):
+        endpoint_path = f"{path}/chat/completions"
+    else:
+        endpoint_path = f"{path}/v1/chat/completions"
+    normalized = SplitResult(parsed.scheme, netloc, endpoint_path, "", "")
+    return urlunsplit(normalized)
+
+
+def _parse_chat_completion(payload: object) -> ChatCompletion:
+    if not isinstance(payload, Mapping):
+        raise _malformed_response()
+    choices = payload.get("choices")
+    if not isinstance(choices, Sequence) or isinstance(choices, str | bytes) or not choices:
+        raise _malformed_response()
+    first = choices[0]
+    if not isinstance(first, Mapping):
+        raise _malformed_response()
+    message = first.get("message")
+    content = message.get("content") if isinstance(message, Mapping) else None
+    text = _content_text(content)
+    if not text:
+        text = _content_text(first.get("text"))
+    if not text:
+        raise _malformed_response()
+    usage = payload.get("usage")
+    return ChatCompletion(
+        text=text,
+        prompt_tokens=_usage_int(usage, "prompt_tokens", "input_tokens"),
+        completion_tokens=_usage_int(
+            usage,
+            "completion_tokens",
+            "output_tokens",
+        ),
+    )
+
+
+def _content_text(content: object) -> str:
+    if isinstance(content, str):
+        return content.strip()
+    if not isinstance(content, Sequence) or isinstance(content, str | bytes):
+        return ""
+    parts: list[str] = []
+    for part in content:
+        if isinstance(part, str) and part.strip():
+            parts.append(part.strip())
+        elif isinstance(part, Mapping):
+            text = part.get("text")
+            if isinstance(text, str) and text.strip():
+                parts.append(text.strip())
+    return "\n".join(parts).strip()
+
+
+def _token_usage(completion: ChatCompletion, *, prompt: str) -> TokenUsage:
+    return TokenUsage(
+        prompt_tokens=(completion.prompt_tokens or approximate_token_count(prompt)),
+        completion_tokens=(
+            completion.completion_tokens or approximate_token_count(completion.text)
+        ),
+    )
+
+
+def _usage_int(usage: object, *names: str) -> int:
+    if not isinstance(usage, Mapping):
+        return 0
+    for name in names:
+        value = usage.get(name)
+        if isinstance(value, bool):
+            continue
+        if isinstance(value, int | float):
+            return max(0, int(value))
+    return 0
+
+
+def _nonblank_model(model: str) -> str:
+    normalized = str(model or "").strip()
+    if not normalized:
+        raise ValueError("chat completions evaluator requires a model")
+    return normalized
+
+
+def _retryable_status(status_code: int) -> bool:
+    return status_code in _RETRYABLE_STATUS_CODES or status_code >= 500
+
+
+def _malformed_response() -> ChatCompletionsMalformedResponseError:
+    return ChatCompletionsMalformedResponseError(
+        "chat completions response did not include assistant text"
+    )