ai

Documentation

Overview

Package ai provides a minimal client for natural-language-to-SQL via a hosted (or self-hosted) chat-completions API. Phase 0 speaks the OpenAI Chat Completions wire format, which is also served by OpenAI-compatible local runtimes (Ollama's /v1 endpoint, LM Studio, vLLM, etc.), so the same client reaches a cloud provider or a localhost model with only a base-URL change. Adding providers with their own wire format (Anthropic, etc.) later is a matter of writing a second implementation of the Completer interface.

Index

• Constants
• Variables
• func AskSQL(ctx context.Context, cfg Config, schema, question string) (string, error)
• func ExtractSQL(reply string) string
• func SchemaContext(conn SchemaIntrospector) (string, error)
• func SystemPrompt(schema string) string
• type Client
  • func New(cfg Config) *Client
  • func (c *Client) Complete(ctx context.Context, messages []Message) (string, error)
  • func (c *Client) CompleteStream(ctx context.Context, messages []Message, onDelta func(StreamDelta)) (string, error)
  • func (c *Client) ListModels(ctx context.Context) ([]string, error)
• type Column
• type Config
• type ForeignKey
• type Message
• type SchemaIntrospector
• type StreamDelta

Constants

const (
	DefaultBaseURL = "https://api.openai.com/v1"
	DefaultModel   = "gpt-4o-mini"
	DefaultTimeout = 60 * time.Second
)

Defaults. BaseURL is the public OpenAI endpoint; pointing it at a local runtime (e.g. http://localhost:11434/v1 for Ollama) makes the whole feature run offline. Model is a cheap, broadly available default; override per use.

Variables

var ErrNoAPIKey = fmt.Errorf("no AI API key set — configure one to use :ai (see :help ai)")

ErrNoAPIKey is returned when no API key is configured. The UI maps this to a helpful, actionable status-bar message.

Functions

func AskSQL

func AskSQL(ctx context.Context, cfg Config, schema, question string) (string, error)

AskSQL is the high-level natural-language-to-SQL helper. It builds a system prompt from the database schema, appends the user's question, and returns just the SQL extracted from the model's reply.

func ExtractSQL

func ExtractSQL(reply string) string

ExtractSQL pulls the SQL statement out of a model reply that may be wrapped in markdown fences and/or surrounded by prose. Models are inconsistent, so this is intentionally lenient: it prefers fenced blocks, then falls back to the first ;-terminated statement, then to the whole trimmed reply.

It strips leading SQL dialect markers ("sql", "SQL", "postgres", …) from a fenced block and trims trailing semicolons only if the statement is a single one — leaving the caller (the editor) in control of execution.

func SchemaContext

func SchemaContext(conn SchemaIntrospector) (string, error)

SchemaContext renders the connected database as compact pseudo-DDL for the model's system prompt. Each table is one block: its columns (with PK / FK annotations inline as SQL comments) so the model sees names, types, and relationships without extra round-trips. Tables with unreadable schemas are listed by name only so the model at least knows they exist.

The returned string is prefixed with a short legend so the model can parse the inline annotations.

func SystemPrompt

func SystemPrompt(schema string) string

SystemPrompt instructs the model to return a single read-only SQL statement for the given schema. It forbids prose and asks for a SELECT so a misunderstood request cannot mutate data — a safety default we can relax per-command once the assistant panel has explicit "run" gating. Exported so the UI layer can assemble multi-turn conversations (system + prior turns).

Types

type Client

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

Client talks to a chat-completions endpoint. The HTTP client is overridable for testing.

func New

func New(cfg Config) *Client

New returns a Client with defaults applied. It performs no network I/O.

func (*Client) Complete

func (c *Client) Complete(ctx context.Context, messages []Message) (string, error)

Complete sends the conversation and returns the assistant's reply text. The context bounds the whole request (dial + headers + body read); callers should derive it from the user's cancel action so esc can abort a slow model.

func (*Client) CompleteStream

func (c *Client) CompleteStream(ctx context.Context, messages []Message, onDelta func(StreamDelta)) (string, error)

CompleteStream is the streaming variant of Complete: it posts with stream:true and invokes onDelta for each token as it arrives (Server-Sent Events) — content in StreamDelta.Content, chain-of-thought (reasoning models) in StreamDelta.Reasoning — returning the full accumulated reply. This lets the UI render the answer (and the model's thinking) forming in real time instead of blocking on the whole response. onDelta may be nil. The same context bounds the call.

func (*Client) ListModels

func (c *Client) ListModels(ctx context.Context) ([]string, error)

ListModels queries the provider's OpenAI-compatible GET /models endpoint and returns the available model ids. The assistant-panel model browser uses it so users can pick a model that is live for their key rather than guessing.

type Column

type Column struct {
	Name string
	Type string
}

Column mirrors db.Column for the introspector interface (name + type only).

type Config

type Config struct {
	APIKey  string
	BaseURL string
	Model   string
	Timeout time.Duration
}

Config holds everything the client needs to reach a provider. Zero values are filled in by New with the package defaults. APIKey is required for hosted providers; local runtimes typically ignore it (pass any non-empty placeholder since some refuse an empty header).

type ForeignKey

type ForeignKey struct {
	Column    string
	RefTable  string
	RefColumn string
}

ForeignKey mirrors db.ForeignKey for the introspector interface.

type Message

type Message struct {
	Role    string `json:"role"`
	Content string `json:"content"`
}

Message is a single chat-completions message. Role is one of "system", "user", or "assistant".

type SchemaIntrospector

type SchemaIntrospector interface {
	Tables() ([]string, error)
	TableSchema(table string) ([]Column, error)
	PrimaryKeys(table string) ([]string, error)
	ForeignKeys(table string) ([]ForeignKey, error)
}

SchemaIntrospector is the slice of db.DB / db.Connection the AI package needs to describe a database to the model. Defining it locally keeps the ai package free of an internal/db import and trivial to fake in tests.

type StreamDelta

type StreamDelta struct {
	Content   string
	Reasoning string
}

StreamDelta is one piece of a streamed reply: either a content token (the visible answer) or a reasoning token (chain-of-thought, from reasoning models). At most one field is non-empty per call.