AI SDK

go get github.com/velocitykode/velocity-ai

import (
    _ "github.com/velocitykode/velocity-ai/provider/anthropic"
    _ "github.com/velocitykode/velocity-ai/provider/cohere"
    _ "github.com/velocitykode/velocity-ai/provider/jina"
    _ "github.com/velocitykode/velocity-ai/provider/openai"
    _ "github.com/velocitykode/velocity-ai/provider/typesafe"
)

package app

import (
    "context"

    "github.com/velocitykode/velocity"

    "github.com/velocitykode/velocity-ai/manager"
)

type AIModule struct {
    ai manager.Lifecycle
}

// Init creates the Manager from the environment and registers it.
func (m *AIModule) Init(s *velocity.Services) error {
    return m.ai.Register(s)
}

// Start hands the app's event dispatcher, queue and cache to the Manager.
func (m *AIModule) Start(s *velocity.Services) error {
    return m.ai.Boot(s)
}

// Shutdown has nothing to release: the framework closes the registered
// Manager when the app stops.
func (m *AIModule) Shutdown(ctx context.Context) error {
    return m.ai.Shutdown(ctx)
}

func Configure(reg *velocity.ModuleRegistry) {
    reg.Add(&AIModule{})
}

func Summary(ctx *router.Context) error {
    s, err := ctx.Services()
    if err != nil {
        return err
    }
    mgr, err := manager.From(s)
    if err != nil {
        return err
    }

    a := agent.New("You write one-paragraph release notes.").Build()

    resp, err := mgr.Agent(a).Prompt(ctx.Request.Context(), "Summarize the changes in v1.4.")
    if err != nil {
        return err
    }
    return ctx.JSON(router.StatusOK, map[string]string{"text": resp.Text})
}

cfg := config.Config{
    DefaultTextProvider: "anthropic",
    Providers: map[string]config.ProviderConfig{
        "anthropic": {APIKey: os.Getenv("ANTHROPIC_API_KEY")},
    },
}
mgr := manager.NewManager(&cfg)

AI_DEFAULT_TEXT_PROVIDER=anthropic
AI_DEFAULT_EMBEDDING_PROVIDER=openai
AI_DEFAULT_RERANKING_PROVIDER=cohere
AI_DEFAULT_CLASSIFICATION_PROVIDER=typesafe

ANTHROPIC_API_KEY=
OPENAI_API_KEY=
COHERE_API_KEY=
TYPESAFE_API_KEY=

a := agent.New("You are a support assistant. Answer in two sentences.").Build()

resp, err := mgr.Agent(a).
    WithProvider("anthropic").
    WithModel("claude-haiku-4-5-20251001").
    WithMaxTokens(512).
    WithTemperature(0.2).
    Prompt(ctx, "How do I rotate an API key?")
if err != nil {
    return err
}

fmt.Println(resp.Text)
fmt.Println(resp.Model, resp.PromptTokens, resp.CompletionTokens, resp.TotalTokens)

// The provider's cheapest model, resolved when the request runs.
resp, err := mgr.Agent(a).UseCheapestModel().Prompt(ctx, "Tag this ticket.")

// The provider's most capable model.
resp, err = mgr.Agent(a).UseSmartestModel().Prompt(ctx, "Review this contract clause.")

// No agent value: instructions inline.
resp, err = mgr.AnonymousAgent("Reply with one word.").Prompt(ctx, "Capital of France?")

type Triage struct{}

func (Triage) Instructions() string { return "You triage incoming support tickets." }

func (Triage) Config() agent.AgentConfig {
    return agent.AgentConfig{Provider: "openai", UseCheapest: true, MaxSteps: 4}
}

stream, err := mgr.Agent(a).Stream(ctx, "Explain context cancellation.")
if err != nil {
    return err
}
defer stream.Close()

for ev := range stream.Events() {
    fmt.Print(ev.Delta)
}

resp, err := stream.Response()
if err != nil {
    return err
}
fmt.Println(resp.Usage.TotalTokens)

stream, err := mgr.Agent(a).PromptStream(ctx, "What is the weather in Lahore?")
if err != nil {
    return err
}
defer stream.Close()

for ev := range stream.Events() {
    switch ev.Kind {
    case manager.StreamTextDelta:
        fmt.Print(ev.Text)
    case manager.StreamToolInvoking:
        fmt.Println("calling", ev.ToolName, ev.ToolParams)
    case manager.StreamToolInvoked:
        fmt.Println("result", ev.ToolResult)
    case manager.StreamDone:
        fmt.Println(ev.Response.TotalTokens)
    case manager.StreamError:
        return ev.Err
    }
}

func Chat(ctx *router.Context) error {
    s, err := ctx.Services()
    if err != nil {
        return err
    }
    mgr, err := manager.From(s)
    if err != nil {
        return err
    }

    a := agent.New("You are a helpful assistant.").Build()

    stream, err := mgr.Agent(a).Stream(ctx.Request.Context(), ctx.Query("q"))
    if err != nil {
        return err
    }
    return provider.WriteVercelDataStream(ctx.Response, stream)
}

weather := tool.NewTool(
    "get_weather",
    "Get the current temperature for a city.",
    agent.Object(map[string]agent.SchemaField{
        "city":  agent.String().Description("City name"),
        "units": agent.Enum("metric", "imperial").Optional(),
    }).Build(),
    func(ctx context.Context, params map[string]any) (string, error) {
        city, _ := params["city"].(string)
        if city == "" {
            return "", errors.New("city is required")
        }
        return "21 degrees and clear in " + city, nil
    },
)

a := agent.New("Answer weather questions with the get_weather tool.").
    WithTools(weather).
    Build()

resp, err := mgr.Agent(a).WithMaxSteps(5).Prompt(ctx, "Is it warm in Lahore?")
if errors.Is(err, manager.ErrMaxStepsExceeded) {
    // The model was still asking for tools after five provider calls.
    return err
}
if err != nil {
    return err
}
fmt.Println(resp.Text)

embedder, err := mgr.EmbeddingProvider()
if err != nil {
    return nil, err
}

a := agent.New("Research the question, then answer with sources.").
    WithTools(
        tool.WebSearch(tool.WebSearchConfig{APIKey: os.Getenv("TAVILY_API_KEY"), MaxResults: 3}),
        tool.WebFetch(tool.WebFetchConfig{Timeout: 10, MaxOutputLen: 20000}),
        tool.SimilaritySearch(tool.SimilaritySearchConfig{
            Provider: embedder,
            TopK:     3,
            Corpus: []tool.Document{
                {ID: "refunds", Text: "Refunds are issued within 14 days of purchase."},
                {ID: "shipping", Text: "Orders ship within two business days."},
            },
        }),
    ).
    Build()

type Sentiment struct {
    Label string   `json:"label" description:"positive, neutral or negative"`
    Score int      `json:"score" description:"Strength from 1 to 10"`
    Tags  []string `json:"tags,omitempty" description:"Topics mentioned"`
}

a := agent.New("Classify the sentiment of the review.").
    WithSchema(agent.SchemaFromStruct(Sentiment{})).
    Build()

resp, err := mgr.Agent(a).Prompt(ctx, "Arrived late, but the build quality is excellent.")
if err != nil {
    return err
}

var s Sentiment
if err := resp.Unmarshal(&s); err != nil {
    return err
}
fmt.Println(s.Label, s.Score, resp.Get("tags"))

schema := agent.Object(map[string]agent.SchemaField{
    "label": agent.Enum("positive", "neutral", "negative"),
    "score": agent.Int().Description("Strength from 1 to 10"),
    "tags":  agent.Array(agent.String()).Optional(),
}).Build()

store := conversation.NewMemoryStore()

// First turn: ForUser opens a new conversation owned by that user.
first, err := mgr.Agent(a).WithStore(store).ForUser("user-42").
    Prompt(ctx, "My order number is 1182.")
if err != nil {
    return err
}

// Later turns: Continue loads the stored messages and appends to them.
next, err := mgr.Agent(a).WithStore(store).Continue(first.ConversationID).
    Prompt(ctx, "What was my order number?")
if err != nil {
    return err
}
fmt.Println(next.Text)

// Resume the most recent conversation of a user.
id, err := store.LatestConversationID(ctx, "user-42")
if err != nil {
    return err
}
fmt.Println(id == first.ConversationID)

import _ "github.com/velocitykode/velocity-ai/conversation/migrations"

store := conversation.NewORMStore(s.DB)

resp, err := mgr.Embeddings().
    For("refund policy", "shipping times").
    WithDimensions(1536).
    Cached(24 * time.Hour).
    Generate(ctx)
if err != nil {
    return err
}
fmt.Println(len(resp.Embeddings), len(resp.Embeddings[0]), resp.Usage.TotalTokens)

docs := []string{
    "Refunds are issued within 14 days of purchase.",
    "Orders ship within two business days.",
    "Gift cards cannot be refunded.",
}

resp, err := mgr.Rerank(docs, "cohere").Query("can I get my money back?").Limit(2).Generate(ctx)
if err != nil {
    return err
}
for _, r := range resp.Results {
    fmt.Printf("%.3f %s\n", r.Score, docs[r.Index])
}

img, err := mgr.Image("A lighthouse at dusk, flat illustration").
    Landscape().
    Quality("high").
    Generate(ctx)
if err != nil {
    return err
}
if err := os.WriteFile("lighthouse.png", img.Data, 0o644); err != nil {
    return err
}

speech, err := mgr.Audio("Your order has shipped.").Voice("alloy").Generate(ctx)
if err != nil {
    return err
}
if err := os.WriteFile("shipped.mp3", speech.Data, 0o644); err != nil {
    return err
}

tr, err := mgr.Transcription(agent.FileFromPath("call.mp3")).Language("en").Generate(ctx)
if err != nil {
    return err
}
fmt.Println(tr.Text)
for _, seg := range tr.Segments {
    fmt.Printf("%.1fs-%.1fs %s\n", seg.Start, seg.End, seg.Text)
}

stored, err := mgr.Files().Upload(ctx, agent.FileFromPath("handbook.pdf"))
if err != nil {
    return err
}

store, err := mgr.Stores().Create(ctx, "handbook", provider.StoreOptions{
    Description:        "Employee handbook",
    ExpiresWhenIdleFor: 30 * 24 * time.Hour,
})
if err != nil {
    return err
}

docID, err := mgr.Stores().AddFile(ctx, store.ID, stored.ID, map[string]any{"year": 2026})
if err != nil {
    return err
}
fmt.Println(docID)

if err := mgr.Stores().Delete(ctx, store.ID); err != nil {
    return err
}
return mgr.Files().Delete(ctx, stored.ID)

urgent, err := classify.Boolean("Is this message urgent?", "It reports an outage or data loss.")
if err != nil {
    return err
}

topic, err := classify.Choice("What is the message about?", []classify.Option{
    {Name: "billing", Description: "Invoices, charges, refunds"},
    {Name: "bug", Description: "Something is broken"},
    {Name: "other"},
})
if err != nil {
    return err
}

anger, err := classify.Score("How angry is the writer?", []classify.Level{
    {Value: 1, Label: "calm"},
    {Value: 2, Label: "annoyed"},
    {Value: 3, Label: "furious"},
})
if err != nil {
    return err
}

resp, err := mgr.Classify(ctx, "Production is down and nobody is answering!",
    urgent.Named("urgent"), topic.Named("topic"), anger.Named("anger"))
if err != nil {
    return err
}

if a, ok := resp.Boolean("urgent"); ok {
    fmt.Println(a.Probability(), a.IsTrue(0.8))
}
if a, ok := resp.Choice("topic"); ok {
    fmt.Println(a.Label(), a.Probability(), a.ProbabilityOf("billing"))
}
if a, ok := resp.Score("anger"); ok {
    fmt.Println(a.Level().Label, a.Score(), a.Normalized())
}

resp, err := mgr.Classifier("typesafe").
    WithTimeout(5*time.Second).
    Classify(ctx, text, topic.Named("topic"))
if err != nil {
    return err
}

a, _ := resp.Choice("topic")
if c, ok := a.Confidence(); ok {
    fmt.Println("confidence", c)
} else {
    fmt.Println("the provider supplied no confidence")
}

summary, err := mgr.Summarize(ctx, article,
    manager.SummarizeSentences(2),
    manager.HelperTimeout(20*time.Second),
)
if err != nil {
    return err
}
fmt.Println(summary)

spam, err := mgr.Decide(ctx, comment, "Is this comment spam?",
    manager.DecideCriteria("It advertises a product or links to one.", "It discusses the article."),
    manager.DecideThreshold(0.9),
    manager.HelperProvider("typesafe"),
)
if err != nil {
    return err
}
fmt.Println(spam)

resp, err := mgr.Agent(a).
    WithProviders("anthropic", "openai").
    UseCheapestModel().
    Prompt(ctx, "Draft a two-line status update.")

handle, err := mgr.Agent(a).UseCheapestModel().Dispatch("Write the weekly digest.", "ai")
if errors.Is(err, manager.ErrQueueNotConfigured) {
    return err
}
if err != nil {
    return err
}

handle.Then(func(resp *agent.AgentResponse, err error) {
    if err != nil {
        return
    }
    fmt.Println(resp.Text)
})

worker := queue.NewWorker(s.Queue, "ai", func(j queue.Job) error {
    return j.Handle()
}, queue.WithConcurrency(2), queue.WithMaxRetries(1))
worker.Start(ctx)

resp, err := mgr.Agent(a).WithBroadcastAs("digest.ready").
    BroadcastNow(ctx, "Write the weekly digest.", "team.42")
if err != nil {
    return err
}
fmt.Println(resp.Text)

type redact struct{}

func (redact) Handle(prompt *agent.AgentPrompt, next func(*agent.AgentPrompt) (any, error)) (any, error) {
    for i, m := range prompt.Messages {
        prompt.Messages[i].Content = scrub(m.Content)
    }

    out, err := next(prompt)
    if err != nil {
        return nil, err
    }
    if resp, ok := out.(*provider.TextResponse); ok {
        resp.Text = scrub(resp.Text)
    }
    return out, nil
}

// Global: runs on every agent prompt made through this Manager.
mgr.Use(middleware.LogPrompts(func(p *agent.AgentPrompt, response any, d time.Duration, err error) {
    logger.Info("ai prompt", "model", p.Model, "messages", len(p.Messages), "duration", d, "error", err)
}))

// Per agent.
a := agent.New("You answer billing questions.").
    WithMiddleware(redact{}).
    Build()

// Per request.
resp, err := mgr.Agent(a).
    WithMiddleware(
        middleware.RateLimit(20, time.Minute, func(p *agent.AgentPrompt) string { return p.Provider }),
        middleware.Cache(10*time.Minute, func(p *agent.AgentPrompt) string {
            return p.Instructions + "\x00" + p.Messages[len(p.Messages)-1].Content
        }),
    ).
    Prompt(ctx, "When is my next invoice?")
if errors.Is(err, middleware.ErrRateLimited) {
    return
}

func Events(logger log.Logger) func(events.Dispatcher) {
    return func(d events.Dispatcher) {
        d.Listen("ai.response.received", listenerFunc(func(_ context.Context, e interface{}) error {
            if ev, ok := e.(*event.ResponseReceived); ok {
                logger.Info("ai response", "provider", ev.Provider, "model", ev.Model, "tokens", ev.Usage.TotalTokens)
            }
            return nil
        }))

        d.Listen("ai.provider.error", listenerFunc(func(_ context.Context, e interface{}) error {
            if ev, ok := e.(*event.ProviderError); ok {
                logger.Error("ai provider error", "provider", ev.Provider, "model", ev.Model, "error", ev.Err)
            }
            return nil
        }))
    }
}

type listenerFunc func(ctx context.Context, e interface{}) error

func (f listenerFunc) Handle(ctx context.Context, e interface{}) error { return f(ctx, e) }
func (f listenerFunc) Async() bool                                     { return false }

resp, err := mgr.Agent(a).Prompt(ctx, "Hello")
switch {
case err == nil:
    fmt.Println(resp.Text)
case errors.Is(err, manager.ErrNoDefaultProvider):
    // AI_DEFAULT_TEXT_PROVIDER is unset and no provider was named.
case errors.Is(err, manager.ErrProviderNotConfigured):
    // The named provider has no API key in the environment.
case errors.Is(err, manager.ErrCapabilityNotSupported):
    // The provider exists but does not implement this capability.
case errors.Is(err, manager.ErrMaxStepsExceeded):
    // The tool loop hit its step limit.
case errors.Is(err, context.DeadlineExceeded):
    // The caller's context expired.
default:
    // A provider or transport error.
}

// Reply drafts an answer to a ticket unless the ticket is spam.
func Reply(ctx context.Context, mgr *manager.Manager, ticket string) (string, error) {
    spam, err := mgr.Decide(ctx, ticket, "Is this ticket spam?")
    if err != nil {
        return "", err
    }
    if spam {
        return "", nil
    }

    a := agent.New("You draft replies to support tickets.").Build()
    resp, err := mgr.Agent(a).Prompt(ctx, ticket)
    if err != nil {
        return "", err
    }
    return resp.Text, nil
}

package support

import (
    "context"
    "testing"

    "github.com/velocitykode/velocity-ai/fake"
)

func TestReply(t *testing.T) {
    text := fake.Text(fake.TextResponse("Thanks, we are on it.")).PreventStray()
    classifier := fake.Classification().WithDecision(false)

    mgr := fake.NewManager(t,
        fake.WithText(text),
        fake.WithClassification(classifier),
    )

    got, err := Reply(context.Background(), mgr, "The export button does nothing.")
    if err != nil {
        t.Fatal(err)
    }
    if got != "Thanks, we are on it." {
        t.Fatalf("got %q", got)
    }

    classifier.AssertDecided(t, "The export button does nothing.", "Is this ticket spam?")
    text.AssertCallCount(t, 1)
    text.AssertCalledWith(t, "The export button does nothing.")
}

func TestReply_Spam(t *testing.T) {
    text := fake.Text().PreventStray()
    classifier := fake.Classification().WithDecision(true)

    mgr := fake.NewManager(t, fake.WithText(text), fake.WithClassification(classifier))

    got, err := Reply(context.Background(), mgr, "Buy cheap watches")
    if err != nil {
        t.Fatal(err)
    }
    if got != "" {
        t.Fatalf("got %q, want no reply", got)
    }
    text.AssertNotCalled(t)
}