ai

package module
v0.2.0 Latest Latest
Warning

This package is not in the latest version of its module.

 Go to latest Published: Jul 26, 2026 License: MIT Imports: 6 Imported by: 0

README

github.com/spectrum-labs-tech/ai

A minimal Go library for structured AI completions. Uses a database/sql-style driver registry so you import only the providers you need.

Install

go get github.com/spectrum-labs-tech/ai

Import the driver(s) you want alongside the root package:

import (
    "github.com/spectrum-labs-tech/ai"
    _ "github.com/spectrum-labs-tech/ai/drivers/openai"
    // _ "github.com/spectrum-labs-tech/ai/drivers/anthropic"
    // _ "github.com/spectrum-labs-tech/ai/drivers/gemini"
    // _ "github.com/spectrum-labs-tech/ai/drivers/otium"
)

The otium driver targets an OpenAI-compatible endpoint and requires Config.BaseURL (it is never defaulted) — set it to your Otium endpoint, e.g. ai.New(&ai.Config{Provider: "otium", APIKey: key, BaseURL: url, Model: otium.ModelOtiumMedium}).

Quick start

p, err := ai.New(&ai.Config{
    Provider: "openai",
    APIKey:   os.Getenv("OPENAI_API_KEY"),
    Model:    "gpt-4o-mini",
})
if err != nil {
    log.Fatal(err)
}
defer p.Close()

const schema = `{
    "type": "object",
    "properties": {
        "summary": {"type": "string"}
    },
    "required": ["summary"],
    "additionalProperties": false
}`

result, err := p.Complete(ctx, "You are a summarizer.", "Summarize: "+text, schema, ai.Options{})

Structured output

Every Complete call accepts a JSON Schema string. The provider enforces it via its native structured-output API (OpenAI response format, Anthropic tool use, Gemini controlled generation). The return value is always a raw JSON string matching your schema.

Token usage

Wrap your context with a UsageRecorder to capture tokens and estimated cost:

rec := &ai.UsageRecorder{}
ctx = ai.WithUsageRecorder(ctx, rec)

result, err := p.Complete(ctx, system, user, schema, ai.Options{})

if u := rec.Usage(); u != nil {
    fmt.Printf("tokens: %d prompt / %d completion, cost: $%.6f\n",
        u.PromptTokens, u.CompletionTokens, u.Cost)
}

Async batch

Drivers implement BatchProvider for provider-managed async batches (typically ~50 % cheaper):

bp := p.(ai.BatchProvider)

job, err := bp.SubmitBatch(ctx, []ai.BatchRequest{
    {CustomID: "row-1", SystemPrompt: system, UserPrompt: user1, JSONSchema: schema},
    {CustomID: "row-2", SystemPrompt: system, UserPrompt: user2, JSONSchema: schema},
}, ai.BatchOptions{})

// Poll until job.Done, then:
results, err := bp.GetBatchResults(ctx, job.ID)

Agent / tool-calling loop

Drivers that implement AgentProvider run a multi-turn tool loop:

ap := p.(ai.AgentProvider)

result, err := ap.CompleteWithTools(ctx, system, user, schema,
    []ai.ToolDefinition{{Name: "lookup", Description: "...", Parameters: paramsSchema}},
    ai.Options{},
    func(ctx context.Context, calls []ai.ToolCallRequest) ([]ai.ToolCallResult, error) {
        // execute each tool call and return results
    },
)

Drivers

Import path Provider Complete Batch Agent (tool loop) Vision
drivers/openai OpenAI ✅ ✅ ✅ ✅
drivers/anthropic Anthropic ✅ ✅ ❌ ❌
drivers/gemini Google Gemini ✅ ✅ ❌ ❌
drivers/otium Otium (OpenAI-compatible) ✅ ✅ ❌ ✅

Agent and vision support for Anthropic and Gemini are not yet implemented.

Dynamic factory

NewProviderFactory returns a ProviderFactory that constructs a BatchProvider on demand, useful when different tasks use different models:

factory := ai.NewProviderFactory("openai", os.Getenv("OPENAI_API_KEY"))

bp, err := factory("gpt-4.1-mini")

Testing

Unit tests require no API keys:

go test ./...

Paid integration tests hit live APIs and require the relevant *_API_KEY environment variable:

go test -tags=paid_integration ./...

Roadmap

  • Streaming (StreamProvider) — a fourth optional interface for token-by-token streaming. Designed to support a split-stream pattern where structured JSON output is embedded in the stream between sentinel tokens (__JSON_START__ / __JSON_END__), letting callers pipe narrative text to a UI while still parsing a structured result at the end — without a second round trip.
  • Agent and vision support for Anthropic and Gemini drivers.
  • Driver support for self-hosted inference (e.g., vLLM, Ollama).

License

MIT — see LICENSE.

Documentation

Index

Constants

This section is empty.

Variables

View Source 
var ErrBatchOutputExpired = errors.New("batch output file expired or unavailable")

ErrBatchOutputExpired is returned by GetBatchResults when the provider's output file is no longer available (deleted after retention window). Callers should treat this the same as an expired batch and reset entities.

Functions

func Drivers 

func Drivers() []string

Drivers returns a list of the names of the registered drivers.

func Register 

func Register(name string, driver DriverFunc)

Register makes a driver available by the provided name. If Register is called twice with the same name or if driver is nil, it panics.

func WithUsageRecorder 

func WithUsageRecorder(ctx context.Context, r *UsageRecorder) context.Context

WithUsageRecorder returns a child context carrying r. The Provider will call r.Record() after a successful completion.

Types

type AgentProvider 

type AgentProvider interface {
	Provider
	// CompleteWithTools runs an agentic loop: the model receives tool definitions
	// and may call them via execTools before returning a final JSON response.
	CompleteWithTools(
		ctx context.Context,
		systemPrompt, userPrompt, jsonSchema string,
		tools []ToolDefinition,
		opts Options,
		execTools func(context.Context, []ToolCallRequest) ([]ToolCallResult, error),
	) (string, error)
}

AgentProvider extends Provider with an agentic tool-calling loop. The model may call tools zero or more times before producing its final JSON response. Usage across all iterations is recorded via the context UsageRecorder.

type BatchJob 

type BatchJob struct {
	ID               string             `json:"id"`
	Provider         string             `json:"provider"`
	Model            string             `json:"model,omitempty"`
	Status           string             `json:"status"`
	InputFileID      string             `json:"input_file_id,omitempty"`
	OutputFileID     string             `json:"output_file_id,omitempty"`
	ErrorFileID      string             `json:"error_file_id,omitempty"`
	RequestCounts    BatchRequestCounts `json:"request_counts"`
	CreatedAt        *time.Time         `json:"created_at,omitempty"`
	StartedAt        *time.Time         `json:"started_at,omitempty"`
	CompletedAt      *time.Time         `json:"completed_at,omitempty"`
	FailedAt         *time.Time         `json:"failed_at,omitempty"`
	CancelledAt      *time.Time         `json:"cancelled_at,omitempty"`
	Metadata         map[string]string  `json:"metadata,omitempty"`
	Done             bool               `json:"done"`
	ResultInline     bool               `json:"result_inline,omitempty"`
	ProviderResponse json.RawMessage    `json:"provider_response,omitempty"`
}

BatchJob describes a provider batch at a normalized level.

type BatchOptions 

type BatchOptions struct {
	// CompletionWindow is the requested provider turnaround window, if supported.
	CompletionWindow string

	// DisplayName is an optional human-readable label for the batch.
	DisplayName string

	// Metadata carries provider-supported batch metadata.
	Metadata map[string]string

	// ForceFile prefers file-backed submission when the provider supports both
	// inline and file input modes.
	ForceFile bool
}

BatchOptions controls provider batch submission behavior.

type BatchProvider 

type BatchProvider interface {
	Provider

	// SubmitBatch submits a provider-managed asynchronous batch job.
	SubmitBatch(ctx context.Context, requests []BatchRequest, opts BatchOptions) (*BatchJob, error)

	// GetBatch retrieves the current status for a submitted batch.
	GetBatch(ctx context.Context, batchID string) (*BatchJob, error)

	// CancelBatch attempts to cancel an in-flight batch.
	CancelBatch(ctx context.Context, batchID string) (*BatchJob, error)

	// GetBatchResults retrieves all currently available batch results. Providers
	// may return partial results for completed, cancelled, or expired batches.
	GetBatchResults(ctx context.Context, batchID string) ([]BatchResult, error)
}

BatchProvider extends Provider with asynchronous batch submission support. Implementations should translate each BatchRequest into the provider's native request shape and preserve CustomID in batch results for reconciliation.

type BatchRequest 

type BatchRequest struct {
	CustomID     string
	SystemPrompt string
	UserPrompt   string
	JSONSchema   string
	Options      Options
}

BatchRequest is one structured completion request inside an asynchronous batch.

type BatchRequestCounts 

type BatchRequestCounts struct {
	Total     int `json:"total"`
	Completed int `json:"completed"`
	Failed    int `json:"failed"`
}

BatchRequestCounts summarizes request completion status inside a batch.

type BatchResult 

type BatchResult struct {
	CustomID         string           `json:"custom_id"`
	Output           string           `json:"output,omitempty"`
	Error            string           `json:"error,omitempty"`
	StatusCode       int              `json:"status_code,omitempty"`
	RequestID        string           `json:"request_id,omitempty"`
	Usage            *CompletionUsage `json:"usage,omitempty"`
	ProviderResponse json.RawMessage  `json:"provider_response,omitempty"`
}

BatchResult is one normalized result row from a batch output.

type CompletionUsage 

type CompletionUsage struct {
	PromptTokens     int
	CompletionTokens int
	TotalTokens      int
	CachedTokens     int
	Cost             float64 // estimated USD; computed by the provider
}

CompletionUsage holds token usage and cost from a provider completion call.

type Config 

type Config struct {
	// Provider is the AI provider name (e.g., "openai", "anthropic").
	Provider string

	// Model is the specific model to use (e.g., "gpt-4o-mini").
	Model string

	// APIKey is the API key for cloud providers.
	APIKey string

	// BaseURL is the base URL for the API (optional, provider-specific).
	BaseURL string

	// Options allows provider-specific configuration.
	Options map[string]interface{}

	// MaxRetries is the number of retry attempts on transient errors (429, 5xx).
	// Zero uses the default (3). Set to -1 to disable retries.
	MaxRetries int
}

Config holds configuration for creating a Provider.

type DriverFunc 

type DriverFunc func(*Config) (Provider, error)

DriverFunc is a function that creates a new Provider from a Config.

type Options 

type Options struct {
	// Temperature controls randomness (nil = use provider default).
	Temperature *float64

	// MaxTokens limits the response length (0 = use provider default).
	MaxTokens int

	// ImageURL is a publicly accessible URL of an image to include in the user
	// message. When set, providers that support vision will send the image
	// alongside the text prompt. Empty string means no image.
	ImageURL string
}

Options contains per-request tuning knobs.

type Provider 

type Provider interface {
	// Complete sends a prompt and returns the raw JSON response string.
	// jsonSchema constrains the response format (passed to the provider's
	// structured output / function calling API).
	Complete(ctx context.Context, systemPrompt, userPrompt, jsonSchema string, opts Options) (string, error)

	// ProviderName returns the name of the AI provider (e.g., "openai").
	ProviderName() string

	// ModelName returns the specific model being used (e.g., "gpt-4o-mini").
	ModelName() string

	// Close releases any resources held by the provider.
	Close() error
}

Provider is the generic interface for AI completions with structured output. Implementations must be safe for concurrent use.

func New 

func New(cfg *Config) (Provider, error)

New creates a new Provider based on the provided config. The provider must be registered via Register.

type ProviderFactory 

type ProviderFactory func(model string) (BatchProvider, error)

ProviderFactory creates a BatchProvider for a given model on demand. Calling the factory at task run time (rather than at startup) means infrequent tasks don't hold a live connection for their entire idle period.

func NewProviderFactory 

func NewProviderFactory(providerName, apiKey string) ProviderFactory

NewProviderFactory returns a ProviderFactory that constructs BatchProviders using the given provider name and API key, with the model supplied per-call.

type Schema 

type Schema struct {
	// Name is a unique identifier for this schema.
	Name string

	// SystemPrompt is the system message with instructions.
	SystemPrompt string

	// UserPromptTemplate is a Go text/template for the user message.
	UserPromptTemplate string

	// JSONSchema is the JSON Schema for structured output.
	JSONSchema string
}

Schema defines the prompt template and JSON schema for a completion task.

type ToolCallRequest 

type ToolCallRequest struct {
	// ID is the provider-assigned call identifier; must be echoed back in ToolCallResult.
	ID string
	// Name is the tool name, matching a ToolDefinition.Name.
	Name string
	// ArgsJSON is the JSON-encoded arguments matching the tool's parameter schema.
	ArgsJSON string
}

ToolCallRequest is one tool invocation requested by the AI model.

type ToolCallResult 

type ToolCallResult struct {
	ID      string // matches ToolCallRequest.ID
	Content string // tool output, typically JSON
}

ToolCallResult is the response to one ToolCallRequest.

type ToolDefinition 

type ToolDefinition struct {
	Name        string
	Description string
	// Parameters is a JSON Schema object describing the function arguments.
	Parameters json.RawMessage
}

ToolDefinition describes a callable function for providers that support tool use.

type UsageRecorder 

type UsageRecorder struct {
	// contains filtered or unexported fields
}

UsageRecorder captures token usage from a Provider.Complete() call via context. Safe for concurrent use.

func UsageRecorderFromContext 

func UsageRecorderFromContext(ctx context.Context) *UsageRecorder

UsageRecorderFromContext retrieves the UsageRecorder from ctx. Returns nil if none was set. Intended for Provider implementations.

func (*UsageRecorder) Record 

func (r *UsageRecorder) Record(u CompletionUsage)

Record is called by Provider implementations to report usage after a completion.

func (*UsageRecorder) Usage 

func (r *UsageRecorder) Usage() *CompletionUsage

Usage returns the recorded usage, or nil if Complete() has not been called.

Source Files

View all Source files

Directories

Path Synopsis
drivers
anthropic
gemini
openai
otium
Package otium is a driver for Otium, a batch-first inference service that runs open models on low-cost compute behind an OpenAI-compatible API.
 Package otium is a driver for Otium, a batch-first inference service that runs open models on low-cost compute behind an OpenAI-compatible API.
internal
httpretry

Jump to

Keyboard shortcuts

? : This menu
/ : Search site
f or F : Jump to
y or Y : Canonical URL
go.dev uses cookies from Google to deliver and enhance the quality of its services and to analyze traffic. Learn more.