polyllmkit 0.1.0__py3-none-any.whl

llmkit/bridge/bridge.py

@@ -0,0 +1,374 @@
+"""The dataclass-first bridge API.
+
+The reusable surface: callers build a :class:`~llmkit.bridge.config.Provider`
+(however they like — from TOML, env, hardcoded) plus a request, and call
+:func:`chat` / :func:`complete`. The flag-based CLI (:mod:`llmkit.bridge.cli`)
+is just one adapter over this; it owns no logic the API doesn't.
+
+Four adapters back :func:`chat`: the default ``openai-compatible`` HTTP path,
+``anthropic`` (native Messages API), ``google`` (native Gemini API), and
+``claude_code`` (Claude Agent SDK, chat-only). A "producer" pumps
+reasoning/content deltas into an :class:`_Emitter`; :func:`_run_chat` owns
+the sink/splitter/status lifecycle so every adapter shares identical
+failure semantics. ``complete`` (FIM) remains ``openai-compatible``-only.
+"""
+
+from __future__ import annotations
+
+import dataclasses
+import io
+import os
+import sys
+from dataclasses import dataclass
+from typing import Any, Callable, Optional, TextIO
+
+from .client import build_client
+from .config import Config, Provider
+from .sinks import close_sink, open_target
+from .stream import StreamSplitter
+
+DEFAULT_ENDPOINT = "http://localhost:11434/v1"
+DEFAULT_MAX_TOKENS = 1024
+DEFAULT_TEMPERATURE = 0.2
+DEFAULT_ADAPTER = "openai-compatible"
+
+# Every transport adapter :func:`chat` can dispatch to.
+ADAPTERS = ("openai-compatible", "anthropic", "google", "claude_code")
+
+# Adapters that can serve ``complete`` (text-completion / FIM). Only the
+# openai-compatible transport exposes a /v1/completions FIM endpoint; the
+# native chat protocols (anthropic, google) and claude_code have no FIM
+# equivalent. Kept as data so the capability is queryable — see
+# :func:`adapter_supports_complete`.
+_COMPLETE_ADAPTERS = frozenset({"openai-compatible"})
+
+
+def adapter_supports_complete(adapter: Optional[str]) -> bool:
+    """Whether ``adapter`` (``None`` → the default) can serve ``complete`` (FIM).
+
+    Queryable up front so a consumer can disable a FIM keybinding / widget for
+    an incompatible provider or profile, instead of relying on the runtime
+    rejection inside :func:`complete`."""
+    return (adapter or DEFAULT_ADAPTER) in _COMPLETE_ADAPTERS
+
+
+def provider_supports_complete(provider: Provider) -> bool:
+    """:func:`adapter_supports_complete` for a provider's configured adapter."""
+    return adapter_supports_complete(provider.adapter)
+
+
+def profile_supports_complete(config: Config[Provider], profile: str, key: str) -> bool:
+    """Whether the provider a ``profile``'s ``key`` (widget) selects can serve
+    ``complete`` (FIM). ``False`` if the key selects no provider. Lets a
+    consumer answer "is FIM available for this profile+widget?" — e.g. zsh-ai
+    disabling its FIM keybind for a profile whose ``fim`` widget points at a
+    chat-only provider — without constructing a request."""
+    name = config.select(profile, key)
+    if name is None:
+        return False
+    return provider_supports_complete(config.resolve(name))
+
+
+@dataclass
+class ChatRequest:
+    """A single chat turn. ``enable_thinking`` (auto|true|false), if set,
+    overrides the provider's configured value for this call.
+
+    ``schema`` requests *structured output*: a JSON Schema the model must fill.
+    Adapters enforce it natively — anthropic via a forced ``tool_use`` (the tool
+    input streams as content), openai-compatible via ``response_format`` — and
+    the filled JSON is what lands in the content sink, so a caller gets a
+    conformant object instead of parsing free-form prose. ``schema_name`` /
+    ``schema_description`` label the tool for the model. Thinking is forced off
+    when a schema is set (forced tool-use and extended thinking don't combine).
+    """
+
+    user: str
+    system: Optional[str] = None
+    enable_thinking: Optional[str] = None
+    schema: Optional[dict] = None
+    schema_name: str = "emit"
+    schema_description: str = ""
+
+
+@dataclass
+class CompleteRequest:
+    """A text-completion (FIM) request. ``openai-compatible`` only."""
+
+    prompt: str
+    suffix: str = ""
+    stop: tuple[str, ...] = ()
+
+
+def with_defaults(provider: Provider) -> Provider:
+    """Fill the hard defaults a producer needs onto an otherwise-partial
+    provider. Returns a base :class:`Provider` (producers read base fields
+    only). Idempotent."""
+    return Provider(
+        model=provider.model,
+        adapter=provider.adapter or DEFAULT_ADAPTER,
+        endpoint=provider.endpoint or DEFAULT_ENDPOINT,
+        api_key=provider.api_key,
+        api_key_env=provider.api_key_env,
+        max_tokens=provider.max_tokens
+        if provider.max_tokens is not None
+        else DEFAULT_MAX_TOKENS,
+        temperature=provider.temperature
+        if provider.temperature is not None
+        else DEFAULT_TEMPERATURE,
+        enable_thinking=provider.enable_thinking or "auto",
+        stop=provider.stop,
+        headers=provider.headers,
+        extra=provider.extra,
+    )
+
+
+# ── status events ───────────────────────────────────────────────────────────
+def _open_status(path: str) -> Optional[TextIO]:
+    """Open a status file (regular file OR fifo) for line-buffered writes.
+    O_RDWR | O_APPEND | O_NONBLOCK so it works against either kind without
+    blocking on fifo connect (and without truncating regular files).
+    Returns None if disabled or open fails — status writes are best-effort."""
+    if not path:
+        return None
+    try:
+        fd = os.open(path, os.O_RDWR | os.O_APPEND | os.O_NONBLOCK)
+        return os.fdopen(fd, "w", buffering=1)
+    except OSError:
+        return None
+
+
+def _write_status(fp: Optional[TextIO], event: str) -> None:
+    """Best-effort status line write. Never raises."""
+    if fp is None:
+        return
+    try:
+        fp.write(event + "\n")
+        fp.flush()
+    except OSError:
+        pass
+
+
+class _Emitter:
+    """Routes a producer's content/reasoning deltas into the splitter and
+    fires the one-shot ``streaming`` status event on the first sign of life
+    from the backend (reasoning OR content). Producers call
+    :meth:`signal_streaming` directly when a chunk arrives that carries
+    neither."""
+
+    def __init__(self, splitter: StreamSplitter, status_fp: Optional[TextIO]) -> None:
+        self._splitter = splitter
+        self._status_fp = status_fp
+        self._started = False
+
+    def signal_streaming(self) -> None:
+        if not self._started:
+            _write_status(self._status_fp, "streaming")
+            self._started = True
+
+    def feed_content(self, text: str) -> None:
+        if text:
+            self.signal_streaming()
+            self._splitter.feed_content_delta(text)
+
+    def feed_reasoning(self, text: str) -> None:
+        if text:
+            self.signal_streaming()
+            self._splitter.feed_reasoning_delta(text)
+
+
+# A producer drives the backend and pumps deltas into the emitter. It raises
+# on error (caught by _run_chat) and returns None on success.
+Producer = Callable[[Provider, ChatRequest, _Emitter], None]
+
+
+def _run_chat(
+    provider: Provider,
+    request: ChatRequest,
+    content: str | TextIO,
+    thinking: str,
+    status_file: str,
+    producer: Producer,
+) -> int:
+    """Own the sink/splitter/status lifecycle and run ``producer`` inside it.
+    Shared by every adapter so the failure semantics (status events, exit
+    codes, guaranteed ``splitter.finish()``) are identical."""
+    content_sink = open_target(content)
+    if content_sink == "inline":
+        print(
+            'llmkit: content sink cannot be "inline" '
+            "(that's a thinking-sink mode only)",
+            file=sys.stderr,
+        )
+        return 2
+    thinking_sink = open_target(thinking)
+    splitter = StreamSplitter(content_sink, thinking_sink)
+    status_fp = _open_status(status_file)
+    emitter = _Emitter(splitter, status_fp)
+
+    rc = 0
+    try:
+        producer(provider, request, emitter)
+        _write_status(status_fp, "complete")
+    except KeyboardInterrupt:
+        _write_status(status_fp, "interrupted")
+        rc = 130
+    except Exception as e:
+        _write_status(status_fp, "error")
+        print(f"llmkit: {e}", file=sys.stderr)
+        rc = 1
+    finally:
+        # ALWAYS flush — interrupts and exceptions still get the close wrap
+        # emitted, so downstream readers see a complete document.
+        try:
+            splitter.finish()
+        except Exception:
+            pass
+        close_sink(content_sink)
+        close_sink(thinking_sink)
+        if status_fp is not None:
+            try:
+                status_fp.close()
+            except OSError:
+                pass
+    return rc
+
+
+def chat(
+    provider: Provider,
+    request: ChatRequest,
+    *,
+    content: str | TextIO = "-",
+    thinking: str = "none",
+    status_file: str = "",
+) -> int:
+    """Run one chat turn. ``content`` / ``thinking`` are sink specs (``-``,
+    a path, ``none``, or — for thinking — ``inline``); ``content`` may also be an
+    already-open writable stream for in-process capture (see :func:`chat_to_str`).
+    ``status_file`` is an optional file/fifo path for
+    ``streaming``/``complete``/``error`` events.
+
+    Adapter is chosen from ``provider.adapter``. The request's
+    ``enable_thinking`` overrides the provider's configured value."""
+    provider = with_defaults(provider)
+    eff_thinking = request.enable_thinking or provider.enable_thinking or "auto"
+    provider = dataclasses.replace(provider, enable_thinking=eff_thinking)
+
+    # Adapters are imported lazily so the default openai path (and --help) never
+    # pays for an optional dependency that may not be installed.
+    producer: Producer
+    if provider.adapter == "claude_code":
+        from .adapters.claude_code import stream_claude_code
+
+        producer = stream_claude_code
+    elif provider.adapter == "anthropic":
+        from .adapters.anthropic import stream_anthropic
+
+        producer = stream_anthropic
+    elif provider.adapter == "google":
+        from .adapters.google import stream_google
+
+        producer = stream_google
+    else:
+        from .adapters.openai import stream_openai
+
+        producer = stream_openai
+    return _run_chat(provider, request, content, thinking, status_file, producer)
+
+
+class _StrSink(io.StringIO):
+    """A ``StringIO`` whose ``close()`` is a no-op, so ``_run_chat``'s
+    finally-block close leaves the buffer readable — :func:`chat_to_str` reads
+    ``getvalue()`` after the run returns."""
+
+    def close(self) -> None:  # noqa: D401 — intentional no-op
+        pass
+
+
+def chat_to_str(
+    provider: Provider,
+    request: ChatRequest,
+    *,
+    thinking: str = "none",
+    status_file: str = "",
+) -> tuple[int, str]:
+    """Run one chat turn and return ``(exit_code, content)`` in-process — the
+    content stream is captured to a string, no temp file. Failure semantics match
+    :func:`chat` (non-zero code on error); the partial content captured so far is
+    still returned."""
+    sink = _StrSink()
+    code = chat(provider, request, content=sink, thinking=thinking, status_file=status_file)
+    return code, sink.getvalue()
+
+
+def chat_structured(
+    provider: Provider,
+    request: ChatRequest,
+    *,
+    status_file: str = "",
+) -> tuple[int, Any]:
+    """Structured output: for a ``request`` carrying a ``schema``, capture the
+    model's schema-filled JSON and return ``(exit_code, parsed)``. Returns
+    ``(code, None)`` on a non-zero exit or unparseable output — the caller
+    decides whether to fall back. Thinking is off (forced tool-use excludes it)."""
+    import json
+
+    # claude_code enforces the schema via the SDK's native output_format and returns
+    # a parsed object in ResultMessage.structured_output — no text round-trip.
+    if provider.adapter == "claude_code":
+        from .adapters.claude_code import structured_claude_code
+        return structured_claude_code(with_defaults(provider), request)
+
+    code, text = chat_to_str(provider, request, thinking="none", status_file=status_file)
+    if code != 0:
+        return code, None
+    try:
+        return code, json.loads(text)
+    except (ValueError, TypeError):
+        return code, None
+
+
+def complete(
+    provider: Provider,
+    request: CompleteRequest,
+    *,
+    out: Optional[TextIO] = None,
+) -> int:
+    """Run a text-completion (FIM) request, streaming text to ``out``
+    (default stdout). ``openai-compatible`` only."""
+    out = out if out is not None else sys.stdout
+    provider = with_defaults(provider)
+    if not provider_supports_complete(provider):
+        print(
+            f"llmkit: FIM (complete) supports only the openai-compatible "
+            f"adapter, not '{provider.adapter}'. Point the provider at an "
+            f"openai-compatible backend. (Query "
+            f"llmkit.bridge.adapter_supports_complete to gate this up front.)",
+            file=sys.stderr,
+        )
+        return 2
+
+    client = build_client(provider)
+    kwargs: dict[str, Any] = dict(
+        model=provider.model,
+        prompt=request.prompt,
+        max_tokens=provider.max_tokens,
+        temperature=provider.temperature,
+    )
+    if request.suffix:
+        kwargs["suffix"] = request.suffix
+    if request.stop:
+        kwargs["stop"] = list(request.stop)
+    kwargs["stream"] = True
+    try:
+        for chunk in client.completions.create(**kwargs):
+            t = chunk.choices[0].text if chunk.choices else ""
+            if t:
+                out.write(t)
+                out.flush()
+    except KeyboardInterrupt:
+        return 130
+    except Exception as e:
+        print(f"llmkit: {e}", file=sys.stderr)
+        return 1
+    return 0

llmkit/bridge/cli.py

@@ -0,0 +1,174 @@
+"""Flag-based CLI — a thin adapter over the dataclass API in
+:mod:`llmkit.bridge.bridge`.
+
+Two subcommands:
+
+  chat       chat-completions (openai-compatible or claude_code)
+  complete   text-completions with FIM ``suffix`` (openai-compatible only)
+
+This is the shape a host like a shell plugin drives: it resolves its own
+config (zstyle, env, …) and passes the result as flags. Everything here just
+maps those flags onto a :class:`Provider` + request and calls ``chat`` /
+``complete``; no logic lives here that the API doesn't already own.
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+import sys
+from typing import List, Optional
+
+from . import bridge
+from .config import Provider, _header_item
+
+
+def _add_common(s: argparse.ArgumentParser) -> None:
+    s.add_argument("--model", required=True)
+    s.add_argument("--max-tokens", type=int, default=1024)
+    s.add_argument("--temperature", type=float, default=0.2)
+    s.add_argument(
+        "--endpoint",
+        default=os.environ.get("LLMKIT_ENDPOINT", "http://localhost:11434/v1"),
+    )
+    s.add_argument("--api-key", default="")
+    s.add_argument("--api-key-env", default="")
+    s.add_argument(
+        "--header",
+        action="append",
+        default=[],
+        metavar="KEY=VALUE",
+        help="extra HTTP request header (repeatable). Sent on every call by the "
+        "HTTP adapters (openai-compatible/anthropic/google); claude_code has no "
+        "HTTP client and ignores it. VALUE supports $VAR/${VAR} env expansion "
+        "plus the ${uuid}/${epoch} dynamic tokens (braces required; there is no "
+        "escape for a literal $). Use for gateway routing keys the SDK won't "
+        "set (e.g. --header x-opencode-request=zsh-ai for OpenCode Zen).",
+    )
+    s.add_argument(
+        "--adapter",
+        choices=["openai-compatible", "claude_code", "anthropic", "google"],
+        default=os.environ.get("LLMKIT_ADAPTER", "openai-compatible"),
+        help="transport backend: openai-compatible (HTTP endpoint, default), "
+        "anthropic (native Messages API), google (native Gemini API), or "
+        "claude_code (Claude Agent SDK; chat only). Point --endpoint at a "
+        "gateway (e.g. OpenCode Zen) to route a native protocol through it. "
+        "complete/FIM is openai-compatible-only.",
+    )
+
+
+def _parse_headers(items: List[str]) -> dict[str, str]:
+    """Parse repeated ``--header KEY=VALUE`` flags into a dict. Splits on the
+    first ``=`` so a value may itself contain ``=``. A missing ``=`` or empty
+    name is a usage error. Names are lowercased and both halves stripped (see
+    ``config._header_item``), so later duplicates win per key regardless of
+    case — which is what makes the zstyle→TOML override order reliable."""
+    out: dict[str, str] = {}
+    for item in items:
+        key, sep, val = item.partition("=")
+        key, val = _header_item(key, val)
+        if not sep or not key:
+            raise SystemExit(f"llmkit: --header expects KEY=VALUE, got {item!r}")
+        out[key] = val
+    return out
+
+
+def _provider_from_args(args: argparse.Namespace) -> Provider:
+    """Map the connection-level flags onto a Provider. enable_thinking is a
+    per-call (request) concern, so it's not set here."""
+    return Provider(
+        model=args.model,
+        adapter=args.adapter,
+        endpoint=args.endpoint,
+        api_key=args.api_key or None,
+        api_key_env=args.api_key_env or None,
+        max_tokens=args.max_tokens,
+        temperature=args.temperature,
+        headers=_parse_headers(args.header),
+    )
+
+
+def build_parser() -> argparse.ArgumentParser:
+    p = argparse.ArgumentParser(
+        prog="llmkit",
+        description="Streaming LLM bridge "
+        "(openai-compatible / anthropic / google / claude_code).",
+    )
+    sub = p.add_subparsers(dest="cmd", required=True)
+
+    chat = sub.add_parser("chat", help="chat-completions endpoint")
+    _add_common(chat)
+    chat.add_argument("--system", default="")
+    chat.add_argument("--user", required=True)
+    chat.add_argument(
+        "--enable-thinking",
+        choices=["auto", "true", "false"],
+        default="auto",
+        help="maps to chat_template_kwargs.enable_thinking (or the SDK's "
+        "thinking config for claude_code)",
+    )
+    chat.add_argument(
+        "--content",
+        default="-",
+        help="content sink: -, PATH, or none (default: -)",
+    )
+    chat.add_argument(
+        "--thinking",
+        default="none",
+        help="thinking sink: -, PATH, inline, or none (default: none). "
+        "inline merges reasoning into the content stream, separated from "
+        "real content by a blank line",
+    )
+    chat.add_argument(
+        "--status-file",
+        default="",
+        help="path (file or fifo) to receive status lines: 'streaming' on "
+        "first chunk, 'complete' on clean exit, 'error' on exception, "
+        "'interrupted' on ^C.",
+    )
+
+    comp = sub.add_parser("complete", help="text-completions endpoint (FIM)")
+    _add_common(comp)
+    comp.add_argument("--prompt", required=True)
+    comp.add_argument("--suffix", default="")
+    comp.add_argument(
+        "--stop",
+        action="append",
+        default=[],
+        help="stop token (repeatable)",
+    )
+
+    return p
+
+
+def main(argv: Optional[List[str]] = None) -> int:
+    args = build_parser().parse_args(argv)
+    provider = _provider_from_args(args)
+    if args.cmd == "chat":
+        request = bridge.ChatRequest(
+            user=args.user,
+            system=args.system or None,
+            enable_thinking=args.enable_thinking,
+        )
+        return bridge.chat(
+            provider,
+            request,
+            content=args.content,
+            thinking=args.thinking,
+            status_file=args.status_file,
+        )
+    if args.cmd == "complete":
+        creq = bridge.CompleteRequest(
+            prompt=args.prompt,
+            suffix=args.suffix,
+            stop=tuple(args.stop),
+        )
+        return bridge.complete(provider, creq)
+    return 2
+
+
+if __name__ == "__main__":
+    try:
+        sys.exit(main())
+    except BrokenPipeError:
+        sys.exit(0)

llmkit/bridge/client.py

@@ -0,0 +1,88 @@
+"""OpenAI client construction + small shared helpers.
+
+Kept tiny on purpose — the SDK does the heavy lifting; this module just
+resolves the endpoint, API key, and the ``enable_thinking`` extra-body
+the same way for every subcommand. Works off a resolved :class:`Provider`
+(see :func:`llmkit.bridge.bridge.with_defaults`).
+"""
+
+from __future__ import annotations
+
+import os
+import re
+import sys
+import time
+import uuid
+from typing import Any, Callable, Dict, Optional
+
+from .config import Provider
+
+# Dynamic ``${...}`` tokens usable in a header VALUE, generated fresh each time
+# headers are resolved (once per bridge process, i.e. per widget invocation — so
+# ${uuid} is unique per call). Env vars are expanded after these.
+_HEADER_TOKENS: Dict[str, Callable[[], str]] = {
+    "uuid": lambda: uuid.uuid4().hex,
+    "epoch": lambda: str(int(time.time())),
+}
+_TOKEN_RE = re.compile(r"\$\{(\w+)\}")
+
+
+def _expand_header_value(value: str) -> str:
+    """Expand a header value: ``${uuid}``/``${epoch}`` dynamic tokens first
+    (braces required — ``$uuid`` is NOT a token), then ``$VAR`` / ``${VAR}``
+    from the environment (unknown vars left intact). There is no escape for a
+    literal ``$``: a value whose ``$…`` happens to name an env var WILL be
+    substituted."""
+    value = _TOKEN_RE.sub(
+        lambda m: _HEADER_TOKENS[m.group(1)]() if m.group(1) in _HEADER_TOKENS else m.group(0),
+        value,
+    )
+    return os.path.expandvars(value)
+
+
+def resolved_headers(provider: Provider) -> Optional[Dict[str, str]]:
+    """The provider's custom headers with each value expanded (see
+    :func:`_expand_header_value`), or ``None`` when there are none — the shape the
+    SDKs' ``default_headers`` want. Shared by all three HTTP adapters."""
+    if not provider.headers:
+        return None
+    return {k: _expand_header_value(v) for k, v in provider.headers.items()}
+
+
+def build_client(provider: Provider) -> Any:
+    """Construct the OpenAI client. Imports inside so ``--help`` works even
+    before the dependency is installed. Return type is ``Any`` to avoid a
+    hard import of ``openai.OpenAI`` at module load."""
+    try:
+        from openai import OpenAI
+    except ImportError:
+        print(
+            "llmkit: the openai package is not installed. "
+            "Install it with `pip install llmkit[bridge]` (or `pip install "
+            "openai`).",
+            file=sys.stderr,
+        )
+        raise SystemExit(2)
+    return OpenAI(
+        base_url=provider.endpoint,
+        api_key=_resolve_api_key(provider),
+        default_headers=resolved_headers(provider),
+    )
+
+
+def _resolve_api_key(provider: Provider) -> str:
+    if provider.api_key_env:
+        return os.environ.get(provider.api_key_env, "") or "placeholder"
+    if provider.api_key:
+        return provider.api_key
+    # The SDK requires a non-empty string; local servers don't check it.
+    return "placeholder"
+
+
+def enable_thinking_extra(flag: Optional[str]) -> Dict[str, Any]:
+    """Map ``enable_thinking`` to the ``chat_template_kwargs`` extra body."""
+    if flag == "true":
+        return {"chat_template_kwargs": {"enable_thinking": True}}
+    if flag == "false":
+        return {"chat_template_kwargs": {"enable_thinking": False}}
+    return {}