pi_llm: LiteLLM stream, tools, credentials, errors

packages/pi_llm/pyproject.toml

@@ -3,7 +3,9 @@ name = "pi_llm"
 version = "0.1.0"
 description = "Thin LiteLLM adapter for the pi-python agent runtime"
 requires-python = ">=3.11"
-dependencies = []
+dependencies = [
+    "litellm>=1.93.0",
+]
 
 [build-system]
 requires = ["hatchling"]

packages/pi_llm/src/pi_llm/__init__.py

@@ -1,3 +1,38 @@
-"""Thin LLM adapter (LiteLLM-backed)."""
+"""Thin LLM adapter (LiteLLM-backed OpenAI Chat Completions)."""
 
-__all__: list[str] = []
+from pi_llm.capabilities import supports_parallel_tools, supports_tools
+from pi_llm.complete import complete
+from pi_llm.credentials import (
+    apply_credentials_to_environ,
+    default_auth_path,
+    resolve_credentials,
+)
+from pi_llm.errors import ErrorKind, LLMError, map_provider_error
+from pi_llm.stream import stream
+from pi_llm.types import (
+    AssistantMessage,
+    StreamEvent,
+    TextDelta,
+    ToolCall,
+    ToolCallDelta,
+    TurnFinished,
+)
+
+__all__ = [
+    "AssistantMessage",
+    "ErrorKind",
+    "LLMError",
+    "StreamEvent",
+    "TextDelta",
+    "ToolCall",
+    "ToolCallDelta",
+    "TurnFinished",
+    "apply_credentials_to_environ",
+    "complete",
+    "default_auth_path",
+    "map_provider_error",
+    "resolve_credentials",
+    "stream",
+    "supports_parallel_tools",
+    "supports_tools",
+]

packages/pi_llm/src/pi_llm/capabilities.py

@@ -0,0 +1,37 @@
+"""Model capability probes (tools / parallel tools)."""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+
+Probe = Callable[[str], bool]
+
+
+def supports_tools(model: str, *, probe: Probe | None = None) -> bool:
+    """Return whether `model` supports structured function/tool calling."""
+    fn = probe or _default_supports_function_calling
+    try:
+        return bool(fn(model))
+    except Exception:
+        return False
+
+
+def supports_parallel_tools(model: str, *, probe: Probe | None = None) -> bool:
+    """Return whether `model` supports parallel tool calls."""
+    fn = probe or _default_supports_parallel_function_calling
+    try:
+        return bool(fn(model))
+    except Exception:
+        return False
+
+
+def _default_supports_function_calling(model: str) -> bool:
+    from litellm import supports_function_calling
+
+    return bool(supports_function_calling(model=model))
+
+
+def _default_supports_parallel_function_calling(model: str) -> bool:
+    from litellm import supports_parallel_function_calling
+
+    return bool(supports_parallel_function_calling(model=model))

packages/pi_llm/src/pi_llm/complete.py

@@ -0,0 +1,63 @@
+"""Non-stream completion turn via LiteLLM `acompletion`."""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from typing import Any
+
+from pi_llm.errors import map_provider_error
+from pi_llm.stream import ACompletion, _default_acompletion
+from pi_llm.types import AssistantMessage, ToolCall, TurnFinished
+
+
+async def complete(
+    request: Mapping[str, Any],
+    *,
+    acompletion: ACompletion | None = None,
+) -> TurnFinished:
+    """One non-stream chat-completion turn (modern `tools` / `tool_choice` ok)."""
+    call = acompletion or _default_acompletion
+    kwargs = dict(request)
+    kwargs["stream"] = False
+
+    try:
+        response = await call(**kwargs)
+    except Exception as exc:
+        raise map_provider_error(exc) from exc
+
+    choices = getattr(response, "choices", None) or []
+    if not choices:
+        return TurnFinished(message=AssistantMessage(content=None), finish_reason=None)
+
+    choice = choices[0]
+    message = getattr(choice, "message", None)
+    content = getattr(message, "content", None) if message is not None else None
+    raw_tools = getattr(message, "tool_calls", None) if message is not None else None
+
+    tool_calls: list[ToolCall] = []
+    for i, tc in enumerate(raw_tools or []):
+        fn = getattr(tc, "function", None)
+        tool_calls.append(
+            ToolCall(
+                id=str(getattr(tc, "id", None) or f"call_{i}"),
+                name=str(getattr(fn, "name", "") if fn is not None else ""),
+                arguments=str(getattr(fn, "arguments", "") if fn is not None else ""),
+            )
+        )
+
+    usage = getattr(response, "usage", None)
+    usage_dict: dict[str, Any] | None = None
+    if isinstance(usage, Mapping):
+        usage_dict = dict(usage)
+    elif usage is not None:
+        usage_dict = {
+            key: getattr(usage, key)
+            for key in ("prompt_tokens", "completion_tokens", "total_tokens")
+            if getattr(usage, key, None) is not None
+        }
+
+    return TurnFinished(
+        message=AssistantMessage(content=content, tool_calls=tool_calls),
+        finish_reason=getattr(choice, "finish_reason", None),
+        usage=usage_dict,
+    )

packages/pi_llm/src/pi_llm/credentials.py

@@ -0,0 +1,77 @@
+"""Credential resolution: process env, then optional auth.json."""
+
+from __future__ import annotations
+
+import json
+import os
+from collections.abc import Mapping, MutableMapping
+from pathlib import Path
+from typing import Any
+
+
+def default_auth_path() -> Path:
+    return Path.home() / ".pi" / "agent" / "auth.json"
+
+
+def resolve_credentials(
+    *,
+    environ: Mapping[str, str] | None = None,
+    auth_path: Path | None = None,
+) -> dict[str, str]:
+    """Merge credentials with env taking precedence over `auth.json`.
+
+    `auth.json` is a flat object of string keys/values (env-var style).
+    Returns only credential-like keys (from the file and/or env).
+    """
+    env = dict(os.environ if environ is None else environ)
+    path = default_auth_path() if auth_path is None else auth_path
+    file_creds = _load_auth_file(path)
+
+    keys = set(file_creds) | {k for k in env if _is_credential_key(k)}
+    out: dict[str, str] = {}
+    for key in keys:
+        if key in env and isinstance(env[key], str):
+            out[key] = env[key]
+        elif key in file_creds:
+            out[key] = file_creds[key]
+    return out
+
+
+def apply_credentials_to_environ(
+    *,
+    environ: MutableMapping[str, str] | None = None,
+    auth_path: Path | None = None,
+) -> dict[str, str]:
+    """Resolve credentials and `setdefault` them into `environ` (process env by default)."""
+    target: MutableMapping[str, str] = os.environ if environ is None else environ
+    resolved = resolve_credentials(environ=dict(target), auth_path=auth_path)
+    for key, value in resolved.items():
+        target.setdefault(key, value)
+    return resolved
+
+
+def _is_credential_key(key: str) -> bool:
+    upper = key.upper()
+    return (
+        upper.endswith("_API_KEY")
+        or upper.endswith("_API_TOKEN")
+        or upper.endswith("_ACCESS_TOKEN")
+        or upper.startswith("LITELLM_")
+        or upper in {"OPENAI_API_BASE", "OPENAI_BASE_URL"}
+    )
+
+
+def _load_auth_file(path: Path) -> dict[str, str]:
+    if not path.is_file():
+        return {}
+    try:
+        raw: Any = json.loads(path.read_text(encoding="utf-8"))
+    except (OSError, json.JSONDecodeError):
+        return {}
+    if not isinstance(raw, dict):
+        return {}
+    out: dict[str, str] = {}
+    for key, value in raw.items():
+        if isinstance(key, str) and isinstance(value, str):
+            out[key] = value
+    return out

packages/pi_llm/src/pi_llm/errors.py

@@ -0,0 +1,92 @@
+"""OpenAI-shaped error surface for provider failures."""
+
+from __future__ import annotations
+
+from enum import Enum
+from typing import Any
+
+
+class ErrorKind(Enum):
+    """Coarse classification for retry / UX decisions."""
+
+    AUTH = "auth"
+    CONTEXT_WINDOW = "context_window"
+    RATE_LIMIT = "rate_limit"
+    RETRYABLE = "retryable"
+    FATAL = "fatal"
+
+
+class LLMError(Exception):
+    """Thin, OpenAI-shaped error for callers (auth, rate limit, bad request, …)."""
+
+    def __init__(
+        self,
+        message: str,
+        *,
+        status_code: int | None = None,
+        type: str | None = None,
+        code: str | None = None,
+        llm_provider: str | None = None,
+        kind: ErrorKind = ErrorKind.FATAL,
+        cause: BaseException | None = None,
+    ) -> None:
+        super().__init__(message)
+        self.message = message
+        self.status_code = status_code
+        self.type = type
+        self.code = code
+        self.llm_provider = llm_provider
+        self.kind = kind
+        self.__cause__ = cause
+
+
+def map_provider_error(exc: BaseException) -> LLMError:
+    """Map a LiteLLM/OpenAI-style exception into `LLMError`."""
+    if isinstance(exc, LLMError):
+        return exc
+
+    message = str(exc) or exc.__class__.__name__
+    status_code = _attr_int(exc, "status_code")
+    code = _attr_str(exc, "code")
+    err_type = _attr_str(exc, "type") or _attr_str(exc, "error_type")
+    return LLMError(
+        message,
+        status_code=status_code,
+        type=err_type,
+        code=code,
+        llm_provider=_attr_str(exc, "llm_provider"),
+        kind=_classify(exc, status_code=status_code, code=code, err_type=err_type),
+        cause=exc,
+    )
+
+
+def _classify(
+    exc: BaseException,
+    *,
+    status_code: int | None,
+    code: str | None,
+    err_type: str | None,
+) -> ErrorKind:
+    name = exc.__class__.__name__.lower()
+    blob = " ".join(x for x in (name, code or "", err_type or "", str(exc).lower()) if x)
+
+    if status_code in {401, 403} or "auth" in blob or "permission" in blob:
+        return ErrorKind.AUTH
+    if "contextwindow" in name or "context_window" in blob or "context length" in blob:
+        return ErrorKind.CONTEXT_WINDOW
+    if status_code == 429 or "ratelimit" in name or "rate_limit" in blob:
+        return ErrorKind.RATE_LIMIT
+    retry_names = ("timeout", "serviceunavailable", "apiconnection", "internalserver")
+    if status_code in {408, 500, 502, 503, 504} or any(token in name for token in retry_names):
+        return ErrorKind.RETRYABLE
+    return ErrorKind.FATAL
+
+
+def _attr_str(exc: BaseException, name: str) -> str | None:
+    value: Any = getattr(exc, name, None)
+    return value if isinstance(value, str) else None
+
+
+def _attr_int(exc: BaseException, name: str) -> int | None:
+    value: Any = getattr(exc, name, None)
+    return value if isinstance(value, int) else None