chimerical/odidere
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Refactor llm package

Browse Source
This commit is contained in:
dwrz
2026-09-20 12:49:42 +00:00
parent eba6c2a028
commit 7224fea99f
11 changed files with 1176 additions and 1627 deletions
Download Patch File Download Diff File Expand all files Collapse all files
internal/llm/client.go
+48 -38
View File
@@ -7,20 +7,34 @@ import (
"time"
)
// Config holds the configuration for an LLM client.
const (
// chatCompletionsPath is the endpoint for chat completion requests.
chatCompletionsPath = "/chat/completions"
// chatTemplateEnableThinking is the chat-template argument that selects
// whether a reasoning model emits reasoning tokens.
chatTemplateEnableThinking = "enable_thinking"
// defaultTimeout limits requests when Config.Timeout is empty.
defaultTimeout = 5 * time.Minute
// sseDone is the sentinel that terminates a chat completion stream.
sseDone = "[DONE]"
)
// Config configures a Client.
type Config struct {
// Key is the API key for authentication.
// Key is sent as a bearer token. An empty key disables authentication.
Key string `yaml:"key"`
// SystemMessage is prepended to all conversations.
SystemMessage string `yaml:"system_message"`
// Timeout is the maximum duration for a query (e.g., "5m").
// Defaults to 5 minutes if empty.
// Provider is returned unchanged by Client.Provider. It is not sent to the
// inference endpoint.
Provider string `yaml:"provider"`
// Timeout limits each request, including reading a streaming response. It
// uses Go duration syntax and defaults to five minutes when empty.
Timeout string `yaml:"timeout"`
// URL is the base URL of the OpenAI-compatible API endpoint.
// URL is the endpoint root. Client appends /chat/completions or /models.
URL string `yaml:"url"`
}
// Validate checks that required configuration values are present and valid.
// Validate reports whether Timeout is a Go duration, when set, and URL is
// non-empty. It does not contact the endpoint.
func (cfg Config) Validate() error {
if cfg.Timeout != "" {
if _, err := time.ParseDuration(cfg.Timeout); err != nil {
@@ -33,30 +47,27 @@ func (cfg Config) Validate() error {
return nil
}
// Client wraps an OpenAI-compatible client with tool execution support.
// Client is an OpenAI-compatible inference client.
type Client struct {
baseURL string
httpc *http.Client
key string
log *slog.Logger
registry *Registry
systemMessage string
timeout time.Duration
tools []APITool
baseURL string
httpc *http.Client
key string
log *slog.Logger
provider string
timeout time.Duration
}
// NewClient creates a new LLM client with the provided configuration.
// The registry is optional; if nil, tool calling is disabled.
func NewClient(
cfg Config,
registry *Registry,
log *slog.Logger,
) (*Client, error) {
// NewClient creates a client from cfg. If cfg.Timeout is empty, requests have a
// five-minute timeout. log must be non-nil.
func NewClient(cfg Config, log *slog.Logger) (*Client, error) {
if err := cfg.Validate(); err != nil {
return nil, fmt.Errorf("invalid config: %w", err)
}
if log == nil {
return nil, fmt.Errorf("missing logger")
}
timeout := 5 * time.Minute
timeout := defaultTimeout
if cfg.Timeout != "" {
d, err := time.ParseDuration(cfg.Timeout)
if err != nil {
@@ -65,18 +76,17 @@ func NewClient(
timeout = d
}
llm := &Client{
baseURL: cfg.URL,
key: cfg.Key,
log: log,
systemMessage: cfg.SystemMessage,
registry: registry,
httpc: &http.Client{},
timeout: timeout,
}
return &Client{
baseURL: cfg.URL,
key: cfg.Key,
log: log,
provider: cfg.Provider,
httpc: &http.Client{},
timeout: timeout,
}, nil
}
// Convert tools from registry.
llm.tools = registry.ToAPI()
return llm, nil
// Provider returns the application-defined provider name supplied in Config.
func (c *Client) Provider() string {
return c.provider
}
internal/llm/client_completions.go
+280 -352
View File
@@ -9,119 +9,70 @@ import (
"io"
"log/slog"
"net/http"
"slices"
"strings"
)
// ErrMaxIterations is returned when the agent loop exceeds the configured
// maximum number of iterations.
var ErrMaxIterations = fmt.Errorf("max iterations exceeded")
// CompletionsResult bundles the messages generated during a completions
// request along with usage and timings from the final API response.
type CompletionsResult struct {
Messages []Message
Usage *Usage
Timings *Timings
}
// CompletionsParams holds parameters for Completions and CompletionsStream.
// CompletionsParams contains the inputs to a chat completion.
type CompletionsParams struct {
Model string
SystemMessage string
MaxIterations int
Messages []Message
// Messages is the conversation presented to the model, in order.
Messages []Message
// Model identifies the model according to the configured endpoint.
Model string
// ThinkingBudgetTokens is sent as the endpoint's reasoning-token budget.
// Nil omits both the budget and enable_thinking; zero sets enable_thinking
// to false; any other value sets it to true.
ThinkingBudgetTokens *int
// Tools describes the functions available for the model to call.
Tools []APITool
}
// Completions sends messages to the LLM using the specified model.
// If systemMessage is non-empty, it overrides the configured system message.
// maxIterations caps the number of agent loop iterations.
// A value of 0 means unlimited (no cap).
// Returns all messages generated during the completions request, including
// tool calls and tool results, along with usage and timings from the final response.
// CompletionResult contains one generated Message and its response metadata.
type CompletionResult struct {
// FinishReason is the endpoint's reason for ending generation.
FinishReason string
// Message is the generated message.
Message Message
// Timings contains llama.cpp performance data, when reported.
Timings *Timings
// Usage contains token counts, when reported.
Usage *Usage
}
// Completions requests a non-streaming chat completion and returns the first
// choice. It returns an error if the response contains no choices.
func (c *Client) Completions(
ctx context.Context, p CompletionsParams,
) (*CompletionsResult, error) {
) (*CompletionResult, error) {
ctx, cancel := context.WithTimeout(ctx, c.timeout)
defer cancel()
model, systemMessage, maxIterations, messages, thinkingBudgetTokens :=
p.Model, p.SystemMessage, p.MaxIterations, p.Messages, p.ThinkingBudgetTokens
// Use request system message if provided, config default otherwise.
smsg := c.systemMessage
if systemMessage != "" {
smsg = systemMessage
cr, err := c.completions(ctx, p)
if err != nil {
return nil, err
}
// Prepend system message, if configured and not already present.
if smsg != "" &&
(len(messages) == 0 || messages[0].Role != RoleSystem) {
messages = append(
[]Message{{
Role: RoleSystem,
ContentParts: []ContentPart{{
Type: ContentTypeText,
Text: smsg,
}},
}},
messages...,
)
}
// Remember the starting length so we can return only the messages
// generated during this call (the suffix of the conversation history).
start := len(messages)
// Loop for tool calls. maxIterations == 0 means unlimited.
var finalUsage *Usage
var finalTimings *Timings
for i := 0; maxIterations == 0 || i < maxIterations; i++ {
cr, err := c.completions(ctx, model, thinkingBudgetTokens, messages)
if err != nil {
return nil, fmt.Errorf("completions: %w", err)
}
// Capture usage and timings from this iteration.
if cr.Usage != nil {
finalUsage = cr.Usage
}
if cr.Timings != nil {
finalTimings = cr.Timings
}
// Use first choice for tool call loop.
message := &cr.Choices[0].Message
messages = append(messages, *message)
if len(message.ToolCalls) == 0 {
return &CompletionsResult{
Messages: messages[start:],
Usage: finalUsage,
Timings: finalTimings,
}, nil
}
toolResults := c.callTools(ctx, *message)
messages = append(messages, toolResults...)
}
return &CompletionsResult{
Messages: messages[start:],
Usage: finalUsage,
Timings: finalTimings,
}, ErrMaxIterations
choice := cr.Choices[0]
return &CompletionResult{
Message: choice.Message,
FinishReason: choice.FinishReason,
Usage: cr.Usage,
Timings: cr.Timings,
}, nil
}
// completions sends a single non-streaming request and returns the
// full API response.
func (c *Client) completions(
ctx context.Context, model string, thinkingBudgetTokens *int, messages []Message,
ctx context.Context, p CompletionsParams,
) (*ChatResponse, error) {
kwargs := map[string]any{}
if thinkingBudgetTokens != nil {
kwargs["enable_thinking"] = *thinkingBudgetTokens != 0
if p.ThinkingBudgetTokens != nil {
kwargs[chatTemplateEnableThinking] = *p.ThinkingBudgetTokens != 0
}
body, err := json.Marshal(ChatRequest{
Model: model,
Messages: messages,
Tools: c.tools,
ThinkingBudgetTokens: thinkingBudgetTokens,
Model: p.Model,
Messages: p.Messages,
Tools: p.Tools,
ThinkingBudgetTokens: p.ThinkingBudgetTokens,
ChatTemplateKWARGS: kwargs,
ParallelToolCalls: new(true),
})
@@ -133,7 +84,7 @@ func (c *Client) completions(
req, err := http.NewRequestWithContext(
ctx,
http.MethodPost,
c.baseURL+"/chat/completions",
c.baseURL+chatCompletionsPath,
bytes.NewReader(body),
)
if err != nil {
@@ -163,11 +114,7 @@ func (c *Client) completions(
)
}
var apiError struct {
Error *struct {
Type string `json:"type"`
Code string `json:"code"`
Message string `json:"message"`
} `json:"error"`
Error *APIError `json:"error"`
}
if err := json.Unmarshal(
body, &apiError,
@@ -182,11 +129,7 @@ func (c *Client) completions(
),
slog.String("type", apiError.Error.Type),
)
return nil, fmt.Errorf(
"completions: %s: %s",
apiError.Error.Code,
apiError.Error.Message,
)
return nil, fmt.Errorf("completions: %w", apiError.Error)
}
c.log.ErrorContext(
ctx,
@@ -212,131 +155,140 @@ func (c *Client) completions(
return cr, nil
}
// StreamEvent wraps a message or delta produced during streaming.
type StreamEvent struct {
// Type identifies the event kind: delta, reasoning_delta, done, or message.
Type string
// Delta is a text fragment from the assistant response.
Delta string
// ReasoningDelta is a reasoning text fragment from the assistant response.
ReasoningDelta string
// Message is the complete chat completion message.
Message Message
// Usage is token usage info from the final chunk, sent on the done event.
Usage *Usage
// Timings is llama.cpp timing info from the final chunk, sent on the done event.
Timings *Timings
// StreamAccumulator reconstructs a Message from the ordered deltas for one
// streaming choice. The zero value is ready for use.
type StreamAccumulator struct {
// Content contains the accumulated message text.
Content strings.Builder
// Reasoning combines fragments received under either supported reasoning
// field name.
Reasoning strings.Builder
// Refusal contains the accumulated refusal text.
Refusal strings.Builder
// Role is the latest non-empty role received.
Role Role
// ToolCalls contains calls assembled by tool-call index.
ToolCalls []ToolCall
}
const (
streamEventDelta = "delta"
streamEventReasoningDelta = "reasoning_delta"
streamEventDone = "done"
streamEventMessage = "message"
)
// Add incorporates the next delta. Content, refusal, reasoning, and tool-call
// fragments are appended in arrival order. Tool-call fragments are grouped by
// Index; a missing index denotes the first call. A non-empty Role replaces the
// previously accumulated role.
func (acc *StreamAccumulator) Add(delta Delta) {
if delta.Role != "" {
acc.Role = delta.Role
}
acc.Content.WriteString(delta.Content)
acc.Refusal.WriteString(delta.Refusal)
acc.Reasoning.WriteString(delta.ReasoningContent)
acc.Reasoning.WriteString(delta.Reasoning)
// CompletionsStream sends messages to the LLM using the specified model and
// streams results. Each complete message (assistant reply, tool call,
// tool result) is sent to the events channel as it becomes available.
// The channel is closed before returning.
// If systemMessage is non-empty, it overrides the configured system message.
// maxIterations caps the number of agent loop iterations; a value of 0
// means unlimited (no cap).
// Returns ErrMaxIterations if the iteration limit is exceeded.
for _, tc := range delta.ToolCalls {
idx := 0
if tc.Index != nil {
idx = *tc.Index
}
for len(acc.ToolCalls) <= idx {
acc.ToolCalls = append(
acc.ToolCalls,
ToolCall{Type: ToolTypeFunction},
)
}
if tc.ID != "" {
acc.ToolCalls[idx].ID = tc.ID
}
if tc.Type != "" {
acc.ToolCalls[idx].Type = tc.Type
}
acc.ToolCalls[idx].Function.Name += tc.Function.Name
acc.ToolCalls[idx].Function.Arguments += tc.Function.Arguments
}
}
// Message returns the accumulated state without resetting the accumulator. If
// the stream contains a refusal, the returned content consists of that refusal
// rather than the accumulated text.
func (acc *StreamAccumulator) Message() Message {
var contentParts []ContentPart
if refusal := acc.Refusal.String(); refusal != "" {
contentParts = []ContentPart{{
Type: ContentTypeRefusal,
Refusal: refusal,
}}
} else {
contentParts = []ContentPart{{
Type: ContentTypeText,
Text: acc.Content.String(),
}}
}
return Message{
ContentParts: contentParts,
ReasoningContent: acc.Reasoning.String(),
Refusal: acc.Refusal.String(),
Role: acc.Role,
ToolCalls: acc.ToolCalls,
}
}
// CompletionStream is a caller-driven chat-completion stream. A stream owns
// its response body and timeout context until it terminates or is closed.
type CompletionStream struct {
ctx context.Context
reader *bufio.Reader
body io.ReadCloser
cancel context.CancelFunc
acc StreamAccumulator
current StreamChunk
finishReason string
usage *Usage
timings *Timings
err error
cleanupErr error
terminated bool
resourcesDone bool
}
// CompletionsStream opens a streaming chat completion. Request setup,
// transport, and non-success response errors are returned before a stream is
// created.
func (c *Client) CompletionsStream(
ctx context.Context, p CompletionsParams, events chan<- StreamEvent,
) error {
defer close(events)
ctx context.Context,
p CompletionsParams,
) (*CompletionStream, error) {
ctx, cancel := context.WithTimeout(ctx, c.timeout)
defer cancel()
model, systemMessage, maxIterations, messages, thinkingBudgetTokens :=
p.Model, p.SystemMessage, p.MaxIterations, p.Messages, p.ThinkingBudgetTokens
// Use request system message if provided, config default otherwise.
smsg := c.systemMessage
if systemMessage != "" {
smsg = systemMessage
}
// Prepend system message, if configured and not already present.
if smsg != "" &&
(len(messages) == 0 || messages[0].Role != RoleSystem) {
messages = append(
[]Message{{
Role: RoleSystem,
ContentParts: []ContentPart{
{Type: ContentTypeText, Text: smsg},
},
}},
messages...,
)
}
// Loop for tool calls. maxIterations == 0 means unlimited.
for i := 0; maxIterations == 0 || i < maxIterations; i++ {
choices, err := c.completionsStream(ctx, messages, model, thinkingBudgetTokens, events)
if err != nil {
return fmt.Errorf("completions stream: %w", err)
}
// Use first choice for tool call loop.
message := choices[0]
if len(message.ToolCalls) == 0 {
return nil
}
toolResults := c.callTools(ctx, message)
for _, tr := range toolResults {
events <- StreamEvent{
Type: streamEventMessage,
Message: tr,
}
}
messages = append(messages, choices...)
messages = append(messages, toolResults...)
}
return ErrMaxIterations
}
// streamAccum accumulates deltas for a single choice in a streaming response.
type streamAccum struct {
content strings.Builder
reasoning strings.Builder
refusal strings.Builder
role string
toolCalls []ToolCall
}
// completionsStream sends a single streaming request, accumulates the
// response, and returns all assistant messages from the response choices.
// Sends each message to the events channel. Usage and timings are included
// on the done event for the last message.
func (c *Client) completionsStream(ctx context.Context, messages []Message, model string, thinkingBudgetTokens *int, events chan<- StreamEvent) ([]Message, error) {
kwargs := map[string]any{}
if thinkingBudgetTokens != nil {
kwargs["enable_thinking"] = *thinkingBudgetTokens != 0
if p.ThinkingBudgetTokens != nil {
kwargs[chatTemplateEnableThinking] = *p.ThinkingBudgetTokens != 0
}
body, err := json.Marshal(ChatRequest{
Model: model,
Messages: messages,
Stream: true,
Tools: c.tools,
ThinkingBudgetTokens: thinkingBudgetTokens,
ChatTemplateKWARGS: kwargs,
Messages: p.Messages,
Model: p.Model,
ParallelToolCalls: new(true),
Stream: true,
StreamOptions: &StreamOptions{IncludeUsage: true},
ThinkingBudgetTokens: p.ThinkingBudgetTokens,
Tools: p.Tools,
})
if err != nil {
cancel()
return nil, fmt.Errorf(
"completions stream: marshal request body: %w", err,
)
}
req, err := http.NewRequestWithContext(
ctx, http.MethodPost,
c.baseURL+"/chat/completions",
ctx,
http.MethodPost,
c.baseURL+chatCompletionsPath,
bytes.NewReader(body),
)
if err != nil {
cancel()
return nil, fmt.Errorf(
"completions stream: create request: %w", err,
)
@@ -349,12 +301,14 @@ func (c *Client) completionsStream(ctx context.Context, messages []Message, mode
res, err := c.httpc.Do(req)
if err != nil {
cancel()
return nil, fmt.Errorf("completions stream: %w", err)
}
if res.StatusCode != http.StatusOK {
defer cancel()
defer res.Body.Close()
body, err := io.ReadAll(res.Body)
res.Body.Close()
if err != nil {
c.log.ErrorContext(
ctx,
@@ -367,11 +321,7 @@ func (c *Client) completionsStream(ctx context.Context, messages []Message, mode
)
}
var apiError struct {
Error *struct {
Type string `json:"type"`
Code string `json:"code"`
Message string `json:"message"`
} `json:"error"`
Error *APIError `json:"error"`
}
if err := json.Unmarshal(
body, &apiError,
@@ -381,15 +331,11 @@ func (c *Client) completionsStream(ctx context.Context, messages []Message, mode
"completions stream: request error",
slog.Int("status", res.StatusCode),
slog.String("code", apiError.Error.Code),
slog.String(
"message", apiError.Error.Message,
),
slog.String("message", apiError.Error.Message),
slog.String("type", apiError.Error.Type),
)
return nil, fmt.Errorf(
"completions stream: %s: %s",
apiError.Error.Code,
apiError.Error.Message,
"completions stream: %w", apiError.Error,
)
}
c.log.ErrorContext(
@@ -400,161 +346,143 @@ func (c *Client) completionsStream(ctx context.Context, messages []Message, mode
)
return nil, fmt.Errorf(
"completions stream: %d: %s",
res.StatusCode, string(body),
res.StatusCode,
string(body),
)
}
// Accumulate the streamed response per choice index.
accums := make(map[int]*streamAccum)
var usage *Usage
var timings *Timings
return &CompletionStream{
ctx: ctx,
reader: bufio.NewReader(res.Body),
body: res.Body,
cancel: cancel,
}, nil
}
scanner := bufio.NewScanner(res.Body)
for scanner.Scan() {
line := scanner.Text()
// Skip empty lines and comments.
if line == "" || line == ":" {
continue
}
// Parse data lines.
if !strings.HasPrefix(line, "data: ") {
continue
}
data := strings.TrimPrefix(line, "data: ")
// [DONE] signals end of stream.
if data == "[DONE]" {
break
// Next advances to the next JSON data event. It returns false after [DONE], an
// error, or Close. On an error, Err reports the cause and Result remains
// available with all chunks accumulated before the error.
func (s *CompletionStream) Next() bool {
if s.terminated {
return false
}
var data []string
for {
line, err := s.reader.ReadString('\n')
if err != nil && err != io.EOF {
s.fail(fmt.Errorf("completions stream: read response: %w", err))
return false
}
var chunk StreamChunk
if err := json.Unmarshal([]byte(data), &chunk); err != nil {
continue
}
for _, ch := range chunk.Choices {
acc, ok := accums[ch.Index]
if !ok {
acc = &streamAccum{}
accums[ch.Index] = acc
}
delta := ch.Delta
if delta.Role != "" {
acc.role = delta.Role
}
if delta.Content != "" {
acc.content.WriteString(delta.Content)
events <- StreamEvent{
Type: streamEventDelta,
Delta: delta.Content,
}
}
if delta.Refusal != "" {
acc.refusal.WriteString(delta.Refusal)
}
if delta.ReasoningContent != "" {
acc.reasoning.WriteString(
delta.ReasoningContent,
)
events <- StreamEvent{
Type: streamEventReasoningDelta,
ReasoningDelta: delta.ReasoningContent,
}
}
if delta.Reasoning != "" {
acc.reasoning.WriteString(delta.Reasoning)
events <- StreamEvent{
Type: streamEventReasoningDelta,
ReasoningDelta: delta.Reasoning,
}
}
// Accumulate tool call deltas by index.
for _, tc := range delta.ToolCalls {
idx := 0
if tc.Index != nil {
idx = *tc.Index
}
// Grow the slice as needed.
for len(acc.toolCalls) <= idx {
acc.toolCalls = append(
acc.toolCalls,
ToolCall{Type: "function"},
)
if len(line) > 0 {
line = strings.TrimSuffix(line, "\n")
line = strings.TrimSuffix(line, "\r")
if line == "" {
if len(data) > 0 {
break
}
if tc.ID != "" {
acc.toolCalls[idx].ID = tc.ID
} else if line[0] != ':' {
field, value, found := strings.Cut(line, ":")
if !found {
field = line
value = ""
}
if tc.Function.Name != "" {
acc.toolCalls[idx].Function.Name +=
tc.Function.Name
value = strings.TrimPrefix(value, " ")
if field == "data" {
data = append(data, value)
}
if tc.Function.Arguments != "" {
acc.toolCalls[idx].Function.Arguments +=
tc.Function.Arguments
}
}
}
// Capture usage from the final chunk (has null choices).
if chunk.Usage != nil {
usage = chunk.Usage
}
if chunk.Timings != nil {
timings = chunk.Timings
if err == io.EOF {
s.fail(io.ErrUnexpectedEOF)
return false
}
}
res.Body.Close()
if err := scanner.Err(); err != nil {
return nil, fmt.Errorf(
"completions stream: read error: %w", err,
)
payload := strings.Join(data, "\n")
if payload == sseDone {
s.terminated = true
s.release()
return false
}
// Build messages in index order.
indices := make([]int, 0, len(accums))
for idx := range accums {
indices = append(indices, idx)
var chunk StreamChunk
if err := json.Unmarshal([]byte(payload), &chunk); err != nil {
s.fail(fmt.Errorf("completions stream: decode chunk: %w", err))
return false
}
for _, choice := range chunk.Choices {
if choice.Index != 0 {
s.fail(fmt.Errorf(
"completions stream: unexpected choice index %d",
choice.Index,
))
return false
}
if choice.FinishReason != "" {
s.finishReason = choice.FinishReason
}
s.acc.Add(choice.Delta)
}
if chunk.Usage != nil {
s.usage = chunk.Usage
}
slices.Sort(indices)
if chunk.Timings != nil {
s.timings = chunk.Timings
}
s.current = chunk
return true
}
msgs := make([]Message, 0, len(accums))
for i, idx := range indices {
acc := accums[idx]
// Current returns the most recently decoded chunk. The value is replaced by
// the next successful call to Next.
func (s *CompletionStream) Current() StreamChunk {
return s.current
}
// Content and refusal are mutually exclusive in the API.
var contentParts []ContentPart
if refusalStr := acc.refusal.String(); refusalStr != "" {
contentParts = []ContentPart{{
Type: ContentTypeRefusal,
Refusal: refusalStr,
}}
} else {
contentParts = []ContentPart{{
Type: ContentTypeText,
Text: acc.content.String(),
}}
}
// Result returns a snapshot of the state accumulated so far. Err reports
// whether the snapshot is a complete response or a partial result.
func (s *CompletionStream) Result() CompletionResult {
message := s.acc.Message()
message.ToolCalls = append([]ToolCall(nil), message.ToolCalls...)
return CompletionResult{
Message: message,
FinishReason: s.finishReason,
Usage: s.usage,
Timings: s.timings,
}
}
// Err returns the stream failure, if any. Closing a stream does not create or
// replace a stream failure.
func (s *CompletionStream) Err() error {
return s.err
}
message := Message{
Role: acc.role,
ContentParts: contentParts,
Refusal: acc.refusal.String(),
ReasoningContent: acc.reasoning.String(),
ToolCalls: acc.toolCalls,
}
// Close releases the response body and timeout context. It is safe to call
// repeatedly and returns the first response-body cleanup error.
func (s *CompletionStream) Close() error {
s.terminated = true
s.release()
return s.cleanupErr
}
// Attach usage and timings to the done event of the last message.
var evUsage *Usage
var evTimings *Timings
if i == len(indices)-1 {
evUsage = usage
evTimings = timings
}
events <- StreamEvent{
Type: streamEventDone,
Message: message,
Usage: evUsage,
Timings: evTimings,
}
msgs = append(msgs, message)
func (s *CompletionStream) fail(err error) {
if ctxErr := s.ctx.Err(); ctxErr != nil {
err = ctxErr
}
s.err = err
s.terminated = true
s.release()
}
return msgs, nil
func (s *CompletionStream) release() {
if s.resourcesDone {
return
}
s.resourcesDone = true
if err := s.body.Close(); err != nil {
s.cleanupErr = err
}
s.cancel()
}
internal/llm/client_completions_test.go
+222 -167
View File
@@ -2,96 +2,159 @@ package llm
import (
"context"
"errors"
"fmt"
"io"
"log/slog"
"net/http"
"net/http/httptest"
"strings"
"testing"
)
func TestCompletionsStreamStreamsReasoningDeltas(t *testing.T) {
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path != "/chat/completions" {
t.Fatalf("unexpected path: %s", r.URL.Path)
}
func newCompletionTestClient(t *testing.T, handler http.HandlerFunc) *Client {
t.Helper()
srv := httptest.NewServer(handler)
t.Cleanup(srv.Close)
client, err := NewClient(Config{Provider: "test", URL: srv.URL}, slog.Default())
if err != nil {
t.Fatalf("NewClient: %v", err)
}
return client
}
func openCompletionTestStream(t *testing.T, client *Client) *CompletionStream {
t.Helper()
stream, err := client.CompletionsStream(
t.Context(), CompletionsParams{Model: "test-model"},
)
if err != nil {
t.Fatalf("CompletionsStream: %v", err)
}
t.Cleanup(func() { _ = stream.Close() })
return stream
}
func TestCompletionsStreamReconstructsReasoning(t *testing.T) {
client := newCompletionTestClient(t, func(w http.ResponseWriter, _ *http.Request) {
w.Header().Set("Content-Type", "text/event-stream")
for _, data := range []string{
`{"choices":[{"index":0,"delta":{"role":"assistant"}}]}`,
`{"choices":[{"index":0,"delta":{"reasoning_content":"thinking"}}]}`,
`{"choices":[{"index":0,"delta":{"reasoning_content":" more"}}]}`,
`{"choices":[{"index":0,"delta":{"content":"answer"}}]}`,
`{"choices":[{"index":0,"delta":{"content":"answer"},"finish_reason":"stop"}]}`,
} {
fmt.Fprintf(w, "data: %s\n\n", data)
}
fmt.Fprint(w, "data: [DONE]\n\n")
}))
defer srv.Close()
})
client, err := NewClient(Config{URL: srv.URL}, nil, nil)
stream := openCompletionTestStream(t, client)
if got := stream.Result(); got.Message.Text() != "" {
t.Fatalf("initial result text = %q, want empty", got.Message.Text())
}
var chunks []StreamChunk
for stream.Next() {
chunks = append(chunks, stream.Current())
}
if err := stream.Err(); err != nil {
t.Fatalf("stream.Err: %v", err)
}
if len(chunks) != 4 {
t.Fatalf("chunks = %d, want 4", len(chunks))
}
result := stream.Result()
if result.Message.ReasoningContent != "thinking more" {
t.Errorf("reasoning = %q", result.Message.ReasoningContent)
}
if result.Message.Text() != "answer" {
t.Errorf("text = %q", result.Message.Text())
}
if result.FinishReason != "stop" {
t.Errorf("finish reason = %q", result.FinishReason)
}
if stream.Next() {
t.Error("Next returned true after termination")
}
}
func TestCompletionsStreamReportsCancellationWithPartialResult(t *testing.T) {
firstSent := make(chan struct{})
client := newCompletionTestClient(t, func(w http.ResponseWriter, r *http.Request) {
flusher := w.(http.Flusher)
fmt.Fprint(w, "data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\"first\"}}]}\n\n")
flusher.Flush()
close(firstSent)
<-r.Context().Done()
})
ctx, cancel := context.WithCancel(t.Context())
stream, err := client.CompletionsStream(ctx, CompletionsParams{Model: "test-model"})
if err != nil {
t.Fatalf("NewClient: %v", err)
}
events := make(chan StreamEvent)
errCh := make(chan error, 1)
go func() {
errCh <- client.CompletionsStream(
context.Background(), CompletionsParams{
Model: "test-model",
SystemMessage: "",
MaxIterations: 0,
Messages: []Message{{
Role: RoleUser,
ContentParts: []ContentPart{{
Type: ContentTypeText,
Text: "hello",
}},
}},
}, events,
)
}()
var got []StreamEvent
for evt := range events {
got = append(got, evt)
}
if err := <-errCh; err != nil {
t.Fatalf("CompletionsStream: %v", err)
}
defer stream.Close()
<-firstSent
if !stream.Next() {
t.Fatalf("first Next = false: %v", stream.Err())
}
cancel()
if stream.Next() {
t.Fatal("Next returned true after cancellation")
}
if !errors.Is(stream.Err(), context.Canceled) {
t.Fatalf("Err = %v, want context.Canceled", stream.Err())
}
result := stream.Result()
if got := result.Message.Text(); got != "first" {
t.Errorf("partial text = %q, want first", got)
}
}
wantTypes := []string{
streamEventReasoningDelta,
streamEventReasoningDelta,
streamEventDelta,
streamEventDone,
}
if len(got) != len(wantTypes) {
t.Fatalf("got %d events, want %d: %#v", len(got), len(wantTypes), got)
}
for i, want := range wantTypes {
if got[i].Type != want {
t.Fatalf("event %d type = %q, want %q", i, got[i].Type, want)
func TestCompletionsStreamReconstructsFragmentedParallelToolCalls(t *testing.T) {
client := newCompletionTestClient(t, func(w http.ResponseWriter, _ *http.Request) {
for _, data := range []string{
`{"choices":[{"index":0,"delta":{"role":"assistant"}}]}`,
`{"choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"id":"call_0","function":{"name":"ec"}},{"index":1,"id":"call_1","function":{"name":"pr"}}]}}]}`,
`{"choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"function":{"name":"ho","arguments":"{\"value\":"}},{"index":1,"function":{"name":"intf","arguments":"{\"format\":"}}]}}]}`,
`{"choices":[{"index":0,"delta":{"tool_calls":[{"index":1,"function":{"arguments":"\"%s\"}"}},{"index":0,"function":{"arguments":"\"one\"}"}}]},"finish_reason":"tool_calls"}]}`,
} {
fmt.Fprintf(w, "data: %s\n\n", data)
}
fmt.Fprint(w, "data: [DONE]\n\n")
})
stream := openCompletionTestStream(t, client)
var chunks []StreamChunk
for stream.Next() {
chunks = append(chunks, stream.Current())
}
if got[0].ReasoningDelta != "thinking" {
t.Errorf("first reasoning delta = %q", got[0].ReasoningDelta)
if err := stream.Err(); err != nil {
t.Fatalf("stream.Err: %v", err)
}
if got[1].ReasoningDelta != " more" {
t.Errorf("second reasoning delta = %q", got[1].ReasoningDelta)
if len(chunks) != 4 || len(chunks[1].Choices[0].Delta.ToolCalls) != 2 {
t.Fatalf("tool-call chunks not exposed unchanged: %#v", chunks)
}
if got[2].Delta != "answer" {
t.Errorf("content delta = %q", got[2].Delta)
result := stream.Result()
toolCalls := result.Message.ToolCalls
if len(toolCalls) != 2 {
t.Fatalf("tool calls = %d, want 2", len(toolCalls))
}
final := got[3].Message
if final.ReasoningContent != "thinking more" {
t.Errorf("final reasoning = %q", final.ReasoningContent)
if toolCalls[0].ID != "call_0" || toolCalls[0].Function.Name != "echo" ||
toolCalls[0].Function.Arguments != `{"value":"one"}` {
t.Errorf("tool call 0 = %#v", toolCalls[0])
}
if final.Text() != "answer" {
t.Errorf("final text = %q", final.Text())
if toolCalls[1].ID != "call_1" || toolCalls[1].Function.Name != "printf" ||
toolCalls[1].Function.Arguments != `{"format":"%s"}` {
t.Errorf("tool call 1 = %#v", toolCalls[1])
}
if result.FinishReason != "tool_calls" {
t.Errorf("finish reason = %q", result.FinishReason)
}
}
func TestCompletionsReturnsUsageAndTimings(t *testing.T) {
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
client := newCompletionTestClient(t, func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path != "/chat/completions" {
t.Fatalf("unexpected path: %s", r.URL.Path)
}
@@ -122,135 +185,127 @@ func TestCompletionsReturnsUsageAndTimings(t *testing.T) {
"predicted_per_second": 50.0
}
}`)
}))
defer srv.Close()
})
client, err := NewClient(Config{URL: srv.URL}, nil, nil)
if err != nil {
t.Fatalf("NewClient: %v", err)
}
res, err := client.Completions(
context.Background(), CompletionsParams{
Model: "test-model",
SystemMessage: "",
MaxIterations: 0,
Messages: []Message{{
Role: RoleUser,
ContentParts: []ContentPart{{
Type: ContentTypeText,
Text: "hi",
}},
}},
},
)
res, err := client.Completions(context.Background(), CompletionsParams{Model: "test-model"})
if err != nil {
t.Fatalf("Completions: %v", err)
}
if len(res.Messages) != 1 {
t.Fatalf("got %d messages, want 1", len(res.Messages))
if res.Message.Text() != "hello" || res.Usage == nil || res.Timings == nil {
t.Fatalf("unexpected result: %#v", res)
}
if res.Usage == nil {
t.Fatal("usage is nil")
}
if res.Usage.PromptTokens != 10 {
t.Errorf("prompt_tokens = %d, want 10", res.Usage.PromptTokens)
}
if res.Usage.CompletionTokens != 5 {
t.Errorf("completion_tokens = %d, want 5", res.Usage.CompletionTokens)
}
if res.Usage.TotalTokens != 15 {
t.Errorf("total_tokens = %d, want 15", res.Usage.TotalTokens)
if res.Usage.PromptTokens != 10 || res.Usage.CompletionTokens != 5 || res.Usage.TotalTokens != 15 {
t.Errorf("usage = %#v", res.Usage)
}
if res.Usage.PromptTokensDetails == nil || res.Usage.PromptTokensDetails.CachedTokens != 3 {
t.Errorf("cached_tokens unexpected: %#v", res.Usage.PromptTokensDetails)
t.Errorf("prompt token details = %#v", res.Usage.PromptTokensDetails)
}
if res.Usage.CompletionTokensDetails == nil || res.Usage.CompletionTokensDetails.ReasoningTokens != 2 {
t.Errorf("reasoning_tokens unexpected: %#v", res.Usage.CompletionTokensDetails)
t.Errorf("completion token details = %#v", res.Usage.CompletionTokensDetails)
}
if res.Timings == nil {
t.Fatal("timings is nil")
}
if res.Timings.PromptMS != 50.5 {
t.Errorf("prompt_ms = %f, want 50.5", res.Timings.PromptMS)
}
if res.Timings.PredictedMS != 100.0 {
t.Errorf("predicted_ms = %f, want 100.0", res.Timings.PredictedMS)
if res.Timings.PromptMS != 50.5 || res.Timings.PredictedMS != 100 {
t.Errorf("timings = %#v", res.Timings)
}
}
func TestCompletionsStreamEmitsUsageEvent(t *testing.T) {
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path != "/chat/completions" {
t.Fatalf("unexpected path: %s", r.URL.Path)
}
w.Header().Set("Content-Type", "text/event-stream")
// Stream some content
fmt.Fprintf(w, "data: %s\n\n",
`{"choices":[{"index":0,"delta":{"role":"assistant"}}]}`)
fmt.Fprintf(w, "data: %s\n\n",
`{"choices":[{"index":0,"delta":{"content":"hello"}}]}`)
// Final chunk with usage
fmt.Fprintf(w, "data: %s\n\n",
`{"choices":[{"index":0,"delta":{},"finish_reason":"stop"}],"usage":{"prompt_tokens":10,"completion_tokens":5,"total_tokens":15}}`)
fmt.Fprint(w, "data: [DONE]\n\n")
}))
defer srv.Close()
func TestCompletionsStreamSSEAndMetadata(t *testing.T) {
large := strings.Repeat("x", 70*1024)
client := newCompletionTestClient(t, func(w http.ResponseWriter, _ *http.Request) {
fmt.Fprint(w, ": comment\n")
fmt.Fprintf(w, "data:{\"choices\":[{\"index\":0,\"delta\":{\"content\":%q}}],\n", large)
fmt.Fprint(w, "data: \"usage\":{\"prompt_tokens\":10,\"total_tokens\":15}}\n\n")
fmt.Fprint(w, "data: {\"choices\":[],\"timings\":{\"prompt_ms\":50.5}}\n\n")
fmt.Fprint(w, "data:[DONE]\n\n")
})
client, err := NewClient(Config{URL: srv.URL}, nil, nil)
if err != nil {
t.Fatalf("NewClient: %v", err)
stream := openCompletionTestStream(t, client)
var count int
for stream.Next() {
count++
}
events := make(chan StreamEvent)
errCh := make(chan error, 1)
go func() {
errCh <- client.CompletionsStream(
context.Background(), CompletionsParams{
Model: "test-model",
SystemMessage: "",
MaxIterations: 0,
Messages: []Message{{
Role: RoleUser,
ContentParts: []ContentPart{{
Type: ContentTypeText,
Text: "hi",
}},
}},
}, events,
)
}()
var got []StreamEvent
for evt := range events {
got = append(got, evt)
if err := stream.Err(); err != nil {
t.Fatalf("stream.Err: %v", err)
}
if err := <-errCh; err != nil {
t.Fatalf("CompletionsStream: %v", err)
result := stream.Result()
if count != 2 || result.Message.Text() != large {
t.Fatalf("count = %d, text length = %d", count, len(result.Message.Text()))
}
if result.Usage == nil || result.Usage.PromptTokens != 10 || result.Timings == nil || result.Timings.PromptMS != 50.5 {
t.Errorf("metadata = usage %#v, timings %#v", result.Usage, result.Timings)
}
}
// Expect: delta, done events (usage attached to done event)
wantTypes := []string{
streamEventDelta,
streamEventDone,
func TestCompletionsStreamFailuresPreservePartialResult(t *testing.T) {
tests := []struct {
name string
response string
wantErr error
}{
{
name: "EOF before done",
response: "data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\"partial\"}}]}\n\n",
wantErr: io.ErrUnexpectedEOF,
},
{
name: "malformed JSON",
response: "data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\"partial\"}}]}\n\ndata: {bad}\n\n",
},
{
name: "unexpected choice",
response: "data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\"partial\"}}]}\n\ndata: {\"choices\":[{\"index\":1,\"delta\":{}}]}\n\n",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
client := newCompletionTestClient(t, func(w http.ResponseWriter, _ *http.Request) {
fmt.Fprint(w, tt.response)
})
stream := openCompletionTestStream(t, client)
for stream.Next() {
}
if stream.Err() == nil {
t.Fatal("Err = nil")
}
if tt.wantErr != nil && !errors.Is(stream.Err(), tt.wantErr) {
t.Errorf("Err = %v, want %v", stream.Err(), tt.wantErr)
}
result := stream.Result()
if got := result.Message.Text(); got != "partial" {
t.Errorf("partial text = %q", got)
}
})
}
if len(got) != len(wantTypes) {
t.Fatalf("got %d events, want %d: %#v", len(got), len(wantTypes), got)
}
func TestCompletionStreamCloseIsIdempotent(t *testing.T) {
client := newCompletionTestClient(t, func(w http.ResponseWriter, r *http.Request) {
w.(http.Flusher).Flush()
<-r.Context().Done()
})
stream := openCompletionTestStream(t, client)
if err := stream.Close(); err != nil {
t.Fatalf("first Close: %v", err)
}
for i, want := range wantTypes {
if got[i].Type != want {
t.Fatalf("event %d type = %q, want %q", i, got[i].Type, want)
}
if err := stream.Close(); err != nil {
t.Fatalf("second Close: %v", err)
}
// Verify usage is on the done event
doneEvt := got[1]
if doneEvt.Usage == nil {
t.Fatal("done event has nil Usage")
if stream.Err() != nil {
t.Errorf("Err after Close = %v", stream.Err())
}
if doneEvt.Usage.PromptTokens != 10 {
t.Errorf("prompt_tokens = %d, want 10", doneEvt.Usage.PromptTokens)
if stream.Next() {
t.Error("Next after Close returned true")
}
}
func TestCompletionsStreamRequiresCompleteDoneEvent(t *testing.T) {
client := newCompletionTestClient(t, func(w http.ResponseWriter, _ *http.Request) {
fmt.Fprint(w, "data: [DONE]")
})
stream := openCompletionTestStream(t, client)
if stream.Next() {
t.Fatal("Next returned true")
}
if doneEvt.Usage.TotalTokens != 15 {
t.Errorf("total_tokens = %d, want 15", doneEvt.Usage.TotalTokens)
if !errors.Is(stream.Err(), io.ErrUnexpectedEOF) {
t.Fatalf("Err = %v, want io.ErrUnexpectedEOF", stream.Err())
}
}
internal/llm/client_models.go
+2 -10
View File
@@ -44,11 +44,7 @@ func (c *Client) ListModels(ctx context.Context) ([]Model, error) {
)
}
var apiError struct {
Error *struct {
Type string `json:"type"`
Code string `json:"code"`
Message string `json:"message"`
} `json:"error"`
Error *APIError `json:"error"`
}
if err := json.Unmarshal(
body, &apiError,
@@ -63,11 +59,7 @@ func (c *Client) ListModels(ctx context.Context) ([]Model, error) {
),
slog.String("type", apiError.Error.Type),
)
return nil, fmt.Errorf(
"list models: %s: %s",
apiError.Error.Code,
apiError.Error.Message,
)
return nil, fmt.Errorf("list models: %w", apiError.Error)
}
c.log.ErrorContext(
ctx,
internal/llm/client_test.go
+41 -22
View File
@@ -1,6 +1,7 @@
package llm
import (
"log/slog"
"testing"
)
@@ -54,10 +55,9 @@ func TestConfigValidate(t *testing.T) {
{
name: "valid full config",
cfg: Config{
Key: "sk-test-key",
SystemMessage: "You are a helpful assistant.",
Timeout: "30m",
URL: "http://localhost:8080",
Key: "sk-test-key",
Timeout: "30m",
URL: "http://localhost:8080",
},
wantErr: false,
},
@@ -77,22 +77,31 @@ func TestConfigValidate(t *testing.T) {
}
func TestNewClient(t *testing.T) {
log := slog.Default()
tests := []struct {
name string
cfg Config
wantErr bool
log *slog.Logger
wantErr string
}{
{
name: "invalid config",
name: "invalid config precedes nil logger",
cfg: Config{},
wantErr: true,
wantErr: "invalid config: missing URL",
},
{
name: "nil logger",
cfg: Config{
URL: "http://localhost:8080",
},
wantErr: "missing logger",
},
{
name: "valid config without timeout",
cfg: Config{
URL: "http://localhost:8080",
},
wantErr: false,
log: log,
},
{
name: "valid config with timeout",
@@ -100,32 +109,42 @@ func TestNewClient(t *testing.T) {
URL: "http://localhost:8080",
Timeout: "10m",
},
wantErr: false,
log: log,
},
{
name: "valid config with all fields",
cfg: Config{
Key: "test-key",
SystemMessage: "Test message",
Timeout: "5m",
URL: "http://localhost:8080",
Key: "test-key",
Timeout: "5m",
URL: "http://localhost:8080",
},
wantErr: false,
log: log,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
client, err := NewClient(tt.cfg, nil, nil)
if (err != nil) != tt.wantErr {
t.Errorf(
"NewClient() error = %v, wantErr %v",
err, tt.wantErr,
)
client, err := NewClient(tt.cfg, tt.log)
if tt.wantErr != "" {
if err == nil {
t.Fatalf("NewClient() error = nil, want %q", tt.wantErr)
}
if err.Error() != tt.wantErr {
t.Errorf(
"NewClient() error = %q, want %q",
err, tt.wantErr,
)
}
if client != nil {
t.Errorf("NewClient() client = %#v, want nil", client)
}
return
}
if !tt.wantErr && client == nil {
t.Error("NewClient() returned nil client")
if err != nil {
t.Fatalf("NewClient() error = %v, want nil", err)
}
if client == nil {
t.Error("NewClient() client = nil, want non-nil")
}
})
}
internal/llm/client_tools.go
-77
View File
@@ -1,77 +0,0 @@
package llm
import (
"context"
"runtime"
"sync"
"log/slog"
)
// callTools executes each tool call in the message concurrently and returns
// the resulting tool result messages in the same order as the input calls.
func (c *Client) callTools(ctx context.Context, msg Message) []Message {
var (
results = make([]Message, len(msg.ToolCalls))
sem = make(chan struct{}, runtime.NumCPU()*2)
wg sync.WaitGroup
)
for i, tc := range msg.ToolCalls {
wg.Add(1)
go func(idx int, tc ToolCall) {
defer wg.Done()
sem <- struct{}{}
defer func() { <-sem }()
c.log.InfoContext(
ctx,
"calling tool",
slog.String("name", tc.Function.Name),
slog.String("args", tc.Function.Arguments),
)
result, err := c.registry.Execute(
ctx, tc.Function.Name, tc.Function.Arguments,
)
if err != nil {
c.log.ErrorContext(
ctx,
"failed to call tool",
slog.Any("error", err),
slog.String("name", tc.Function.Name),
)
results[idx] = Message{
Role: RoleTool,
ContentParts: []ContentPart{{
Type: ContentTypeText,
Text: err.Error(),
}},
Name: tc.Function.Name,
ToolCallID: tc.ID,
}
return
}
c.log.InfoContext(
ctx,
"called tool",
slog.String("name", tc.Function.Name),
)
formatted := result.FormatResult()
results[idx] = Message{
Role: RoleTool,
ContentParts: []ContentPart{{
Type: ContentTypeText,
Text: formatted,
}},
Name: tc.Function.Name,
ToolCallID: tc.ID,
}
}(i, tc)
}
wg.Wait()
return results
}
internal/llm/llm.go
+373 -164
View File
@@ -1,16 +1,5 @@
// Package llm provides an OpenAI-compatible client for LLM interactions.
// It handles chat completions with automatic tool call execution, where the
// client iteratively sends messages to the LLM, executes any requested tool
// calls, and continues until the LLM produces a final text response.
//
// The package defines the core types for constructing requests and parsing
// responses: Message, ChatRequest, ChatResponse, and streaming types
// such as StreamChunk and Delta.
//
// Tools are defined declaratively in YAML configuration, stored as Tool
// structs in a Registry, and executed as subprocesses when the LLM invokes
// them. The Registry also converts tool definitions into the API format for
// inclusion in ChatRequest bodies.
// Package llm implements an OpenAI-compatible inference client and its
// protocol types.
package llm
@@ -19,195 +8,350 @@ import (
"strings"
)
// Role constants for messages.
const (
RoleSystem = "system"
RoleUser = "user"
RoleAssistant = "assistant"
RoleTool = "tool"
RoleDeveloper = "developer"
)
// APIError contains the structured error details returned by an
// OpenAI-compatible API. Error values are defined by the server. Compatible
// servers need not use the same set of types or codes. Callers should use
// Message for diagnostics and Type and Code for programmatic error handling.
//
// See [OpenAI API error codes].
//
// [OpenAI API error codes]: https://developers.openai.com/api/docs/guides/error-codes
type APIError struct {
// Code identifies the precise error condition reported by the server.
// Callers may use it to distinguish errors that have the same Type.
Code string `json:"code"`
// APITool defines a tool available to the LLM in the API request format.
// Message explains what went wrong. Its wording is server-defined and
// should not be parsed by callers.
Message string `json:"message"`
// Type identifies the general kind of failure reported by the server.
// Callers may use it to handle related error conditions as a group.
Type string `json:"type"`
}
// Error returns the server-defined code and diagnostic message.
func (e *APIError) Error() string {
return e.Code + ": " + e.Message
}
// APITool describes a tool that a model may call. Function must be set when
// Type is ToolTypeFunction.
type APITool struct {
Type string `json:"type"` // "function"
// Function defines the callable function when Type is ToolTypeFunction.
Function *FunctionDef `json:"function,omitempty"`
// Type selects the representation of the tool.
Type ToolType `json:"type"`
}
// ChatRequest is the request body for chat completions.
// ChatRequest is the request body of POST /chat/completions.
//
// See [OpenAI Chat Completions API].
//
// [OpenAI Chat Completions API]: https://platform.openai.com/docs/api-reference/chat/create
type ChatRequest struct {
Model string `json:"model"`
Messages []Message `json:"messages"`
Temperature *float32 `json:"temperature,omitempty"`
TopP *float32 `json:"top_p,omitempty"`
N int `json:"n,omitempty"`
Stream bool `json:"stream,omitempty"`
Stop []string `json:"stop,omitempty"`
MaxTokens int `json:"max_tokens,omitempty"`
MaxCompletionTokens int `json:"max_completion_tokens,omitempty"`
ResponseFormat *ResponseFormat `json:"response_format,omitempty"`
Seed *int `json:"seed,omitempty"`
PresencePenalty *float32 `json:"presence_penalty,omitempty"`
FrequencyPenalty *float32 `json:"frequency_penalty,omitempty"`
LogProbs bool `json:"logprobs,omitempty"`
TopLogProbs int `json:"top_logprobs,omitempty"`
Tools []APITool `json:"tools,omitempty"`
ToolChoice any `json:"tool_choice,omitempty"`
StreamOptions *StreamOptions `json:"stream_options,omitempty"`
ParallelToolCalls *bool `json:"parallel_tool_calls,omitempty"`
ChatTemplateKWARGS map[string]any `json:"chat_template_kwargs,omitempty"`
ThinkingBudgetTokens *int `json:"thinking_budget_tokens,omitempty"`
ServiceTier string `json:"service_tier,omitempty"`
// ChatTemplateKWARGS contains arguments made available to the
// endpoint's chat template.
// It is a llama.cpp extension to the OpenAI API.
ChatTemplateKWARGS map[string]any `json:"chat_template_kwargs,omitempty"`
// FrequencyPenalty adjusts token likelihood according to frequency in
// the generated text. Nil leaves the choice to the endpoint.
FrequencyPenalty *float32 `json:"frequency_penalty,omitempty"`
// LogProbs requests log probabilities for generated tokens.
LogProbs bool `json:"logprobs,omitempty"`
// MaxCompletionTokens bounds generated tokens, including reasoning
// tokens.
MaxCompletionTokens int `json:"max_completion_tokens,omitempty"`
// MaxTokens is the deprecated predecessor of MaxCompletionTokens.
MaxTokens int `json:"max_tokens,omitempty"`
// Messages is the conversation presented to the model, in order.
Messages []Message `json:"messages"`
// Model identifies the model to use according to the endpoint.
Model string `json:"model"`
// N is the number of alternatives to generate.
N int `json:"n,omitempty"`
// ParallelToolCalls controls whether the model may request multiple
// function calls in one turn. Nil leaves the choice to the endpoint.
ParallelToolCalls *bool `json:"parallel_tool_calls,omitempty"`
// PresencePenalty adjusts token likelihood according to whether a
// token has already appeared. Nil leaves the choice to the endpoint.
PresencePenalty *float32 `json:"presence_penalty,omitempty"`
// ResponseFormat constrains the representation of generated content.
ResponseFormat *ResponseFormat `json:"response_format,omitempty"`
// Seed requests deterministic sampling on endpoints that support it.
// Equal seeds and parameters do not guarantee equal results.
Seed *int `json:"seed,omitempty"`
// ServiceTier selects the endpoint's request-processing tier.
ServiceTier string `json:"service_tier,omitempty"`
// Stop contains sequences that end generation.
// The returned text excludes the matched sequence.
Stop []string `json:"stop,omitempty"`
// Stream requests a sequence of StreamChunk values instead of
// ChatResponse.
Stream bool `json:"stream,omitempty"`
// StreamOptions configures a streaming response.
// It is meaningful only when Stream is true.
StreamOptions *StreamOptions `json:"stream_options,omitempty"`
// Temperature controls sampling randomness.
// Nil leaves the choice to the endpoint.
Temperature *float32 `json:"temperature,omitempty"`
// ThinkingBudgetTokens is a llama.cpp extension controlling the token
// budget for reasoning. Nil omits it, a negative value requests an
// unrestricted budget, zero requests an immediate end, and a positive
// value sets the budget.
ThinkingBudgetTokens *int `json:"thinking_budget_tokens,omitempty"`
// ToolChoice controls whether the model may or must call a tool.
// It accepts a mode string or an endpoint-defined object selecting a
// particular tool.
ToolChoice any `json:"tool_choice,omitempty"`
// Tools describes the tools available to the model.
Tools []APITool `json:"tools,omitempty"`
// TopLogProbs is the number of most likely tokens reported at each
// position when LogProbs is true.
TopLogProbs int `json:"top_logprobs,omitempty"`
// TopP controls nucleus sampling.
// Nil leaves the choice to the endpoint.
TopP *float32 `json:"top_p,omitempty"`
}
// Usage represents token usage from the API response.
// Usage reports input, output, and total token counts for a completion.
type Usage struct {
PromptTokens int `json:"prompt_tokens"`
CompletionTokens int `json:"completion_tokens"`
TotalTokens int `json:"total_tokens"`
PromptTokensDetails *TokenDetails `json:"prompt_tokens_details,omitempty"`
// CompletionTokens is the number of tokens generated by the model.
CompletionTokens int `json:"completion_tokens"`
// CompletionTokensDetails breaks down CompletionTokens by category.
CompletionTokensDetails *TokenDetails `json:"completion_tokens_details,omitempty"`
// PromptTokens is the number of tokens in the input.
PromptTokens int `json:"prompt_tokens"`
// PromptTokensDetails breaks down PromptTokens by category.
PromptTokensDetails *TokenDetails `json:"prompt_tokens_details,omitempty"`
// TotalTokens is the sum of PromptTokens and CompletionTokens.
TotalTokens int `json:"total_tokens"`
}
// TokenDetails breaks down token usage by category.
// TokenDetails reports token categories within a Usage count. A zero value may
// mean either that no tokens belong to the category or that it was not
// reported.
type TokenDetails struct {
// CachedTokens is the number of input tokens served from a cache.
CachedTokens int `json:"cached_tokens,omitempty"`
// ReasoningTokens is the number of output tokens used for reasoning.
ReasoningTokens int `json:"reasoning_tokens,omitempty"`
CachedTokens int `json:"cached_tokens,omitempty"`
}
// Timings is a llama.cpp-specific extension to the chat completion
// response. It is not part of the OpenAI API spec.
// Timings reports llama.cpp prompt-processing and token-generation performance.
// The N fields are token counts, MS fields are elapsed milliseconds,
// PerTokenMS fields are milliseconds per token, and PerSecond fields are tokens
// per second. Timings is a llama.cpp extension to the OpenAI API.
//
// See [llama.cpp API documentation].
//
// [llama.cpp API documentation]: https://llama.app/docs/api
type Timings struct {
PromptN int `json:"prompt_n"`
PromptMS float64 `json:"prompt_ms"`
PromptPerTokenMS float64 `json:"prompt_per_token_ms"`
PromptPerSecond float64 `json:"prompt_per_second"`
PredictedN int `json:"predicted_n"`
PredictedMS float64 `json:"predicted_ms"`
// PredictedMS is the time spent generating output, in milliseconds.
PredictedMS float64 `json:"predicted_ms"`
// PredictedN is the number of generated tokens.
PredictedN int `json:"predicted_n"`
// PredictedPerSecond is the output-generation rate in tokens per
// second.
PredictedPerSecond float64 `json:"predicted_per_second"`
// PredictedPerTokenMS is the mean output-generation time per token, in
// milliseconds.
PredictedPerTokenMS float64 `json:"predicted_per_token_ms"`
PredictedPerSecond float64 `json:"predicted_per_second"`
// PromptMS is the time spent processing the prompt, in milliseconds.
PromptMS float64 `json:"prompt_ms"`
// PromptN is the number of prompt tokens processed.
PromptN int `json:"prompt_n"`
// PromptPerSecond is the prompt-processing rate in tokens per second.
PromptPerSecond float64 `json:"prompt_per_second"`
// PromptPerTokenMS is the mean prompt-processing time per token, in
// milliseconds.
PromptPerTokenMS float64 `json:"prompt_per_token_ms"`
}
// ChatResponse is the non-streaming response from the API.
// ChatResponse is the body of a non-streaming chat completion
// response.
type ChatResponse struct {
ID string `json:"id"`
Model string `json:"model"`
// Choices contains the generated alternatives.
Choices []Choice `json:"choices"`
Usage *Usage `json:"usage,omitempty"`
// ID identifies the completion.
ID string `json:"id"`
// Model identifies the model that generated the response.
Model string `json:"model"`
// Timings contains llama.cpp performance data, when reported.
Timings *Timings `json:"timings,omitempty"`
// Usage contains token counts, when reported.
Usage *Usage `json:"usage,omitempty"`
}
// Choice is a single choice in a ChatResponse.
// Choice is one generated alternative in a non-streaming response. Index is its
// position among the alternatives, and FinishReason reports why generation
// ended.
type Choice struct {
Index int `json:"index"`
Message Message `json:"message"`
FinishReason string `json:"finish_reason"`
// FinishReason is the endpoint's reason for ending generation.
FinishReason string `json:"finish_reason"`
// Index is the position of this choice in the response.
Index int `json:"index"`
// Message is the generated message.
Message Message `json:"message"`
}
// ContentType identifies the kind of data in a ContentPart.
// ContentType identifies the representation of a ContentPart. Which content
// types are valid depends on the message role and the endpoint.
//
// User messages (ChatCompletionContentPart) support: text, image_url,
// input_audio, file.
// See [OpenAI Chat Completions API].
//
// Assistant messages (ChatCompletionAssistantMessageParam) support:
// text, refusal.
//
// See: https://platform.openai.com/docs/api-reference/chat/create
// [OpenAI Chat Completions API]: https://platform.openai.com/docs/api-reference/chat/create
type ContentType string
const (
ContentTypeText ContentType = "text"
ContentTypeFile ContentType = "file"
ContentTypeImageURL ContentType = "image_url"
ContentTypeInputAudio ContentType = "input_audio"
ContentTypeFile ContentType = "file"
ContentTypeRefusal ContentType = "refusal"
ContentTypeText ContentType = "text"
)
// ContentPart represents a single part of a multi-content message.
// ContentPart is one element of a message's content array. Type determines
// which payload field is meaningful: File, ImageURL, InputAudio, Refusal, or
// Text. The package does not validate that correspondence.
type ContentPart struct {
Type ContentType `json:"type"`
Text string `json:"text"`
ImageURL *ImageURL `json:"image_url,omitempty"`
// File contains the payload for ContentTypeFile.
File *FileInput `json:"file,omitempty"`
// ImageURL contains the payload for ContentTypeImageURL.
ImageURL *ImageURL `json:"image_url,omitempty"`
// InputAudio contains the payload for ContentTypeInputAudio.
InputAudio *InputAudio `json:"input_audio,omitempty"`
File *FileInput `json:"file,omitempty"`
Refusal string `json:"refusal"`
// Refusal contains the payload for ContentTypeRefusal.
Refusal string `json:"refusal,omitempty"`
// Text contains the payload for ContentTypeText.
Text string `json:"text,omitempty"`
// Type identifies which payload field is meaningful.
Type ContentType `json:"type"`
}
// Delta is a partial message update in a streaming response.
// Delta is an incremental message update in a StreamChunk. Reasoning and
// ReasoningContent are alternative field names used by compatible endpoints;
// StreamAccumulator preserves fragments from both.
type Delta struct {
Role string `json:"role,omitempty"`
Content string `json:"content,omitempty"`
Refusal string `json:"refusal,omitempty"`
ReasoningContent string `json:"reasoning_content,omitempty"`
Reasoning string `json:"reasoning,omitempty"`
ToolCalls []ToolCall `json:"tool_calls,omitempty"`
// Content is the next fragment of message text.
Content string `json:"content,omitempty"`
// Reasoning is the next reasoning fragment under the reasoning field
// name.
Reasoning string `json:"reasoning,omitempty"`
// ReasoningContent is the next reasoning fragment under the
// reasoning_content field name.
ReasoningContent string `json:"reasoning_content,omitempty"`
// Refusal is the next fragment of refusal text.
Refusal string `json:"refusal,omitempty"`
// Role identifies the message role when supplied by the endpoint.
Role Role `json:"role,omitempty"`
// ToolCalls contains the next fragments of one or more tool calls.
ToolCalls []ToolCall `json:"tool_calls,omitempty"`
}
// FileInput holds file data for file content parts.
// FileInput identifies file content either inline through FileData or by a
// server-assigned FileID.
type FileInput struct {
// FileData is the base64-encoded contents of an inline file.
FileData string `json:"file_data,omitempty"`
FileID string `json:"file_id,omitempty"`
// FileID references a file previously uploaded to the server.
FileID string `json:"file_id,omitempty"`
// Filename names an inline file supplied through FileData.
Filename string `json:"filename,omitempty"`
}
// Function holds the name and arguments of a tool call.
// Function identifies a called function and carries its arguments as a JSON
// string. The model may emit invalid JSON, so callers must validate Arguments
// before invoking the function.
type Function struct {
Name string `json:"name,omitempty"`
// Arguments is the JSON-encoded argument object generated by the model.
Arguments string `json:"arguments,omitempty"`
// Name identifies the function to call.
Name string `json:"name,omitempty"`
}
// FunctionDef defines a function that can be called by the LLM.
// FunctionDef declares a function the model may call.
type FunctionDef struct {
Name string `json:"name"`
// Description states what the function does and when to call it.
Description string `json:"description,omitempty"`
Parameters any `json:"parameters"` // JSON schema
Strict bool `json:"strict,omitempty"`
// Name identifies the function.
Name string `json:"name"`
// Parameters is a JSON Schema describing the arguments accepted by the
// function.
Parameters any `json:"parameters"`
// Strict requests exact conformance to Parameters, where the server
// supports it.
Strict bool `json:"strict,omitempty"`
}
// ImageURL holds an image reference for multi-modal messages.
// ImageDetail selects the fidelity used when processing an image. Auto lets the
// endpoint choose.
type ImageDetail string
const (
ImageDetailAuto ImageDetail = "auto"
ImageDetailHigh ImageDetail = "high"
ImageDetailLow ImageDetail = "low"
)
// ImageURL carries an image content part.
type ImageURL struct {
URL string `json:"url"`
Detail string `json:"detail,omitempty"` // "auto", "low", "high"
// Detail is optional; when omitted the server chooses.
Detail ImageDetail `json:"detail,omitempty"`
// URL is a public image URL or a base64 data URI.
URL string `json:"url"`
}
// InputAudio holds audio data for multi-modal messages.
// AudioFormat names an audio encoding accepted in an input_audio
// part.
type AudioFormat string
const (
AudioFormatWAV AudioFormat = "wav"
AudioFormatMP3 AudioFormat = "mp3"
)
// InputAudio carries base64-encoded WAV or MP3 audio.
type InputAudio struct {
Data string `json:"data"`
Format string `json:"format"` // "wav", "mp3"
// Data is base64-encoded audio data.
Data string `json:"data"`
// Format identifies the encoding of Data.
Format AudioFormat `json:"format"`
}
// Message represents a chat message in a conversation.
// Content is always stored as []ContentPart internally.
// A custom UnmarshalJSON handles the JSON polymorphism where content
// can arrive as either a string or an array of parts.
// Message is one message in a conversation or completion response.
//
// Refusal handling:
// The JSON content field may be a string or an array. UnmarshalJSON normalizes
// either representation into ContentParts; marshaling a Message encodes
// ContentParts as an array.
//
// In API responses (ChatCompletionMessage), the model sends refusal as a
// top-level field alongside content:
// See [OpenAI Chat Completions API].
//
// "content": "text response"
// "refusal": "I can't help with that"
//
// Only one of content or refusal will be non-empty.
//
// In API requests (ChatCompletionAssistantMessageParam), the model can
// receive refusal as a content part with type "refusal" in the content
// array. Our ContentPart type supports this for completeness.
//
// See: https://platform.openai.com/docs/api-reference/chat/create
// [OpenAI Chat Completions API]: https://platform.openai.com/docs/api-reference/chat/create
type Message struct {
Role string `json:"role"`
ContentParts []ContentPart `json:"content,omitempty"`
Refusal string `json:"refusal,omitempty"`
ReasoningContent string `json:"reasoning_content,omitempty"`
Reasoning string `json:"reasoning,omitempty"`
Name string `json:"name,omitempty"`
ToolCalls []ToolCall `json:"tool_calls,omitempty"`
ToolCallID string `json:"tool_call_id,omitempty"`
// ContentParts holds the message content in parts; it is the
// in-memory form of the JSON "content" field.
ContentParts []ContentPart `json:"content,omitempty"`
// Name optionally identifies the sender.
Name string `json:"name,omitempty"`
// Reasoning and ReasoningContent are alternative wire fields used by
// compatible servers. After unmarshalling, ReasoningContent contains
// either representation, with reasoning_content taking precedence when
// both occur.
Reasoning string `json:"reasoning,omitempty"`
ReasoningContent string `json:"reasoning_content,omitempty"`
// Refusal contains the endpoint's refusal text, when present.
Refusal string `json:"refusal,omitempty"`
// Role identifies the author or purpose of the message.
Role Role `json:"role"`
// ToolCallID links a tool response to the call it answers.
ToolCallID string `json:"tool_call_id,omitempty"`
// ToolCalls contains calls requested by an assistant message.
ToolCalls []ToolCall `json:"tool_calls,omitempty"`
}
// Text returns the concatenated text from all text content parts.
// Text concatenates the Text payload of every text content part in order. It
// ignores non-text parts and does not fall back to Refusal or reasoning text.
func (m *Message) Text() string {
var b strings.Builder
for _, p := range m.ContentParts {
@@ -218,8 +362,10 @@ func (m *Message) Text() string {
return b.String()
}
// UnmarshalJSON handles the polymorphic content field which can be
// either a string or an array of ContentPart objects.
// UnmarshalJSON decodes a message while normalizing a string-valued content
// field to one text ContentPart. An array-valued content field is decoded in
// order. If reasoning_content is absent, reasoning is copied into
// ReasoningContent so callers can use the latter for either server spelling.
func (m *Message) UnmarshalJSON(data []byte) error {
type Alias Message
aux := &struct {
@@ -235,6 +381,8 @@ func (m *Message) UnmarshalJSON(data []byte) error {
case string:
m.ContentParts = []ContentPart{{Type: ContentTypeText, Text: c}}
case []any:
// Array content replaces any content already held by the Message.
parts := make([]ContentPart, 0, len(c))
for _, raw := range c {
partBytes, err := json.Marshal(raw)
if err != nil {
@@ -244,68 +392,129 @@ func (m *Message) UnmarshalJSON(data []byte) error {
if err := json.Unmarshal(partBytes, &part); err != nil {
return err
}
m.ContentParts = append(m.ContentParts, part)
parts = append(parts, part)
}
m.ContentParts = parts
}
// Normalize: if ReasoningContent is empty but Reasoning is set, use Reasoning
// If ReasoningContent is empty but Reasoning is set, use Reasoning.
if m.ReasoningContent == "" && m.Reasoning != "" {
m.ReasoningContent = m.Reasoning
}
return nil
}
// Model represents an available model from the API.
// Model describes one model exposed by the server's /models
// endpoint.
//
// See [OpenAI Models API].
//
// [OpenAI Models API]: https://platform.openai.com/docs/api-reference/models/list
type Model struct {
ID string `json:"id"`
Created int64 `json:"created"`
Object string `json:"object"`
OwnedBy string `json:"owned_by"`
Meta *ModelMeta `json:"meta,omitempty"`
ContextLength int `json:"context_length,omitempty"`
// ContextLength is the model's context window in tokens, when reported.
ContextLength int `json:"context_length,omitempty"`
// Created is the model's creation time as Unix seconds.
Created int64 `json:"created"`
// ID identifies the model in inference requests.
ID string `json:"id"`
// Meta contains endpoint-specific model metadata, when reported.
Meta *ModelMeta `json:"meta,omitempty"`
// Object is the endpoint's object type for the model entry.
Object string `json:"object"`
// OwnedBy identifies the organization that owns the model, when
//reported.
OwnedBy string `json:"owned_by"`
}
// ModelMeta contains optional model metadata from the provider.
// llama.cpp includes n_ctx (configured context size) and
// n_ctx_train (model's native training context).
// ModelMeta contains llama.cpp model metadata. Context-size fields are measured
// in tokens and are zero when omitted.
//
// See [llama.cpp API documentation].
//
// [llama.cpp API documentation]: https://llama.app/docs/api
type ModelMeta struct {
NCtx int `json:"n_ctx,omitempty"`
// NCtx is the context size with which llama.cpp loaded the model.
NCtx int `json:"n_ctx,omitempty"`
// NCtxTrain is the model's training context size.
NCtxTrain int `json:"n_ctx_train,omitempty"`
}
// ModelsResponse is the response for GET /models.
// ModelsResponse is the response body of GET /models.
type ModelsResponse struct {
Object string `json:"object"`
Data []Model `json:"data"`
// Data contains the models exposed by the endpoint.
Data []Model `json:"data"`
// Object is the endpoint's object type for the response.
Object string `json:"object"`
}
// ResponseFormat specifies the format of the response.
// Role identifies the author or purpose of a message.
type Role string
const (
RoleAssistant Role = "assistant"
RoleDeveloper Role = "developer"
RoleSystem Role = "system"
RoleTool Role = "tool"
RoleUser Role = "user"
)
// ResponseFormat constrains the format of the generated response.
type ResponseFormat struct {
Type string `json:"type"`
JSONSchema any `json:"json_schema,omitempty"`
// JSONSchema defines the required output schema when Type is
// json_schema.
JSONSchema any `json:"json_schema,omitempty"`
// Type selects text, json_object, or json_schema output.
Type string `json:"type"`
}
// StreamChunk is a single SSE chunk from a streaming response.
// StreamChunk is one data event in a streaming response. Delta fields are
// incremental and must be combined in order to reconstruct a message. A chunk
// may contain Usage or Timings without containing a choice.
type StreamChunk struct {
ID string `json:"id"`
Model string `json:"model"`
// Choices contains incremental updates for generated alternatives.
// It may be empty in a metadata-only chunk.
Choices []struct {
Index int `json:"index"`
Delta Delta `json:"delta"`
// Index identifies the choice updated by this delta.
Index int `json:"index"`
// Delta contains the incremental message update.
Delta Delta `json:"delta"`
// FinishReason is set when generation for the choice ends.
FinishReason string `json:"finish_reason,omitempty"`
} `json:"choices"`
Usage *Usage `json:"usage,omitempty"`
// ID identifies the completion to which the chunk belongs.
ID string `json:"id"`
// Model identifies the model generating the response.
Model string `json:"model"`
// Timings contains llama.cpp performance data, when reported.
Timings *Timings `json:"timings,omitempty"`
// Usage contains aggregate token counts, when reported.
Usage *Usage `json:"usage,omitempty"`
}
// StreamOptions controls streaming behavior.
// StreamOptions configures a streaming response.
type StreamOptions struct {
// IncludeUsage requests an additional chunk containing aggregate token
// usage before the terminal data marker. If the stream is interrupted,
// that chunk may not arrive.
IncludeUsage bool `json:"include_usage,omitempty"`
}
// ToolCall represents a tool call made by the LLM.
// ToolCall is one tool call requested by the model.
type ToolCall struct {
Index *int `json:"index,omitempty"` // streaming only
ID string `json:"id,omitempty"`
Type string `json:"type"` // "function"
// Function identifies the function and its generated arguments.
Function Function `json:"function"`
// ID identifies the call. The corresponding tool response carries the
// same value in Message.ToolCallID.
ID string `json:"id,omitempty"`
// Index identifies the call whose fragments a streaming delta
// contains. It is omitted from non-streaming responses.
Index *int `json:"index,omitempty"`
// Type identifies the kind of tool being called.
Type ToolType `json:"type"`
}
// ToolType identifies the kind of tool in a declaration or tool call.
type ToolType string
const (
ToolTypeFunction ToolType = "function"
)
internal/llm/llm_integration_test.go
+51 -73
View File
@@ -5,8 +5,6 @@ import (
"log/slog"
"os"
"testing"
"code.chimeric.al/chimerical/odidere/internal/config"
)
// testURL returns the LLM endpoint for integration tests.
@@ -45,7 +43,7 @@ func TestIntegration_ListModels(t *testing.T) {
Key: testKey(t),
Timeout: "30s",
}
client, err := NewClient(cfg, nil, slog.Default())
client, err := NewClient(cfg, slog.Default())
if err != nil {
t.Fatalf("NewClient: %v", err)
}
@@ -71,16 +69,14 @@ func TestIntegration_Completions(t *testing.T) {
Key: testKey(t),
Timeout: "60s",
}
client, err := NewClient(cfg, nil, slog.Default())
client, err := NewClient(cfg, slog.Default())
if err != nil {
t.Fatalf("NewClient: %v", err)
}
res, err := client.Completions(
context.Background(), CompletionsParams{
Model: testModel(t),
SystemMessage: "",
MaxIterations: 0,
Model: testModel(t),
Messages: []Message{{
Role: RoleUser,
ContentParts: []ContentPart{{
@@ -93,10 +89,7 @@ func TestIntegration_Completions(t *testing.T) {
if err != nil {
t.Fatalf("Completions: %v", err)
}
if len(res.Messages) == 0 {
t.Fatal("expected at least one message")
}
text := res.Messages[len(res.Messages)-1].Text()
text := res.Message.Text()
if text == "" {
t.Error("final message has no text")
}
@@ -109,84 +102,54 @@ func TestIntegration_CompletionsStream(t *testing.T) {
Key: testKey(t),
Timeout: "60s",
}
client, err := NewClient(cfg, nil, slog.Default())
client, err := NewClient(cfg, slog.Default())
if err != nil {
t.Fatalf("NewClient: %v", err)
}
events := make(chan StreamEvent)
errCh := make(chan error, 1)
go func() {
errCh <- client.CompletionsStream(
context.Background(), CompletionsParams{
Model: testModel(t),
SystemMessage: "",
MaxIterations: 0,
Messages: []Message{{
Role: RoleUser,
ContentParts: []ContentPart{{
Type: "text",
Text: "Say hello in three words.",
}},
stream, err := client.CompletionsStream(
context.Background(),
CompletionsParams{
Model: testModel(t),
Messages: []Message{{
Role: RoleUser,
ContentParts: []ContentPart{{
Type: "text",
Text: "Say hello in three words.",
}},
}, events,
)
}()
var messages []Message
for evt := range events {
messages = append(messages, evt.Message)
}
if err := <-errCh; err != nil {
}},
},
)
if err != nil {
t.Fatalf("CompletionsStream: %v", err)
}
if len(messages) == 0 {
t.Fatal("expected at least one message")
defer stream.Close()
for stream.Next() {
}
text := messages[len(messages)-1].Text()
if text == "" {
if err := stream.Err(); err != nil {
t.Fatalf("stream.Err: %v", err)
}
result := stream.Result()
if result.Message.Text() == "" {
t.Error("final message has no text")
}
}
func TestIntegration_ToolCalling(t *testing.T) {
_ = testURL(t)
// Build a minimal tool registry with a single echo tool.
reg, err := NewRegistry([]config.ToolConfig{
{
Name: "echo",
Description: "Echo back the input string",
Command: "echo",
Arguments: []string{"{{.input}}"},
Parameters: map[string]any{
"type": "object",
"properties": map[string]any{
"input": map[string]any{"type": "string"},
},
"required": []string{"input"},
},
},
})
if err != nil {
t.Fatalf("NewRegistry: %v", err)
}
cfg := Config{
client, err := NewClient(Config{
URL: testURL(t),
Key: testKey(t),
Timeout: "60s",
}
client, err := NewClient(cfg, reg, slog.Default())
}, slog.Default())
if err != nil {
t.Fatalf("NewClient: %v", err)
}
res, err := client.Completions(
context.Background(), CompletionsParams{
Model: testModel(t),
SystemMessage: "",
MaxIterations: 10,
context.Background(),
CompletionsParams{
Model: testModel(t),
Messages: []Message{{
Role: RoleUser,
ContentParts: []ContentPart{{
@@ -194,17 +157,32 @@ func TestIntegration_ToolCalling(t *testing.T) {
Text: "Use the echo tool to say 'hello'",
}},
}},
Tools: []APITool{{
Type: "function",
Function: &FunctionDef{
Name: "echo",
Description: "Echo back the input string",
Parameters: map[string]any{
"type": "object",
"properties": map[string]any{
"input": map[string]any{"type": "string"},
},
"required": []string{"input"},
},
},
}},
},
)
if err != nil {
t.Fatalf("Completions: %v", err)
}
if len(res.Messages) == 0 {
t.Fatal("expected at least one message")
if len(res.Message.ToolCalls) == 0 {
t.Fatal("expected a tool call")
}
// The final message should be a non-tool-call assistant message.
final := res.Messages[len(res.Messages)-1]
if len(final.ToolCalls) > 0 {
t.Error("final message still has tool calls")
if res.Message.ToolCalls[0].Function.Name != "echo" {
t.Errorf(
"tool name = %q, want echo",
res.Message.ToolCalls[0].Function.Name,
)
}
}
internal/llm/llm_test.go
+159
View File
@@ -0,0 +1,159 @@
package llm
import (
"context"
"encoding/json"
"errors"
"fmt"
"io"
"log/slog"
"net/http"
"net/http/httptest"
"testing"
)
func TestContentPartMarshalJSON(t *testing.T) {
tests := []struct {
name string
part ContentPart
want string
}{
{
name: "text",
part: ContentPart{
Type: ContentTypeText,
Text: "hello",
},
want: `{"text":"hello","type":"text"}`,
},
{
name: "refusal",
part: ContentPart{
Type: ContentTypeRefusal,
Refusal: "cannot comply",
},
want: `{"refusal":"cannot comply","type":"refusal"}`,
},
{
name: "image URL",
part: ContentPart{
Type: ContentTypeImageURL,
ImageURL: &ImageURL{
URL: "https://example.com/image.png",
},
},
want: `{"image_url":{"url":"https://example.com/image.png"},"type":"image_url"}`,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got, err := json.Marshal(tt.part)
if err != nil {
t.Fatalf("json.Marshal() error = %v", err)
}
if string(got) != tt.want {
t.Errorf("json.Marshal() = %s, want %s", got, tt.want)
}
})
}
}
func TestClientReturnsAPIError(t *testing.T) {
tests := []struct {
name string
call func(*Client) error
wantErr string
}{
{
name: "completions",
call: func(client *Client) error {
_, err := client.Completions(t.Context(), CompletionsParams{})
return err
},
wantErr: "completions: model_not_found: model unavailable",
},
{
name: "completions stream",
call: func(client *Client) error {
_, err := client.CompletionsStream(
t.Context(), CompletionsParams{},
)
return err
},
wantErr: "completions stream: model_not_found: model unavailable",
},
{
name: "list models",
call: func(client *Client) error {
_, err := client.ListModels(context.Background())
return err
},
wantErr: "list models: model_not_found: model unavailable",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(
w http.ResponseWriter, _ *http.Request,
) {
w.WriteHeader(http.StatusNotFound)
fmt.Fprint(w, `{"error":{"code":"model_not_found","message":"model unavailable","type":"invalid_request_error"}}`)
}))
defer server.Close()
log := slog.New(slog.NewTextHandler(io.Discard, nil))
client, err := NewClient(Config{URL: server.URL}, log)
if err != nil {
t.Fatalf("NewClient() error = %v", err)
}
err = tt.call(client)
if err == nil {
t.Fatal("call error = nil, want API error")
}
if err.Error() != tt.wantErr {
t.Errorf("call error = %q, want %q", err, tt.wantErr)
}
var apiErr *APIError
if !errors.As(err, &apiErr) {
t.Fatalf("errors.As(%T) = false, want *APIError", err)
}
if apiErr.Code != "model_not_found" {
t.Errorf("APIError.Code = %q, want %q", apiErr.Code, "model_not_found")
}
if apiErr.Message != "model unavailable" {
t.Errorf("APIError.Message = %q, want %q", apiErr.Message, "model unavailable")
}
if apiErr.Type != "invalid_request_error" {
t.Errorf(
"APIError.Type = %q, want %q",
apiErr.Type,
"invalid_request_error",
)
}
})
}
}
func TestMessageUnmarshalJSONReplacesArrayContent(t *testing.T) {
message := Message{
ContentParts: []ContentPart{{
Type: ContentTypeText,
Text: "old",
}},
}
err := json.Unmarshal(
[]byte(`{"content":[{"type":"text","text":"new"}]}`),
&message,
)
if err != nil {
t.Fatalf("json.Unmarshal() error = %v", err)
}
if got := message.Text(); got != "new" {
t.Errorf("Message.Text() = %q, want %q", got, "new")
}
}
internal/llm/tool.go
-335
View File
@@ -1,335 +0,0 @@
package llm
import (
"bytes"
"context"
"encoding/json"
"fmt"
"net/http"
"os"
"os/exec"
"strings"
"text/template"
"time"
"code.chimeric.al/chimerical/odidere/internal/config"
)
// fn provides template functions available to argument templates.
var fn = template.FuncMap{
"json": func(v any) string {
b, _ := json.Marshal(v)
return string(b)
},
}
// Tool represents an external tool that can be invoked by LLMs.
// Tools are executed as subprocesses with templated arguments,
// and described to the LLM via their JSON Schema parameters.
type Tool struct {
// Name uniquely identifies the tool within a registry.
Name string
// Description explains the tool's purpose for the LLM.
Description string
// Command is the executable path or name.
Command string
// Arguments are Go templates expanded with LLM-provided parameters.
Arguments []string
// Parameters is a JSON Schema describing expected input from the LLM.
Parameters map[string]any
// timeout is the parsed execution duration limit.
timeout time.Duration
}
// NewTool creates a Tool from a ToolConfig.
func NewTool(cfg config.ToolConfig) (*Tool, error) {
var timeout time.Duration
if cfg.Timeout != "" {
var err error
timeout, err = time.ParseDuration(cfg.Timeout)
if err != nil {
return nil, fmt.Errorf(
"invalid timeout %q: %w", cfg.Timeout, err,
)
}
}
return &Tool{
Name: cfg.Name,
Description: cfg.Description,
Command: cfg.Command,
Arguments: cfg.Arguments,
Parameters: cfg.Parameters,
timeout: timeout,
}, nil
}
// ParseArguments expands argument templates with the provided JSON data.
// The args parameter should be a JSON object string; empty string or "{}"
// results in an empty data map. Templates producing empty strings are
// filtered from the result, allowing conditional arguments.
func (t *Tool) ParseArguments(args string) ([]string, error) {
var data = map[string]any{}
if args != "" && args != "{}" {
if err := json.Unmarshal([]byte(args), &data); err != nil {
return nil, fmt.Errorf(
"invalid arguments JSON: %w", err,
)
}
}
var result []string
for _, v := range t.Arguments {
tmpl, err := template.New("").Funcs(fn).Parse(v)
if err != nil {
return nil, fmt.Errorf(
"invalid template %q: %w", v, err,
)
}
var buf bytes.Buffer
if err := tmpl.Execute(&buf, data); err != nil {
return nil, fmt.Errorf(
"execute template %q: %w", v, err,
)
}
// Filter out empty strings (unused conditional arguments).
if s := buf.String(); s != "" {
result = append(result, s)
}
}
return result, nil
}
// ToAPI converts the tool definition into the API format suitable for
// inclusion in a ChatRequest.
func (t *Tool) ToAPI() APITool {
return APITool{
Type: "function",
Function: &FunctionDef{
Name: t.Name,
Description: t.Description,
Parameters: t.Parameters,
},
}
}
// Registry holds tools indexed by name and handles their execution.
// It validates all tool definitions at construction time to fail fast on
// configuration errors.
type Registry struct {
tools map[string]*Tool
}
// NewRegistry creates a registry from the provided tool definitions.
// Returns an error if any tool fails validation or if duplicate names exist.
func NewRegistry(tools []config.ToolConfig) (*Registry, error) {
var r = &Registry{
tools: make(map[string]*Tool),
}
for _, tc := range tools {
t, err := NewTool(tc)
if err != nil {
return nil, fmt.Errorf(
"invalid tool %q: %w", tc.Name, err,
)
}
if _, exists := r.tools[t.Name]; exists {
return nil, fmt.Errorf(
"duplicate tool name: %s", t.Name,
)
}
r.tools[t.Name] = t
}
return r, nil
}
// Get returns a tool by name and a boolean indicating if it was found.
func (r *Registry) Get(name string) (*Tool, bool) {
tool, ok := r.tools[name]
return tool, ok
}
// List returns all registered tool names in arbitrary order.
func (r *Registry) List() []string {
names := make([]string, 0, len(r.tools))
for name := range r.tools {
names = append(names, name)
}
return names
}
// ToAPI converts all registered tools into API format suitable for
// inclusion in a ChatRequest.
func (r *Registry) ToAPI() []APITool {
if r == nil {
return nil
}
tools := make([]APITool, 0, len(r.tools))
for _, t := range r.tools {
tools = append(tools, t.ToAPI())
}
return tools
}
// binaryOutputPrefix is the message shown when tool output is detected as binary.
const binaryOutputPrefix = "error: binary output detected: "
// ToolResult holds the structured output from a tool execution.
type ToolResult struct {
Stdout []byte
Stderr []byte
ProcessState *os.ProcessState
}
// FormatResult formats the tool result using the structured label-based
// format with stdout, stderr, and exit code sections.
//
// http.DetectContentType implements the WHATWG MIME-sniffing spec and
// returns text/plain for any content consisting entirely of printable
// characters and whitespace. Text-based formats like JSON, JavaScript,
// and XML will be detected as text/plain; charset=utf-8. We assume
// checking for the "text/" prefix is sufficient.
func (r *ToolResult) FormatResult() string {
var buf strings.Builder
fmt.Fprintln(&buf, "▸ stdout:")
if len(r.Stdout) > 0 {
ctype := http.DetectContentType(r.Stdout)
if !strings.HasPrefix(ctype, "text/") {
fmt.Fprintf(&buf, binaryOutputPrefix+"%s\n", ctype)
} else {
buf.Write(r.Stdout)
if r.Stdout[len(r.Stdout)-1] != '\n' {
fmt.Fprint(&buf, "\n")
}
}
}
fmt.Fprintln(&buf, "▸ stderr:")
if len(r.Stderr) > 0 {
ctype := http.DetectContentType(r.Stderr)
if !strings.HasPrefix(ctype, "text/") {
fmt.Fprintf(&buf, binaryOutputPrefix+"%s\n", ctype)
} else {
buf.Write(r.Stderr)
if r.Stderr[len(r.Stderr)-1] != '\n' {
fmt.Fprint(&buf, "\n")
}
}
}
exitCode := r.ProcessState.ExitCode()
fmt.Fprintf(&buf, "┗ exit: %d", exitCode)
// Include signal info when the process was killed by a signal.
if exitCode == -1 {
fmt.Fprintf(&buf, " (%s)", r.ProcessState.String())
}
fmt.Fprintln(&buf)
return buf.String()
}
// Execute runs a tool by name with the provided JSON arguments.
// It expands argument templates, executes the command as a subprocess, and
// returns a structured result with stdout, stderr, and exit code.
// The context can be used for cancellation; tool-specific timeouts are
// applied on top of any context deadline.
//
// Start failures (process never ran) return an error directly.
func (r *Registry) Execute(
ctx context.Context, name string, args string,
) (*ToolResult, error) {
tool, ok := r.tools[name]
if !ok {
return nil, fmt.Errorf("unknown tool: %q", name)
}
// Evaluate argument templates.
cmdArgs, err := tool.ParseArguments(args)
if err != nil {
return nil, fmt.Errorf(
"tool %q: parse arguments: %w", name, err,
)
}
// If defined, use the timeout.
if tool.timeout > 0 {
var cancel context.CancelFunc
ctx, cancel = context.WithTimeout(ctx, tool.timeout)
defer cancel()
}
// Setup and run the command.
var (
stdout, stderr bytes.Buffer
cmd = exec.CommandContext(
ctx, tool.Command, cmdArgs...,
)
)
cmd.Stdout = &stdout
cmd.Stderr = &stderr
err = cmd.Run()
// Extract exit code and handle errors.
// After cmd.Run() returns, the possible error cases are:
//
// 1. Context error (ctx.Err() != nil): the context was cancelled or
// timed out. If the process had started, it was killed and
// ProcessState is set. If it had not, Process is nil.
//
// 2. Start failure (cmd.Process == nil): the executable was not found,
// the fork failed, or the context was already done before Start.
// No stdout/stderr to capture.
//
// 3. Process exited with error (err is *exec.ExitError): the process
// ran but exited with a non-zero status or was killed by a signal.
// ProcessState is set and contains the exit code.
//
// In all cases where cmd.Process != nil, ProcessState is guaranteed
// to be set because Start() succeeded.
if err != nil {
switch {
// Context was cancelled or timed out. The process was killed if
// it had started. Distinguish tool timeout from other context
// errors so we can return an appropriate message.
case ctx.Err() != nil:
if ctx.Err() == context.DeadlineExceeded &&
tool.timeout > 0 {
// Tool-specific timeout. Return a clear
// message with the configured duration.
return nil, fmt.Errorf(
"tool %q: timed out after %v",
name, tool.timeout,
)
}
// Context cancelled or parent deadline exceeded.
return nil, fmt.Errorf("tool %q: %w", name, ctx.Err())
// Process never started. No ProcessState, no output.
// The error is the raw start failure (e.g., *PathError).
case cmd.Process == nil:
return nil, fmt.Errorf("tool %q: %w", name, err)
// Process ran and exited with non-zero status or was killed by a signal.
// ProcessState is set; return it along with captured output.
default:
return &ToolResult{
Stdout: stdout.Bytes(),
Stderr: stderr.Bytes(),
ProcessState: cmd.ProcessState,
}, nil
}
}
// No error. Process exited with status 0.
return &ToolResult{
Stdout: stdout.Bytes(),
Stderr: stderr.Bytes(),
ProcessState: cmd.ProcessState,
}, nil
}
internal/llm/tool_test.go
-389
View File
@@ -1,389 +0,0 @@
package llm
import (
"bytes"
"context"
"os/exec"
"strings"
"testing"
"time"
"code.chimeric.al/chimerical/odidere/internal/config"
)
func TestTool_ParseArguments(t *testing.T) {
tool := &Tool{
Arguments: []string{"--name {{.name}}", "{{if .flag}}--enabled{{end}}"},
}
tests := []struct {
name string
args string
expected []string
}{
{
name: "basic",
args: `{"name": "test"}`,
expected: []string{"--name test"},
},
{
name: "with flag",
args: `{"name": "test", "flag": true}`,
expected: []string{"--name test", "--enabled"},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got, err := tool.ParseArguments(tt.args)
if err != nil {
t.Fatalf("ParseArguments failed: %v", err)
}
if len(got) != len(tt.expected) {
t.Errorf("expected len %d, got %d", len(tt.expected), len(got))
}
for i := range got {
if got[i] != tt.expected[i] {
t.Errorf("expected %q, got %q", tt.expected[i], got[i])
}
}
})
}
}
func TestNewTool(t *testing.T) {
tests := []struct {
name string
cfg config.ToolConfig
wantErr bool
}{
{
name: "valid",
cfg: config.ToolConfig{
Name: "test",
Description: "desc",
Command: "ls",
Timeout: "10s",
},
wantErr: false,
},
{
name: "invalid timeout",
cfg: config.ToolConfig{
Name: "test",
Description: "desc",
Command: "ls",
Timeout: "invalid",
},
wantErr: true,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
tool, err := NewTool(tt.cfg)
if (err != nil) != tt.wantErr {
t.Errorf("NewTool() error = %v, wantErr %v", err, tt.wantErr)
}
if err == nil {
if tool.Name != tt.cfg.Name {
t.Errorf("expected name %s, got %s", tt.cfg.Name, tool.Name)
}
if tt.cfg.Timeout == "10s" && tool.timeout != 10*time.Second {
t.Errorf("expected timeout 10s, got %v", tool.timeout)
}
}
})
}
}
func TestRegistry_NewRegistry(t *testing.T) {
tests := []struct {
name string
tools []config.ToolConfig
wantErr bool
}{
{
name: "valid",
tools: []config.ToolConfig{
{Name: "t1", Description: "d1", Command: "c1"},
{Name: "t2", Description: "d2", Command: "c2"},
},
wantErr: false,
},
{
name: "duplicate",
tools: []config.ToolConfig{
{Name: "t1", Description: "d1", Command: "c1"},
{Name: "t1", Description: "d2", Command: "c2"},
},
wantErr: true,
},
{
name: "invalid tool",
tools: []config.ToolConfig{
{Name: "t1", Description: "d1", Command: "c1", Timeout: "invalid"},
},
wantErr: true,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
_, err := NewRegistry(tt.tools)
if (err != nil) != tt.wantErr {
t.Errorf("NewRegistry() error = %v, wantErr %v", err, tt.wantErr)
}
})
}
}
func TestRegistry_ToAPI_Nil(t *testing.T) {
var reg *Registry
got := reg.ToAPI()
if got != nil {
t.Errorf("Registry.ToAPI(nil) = %v, want nil", got)
}
}
func TestRegistry_ToAPI_Empty(t *testing.T) {
reg, err := NewRegistry(nil)
if err != nil {
t.Fatalf("NewRegistry(nil) error = %v", err)
}
got := reg.ToAPI()
if len(got) != 0 {
t.Errorf("Registry.ToAPI(empty) len = %d, want 0", len(got))
}
}
func TestRegistry_ToAPI_Single(t *testing.T) {
reg, err := NewRegistry([]config.ToolConfig{
{
Name: "weather",
Description: "Get the weather",
Parameters: map[string]any{"type": "object"},
},
})
if err != nil {
t.Fatalf("NewRegistry error = %v", err)
}
got := reg.ToAPI()
if len(got) != 1 {
t.Fatalf("len = %d, want 1", len(got))
}
if got[0].Type != "function" {
t.Errorf("type = %q, want %q", got[0].Type, "function")
}
if got[0].Function.Name != "weather" {
t.Errorf("name = %q, want %q", got[0].Function.Name, "weather")
}
if got[0].Function.Description != "Get the weather" {
t.Errorf("description = %q, want %q", got[0].Function.Description, "Get the weather")
}
}
func TestTool_ToAPI(t *testing.T) {
tool := &Tool{
Name: "test",
Description: "A test tool",
Parameters: map[string]any{"type": "object"},
Command: "echo",
Arguments: []string{"hello"},
}
api := tool.ToAPI()
if api.Type != "function" {
t.Errorf("type = %q, want %q", api.Type, "function")
}
if api.Function.Name != "test" {
t.Errorf("name = %q, want %q", api.Function.Name, "test")
}
if api.Function.Description != "A test tool" {
t.Errorf("description = %q, want %q", api.Function.Description, "A test tool")
}
if api.Function.Parameters == nil {
t.Error("parameters is nil")
}
}
func TestToolResult_FormatResult(t *testing.T) {
// Get real ProcessState values from actual commands.
cmdOK := exec.Command("sh", "-c", "exit 0")
if err := cmdOK.Run(); err != nil {
t.Fatalf("cmdOK: %v", err)
}
cmdFail := exec.Command("sh", "-c", "exit 3")
cmdFail.Run() // expect non-nil error
tests := []struct {
name string
result ToolResult
wantContains []string
}{
{
name: "success with stdout",
result: ToolResult{Stdout: []byte("hello\n"), Stderr: nil, ProcessState: cmdOK.ProcessState},
wantContains: []string{"▸ stdout:", "hello", "▸ stderr:", "┗ exit: 0"},
},
{
name: "failure with stderr",
result: ToolResult{Stdout: nil, Stderr: []byte("error occurred\n"), ProcessState: cmdFail.ProcessState},
wantContains: []string{"▸ stdout:", "▸ stderr:", "error occurred", "┗ exit: 3"},
},
{
name: "both stdout and stderr",
result: ToolResult{Stdout: []byte("output"), Stderr: []byte("warning"), ProcessState: cmdOK.ProcessState},
wantContains: []string{"▸ stdout:", "output", "▸ stderr:", "warning", "┗ exit: 0"},
},
{
name: "empty output",
result: ToolResult{Stdout: nil, Stderr: nil, ProcessState: cmdOK.ProcessState},
wantContains: []string{"▸ stdout:", "▸ stderr:", "┗ exit: 0"},
},
{
name: "binary stdout (PNG)",
result: ToolResult{Stdout: []byte("\x89PNG\r\n\x1a\n"), Stderr: nil, ProcessState: cmdOK.ProcessState},
wantContains: []string{"▸ stdout:", binaryOutputPrefix, "image/png", "▸ stderr:", "┗ exit: 0"},
},
{
name: "binary stderr",
result: ToolResult{Stdout: nil, Stderr: []byte("\x89PNG\r\n\x1a\n"), ProcessState: cmdOK.ProcessState},
wantContains: []string{"▸ stdout:", "▸ stderr:", binaryOutputPrefix, "image/png", "┗ exit: 0"},
},
{
name: "HTML",
result: ToolResult{Stdout: []byte("<HTML><BODY>Hello</BODY></HTML>\n"), Stderr: nil, ProcessState: cmdOK.ProcessState},
wantContains: []string{"▸ stdout:", "<HTML>", "▸ stderr:", "┗ exit: 0"},
},
{
name: "XML",
result: ToolResult{Stdout: []byte("<?xml version=\"1.0\"?><root/>\n"), Stderr: nil, ProcessState: cmdOK.ProcessState},
wantContains: []string{"▸ stdout:", "<?xml", "▸ stderr:", "┗ exit: 0"},
},
{
name: "Go source",
result: ToolResult{Stdout: []byte("package main\n\nfunc main() {}\n"), Stderr: nil, ProcessState: cmdOK.ProcessState},
wantContains: []string{"▸ stdout:", "package main", "▸ stderr:", "┗ exit: 0"},
},
{
name: "JSON",
result: ToolResult{Stdout: []byte(`{"key":"value"}`), Stderr: nil, ProcessState: cmdOK.ProcessState},
wantContains: []string{"▸ stdout:", `{"key":"value"}`, "▸ stderr:", "┗ exit: 0"},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got := tt.result.FormatResult()
for _, want := range tt.wantContains {
if !strings.Contains(got, want) {
t.Errorf("FormatResult() = %q, want to contain %q", got, want)
}
}
})
}
}
func TestRegistry_Execute_Success(t *testing.T) {
reg, err := NewRegistry([]config.ToolConfig{
{Name: "echo", Description: "echo", Command: "echo", Arguments: []string{"hello"}},
})
if err != nil {
t.Fatalf("NewRegistry: %v", err)
}
result, err := reg.Execute(context.Background(), "echo", "{}")
if err != nil {
t.Fatalf("Execute: %v", err)
}
if result.ProcessState.ExitCode() != 0 {
t.Errorf("ExitCode = %d, want 0", result.ProcessState.ExitCode())
}
if !bytes.Contains(result.Stdout, []byte("hello")) {
t.Errorf("Stdout = %q, want to contain hello", result.Stdout)
}
}
func TestRegistry_Execute_NonZeroExit(t *testing.T) {
reg, err := NewRegistry([]config.ToolConfig{
{Name: "fail", Description: "fail", Command: "sh", Arguments: []string{"-c", "exit 42"}},
})
if err != nil {
t.Fatalf("NewRegistry: %v", err)
}
result, err := reg.Execute(context.Background(), "fail", "{}")
if err != nil {
t.Fatalf("Execute: %v", err)
}
if result.ProcessState.ExitCode() != 42 {
t.Errorf("ExitCode = %d, want 42", result.ProcessState.ExitCode())
}
}
func TestRegistry_Execute_UnknownCommand(t *testing.T) {
reg, err := NewRegistry([]config.ToolConfig{
{Name: "nope", Description: "nope", Command: "nonexistent_binary_xyz", Arguments: []string{}},
})
if err != nil {
t.Fatalf("NewRegistry: %v", err)
}
_, err = reg.Execute(context.Background(), "nope", "{}")
if err == nil {
t.Fatal("Execute: want error, got nil")
}
}
func TestRegistry_Execute_UnknownTool(t *testing.T) {
reg, err := NewRegistry([]config.ToolConfig{})
if err != nil {
t.Fatalf("NewRegistry: %v", err)
}
_, err = reg.Execute(context.Background(), "missing", "{}")
if err == nil {
t.Fatal("Execute: want error, got nil")
}
if !strings.Contains(err.Error(), "unknown tool") {
t.Errorf("err = %v, want unknown tool error", err)
}
}
func TestRegistry_Execute_Timeout(t *testing.T) {
reg, err := NewRegistry([]config.ToolConfig{
{Name: "sleep", Description: "sleep", Command: "sleep", Arguments: []string{"10"}, Timeout: "50ms"},
})
if err != nil {
t.Fatalf("NewRegistry: %v", err)
}
_, err = reg.Execute(context.Background(), "sleep", "{}")
if err == nil {
t.Fatal("Execute: want timeout error, got nil")
}
if !strings.Contains(err.Error(), "timed out") {
t.Errorf("err = %v, want timed out error", err)
}
}
func TestRegistry_Execute_SignalKilled(t *testing.T) {
reg, err := NewRegistry([]config.ToolConfig{
{Name: "kill", Description: "kill", Command: "sh", Arguments: []string{"-c", "sleep 100 & kill -9 $!; wait"}},
})
if err != nil {
t.Fatalf("NewRegistry: %v", err)
}
result, err := reg.Execute(context.Background(), "kill", "{}")
if err != nil {
t.Fatalf("Execute: %v", err)
}
// The process we kill exits with signal, but the shell itself may exit 0.
// On some systems the exit code is 137 (128+9). Check ProcessState for signal info.
if result.ProcessState.ExitCode() == -1 {
if !strings.Contains(result.ProcessState.String(), "signal:") {
t.Errorf("ProcessState.String() = %q, want signal info", result.ProcessState.String())
}
}
}