python-infrakit-dev 0.1.0__py3-none-any.whl

infrakit/llm/llm_readme.md

@@ -0,0 +1,306 @@
+# infrakit.llm
+
+Unified LLM client for OpenAI and Gemini with key rotation, quota tracking,
+rate limiting, and batch processing. Designed for free-tier API keys where
+RPM/daily limits matter.
+
+---
+
+## Installation
+
+```bash
+# core
+pip install pydantic tqdm
+
+# for OpenAI
+pip install openai
+
+# for Gemini
+pip install google-generativeai
+```
+
+---
+
+## File structure
+
+```
+infrakit/llm/
+├── __init__.py          Public API surface
+├── client.py            LLMClient — main entry point
+├── key_manager.py       Key rotation, quota, persistence
+├── rate_limiter.py      RPM / TPM async and sync gates
+├── batch.py             Async + threaded batch engine
+├── models.py            Shared types (Prompt, LLMResponse, etc.)
+└── providers/
+    ├── __init__.py
+    ├── base.py          AbstractProvider + schema validation
+    ├── openai.py        OpenAI provider
+    └── gemini.py        Gemini provider
+
+infrakit/cli/
+└── llm_cmd.py           CLI commands (ik llm status, ik llm quota set)
+```
+
+---
+
+## Quick start
+
+```python
+from infrakit.llm import LLMClient, Prompt
+
+client = LLMClient(
+    keys={
+        "openai_keys": ["sk-key1", "sk-key2"],
+        "gemini_keys": ["AIza-key1"],
+    },
+    storage_dir="./logs",   # key state persisted here
+)
+
+# simple prompt
+response = client.generate(Prompt(user="What is 2 + 2?"), provider="openai")
+print(response.content)
+
+# system + user split
+response = client.generate(
+    Prompt(system="You are a maths tutor.", user="Explain derivatives."),
+    provider="gemini",
+)
+print(response.content)
+```
+
+---
+
+## Structured output (Pydantic)
+
+Pass a Pydantic model as `response_model`. The system will try to parse the
+model's JSON response and validate it against your schema. If validation fails
+after retries, you still get the raw `content` back with `schema_matched=False`.
+
+```python
+from pydantic import BaseModel
+
+class Sentiment(BaseModel):
+    label: str        # "positive" | "negative" | "neutral"
+    confidence: float
+
+# Your prompt must instruct the model to return JSON — infrakit does not
+# inject instructions automatically.
+response = client.generate(
+    Prompt(
+        system='Respond ONLY with valid JSON matching: {"label": str, "confidence": float}',
+        user="I love this product!",
+    ),
+    provider="openai",
+    response_model=Sentiment,
+)
+
+if response.schema_matched:
+    print(response.parsed.label)       # "positive"
+    print(response.parsed.confidence)  # 0.97
+else:
+    print("Schema mismatch — raw response:", response.content)
+```
+
+---
+
+## Batch processing
+
+```python
+from infrakit.llm import Prompt
+
+words = ["cat", "dog", "bird", "fish"]
+prompts = [
+    Prompt(
+        system="Translate to French. Reply with only the translation.",
+        user=word,
+    )
+    for word in words
+]
+
+batch = client.batch_generate(prompts, provider="gemini")
+
+# results are in the same order as prompts
+for word, result in zip(words, batch.results):
+    if result.error:
+        print(f"{word}: ERROR — {result.error}")
+    else:
+        print(f"{word}: {result.content}")
+
+print(f"\nTotal tokens: {batch.total_tokens}")
+print(f"Success: {batch.success_count} / Failure: {batch.failure_count}")
+```
+
+---
+
+## LLMClient parameters
+
+| Parameter | Default | Description |
+|---|---|---|
+| `keys` | required | `{"openai_keys": [...], "gemini_keys": [...]}` |
+| `storage_dir` | required | Folder for persistent key state |
+| `mode` | `"async"` | `"async"` or `"threaded"` |
+| `max_concurrent` | `3` | Max simultaneous batch requests |
+| `key_retries` | `2` | Retries on same key before rotating |
+| `schema_retries` | `2` | JSON parse/validate retries |
+| `meta_window` | `50` | Recent request metadata records per key |
+| `openai_model` | `"gpt-4o-mini"` | Default OpenAI model |
+| `gemini_model` | `"gemini-1.5-flash"` | Default Gemini model |
+| `show_progress` | `True` | tqdm progress bar for batch calls |
+
+---
+
+## Quota management
+
+### Set quota for a key
+
+```python
+from infrakit.llm import QuotaConfig
+
+client.set_quota(
+    provider="openai",
+    key_id="sk-abc123",          # first 8 chars of the key
+    quota=QuotaConfig(
+        rpm_limit=60,            # 60 requests per minute
+        tpm_limit=90_000,        # 90k tokens per minute
+        daily_token_limit=1_000_000,
+        reset_hour_utc=0,        # resets at midnight UTC
+    ),
+)
+```
+
+Only the fields you set are enforced. Unset fields are unconstrained.
+
+### Key lifecycle
+
+- **active** — key is in normal use.
+- **inactive** — key hit a quota limit (daily tokens, or a hard 401/429 from the API).
+  Automatically reactivates at `reset_hour_utc` the following day.
+- Keys are rotated round-robin among all active keys. When a key is deactivated,
+  remaining active keys absorb the load.
+
+---
+
+## Check status (Python)
+
+```python
+# all keys
+client.print_status()
+
+# one provider
+client.print_status(provider="openai")
+
+# one key
+client.print_status(provider="openai", key_id="sk-abc123")
+
+# raw dict (for programmatic use)
+rows = client.status(provider="openai")
+```
+
+---
+
+## CLI
+
+Register `llm_cmd` in your main CLI (add to `infrakit/cli/__init__.py`):
+
+```python
+from infrakit.cli.llm_cmd import register as register_llm
+register_llm(cli)
+```
+
+### Commands
+
+```bash
+# show all keys
+ik llm status --storage-dir ./logs
+
+# filter by provider
+ik llm status --provider openai --storage-dir ./logs
+
+# filter by key (first 8 chars)
+ik llm status --key sk-abc123 --storage-dir ./logs
+
+# JSON output
+ik llm status --json --storage-dir ./logs
+
+# set quota for a key
+ik llm quota set \
+  --provider openai \
+  --key sk-abc123 \
+  --rpm 60 \
+  --tpm 90000 \
+  --daily 1000000 \
+  --reset-hour 0 \
+  --storage-dir ./logs
+```
+
+If your keys are stored in a JSON file:
+```bash
+ik llm status --keys-file keys.json --storage-dir ./logs
+```
+
+`keys.json` format:
+```json
+{
+  "openai_keys": ["sk-key1", "sk-key2"],
+  "gemini_keys": ["AIza-key1"]
+}
+```
+
+---
+
+## Error handling
+
+Every `generate()` call returns an `LLMResponse` — it never raises. Check `.error`:
+
+```python
+result = client.generate(Prompt(user="Hello"), provider="openai")
+if result.error:
+    print(f"Request failed: {result.error}")
+else:
+    print(result.content)
+```
+
+### What infrakit handles automatically
+
+| Situation | Behaviour |
+|---|---|
+| Transient error (network, 5xx) | Retry same key up to `key_retries` times with backoff |
+| Quota / auth error (401, 402, 429+quota) | Deactivate key immediately, rotate to next |
+| All keys exhausted | Return `LLMResponse(error="All keys exhausted.")` |
+| RPM limit reached | Async/sync sleep until slot opens |
+| Daily token limit reached | Deactivate key, auto-reactivate at reset hour |
+| Schema validation fails | Return raw content + `schema_matched=False` |
+
+---
+
+## LLMResponse fields
+
+| Field | Type | Description |
+|---|---|---|
+| `content` | `str` | Raw text from the model |
+| `parsed` | `BaseModel \| None` | Validated Pydantic instance (if `response_model` given and matched) |
+| `schema_matched` | `bool` | True if structured output validated successfully |
+| `provider` | `str` | `"openai"` or `"gemini"` |
+| `model` | `str` | Model string used |
+| `key_id` | `str` | First 8 chars of the key used |
+| `input_tokens` | `int` | Prompt token count |
+| `output_tokens` | `int` | Completion token count |
+| `total_tokens` | `int` | input + output |
+| `latency_ms` | `float` | Wall-clock API call time |
+| `error` | `str \| None` | Set on failure |
+
+---
+
+## What is and isn't stored
+
+infrakit stores **no prompt or response content**. Per key, it persists:
+
+- Status (active/inactive), deactivation timestamp
+- Quota config (RPM, TPM, daily limit, reset hour)
+- Lifetime totals (requests, tokens in/out, errors)
+- Daily token total + day start epoch
+- RPM window (timestamps of last 60 s of requests)
+- TPM window (timestamps + token counts for last 60 s)
+- Rolling metadata for last N requests: timestamp, model, token counts, latency, success/error code
+
+All stored in `{storage_dir}/llm_key_state.json`.

infrakit/llm/models.py

@@ -0,0 +1,148 @@
+"""
+infrakit.llm.models
+-------------------
+Shared data structures for the LLM subsystem.
+"""
+
+from __future__ import annotations
+
+import time
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Any, Optional, Type
+
+from pydantic import BaseModel
+
+
+# ── enums ──────────────────────────────────────────────────────────────────
+
+class Provider(str, Enum):
+    OPENAI = "openai"
+    GEMINI = "gemini"
+
+
+class KeyStatus(str, Enum):
+    ACTIVE   = "active"
+    INACTIVE = "inactive"   # all models exhausted — auto-reactivates after reset
+
+
+class ModelStatus(str, Enum):
+    ACTIVE   = "active"
+    INACTIVE = "inactive"   # this model's quota exhausted on this key
+
+
+# ── prompt input ───────────────────────────────────────────────────────────
+
+@dataclass
+class Prompt:
+    """
+    Represents a single LLM prompt.
+
+    Usage::
+
+        # combined prompt
+        Prompt(user="Tell me about Python.")
+
+        # system + user split
+        Prompt(system="You are a helpful assistant.", user="Tell me about Python.")
+    """
+    user: str
+    system: Optional[str] = None
+
+
+# ── response ───────────────────────────────────────────────────────────────
+
+@dataclass
+class LLMResponse:
+    """
+    Returned by every generate() call.
+
+    Attributes
+    ----------
+    content         Raw text from the model.
+    parsed          Populated when a response_model is given and validation
+                    succeeds; None otherwise.
+    schema_matched  True if parsed is not None. False means validation failed
+                    after all retries — content still contains the raw reply.
+    provider        Which provider handled this request.
+    model           Model string used (e.g. "gpt-4o-mini").
+    key_id          Truncated key identifier (first 8 chars).
+    input_tokens    Prompt token count.
+    output_tokens   Completion token count.
+    total_tokens    input + output.
+    latency_ms      Wall-clock time for the API call in milliseconds.
+    error           Set when the request ultimately failed.
+    """
+    content: str
+    parsed: Optional[Any]
+    schema_matched: bool
+    provider: str
+    model: str
+    key_id: str
+    input_tokens: int
+    output_tokens: int
+    total_tokens: int
+    latency_ms: float
+    error: Optional[str] = None
+
+
+# ── request metadata (stored for transparency) ────────────────────────────
+
+@dataclass
+class RequestMeta:
+    """
+    Lightweight record of one API call stored in the rolling window.
+    NO prompt or response content is kept here.
+    """
+    timestamp: float = field(default_factory=time.time)
+    provider: str = ""
+    key_id: str = ""
+    model: str = ""
+    input_tokens: int = 0
+    output_tokens: int = 0
+    total_tokens: int = 0
+    latency_ms: float = 0.0
+    success: bool = True
+    error: Optional[str] = None
+
+
+# ── quota config ───────────────────────────────────────────────────────────
+
+@dataclass
+class QuotaConfig:
+    """
+    Quota limits for a key, optionally scoped to a specific model.
+
+    Fields
+    ------
+    model               Model this config applies to.  ``None`` means it is a
+                        default that applies to any model without an explicit
+                        entry.  When both a model-specific config and a default
+                        exist, the model-specific one wins.
+    rpm_limit           Max requests per minute (key-level, shared across models).
+    tpm_limit           Max tokens per minute for this model.
+    daily_token_limit   Max tokens per calendar day for this model.
+    reset_hour_utc      UTC hour (0-23) at which the daily quota resets.
+    """
+    model: Optional[str] = None          # None = default / applies to all
+    rpm_limit: Optional[int] = None      # key-level
+    tpm_limit: Optional[int] = None      # model-level
+    daily_token_limit: Optional[int] = None   # model-level
+    reset_hour_utc: int = 0
+
+
+# ── batch result ───────────────────────────────────────────────────────────
+
+@dataclass
+class BatchResult:
+    """
+    Container for a batch generate() call. Results are in the same
+    order as the input prompts.
+    """
+    results: list[Optional[LLMResponse]]
+    total_input_tokens: int
+    total_output_tokens: int
+    total_tokens: int
+    total_latency_ms: float
+    success_count: int
+    failure_count: int

infrakit/llm/providers/__init__.py

@@ -0,0 +1,5 @@
+from .base import BaseProvider
+from .openai import OpenAIProvider
+from .gemini import GeminiProvider
+
+__all__ = ["BaseProvider", "OpenAIProvider", "GeminiProvider"]

infrakit/llm/providers/base.py

@@ -0,0 +1,112 @@
+"""
+infrakit.llm.providers.base
+---------------------------
+Abstract base class that every provider must implement.
+"""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import Any, Optional, Type
+
+from pydantic import BaseModel
+
+from ..models import LLMResponse, Prompt
+
+
+class BaseProvider(ABC):
+    """
+    All provider-specific logic (API call, token counting, error
+    classification) lives in subclasses.  The client only talks to
+    this interface.
+    """
+
+    # ── configuration ──────────────────────────────────────────────────────
+
+    #: Override in subclass with a sensible default (e.g. "gpt-4o-mini").
+    DEFAULT_MODEL: str = ""
+
+    def __init__(self, model: Optional[str] = None) -> None:
+        self.model = model or self.DEFAULT_MODEL
+
+    # ── abstract interface ─────────────────────────────────────────────────
+
+    @abstractmethod
+    async def async_generate(
+        self,
+        prompt: Prompt,
+        api_key: str,
+        response_model: Optional[Type[BaseModel]] = None,
+        schema_retries: int = 2,
+        **kwargs: Any,
+    ) -> LLMResponse:
+        """
+        Async generate.  Must return an LLMResponse even on soft failures
+        (schema mismatch).  Hard failures (network, auth) should raise.
+        """
+
+    @abstractmethod
+    def sync_generate(
+        self,
+        prompt: Prompt,
+        api_key: str,
+        response_model: Optional[Type[BaseModel]] = None,
+        schema_retries: int = 2,
+        **kwargs: Any,
+    ) -> LLMResponse:
+        """Sync wrapper around async_generate (or a native sync call)."""
+
+    # ── helpers shared by subclasses ───────────────────────────────────────
+
+    @staticmethod
+    def _validate_schema(
+        content: str,
+        response_model: Type[BaseModel],
+        retries: int,
+    ) -> tuple[Optional[BaseModel], bool]:
+        """
+        Try to parse *content* as JSON and validate against *response_model*.
+
+        Returns
+        -------
+        (parsed_instance, matched)
+            matched is True on success, False after all retries exhausted.
+        """
+        import json
+
+        for attempt in range(retries + 1):
+            try:
+                # strip common markdown fences
+                text = content.strip()
+                if text.startswith("```"):
+                    lines = text.splitlines()
+                    # drop first (```json) and last (```) lines
+                    text = "\n".join(lines[1:-1]) if len(lines) > 2 else text
+                data = json.loads(text)
+                instance = response_model.model_validate(data)
+                return instance, True
+            except Exception:
+                if attempt == retries:
+                    return None, False
+        return None, False
+
+    @staticmethod
+    def _is_quota_error(exc: Exception) -> bool:
+        """
+        Return True if *exc* indicates a hard quota / auth failure
+        (should deactivate the key, not retry).
+        Subclasses may override for provider-specific error codes.
+        """
+        msg = str(exc).lower()
+        return any(
+            kw in msg
+            for kw in (
+                "quota",
+                "rate_limit_exceeded",
+                "billing",
+                "insufficient_quota",
+                "resource_exhausted",
+                "invalid_api_key",
+                "permission_denied",
+            )
+        )