Skip to content

Large Language Models (LLM)

Overview

The athomic.ai.llm module provides a unified, provider-agnostic interface for interacting with Large Language Models (LLMs) such as OpenAI (GPT-4) and Google Vertex AI (Gemini). It is designed to solve common challenges in production AI applications:

  • Vendor Lock-in: Switch providers via configuration without changing code.
  • Real-Time Streaming: Native support for asynchronous token streaming with standardized chunks.
  • Structured Data: Reliably extract Pydantic models from LLM outputs using a standardized API (generate_structured).
  • Observability: Automatic tracing (OpenTelemetry) and metrics (Prometheus) for every call (unary or streaming).
  • Embedded Governance: Enforce rate limits and safety checks directly within the provider pipeline.
  • Resilient Routing: Dynamically route requests across fallback chains using transparent wrappers, ensuring high availability without changing business logic.

Core Concepts

LLMProviderProtocol

The contract that all LLM providers must implement. It defines three primary operations:

  • generate_content(prompt, ...): Generates unstructured text (blocking).
  • stream_content(prompt, ...): Generates unstructured text incrementally (async iterator).
  • generate_structured(prompt, response_model, ...): Generates a structured object adhering to a specific Pydantic schema.

BaseLLM & Governance

The abstract base class now handles AI Governance automatically: 1. Input Guards: Executed before the request is sent to the provider. These are blocking (e.g., Rate Limiting, Prompt Injection check). 2. Output Guards: Executed after the response is received. For streaming, these currently act in audit mode.

Usage Example

Basic Text Generation (Blocking)

from nala.athomic.ai.llm.factory import LLMFactory
from nala.athomic.config import get_settings

async def chat_with_ai(user_input: str):
    # 1. Create the provider from global settings
    settings = get_settings().ai.llms.connections["default"]
    llm = LLMFactory.create(settings)

    # 2. Generate content (waits for full completion)
    response = await llm.generate_content(
        prompt=user_input,
        system_message="You are a helpful assistant."
    )
    return response.content

Real-Time Streaming

import sys
from nala.athomic.ai.llm.factory import LLMFactory

async def stream_chat(user_input: str):
    llm = LLMFactory.create_default()

    # stream_content yields LLMResponseChunk objects
    async for chunk in llm.stream_content(prompt=user_input):

        # 1. Process incremental text (Token Delta)
        if chunk.content_delta:
            sys.stdout.write(chunk.content_delta)
            sys.stdout.flush()

        # 2. Handle metadata on finish (Usage, Stop Reason)
        if chunk.is_final:
            print(f"\n[Meta] Finished. Tokens: {chunk.usage.total_tokens}")

Structured Output (JSON Extraction)

from pydantic import BaseModel, Field
from nala.athomic.ai.llm.factory import LLMFactory

class UserIntent(BaseModel):
    intent: str = Field(..., description="The user's intention (buy, sell, support)")
    confidence: float

async def analyze_intent(text: str):
    llm = LLMFactory.create_default()

    # Guaranteed to return a UserIntent instance or raise StructuredOutputError
    intent: UserIntent = await llm.generate_structured(
        prompt=f"Analyze this text: {text}",
        response_model=UserIntent
    )

    return intent

Configuration

LLM connections are configured in settings.toml under the [ai.llms] section. You can define resilient routing by attaching wrappers to a connection.

[default.ai.llms]
# The default connection to use
default_connection_name = "gpt4_main"

  # 1. Primary Connection (with Router Wrapper)
  [default.ai.llms.connections.gpt4_main]
  enabled = true
  backend = "openai"
  default_model = "gpt-4-turbo"
  timeout = 10.0

    [default.ai.llms.connections.gpt4_main.provider]
    # OpenAI specific settings
    api_key = { path = "ai/openai", key = "api_key" }
    organization_id = "org-123"

    # Transparent fallback to a local or secondary model on failure
    [[default.ai.llms.connections.gpt4_main.wrappers]]
    name = "router"
    enabled = true
    fallback_chain = ["local_fallback"]
    strategy = "tier"

  # 2. Fallback Connection (e.g., Local Ollama)
  [default.ai.llms.connections.local_fallback]
  enabled = true
  backend = "openai"
  default_model = "llama3.2:1b"
  timeout = 30.0

    [default.ai.llms.connections.local_fallback.provider]
    base_url = "http://localhost:11434/v1"
    api_key = "ollama" # pragma: allowlist secret

API Reference

nala.athomic.ai.llm.protocol.LLMProviderProtocol

Bases: Protocol

Protocol defining the contract for Large Language Model (LLM) interactions. Focused on text generation and structured data extraction.

generate_content(prompt=None, system_message=None, **kwargs) async

Generates unstructured text content based on a given prompt.

generate_structured(prompt, response_model, system_message=None, max_retries=3, **kwargs) async

Generates structured data conforming to a specific Pydantic schema.

stream_content(prompt=None, system_message=None, **kwargs) async

Generates unstructured text content incrementally as an asynchronous stream.

This method is critical for low-latency user experiences (UX).

Parameters:

Name Type Description Default
prompt Optional[str]

The user input text.

 None
system_message Optional[str]

Optional system instruction.

 None
**kwargs Any

Additional generation parameters.

 {}

Returns:

Type Description
AsyncIterator[LLMResponseChunk]

An asynchronous iterator yielding LLMResponseChunk objects.

nala.athomic.ai.llm.base.BaseLLM

Bases: BaseService, LLMProviderProtocol, Generic[S], ABC

Abstract base class for all LLM providers (Vertex, OpenAI, etc.).

This class implements the Template Method pattern to enforce: 1. Lifecycle Management (Startup/Shutdown via BaseService). 2. Unified Observability (Tracing, Metrics, Logging). 3. Error Handling boundaries.

Note: Governance (Guards) is handled by the LLMManager/GuardPipeline before calling this provider.

generate(prompt=None, system_message=None, tools=None, **kwargs) async

Public method for generation (Text or Tool Calls) with observability.

generate_structured(prompt, response_model, system_message=None, max_retries=3, **kwargs) async

Public method for structured generation with observability.

get_output_max_tokens()

Helper to get max output tokens from generation settings.

record_token_usage(prompt_tokens, completion_tokens)

Helper method to record token usage metrics to Prometheus.

stream_content(prompt=None, system_message=None, tools=None, **kwargs) async

Executes an asynchronous streaming generation request and applies output governance. Validates the input prompt against configured guards before initiating the stream. Intercepts each generated chunk to enforce output guards, such as PII sanitization, by wrapping the delta in a standard response schema before yielding to the client.

nala.athomic.ai.schemas.llms.LLMResponseChunk

Bases: BaseModel

Represents a single, incremental chunk of the response during streaming.

nala.athomic.ai.llm.factory.LLMFactory

Bases: FactoryProtocol[LLMProviderProtocol, LLMConnectionSettings]

Factory responsible for constructing LLM providers and transparently applying configured decorator wrappers.

Adheres to the Open/Closed Principle by utilizing registries to resolve both base providers and decorator layers dynamically. It delegates the iteration of configured wrappers to a dedicated internal method, maintaining high cohesion and decoupling from specific routing logic.

create(settings, **kwargs) classmethod

Builds the base LLM client and delegates the application of enabled wrappers.

Parameters:

Name Type Description Default
settings LLMConnectionSettings

The configuration object for the connection.

 required
**kwargs Any

Additional unbound dependencies (such as managers or routers) to be passively injected into the custom wrapper creators.

 {}

Returns:

Name Type Description
LLMProviderProtocol LLMProviderProtocol

The fully composed LLM client instance.

Raises:

Type Description
ValueError

If the required base provider is not registered.