Skip to content

AI Module

forge.ai — Unified AI model interface across OpenAI, Anthropic, Gemini, and Ollama.

Overview

The AI module provides a single, consistent API for completions and streaming across all major LLM providers. It handles provider-specific differences internally, so you write code once and switch providers by changing a config value.

Installation

pip install forge-kaif[openai]     # OpenAI support
pip install forge-kaif[anthropic]  # Anthropic support
pip install forge-kaif[all]        # All providers

Quick Start

import asyncio
from forge.ai import complete, Message

async def main():
    response = await complete(
        messages=[Message.user("What is the capital of France?")],
        model="gpt-4o",
    )
    print(response.content)

asyncio.run(main())

Key Features

Structured Output

from pydantic import BaseModel
from forge.ai import complete, Message

class SentimentResult(BaseModel):
    sentiment: str
    confidence: float

result = await complete(
    messages=[Message.user("Analyze: I love this product!")],
    output_schema=SentimentResult,
)
# result is a SentimentResult instance with validated fields

Streaming

from forge.ai import stream, Message

async for chunk in stream(
    messages=[Message.user("Tell me a story")],
    model="gpt-4o",
):
    print(chunk.delta, end="")

Multi-Provider Fallback

response = await complete(
    messages=[Message.user("Hello")],
    model="gpt-4o",
    fallback_models=["claude-3-5-sonnet", "gemini-1.5-pro"],
)

Token Counting & Cost Estimation

from forge.ai import TokenCounter

counter = TokenCounter()
tokens = counter.count("Hello, world")
cost = counter.estimate_cost(tokens, model="gpt-4o")

Adapter Architecture

The module uses an adapter pattern to abstract provider differences:

  • OpenAIAdapter — OpenAI GPT-4, GPT-4o, GPT-3.5
  • AnthropicAdapter — Claude 3.5 Sonnet, Claude 3 Opus
  • GeminiAdapter — Gemini 1.5 Pro, Gemini 1.5 Flash
  • OllamaAdapter — Local models via Ollama
  • MockAdapter — Fully offline for testing

API Reference

AI model abstraction module — unified interface across LLM providers.

Provides a single API surface for completions and streaming across OpenAI, Anthropic, and a fully offline MockAdapter for testing. Includes token estimation, cost calculation, fallback routing, and pre-request budget checking.

Classes

AIError

Base exception for all AI module errors.

AIModule

Methods:

get_metrics 
get_metrics() -> dict[str, Any]

Return the collected observability metrics.

AIProviderError

Raised when the provider API returns an error (auth, billing, bad request).

AllModelsFailedError

Raised when all models in the fallback chain have failed to execute.

AnthropicAdapter

Thin wrapper around the Anthropic Python SDK.

Requires anthropic to be installed. Falls back to mock behaviour when the library is unavailable.

BaseAdapter

Interface that every provider adapter must implement.

Attributes

provider abstractmethod property 
provider: str

Short label such as "openai" or "anthropic".

Methods:

complete abstractmethod async 
complete(request: CompletionRequest) -> CompletionResponse

Synchronous (non-streaming) completion.

count_tokens abstractmethod 
count_tokens(text: str) -> int

Return an estimated token count for text.

estimate_cost 
estimate_cost(usage: Usage, model: str) -> float

Return the estimated USD cost for the given usage.

stream async 
stream(_request: CompletionRequest) -> AsyncIterator[StreamChunk]

Yield content chunks as they arrive.

BudgetExceededError

Raised when the estimated token count exceeds the configured limit.

GeminiAdapter

Adapter for Google Gemini models via google-generativeai.

MockAdapter

Fully offline adapter for testing — no external API calls.

Returns a canned response and never raises rate-limit or auth errors.

ModelNotFoundError

Raised when the requested model is not available or registered.

ModelRouter

Selects an adapter based on model name and supports fallback chains.

Usage::

router = ModelRouter()
router.register("gpt-4o*", openai_adapter)
router.register("claude*", anthropic_adapter)
router.register("*", mock_adapter)

resp = await router.complete(request)

Methods:

complete async 
complete(request: CompletionRequest, fallback_models: list[str] | None = None) -> CompletionResponse

Execute the request with automatic fallback on failure.

fallback_models lists model names to try in order after the primary model fails. If None, uses the router's internal fallback adapter list.

register 
register(model_pattern: str, adapter: BaseAdapter, *, is_fallback: bool = False) -> None

Register adapter for models matching model_pattern (glob-style).

Use "*" as the pattern to match all models. When is_fallback=True the adapter is added to the fallback chain.

resolve 
resolve(model: str) -> BaseAdapter | None

Return the first adapter whose pattern matches model, or None.

OllamaAdapter

Adapter for local Ollama models via the Ollama HTTP API.

Requires httpx to be installed. Falls back to mock behaviour when the library is unavailable.

Methods:

list_models async 
list_models() -> list[str]

Fetch available model names from Ollama's /api/tags endpoint.

OpenAIAdapter

Thin wrapper around the OpenAI Python SDK.

Requires openai to be installed. Falls back to mock behaviour when the library is unavailable.

RateLimitError

Raised when the provider rate-limits the client.

StreamInterruptedError

Raised when an active stream is cut short or disconnected.

StructuredOutputError

Raised when the model fails to produce output conforming to the Pydantic schema after the maximum retry attempts have been exhausted.

TokenCounter

Estimates token usage and enforces pre-request budget limits.

Methods:

check_budget staticmethod 
check_budget(request: CompletionRequest, max_tokens_limit: int) -> None

Raise :class:TokenLimitError if the estimated request exceeds max_tokens_limit.

Catches runaway prompts before calling the adapter.

count_messages staticmethod 
count_messages(messages: list[Message]) -> int

Estimate total tokens across a list of messages.

count_tokens staticmethod 
count_tokens(text: str) -> int

Estimate the number of tokens in text.

estimate_cost staticmethod 
estimate_cost(usage: Usage, model: str) -> float

Calculate cost in USD from token counts using the pricing table.

Returns 0.0 for unknown models.

TokenLimitError

Raised when the requested prompt exceeds the model's budget or maximum context.

Functions:

complete async 

complete(messages: list[Message], model: str | None = None, max_tokens: int | None = None, temperature: float | None = None, output_schema: type[BaseModel] | None = None, fallback_models: list[str] | None = None, max_retries: int | None = None, **kwargs: Any) -> CompletionResponse | BaseModel

Convenience function to perform an AI completion request.

Uses the active ForgeRuntime context.

stream async 

stream(messages: list[Message], model: str | None = None, max_tokens: int | None = None, temperature: float | None = None, **kwargs: Any) -> AsyncIterator[StreamChunk]

Convenience function to perform an AI streaming completion request.

Uses the active ForgeRuntime context.