aituner/src/aituner/http_client.py

from __future__ import annotations

import json
import os
import time
import tomllib
import urllib.error
import urllib.request
from dataclasses import dataclass
from ipaddress import ip_address
from pathlib import Path
from typing import Any, Iterable
from urllib.parse import urlparse


class HttpClientError(RuntimeError):
    """Raised for HTTP client failures."""


def _should_bypass_proxy(url: str) -> bool:
    host = (urlparse(url).hostname or "").strip()
    if not host:
        return False
    if host == "localhost":
        return True
    try:
        return ip_address(host).is_loopback
    except ValueError:
        return False


def _urlopen(request: urllib.request.Request, *, timeout: float):
    if _should_bypass_proxy(request.full_url):
        opener = urllib.request.build_opener(urllib.request.ProxyHandler({}))
        return opener.open(request, timeout=timeout)
    return urllib.request.urlopen(request, timeout=timeout)


def _find_dotenv(start: Path | None = None) -> Path | None:
    current = (start or Path.cwd()).resolve()
    for candidate_dir in (current, *current.parents):
        candidate = candidate_dir / ".env"
        if candidate.is_file():
            return candidate
    return None


def _load_dotenv() -> None:
    path = _find_dotenv()
    if path is None:
        return
    for raw_line in path.read_text(encoding="utf-8").splitlines():
        line = raw_line.strip()
        if not line or line.startswith("#"):
            continue
        if line.startswith("export "):
            line = line[len("export ") :].strip()
        key, separator, value = line.partition("=")
        if not separator:
            continue
        key = key.strip()
        if not key:
            continue
        value = value.strip()
        if len(value) >= 2 and value[0] == value[-1] and value[0] in {'"', "'"}:
            value = value[1:-1]
        os.environ.setdefault(key, value)


def _load_codex_network_env() -> None:
    config_path = Path.home() / ".codex" / "config.toml"
    if not config_path.is_file():
        return
    try:
        payload = tomllib.loads(config_path.read_text(encoding="utf-8"))
    except tomllib.TOMLDecodeError:
        return
    network = payload.get("network")
    if not isinstance(network, dict):
        return
    for key, value in network.items():
        if not isinstance(value, str) or not value.strip():
            continue
        normalized = value.strip()
        os.environ.setdefault(str(key), normalized)
        if str(key).endswith("_proxy"):
            os.environ.setdefault(str(key).upper(), normalized)


def _resolve_api_key(api_key_env: str | None, *, provider: str) -> str | None:
    _load_dotenv()
    if provider == "codex":
        _load_codex_network_env()
    if api_key_env:
        api_key = os.environ.get(api_key_env)
        if api_key:
            return api_key
    if provider != "codex" and api_key_env != "OPENAI_API_KEY":
        return None
    auth_path = Path.home() / ".codex" / "auth.json"
    if not auth_path.is_file():
        return None
    try:
        payload = json.loads(auth_path.read_text(encoding="utf-8"))
    except json.JSONDecodeError:
        return None
    api_key = payload.get("OPENAI_API_KEY")
    if isinstance(api_key, str) and api_key.strip():
        return api_key.strip()
    return None


def _auth_headers(api_key_env: str | None, provider: str = "custom") -> dict[str, str]:
    headers = {"Content-Type": "application/json"}
    api_key = _resolve_api_key(api_key_env, provider=provider)
    if api_key:
        headers["Authorization"] = f"Bearer {api_key}"
    return headers


def _openai_url(base_url: str, path: str) -> str:
    root = base_url.rstrip("/")
    normalized_path = "/" + path.lstrip("/")
    if root.endswith("/v1") and normalized_path.startswith("/v1/"):
        normalized_path = normalized_path[len("/v1") :]
    return root + normalized_path


def wait_for_server(base_url: str, path: str, timeout_s: float) -> None:
    deadline = time.monotonic() + timeout_s
    url = f"{base_url.rstrip('/')}{path}"
    last_error = "server_not_ready"
    while time.monotonic() < deadline:
        try:
            request = urllib.request.Request(url=url, headers=_auth_headers(None), method="GET")
            with _urlopen(request, timeout=5) as response:
                if 200 <= response.status < 500:
                    return
        except Exception as exc:  # noqa: BLE001
            last_error = str(exc)
        time.sleep(1.0)
    raise HttpClientError(f"Timed out waiting for {url}: {last_error}")


def chat_completion(
    *,
    base_url: str,
    api_key_env: str | None,
    provider: str = "custom",
    wire_api: str = "chat.completions",
    model: str,
    messages: list[dict[str, Any]],
    timeout_s: float,
    system_prompt: str = "",
    reasoning_effort: str | None = None,
) -> dict[str, Any]:
    if wire_api == "responses":
        payload = {"model": model, "input": messages}
        if system_prompt:
            payload["instructions"] = system_prompt
        path = "/v1/responses"
    else:
        payload = {"model": model, "messages": messages}
        if system_prompt:
            payload["messages"] = [{"role": "system", "content": system_prompt}, *messages]
        if reasoning_effort:
            payload["reasoning_effort"] = reasoning_effort
        path = "/v1/chat/completions"
    data = json.dumps(payload).encode("utf-8")
    request = urllib.request.Request(
        url=_openai_url(base_url, path),
        headers=_auth_headers(api_key_env, provider),
        data=data,
        method="POST",
    )
    try:
        with _urlopen(request, timeout=timeout_s) as response:
            return json.loads(response.read().decode("utf-8"))
    except urllib.error.HTTPError as exc:
        detail = exc.read().decode("utf-8", errors="replace")
        raise HttpClientError(f"llm_completion failed: {exc.code} {detail}") from exc
    except OSError as exc:
        # TimeoutError (socket.timeout), URLError, ConnectionError all subclass OSError.
        raise HttpClientError(f"llm_completion failed: {exc}") from exc


def stream_text_completion(
    *,
    base_url: str,
    api_key_env: str | None,
    provider: str = "custom",
    wire_api: str = "chat.completions",
    model: str,
    messages: list[dict[str, Any]],
    timeout_s: float,
    system_prompt: str = "",
    reasoning_effort: str | None = None,
) -> str:
    if wire_api != "chat.completions":
        raise HttpClientError("stream_text_completion currently supports only chat.completions")
    payload: dict[str, Any] = {
        "model": model,
        "messages": messages,
        "stream": True,
    }
    if system_prompt:
        payload["messages"] = [{"role": "system", "content": system_prompt}, *messages]
    if reasoning_effort:
        payload["reasoning_effort"] = reasoning_effort
    data = json.dumps(payload).encode("utf-8")
    request = urllib.request.Request(
        url=_openai_url(base_url, "/v1/chat/completions"),
        headers=_auth_headers(api_key_env, provider),
        data=data,
        method="POST",
    )
    parts: list[str] = []
    try:
        with _urlopen(request, timeout=timeout_s) as response:
            for raw in _iter_sse_lines(response):
                if raw == "[DONE]":
                    break
                payload = json.loads(raw)
                if not isinstance(payload, dict):
                    continue
                choices = payload.get("choices")
                if not isinstance(choices, list) or not choices:
                    continue
                delta = choices[0].get("delta", {})
                if not isinstance(delta, dict):
                    continue
                content = delta.get("content")
                if isinstance(content, str):
                    parts.append(content)
    except urllib.error.HTTPError as exc:
        detail = exc.read().decode("utf-8", errors="replace")
        raise HttpClientError(f"stream_text_completion failed: {exc.code} {detail}") from exc
    except OSError as exc:
        raise HttpClientError(f"stream_text_completion failed: {exc}") from exc
    return "".join(parts)


@dataclass(frozen=True)
class StreamMetrics:
    ttft_ms: float | None
    tpot_ms: float | None
    completion_tokens: int | None
    completion_tokens_source: str = "usage"
    streamed_chunk_count: int = 0


def stream_chat_completion(
    *,
    base_url: str,
    body: dict[str, Any],
    timeout_s: float,
) -> StreamMetrics:
    data = json.dumps(body).encode("utf-8")
    request = urllib.request.Request(
        url=_openai_url(base_url, "/v1/chat/completions"),
        headers=_auth_headers(None),
        data=data,
        method="POST",
    )
    start = time.monotonic()
    first_token_at: float | None = None
    last_token_at: float | None = None
    chunk_token_count = 0
    completion_tokens: int | None = None
    completion_tokens_source = "none"
    try:
        with _urlopen(request, timeout=timeout_s) as response:
            for raw in _iter_sse_lines(response):
                if raw == "[DONE]":
                    break
                payload = json.loads(raw)
                if not isinstance(payload, dict):
                    continue
                usage = payload.get("usage")
                if isinstance(usage, dict):
                    comp = usage.get("completion_tokens")
                    if isinstance(comp, int) and comp >= 0:
                        completion_tokens = comp
                        completion_tokens_source = "usage"
                choices = payload.get("choices")
                if not isinstance(choices, list) or not choices:
                    continue
                delta = choices[0].get("delta", {})
                if not isinstance(delta, dict):
                    continue
                content = delta.get("content")
                if isinstance(content, str) and content:
                    now = time.monotonic()
                    if first_token_at is None:
                        first_token_at = now
                    last_token_at = now
                    chunk_token_count += 1
    except urllib.error.HTTPError as exc:
        detail = exc.read().decode("utf-8", errors="replace")
        raise HttpClientError(f"stream_chat_completion failed: {exc.code} {detail}") from exc
    except OSError as exc:
        # A request that exceeds request_timeout_s raises TimeoutError mid-stream;
        # treat it as a failed request (SLO miss), not a crashed trial.
        raise HttpClientError(f"stream_chat_completion failed: {exc}") from exc
    ttft_ms = None if first_token_at is None else (first_token_at - start) * 1000.0
    if completion_tokens is None and chunk_token_count > 0:
        completion_tokens = chunk_token_count
        completion_tokens_source = "stream_chunks"
    used_tokens = completion_tokens
    if (
        first_token_at is None
        or last_token_at is None
        or used_tokens is None
        or used_tokens <= 1
    ):
        tpot_ms = None
    else:
        tpot_ms = ((last_token_at - first_token_at) / max(used_tokens - 1, 1)) * 1000.0
    return StreamMetrics(
        ttft_ms=ttft_ms,
        tpot_ms=tpot_ms,
        completion_tokens=used_tokens if used_tokens is not None and used_tokens > 0 else None,
        completion_tokens_source=completion_tokens_source,
        streamed_chunk_count=chunk_token_count,
    )


def _iter_sse_lines(response: Any) -> Iterable[str]:
    for raw in response:
        line = raw.decode("utf-8", errors="replace").strip()
        if not line.startswith("data:"):
            continue
        payload = line[len("data:") :].strip()
        if payload:
            yield payload