limbo-code 0.1.0__py3-none-any.whl

limbo/llm/factory.py

@@ -0,0 +1,46 @@
+"""LLM client factory.
+
+Selects the client implementation for the model's provider API dialect
+(``ProviderSpec.api``). New dialects (e.g. ``anthropic-messages`` for Kimi's
+coding endpoint) plug in via ``register_client`` without touching call sites.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+
+from limbo.config import Config
+from limbo.llm.anthropic_client import AnthropicMessagesClient
+from limbo.llm.catalog import (
+    API_ANTHROPIC_MESSAGES,
+    API_OPENAI_COMPLETIONS,
+    resolve_model,
+)
+from limbo.llm.client import LLMClient
+from limbo.llm.openai_client import OpenAICompatibleClient
+
+ClientFactory = Callable[[Config], LLMClient]
+
+_FACTORIES: dict[str, ClientFactory] = {
+    API_OPENAI_COMPLETIONS: OpenAICompatibleClient,
+    API_ANTHROPIC_MESSAGES: AnthropicMessagesClient,
+}
+
+
+def register_client(api: str, factory: ClientFactory) -> None:
+    """Register a client factory for a provider API dialect."""
+    _FACTORIES[api] = factory
+
+
+def create_llm_client(config: Config) -> LLMClient:
+    """Create the LLM client for the configured model's provider."""
+    spec = resolve_model(config.llm.model)
+    factory = _FACTORIES.get(spec.provider.api)
+    if factory is None:
+        supported = ", ".join(sorted(_FACTORIES))
+        raise ValueError(
+            f"No LLM client implementation for API dialect "
+            f"{spec.provider.api!r} (model {spec.id!r}, provider "
+            f"{spec.provider.id!r}). Supported dialects: {supported}."
+        )
+    return factory(config)

limbo/llm/openai_client.py

@@ -0,0 +1,261 @@
+"""OpenAI-compatible LLM client."""
+
+from __future__ import annotations
+
+import json
+import time
+import warnings
+from collections.abc import AsyncIterator
+from typing import Any
+
+from openai import AsyncOpenAI, BadRequestError
+
+from limbo.config import Config
+from limbo.llm.catalog import (
+    ModelSpec,
+    resolve_api_key,
+    resolve_base_url,
+    resolve_model,
+)
+from limbo.llm.client import RequestHook
+from limbo.models import (
+    CompletionMeta,
+    LLMEvent,
+    Message,
+    TextChunk,
+    ThinkingChunk,
+    ToolCallEvent,
+)
+
+
+class OpenAICompatibleClient:
+    """Client for OpenAI-compatible chat completions with streaming.
+
+    Adapts request parameters and stream parsing to the model's catalog spec
+    (see ``limbo.llm.catalog``): per-model ``max_tokens``, thinking-parameter
+    format (``reasoning_effort`` vs ``thinking: {type: ...}``), and
+    ``reasoning_content`` deltas/replay for reasoning models such as Kimi K3.
+    """
+
+    def __init__(self, config: Config):
+        self.config = config
+        self.spec: ModelSpec = resolve_model(config.llm.model)
+        self._client: AsyncOpenAI | None = None
+        # Whether to ask for streamed token usage. Disabled permanently after
+        # a provider rejects the stream_options parameter.
+        self._include_usage = True
+
+    @property
+    def client(self) -> AsyncOpenAI:
+        if self._client is None:
+            api_key = resolve_api_key(self.spec, self.config.llm.api_key)
+            if not api_key:
+                env = self.spec.provider.api_key_env
+                raise ValueError(
+                    f"No API key for provider {self.spec.provider.id!r}. "
+                    f"Set [llm] api_key in ~/.limbo/config.toml"
+                    + (f" or the ${env} environment variable." if env else ".")
+                )
+            self._client = AsyncOpenAI(
+                api_key=api_key,
+                base_url=resolve_base_url(self.spec, self.config.llm.base_url),
+            )
+        return self._client
+
+    async def close(self) -> None:
+        """Close the underlying HTTP client if it has been created."""
+        if self._client is not None:
+            await self._client.close()
+            self._client = None
+
+    def _thinking_params(self) -> dict[str, Any]:
+        """Build thinking-control parameters for the model's dialect."""
+        spec = self.spec
+        if not spec.reasoning or spec.thinking_format is None:
+            return {}
+        effort = self.config.llm.thinking_effort
+        if effort is None:
+            return {}
+        if spec.thinking_format == "openai":
+            value = spec.thinking_levels.get(effort)
+            if value is None:
+                warnings.warn(
+                    f"thinking_effort={effort!r} is not supported by {spec.id} "
+                    f"(supported: {sorted(spec.thinking_levels)}); ignoring.",
+                    stacklevel=2,
+                )
+                return {}
+            return {"reasoning_effort": value}
+        if spec.thinking_format == "deepseek":
+            if effort == "off":
+                if not spec.thinking_can_disable:
+                    warnings.warn(
+                        f"{spec.id} does not support disabling thinking; ignoring.",
+                        stacklevel=2,
+                    )
+                    return {}
+                return {"thinking": {"type": "disabled"}}
+            return {"thinking": {"type": "enabled"}}
+        return {}
+
+    async def chat(
+        self,
+        messages: list[Message],
+        tools: list[dict[str, Any]],
+        on_request: RequestHook | None = None,
+    ) -> AsyncIterator[LLMEvent]:
+        include_reasoning = self.spec.requires_reasoning_content
+        request_messages = [
+            _message_to_openai(m, include_reasoning=include_reasoning)
+            for m in messages
+        ]
+        kwargs: dict[str, Any] = {
+            "model": self.config.llm.model,
+            "messages": request_messages,
+            "temperature": self.config.llm.temperature,
+            "max_tokens": self.config.llm.max_tokens or self.spec.max_tokens,
+            "stream": True,
+        }
+        thinking = self._thinking_params()
+        if self.spec.thinking_format == "deepseek" and thinking:
+            # `thinking` is not an OpenAI SDK parameter; pass it through the
+            # request body untouched (Kimi K2 thinking dialect).
+            kwargs["extra_body"] = thinking
+        else:
+            kwargs.update(thinking)
+        if tools:
+            kwargs["tools"] = tools
+        if self._include_usage:
+            kwargs["stream_options"] = {"include_usage": True}
+
+        start = time.monotonic()
+        try:
+            stream = await self._create(kwargs, on_request)
+        except BadRequestError as e:
+            # Some providers reject stream_options; retry once without it and
+            # stop asking for usage on later requests.
+            if not self._include_usage or not _is_stream_options_error(e):
+                raise
+            self._include_usage = False
+            kwargs.pop("stream_options", None)
+            stream = await self._create(kwargs, on_request)
+
+        usage: dict[str, Any] | None = None
+        finish_reason: str | None = None
+        ttft: float | None = None
+        active_calls: dict[int, dict[str, Any]] = {}
+        async for chunk in stream:
+            if ttft is None:
+                ttft = time.monotonic() - start
+            if chunk.usage is not None:
+                # The final chunk carries usage when include_usage is on.
+                # model_dump keeps provider extras such as DeepSeek's
+                # prompt_cache_hit_tokens / prompt_cache_miss_tokens.
+                try:
+                    usage = chunk.usage.model_dump()
+                except Exception:  # noqa: BLE001
+                    usage = {"raw": str(chunk.usage)}
+            if not chunk.choices:
+                continue
+            if chunk.choices[0].finish_reason:
+                finish_reason = chunk.choices[0].finish_reason
+            delta = chunk.choices[0].delta
+            # Reasoning models stream thinking in a separate field; the name
+            # varies by provider (Moonshot/DeepSeek: reasoning_content).
+            reasoning = None
+            for attr in ("reasoning_content", "reasoning"):
+                value = getattr(delta, attr, None)
+                if isinstance(value, str) and value:
+                    reasoning = value
+                    break
+            if reasoning:
+                yield ThinkingChunk(text=reasoning)
+            if delta.content:
+                yield TextChunk(text=delta.content)
+
+            if delta.tool_calls:
+                for tc in delta.tool_calls:
+                    idx = tc.index
+                    if idx not in active_calls:
+                        active_calls[idx] = {
+                            "id": tc.id or "",
+                            "name": "",
+                            "arguments": "",
+                        }
+                    current = active_calls[idx]
+                    if tc.id:
+                        current["id"] = tc.id
+                    if tc.function:
+                        if tc.function.name:
+                            current["name"] += tc.function.name
+                        if tc.function.arguments:
+                            current["arguments"] += tc.function.arguments
+
+        for idx in sorted(active_calls):
+            call = active_calls[idx]
+            raw_arguments = call["arguments"]
+            try:
+                args = json.loads(raw_arguments) if raw_arguments else {}
+            except json.JSONDecodeError as e:
+                args = {
+                    "raw_arguments": raw_arguments,
+                    "parse_error": str(e),
+                }
+            yield ToolCallEvent(
+                id=call["id"] or f"call_{idx}",
+                name=call["name"],
+                arguments=args,
+            )
+
+        yield CompletionMeta(
+            usage=usage,
+            finish_reason=finish_reason,
+            ttft=ttft,
+            duration=time.monotonic() - start,
+        )
+
+    async def _create(
+        self, kwargs: dict[str, Any], on_request: RequestHook | None
+    ):
+        """Fire one HTTP attempt, reporting the exact body via the hook."""
+        if on_request is not None:
+            on_request(kwargs)
+        return await self.client.chat.completions.create(**kwargs)
+
+
+def _is_stream_options_error(error: BadRequestError) -> bool:
+    message = str(error).lower()
+    return "stream_options" in message or "include_usage" in message
+
+
+def _message_to_openai(
+    message: Message, *, include_reasoning: bool = False
+) -> dict[str, Any]:
+    m: dict[str, Any] = {"role": message.role}
+    # OpenAI requires `content` on assistant and tool messages, even when empty.
+    m["content"] = message.content or ""
+    if include_reasoning and message.role == "assistant":
+        # Kimi K3 requires reasoning_content on replayed assistant messages;
+        # send the stored thinking or an empty string.
+        m["reasoning_content"] = message.reasoning or ""
+    if message.tool_calls:
+        # The OpenAI SDK requires function.arguments to be a JSON string.
+        # Internal messages keep arguments as a dict for validation; serialize
+        # only when building the API request.
+        m["tool_calls"] = [
+            {
+                **tc,
+                "function": {
+                    **tc["function"],
+                    "arguments": (
+                        json.dumps(tc["function"]["arguments"])
+                        if isinstance(tc["function"]["arguments"], dict)
+                        else tc["function"]["arguments"]
+                    ),
+                },
+            }
+            for tc in message.tool_calls
+        ]
+    if message.tool_call_id:
+        m["tool_call_id"] = message.tool_call_id
+    return m

limbo/models.py

@@ -0,0 +1,81 @@
+"""Shared data models for Limbo."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+from pydantic import BaseModel
+
+
+class Message(BaseModel):
+    """A chat message in the conversation."""
+
+    role: str  # system | user | assistant | tool
+    content: str | None = None
+    tool_calls: list[dict[str, Any]] | None = None
+    tool_call_id: str | None = None
+    # Reasoning/thinking text produced by reasoning models. Stored so it can
+    # be replayed to APIs that require it (e.g. Kimi K3) on later turns.
+    reasoning: str | None = None
+    # Signature authenticating a replayed thinking block (Anthropic dialect).
+    reasoning_signature: str | None = None
+
+
+class ToolCall(BaseModel):
+    """A tool call requested by the LLM."""
+
+    id: str
+    name: str
+    arguments: dict[str, Any]
+
+
+class ToolResult(BaseModel):
+    """Result of executing a tool."""
+
+    success: bool
+    output: str | None = None
+    error: str | None = None
+
+
+@dataclass(frozen=True)
+class TextChunk:
+    text: str
+
+
+@dataclass(frozen=True)
+class ThinkingChunk:
+    """A streamed reasoning/thinking delta from a reasoning model."""
+
+    text: str
+    # Thinking-block signature (Anthropic dialect); carried on an empty-text
+    # chunk at the end of a thinking block.
+    signature: str | None = None
+
+
+@dataclass(frozen=True)
+class ToolCallEvent:
+    id: str
+    name: str
+    arguments: dict[str, Any]
+
+
+@dataclass(frozen=True)
+class CompletionMeta:
+    """Response-level metadata emitted once per LLM call, after all chunks.
+
+    Carries whatever the provider returned: token usage (including cache-hit
+    counters where the provider exposes them), the finish/stop reason, and
+    client-measured timing. All fields may be None — providers differ widely
+    in what they report.
+    """
+
+    usage: dict[str, Any] | None = None
+    finish_reason: str | None = None
+    # Seconds from request start to the first streamed chunk.
+    ttft: float | None = None
+    # Seconds from request start to the end of the stream.
+    duration: float | None = None
+
+
+LLMEvent = TextChunk | ThinkingChunk | ToolCallEvent | CompletionMeta

limbo/sessions.py

@@ -0,0 +1,258 @@
+"""Session storage: save, load, list, find, and export sessions.
+
+Sessions are stored as JSONL files. The first line is a metadata record
+(``{"type": "meta", ...}``) followed by one message per line. Files written
+by older versions (no meta line) still load — the id falls back to the file
+stem and the title is empty.
+
+This module has no UI dependencies.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any
+
+from pydantic import BaseModel
+
+from limbo.models import Message
+from limbo.trace import read_trace
+
+META_TYPE = "meta"
+SNAPSHOT_TYPE = "messages_snapshot"
+
+
+class SessionNotFoundError(Exception):
+    """Raised when no session matches a lookup."""
+
+
+class AmbiguousSessionError(Exception):
+    """Raised when an id prefix matches multiple sessions."""
+
+    def __init__(self, prefix: str, matches: list[str]):
+        self.prefix = prefix
+        self.matches = matches
+        super().__init__(
+            f"Session id prefix '{prefix}' is ambiguous, matches: "
+            + ", ".join(matches)
+        )
+
+
+class SessionMeta(BaseModel):
+    """Metadata describing a session (stored as the first JSONL line)."""
+
+    id: str
+    workdir: str = ""
+    model: str = ""
+    title: str = ""
+    created_at: str = ""
+    updated_at: str = ""
+    # Populated by list/load; not written to disk.
+    path: Path | None = None
+
+
+def _utc_now_iso() -> str:
+    return datetime.now(timezone.utc).isoformat(timespec="seconds")
+
+
+def derive_title(messages: list[Message], max_length: int = 50) -> str:
+    """Derive a session title from the first user message."""
+    for msg in messages:
+        if msg.role == "user" and msg.content:
+            title = " ".join(msg.content.split())
+            if len(title) > max_length:
+                title = title[: max_length - 1].rstrip() + "…"
+            return title
+    return ""
+
+
+def save_session(path: Path, meta: SessionMeta, messages: list[Message]) -> None:
+    """Atomically write the whole session (meta line + message lines)."""
+    path.parent.mkdir(parents=True, exist_ok=True, mode=0o700)
+    meta.updated_at = _utc_now_iso()
+    meta_line = json.dumps(
+        {"type": META_TYPE, **meta.model_dump(exclude={"path"})},
+        ensure_ascii=False,
+    )
+    tmp_file = path.with_suffix(".tmp")
+    file_existed = path.exists()
+    with tmp_file.open("w", encoding="utf-8") as f:
+        f.write(meta_line + "\n")
+        for msg in messages:
+            f.write(msg.model_dump_json() + "\n")
+    os.replace(tmp_file, path)
+    if not file_existed:
+        path.chmod(0o600)
+
+
+def _parse_meta(raw: dict[str, Any], fallback_id: str, path: Path) -> SessionMeta:
+    data = {k: v for k, v in raw.items() if k != "type"}
+    data.setdefault("id", fallback_id)
+    return SessionMeta(path=path, **data)
+
+
+def load_session(path: Path) -> tuple[SessionMeta, list[Message]]:
+    """Load a session file, tolerating legacy files and malformed lines."""
+    meta: SessionMeta | None = None
+    messages: list[Message] = []
+    with path.open("r", encoding="utf-8") as f:
+        for line in f:
+            line = line.strip()
+            if not line:
+                continue
+            try:
+                record = json.loads(line)
+            except json.JSONDecodeError:
+                continue
+            if record.get("type") == META_TYPE:
+                meta = _parse_meta(record, path.stem, path)
+                continue
+            try:
+                messages.append(Message(**record))
+            except ValueError:
+                continue
+    if meta is None:
+        # Legacy file without a meta line.
+        meta = SessionMeta(id=path.stem, path=path)
+    return meta, messages
+
+
+def _read_meta_line(path: Path) -> SessionMeta:
+    """Read only the first line of a session file (O(1) listing)."""
+    try:
+        with path.open("r", encoding="utf-8") as f:
+            first = f.readline().strip()
+    except OSError:
+        first = ""
+    if first:
+        try:
+            record = json.loads(first)
+            if record.get("type") == META_TYPE:
+                return _parse_meta(record, path.stem, path)
+        except (json.JSONDecodeError, ValueError):
+            pass
+    # Legacy file: fall back to stem + mtime for sorting.
+    mtime = path.stat().st_mtime
+    updated = datetime.fromtimestamp(mtime, tz=timezone.utc).isoformat(
+        timespec="seconds"
+    )
+    return SessionMeta(id=path.stem, updated_at=updated, path=path)
+
+
+def list_sessions(
+    session_dir: Path, workdir: Path | None = None
+) -> list[SessionMeta]:
+    """List sessions, most recently updated first.
+
+    When ``workdir`` is given, only sessions recorded for that directory are
+    returned.
+    """
+    if not session_dir.is_dir():
+        return []
+    wanted = str(workdir.resolve()) if workdir is not None else None
+    sessions: list[SessionMeta] = []
+    for path in session_dir.glob("*.jsonl"):
+        try:
+            meta = _read_meta_line(path)
+        except OSError:
+            continue
+        if wanted is not None and meta.workdir != wanted:
+            continue
+        sessions.append(meta)
+    sessions.sort(key=lambda m: m.updated_at, reverse=True)
+    return sessions
+
+
+def latest_session(
+    session_dir: Path, workdir: Path | None = None
+) -> Path | None:
+    """Return the path of the most recently updated session, if any."""
+    sessions = list_sessions(session_dir, workdir=workdir)
+    if not sessions:
+        return None
+    return sessions[0].path
+
+
+def find_session(session_dir: Path, id_prefix: str) -> Path:
+    """Find a session by exact id or unique id prefix."""
+    if not session_dir.is_dir():
+        raise SessionNotFoundError(f"No sessions in {session_dir}")
+    matches = sorted(p.stem for p in session_dir.glob(f"{id_prefix}*.jsonl"))
+    if not matches:
+        raise SessionNotFoundError(f"No session matching '{id_prefix}'")
+    if len(matches) > 1:
+        raise AmbiguousSessionError(id_prefix, matches)
+    return session_dir / f"{matches[0]}.jsonl"
+
+
+def export_markdown(
+    meta: SessionMeta, messages: list[Message], path: Path
+) -> None:
+    """Export the conversation (user/assistant text) as a Markdown file."""
+    lines = [
+        f"# {meta.title or meta.id}",
+        "",
+        f"- workdir: `{meta.workdir}`",
+        f"- model: `{meta.model}`",
+        f"- updated: {meta.updated_at}",
+        "",
+    ]
+    for msg in messages:
+        if msg.role == "user" and msg.content:
+            lines += ["## User", "", msg.content, ""]
+        elif msg.role == "assistant" and msg.content:
+            lines += ["## Assistant", "", msg.content, ""]
+    path.parent.mkdir(parents=True, exist_ok=True)
+    path.write_text("\n".join(lines), encoding="utf-8")
+
+
+def export_jsonl(
+    meta: SessionMeta,
+    messages: list[Message],
+    path: Path,
+    trace_path: Path | None = None,
+) -> None:
+    """Export the full-fidelity session log as a single JSONL file.
+
+    Layout:
+
+    1. a ``meta`` record (session metadata + export timestamp),
+    2. every trace record in chronological order (LLM request bodies,
+       usage, tool calls/results, errors) when the trace
+       file exists — otherwise the raw conversation messages,
+    3. a final ``messages_snapshot`` record with the exact message history
+       as persisted.
+    """
+    meta_line = json.dumps(
+        {
+            "type": META_TYPE,
+            **meta.model_dump(exclude={"path"}),
+            "exported_at": _utc_now_iso(),
+        },
+        ensure_ascii=False,
+    )
+    path.parent.mkdir(parents=True, exist_ok=True)
+    with path.open("w", encoding="utf-8") as f:
+        f.write(meta_line + "\n")
+        trace_records = read_trace(trace_path) if trace_path is not None else []
+        if trace_records:
+            for record in trace_records:
+                f.write(json.dumps(record, ensure_ascii=False) + "\n")
+        else:
+            # Sessions from before tracing existed still export their messages.
+            for msg in messages:
+                f.write(msg.model_dump_json() + "\n")
+        snapshot = json.dumps(
+            {
+                "type": SNAPSHOT_TYPE,
+                "count": len(messages),
+                "messages": [
+                    json.loads(msg.model_dump_json()) for msg in messages
+                ],
+            },
+            ensure_ascii=False,
+        )
+        f.write(snapshot + "\n")