Skip to content

Latest commit

 

History

History
203 lines (164 loc) · 9.03 KB

File metadata and controls

203 lines (164 loc) · 9.03 KB

rig — design contract

Minimal agent loop for OpenAI-compatible self-hosted models. Extracted from and first consumed by github.com/RandomCodeSpace/kb (see internal/server/ai.go there for the battle-tested reference behavior).

Single package rig at the module root. Go 1.25. Dependency: github.com/openai/openai-go/v3 v3.50.0 only (plus its transitive deps). No other third-party dependencies.

Design goals, in priority order:

  1. Survive self-hosted reasoning models (Qwen, DeepSeek-style) whose serving templates leak chain-of-thought into message content.
  2. Hardened by default against user-configured endpoints: response size cap, ambient credential scrubbing, explicit token budgets on every call.
  3. Small, boring API. This is a loop and a toolbox, not a framework.

File layout

File Contents
client.go Client, NewClient, options, probe
content.go StripReasoning, DecodeJSONObject
tool.go Tool type, validation
loop.go RunRequest/RunResult, Client.Run loop
skill.go Skill type, LoadSkills, Advertise, LoadSkillTool
errors.go sentinel errors
*_test.go table tests per file; httptest fake backends for client/loop

client.go

type Client struct { /* wraps openai.Client + config */ }

type Option func(*config)

func WithHTTPClient(hc *http.Client) Option      // caller supplies hardened transport (SSRF guard lives in the caller)
func WithMaxResponseBytes(n int64) Option        // default 1 << 20; cap applies to the decompressed body
func WithUserAgent(ua string) Option

func NewClient(baseURL, apiKey string, opts ...Option) (*Client, error)

Behavior (port from kb's aiSDKOptions/newAIClient):

  • Always append option.WithHeaderDel for OpenAI-Organization, OpenAI-Project, and any ambient custom headers; never let OPENAI_API_KEY/OPENAI_ORG_ID/OPENAI_PROJECT_ID/ OPENAI_CUSTOM_HEADERS env values reach the wire — construct the SDK client with explicit values only (option.WithAPIKey, option.WithBaseURL), not environment fallbacks.
  • Response cap: SDK middleware wrapping the response body in a limited reader over the decompressed stream; exceeding the cap returns ErrResponseTooLarge.
  • baseURL must be http(s) and non-empty; apiKey may be empty (Ollama).
func (c *Client) ProbeToolCalling(ctx context.Context, model string) error

Sends a trivial single-tool request (max_tokens ~256) and returns nil only if the reply contains a tool call. Error otherwise (ErrNoToolCalling wrapped with detail). Port of kb's /api/ai/test probe semantics. A reply with finish_reason length is ErrOutputLimit, not ErrNoToolCalling: a reasoning model can spend the whole probe budget before its first call. No reply text -- content or finish_reason, both endpoint-controlled -- is interpolated into a returned error.

content.go

Direct port of kb's reasoning handling (internal/server/ai.go: thinkBlockRe, decodeJSONObject) with identical semantics — this code is regression-tested against three real-world failure modes; do not "improve" it:

func StripReasoning(s string) string
// removes closed <think>...</think> blocks; if a bare closing </think> remains
// (template consumed the opening tag), drops everything through the last one; an
// opening <think> left unclosed drops everything from it on (a reply that stopped
// inside the reasoning has no answer); trims space.

func DecodeJSONObject(s string) (map[string]any, error)
// StripReasoning first; try whole-string unmarshal; else scan every '{' offset,
// decode candidates, return the LONGEST parseable object (json.Decoder InputOffset).
// Errors: ErrNotJSON if no '{' present, ErrInvalidJSON if candidates all fail.

The scan resumes past a candidate that parsed (a nested object is shorter than the one holding it) and stops after a fixed total of input handed to json.Decoder. Both bound a scan that is otherwise quadratic on a reply the endpoint chooses -- a megabyte of {"a": fits the response cap and cost minutes of one core.

tool.go

type Tool struct {
    Name        string
    Description string
    InputSchema map[string]any                                        // JSON Schema object
    Run         func(ctx context.Context, input json.RawMessage) (string, error)
}

Validation helper func validateTools(tools []Tool) error: non-empty unique names matching ^[a-zA-Z0-9_-]{1,64}$, non-nil Run, non-empty Description.

loop.go

type RunRequest struct {
    Model         string
    System        string   // system prompt; skills advertisement appended by caller
    Prompt        string   // initial user message
    Tools         []Tool
    MaxTokens     int64    // required, >0; no silent defaults
    MaxIterations int      // default 8 when 0; hard cap 32
    Temperature   float64  // omitted from the wire when 0
}

type RunResult struct {
    Text       string        // final assistant text, StripReasoning applied
    Iterations int           // completed request/response rounds
    ToolCalls  int           // total tool invocations executed
}

func (c *Client) Run(ctx context.Context, req RunRequest) (*RunResult, error)

Loop semantics:

  • Build openai-go params (tools converted via InputSchema), send, inspect choice 0.
  • finish_reason length → ErrOutputLimit (truncated replies are errors, not results).
  • If the message has tool calls: execute each sequentially in order via Tool.Run; a tool error becomes the tool result string "error: <msg>" fed back to the model (the loop does not abort on tool errors); append assistant message + tool results; next iteration.
  • No tool calls → terminal: return RunResult with stripped text. A reply that carried tool calls the loop cannot execute (a custom-tool dialect, a call without a name) is an error instead: an empty answer with a nil error tells the caller nothing.
  • More than 32 tool calls in one reply → ErrIterationLimit, none of them executed. The iteration cap bounds rounds, not work, and one reply inside the response cap has room for thousands of side effects.
  • The echoed assistant turn carries what the loop acted on: StripReasoning applied to content, blank arguments normalized to {}, and empty or duplicate tool call ids replaced (an upstream pairs results to calls by that id).
  • Iteration cap reached with the model still calling tools → ErrIterationLimit.
  • Transport/API errors are wrapped, never exposing the raw endpoint URL or key (opaque mapping: fmt.Errorf("ai backend: %w", ...) is fine, no URL interpolation).
  • ctx cancellation respected between iterations and passed to Tool.Run.

skill.go

Skill file format: markdown with YAML-ish frontmatter delimited by --- lines, keys name: and description: (single-line values, no YAML dependency — parse the two keys manually). Body is everything after the closing delimiter.

type Skill struct {
    Name        string
    Description string
    Body        string
}

func LoadSkills(fsys fs.FS, dir string) ([]Skill, error)
// loads *.md under dir; skips files with missing/duplicate names (collect into error? no —
// return error listing offenders; a broken skill file should be loud).

func MergeSkills(base, overrides []Skill) []Skill
// override wins by name; result sorted by name. Callers load embedded defaults as base
// and a user dir (e.g. KB_DATA/skills) as overrides.

func Advertise(skills []Skill) string
// deterministic block for the system prompt: one "- name: description" line per skill,
// preceded by a fixed instruction to call load_skill before acting on one.

func LoadSkillTool(skills []Skill) Tool
// built-in tool `load_skill` with input {"name": string}; returns the skill body,
// or an error string listing available names on a miss.

errors.go

var (
    ErrNoToolCalling    = errors.New("backend did not return a tool call")
    ErrOutputLimit      = errors.New("reply hit the output token limit")
    ErrIterationLimit   = errors.New("tool loop hit the iteration limit")
    ErrResponseTooLarge = errors.New("response body exceeds the configured cap")
    ErrNotJSON          = errors.New("reply is not JSON")
    ErrInvalidJSON      = errors.New("reply is not valid JSON")
)

All returned wrapped (%w) so callers use errors.Is.

Testing requirements

  • content: port kb's decodeJSONObject test table (closed think blocks, consumed opening tag, schema sketch before longer reply, braces in prose, fenced JSON) — see kb's internal/server/ai_test.go for the cases; same expectations.
  • client: httptest server asserting scrubbed headers absent, response over cap → ErrResponseTooLarge, probe with/without tool_calls.
  • loop: httptest fake OpenAI backend scripted per-iteration; cover: no-tool terminal reply, single tool round-trip, tool error fed back, iteration cap, finish_reason length, reasoning-leaking final text stripped.
  • skill: frontmatter parsing, override merge, advertisement determinism, load_skill hit/miss.
  • go vet ./... clean; no test may hit the network.

Style

Standard library first. Table tests. No stuttering names (rig.RigClient — no). Comments state constraints, not narration. Plain ASCII.