aiutil

package module
v1.2.5 Latest Latest
Warning

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

 Go to latest Published: Jan 7, 2026 License: MIT Imports: 15 Imported by: 2

README

Latest Release

AI Util

Tools for building with AI.

Features

  • Supported AI Providers:
  • Shared client interface across providers:
    • Complete - Single completion requests
    • Stream - Streaming completion requests
    • GetModels - List available models
  • Conversation Management:
    • Manage message history and token counts with auto-truncation
    • Support for system prompts and role-based messaging
  • Tool Calling:
    • Invoke backend tools and APIs from within conversations
    • Available only for supported models.

Installation

go get github.com/ztkent/ai-util@latest

Configuration Options

Create a Client:

client, err := NewAIClient().
    WithOpenAI("api-key").                    // Add OpenAI provider
    WithDefaultProvider("openai").            // Set default provider
    WithDefaultModel("gpt-4o").               // Set default model
    WithDefaultTemperature(0.7).              // Set default temperature
    WithDefaultMaxTokens(4096).               // Set default max tokens
    Build()

Request Options:

  • Temperature(float64): Sampling temperature (0.0 to 2.0)
  • MaxTokens(int): Maximum tokens to generate
  • TopP(float64): Nucleus sampling probability
  • FrequencyPenalty(float64): Penalize frequent tokens (OpenAI)
  • PresencePenalty(float64): Penalize present tokens (OpenAI)
  • Stop([]string): Stop sequences

Conversation Options:

  • SystemPrompt: Initial system message
  • MaxTokens: Token limit for conversation
  • AutoTruncate: Automatically remove old messages when limit reached
  • PreserveSystem: Keep system message during truncation

API Keys

API keys are loaded from environment variables by default:

  • OpenAI: OPENAI_API_KEY
  • Google AI: GOOGLE_API_KEY and GOOGLE_PROJECT_ID
  • Replicate: REPLICATE_API_TOKEN

Or explicitly provided via the builder pattern.

Examples

The repository includes examples demonstrating features for each supported provider.

  • OpenAI Provider Example (examples/openai/openai_provider_example.go)

  • Google AI Provider Example (examples/google/google_provider_example.go)

  • Features:

    • Basic chat completions
    • Streaming responses
    • Tool/function calling
    • Token estimation
    • Model listing
    • Error handling

Usage

// Basic completion request
resp, err := client.Complete(ctx, &types.CompletionRequest{
    Messages: []*types.Message{
        types.NewTextMessage(types.RoleUser, "What is the capital of France?"),
    },
    Model:       "gpt-4o",
    MaxTokens:   100,
    Temperature: 0.7,
})
if err != nil {
    log.Fatal(err)
}

fmt.Printf("Response: %s\n", resp.Message.TextData)
fmt.Printf("Usage: %+v\n", resp.Usage)

// Multi-message conversation
resp, err := client.Complete(ctx, &types.CompletionRequest{
    Messages: []*types.Message{
        types.NewTextMessage(types.RoleSystem, "You are a helpful assistant."),
        types.NewTextMessage(types.RoleUser, "Explain quantum computing in simple terms."),
    },
    Model:       "gpt-4o",
    MaxTokens:   500,
    Temperature: 0.5,
})
if err != nil {
    log.Fatal(err)
}

fmt.Printf("Assistant: %s\n", resp.Message.TextData)
Streaming
// Stream responses
err := client.Stream(ctx, &types.CompletionRequest{
    Messages: []*types.Message{
        types.NewTextMessage(types.RoleUser, "Tell me a story"),
    },
    MaxTokens: 1000,
}, func(ctx context.Context, response *types.StreamResponse) error {
    if response.Delta != nil && response.Delta.TextData != "" {
        fmt.Print(response.Delta.TextData)
    }
    return nil
})
if err != nil {
    log.Fatal(err)
}
Error Handling

The API provides structured error handling:

resp, err := client.Complete(ctx, req)
if err != nil {
    if aiErr, ok := err.(*types.Error); ok {
        fmt.Printf("Provider: %s, Code: %s, Message: %s\n", 
            aiErr.Provider, aiErr.Code, aiErr.Message)
    }
}

Available Models

OpenAI Models
Model Name Model Identifier
GPT-5 gpt-5
O3 Preview o3-preview
O3 Mini o3-mini
GPT-4o gpt-4o
GPT-4o Mini gpt-4o-mini
GPT-4 Turbo gpt-4-turbo
GPT-4 gpt-4
O1 Preview o1-preview
O1 Mini o1-mini
Google AI Models
Model Name Model Identifier Capabilities
Gemini 3.0 Series (Next Gen)
Gemini 3 Pro gemini-3-pro-preview Chat, Streaming, Tools, Vision, Audio, Video, Thinking
Gemini 3 Flash Preview gemini-3-flash-preview Chat, Streaming, Tools, Vision, Audio, Video, Thinking
Gemini 2.5 Series (Latest)
Gemini 2.5 Pro gemini-2.5-pro Chat, Streaming, Tools, Vision, Audio, Video, Thinking
Gemini 2.5 Flash gemini-2.5-flash Chat, Streaming, Tools, Vision, Audio, Video, Thinking
Gemini 2.5 Flash-Lite gemini-2.5-flash-lite Chat, Streaming, Tools, Vision, Audio, Video
Gemini 2.5 Flash Preview TTS gemini-2.5-flash-preview-tts Text-to-Speech
Gemini 2.5 Pro Preview TTS gemini-2.5-pro-preview-tts Text-to-Speech
Live Interaction Models
Gemini 2.5 Flash Live gemini-2.5-flash-live Live Audio/Video, Streaming
Gemini 2.0 Flash Live gemini-2.0-flash-live Live Audio/Video, Streaming
Embedding Models
Text Embedding 004 text-embedding-004 Text Embeddings
Gemini Embedding Experimental gemini-embedding-exp Text Embeddings
Generation Models
Imagen 4 imagen-4.0-generate-preview Image Generation
Imagen 3 imagen-3.0-generate-002 Image Generation
Veo 2 veo-2.0-generate-001 Video Generation
Veo 3 veo-3.0-generate-001 Video Generation
Replicate Models
Model Name Model Identifier
Meta Llama 3.1 8B Instruct meta/meta-llama-3.1-8b-instruct
Meta Llama 3.1 70B Instruct meta/meta-llama-3.1-70b-instruct
Meta Llama 3.1 405B Instruct meta/meta-llama-3.1-405b-instruct
Meta Llama 3 8B Instruct meta/meta-llama-3-8b-instruct
Meta Llama 3 70B Instruct meta/meta-llama-3-70b-instruct
Mistral 7B Instruct mistralai/mistral-7b-instruct-v0.2
Mixtral 8x7B Instruct mistralai/mixtral-8x7b-instruct-v0.1

Documentation

Index

Constants

This section is empty.

Variables

This section is empty.

Functions

func IsQuotaExceededError added in v1.2.4 

func IsQuotaExceededError(err error) bool

IsQuotaExceededError checks if an error indicates quota has been exceeded

func IsRetryableError added in v1.2.4 

func IsRetryableError(err error) bool

IsRetryableError determines if an error is worth retrying

func ParseRateLimitDelay added in v1.2.4 

func ParseRateLimitDelay(err error) time.Duration

ParseRateLimitDelay extracts the suggested retry delay from rate limit errors Looks for patterns like "Please retry in 34.42245165s"

func TestBuilder added in v1.2.0 

func TestBuilder(t *testing.T)

func TestConversationConfig added in v1.2.0 

func TestConversationConfig(t *testing.T)

func TestConversationMessages added in v1.2.0 

func TestConversationMessages(t *testing.T)

func TestError added in v1.2.0 

func TestError(t *testing.T)

func TestMessage added in v1.2.0 

func TestMessage(t *testing.T)

func TestModel added in v1.2.0 

func TestModel(t *testing.T)

func TestModelRegistry added in v1.2.0 

func TestModelRegistry(t *testing.T)

func WithRetry added in v1.2.4 

func WithRetry(ctx context.Context, req *types.CompletionRequest, config *RetryConfig, fn CompletionFunc) (*types.CompletionResponse, error)

WithRetry executes a completion request with smart retry logic - Parses rate limit errors for suggested retry delays - Uses exponential backoff for other transient errors - Skips retries for non-retryable errors (auth, invalid request) - Falls back through models in FallbackModels on quota errors (if provided)

Types

type AIClient added in v1.2.0 

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

AIClient provides a fluent interface for configuring and creating a Client

func NewAIClient 

func NewAIClient() *AIClient

NewAIClient creates a new client builder

func (*AIClient) Build added in v1.2.0 

func (b *AIClient) Build() (*Client, error)

Build creates and configures the client

func (*AIClient) WithDefaultMaxTokens added in v1.2.0 

func (b *AIClient) WithDefaultMaxTokens(maxTokens int) *AIClient

WithDefaultMaxTokens sets the default max tokens

func (*AIClient) WithDefaultModel added in v1.2.0 

func (b *AIClient) WithDefaultModel(model string) *AIClient

WithDefaultModel sets the default model

func (*AIClient) WithDefaultProvider added in v1.2.0 

func (b *AIClient) WithDefaultProvider(provider string) *AIClient

WithDefaultProvider sets the default provider

func (*AIClient) WithDefaultTemperature added in v1.2.0 

func (b *AIClient) WithDefaultTemperature(temperature float64) *AIClient

WithDefaultTemperature sets the default temperature

func (*AIClient) WithGoogle added in v1.2.0 

func (b *AIClient) WithGoogle(apiKey, projectID string, options ...GoogleOption) *AIClient

WithGoogle configures Google AI provider

func (*AIClient) WithMiddleware added in v1.2.0 

func (b *AIClient) WithMiddleware(middleware ...Middleware) *AIClient

WithMiddleware adds middleware to the client

func (*AIClient) WithOpenAI added in v1.2.0 

func (b *AIClient) WithOpenAI(apiKey string, options ...OpenAIOption) *AIClient

WithOpenAI configures OpenAI provider

func (*AIClient) WithReplicate added in v1.2.0 

func (b *AIClient) WithReplicate(apiKey string, options ...ReplicateOption) *AIClient

WithReplicate configures Replicate provider

type Client 

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

Client is the main interface for interacting with AI providers

func NewClient added in v1.2.0 

func NewClient(config *ClientConfig) *Client

NewClient creates a new AI client

func NewGoogle added in v1.2.0 

func NewGoogle(apiKey, projectID string) (*Client, error)

func NewOpenAI added in v1.2.0 

func NewOpenAI(apiKey string) (*Client, error)

Simple Client Connections

func NewReplicate added in v1.2.0 

func NewReplicate(apiKey string) (*Client, error)

func (*Client) Close added in v1.2.0 

func (c *Client) Close() error

Close closes all providers and cleans up resources

func (*Client) Complete added in v1.2.0 

func (c *Client) Complete(ctx context.Context, req *types.CompletionRequest) (*types.CompletionResponse, error)

Complete performs a completion request

func (*Client) EstimateTokens added in v1.2.0 

func (c *Client) EstimateTokens(ctx context.Context, messages []*types.Message, model string) (int, error)

EstimateTokens estimates token count for messages and model

func (*Client) GetModel 

func (c *Client) GetModel(provider, id string) (*types.Model, error)

GetModel returns a model by provider and ID

func (*Client) GetProvider added in v1.2.0 

func (c *Client) GetProvider(name string) (types.Provider, error)

GetProvider returns a provider by name

func (*Client) ListModels 

func (c *Client) ListModels() []*types.Model

ListModels returns all available models

func (*Client) ListModelsByProvider added in v1.2.0 

func (c *Client) ListModelsByProvider(provider string) []*types.Model

ListModelsByProvider returns models for a specific provider

func (*Client) NewConversation added in v1.2.0 

func (c *Client) NewConversation(config *ConversationConfig) *Conversation

NewConversation creates a new conversation with optional system prompt

func (*Client) RegisterProvider added in v1.2.0 

func (c *Client) RegisterProvider(provider types.Provider) error

RegisterProvider registers a new provider with the client

func (*Client) Stream added in v1.2.0 

func (c *Client) Stream(ctx context.Context, req *types.CompletionRequest, callback types.StreamCallback) error

Stream performs a streaming completion request

type ClientConfig added in v1.0.0 

type ClientConfig struct {
	DefaultProvider    string                  `json:"default_provider,omitempty"`
	DefaultModel       string                  `json:"default_model,omitempty"`
	DefaultMaxTokens   int                     `json:"default_max_tokens,omitempty"`
	DefaultTemperature float64                 `json:"default_temperature,omitempty"`
	ProviderConfigs    map[string]types.Config `json:"provider_configs,omitempty"`
	Middleware         []Middleware            `json:"-"`
}

ClientConfig holds global client configuration

type CompletionFunc added in v1.2.4 

type CompletionFunc func(ctx context.Context, req *types.CompletionRequest) (*types.CompletionResponse, error)

CompletionFunc is the function signature for AI completion calls

type Conversation 

type Conversation struct {
	ID            string                 `json:"id"`
	Messages      []*types.Message       `json:"messages"`
	MaxTokens     int                    `json:"max_tokens"`
	CurrentTokens int                    `json:"current_tokens"`
	CreatedAt     time.Time              `json:"created_at"`
	UpdatedAt     time.Time              `json:"updated_at"`
	Metadata      map[string]interface{} `json:"metadata,omitempty"`
	// contains filtered or unexported fields
}

Conversation represents a conversation with message history and management

func (*Conversation) AddAssistantMessage added in v1.2.0 

func (c *Conversation) AddAssistantMessage(text string) error

AddAssistantMessage adds an assistant message to the conversation

func (*Conversation) AddMessage added in v1.2.0 

func (c *Conversation) AddMessage(message *types.Message) error

AddMessage adds a message to the conversation

func (*Conversation) AddSystemMessage added in v1.2.0 

func (c *Conversation) AddSystemMessage(text string) error

AddSystemMessage adds a system message to the conversation

func (*Conversation) AddUserMessage added in v1.2.0 

func (c *Conversation) AddUserMessage(text string) error

AddUserMessage adds a user message to the conversation

func (*Conversation) Clear added in v1.2.0 

func (c *Conversation) Clear()

Clear removes all messages from the conversation

func (*Conversation) Clone added in v1.2.0 

func (c *Conversation) Clone() *Conversation

Clone creates a copy of the conversation

func (*Conversation) EstimateTokens added in v1.2.0 

func (c *Conversation) EstimateTokens(ctx context.Context, model string) (int, error)

EstimateTokens estimates the current token count of the conversation

func (*Conversation) Export added in v1.2.0 

func (c *Conversation) Export() map[string]interface{}

Export exports the conversation to a JSON-serializable format

func (*Conversation) GetLastMessage added in v1.2.0 

func (c *Conversation) GetLastMessage() *types.Message

GetLastMessage returns the last message in the conversation

func (*Conversation) GetMessages added in v1.2.0 

func (c *Conversation) GetMessages() []*types.Message

GetMessages returns a copy of all messages

func (*Conversation) GetMessagesByRole added in v1.2.0 

func (c *Conversation) GetMessagesByRole(role types.Role) []*types.Message

GetMessagesByRole returns messages filtered by role

func (*Conversation) GetTokenCount added in v1.2.0 

func (c *Conversation) GetTokenCount() int

GetTokenCount returns the last estimated token count

func (*Conversation) Send added in v1.2.0 

func (c *Conversation) Send(ctx context.Context, userMessage string, model string) (*types.CompletionResponse, error)

Send sends a user message and gets a response

func (*Conversation) SendStream added in v1.2.0 

func (c *Conversation) SendStream(ctx context.Context, userMessage string, model string, callback types.StreamCallback) error

SendStream sends a user message and streams the response

func (*Conversation) TruncateToFit added in v1.2.0 

func (c *Conversation) TruncateToFit(ctx context.Context, model string, preserveSystem bool) error

TruncateToFit ensures the conversation fits within token limits

type ConversationConfig added in v1.2.0 

type ConversationConfig struct {
	SystemPrompt   string                 `json:"system_prompt,omitempty"`
	MaxTokens      int                    `json:"max_tokens,omitempty"`
	Model          string                 `json:"model,omitempty"`
	Metadata       map[string]interface{} `json:"metadata,omitempty"`
	AutoTruncate   bool                   `json:"auto_truncate,omitempty"`
	PreserveSystem bool                   `json:"preserve_system,omitempty"` // Keep system message when truncating
}

ConversationConfig holds configuration for creating a conversation

type GoogleOption added in v1.2.0 

type GoogleOption func(*google.Config)

GoogleOption configures Google AI-specific settings

func WithGoogleBaseURL added in v1.2.0 

func WithGoogleBaseURL(baseURL string) GoogleOption

WithGoogleBaseURL sets a custom base URL for Google AI

func WithGoogleLocation added in v1.2.0 

func WithGoogleLocation(location string) GoogleOption

WithGoogleLocation sets the Google Cloud location

type LoggingMiddleware added in v1.2.0 

type LoggingMiddleware struct{}

LoggingMiddleware is an example middleware that logs requests and responses

func (*LoggingMiddleware) ProcessRequest added in v1.2.0 

func (m *LoggingMiddleware) ProcessRequest(ctx context.Context, req *types.CompletionRequest) (*types.CompletionRequest, error)

func (*LoggingMiddleware) ProcessResponse added in v1.2.0 

func (m *LoggingMiddleware) ProcessResponse(ctx context.Context, resp *types.CompletionResponse) (*types.CompletionResponse, error)

type Middleware added in v1.2.0 

type Middleware interface {
	ProcessRequest(ctx context.Context, req *types.CompletionRequest) (*types.CompletionRequest, error)
	ProcessResponse(ctx context.Context, resp *types.CompletionResponse) (*types.CompletionResponse, error)
}

Middleware defines the interface for request/response middleware

type OpenAIOption added in v1.2.0 

type OpenAIOption func(*openai.Config)

OpenAIOption configures OpenAI-specific settings

func WithOpenAIBaseURL added in v1.2.0 

func WithOpenAIBaseURL(baseURL string) OpenAIOption

WithOpenAIBaseURL sets a custom base URL for OpenAI

func WithOpenAIOrg added in v1.2.0 

func WithOpenAIOrg(orgID string) OpenAIOption

WithOpenAIOrg sets the OpenAI organization ID

func WithOpenAIProject added in v1.2.0 

func WithOpenAIProject(project string) OpenAIOption

WithOpenAIProject sets the OpenAI project

func WithOpenAIUser added in v1.0.0 

func WithOpenAIUser(user string) OpenAIOption

WithOpenAIUser sets the user identifier for OpenAI

type ReplicateOption added in v1.2.0 

type ReplicateOption func(*replicate.Config)

ReplicateOption configures Replicate-specific settings

func WithReplicateBaseURL added in v1.2.0 

func WithReplicateBaseURL(baseURL string) ReplicateOption

WithReplicateBaseURL sets a custom base URL for Replicate

func WithReplicateExtraInput added in v1.2.0 

func WithReplicateExtraInput(key string, value interface{}) ReplicateOption

WithReplicateExtraInput adds extra input parameters for Replicate

func WithReplicateWebhook added in v1.0.0 

func WithReplicateWebhook(webhookURL string) ReplicateOption

WithReplicateWebhook sets the webhook URL for Replicate

type RetryConfig added in v1.2.4 

type RetryConfig struct {
	MaxAttempts    int           // Maximum number of retry attempts (default: 5)
	BaseDelay      time.Duration // Initial delay for exponential backoff (default: 2s)
	MaxDelay       time.Duration // Maximum delay between retries (default: 30s)
	FallbackModels []string      // Models to try in order on quota errors (optional)
}

RetryConfig holds configuration for the retry logic

func DefaultRetryConfig added in v1.2.4 

func DefaultRetryConfig() *RetryConfig

DefaultRetryConfig returns the default retry configuration

Source Files

View all Source files

Directories

Path Synopsis
examples
google command
openai command
providers
google
openai
replicate

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.