AI SDK

A comprehensive AI backend framework for Go — providing unified abstractions for text generation, streaming, structured output, agents, RAG, embeddings, workflows, image generation, observability, and more. Inspired by Laravel's AI SDK and Vercel AI SDK.

Try the AI demo to see it in action.

Installation

Add the plugin in your bin/server.go:

import (
    "github.com/CodeSyncr/nimbus"
    "github.com/CodeSyncr/nimbus/plugins/ai"
)

func main() {
    app := nimbus.New()
    app.Use(ai.New())
    // ...
}

Configuration

Set environment variables in your .env:

AI_PROVIDER=openai
AI_MODEL=gpt-4o
OPENAI_API_KEY=sk-...

For local development without an API key, use AI_PROVIDER=ollama with Ollama running locally.

Environment Variables

Variable	Description
AI_PROVIDER	openai, xai, ollama, anthropic, gemini, mistral, cohere
AI_MODEL	Model name (e.g. gpt-4o, llama3.2)
OPENAI_API_KEY	Required for OpenAI
ANTHROPIC_API_KEY	Required for Anthropic
GEMINI_API_KEY	Required for Gemini
MISTRAL_API_KEY	Required for Mistral
COHERE_API_KEY	Required for Cohere
XAI_API_KEY	Required for xAI (Grok)
OLLAMA_HOST	Ollama server URL (default: http://localhost:11434)

Text Generation

Use ai.Generate for simple completions:

response, err := ai.Generate(c.Request().Context(), "Explain quantum computing")
if err != nil {
    return err
}
return c.JSON(200, map[string]string{"answer": response.Text})

Options

Customize generation with functional options:

response, err := ai.Generate(ctx, "Summarize this article",
    ai.WithModel("gpt-4o-mini"),
    ai.WithMaxTokens(500),
    ai.WithTemperature(0.7),
    ai.WithSystem("You are a concise summarizer. Use bullet points."),
)

Streaming

Use ai.Stream for SSE streaming responses:

func (ctrl *AI) StreamChat(c *http.Context) error {
    c.Response().Header().Set("Content-Type", "text/event-stream")
    c.Response().Header().Set("Cache-Control", "no-cache")

    stream, errCh := ai.Stream(c.Request().Context(), c.Request.FormValue("prompt"))
    flusher := c.Response().(http.Flusher)

    for chunk := range stream {
        fmt.Fprintf(c.Response(), "data: %s\n\n", chunk)
        flusher.Flush()
    }
    if err := <-errCh; err != nil {
        return err
    }
    return nil
}

Structured Output (Extract)

Extract typed Go structs from unstructured text. The AI generates JSON matching your struct's shape, with automatic retries.

Example: Receipt Parser

type Receipt struct {
    Merchant string
    Amount   float64
    Date     string
    Items    []Item
}

type Item struct {
    Name  string
    Price float64
    Qty   int
}

func (ctrl *Receipts) Parse(c *http.Context) error {
    text := c.Request.FormValue("receipt_text")

    receipt, err := ai.Extract[Receipt](c.Request().Context(), text)
    if err != nil {
        return err
    }

    return c.JSON(200, receipt)
    // { "merchant": "Starbucks", "amount": 12.50, "date": "2025-03-13", "items": [...] }
}

Example: Sentiment Classifier

func (ctrl *Reviews) Classify(c *http.Context) error {
    text := c.Request.FormValue("review")

    label, err := ai.Classify(c.Request().Context(), text,
        []string{"positive", "negative", "neutral"},
    )
    if err != nil {
        return err
    }

    return c.JSON(200, map[string]string{"sentiment": label})
}

Agents

Agents combine instructions, tools, memory, and a reasoning loop. They can autonomously call tools and iterate until they have an answer.

Example: Customer Support Agent

func init() {
    // Register tools the agent can use.
    ai.RegisterTool(ai.Tool{
        Name:        "lookup_order",
        Description: "Look up a customer order by ID",
        Run: func(ctx context.Context, in struct {
            OrderID string `description:"The order ID"`
        }) (map[string]any, error) {
            order := db.FindOrder(in.OrderID)
            return map[string]any{
                "status":  order.Status,
                "total":   order.Total,
                "shipped": order.ShippedAt,
            }, nil
        },
    })

    ai.RegisterTool(ai.Tool{
        Name:        "create_refund",
        Description: "Create a refund for an order",
        Run: func(ctx context.Context, in struct {
            OrderID string
            Amount  float64
            Reason  string
        }) (map[string]string, error) {
            refundID := billing.CreateRefund(in.OrderID, in.Amount, in.Reason)
            return map[string]string{"refund_id": refundID, "status": "processed"}, nil
        },
    })
}

func (ctrl *Support) Chat(c *http.Context) error {
    userID := c.Session().GetString("user_id")

    agent := ai.NewAgent("You are a helpful customer support agent for an e-commerce store. "+
        "Always look up the order before taking action. Be polite and concise.").
        WithTools("lookup_order", "create_refund").
        WithMemory(ai.RedisMemory(rdb), "support:"+userID).
        MaxSteps(5)

    response, err := agent.Prompt(c.Request().Context(), c.Request.FormValue("message"))
    if err != nil {
        return err
    }

    return c.JSON(200, map[string]string{"reply": response.Text})
}

Example: Code Review Agent

func (ctrl *CodeReview) Review(c *http.Context) error {
    code := c.Request.FormValue("code")

    agent := ai.NewAgent("You are an expert Go code reviewer. "+
        "Analyze code for bugs, performance issues, and style improvements. "+
        "Rate severity as: critical, warning, or info.").
        WithModel("gpt-4o")

    response, err := agent.Prompt(c.Request().Context(),
        "Review this Go code:\n\n```go\n"+code+"\n```",
    )
    if err != nil {
        return err
    }

    return c.JSON(200, map[string]string{"review": response.Text})
}

Tools (Function Calling)

Register Go functions as tools that agents can call. Input types generate JSON schemas automatically via reflection.

package tools

import (
    "context"
    "github.com/CodeSyncr/nimbus/plugins/ai"
)

type WeatherInput struct {
    City    string `description:"City name"`
    Country string `description:"ISO country code"`
}

type WeatherOutput struct {
    Temp       int
    Humidity   int
    Conditions string
}

func init() {
    ai.NewTool("get_weather").
        Desc("Get current weather for a city").
        Handler(func(ctx context.Context, in WeatherInput) (WeatherOutput, error) {
            // Call your weather API here.
            w := weatherAPI.Get(in.City, in.Country)
            return WeatherOutput{
                Temp:       w.TempF,
                Humidity:   w.Humidity,
                Conditions: w.Description,
            }, nil
        }).
        Register()
}

Embeddings & Vector Store

Generate embeddings and perform similarity search with pluggable backends.

func (ctrl *Search) Semantic(c *http.Context) error {
    query := c.Request.FormValue("q")

    // Search the knowledge base.
    store := ai.VectorStoreInstance("knowledge")
    results, err := store.Search(c.Request().Context(), query, 5)
    if err != nil {
        return err
    }

    items := make([]map[string]any, len(results))
    for i, doc := range results {
        items[i] = map[string]any{
            "id":    doc.ID,
            "text":  doc.Text,
            "score": doc.Score,
        }
    }

    return c.JSON(200, map[string]any{"results": items})
}

Vector Store Backends

Backend	Constructor	Use Case
In-memory	ai.NewMemoryVectorStore()	Development / testing
pgvector	ai.NewPgvectorStore(db)	PostgreSQL production
Qdrant	ai.NewQdrantStore(url, collection)	Dedicated vector DB
Pinecone	ai.NewPineconeStore(host, key)	Managed cloud

Example: pgvector Setup

import (
    "github.com/CodeSyncr/nimbus/plugins/ai"
    "gorm.io/driver/postgres"
    "github.com/CodeSyncr/nimbus/lucid"
)

func Boot(app *nimbus.App) {
    db := app.Container.MustMake("db").(*nimbus.DB)

    // Use pgvector for production vector storage.
    backend := ai.NewPgvectorStore(db.DB,
        ai.WithPgvectorTable("document_embeddings"),
        ai.WithPgvectorDimension(1536),
    )
    store := ai.VectorStoreInstance("knowledge", backend)
    app.Container.Singleton("vectorstore", func() any { return store })
}

RAG (Retrieval-Augmented Generation)

Build knowledge-grounded Q&A systems. The RAG engine embeds the query, searches the vector store, and augments the prompt with relevant context.

Example: Documentation Q&A

func (ctrl *DocsQA) Ask(c *http.Context) error {
    store := app.Container.MustMake("vectorstore").(*ai.Store)

    rag := ai.NewRAG(store).
        TopK(5).
        MinScore(0.7).
        WithCitations(true).
        WithSystem("You are a Nimbus documentation assistant. " +
            "Answer ONLY based on the provided context.")

    answer, err := rag.Ask(c.Request().Context(), c.Request.FormValue("question"))
    if err != nil {
        return err
    }

    sources := make([]map[string]any, len(answer.Sources))
    for i, src := range answer.Sources {
        sources[i] = map[string]any{"id": src.ID, "text": src.Text, "score": src.Score}
    }

    return c.JSON(200, map[string]any{
        "answer":  answer.Answer,
        "sources": sources,
    })
}

Example: Ingest Documents

func ingestDocs(ctx context.Context, store *ai.Store) error {
    docs := map[string]string{
        "routing":    routingDocText,
        "database":   databaseDocText,
        "middleware": middlewareDocText,
    }
    return ai.IngestDocuments(ctx, store, docs)
}

Prompt Templates

Reusable, composable prompt templates with Go's text/template syntax.

// Simple template
prompt := ai.Template("Translate the following to {{.language}}:\n\n{{.text}}")
resp, err := prompt.Generate(ctx, map[string]any{
    "language": "Spanish",
    "text":     "Hello, how are you?",
})

// Few-shot classifier
classifier := ai.FewShot("Classify the support ticket priority.").
    Add("My account is hacked!", "critical").
    Add("Can't change my profile picture", "low").
    Add("Payment failed for my order", "high")

resp, err := classifier.Generate(ctx, "I want to update my shipping address")

// Composable chain
resp, err := ai.Chain(
    ai.SystemTemplate("You are a {{.role}} expert."),
    ai.Template("Explain {{.topic}} for a {{.audience}} audience."),
).Generate(ctx, map[string]any{
    "role":     "Go",
    "topic":    "goroutines and channels",
    "audience": "beginner",
})

Workflows (Multi-Step Pipelines)

Orchestrate complex AI pipelines with sequential, parallel, and branching steps.

Example: Blog Post Generator

func generateBlogPost(c *http.Context) error {
    wf := ai.NewWorkflow("blog-post").
        Step("research", func(ctx context.Context, wc *ai.WorkflowContext) error {
            resp, _ := ai.Generate(ctx, "Research key points about: "+wc.GetString("topic"),
                ai.WithModel("gpt-4o"))
            wc.Set("research", resp.Text)
            return nil
        }).
        Step("outline", func(ctx context.Context, wc *ai.WorkflowContext) error {
            resp, _ := ai.Generate(ctx,
                "Create a blog outline from this research:\n"+wc.GetString("research"),
            )
            wc.Set("outline", resp.Text)
            return nil
        }).
        Parallel("media",
            ai.StepFunc("hero_image", func(ctx context.Context, wc *ai.WorkflowContext) error {
                img, _ := ai.Image().
                    Prompt("Blog hero image: "+wc.GetString("topic")).
                    Size("1792x1024").
                    Generate(ctx)
                wc.Set("hero_url", img.Images[0].URL)
                return nil
            }),
            ai.StepFunc("social_copy", func(ctx context.Context, wc *ai.WorkflowContext) error {
                resp, _ := ai.Generate(ctx,
                    "Write a Twitter thread for: "+wc.GetString("outline"))
                wc.Set("social", resp.Text)
                return nil
            }),
        ).
        Step("draft", func(ctx context.Context, wc *ai.WorkflowContext) error {
            resp, _ := ai.Generate(ctx,
                "Write a 1500-word blog post from this outline:\n"+wc.GetString("outline"),
                ai.WithModel("gpt-4o"),
            )
            wc.Set("post", resp.Text)
            return nil
        })

    result, err := wf.Run(c.Request().Context(), ai.WorkflowInput{"topic": c.Request.FormValue("topic")})
    if err != nil {
        return err
    }

    return c.JSON(200, result.Data)
}

Document Processing

Chunk, summarize, and query documents with auto map-reduce for long texts.

// Summarize (auto map-reduce for texts exceeding token limits)
summary, err := ai.Summarize(ctx, longContractText,
    ai.SummarizeStyle("bullet-points"),
    ai.SummarizeMaxLength(500),
)

// Chunk text for embedding ingestion
chunks := ai.ChunkText(text, ai.ChunkSize(500), ai.ChunkOverlap(50))

// Q&A over a single document
qa := ai.DocumentQA(contractText)
answer, err := qa.Ask(ctx, "What are the payment terms?")

Image Generation

func (ctrl *Images) Generate(c *http.Context) error {
    img, err := ai.Image().
        Model("dall-e-3").
        Prompt(c.Request.FormValue("prompt")).
        Size("1024x1024").
        Style("natural").
        Generate(c.Request().Context())
    if err != nil {
        return err
    }

    return c.JSON(200, map[string]string{"url": img.Images[0].URL})
}

Guardrails

Validate AI outputs for length, blocked patterns, PII, and custom rules.

func (ctrl *Chat) SafeGenerate(c *http.Context) error {
    g := ai.NewGuardrails().
        MaxLength(5000).
        BlockPatterns(`(?i)password`, `\b\d{16}\b`). // block passwords & credit cards
        SetContentFilter(ai.FilterPII).
        CustomCheck(func(output string) error {
            if strings.Contains(output, "TODO") {
                return fmt.Errorf("response contains unfinished content")
            }
            return nil
        })

    resp, err := ai.Generate(c.Request().Context(), c.Request.FormValue("prompt"))
    if err != nil {
        return err
    }

    if err := g.Validate(resp.Text); err != nil {
        return c.JSON(400, map[string]string{"error": "Response filtered: " + err.Error()})
    }

    return c.JSON(200, map[string]string{"text": resp.Text})
}

Memory

Persist agent conversations across requests with pluggable backends.

Backend	Constructor	Use Case
In-memory	ai.MemoryStore()	Development / testing
Redis	ai.RedisMemory(rdb)	Production, multi-instance
Database	ai.DatabaseMemory(db)	Persistent, auditable
Sliding window	ai.SlidingWindowMemory(inner, n)	Keep last N messages
Summary	ai.SummaryMemory(inner, n)	Auto-compress history

Example: Redis Memory with Sliding Window

import "github.com/CodeSyncr/nimbus/redis"

func (ctrl *Chat) Message(c *http.Context) error {
    rdb := redis.NewClient(&redis.Options{Addr: "localhost:6379"})

    // Redis memory with TTL, wrapped in sliding window (keep last 100 messages)
    mem := ai.SlidingWindowMemory(
        ai.RedisMemory(rdb, ai.WithRedisTTL(24*time.Hour)),
        100,
    )

    userID := c.Session().GetString("user_id")
    agent := ai.NewAgent("You are a helpful assistant.").
        WithMemory(mem, "chat:"+userID)

    resp, err := agent.Prompt(c.Request().Context(), c.Request.FormValue("message"))
    if err != nil {
        return err
    }

    return c.JSON(200, map[string]string{"reply": resp.Text})
}

Observability & Tracing

Monitor every AI operation with hooks, metrics, and OpenTelemetry-compatible tracing.

Observability Hooks

// Log every AI request
ai.OnRequest(func(e ai.RequestEvent) {
    log.Printf("[AI] model=%s tokens=%d latency=%s",
        e.Model, e.Usage.TotalTokens, e.Latency)
})

// Track errors (send to Sentry, DataDog, etc.)
ai.OnError(func(e ai.RequestEvent) {
    sentry.CaptureException(e.Error)
})

// Aggregate metrics
fmt.Println(ai.UsageReport())
// → "AI Usage: requests=142 tokens=58320 errors=2 avg_latency=1.2s"

OpenTelemetry Tracing

// Enable tracing for all AI operations
ai.EnableTracing(ai.TracingConfig{
    ServiceName:   "my-app",
    RecordPrompts: false, // disable in prod for privacy
})

// Log spans to console (dev)
ai.AddSpanExporter(&ai.LogExporter{})

// Send to OTLP collector (prod)
ai.AddSpanExporter(&ai.OTLPExporter{
    Endpoint: "http://localhost:4318/v1/traces",
})

// Inspect recent spans
spans := ai.GetTraceSpans(10)
for _, span := range spans {
    fmt.Printf("%s %s %s\n", span.Name, span.Attributes["ai.model"], span.Duration)
}

Cost Tracking

Track per-model, per-provider costs with built-in pricing for 20+ models and budget alerts.

Example: Budget Monitoring

ai.EnableCostTracking(ai.CostConfig{
    MonthlyBudget:   500.00, // $500/month
    AlertThresholds: []float64{50, 75, 90, 100},
    OnBudgetAlert: func(usage ai.CostSummary) {
        log.Printf("⚠️  AI budget at %.0f%% ($%.2f / $%.2f)",
            usage.BudgetPercent, usage.TotalCost, usage.MonthlyBudget)
        // Send Slack/email notification here.
    },
})

Example: Cost Dashboard API

func (ctrl *Admin) CostDashboard(c *http.Context) error {
    dashboard := ai.GetCostDashboard()

    return c.JSON(200, map[string]any{
        "total_cost":       dashboard.TotalCost,
        "total_requests":   dashboard.TotalRequests,
        "budget_percent":   dashboard.BudgetPercent,
        "budget_remaining": dashboard.BudgetRemaining,
        "by_model":         dashboard.CostByModel,
        "by_provider":      dashboard.CostByProvider,
        "hourly":           dashboard.HourlyCosts,
    })
}

// CLI: print cost report
// fmt.Println(ai.CostReport())
// → AI Cost Report
//   Total Cost:     $12.3456
//   Total Requests: 423
//   Budget:         $500.00 (2.5% used, $487.65 remaining)
//   Per Model:
//     gpt-4o                          312 reqs  $11.2345  (avg $0.036013/req)
//     gpt-4o-mini                     111 reqs  $1.1111   (avg $0.010010/req)

Model Evaluation & Benchmarking

Test AI quality with automated benchmarks, compare models, and use LLM-as-judge for subjective scoring.

Example: Quality Test Suite

func TestAIQuality(t *testing.T) {
    suite := ai.NewEvalSuite("qa-quality").
        AddCase("greeting", "Say hello in a friendly way",
            ai.ExpectContains("hello"),
            ai.ExpectMinLength(10),
        ).
        AddCase("math", "What is 15 * 23?",
            ai.ExpectContains("345"),
        ).
        AddCase("json_output", "Return a JSON object with name and age fields",
            ai.ExpectJSON(),
            ai.ExpectContains("name"),
        ).
        AddCaseWithSystem("pirate",
            "You are a pirate", "Introduce yourself",
            ai.ExpectContains("arr"),
            ai.LLMJudge("creativity and pirate speech"),
        ).
        Parallel(3) // run 3 cases concurrently

    report := suite.Run(context.Background())

    if report.AvgScore < 0.8 {
        t.Errorf("AI quality score %.2f is below threshold 0.80", report.AvgScore)
    }
    t.Log(report.Summary())
}

Example: Model Comparison

func (ctrl *Admin) CompareModels(c *http.Context) error {
    suite := ai.NewEvalSuite("model-compare").
        AddCase("summarize", "Summarize Go's concurrency model in 3 sentences",
            ai.ExpectMinLength(50),
            ai.ExpectMaxLength(500),
            ai.LLMJudge("accuracy and conciseness"),
        ).
        AddCase("code", "Write a Go function to reverse a string",
            ai.ExpectContains("func"),
            ai.ExpectContains("string"),
        )

    comparison := ai.CompareModels(c.Request().Context(), suite,
        "gpt-4o", "gpt-4o-mini", "claude-3-5-sonnet-20241022",
    )

    return c.JSON(200, map[string]any{
        "suite":   comparison.Suite,
        "reports": comparison.Reports,
        "summary": comparison.Summary(),
    })
}

HTTP Middleware

// Rate-limit AI endpoints (10 requests per second per IP)
aiGroup := app.Router.Group("/api/ai")
aiGroup.Use(ai.RateLimit(10, time.Second))
aiGroup.Use(ai.CostGuard(1.00)) // max $1 per request
aiGroup.Use(ai.Logger())
aiGroup.Post("/chat", chatHandler)
aiGroup.Post("/generate", generateHandler)

Providers

Provider	Generate	Stream	Embeddings	Image
OpenAI	✓	✓	✓	✓
Anthropic	✓	✓	—	—
Gemini	✓	✓	—	—
Mistral	✓	✓	—	—
Cohere	✓	✓	—	—
xAI (Grok)	✓	✓	—	—
Ollama	✓	✓	—	—

Full Example: AI-Powered SaaS Feature

This example combines multiple AI features to build a complete contract analysis endpoint — extract data, classify risk, summarize, and track costs.

type ContractAnalysis struct {
    Parties       []string
    EffectiveDate string
    TermMonths    int
    TotalValue    float64
    KeyClauses    []string
}

func (ctrl *Contracts) Analyze(c *http.Context) error {
    ctx := c.Request().Context()
    text := c.Request.FormValue("contract_text")

    // 1. Extract structured data from the contract.
    analysis, err := ai.Extract[ContractAnalysis](ctx, text)
    if err != nil {
        return err
    }

    // 2. Classify risk level.
    risk, err := ai.Classify(ctx, text, []string{"low", "medium", "high"})
    if err != nil {
        return err
    }

    // 3. Generate executive summary.
    summary, err := ai.Summarize(ctx, text,
        ai.SummarizeStyle("executive-summary"),
        ai.SummarizeMaxLength(200),
    )
    if err != nil {
        return err
    }

    return c.JSON(200, map[string]any{
        "analysis": analysis,
        "risk":     risk,
        "summary":  summary,
    })
}

Related

• AI Video Pipeline — Cinematic video production: scene planning, keyframes, video synthesis, stitching, social packaging
• MCP — AI client integrations (tools, resources, prompts)