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:
- Survive self-hosted reasoning models (Qwen, DeepSeek-style) whose serving templates leak chain-of-thought into message content.
- Hardened by default against user-configured endpoints: response size cap, ambient credential scrubbing, explicit token budgets on every call.
- Small, boring API. This is a loop and a toolbox, not a framework.
| 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 |
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.WithHeaderDelforOpenAI-Organization,OpenAI-Project, and any ambient custom headers; never letOPENAI_API_KEY/OPENAI_ORG_ID/OPENAI_PROJECT_ID/OPENAI_CUSTOM_HEADERSenv 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) errorSends 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.
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.
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.
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 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.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.
- 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.
Standard library first. Table tests. No stuttering names (rig.RigClient — no). Comments state constraints, not narration. Plain ASCII.