victorvargas/rony-llm-agent
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Merge pull request #8 from VictorVargas/fix/providers-agent-loop

Browse source 
Fix providers (OpenAI/llamacpp/Anthropic) and agent loop correctness
This commit is contained in:
Victor Hugo Vargas Servin 2026-07-12 16:22:19 -07:00 committed by GitHub
commit ee9319b823
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
19 changed files with 410 additions and 254 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

14
pkg/agent/integration_test.go
View file

@ -96,10 +96,10 @@ func TestIntegration_AgentLoop_Generate(t *testing.T) {
}) })
loop := agent.New(agent.Config{ loop := agent.New(agent.Config{
LLM: client, LLM: client,
Persona: persona.DefaultPersona(), Persona: persona.DefaultPersona(),
Tools: registry, Tools: registry,
Sandbox: &mockSandbox{}, Sandbox: &mockSandbox{},
MaxIters: 3, MaxIters: 3,
}) })
@ -124,9 +124,9 @@ func TestIntegration_AgentLoop_Stream(t *testing.T) {
} }
loop := agent.New(agent.Config{ loop := agent.New(agent.Config{
LLM: client, LLM: client,
Persona: persona.DefaultPersona(), Persona: persona.DefaultPersona(),
Tools: tools.NewRegistry(), Tools: tools.NewRegistry(),
MaxIters: 3, MaxIters: 3,
}) })

59
pkg/agent/loop.go
View file

@ -47,19 +47,19 @@ type Config struct {
// Iteration represents a single cycle of the agent loop. // Iteration represents a single cycle of the agent loop.
type Iteration struct { type Iteration struct {
Number int Number int
ToolCalls []llm.ToolCall ToolCalls []llm.ToolCall
ToolsUsed int ToolsUsed int
Duration time.Duration Duration time.Duration
} }
// Response is the final output of the agent loop. // Response is the final output of the agent loop.
type Response struct { type Response struct {
Content string Content string
ToolCalls []llm.ToolCall ToolCalls []llm.ToolCall
Iterations int Iterations int
Duration time.Duration Duration time.Duration
TokenUsage llm.TokenUsage TokenUsage llm.TokenUsage
} }
// Loop is the main agent loop that orchestrates LLM calls and tool execution. // Loop is the main agent loop that orchestrates LLM calls and tool execution.
@ -82,17 +82,21 @@ func (l *Loop) Run(ctx context.Context, input string, history ...llm.Message) (R
start := time.Now() start := time.Now()
messages := l.buildInitialMessages(input, history) messages := l.buildInitialMessages(input, history)
// Tool schemas don't change between iterations, so build the JSON once
// per Run instead of re-marshaling every tool on every loop pass.
toolSchemas := l.getToolSchemas()
var finalContent string var finalContent string
var allToolCalls []llm.ToolCall var allToolCalls []llm.ToolCall
var totalUsage llm.TokenUsage var totalUsage llm.TokenUsage
iterations := 0 iterations := 0
completed := false
for iterations < l.cfg.MaxIters { for iterations < l.cfg.MaxIters {
iterations++ iterations++
resp, err := l.cfg.LLM.Generate(ctx, llm.CompletionRequest{ resp, err := l.cfg.LLM.Generate(ctx, llm.CompletionRequest{
Messages: messages, Messages: messages,
Tools: l.getToolSchemas(), Tools: toolSchemas,
ChatTemplateKwargs: l.cfg.ChatTemplateKwargs, ChatTemplateKwargs: l.cfg.ChatTemplateKwargs,
}) })
if err != nil { if err != nil {
@ -105,6 +109,7 @@ func (l *Loop) Run(ctx context.Context, input string, history ...llm.Message) (R
if len(resp.ToolCalls) == 0 { if len(resp.ToolCalls) == 0 {
finalContent = resp.Content finalContent = resp.Content
completed = true
break break
} }
@ -138,16 +143,20 @@ func (l *Loop) Run(ctx context.Context, input string, history ...llm.Message) (R
} }
duration := time.Since(start) duration := time.Since(start)
if iterations >= l.cfg.MaxIters { // Only report the max-iterations failure when the loop actually ran out
// of budget without producing a final answer — an answer that arrives
// exactly on the last allowed iteration is still a success (the old
// `iterations >= MaxIters` check threw that valid response away).
if !completed {
return Response{}, fmt.Errorf("max iterations (%d) reached", l.cfg.MaxIters) return Response{}, fmt.Errorf("max iterations (%d) reached", l.cfg.MaxIters)
} }
return Response{ return Response{
Content: finalContent, Content: finalContent,
ToolCalls: allToolCalls, ToolCalls: allToolCalls,
Iterations: iterations, Iterations: iterations,
Duration: duration, Duration: duration,
TokenUsage: totalUsage, TokenUsage: totalUsage,
}, nil }, nil
} }
@ -157,6 +166,8 @@ func (l *Loop) Run(ctx context.Context, input string, history ...llm.Message) (R
func (l *Loop) RunStream(ctx context.Context, input string, history ...llm.Message) iter.Seq2[llm.StreamChunk, error] { func (l *Loop) RunStream(ctx context.Context, input string, history ...llm.Message) iter.Seq2[llm.StreamChunk, error] {
return func(yield func(llm.StreamChunk, error) bool) { return func(yield func(llm.StreamChunk, error) bool) {
messages := l.buildInitialMessages(input, history) messages := l.buildInitialMessages(input, history)
// Same as Run: the schemas are identical on every iteration.
toolSchemas := l.getToolSchemas()
iterations := 0 iterations := 0
for iterations < l.cfg.MaxIters { for iterations < l.cfg.MaxIters {
@ -164,7 +175,7 @@ func (l *Loop) RunStream(ctx context.Context, input string, history ...llm.Messa
stream := l.cfg.LLM.Stream(ctx, llm.CompletionRequest{ stream := l.cfg.LLM.Stream(ctx, llm.CompletionRequest{
Messages: messages, Messages: messages,
Tools: l.getToolSchemas(), Tools: toolSchemas,
ChatTemplateKwargs: l.cfg.ChatTemplateKwargs, ChatTemplateKwargs: l.cfg.ChatTemplateKwargs,
}) })
@ -226,11 +237,23 @@ func (l *Loop) RunStream(ctx context.Context, input string, history ...llm.Messa
// real token counts. // real token counts.
hasContent := chunk.Delta != "" || chunk.ReasoningDelta != "" hasContent := chunk.Delta != "" || chunk.ReasoningDelta != ""
hasUsage := chunk.Usage.TotalTokens > 0 hasUsage := chunk.Usage.TotalTokens > 0
if !hasToolCalls && (hasContent || hasUsage) { switch {
case !hasToolCalls && (hasContent || hasUsage):
responseBuilder.WriteString(chunk.Delta) responseBuilder.WriteString(chunk.Delta)
if !yield(chunk, nil) { if !yield(chunk, nil) {
return return
} }
case hasToolCalls && hasUsage:
// The content gate above exists to hide raw provider
// deltas during a tool-call round, but it also swallowed
// that round's token usage — so callers tracking context
// occupancy (e.g. a UI's context bar deciding when to
// compact) only ever saw the usage of the final,
// tool-free round. Forward the usage on its own,
// without the content.
if !yield(llm.StreamChunk{Usage: chunk.Usage}, nil) {
return
}
} }
} }

22
pkg/agent/loop_test.go
View file

@ -28,8 +28,8 @@ func (m *mockLLM) Stream(ctx context.Context, req llm.CompletionRequest) iter.Se
return m.streamFunc(ctx, req) return m.streamFunc(ctx, req)
} }
func (m *mockLLM) Name() string { return "mock" } func (m *mockLLM) Name() string { return "mock" }
func (m *mockLLM) Capabilities() llm.ProviderCapabilities { return llm.ProviderCapabilities{} } func (m *mockLLM) Capabilities() llm.ProviderCapabilities { return llm.ProviderCapabilities{} }
type mockSandbox struct { type mockSandbox struct {
validateFunc func(tool tools.Tool, call llm.ToolCall) error validateFunc func(tool tools.Tool, call llm.ToolCall) error
@ -71,9 +71,9 @@ func TestNew_CustomMaxIters(t *testing.T) {
} }
loop := New(Config{ loop := New(Config{
LLM: mockClient, LLM: mockClient,
Persona: persona.DefaultPersona(), Persona: persona.DefaultPersona(),
Tools: tools.NewRegistry(), Tools: tools.NewRegistry(),
MaxIters: 10, MaxIters: 10,
}) })
@ -373,9 +373,9 @@ func TestRun_MaxIterations(t *testing.T) {
}) })
loop := New(Config{ loop := New(Config{
LLM: mockClient, LLM: mockClient,
Persona: persona.DefaultPersona(), Persona: persona.DefaultPersona(),
Tools: registry, Tools: registry,
MaxIters: 3, MaxIters: 3,
}) })
@ -706,9 +706,9 @@ func TestRun_Stream_MaxIterations(t *testing.T) {
}) })
loop := New(Config{ loop := New(Config{
LLM: mockClient, LLM: mockClient,
Persona: persona.DefaultPersona(), Persona: persona.DefaultPersona(),
Tools: registry, Tools: registry,
MaxIters: 1, MaxIters: 1,
}) })

30
pkg/config/config.go
View file

@ -7,34 +7,34 @@ import (
// ProviderConfig holds provider-specific settings. // ProviderConfig holds provider-specific settings.
type ProviderConfig struct { type ProviderConfig struct {
Type string `yaml:"type"` Type string `yaml:"type"`
Model string `yaml:"model"` Model string `yaml:"model"`
APIKey string `yaml:"api_key"` APIKey string `yaml:"api_key"`
BaseURL string `yaml:"base_url,omitempty"` BaseURL string `yaml:"base_url,omitempty"`
MaxTokens int `yaml:"max_tokens,omitempty"` MaxTokens int `yaml:"max_tokens,omitempty"`
Temperature float32 `yaml:"temperature,omitempty"` Temperature float32 `yaml:"temperature,omitempty"`
} }
// ToolPolicy controls which tools are available and their permissions. // ToolPolicy controls which tools are available and their permissions.
type ToolPolicy struct { type ToolPolicy struct {
DefaultPermission string `yaml:"default_permission"` DefaultPermission string `yaml:"default_permission"`
AllowList []string `yaml:"allow_list,omitempty"` AllowList []string `yaml:"allow_list,omitempty"`
DenyList []string `yaml:"deny_list,omitempty"` DenyList []string `yaml:"deny_list,omitempty"`
} }
// LoggingConfig controls logging output. // LoggingConfig controls logging output.
type LoggingConfig struct { type LoggingConfig struct {
Level string `yaml:"level"` Level string `yaml:"level"`
Format string `yaml:"format"` Format string `yaml:"format"`
Output string `yaml:"output"` Output string `yaml:"output"`
} }
// Config is the top-level configuration for the agent. // Config is the top-level configuration for the agent.
type Config struct { type Config struct {
Model string `yaml:"model"` Model string `yaml:"model"`
Provider ProviderConfig `yaml:"provider"` Provider ProviderConfig `yaml:"provider"`
Tools ToolPolicy `yaml:"tools"` Tools ToolPolicy `yaml:"tools"`
Logging LoggingConfig `yaml:"logging"` Logging LoggingConfig `yaml:"logging"`
} }
// Loader is responsible for loading configuration from various sources. // Loader is responsible for loading configuration from various sources.

6
pkg/llm/mock/mock.go
View file

@ -10,9 +10,9 @@ import (
// MockLLMClient is a deterministic implementation of llm.LLMClient for testing. // MockLLMClient is a deterministic implementation of llm.LLMClient for testing.
type MockLLMClient struct { type MockLLMClient struct {
GenerateFunc func(ctx context.Context, req llm.CompletionRequest) (llm.CompletionResponse, error) GenerateFunc func(ctx context.Context, req llm.CompletionRequest) (llm.CompletionResponse, error)
StreamFunc func(ctx context.Context, req llm.CompletionRequest) iter.Seq2[llm.StreamChunk, error] StreamFunc func(ctx context.Context, req llm.CompletionRequest) iter.Seq2[llm.StreamChunk, error]
NameFunc func() string NameFunc func() string
CapabilitiesFunc func() llm.ProviderCapabilities CapabilitiesFunc func() llm.ProviderCapabilities
} }

2
pkg/llm/mock/mock_test.go
View file

@ -167,7 +167,7 @@ func TestMockLLMClient_MatchResponse(t *testing.T) {
client := NewWithMatch([]MatchResponse{ client := NewWithMatch([]MatchResponse{
{Match: "hello", Response: "hi there!"}, {Match: "hello", Response: "hi there!"},
{Match: "world", Response: "earth"}, {Match: "world", Response: "earth"},
{Match: "*", Response: "default"}, {Match: "*", Response: "default"},
}) })
resp, _ := client.Generate(context.Background(), llm.CompletionRequest{ resp, _ := client.Generate(context.Background(), llm.CompletionRequest{

93
pkg/llm/providers/anthropic/client.go
View file

@ -7,39 +7,48 @@ import (
"encoding/json" "encoding/json"
"fmt" "fmt"
"io" "io"
"net/http"
"iter" "iter"
"net/http"
"strings" "strings"
"github.com/VictorVargas/rony-llm-agent/pkg/llm" "github.com/VictorVargas/rony-llm-agent/pkg/llm"
) )
const ( const (
defaultBaseURL = "https://api.anthropic.com/v1" defaultBaseURL = "https://api.anthropic.com/v1"
defaultModel = "claude-opus-4-8" defaultModel = "claude-opus-4-8"
defaultMaxTokens = 8192 defaultMaxTokens = 8192
anthropicVersion = "2023-06-01" anthropicVersion = "2023-06-01"
// defaultContextWindow is what Capabilities() reports when
// Config.ContextWindow is unset: 200k tokens, the standard window for
// Claude models. Callers use this number to decide when to compact
// their conversation, so over-reporting it (the old code assumed 1M
// for anything that wasn't haiku) meant compaction fired far too late
// and requests started overflowing the real window.
defaultContextWindow = 200000
) )
// Config holds the settings needed to create an Anthropic client. // Config holds the settings needed to create an Anthropic client.
type Config struct { type Config struct {
APIKey string APIKey string
Model string // defaults to claude-opus-4-8 Model string // defaults to claude-opus-4-8
BaseURL string // defaults to https://api.anthropic.com/v1 BaseURL string // defaults to https://api.anthropic.com/v1
MaxTokens int // default max_tokens sent on every request (Anthropic requires one); 0 = defaultMaxTokens MaxTokens int // default max_tokens sent on every request (Anthropic requires one); 0 = defaultMaxTokens
Temperature *float32 ContextWindow int // model's context window in tokens (0 = defaultContextWindow); raise it only for models/plans with an extended window
TopP *float32 Temperature *float32
TopP *float32
} }
// Client implements llm.LLMClient for Anthropic. // Client implements llm.LLMClient for Anthropic.
type Client struct { type Client struct {
apiKey string apiKey string
baseURL string baseURL string
model string model string
maxTokens int maxTokens int
temperature *float32 contextWindow int
topP *float32 temperature *float32
http *http.Client topP *float32
http *http.Client
} }
// New returns a new Anthropic client. // New returns a new Anthropic client.
@ -63,14 +72,20 @@ func New(cfg Config) (*Client, error) {
maxTokens = defaultMaxTokens maxTokens = defaultMaxTokens
} }
contextWindow := cfg.ContextWindow
if contextWindow == 0 {
contextWindow = defaultContextWindow
}
return &Client{ return &Client{
apiKey: cfg.APIKey, apiKey: cfg.APIKey,
baseURL: baseURL, baseURL: baseURL,
model: model, model: model,
maxTokens: maxTokens, maxTokens: maxTokens,
temperature: cfg.Temperature, contextWindow: contextWindow,
topP: cfg.TopP, temperature: cfg.Temperature,
http: http.DefaultClient, topP: cfg.TopP,
http: http.DefaultClient,
}, nil }, nil
} }
@ -257,15 +272,11 @@ func (c *Client) Name() string {
} }
func (c *Client) Capabilities() llm.ProviderCapabilities { func (c *Client) Capabilities() llm.ProviderCapabilities {
maxContext := 1000000
if strings.Contains(c.model, "haiku") {
maxContext = 200000
}
return llm.ProviderCapabilities{ return llm.ProviderCapabilities{
SupportsTools: true, SupportsTools: true,
SupportsVision: true, SupportsVision: true,
SupportsJSON: true, SupportsJSON: true,
MaxContextWindow: maxContext, MaxContextWindow: c.contextWindow,
} }
} }
@ -489,20 +500,20 @@ func mapStopReason(reason string) string {
// Anthropic API types // Anthropic API types
type anthropicRequest struct { type anthropicRequest struct {
Model string `json:"model"` Model string `json:"model"`
Messages []anthropicMessage `json:"messages"` Messages []anthropicMessage `json:"messages"`
System string `json:"system,omitempty"` System string `json:"system,omitempty"`
MaxTokens int `json:"max_tokens"` MaxTokens int `json:"max_tokens"`
Tools []anthropicTool `json:"tools,omitempty"` Tools []anthropicTool `json:"tools,omitempty"`
ToolChoice json.RawMessage `json:"tool_choice,omitempty"` ToolChoice json.RawMessage `json:"tool_choice,omitempty"`
Temperature *float32 `json:"temperature,omitempty"` Temperature *float32 `json:"temperature,omitempty"`
TopP *float32 `json:"top_p,omitempty"` TopP *float32 `json:"top_p,omitempty"`
StopSequences []string `json:"stop_sequences,omitempty"` StopSequences []string `json:"stop_sequences,omitempty"`
Stream bool `json:"stream,omitempty"` Stream bool `json:"stream,omitempty"`
} }
type anthropicMessage struct { type anthropicMessage struct {
Role string `json:"role"` Role string `json:"role"`
Content []anthropicContentBlock `json:"content"` Content []anthropicContentBlock `json:"content"`
} }

115
pkg/llm/providers/llamacpp/client.go
View file

@ -2,12 +2,13 @@ package llamacpp
import ( import (
"bufio" "bufio"
"bytes"
"context" "context"
"encoding/json" "encoding/json"
"fmt" "fmt"
"io" "io"
"net/http"
"iter" "iter"
"net/http"
"strings" "strings"
"time" "time"
@ -24,15 +25,15 @@ const defaultContextWindow = 32768
// Config holds the settings needed to create a llama.cpp client. // Config holds the settings needed to create a llama.cpp client.
type Config struct { type Config struct {
BaseURL string // defaults to http://localhost:8080/v1 BaseURL string // defaults to http://localhost:8080/v1
Model string Model string
Timeout int // request timeout in seconds (0 = default, no timeout) Timeout int // request timeout in seconds (0 = default, no timeout)
ContextWindow int // model's context window in tokens (0 = defaultContextWindow) ContextWindow int // model's context window in tokens (0 = defaultContextWindow)
MaxTokens int // default max_tokens (0 = defaultMaxTokens) MaxTokens int // default max_tokens (0 = defaultMaxTokens)
TopK int // top-k sampling (0 = model/server default) TopK int // top-k sampling (0 = model/server default)
TopP float32 // nucleus sampling (0 = model/server default) TopP float32 // nucleus sampling (0 = model/server default)
Temperature float32 Temperature float32
MinP float32 // min-p sampling (llama.cpp extension) MinP float32 // min-p sampling (llama.cpp extension)
PresencePenalty float32 PresencePenalty float32
RepetitionPenalty float32 // sent as the server's `repeat_penalty` field RepetitionPenalty float32 // sent as the server's `repeat_penalty` field
MaxThinkingTokens int // best-effort cap on reasoning tokens; ignored by servers that don't support it MaxThinkingTokens int // best-effort cap on reasoning tokens; ignored by servers that don't support it
@ -123,7 +124,7 @@ func (c *Client) Generate(ctx context.Context, req llm.CompletionRequest) (llm.C
return llm.CompletionResponse{}, fmt.Errorf("decoding response: %w", err) return llm.CompletionResponse{}, fmt.Errorf("decoding response: %w", err)
} }
return c.toResponse(apiResp), nil return c.toResponse(apiResp)
} }
func (c *Client) Stream(ctx context.Context, req llm.CompletionRequest) iter.Seq2[llm.StreamChunk, error] { func (c *Client) Stream(ctx context.Context, req llm.CompletionRequest) iter.Seq2[llm.StreamChunk, error] {
@ -191,6 +192,11 @@ func (c *Client) Stream(ctx context.Context, req llm.CompletionRequest) iter.Seq
} }
scanner := bufio.NewScanner(resp.Body) scanner := bufio.NewScanner(resp.Body)
// A single SSE line can exceed bufio.Scanner's 64KB default cap
// (e.g. a large tool-call arguments delta or a long reasoning
// event), which would kill the stream mid-turn with "token too
// long" — same headroom the anthropic client already reserves.
scanner.Buffer(make([]byte, 0, 64*1024), 4*1024*1024)
for scanner.Scan() { for scanner.Scan() {
line := scanner.Text() line := scanner.Text()
if !strings.HasPrefix(line, "data: ") { if !strings.HasPrefix(line, "data: ") {
@ -370,11 +376,16 @@ func (c *Client) buildRequest(req llm.CompletionRequest, stream bool) (io.Reader
if err != nil { if err != nil {
return nil, fmt.Errorf("marshaling request: %w", err) return nil, fmt.Errorf("marshaling request: %w", err)
} }
return strings.NewReader(string(data)), nil return bytes.NewReader(data), nil
} }
// toResponse converts a llama.cpp API response to our CompletionResponse. // toResponse converts a llama.cpp API response to our CompletionResponse.
func (c *Client) toResponse(resp llamaChatResponse) llm.CompletionResponse { func (c *Client) toResponse(resp llamaChatResponse) (llm.CompletionResponse, error) {
// Guard against a 200 response with no choices (e.g. a misbehaving
// server or proxy) — indexing Choices[0] blindly panics the whole app.
if len(resp.Choices) == 0 {
return llm.CompletionResponse{}, fmt.Errorf("llamacpp: response contained no choices")
}
choice := resp.Choices[0] choice := resp.Choices[0]
result := llm.CompletionResponse{ result := llm.CompletionResponse{
ID: resp.ID, ID: resp.ID,
@ -398,32 +409,32 @@ func (c *Client) toResponse(resp llamaChatResponse) llm.CompletionResponse {
TotalTokens: resp.Usage.TotalTokens, TotalTokens: resp.Usage.TotalTokens,
} }
return result return result, nil
} }
// llama.cpp API types // llama.cpp API types
type llamaChatRequest struct { type llamaChatRequest struct {
Model string `json:"model"` Model string `json:"model"`
Messages []llamaMessage `json:"messages"` Messages []llamaMessage `json:"messages"`
Tools []llamaTool `json:"tools,omitempty"` Tools []llamaTool `json:"tools,omitempty"`
ToolChoice interface{} `json:"tool_choice,omitempty"` ToolChoice interface{} `json:"tool_choice,omitempty"`
Temperature float32 `json:"temperature,omitempty"` Temperature float32 `json:"temperature,omitempty"`
MaxTokens int `json:"max_tokens,omitempty"` MaxTokens int `json:"max_tokens,omitempty"`
TopK int `json:"top_k,omitempty"` TopK int `json:"top_k,omitempty"`
TopP float32 `json:"top_p,omitempty"` TopP float32 `json:"top_p,omitempty"`
MinP float32 `json:"min_p,omitempty"` MinP float32 `json:"min_p,omitempty"`
PresencePenalty float32 `json:"presence_penalty,omitempty"` PresencePenalty float32 `json:"presence_penalty,omitempty"`
RepeatPenalty float32 `json:"repeat_penalty,omitempty"` RepeatPenalty float32 `json:"repeat_penalty,omitempty"`
// MaxThinkingTokens is a best-effort reasoning-token cap: not part of // MaxThinkingTokens is a best-effort reasoning-token cap: not part of
// upstream llama.cpp's server API, but harmless to send since JSON // upstream llama.cpp's server API, but harmless to send since JSON
// servers ignore unrecognized fields, and some front-ends (e.g. the // servers ignore unrecognized fields, and some front-ends (e.g. the
// proxy this model's config was written for) do honor it. // proxy this model's config was written for) do honor it.
MaxThinkingTokens int `json:"max_thinking_tokens,omitempty"` MaxThinkingTokens int `json:"max_thinking_tokens,omitempty"`
Stop []string `json:"stop,omitempty"` Stop []string `json:"stop,omitempty"`
Stream bool `json:"stream"` Stream bool `json:"stream"`
StreamOptions *llamaStreamOptions `json:"stream_options,omitempty"` StreamOptions *llamaStreamOptions `json:"stream_options,omitempty"`
ChatTemplateKwargs map[string]any `json:"chat_template_kwargs,omitempty"` ChatTemplateKwargs map[string]any `json:"chat_template_kwargs,omitempty"`
} }
type llamaStreamOptions struct { type llamaStreamOptions struct {
@ -444,16 +455,16 @@ type llamaTool struct {
} }
type llamaChatResponse struct { type llamaChatResponse struct {
ID string `json:"id"` ID string `json:"id"`
Model string `json:"model"` Model string `json:"model"`
Choices []llamaChoice `json:"choices"` Choices []llamaChoice `json:"choices"`
Usage llamaUsage `json:"usage"` Usage llamaUsage `json:"usage"`
} }
type llamaChoice struct { type llamaChoice struct {
Index int `json:"index"` Index int `json:"index"`
Message llamaMessageResult `json:"message"` Message llamaMessageResult `json:"message"`
FinishReason string `json:"finish_reason"` FinishReason string `json:"finish_reason"`
} }
type llamaMessageResult struct { type llamaMessageResult struct {
@ -464,14 +475,14 @@ type llamaMessageResult struct {
} }
type llamaToolCall struct { type llamaToolCall struct {
ID string `json:"id"` ID string `json:"id"`
Type string `json:"type"` Type string `json:"type"`
Function llamaFunction `json:"function"` Function llamaFunction `json:"function"`
} }
type llamaFunction struct { type llamaFunction struct {
Name string `json:"name"` Name string `json:"name"`
Arguments string `json:"arguments"` Arguments string `json:"arguments"`
} }
type llamaUsage struct { type llamaUsage struct {
@ -483,32 +494,32 @@ type llamaUsage struct {
// Stream event types // Stream event types
type llamaStreamEvent struct { type llamaStreamEvent struct {
ID string `json:"id"` ID string `json:"id"`
Choices []llamaStreamChoice `json:"choices"` Choices []llamaStreamChoice `json:"choices"`
Usage *llamaUsage `json:"usage"` Usage *llamaUsage `json:"usage"`
} }
type llamaStreamChoice struct { type llamaStreamChoice struct {
Index int `json:"index"` Index int `json:"index"`
Delta llamaStreamDelta `json:"delta"` Delta llamaStreamDelta `json:"delta"`
FinishReason string `json:"finish_reason"` FinishReason string `json:"finish_reason"`
} }
type llamaStreamDelta struct { type llamaStreamDelta struct {
Content string `json:"content"` Content string `json:"content"`
ReasoningContent string `json:"reasoning_content"` ReasoningContent string `json:"reasoning_content"`
Role string `json:"role"` Role string `json:"role"`
ToolCalls []llamaStreamToolCall `json:"tool_calls"` ToolCalls []llamaStreamToolCall `json:"tool_calls"`
} }
type llamaStreamToolCall struct { type llamaStreamToolCall struct {
Index int `json:"index"` Index int `json:"index"`
ID string `json:"id"` ID string `json:"id"`
Type string `json:"type"` Type string `json:"type"`
Function llamaStreamFunction `json:"function"` Function llamaStreamFunction `json:"function"`
} }
type llamaStreamFunction struct { type llamaStreamFunction struct {
Name string `json:"name"` Name string `json:"name"`
Arguments string `json:"arguments"` Arguments string `json:"arguments"`
} }

8
pkg/llm/providers/llamacpp/client_test.go
View file

@ -34,8 +34,8 @@ func TestClient_Generate(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "application/json") w.Header().Set("Content-Type", "application/json")
json.NewEncoder(w).Encode(llamaChatResponse{ json.NewEncoder(w).Encode(llamaChatResponse{
ID: "llama-123", ID: "llama-123",
Model: "llama3", Model: "llama3",
Choices: []llamaChoice{ Choices: []llamaChoice{
{ {
Index: 0, Index: 0,
@ -189,8 +189,8 @@ func TestClient_Generate_ToolCall(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "application/json") w.Header().Set("Content-Type", "application/json")
json.NewEncoder(w).Encode(llamaChatResponse{ json.NewEncoder(w).Encode(llamaChatResponse{
ID: "llama-tool-1", ID: "llama-tool-1",
Model: "llama3", Model: "llama3",
Choices: []llamaChoice{ Choices: []llamaChoice{
{ {
Index: 0, Index: 0,

219
pkg/llm/providers/openai/client.go
View file

@ -2,12 +2,13 @@ package openai
import ( import (
"bufio" "bufio"
"bytes"
"context" "context"
"encoding/json" "encoding/json"
"fmt" "fmt"
"io" "io"
"net/http"
"iter" "iter"
"net/http"
"strings" "strings"
"github.com/VictorVargas/rony-llm-agent/pkg/llm" "github.com/VictorVargas/rony-llm-agent/pkg/llm"
@ -15,14 +16,15 @@ import (
// Config holds the settings needed to create an OpenAI client. // Config holds the settings needed to create an OpenAI client.
type Config struct { type Config struct {
APIKey string APIKey string
Model string Model string
BaseURL string // defaults to https://api.openai.com/v1 BaseURL string // defaults to https://api.openai.com/v1
} }
// Client implements llm.LLMClient for OpenAI. // Client implements llm.LLMClient for OpenAI.
type Client struct { type Client struct {
apiKey string apiKey string
model string
baseURL string baseURL string
http *http.Client http *http.Client
} }
@ -40,6 +42,7 @@ func New(cfg Config) (*Client, error) {
return &Client{ return &Client{
apiKey: cfg.APIKey, apiKey: cfg.APIKey,
model: cfg.Model,
baseURL: baseURL, baseURL: baseURL,
http: http.DefaultClient, http: http.DefaultClient,
}, nil }, nil
@ -48,7 +51,7 @@ func New(cfg Config) (*Client, error) {
func (c *Client) Generate(ctx context.Context, req llm.CompletionRequest) (llm.CompletionResponse, error) { func (c *Client) Generate(ctx context.Context, req llm.CompletionRequest) (llm.CompletionResponse, error) {
endpoint := c.baseURL + "/chat/completions" endpoint := c.baseURL + "/chat/completions"
payload, err := c.buildRequest(req) payload, err := c.buildRequest(req, false)
if err != nil { if err != nil {
return llm.CompletionResponse{}, fmt.Errorf("building request: %w", err) return llm.CompletionResponse{}, fmt.Errorf("building request: %w", err)
} }
@ -76,14 +79,14 @@ func (c *Client) Generate(ctx context.Context, req llm.CompletionRequest) (llm.C
return llm.CompletionResponse{}, fmt.Errorf("decoding response: %w", err) return llm.CompletionResponse{}, fmt.Errorf("decoding response: %w", err)
} }
return c.toResponse(apiResp), nil return c.toResponse(apiResp)
} }
func (c *Client) Stream(ctx context.Context, req llm.CompletionRequest) iter.Seq2[llm.StreamChunk, error] { func (c *Client) Stream(ctx context.Context, req llm.CompletionRequest) iter.Seq2[llm.StreamChunk, error] {
return func(yield func(llm.StreamChunk, error) bool) { return func(yield func(llm.StreamChunk, error) bool) {
endpoint := c.baseURL + "/chat/completions" endpoint := c.baseURL + "/chat/completions"
payload, err := c.buildRequest(req) payload, err := c.buildRequest(req, true)
if err != nil { if err != nil {
yield(llm.StreamChunk{}, fmt.Errorf("building request: %w", err)) yield(llm.StreamChunk{}, fmt.Errorf("building request: %w", err))
return return
@ -111,7 +114,45 @@ func (c *Client) Stream(ctx context.Context, req llm.CompletionRequest) iter.Seq
return return
} }
// toolCallAccum buffers one tool call's fragments as they stream in:
// the SSE format sends the id/name in the first delta for a given
// tool-call index and the (potentially large) arguments JSON in
// pieces across many subsequent deltas, so it can't be handed to a
// tool handler until it's fully assembled. Same accumulation the
// llamacpp client does — without it, tool calls made over a stream
// were silently dropped and the agent loop never executed them.
type toolCallAccum struct {
id string
name string
args strings.Builder
}
toolCallFrags := map[int]*toolCallAccum{}
var toolCallOrder []int
flushToolCalls := func() []llm.ToolCall {
if len(toolCallOrder) == 0 {
return nil
}
calls := make([]llm.ToolCall, 0, len(toolCallOrder))
for _, idx := range toolCallOrder {
frag := toolCallFrags[idx]
calls = append(calls, llm.ToolCall{
ID: frag.id,
Name: frag.name,
Arguments: json.RawMessage(frag.args.String()),
})
}
toolCallFrags = map[int]*toolCallAccum{}
toolCallOrder = nil
return calls
}
scanner := bufio.NewScanner(resp.Body) scanner := bufio.NewScanner(resp.Body)
// A single SSE line can exceed bufio.Scanner's 64KB default cap
// (e.g. a large tool-call arguments delta), which would kill the
// stream with "token too long" — same headroom the anthropic
// client already reserves.
scanner.Buffer(make([]byte, 0, 64*1024), 4*1024*1024)
for scanner.Scan() { for scanner.Scan() {
line := scanner.Text() line := scanner.Text()
if !strings.HasPrefix(line, "data: ") { if !strings.HasPrefix(line, "data: ") {
@ -128,13 +169,61 @@ func (c *Client) Stream(ctx context.Context, req llm.CompletionRequest) iter.Seq
return return
} }
var usage llm.TokenUsage
if event.Usage != nil {
usage = llm.TokenUsage{
InputTokens: event.Usage.PromptTokens,
OutputTokens: event.Usage.CompletionTokens,
TotalTokens: event.Usage.TotalTokens,
}
}
if len(event.Choices) == 0 {
// The usage-only event (per stream_options.include_usage)
// carries no choices, so it needs its own chunk.
if event.Usage != nil {
if !yield(llm.StreamChunk{Usage: usage}, nil) {
return
}
}
continue
}
for _, choice := range event.Choices { for _, choice := range event.Choices {
hasFragment := len(choice.Delta.ToolCalls) > 0
for _, tc := range choice.Delta.ToolCalls {
frag, ok := toolCallFrags[tc.Index]
if !ok {
frag = &toolCallAccum{}
toolCallFrags[tc.Index] = frag
toolCallOrder = append(toolCallOrder, tc.Index)
}
if tc.ID != "" {
frag.id = tc.ID
}
if tc.Function.Name != "" {
frag.name = tc.Function.Name
}
frag.args.WriteString(tc.Function.Arguments)
}
chunk := llm.StreamChunk{ chunk := llm.StreamChunk{
Delta: choice.Delta.Content, Delta: choice.Delta.Content,
Usage: usage,
} }
if choice.FinishReason != "" { if choice.FinishReason != "" {
chunk.FinishReason = choice.FinishReason chunk.FinishReason = choice.FinishReason
chunk.ToolCalls = flushToolCalls()
} }
// A fragment-only event (a piece of a tool call's streamed
// arguments, with nothing else in this delta) has nothing
// yet for the agent loop to act on: it was buffered above,
// so skip yielding an empty chunk for it.
if hasFragment && chunk.Delta == "" && chunk.FinishReason == "" {
continue
}
if !yield(chunk, nil) { if !yield(chunk, nil) {
return return
} }
@ -160,7 +249,7 @@ func (c *Client) Capabilities() llm.ProviderCapabilities {
} }
// buildRequest converts an llm.CompletionRequest to the OpenAI API format. // buildRequest converts an llm.CompletionRequest to the OpenAI API format.
func (c *Client) buildRequest(req llm.CompletionRequest) (io.Reader, error) { func (c *Client) buildRequest(req llm.CompletionRequest, stream bool) (io.Reader, error) {
// Convert messages to OpenAI format // Convert messages to OpenAI format
messages := make([]openaiMessage, len(req.Messages)) messages := make([]openaiMessage, len(req.Messages))
for i, m := range req.Messages { for i, m := range req.Messages {
@ -195,10 +284,25 @@ func (c *Client) buildRequest(req llm.CompletionRequest) (io.Reader, error) {
tools[i] = tool tools[i] = tool
} }
// The per-request model wins when set; otherwise fall back to the
// client's configured one (Config.Model used to be discarded entirely,
// so every request went out with an empty model — a hard API error on
// OpenAI, and the agent loop never sets req.Model).
model := req.Model
if model == "" {
model = c.model
}
openaiReq := openaiChatRequest{ openaiReq := openaiChatRequest{
Model: req.Model, Model: model,
Messages: messages, Messages: messages,
Stream: false, Stream: stream,
}
if stream {
// Ask for a final SSE event carrying token usage (OpenAI-style
// streaming omits it otherwise), so callers can track real token
// counts per turn instead of always seeing zero.
openaiReq.StreamOptions = &openaiStreamOptions{IncludeUsage: true}
} }
if len(tools) > 0 { if len(tools) > 0 {
openaiReq.Tools = tools openaiReq.Tools = tools
@ -222,11 +326,14 @@ func (c *Client) buildRequest(req llm.CompletionRequest) (io.Reader, error) {
if err != nil { if err != nil {
return nil, fmt.Errorf("marshaling request: %w", err) return nil, fmt.Errorf("marshaling request: %w", err)
} }
return strings.NewReader(string(data)), nil return bytes.NewReader(data), nil
} }
// toResponse converts an OpenAI API response to our CompletionResponse. // toResponse converts an OpenAI API response to our CompletionResponse.
func (c *Client) toResponse(resp openaiChatResponse) llm.CompletionResponse { func (c *Client) toResponse(resp openaiChatResponse) (llm.CompletionResponse, error) {
if len(resp.Choices) == 0 {
return llm.CompletionResponse{}, fmt.Errorf("openai: response contained no choices")
}
choice := resp.Choices[0] choice := resp.Choices[0]
result := llm.CompletionResponse{ result := llm.CompletionResponse{
ID: resp.ID, ID: resp.ID,
@ -249,63 +356,68 @@ func (c *Client) toResponse(resp openaiChatResponse) llm.CompletionResponse {
TotalTokens: resp.Usage.TotalTokens, TotalTokens: resp.Usage.TotalTokens,
} }
return result return result, nil
} }
// OpenAI API types // OpenAI API types
type openaiChatRequest struct { type openaiChatRequest struct {
Model string `json:"model"` Model string `json:"model"`
Messages []openaiMessage `json:"messages"` Messages []openaiMessage `json:"messages"`
Tools []openaiTool `json:"tools,omitempty"` Tools []openaiTool `json:"tools,omitempty"`
ToolChoice interface{} `json:"tool_choice,omitempty"` ToolChoice interface{} `json:"tool_choice,omitempty"`
Temperature *float32 `json:"temperature,omitempty"` Temperature *float32 `json:"temperature,omitempty"`
MaxTokens *int `json:"max_tokens,omitempty"` MaxTokens *int `json:"max_tokens,omitempty"`
Stop []string `json:"stop,omitempty"` Stop []string `json:"stop,omitempty"`
Stream bool `json:"stream"` Stream bool `json:"stream"`
StreamOptions *openaiStreamOptions `json:"stream_options,omitempty"`
}
type openaiStreamOptions struct {
IncludeUsage bool `json:"include_usage"`
} }
type openaiMessage struct { type openaiMessage struct {
Role string `json:"role"` Role string `json:"role"`
Content string `json:"content"` Content string `json:"content"`
ToolCallID string `json:"tool_call_id,omitempty"` ToolCallID string `json:"tool_call_id,omitempty"`
Name string `json:"name,omitempty"` Name string `json:"name,omitempty"`
ToolCalls []openaiToolCall `json:"tool_calls,omitempty"` ToolCalls []openaiToolCall `json:"tool_calls,omitempty"`
} }
type openaiTool struct { type openaiTool struct {
Type string `json:"type"` Type string `json:"type"`
Function json.RawMessage `json:"function"` Function json.RawMessage `json:"function"`
} }
type openaiChatResponse struct { type openaiChatResponse struct {
ID string `json:"id"` ID string `json:"id"`
Model string `json:"model"` Model string `json:"model"`
Choices []openaiChoice `json:"choices"` Choices []openaiChoice `json:"choices"`
Usage openaiUsage `json:"usage"` Usage openaiUsage `json:"usage"`
} }
type openaiChoice struct { type openaiChoice struct {
Index int `json:"index"` Index int `json:"index"`
Message openaiMessageResult `json:"message"` Message openaiMessageResult `json:"message"`
FinishReason string `json:"finish_reason"` FinishReason string `json:"finish_reason"`
} }
type openaiMessageResult struct { type openaiMessageResult struct {
Role string `json:"role"` Role string `json:"role"`
Content string `json:"content"` Content string `json:"content"`
ToolCalls []openaiToolCall `json:"tool_calls"` ToolCalls []openaiToolCall `json:"tool_calls"`
} }
type openaiToolCall struct { type openaiToolCall struct {
ID string `json:"id"` ID string `json:"id"`
Type string `json:"type"` Type string `json:"type"`
Function openaiFunction `json:"function"` Function openaiFunction `json:"function"`
} }
type openaiFunction struct { type openaiFunction struct {
Name string `json:"name"` Name string `json:"name"`
Arguments string `json:"arguments"` Arguments string `json:"arguments"`
} }
type openaiUsage struct { type openaiUsage struct {
@ -317,30 +429,31 @@ type openaiUsage struct {
// Stream event types // Stream event types
type openaiStreamEvent struct { type openaiStreamEvent struct {
ID string `json:"id"` ID string `json:"id"`
Choices []openaiStreamChoice `json:"choices"` Choices []openaiStreamChoice `json:"choices"`
Usage *openaiUsage `json:"usage"`
} }
type openaiStreamChoice struct { type openaiStreamChoice struct {
Index int `json:"index"` Index int `json:"index"`
Delta openaiStreamDelta `json:"delta"` Delta openaiStreamDelta `json:"delta"`
FinishReason string `json:"finish_reason"` FinishReason string `json:"finish_reason"`
} }
type openaiStreamDelta struct { type openaiStreamDelta struct {
Content string `json:"content"` Content string `json:"content"`
Role string `json:"role"` Role string `json:"role"`
ToolCalls []openaiStreamToolCall `json:"tool_calls"` ToolCalls []openaiStreamToolCall `json:"tool_calls"`
} }
type openaiStreamToolCall struct { type openaiStreamToolCall struct {
Index int `json:"index"` Index int `json:"index"`
ID string `json:"id"` ID string `json:"id"`
Type string `json:"type"` Type string `json:"type"`
Function openaiStreamFunction `json:"function"` Function openaiStreamFunction `json:"function"`
} }
type openaiStreamFunction struct { type openaiStreamFunction struct {
Name string `json:"name"` Name string `json:"name"`
Arguments string `json:"arguments"` Arguments string `json:"arguments"`
} }

16
pkg/llm/providers/openai/client_test.go
View file

@ -33,8 +33,8 @@ func TestClient_Generate(t *testing.T) {
w.Header().Set("Content-Type", "application/json") w.Header().Set("Content-Type", "application/json")
w.WriteHeader(http.StatusOK) w.WriteHeader(http.StatusOK)
json.NewEncoder(w).Encode(openaiChatResponse{ json.NewEncoder(w).Encode(openaiChatResponse{
ID: "test-123", ID: "test-123",
Model: "gpt-4", Model: "gpt-4",
Choices: []openaiChoice{ Choices: []openaiChoice{
{ {
Index: 0, Index: 0,
@ -55,8 +55,8 @@ func TestClient_Generate(t *testing.T) {
defer server.Close() defer server.Close()
client, err := New(Config{ client, err := New(Config{
APIKey: "test-key", APIKey: "test-key",
BaseURL: server.URL, BaseURL: server.URL,
}) })
if err != nil { if err != nil {
t.Fatalf("unexpected error: %v", err) t.Fatalf("unexpected error: %v", err)
@ -90,8 +90,8 @@ func TestClient_Generate_Error(t *testing.T) {
defer server.Close() defer server.Close()
client, err := New(Config{ client, err := New(Config{
APIKey: "bad-key", APIKey: "bad-key",
BaseURL: server.URL, BaseURL: server.URL,
}) })
if err != nil { if err != nil {
t.Fatalf("unexpected error: %v", err) t.Fatalf("unexpected error: %v", err)
@ -113,8 +113,8 @@ func TestClient_Stream(t *testing.T) {
defer server.Close() defer server.Close()
client, err := New(Config{ client, err := New(Config{
APIKey: "test-key", APIKey: "test-key",
BaseURL: server.URL + "/v1", BaseURL: server.URL + "/v1",
}) })
if err != nil { if err != nil {
t.Fatalf("unexpected error: %v", err) t.Fatalf("unexpected error: %v", err)

38
pkg/llm/types.go
View file

@ -10,18 +10,18 @@ import (
type Role string type Role string
const ( const (
RoleSystem Role = "system" RoleSystem Role = "system"
RoleUser Role = "user" RoleUser Role = "user"
RoleAssistant Role = "assistant" RoleAssistant Role = "assistant"
RoleTool Role = "tool" RoleTool Role = "tool"
) )
// Message is a single message in a conversation. // Message is a single message in a conversation.
type Message struct { type Message struct {
Role Role `json:"role"` Role Role `json:"role"`
Content string `json:"content"` Content string `json:"content"`
ToolCallID string `json:"tool_call_id,omitempty"` ToolCallID string `json:"tool_call_id,omitempty"`
Name string `json:"name,omitempty"` Name string `json:"name,omitempty"`
// ToolCalls records the calls an assistant message requested, so the // ToolCalls records the calls an assistant message requested, so the
// agent loop can replay them on the next request: without this, the // agent loop can replay them on the next request: without this, the
// conversation sent back to the model has tool-result messages with no // conversation sent back to the model has tool-result messages with no
@ -50,8 +50,8 @@ type TokenUsage struct {
type ToolChoice string type ToolChoice string
const ( const (
ToolChoiceAuto ToolChoice = "auto" ToolChoiceAuto ToolChoice = "auto"
ToolChoiceNone ToolChoice = "none" ToolChoiceNone ToolChoice = "none"
ToolChoiceRequired ToolChoice = "required" ToolChoiceRequired ToolChoice = "required"
) )
@ -73,14 +73,14 @@ func (t *ToolRef) MarshalJSON() ([]byte, error) {
// CompletionRequest is sent to an LLM provider. // CompletionRequest is sent to an LLM provider.
type CompletionRequest struct { type CompletionRequest struct {
Model string `json:"model"` Model string `json:"model"`
Messages []Message `json:"messages"` Messages []Message `json:"messages"`
Tools []json.RawMessage `json:"tools,omitempty"` Tools []json.RawMessage `json:"tools,omitempty"`
ToolChoice interface{} `json:"tool_choice,omitempty"` // ToolChoice, ToolRef, or null ToolChoice interface{} `json:"tool_choice,omitempty"` // ToolChoice, ToolRef, or null
Temperature *float32 `json:"temperature,omitempty"` Temperature *float32 `json:"temperature,omitempty"`
MaxTokens *int `json:"max_tokens,omitempty"` MaxTokens *int `json:"max_tokens,omitempty"`
Stop []string `json:"stop,omitempty"` Stop []string `json:"stop,omitempty"`
Metadata map[string]string `json:"metadata,omitempty"` Metadata map[string]string `json:"metadata,omitempty"`
ChatTemplateKwargs map[string]any `json:"chat_template_kwargs,omitempty"` // model-specific chat template params, e.g. Qwen enable_thinking ChatTemplateKwargs map[string]any `json:"chat_template_kwargs,omitempty"` // model-specific chat template params, e.g. Qwen enable_thinking
} }
@ -98,9 +98,9 @@ type CompletionResponse struct {
// StopReason values. // StopReason values.
const ( const (
StopReasonEndTurn = "end_turn" StopReasonEndTurn = "end_turn"
StopReasonToolUse = "tool_use" StopReasonToolUse = "tool_use"
StopReasonMaxTokens = "max_tokens" StopReasonMaxTokens = "max_tokens"
StopReasonStopSeq = "stop_sequence" StopReasonStopSeq = "stop_sequence"
) )
// ToolCall represents a function invocation requested by the model. // ToolCall represents a function invocation requested by the model.

6
pkg/llm/types_test.go
View file

@ -121,7 +121,7 @@ func TestToolCall_JSON(t *testing.T) {
func TestStreamChunk_JSON(t *testing.T) { func TestStreamChunk_JSON(t *testing.T) {
chunk := StreamChunk{ chunk := StreamChunk{
Delta: "hello", Delta: "hello",
ToolCalls: []ToolCall{}, ToolCalls: []ToolCall{},
} }
@ -139,7 +139,5 @@ func TestStreamChunk_JSON(t *testing.T) {
} }
} }
func float32Ptr(f float32) *float32 { return &f }
func float32Ptr(f float32) *float32 { return &f }
func intPtr(i int) *int { return &i } func intPtr(i int) *int { return &i }

2
pkg/persona/persona.go
View file

@ -18,7 +18,7 @@ type Persona struct {
Style string Style string
Language string Language string
Constraints []string Constraints []string
FewShot []llm.Message FewShot []llm.Message
} }
// Loader loads personas from files. // Loader loads personas from files.

8
pkg/rag/backends/chroma/chroma.go
View file

@ -224,13 +224,13 @@ func stringifyMap(m map[string]interface{}) map[string]string {
// chromaQueryResponse represents the structure of a ChromaDB query response. // chromaQueryResponse represents the structure of a ChromaDB query response.
type chromaQueryResponse struct { type chromaQueryResponse struct {
Names []string `json:"names"` Names []string `json:"names"`
Results []chromaQueryResults `json:"results"` Results []chromaQueryResults `json:"results"`
} }
type chromaQueryResults struct { type chromaQueryResults struct {
IDs [][]string `json:"ids"` IDs [][]string `json:"ids"`
Documents [][]string `json:"documents"` Documents [][]string `json:"documents"`
Distances [][]float64 `json:"distances"` Distances [][]float64 `json:"distances"`
Metadatas [][]map[string]interface{} `json:"metadatas"` Metadatas [][]map[string]interface{} `json:"metadatas"`
} }

2
pkg/rag/backends/chroma/chroma_test.go
View file

@ -54,7 +54,7 @@ func TestBackend_Search(t *testing.T) {
meta := []map[string]interface{}{{"key": "value"}} meta := []map[string]interface{}{{"key": "value"}}
metaNested := [][]map[string]interface{}{meta} metaNested := [][]map[string]interface{}{meta}
mockResponse := map[string]interface{}{ mockResponse := map[string]interface{}{
"names": []string{"rony-memory"}, "names": []string{"rony-memory"},
"results": []map[string]interface{}{ "results": []map[string]interface{}{
{ {
"ids": [][]string{{"test-id"}}, "ids": [][]string{{"test-id"}},

6
pkg/rag/memory_test.go
View file

@ -193,9 +193,9 @@ func TestMemory_Search_EmbeddingErrorFallsBackToLexicalSearch(t *testing.T) {
// mockBackend implements chroma.Backend for testing. // mockBackend implements chroma.Backend for testing.
type mockBackend struct { type mockBackend struct {
upsertFunc func(ctx context.Context, id string, vector []float32, content string, metadata map[string]string) error upsertFunc func(ctx context.Context, id string, vector []float32, content string, metadata map[string]string) error
searchFunc func(ctx context.Context, query string, queryVector []float32, topK int) ([]rag.SearchResult, error) searchFunc func(ctx context.Context, query string, queryVector []float32, topK int) ([]rag.SearchResult, error)
forgetAllFunc func(ctx context.Context) error forgetAllFunc func(ctx context.Context) error
} }
func (m *mockBackend) Upsert(ctx context.Context, id string, vector []float32, content string, metadata map[string]string) error { func (m *mockBackend) Upsert(ctx context.Context, id string, vector []float32, content string, metadata map[string]string) error {

6
pkg/tools/registry.go
View file

@ -6,9 +6,9 @@ import (
// registry is the default implementation of Registry. // registry is the default implementation of Registry.
type registry struct { type registry struct {
mu sync.RWMutex mu sync.RWMutex
tools map[string]Tool tools map[string]Tool
order []string order []string
} }
// NewRegistry returns a new empty registry. // NewRegistry returns a new empty registry.

12
pkg/tools/types.go
View file

@ -45,10 +45,10 @@ type ToolHandler func(ctx context.Context, args json.RawMessage) (ToolResult, er
// ToolResult is returned by a ToolHandler. // ToolResult is returned by a ToolHandler.
type ToolResult struct { type ToolResult struct {
Content string Content string
IsError bool IsError bool
Metadata map[string]string Metadata map[string]string
Artifacts []Artifact Artifacts []Artifact
} }
// Artifact represents a file or data artifact produced by a tool. // Artifact represents a file or data artifact produced by a tool.
@ -60,8 +60,8 @@ type Artifact struct {
// ToolExample provides few-shot examples for the LLM to improve tool usage. // ToolExample provides few-shot examples for the LLM to improve tool usage.
type ToolExample struct { type ToolExample struct {
Input map[string]interface{} Input map[string]interface{}
Output string Output string
} }
// Registry manages tool registration and lookup. // Registry manages tool registration and lookup.