From 25090297772a9ce95d980bbd765be20a0f05ad1f Mon Sep 17 00:00:00 2001 From: Victor Vargas Date: Tue, 30 Jun 2026 23:53:15 -0700 Subject: [PATCH 01/11] chore: add YAML config loading and UUID dependencies --- go.mod | 6 +++++- go.sum | 6 ++++++ 2 files changed, 11 insertions(+), 1 deletion(-) create mode 100644 go.sum diff --git a/go.mod b/go.mod index 2cc601f..f4ccdc1 100644 --- a/go.mod +++ b/go.mod @@ -1,3 +1,7 @@ module github.com/VictorVargas/rony-llm-agent -go 1.26 \ No newline at end of file +go 1.26 + +require gopkg.in/yaml.v3 v3.0.1 + +require github.com/google/uuid v1.6.0 // indirect diff --git a/go.sum b/go.sum new file mode 100644 index 0000000..b4c5744 --- /dev/null +++ b/go.sum @@ -0,0 +1,6 @@ +github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= +github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM= +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= +gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= +gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= From ce64b68f1268f16c503219558383614329dd36e0 Mon Sep 17 00:00:00 2001 From: Victor Vargas Date: Tue, 30 Jun 2026 23:53:20 -0700 Subject: [PATCH 02/11] feat(config): add YAML config loader with hierarchical precedence --- pkg/config/config.go | 70 +++++++++++++++++++++++++ pkg/config/config_test.go | 82 +++++++++++++++++++++++++++++ pkg/config/loader.go | 105 ++++++++++++++++++++++++++++++++++++++ 3 files changed, 257 insertions(+) create mode 100644 pkg/config/config.go create mode 100644 pkg/config/config_test.go create mode 100644 pkg/config/loader.go diff --git a/pkg/config/config.go b/pkg/config/config.go new file mode 100644 index 0000000..a663a54 --- /dev/null +++ b/pkg/config/config.go @@ -0,0 +1,70 @@ +package config + +import ( + "context" + "fmt" +) + +// ProviderConfig holds provider-specific settings. +type ProviderConfig struct { + Type string `yaml:"type"` + Model string `yaml:"model"` + APIKey string `yaml:"api_key"` + BaseURL string `yaml:"base_url,omitempty"` + MaxTokens int `yaml:"max_tokens,omitempty"` + Temperature float32 `yaml:"temperature,omitempty"` +} + +// ToolPolicy controls which tools are available and their permissions. +type ToolPolicy struct { + DefaultPermission string `yaml:"default_permission"` + AllowList []string `yaml:"allow_list,omitempty"` + DenyList []string `yaml:"deny_list,omitempty"` +} + +// LoggingConfig controls logging output. +type LoggingConfig struct { + Level string `yaml:"level"` + Format string `yaml:"format"` + Output string `yaml:"output"` +} + +// Config is the top-level configuration for the agent. +type Config struct { + Model string `yaml:"model"` + Provider ProviderConfig `yaml:"provider"` + Tools ToolPolicy `yaml:"tools"` + Logging LoggingConfig `yaml:"logging"` +} + +// Loader is responsible for loading configuration from various sources. +type Loader interface { + Load(ctx context.Context, workdir string) (Config, error) +} + +// ErrConfigNotFound is returned when no configuration file is found. +var ErrConfigNotFound = fmt.Errorf("config file not found") + +// ErrInvalidConfig is returned when the configuration is invalid. +var ErrInvalidConfig = fmt.Errorf("invalid configuration") + +// LoadDefault returns a Config with sensible defaults. +func LoadDefault() Config { + return Config{ + Model: "gpt-4o", + Provider: ProviderConfig{ + Type: "openai", + Model: "gpt-4o", + Temperature: 0.7, + MaxTokens: 4096, + }, + Tools: ToolPolicy{ + DefaultPermission: "allow", + }, + Logging: LoggingConfig{ + Level: "info", + Format: "text", + Output: "stderr", + }, + } +} diff --git a/pkg/config/config_test.go b/pkg/config/config_test.go new file mode 100644 index 0000000..f7878b1 --- /dev/null +++ b/pkg/config/config_test.go @@ -0,0 +1,82 @@ +package config + +import ( + "context" + "os" + "path/filepath" + "testing" +) + +func TestLoadDefault(t *testing.T) { + cfg := LoadDefault() + if cfg.Model != "gpt-4o" { + t.Errorf("expected default model 'gpt-4o', got %q", cfg.Model) + } + if cfg.Provider.Type != "openai" { + t.Errorf("expected default provider 'openai', got %q", cfg.Provider.Type) + } +} + +func TestYAMLLoader_NoFile(t *testing.T) { + loader := NewYAMLLoader() + ctx := context.Background() + cfg, err := loader.Load(ctx, "/tmp/nonexistent_workdir_12345") + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if cfg.Model != "gpt-4o" { + t.Errorf("expected default model, got %q", cfg.Model) + } +} + +func TestYAMLLoader_WithFile(t *testing.T) { + tmpDir := t.TempDir() + yaml := ` +provider: + type: anthropic + model: claude-sonnet-4.5 + api_key: test-key +` + if err := os.WriteFile(filepath.Join(tmpDir, "rony.yaml"), []byte(yaml), 0644); err != nil { + t.Fatal(err) + } + + loader := NewYAMLLoader() + ctx := context.Background() + cfg, err := loader.Load(ctx, tmpDir) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if cfg.Provider.Type != "anthropic" { + t.Errorf("expected provider 'anthropic', got %q", cfg.Provider.Type) + } + if cfg.Provider.Model != "claude-sonnet-4.5" { + t.Errorf("expected model 'claude-sonnet-4.5', got %q", cfg.Provider.Model) + } + if cfg.Provider.APIKey != "test-key" { + t.Errorf("expected api_key 'test-key', got %q", cfg.Provider.APIKey) + } +} + +func TestYAMLLoader_PreservesDefaults(t *testing.T) { + tmpDir := t.TempDir() + yaml := ` +provider: + type: openai + model: gpt-4o-mini +` + if err := os.WriteFile(filepath.Join(tmpDir, "rony.yaml"), []byte(yaml), 0644); err != nil { + t.Fatal(err) + } + + loader := NewYAMLLoader() + ctx := context.Background() + cfg, err := loader.Load(ctx, tmpDir) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + // Should still have default logging config + if cfg.Logging.Level != "info" { + t.Errorf("expected default logging level 'info', got %q", cfg.Logging.Level) + } +} diff --git a/pkg/config/loader.go b/pkg/config/loader.go new file mode 100644 index 0000000..d0f2e2f --- /dev/null +++ b/pkg/config/loader.go @@ -0,0 +1,105 @@ +package config + +import ( + "context" + "fmt" + "os" + "path/filepath" + + "gopkg.in/yaml.v3" +) + +// yamlLoader loads config from YAML files with hierarchical precedence. +type yamlLoader struct{} + +// NewYAMLLoader returns a Loader that reads from YAML files. +func NewYAMLLoader() Loader { + return &yamlLoader{} +} + +func (l *yamlLoader) Load(ctx context.Context, workdir string) (Config, error) { + defaults := LoadDefault() + + // Load from workdir + path := filepath.Join(workdir, "rony.yaml") + if _, err := os.Stat(path); err == nil { + cfg, err := loadFromFile(path) + if err != nil { + return Config{}, fmt.Errorf("loading %s: %w", path, err) + } + // Merge with defaults + cfg = merge(defaults, cfg) + return cfg, nil + } + + // Load from home directory + home, err := os.UserHomeDir() + if err == nil { + homePath := filepath.Join(home, ".config", "rony", "config.yaml") + if _, err := os.Stat(homePath); err == nil { + cfg, err := loadFromFile(homePath) + if err == nil { + return merge(defaults, cfg), nil + } + } + } + + return defaults, nil +} + +func loadFromFile(path string) (Config, error) { + data, err := os.ReadFile(path) + if err != nil { + return Config{}, err + } + + var cfg Config + if err := yaml.Unmarshal(data, &cfg); err != nil { + return Config{}, fmt.Errorf("parsing YAML: %w", err) + } + return cfg, nil +} + +// merge combines two configs, with b taking precedence over a. +func merge(a, b Config) Config { + if b.Model != "" { + a.Model = b.Model + } + if b.Provider.Type != "" { + a.Provider.Type = b.Provider.Type + } + if b.Provider.Model != "" { + a.Provider.Model = b.Provider.Model + } + if b.Provider.APIKey != "" { + a.Provider.APIKey = b.Provider.APIKey + } + if b.Provider.BaseURL != "" { + a.Provider.BaseURL = b.Provider.BaseURL + } + if b.Provider.MaxTokens > 0 { + a.Provider.MaxTokens = b.Provider.MaxTokens + } + if b.Provider.Temperature > 0 { + a.Provider.Temperature = b.Provider.Temperature + } + if b.Tools.DefaultPermission != "" { + a.Tools.DefaultPermission = b.Tools.DefaultPermission + } + if len(b.Tools.AllowList) > 0 { + a.Tools.AllowList = b.Tools.AllowList + } + if len(b.Tools.DenyList) > 0 { + a.Tools.DenyList = b.Tools.DenyList + } + if b.Logging.Level != "" { + a.Logging.Level = b.Logging.Level + } + if b.Logging.Format != "" { + a.Logging.Format = b.Logging.Format + } + if b.Logging.Output != "" { + a.Logging.Output = b.Logging.Output + } + return a +} From 641481022fbe5082b1048c0100d0403b85469b08 Mon Sep 17 00:00:00 2001 From: Victor Vargas Date: Tue, 30 Jun 2026 23:53:22 -0700 Subject: [PATCH 03/11] feat(llm): add LLM client interface, types, providers (openai, llamacpp), and mock --- pkg/llm/mock/mock.go | 121 ++++++++ pkg/llm/providers/llamacpp/client.go | 322 +++++++++++++++++++++ pkg/llm/providers/llamacpp/client_test.go | 232 +++++++++++++++ pkg/llm/providers/openai/client.go | 327 ++++++++++++++++++++++ pkg/llm/providers/openai/client_test.go | 141 ++++++++++ pkg/llm/types.go | 120 ++++++++ 6 files changed, 1263 insertions(+) create mode 100644 pkg/llm/mock/mock.go create mode 100644 pkg/llm/providers/llamacpp/client.go create mode 100644 pkg/llm/providers/llamacpp/client_test.go create mode 100644 pkg/llm/providers/openai/client.go create mode 100644 pkg/llm/providers/openai/client_test.go create mode 100644 pkg/llm/types.go diff --git a/pkg/llm/mock/mock.go b/pkg/llm/mock/mock.go new file mode 100644 index 0000000..f34f625 --- /dev/null +++ b/pkg/llm/mock/mock.go @@ -0,0 +1,121 @@ +package mock + +import ( + "context" + "iter" + "strings" + + "github.com/VictorVargas/rony-llm-agent/pkg/llm" +) + +// MockLLMClient is a deterministic implementation of llm.LLMClient for testing. +type MockLLMClient struct { + GenerateFunc func(ctx context.Context, req llm.CompletionRequest) (llm.CompletionResponse, error) + StreamFunc func(ctx context.Context, req llm.CompletionRequest) iter.Seq2[llm.StreamChunk, error] + NameFunc func() string + CapabilitiesFunc func() llm.ProviderCapabilities +} + +func (m *MockLLMClient) Generate(ctx context.Context, req llm.CompletionRequest) (llm.CompletionResponse, error) { + if m.GenerateFunc != nil { + return m.GenerateFunc(ctx, req) + } + return llm.CompletionResponse{ + Content: "default response", + StopReason: llm.StopReasonEndTurn, + }, nil +} + +func (m *MockLLMClient) Stream(ctx context.Context, req llm.CompletionRequest) iter.Seq2[llm.StreamChunk, error] { + if m.StreamFunc != nil { + return m.StreamFunc(ctx, req) + } + return func(yield func(llm.StreamChunk, error) bool) { + yield(llm.StreamChunk{ + Delta: "default response", + }, nil) + } +} + +func (m *MockLLMClient) Name() string { + if m.NameFunc != nil { + return m.NameFunc() + } + return "mock" +} + +func (m *MockLLMClient) Capabilities() llm.ProviderCapabilities { + if m.CapabilitiesFunc != nil { + return m.CapabilitiesFunc() + } + return llm.ProviderCapabilities{ + SupportsTools: true, + SupportsVision: false, + SupportsJSON: true, + MaxContextWindow: 128000, + } +} + +// New returns a MockLLMClient with default behaviors. +func New() *MockLLMClient { + return &MockLLMClient{} +} + +// NewWithGenerate returns a MockLLMClient that returns the given response. +func NewWithGenerate(resp llm.CompletionResponse) *MockLLMClient { + return &MockLLMClient{ + GenerateFunc: func(ctx context.Context, req llm.CompletionRequest) (llm.CompletionResponse, error) { + return resp, nil + }, + } +} + +// NewWithStream returns a MockLLMClient that streams the given chunks. +func NewWithStream(chunks []llm.StreamChunk) *MockLLMClient { + return &MockLLMClient{ + StreamFunc: func(ctx context.Context, req llm.CompletionRequest) iter.Seq2[llm.StreamChunk, error] { + return func(yield func(llm.StreamChunk, error) bool) { + for _, c := range chunks { + if !yield(c, nil) { + return + } + } + } + }, + } +} + +// NewWithMatch returns a MockLLMClient that matches input messages against patterns. +func NewWithMatch(responses []MatchResponse) *MockLLMClient { + return &MockLLMClient{ + GenerateFunc: func(ctx context.Context, req llm.CompletionRequest) (llm.CompletionResponse, error) { + for _, r := range responses { + if r.Match == "*" || matchesAny(r.Match, req.Messages) { + return llm.CompletionResponse{ + Content: r.Response, + StopReason: llm.StopReasonEndTurn, + }, nil + } + } + return llm.CompletionResponse{ + Content: "no match", + StopReason: llm.StopReasonEndTurn, + }, nil + }, + } +} + +// MatchResponse pairs a pattern with a response string. +type MatchResponse struct { + Match string // exact string or "*" for wildcard + Response string +} + +func matchesAny(pattern string, msgs []llm.Message) bool { + for _, m := range msgs { + if m.Role == llm.RoleUser && strings.Contains(m.Content, pattern) { + return true + } + } + return false +} diff --git a/pkg/llm/providers/llamacpp/client.go b/pkg/llm/providers/llamacpp/client.go new file mode 100644 index 0000000..314803e --- /dev/null +++ b/pkg/llm/providers/llamacpp/client.go @@ -0,0 +1,322 @@ +package llamacpp + +import ( + "bufio" + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "iter" + "strings" + + "github.com/VictorVargas/rony-llm-agent/pkg/llm" +) + +// Config holds the settings needed to create a llama.cpp client. +type Config struct { + BaseURL string // defaults to http://localhost:8080/v1 + Model string + Timeout int // request timeout in seconds (0 = default) + TopK int // top-k sampling (0 = default) + TopP float32 + Temperature float32 +} + +// Client implements llm.LLMClient for llama.cpp. +type Client struct { + baseURL string + http *http.Client +} + +// New returns a new llama.cpp client. +func New(cfg Config) (*Client, error) { + baseURL := cfg.BaseURL + if baseURL == "" { + baseURL = "http://localhost:8080/v1" + } + + return &Client{ + baseURL: baseURL, + http: http.DefaultClient, + }, nil +} + +func (c *Client) Generate(ctx context.Context, req llm.CompletionRequest) (llm.CompletionResponse, error) { + endpoint := c.baseURL + "/chat/completions" + + payload, err := c.buildRequest(req, false) + if err != nil { + return llm.CompletionResponse{}, fmt.Errorf("building request: %w", err) + } + + httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, payload) + if err != nil { + return llm.CompletionResponse{}, fmt.Errorf("creating request: %w", err) + } + httpReq.Header.Set("Content-Type", "application/json") + + resp, err := c.http.Do(httpReq) + if err != nil { + return llm.CompletionResponse{}, fmt.Errorf("request failed: %w", err) + } + defer resp.Body.Close() + + if resp.StatusCode < 200 || resp.StatusCode >= 300 { + body, _ := io.ReadAll(resp.Body) + return llm.CompletionResponse{}, fmt.Errorf("API error %d: %s", resp.StatusCode, string(body)) + } + + var apiResp llamaChatResponse + if err := json.NewDecoder(resp.Body).Decode(&apiResp); err != nil { + return llm.CompletionResponse{}, fmt.Errorf("decoding response: %w", err) + } + + return c.toResponse(apiResp), nil +} + +func (c *Client) Stream(ctx context.Context, req llm.CompletionRequest) iter.Seq2[llm.StreamChunk, error] { + return func(yield func(llm.StreamChunk, error) bool) { + endpoint := c.baseURL + "/chat/completions" + + payload, err := c.buildRequest(req, true) + if err != nil { + yield(llm.StreamChunk{}, fmt.Errorf("building request: %w", err)) + return + } + + httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, payload) + if err != nil { + yield(llm.StreamChunk{}, fmt.Errorf("creating request: %w", err)) + return + } + httpReq.Header.Set("Content-Type", "application/json") + + resp, err := c.http.Do(httpReq) + if err != nil { + yield(llm.StreamChunk{}, fmt.Errorf("request failed: %w", err)) + return + } + defer resp.Body.Close() + + if resp.StatusCode < 200 || resp.StatusCode >= 300 { + body, _ := io.ReadAll(resp.Body) + yield(llm.StreamChunk{}, fmt.Errorf("API error %d: %s", resp.StatusCode, string(body))) + return + } + + scanner := bufio.NewScanner(resp.Body) + for scanner.Scan() { + line := scanner.Text() + if !strings.HasPrefix(line, "data: ") { + continue + } + data := strings.TrimPrefix(line, "data: ") + if data == "[DONE]" { + return + } + + var event llamaStreamEvent + if err := json.Unmarshal([]byte(data), &event); err != nil { + yield(llm.StreamChunk{}, fmt.Errorf("decoding event: %w", err)) + return + } + + for _, choice := range event.Choices { + chunk := llm.StreamChunk{ + Delta: choice.Delta.Content, + } + if choice.FinishReason != "" { + chunk.FinishReason = choice.FinishReason + } + if !yield(chunk, nil) { + return + } + } + } + if err := scanner.Err(); err != nil { + yield(llm.StreamChunk{}, fmt.Errorf("stream error: %w", err)) + } + } +} + +func (c *Client) Name() string { + return "llama.cpp" +} + +func (c *Client) Capabilities() llm.ProviderCapabilities { + return llm.ProviderCapabilities{ + SupportsTools: true, + SupportsVision: false, + SupportsJSON: true, + MaxContextWindow: 32768, + } +} + +// buildRequest converts an llm.CompletionRequest to the llama.cpp API format. +func (c *Client) buildRequest(req llm.CompletionRequest, stream bool) (io.Reader, error) { + messages := make([]llamaMessage, len(req.Messages)) + for i, m := range req.Messages { + messages[i] = llamaMessage{ + Role: string(m.Role), + Content: m.Content, + } + } + + var tools []llamaTool + for _, t := range req.Tools { + var tool llamaTool + if err := json.Unmarshal(t, &tool); err != nil { + return nil, fmt.Errorf("parsing tool: %w", err) + } + tools = append(tools, tool) + } + + openReq := llamaChatRequest{ + Model: req.Model, + Messages: messages, + Stream: stream, + } + if len(tools) > 0 { + openReq.Tools = tools + } + if req.ToolChoice != nil { + openReq.ToolChoice = req.ToolChoice + } + if req.Temperature != nil { + tmp := *req.Temperature + openReq.Temperature = tmp + } + if req.MaxTokens != nil { + tmp := *req.MaxTokens + openReq.MaxTokens = tmp + } + if len(req.Stop) > 0 { + openReq.Stop = req.Stop + } + + data, err := json.Marshal(openReq) + if err != nil { + return nil, fmt.Errorf("marshaling request: %w", err) + } + return strings.NewReader(string(data)), nil +} + +// toResponse converts a llama.cpp API response to our CompletionResponse. +func (c *Client) toResponse(resp llamaChatResponse) llm.CompletionResponse { + choice := resp.Choices[0] + result := llm.CompletionResponse{ + ID: resp.ID, + Model: resp.Model, + Content: choice.Message.Content, + StopReason: choice.FinishReason, + } + + for _, tc := range choice.Message.ToolCalls { + result.ToolCalls = append(result.ToolCalls, llm.ToolCall{ + ID: tc.ID, + Name: tc.Function.Name, + Arguments: json.RawMessage(tc.Function.Arguments), + }) + } + + result.Usage = llm.TokenUsage{ + InputTokens: resp.Usage.PromptTokens, + OutputTokens: resp.Usage.CompletionTokens, + TotalTokens: resp.Usage.TotalTokens, + } + + return result +} + +// llama.cpp API types + +type llamaChatRequest struct { + Model string `json:"model"` + Messages []llamaMessage `json:"messages"` + Tools []llamaTool `json:"tools,omitempty"` + ToolChoice interface{} `json:"tool_choice,omitempty"` + Temperature float32 `json:"temperature,omitempty"` + MaxTokens int `json:"max_tokens,omitempty"` + TopK int `json:"top_k,omitempty"` + TopP float32 `json:"top_p,omitempty"` + Stop []string `json:"stop,omitempty"` + Stream bool `json:"stream"` +} + +type llamaMessage struct { + Role string `json:"role"` + Content string `json:"content"` +} + +type llamaTool struct { + Type string `json:"type"` + Function json.RawMessage `json:"function"` +} + +type llamaChatResponse struct { + ID string `json:"id"` + Model string `json:"model"` + Choices []llamaChoice `json:"choices"` + Usage llamaUsage `json:"usage"` +} + +type llamaChoice struct { + Index int `json:"index"` + Message llamaMessageResult `json:"message"` + FinishReason string `json:"finish_reason"` +} + +type llamaMessageResult struct { + Role string `json:"role"` + Content string `json:"content"` + ToolCalls []llamaToolCall `json:"tool_calls"` +} + +type llamaToolCall struct { + ID string `json:"id"` + Type string `json:"type"` + Function llamaFunction `json:"function"` +} + +type llamaFunction struct { + Name string `json:"name"` + Arguments string `json:"arguments"` +} + +type llamaUsage struct { + PromptTokens int `json:"prompt_tokens"` + CompletionTokens int `json:"completion_tokens"` + TotalTokens int `json:"total_tokens"` +} + +// Stream event types + +type llamaStreamEvent struct { + ID string `json:"id"` + Choices []llamaStreamChoice `json:"choices"` +} + +type llamaStreamChoice struct { + Index int `json:"index"` + Delta llamaStreamDelta `json:"delta"` + FinishReason string `json:"finish_reason"` +} + +type llamaStreamDelta struct { + Content string `json:"content"` + Role string `json:"role"` + ToolCalls []llamaStreamToolCall `json:"tool_calls"` +} + +type llamaStreamToolCall struct { + Index int `json:"index"` + ID string `json:"id"` + Type string `json:"type"` + Function llamaStreamFunction `json:"function"` +} + +type llamaStreamFunction struct { + Name string `json:"name"` + Arguments string `json:"arguments"` +} diff --git a/pkg/llm/providers/llamacpp/client_test.go b/pkg/llm/providers/llamacpp/client_test.go new file mode 100644 index 0000000..7b98377 --- /dev/null +++ b/pkg/llm/providers/llamacpp/client_test.go @@ -0,0 +1,232 @@ +package llamacpp + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + + "github.com/VictorVargas/rony-llm-agent/pkg/llm" +) + +func TestClient_Name(t *testing.T) { + client := &Client{} + if client.Name() != "llama.cpp" { + t.Errorf("expected 'llama.cpp', got %q", client.Name()) + } +} + +func TestClient_Capabilities(t *testing.T) { + client := &Client{} + caps := client.Capabilities() + if !caps.SupportsTools { + t.Error("expected SupportsTools to be true") + } + if caps.SupportsVision { + t.Error("expected SupportsVision to be false") + } +} + +func TestClient_Generate(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(llamaChatResponse{ + ID: "llama-123", + Model: "llama3", + Choices: []llamaChoice{ + { + Index: 0, + FinishReason: "stop", + Message: llamaMessageResult{ + Role: "assistant", + Content: "Hello from llama.cpp!", + }, + }, + }, + Usage: llamaUsage{ + PromptTokens: 10, + CompletionTokens: 5, + TotalTokens: 15, + }, + }) + })) + defer server.Close() + + client, err := New(Config{ + BaseURL: server.URL + "/v1", + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + resp, err := client.Generate(context.Background(), llm.CompletionRequest{ + Model: "llama3", + Messages: []llm.Message{ + {Role: llm.RoleUser, Content: "Hi"}, + }, + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if resp.Content != "Hello from llama.cpp!" { + t.Errorf("expected 'Hello from llama.cpp!', got %q", resp.Content) + } + if resp.ID != "llama-123" { + t.Errorf("expected 'llama-123', got %q", resp.ID) + } + if resp.Usage.InputTokens != 10 { + t.Errorf("expected 10 input tokens, got %d", resp.Usage.InputTokens) + } +} + +func TestClient_Generate_Error(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusInternalServerError) + w.Write([]byte("server error")) + })) + defer server.Close() + + client, err := New(Config{ + BaseURL: server.URL + "/v1", + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + _, err = client.Generate(context.Background(), llm.CompletionRequest{}) + if err == nil { + t.Fatal("expected error, got nil") + } +} + +func TestClient_Generate_ToolCall(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(llamaChatResponse{ + ID: "llama-tool-1", + Model: "llama3", + Choices: []llamaChoice{ + { + Index: 0, + FinishReason: "tool_calls", + Message: llamaMessageResult{ + Role: "assistant", + Content: "", + ToolCalls: []llamaToolCall{ + { + ID: "call-1", + Type: "function", + Function: llamaFunction{ + Name: "calculate", + Arguments: `{"a":1,"b":2}`, + }, + }, + }, + }, + }, + }, + Usage: llamaUsage{ + PromptTokens: 20, + CompletionTokens: 10, + TotalTokens: 30, + }, + }) + })) + defer server.Close() + + client, err := New(Config{ + BaseURL: server.URL + "/v1", + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + resp, err := client.Generate(context.Background(), llm.CompletionRequest{ + Model: "llama3", + Messages: []llm.Message{ + {Role: llm.RoleUser, Content: "What is 1 + 2?"}, + }, + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if len(resp.ToolCalls) != 1 { + t.Fatalf("expected 1 tool call, got %d", len(resp.ToolCalls)) + } + if resp.ToolCalls[0].Name != "calculate" { + t.Errorf("expected 'calculate', got %q", resp.ToolCalls[0].Name) + } +} + +func TestClient_Stream(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + w.Write([]byte("data: {\"choices\":[{\"delta\":{\"content\":\"Hello\"},\"finish_reason\":null}]}\n")) + w.Write([]byte("data: {\"choices\":[{\"delta\":{\"content\":\" world\"},\"finish_reason\":null}]}\n")) + w.Write([]byte("data: {\"choices\":[{\"delta\":{\"content\":\"!\"},\"finish_reason\":\"stop\"}]}\n")) + w.Write([]byte("data: [DONE]\n")) + })) + defer server.Close() + + client, err := New(Config{ + BaseURL: server.URL + "/v1", + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + var chunks []llm.StreamChunk + stream := client.Stream(context.Background(), llm.CompletionRequest{}) + for chunk, err := range stream { + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + chunks = append(chunks, chunk) + } + + if len(chunks) != 3 { + t.Errorf("expected 3 chunks, got %d", len(chunks)) + } + if chunks[0].Delta != "Hello" { + t.Errorf("expected 'Hello', got %q", chunks[0].Delta) + } + if chunks[1].Delta != " world" { + t.Errorf("expected ' world', got %q", chunks[1].Delta) + } + if chunks[2].Delta != "!" { + t.Errorf("expected '!', got %q", chunks[2].Delta) + } +} + +func TestClient_Stream_FinishReason(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + w.Write([]byte("data: {\"choices\":[{\"delta\":{\"content\":\"done\"},\"finish_reason\":\"stop\"}]}\n")) + w.Write([]byte("data: [DONE]\n")) + })) + defer server.Close() + + client, err := New(Config{ + BaseURL: server.URL + "/v1", + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + var chunks []llm.StreamChunk + stream := client.Stream(context.Background(), llm.CompletionRequest{}) + for chunk, err := range stream { + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + chunks = append(chunks, chunk) + } + + if len(chunks) != 1 { + t.Fatalf("expected 1 chunk, got %d", len(chunks)) + } + if chunks[0].FinishReason != "stop" { + t.Errorf("expected 'stop' finish reason, got %q", chunks[0].FinishReason) + } +} diff --git a/pkg/llm/providers/openai/client.go b/pkg/llm/providers/openai/client.go new file mode 100644 index 0000000..d3967cf --- /dev/null +++ b/pkg/llm/providers/openai/client.go @@ -0,0 +1,327 @@ +package openai + +import ( + "bufio" + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "iter" + "strings" + + "github.com/VictorVargas/rony-llm-agent/pkg/llm" +) + +// Config holds the settings needed to create an OpenAI client. +type Config struct { + APIKey string + Model string + BaseURL string // defaults to https://api.openai.com/v1 +} + +// Client implements llm.LLMClient for OpenAI. +type Client struct { + apiKey string + baseURL string + http *http.Client +} + +// New returns a new OpenAI client. +func New(cfg Config) (*Client, error) { + if cfg.APIKey == "" { + return nil, fmt.Errorf("openai: API key is required") + } + + baseURL := cfg.BaseURL + if baseURL == "" { + baseURL = "https://api.openai.com/v1" + } + + return &Client{ + apiKey: cfg.APIKey, + baseURL: baseURL, + http: http.DefaultClient, + }, nil +} + +func (c *Client) Generate(ctx context.Context, req llm.CompletionRequest) (llm.CompletionResponse, error) { + endpoint := c.baseURL + "/chat/completions" + + payload, err := c.buildRequest(req) + if err != nil { + return llm.CompletionResponse{}, fmt.Errorf("building request: %w", err) + } + + httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, payload) + if err != nil { + return llm.CompletionResponse{}, fmt.Errorf("creating request: %w", err) + } + httpReq.Header.Set("Authorization", "Bearer "+c.apiKey) + httpReq.Header.Set("Content-Type", "application/json") + + resp, err := c.http.Do(httpReq) + if err != nil { + return llm.CompletionResponse{}, fmt.Errorf("request failed: %w", err) + } + defer resp.Body.Close() + + if resp.StatusCode < 200 || resp.StatusCode >= 300 { + body, _ := io.ReadAll(resp.Body) + return llm.CompletionResponse{}, fmt.Errorf("API error %d: %s", resp.StatusCode, string(body)) + } + + var apiResp openaiChatResponse + if err := json.NewDecoder(resp.Body).Decode(&apiResp); err != nil { + return llm.CompletionResponse{}, fmt.Errorf("decoding response: %w", err) + } + + return c.toResponse(apiResp), nil +} + +func (c *Client) Stream(ctx context.Context, req llm.CompletionRequest) iter.Seq2[llm.StreamChunk, error] { + return func(yield func(llm.StreamChunk, error) bool) { + endpoint := c.baseURL + "/chat/completions" + + payload, err := c.buildRequest(req) + if err != nil { + yield(llm.StreamChunk{}, fmt.Errorf("building request: %w", err)) + return + } + + httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, payload) + if err != nil { + yield(llm.StreamChunk{}, fmt.Errorf("creating request: %w", err)) + return + } + httpReq.Header.Set("Authorization", "Bearer "+c.apiKey) + httpReq.Header.Set("Content-Type", "application/json") + httpReq.Header.Set("Accept", "text/event-stream") + + resp, err := c.http.Do(httpReq) + if err != nil { + yield(llm.StreamChunk{}, fmt.Errorf("request failed: %w", err)) + return + } + defer resp.Body.Close() + + if resp.StatusCode < 200 || resp.StatusCode >= 300 { + body, _ := io.ReadAll(resp.Body) + yield(llm.StreamChunk{}, fmt.Errorf("API error %d: %s", resp.StatusCode, string(body))) + return + } + + scanner := bufio.NewScanner(resp.Body) + for scanner.Scan() { + line := scanner.Text() + if !strings.HasPrefix(line, "data: ") { + continue + } + data := strings.TrimPrefix(line, "data: ") + if data == "[DONE]" { + return + } + + var event openaiStreamEvent + if err := json.Unmarshal([]byte(data), &event); err != nil { + yield(llm.StreamChunk{}, fmt.Errorf("decoding event: %w", err)) + return + } + + for _, choice := range event.Choices { + chunk := llm.StreamChunk{ + Delta: choice.Delta.Content, + } + if choice.FinishReason != "" { + chunk.FinishReason = choice.FinishReason + } + if !yield(chunk, nil) { + return + } + } + } + if err := scanner.Err(); err != nil { + yield(llm.StreamChunk{}, fmt.Errorf("stream error: %w", err)) + } + } +} + +func (c *Client) Name() string { + return "openai" +} + +func (c *Client) Capabilities() llm.ProviderCapabilities { + return llm.ProviderCapabilities{ + SupportsTools: true, + SupportsVision: true, + SupportsJSON: true, + MaxContextWindow: 128000, + } +} + +// buildRequest converts an llm.CompletionRequest to the OpenAI API format. +func (c *Client) buildRequest(req llm.CompletionRequest) (io.Reader, error) { + // Convert messages to OpenAI format + messages := make([]openaiMessage, len(req.Messages)) + for i, m := range req.Messages { + messages[i] = openaiMessage{ + Role: string(m.Role), + Content: m.Content, + } + } + + tools := make([]openaiTool, len(req.Tools)) + for i, t := range req.Tools { + var tool openaiTool + if err := json.Unmarshal(t, &tool); err != nil { + return nil, fmt.Errorf("parsing tool %d: %w", i, err) + } + tools[i] = tool + } + + openaiReq := openaiChatRequest{ + Model: req.Model, + Messages: messages, + Stream: false, + } + if len(tools) > 0 { + openaiReq.Tools = tools + } + if req.ToolChoice != nil { + openaiReq.ToolChoice = req.ToolChoice + } + if req.Temperature != nil { + tmp := *req.Temperature + openaiReq.Temperature = &tmp + } + if req.MaxTokens != nil { + max := *req.MaxTokens + openaiReq.MaxTokens = &max + } + if len(req.Stop) > 0 { + openaiReq.Stop = req.Stop + } + + data, err := json.Marshal(openaiReq) + if err != nil { + return nil, fmt.Errorf("marshaling request: %w", err) + } + return strings.NewReader(string(data)), nil +} + +// toResponse converts an OpenAI API response to our CompletionResponse. +func (c *Client) toResponse(resp openaiChatResponse) llm.CompletionResponse { + choice := resp.Choices[0] + result := llm.CompletionResponse{ + ID: resp.ID, + Model: resp.Model, + Content: choice.Message.Content, + StopReason: choice.FinishReason, + } + + for _, tc := range choice.Message.ToolCalls { + result.ToolCalls = append(result.ToolCalls, llm.ToolCall{ + ID: tc.ID, + Name: tc.Function.Name, + Arguments: json.RawMessage(tc.Function.Arguments), + }) + } + + result.Usage = llm.TokenUsage{ + InputTokens: resp.Usage.PromptTokens, + OutputTokens: resp.Usage.CompletionTokens, + TotalTokens: resp.Usage.TotalTokens, + } + + return result +} + +// OpenAI API types + +type openaiChatRequest struct { + Model string `json:"model"` + Messages []openaiMessage `json:"messages"` + Tools []openaiTool `json:"tools,omitempty"` + ToolChoice interface{} `json:"tool_choice,omitempty"` + Temperature *float32 `json:"temperature,omitempty"` + MaxTokens *int `json:"max_tokens,omitempty"` + Stop []string `json:"stop,omitempty"` + Stream bool `json:"stream"` +} + +type openaiMessage struct { + Role string `json:"role"` + Content string `json:"content"` +} + +type openaiTool struct { + Type string `json:"type"` + Function json.RawMessage `json:"function"` +} + +type openaiChatResponse struct { + ID string `json:"id"` + Model string `json:"model"` + Choices []openaiChoice `json:"choices"` + Usage openaiUsage `json:"usage"` +} + +type openaiChoice struct { + Index int `json:"index"` + Message openaiMessageResult `json:"message"` + FinishReason string `json:"finish_reason"` +} + +type openaiMessageResult struct { + Role string `json:"role"` + Content string `json:"content"` + ToolCalls []openaiToolCall `json:"tool_calls"` +} + +type openaiToolCall struct { + ID string `json:"id"` + Type string `json:"type"` + Function openaiFunction `json:"function"` +} + +type openaiFunction struct { + Name string `json:"name"` + Arguments string `json:"arguments"` +} + +type openaiUsage struct { + PromptTokens int `json:"prompt_tokens"` + CompletionTokens int `json:"completion_tokens"` + TotalTokens int `json:"total_tokens"` +} + +// Stream event types + +type openaiStreamEvent struct { + ID string `json:"id"` + Choices []openaiStreamChoice `json:"choices"` +} + +type openaiStreamChoice struct { + Index int `json:"index"` + Delta openaiStreamDelta `json:"delta"` + FinishReason string `json:"finish_reason"` +} + +type openaiStreamDelta struct { + Content string `json:"content"` + Role string `json:"role"` + ToolCalls []openaiStreamToolCall `json:"tool_calls"` +} + +type openaiStreamToolCall struct { + Index int `json:"index"` + ID string `json:"id"` + Type string `json:"type"` + Function openaiStreamFunction `json:"function"` +} + +type openaiStreamFunction struct { + Name string `json:"name"` + Arguments string `json:"arguments"` +} diff --git a/pkg/llm/providers/openai/client_test.go b/pkg/llm/providers/openai/client_test.go new file mode 100644 index 0000000..dd838f6 --- /dev/null +++ b/pkg/llm/providers/openai/client_test.go @@ -0,0 +1,141 @@ +package openai + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + + "github.com/VictorVargas/rony-llm-agent/pkg/llm" +) + +func TestClient_Name(t *testing.T) { + client := &Client{} + if client.Name() != "openai" { + t.Errorf("expected 'openai', got %q", client.Name()) + } +} + +func TestClient_Capabilities(t *testing.T) { + client := &Client{} + caps := client.Capabilities() + if !caps.SupportsTools { + t.Error("expected SupportsTools to be true") + } + if !caps.SupportsVision { + t.Error("expected SupportsVision to be true") + } +} + +func TestClient_Generate(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + json.NewEncoder(w).Encode(openaiChatResponse{ + ID: "test-123", + Model: "gpt-4", + Choices: []openaiChoice{ + { + Index: 0, + FinishReason: "stop", + Message: openaiMessageResult{ + Role: "assistant", + Content: "Hello!", + }, + }, + }, + Usage: openaiUsage{ + PromptTokens: 10, + CompletionTokens: 5, + TotalTokens: 15, + }, + }) + })) + defer server.Close() + + client, err := New(Config{ + APIKey: "test-key", + BaseURL: server.URL, + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + resp, err := client.Generate(context.Background(), llm.CompletionRequest{ + Model: "gpt-4", + Messages: []llm.Message{ + {Role: llm.RoleUser, Content: "Hi"}, + }, + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if resp.Content != "Hello!" { + t.Errorf("expected 'Hello!', got %q", resp.Content) + } + if resp.ID != "test-123" { + t.Errorf("expected 'test-123', got %q", resp.ID) + } + if resp.Usage.InputTokens != 10 { + t.Errorf("expected 10 input tokens, got %d", resp.Usage.InputTokens) + } +} + +func TestClient_Generate_Error(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusUnauthorized) + w.Write([]byte("invalid api key")) + })) + defer server.Close() + + client, err := New(Config{ + APIKey: "bad-key", + BaseURL: server.URL, + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + _, err = client.Generate(context.Background(), llm.CompletionRequest{}) + if err == nil { + t.Fatal("expected error, got nil") + } +} + +func TestClient_Stream(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + w.Write([]byte("data: {\"choices\":[{\"delta\":{\"content\":\"Hello\"},\"finish_reason\":null}]}\n")) + w.Write([]byte("data: {\"choices\":[{\"delta\":{\"content\":\" world\"},\"finish_reason\":null}]}\n")) + w.Write([]byte("data: [DONE]\n")) + })) + defer server.Close() + + client, err := New(Config{ + APIKey: "test-key", + BaseURL: server.URL + "/v1", + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + var chunks []llm.StreamChunk + stream := client.Stream(context.Background(), llm.CompletionRequest{}) + for chunk, err := range stream { + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + chunks = append(chunks, chunk) + } + + if len(chunks) != 2 { + t.Errorf("expected 2 chunks, got %d", len(chunks)) + } + if chunks[0].Delta != "Hello" { + t.Errorf("expected 'Hello', got %q", chunks[0].Delta) + } + if chunks[1].Delta != " world" { + t.Errorf("expected ' world', got %q", chunks[1].Delta) + } +} diff --git a/pkg/llm/types.go b/pkg/llm/types.go new file mode 100644 index 0000000..f7d9fa8 --- /dev/null +++ b/pkg/llm/types.go @@ -0,0 +1,120 @@ +package llm + +import ( + "context" + "encoding/json" + "iter" +) + +// Role represents the role of a message participant. +type Role string + +const ( + RoleSystem Role = "system" + RoleUser Role = "user" + RoleAssistant Role = "assistant" + RoleTool Role = "tool" +) + +// Message is a single message in a conversation. +type Message struct { + Role Role `json:"role"` + Content string `json:"content"` + ToolCallID string `json:"tool_call_id,omitempty"` + Name string `json:"name,omitempty"` +} + +// ProviderCapabilities describes what a model supports. +type ProviderCapabilities struct { + SupportsTools bool `json:"supports_tools"` + SupportsVision bool `json:"supports_vision"` + SupportsJSON bool `json:"supports_json"` + MaxContextWindow int `json:"max_context_window"` +} + +// TokenUsage tracks token consumption for a response. +type TokenUsage struct { + InputTokens int `json:"input_tokens"` + OutputTokens int `json:"output_tokens"` + TotalTokens int `json:"total_tokens"` +} + +// ToolChoice controls how the model chooses which tools to call. +type ToolChoice string + +const ( + ToolChoiceAuto ToolChoice = "auto" + ToolChoiceNone ToolChoice = "none" + ToolChoiceRequired ToolChoice = "required" +) + +// ToolRef is used when ToolChoice must invoke a specific tool by name. +type ToolRef struct { + Name string `json:"name"` +} + +func (t *ToolRef) MarshalJSON() ([]byte, error) { + type Alias ToolRef + return json.Marshal(&struct { + *Alias + Type string `json:"type"` + }{ + Alias: (*Alias)(t), + Type: "function", + }) +} + +// CompletionRequest is sent to an LLM provider. +type CompletionRequest struct { + Model string `json:"model"` + Messages []Message `json:"messages"` + Tools []json.RawMessage `json:"tools,omitempty"` + ToolChoice interface{} `json:"tool_choice,omitempty"` // ToolChoice, ToolRef, or null + Temperature *float32 `json:"temperature,omitempty"` + MaxTokens *int `json:"max_tokens,omitempty"` + Stop []string `json:"stop,omitempty"` + Metadata map[string]string `json:"metadata,omitempty"` +} + +// CompletionResponse is returned from an LLM provider. +type CompletionResponse struct { + ID string `json:"id"` + Model string `json:"model"` + Content string `json:"content"` + ToolCalls []ToolCall `json:"tool_calls,omitempty"` + StopReason string `json:"stop_reason"` + Usage TokenUsage `json:"usage,omitempty"` +} + +// StopReason values. +const ( + StopReasonEndTurn = "end_turn" + StopReasonToolUse = "tool_use" + StopReasonMaxTokens = "max_tokens" + StopReasonStopSeq = "stop_sequence" +) + +// ToolCall represents a function invocation requested by the model. +type ToolCall struct { + ID string `json:"id"` + Name string `json:"name"` + Arguments json.RawMessage `json:"arguments"` + Thought string `json:"thought,omitempty"` +} + +// StreamChunk is emitted by the iterator returned from Stream(). +type StreamChunk struct { + Delta string `json:"delta"` + ToolCalls []ToolCall `json:"tool_calls,omitempty"` + FinishReason string `json:"finish_reason,omitempty"` + Usage TokenUsage `json:"usage,omitempty"` // only present on final chunk +} + +// LLMClient is the interface that every provider implements. +// Products must never call a provider SDK directly; they go through this interface. +type LLMClient interface { + Generate(ctx context.Context, req CompletionRequest) (CompletionResponse, error) + Stream(ctx context.Context, req CompletionRequest) iter.Seq2[StreamChunk, error] + Name() string + Capabilities() ProviderCapabilities +} From 4b39f52081da98cbd8f64df4bf8ccac3b9d5a60c Mon Sep 17 00:00:00 2001 From: Victor Vargas Date: Tue, 30 Jun 2026 23:53:26 -0700 Subject: [PATCH 04/11] feat(agent): add Agent loop with iteration control, tool handling, and streaming support --- pkg/agent/integration_test.go | 157 ++++++++++ pkg/agent/loop.go | 239 +++++++++++++++ pkg/agent/loop_test.go | 527 ++++++++++++++++++++++++++++++++++ 3 files changed, 923 insertions(+) create mode 100644 pkg/agent/integration_test.go create mode 100644 pkg/agent/loop.go create mode 100644 pkg/agent/loop_test.go diff --git a/pkg/agent/integration_test.go b/pkg/agent/integration_test.go new file mode 100644 index 0000000..98cfd49 --- /dev/null +++ b/pkg/agent/integration_test.go @@ -0,0 +1,157 @@ +package agent_test + +import ( + "context" + "encoding/json" + "os" + "testing" + + "github.com/VictorVargas/rony-llm-agent/pkg/agent" + "github.com/VictorVargas/rony-llm-agent/pkg/llm" + "github.com/VictorVargas/rony-llm-agent/pkg/llm/providers/llamacpp" + "github.com/VictorVargas/rony-llm-agent/pkg/persona" + "github.com/VictorVargas/rony-llm-agent/pkg/tools" +) + +// TestIntegration_LlamaCPP_Generate is an integration test that requires llama.cpp running on localhost:8080. +// Run with: go test ./pkg/agent/ -run TestIntegration_LlamaCPP_Generate -tags=integration +func TestIntegration_LlamaCPP_Generate(t *testing.T) { + client, err := llamacpp.New(llamacpp.Config{ + BaseURL: "http://localhost:8080/v1", + }) + if err != nil { + t.Fatalf("failed to create client: %v", err) + } + + resp, err := client.Generate(context.Background(), llm.CompletionRequest{ + Messages: []llm.Message{ + {Role: llm.RoleUser, Content: "What is 2+2? Answer with just the number."}, + }, + }) + if err != nil { + t.Fatalf("generate failed: %v", err) + } + + if resp.Content == "" { + t.Fatal("expected non-empty response") + } +} + +// TestIntegration_LlamaCPP_Stream is an integration test that requires llama.cpp running on localhost:8080. +func TestIntegration_LlamaCPP_Stream(t *testing.T) { + client, err := llamacpp.New(llamacpp.Config{ + BaseURL: "http://localhost:8080/v1", + }) + if err != nil { + t.Fatalf("failed to create client: %v", err) + } + + stream := client.Stream(context.Background(), llm.CompletionRequest{ + Messages: []llm.Message{ + {Role: llm.RoleUser, Content: "Say hello in 5 words."}, + }, + }) + + var chunks []llm.StreamChunk + for chunk, err := range stream { + if err != nil { + t.Fatalf("stream error: %v", err) + } + chunks = append(chunks, chunk) + } + + if len(chunks) == 0 { + t.Fatal("expected at least one chunk") + } +} + +// TestIntegration_AgentLoop_Generate is an integration test for the agent loop with llama.cpp. +func TestIntegration_AgentLoop_Generate(t *testing.T) { + client, err := llamacpp.New(llamacpp.Config{ + BaseURL: "http://localhost:8080/v1", + }) + if err != nil { + t.Fatalf("failed to create client: %v", err) + } + + registry := tools.NewRegistry() + registry.Register(tools.Tool{ + Name: "add_numbers", + Description: "Add two numbers together", + InputSchema: json.RawMessage(`{"type": "function", "function": {"name": "add_numbers", "description": "Add two numbers together", "parameters": {"type": "object", "properties": {"a": {"type": "number"}, "b": {"type": "number"}}}}}`), + Handler: func(ctx context.Context, args json.RawMessage) (tools.ToolResult, error) { + var input struct { + A float64 `json:"a"` + B float64 `json:"b"` + } + if err := json.Unmarshal(args, &input); err != nil { + return tools.ToolResult{IsError: true}, err + } + return tools.ToolResult{Content: "42"}, nil + }, + Permission: tools.Allow, + }) + + loop := agent.New(agent.Config{ + LLM: client, + Persona: persona.DefaultPersona(), + Tools: registry, + Sandbox: &mockSandbox{}, + MaxIters: 3, + }) + + resp, err := loop.Run(context.Background(), "What is 20+22? Use the add_numbers tool.") + if err != nil { + t.Fatalf("run failed: %v", err) + } + + if resp.Content == "" { + t.Fatal("expected non-empty response") + } +} + +// TestIntegration_AgentLoop_Stream is an integration test for the agent loop with streaming. +func TestIntegration_AgentLoop_Stream(t *testing.T) { + client, err := llamacpp.New(llamacpp.Config{ + BaseURL: "http://localhost:8080/v1", + }) + if err != nil { + t.Fatalf("failed to create client: %v", err) + } + + loop := agent.New(agent.Config{ + LLM: client, + Persona: persona.DefaultPersona(), + Tools: tools.NewRegistry(), + MaxIters: 3, + }) + + stream := loop.RunStream(context.Background(), "Say something interesting.") + + var chunks []llm.StreamChunk + for chunk, err := range stream { + if err != nil { + t.Fatalf("stream error: %v", err) + } + chunks = append(chunks, chunk) + } + + if len(chunks) == 0 { + t.Fatal("expected at least one chunk") + } +} + +// mockSandbox is a simple sandbox that allows all calls. +type mockSandbox struct{} + +func (m *mockSandbox) ValidateToolCall(tool tools.Tool, call llm.ToolCall) error { + return nil +} + +func TestMain(m *testing.M) { + // Skip integration tests unless explicitly enabled + if os.Getenv("INTEGRATION_TESTS") != "1" { + return + } + os.Exit(m.Run()) +} diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go new file mode 100644 index 0000000..aa988b1 --- /dev/null +++ b/pkg/agent/loop.go @@ -0,0 +1,239 @@ +package agent + +import ( + "context" + "encoding/json" + "fmt" + "iter" + "time" + + "github.com/VictorVargas/rony-llm-agent/pkg/llm" + "github.com/VictorVargas/rony-llm-agent/pkg/persona" + "github.com/VictorVargas/rony-llm-agent/pkg/tools" +) + +// MaxIterations is the default maximum number of agent iterations. +const DefaultMaxIterations = 50 + +// MaxToolOutputBytes is the default maximum size for tool output in bytes. +const DefaultMaxToolOutputBytes = 50 * 1024 // 50KB + +// Sandbox defines the interface for filesystem sandbox operations. +type Sandbox interface { + ValidateToolCall(tool tools.Tool, call llm.ToolCall) error +} + +// Approver is called before executing tools with Ask permission. +// Return true to allow, false to deny. +type Approver func(tool tools.Tool, call llm.ToolCall) bool + +// OnIterationHook is called after each iteration. +type OnIterationHook func(Iteration) + +// Config holds the dependencies and settings for the agent loop. +type Config struct { + LLM llm.LLMClient + Persona persona.Persona + Tools tools.Registry + Sandbox Sandbox + MaxIters int + Approver Approver + OnIteration OnIterationHook + ToolTimeout time.Duration +} + +// Iteration represents a single cycle of the agent loop. +type Iteration struct { + Number int + ToolCalls []llm.ToolCall + ToolsUsed int + Duration time.Duration +} + +// Response is the final output of the agent loop. +type Response struct { + Content string + ToolCalls []llm.ToolCall + Iterations int + Duration time.Duration + TokenUsage llm.TokenUsage +} + +// Loop is the main agent loop that orchestrates LLM calls and tool execution. +type Loop struct { + cfg Config +} + +// New creates a new Loop with the given configuration. +func New(cfg Config) *Loop { + if cfg.MaxIters == 0 { + cfg.MaxIters = DefaultMaxIterations + } + return &Loop{cfg: cfg} +} + +// Run executes the agent loop and returns the final response. +func (l *Loop) Run(ctx context.Context, input string) (Response, error) { + start := time.Now() + + messages := l.buildInitialMessages(input) + var finalContent string + var allToolCalls []llm.ToolCall + var totalUsage llm.TokenUsage + iterations := 0 + + for iterations < l.cfg.MaxIters { + iterations++ + + resp, err := l.cfg.LLM.Generate(ctx, llm.CompletionRequest{ + Messages: messages, + Tools: l.getToolSchemas(), + }) + if err != nil { + return Response{}, fmt.Errorf("LLM generate failed: %w", err) + } + + totalUsage.InputTokens += resp.Usage.InputTokens + totalUsage.OutputTokens += resp.Usage.OutputTokens + totalUsage.TotalTokens += resp.Usage.TotalTokens + + if len(resp.ToolCalls) == 0 { + finalContent = resp.Content + break + } + + for _, call := range resp.ToolCalls { + result, err := l.executeTool(ctx, call) + if err != nil { + messages = append(messages, llm.Message{ + Role: llm.RoleTool, + Content: fmt.Sprintf("Error: %v", err), + }) + continue + } + messages = append(messages, llm.Message{ + Role: llm.RoleTool, + Content: result.Content, + }) + allToolCalls = append(allToolCalls, call) + } + } + + duration := time.Since(start) + if iterations >= l.cfg.MaxIters { + return Response{}, fmt.Errorf("max iterations (%d) reached", l.cfg.MaxIters) + } + + return Response{ + Content: finalContent, + ToolCalls: allToolCalls, + Iterations: iterations, + Duration: duration, + TokenUsage: totalUsage, + }, nil +} + +// RunStream executes the agent loop with streaming output. +func (l *Loop) RunStream(ctx context.Context, input string) iter.Seq2[llm.StreamChunk, error] { + return func(yield func(llm.StreamChunk, error) bool) { + messages := l.buildInitialMessages(input) + iterations := 0 + + for iterations < l.cfg.MaxIters { + iterations++ + + stream := l.cfg.LLM.Stream(ctx, llm.CompletionRequest{ + Messages: messages, + Tools: l.getToolSchemas(), + }) + + var hasToolCalls bool + for chunk, err := range stream { + if err != nil { + yield(llm.StreamChunk{}, err) + return + } + + if len(chunk.ToolCalls) > 0 { + hasToolCalls = true + for _, tc := range chunk.ToolCalls { + result, err := l.executeTool(ctx, tc) + if err != nil { + messages = append(messages, llm.Message{ + Role: llm.RoleTool, + Content: fmt.Sprintf("Error: %v", err), + }) + continue + } + messages = append(messages, llm.Message{ + Role: llm.RoleTool, + Content: result.Content, + }) + } + } + + if !hasToolCalls && chunk.Delta != "" { + if !yield(chunk, nil) { + return + } + } + } + + if !hasToolCalls { + return + } + } + + if iterations >= l.cfg.MaxIters { + yield(llm.StreamChunk{}, fmt.Errorf("max iterations (%d) reached", l.cfg.MaxIters)) + } + } +} + +func (l *Loop) buildInitialMessages(input string) []llm.Message { + systemPrompt := persona.AssembleSystemPrompt(l.cfg.Persona, "") + return []llm.Message{ + {Role: llm.RoleSystem, Content: systemPrompt}, + {Role: llm.RoleUser, Content: input}, + } +} + +func (l *Loop) getToolSchemas() []json.RawMessage { + var schemas []json.RawMessage + for _, tool := range l.cfg.Tools.List() { + schemas = append(schemas, tool.InputSchema) + } + return schemas +} + +func (l *Loop) executeTool(ctx context.Context, call llm.ToolCall) (tools.ToolResult, error) { + tool, found := l.cfg.Tools.Get(call.Name) + if !found { + return tools.ToolResult{IsError: true}, fmt.Errorf("tool %q not found", call.Name) + } + + if tool.Permission == tools.Ask && l.cfg.Approver != nil && !l.cfg.Approver(tool, call) { + return tools.ToolResult{Content: "Tool execution denied by user"}, nil + } + + if l.cfg.Sandbox != nil { + if err := l.cfg.Sandbox.ValidateToolCall(tool, call); err != nil { + return tools.ToolResult{IsError: true}, fmt.Errorf("sandbox violation: %w", err) + } + } + + if l.cfg.OnIteration != nil { + l.cfg.OnIteration(Iteration{ + Number: 1, + ToolCalls: []llm.ToolCall{call}, + ToolsUsed: 1, + }) + } + + result, err := tool.Handler(ctx, call.Arguments) + if err != nil { + return tools.ToolResult{IsError: true, Content: fmt.Sprintf("Tool error: %v", err)}, nil + } + + return result, nil +} diff --git a/pkg/agent/loop_test.go b/pkg/agent/loop_test.go new file mode 100644 index 0000000..a4d5c86 --- /dev/null +++ b/pkg/agent/loop_test.go @@ -0,0 +1,527 @@ +package agent + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "iter" + "testing" + "time" + + llm "github.com/VictorVargas/rony-llm-agent/pkg/llm" + "github.com/VictorVargas/rony-llm-agent/pkg/persona" + "github.com/VictorVargas/rony-llm-agent/pkg/tools" +) + +type mockLLM struct { + generateFunc func(ctx context.Context, req llm.CompletionRequest) (llm.CompletionResponse, error) + streamFunc func(ctx context.Context, req llm.CompletionRequest) iter.Seq2[llm.StreamChunk, error] +} + +func (m *mockLLM) Generate(ctx context.Context, req llm.CompletionRequest) (llm.CompletionResponse, error) { + return m.generateFunc(ctx, req) +} + +func (m *mockLLM) Stream(ctx context.Context, req llm.CompletionRequest) iter.Seq2[llm.StreamChunk, error] { + return m.streamFunc(ctx, req) +} + +func (m *mockLLM) Name() string { return "mock" } +func (m *mockLLM) Capabilities() llm.ProviderCapabilities { return llm.ProviderCapabilities{} } + +type mockSandbox struct { + validateFunc func(tool tools.Tool, call llm.ToolCall) error +} + +func (m *mockSandbox) ValidateToolCall(tool tools.Tool, call llm.ToolCall) error { + if m.validateFunc != nil { + return m.validateFunc(tool, call) + } + return nil +} + +func TestNew(t *testing.T) { + mockClient := &mockLLM{ + generateFunc: func(ctx context.Context, req llm.CompletionRequest) (llm.CompletionResponse, error) { + return llm.CompletionResponse{Content: "done"}, nil + }, + } + + loop := New(Config{ + LLM: mockClient, + Persona: persona.DefaultPersona(), + Tools: tools.NewRegistry(), + }) + + if loop == nil { + t.Fatal("expected non-nil loop") + } + if loop.cfg.MaxIters != DefaultMaxIterations { + t.Errorf("expected default max iters %d, got %d", DefaultMaxIterations, loop.cfg.MaxIters) + } +} + +func TestNew_CustomMaxIters(t *testing.T) { + mockClient := &mockLLM{ + generateFunc: func(ctx context.Context, req llm.CompletionRequest) (llm.CompletionResponse, error) { + return llm.CompletionResponse{Content: "done"}, nil + }, + } + + loop := New(Config{ + LLM: mockClient, + Persona: persona.DefaultPersona(), + Tools: tools.NewRegistry(), + MaxIters: 10, + }) + + if loop.cfg.MaxIters != 10 { + t.Errorf("expected max iters 10, got %d", loop.cfg.MaxIters) + } +} + +func TestRun_NoToolCalls(t *testing.T) { + mockClient := &mockLLM{ + generateFunc: func(ctx context.Context, req llm.CompletionRequest) (llm.CompletionResponse, error) { + return llm.CompletionResponse{ + Content: "I understand.", + Usage: llm.TokenUsage{ + InputTokens: 10, + OutputTokens: 5, + TotalTokens: 15, + }, + }, nil + }, + } + + loop := New(Config{ + LLM: mockClient, + Persona: persona.DefaultPersona(), + Tools: tools.NewRegistry(), + }) + + resp, err := loop.Run(context.Background(), "Hello") + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if resp.Content != "I understand." { + t.Errorf("expected 'I understand.', got %q", resp.Content) + } + if resp.Iterations != 1 { + t.Errorf("expected 1 iteration, got %d", resp.Iterations) + } + if resp.TokenUsage.InputTokens != 10 { + t.Errorf("expected 10 input tokens, got %d", resp.TokenUsage.InputTokens) + } +} + +func TestRun_ToolCalls(t *testing.T) { + registry := tools.NewRegistry() + registry.Register(tools.Tool{ + Name: "greet", + Description: "Greet someone", + InputSchema: json.RawMessage(`{"type":"object","properties":{"name":{"type":"string"}}}`), + Handler: func(ctx context.Context, args json.RawMessage) (tools.ToolResult, error) { + return tools.ToolResult{Content: "Hello!"}, nil + }, + Permission: tools.Allow, + }) + + callCount := 0 + mockClient := &mockLLM{ + generateFunc: func(ctx context.Context, req llm.CompletionRequest) (llm.CompletionResponse, error) { + callCount++ + if callCount == 1 { + return llm.CompletionResponse{ + ToolCalls: []llm.ToolCall{ + {ID: "call-1", Name: "greet", Arguments: json.RawMessage(`{"name":"World"}`)}, + }, + }, nil + } + return llm.CompletionResponse{ + Content: "Done!", + }, nil + }, + } + + loop := New(Config{ + LLM: mockClient, + Persona: persona.DefaultPersona(), + Tools: registry, + }) + + resp, err := loop.Run(context.Background(), "Say hi") + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if resp.Content != "Done!" { + t.Errorf("expected 'Done!', got %q", resp.Content) + } + if resp.Iterations != 2 { + t.Errorf("expected 2 iterations, got %d", resp.Iterations) + } + if len(resp.ToolCalls) != 1 { + t.Errorf("expected 1 tool call, got %d", len(resp.ToolCalls)) + } +} + +func TestRun_MaxIterations(t *testing.T) { + mockClient := &mockLLM{ + generateFunc: func(ctx context.Context, req llm.CompletionRequest) (llm.CompletionResponse, error) { + return llm.CompletionResponse{ + ToolCalls: []llm.ToolCall{{ID: "1", Name: "x", Arguments: json.RawMessage("{}")}}, + }, nil + }, + } + + registry := tools.NewRegistry() + registry.Register(tools.Tool{ + Name: "x", + Description: "x", + InputSchema: json.RawMessage(`{}`), + Handler: func(ctx context.Context, args json.RawMessage) (tools.ToolResult, error) { + return tools.ToolResult{Content: "ok"}, nil + }, + Permission: tools.Allow, + }) + + loop := New(Config{ + LLM: mockClient, + Persona: persona.DefaultPersona(), + Tools: registry, + MaxIters: 3, + }) + + _, err := loop.Run(context.Background(), "test") + if err == nil { + t.Fatal("expected error, got nil") + } + if err.Error() != "max iterations (3) reached" { + t.Errorf("expected 'max iterations (3) reached', got %q", err.Error()) + } +} + +func TestRun_ToolNotFound(t *testing.T) { + callCount := 0 + mockClient := &mockLLM{ + generateFunc: func(ctx context.Context, req llm.CompletionRequest) (llm.CompletionResponse, error) { + callCount++ + if callCount == 1 { + return llm.CompletionResponse{ + ToolCalls: []llm.ToolCall{{ID: "1", Name: "nonexistent", Arguments: json.RawMessage("{}")}}, + }, nil + } + return llm.CompletionResponse{Content: "done"}, nil + }, + } + + loop := New(Config{ + LLM: mockClient, + Persona: persona.DefaultPersona(), + Tools: tools.NewRegistry(), + }) + + resp, err := loop.Run(context.Background(), "test") + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if resp.Content != "done" { + t.Errorf("expected 'done', got %q", resp.Content) + } +} + +func TestRun_ApprovalDenied(t *testing.T) { + registry := tools.NewRegistry() + registry.Register(tools.Tool{ + Name: "dangerous", + Description: "Do something dangerous", + InputSchema: json.RawMessage(`{}`), + Handler: func(ctx context.Context, args json.RawMessage) (tools.ToolResult, error) { + return tools.ToolResult{Content: "executed"}, nil + }, + Permission: tools.Ask, + }) + + callCount := 0 + mockClient := &mockLLM{ + generateFunc: func(ctx context.Context, req llm.CompletionRequest) (llm.CompletionResponse, error) { + callCount++ + if callCount == 1 { + return llm.CompletionResponse{ + ToolCalls: []llm.ToolCall{{ID: "1", Name: "dangerous", Arguments: json.RawMessage("{}")}}, + }, nil + } + return llm.CompletionResponse{Content: "done"}, nil + }, + } + + loop := New(Config{ + LLM: mockClient, + Persona: persona.DefaultPersona(), + Tools: registry, + Approver: func(tool tools.Tool, call llm.ToolCall) bool { + return false // deny all + }, + }) + + resp, err := loop.Run(context.Background(), "test") + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if len(resp.ToolCalls) != 1 { + t.Errorf("expected 1 tool call, got %d", len(resp.ToolCalls)) + } +} + +func TestRun_SandboxViolation(t *testing.T) { + registry := tools.NewRegistry() + registry.Register(tools.Tool{ + Name: "restricted", + Description: "Restricted tool", + InputSchema: json.RawMessage(`{}`), + Handler: func(ctx context.Context, args json.RawMessage) (tools.ToolResult, error) { + return tools.ToolResult{Content: "executed"}, nil + }, + Permission: tools.Allow, + }) + + callCount := 0 + mockClient := &mockLLM{ + generateFunc: func(ctx context.Context, req llm.CompletionRequest) (llm.CompletionResponse, error) { + callCount++ + if callCount == 1 { + return llm.CompletionResponse{ + ToolCalls: []llm.ToolCall{{ID: "1", Name: "restricted", Arguments: json.RawMessage("{}")}}, + }, nil + } + return llm.CompletionResponse{Content: "done"}, nil + }, + } + + loop := New(Config{ + LLM: mockClient, + Persona: persona.DefaultPersona(), + Tools: registry, + Sandbox: &mockSandbox{ + validateFunc: func(tool tools.Tool, call llm.ToolCall) error { + return fmt.Errorf("path traversal detected") + }, + }, + }) + + resp, err := loop.Run(context.Background(), "test") + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if len(resp.ToolCalls) != 0 { + t.Errorf("expected 0 tool calls (sandbox rejected), got %d", len(resp.ToolCalls)) + } +} + +func TestRun_OnIterationHook(t *testing.T) { + registry := tools.NewRegistry() + registry.Register(tools.Tool{ + Name: "greet", + Description: "Greet", + InputSchema: json.RawMessage(`{}`), + Handler: func(ctx context.Context, args json.RawMessage) (tools.ToolResult, error) { + return tools.ToolResult{Content: "hello"}, nil + }, + Permission: tools.Allow, + }) + + callCount := 0 + mockClient := &mockLLM{ + generateFunc: func(ctx context.Context, req llm.CompletionRequest) (llm.CompletionResponse, error) { + callCount++ + if callCount == 1 { + return llm.CompletionResponse{ + ToolCalls: []llm.ToolCall{{ID: "1", Name: "greet", Arguments: json.RawMessage("{}")}}, + }, nil + } + return llm.CompletionResponse{Content: "done"}, nil + }, + } + + var iterations []Iteration + loop := New(Config{ + LLM: mockClient, + Persona: persona.DefaultPersona(), + Tools: registry, + OnIteration: func(iter Iteration) { + iterations = append(iterations, iter) + }, + }) + + _, err := loop.Run(context.Background(), "test") + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if len(iterations) != 1 { + t.Errorf("expected 1 iteration hook call, got %d", len(iterations)) + } + if iterations[0].ToolsUsed != 1 { + t.Errorf("expected 1 tool used, got %d", iterations[0].ToolsUsed) + } +} + +func TestRun_Stream_NoToolCalls(t *testing.T) { + mockClient := &mockLLM{ + streamFunc: func(ctx context.Context, req llm.CompletionRequest) iter.Seq2[llm.StreamChunk, error] { + return func(yield func(llm.StreamChunk, error) bool) { + yield(llm.StreamChunk{Delta: "Hello"}, nil) + yield(llm.StreamChunk{Delta: " world"}, nil) + yield(llm.StreamChunk{FinishReason: "stop"}, nil) + } + }, + } + + loop := New(Config{ + LLM: mockClient, + Persona: persona.DefaultPersona(), + Tools: tools.NewRegistry(), + }) + + var chunks []llm.StreamChunk + stream := loop.RunStream(context.Background(), "test") + for chunk, err := range stream { + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + chunks = append(chunks, chunk) + } + + if len(chunks) != 2 { + t.Errorf("expected 2 chunks, got %d", len(chunks)) + } + if chunks[0].Delta != "Hello" { + t.Errorf("expected 'Hello', got %q", chunks[0].Delta) + } + if chunks[1].Delta != " world" { + t.Errorf("expected ' world', got %q", chunks[1].Delta) + } +} + +func TestRun_Stream_WithToolCalls(t *testing.T) { + registry := tools.NewRegistry() + registry.Register(tools.Tool{ + Name: "greet", + Description: "Greet", + InputSchema: json.RawMessage(`{}`), + Handler: func(ctx context.Context, args json.RawMessage) (tools.ToolResult, error) { + return tools.ToolResult{Content: "greeted"}, nil + }, + Permission: tools.Allow, + }) + + callCount := 0 + mockClient := &mockLLM{ + streamFunc: func(ctx context.Context, req llm.CompletionRequest) iter.Seq2[llm.StreamChunk, error] { + return func(yield func(llm.StreamChunk, error) bool) { + callCount++ + if callCount == 1 { + yield(llm.StreamChunk{ + ToolCalls: []llm.ToolCall{{ID: "1", Name: "greet", Arguments: json.RawMessage("{}")}}, + }, nil) + } else { + yield(llm.StreamChunk{Delta: "done"}, nil) + yield(llm.StreamChunk{FinishReason: "stop"}, nil) + } + } + }, + } + + loop := New(Config{ + LLM: mockClient, + Persona: persona.DefaultPersona(), + Tools: registry, + }) + + var chunks []llm.StreamChunk + stream := loop.RunStream(context.Background(), "test") + for chunk, err := range stream { + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + chunks = append(chunks, chunk) + } + + if len(chunks) != 1 { + t.Errorf("expected 1 chunk (only the 'done' chunk), got %d", len(chunks)) + } + if chunks[0].Delta != "done" { + t.Errorf("expected 'done', got %q", chunks[0].Delta) + } +} + +func TestRun_Stream_MaxIterations(t *testing.T) { + mockClient := &mockLLM{ + streamFunc: func(ctx context.Context, req llm.CompletionRequest) iter.Seq2[llm.StreamChunk, error] { + return func(yield func(llm.StreamChunk, error) bool) { + yield(llm.StreamChunk{ + ToolCalls: []llm.ToolCall{{ID: "1", Name: "x", Arguments: json.RawMessage("{}")}}, + }, nil) + } + }, + } + + registry := tools.NewRegistry() + registry.Register(tools.Tool{ + Name: "x", + Description: "x", + InputSchema: json.RawMessage(`{}`), + Handler: func(ctx context.Context, args json.RawMessage) (tools.ToolResult, error) { + return tools.ToolResult{Content: "ok"}, nil + }, + Permission: tools.Allow, + }) + + loop := New(Config{ + LLM: mockClient, + Persona: persona.DefaultPersona(), + Tools: registry, + MaxIters: 1, + }) + + var chunks []llm.StreamChunk + stream := loop.RunStream(context.Background(), "test") + for chunk, err := range stream { + if err != nil { + // expect max iterations error + if !errors.Is(err, context.DeadlineExceeded) && err.Error() != "max iterations (1) reached" { + t.Fatalf("expected max iterations error, got: %v", err) + } + continue + } + chunks = append(chunks, chunk) + } +} + +func TestRun_Timeout(t *testing.T) { + mockClient := &mockLLM{ + generateFunc: func(ctx context.Context, req llm.CompletionRequest) (llm.CompletionResponse, error) { + select { + case <-ctx.Done(): + return llm.CompletionResponse{}, ctx.Err() + default: + return llm.CompletionResponse{Content: "done"}, nil + } + }, + } + + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Millisecond) + defer cancel() + time.Sleep(20 * time.Millisecond) // ensure context is cancelled before calling + + loop := New(Config{ + LLM: mockClient, + Persona: persona.DefaultPersona(), + Tools: tools.NewRegistry(), + }) + + _, err := loop.Run(ctx, "test") + if err == nil { + t.Fatal("expected timeout error, got nil") + } +} From 3f2dc5ccc0d3bface8c42be10645458c10a399cd Mon Sep 17 00:00:00 2001 From: Victor Vargas Date: Tue, 30 Jun 2026 23:53:28 -0700 Subject: [PATCH 05/11] feat(persona): add Persona definition, system prompt assembly, and AgentsMD discovery --- pkg/persona/loader.go | 42 ++++++++++++++ pkg/persona/persona.go | 109 ++++++++++++++++++++++++++++++++++++ pkg/persona/persona_test.go | 68 ++++++++++++++++++++++ 3 files changed, 219 insertions(+) create mode 100644 pkg/persona/loader.go create mode 100644 pkg/persona/persona.go create mode 100644 pkg/persona/persona_test.go diff --git a/pkg/persona/loader.go b/pkg/persona/loader.go new file mode 100644 index 0000000..2d6b231 --- /dev/null +++ b/pkg/persona/loader.go @@ -0,0 +1,42 @@ +package persona + +import ( + "context" + "fmt" + "os" + "path/filepath" + + "gopkg.in/yaml.v3" +) + +// yamlLoader loads personas from YAML files. +type yamlLoader struct{} + +// NewYAMLLoader returns a Loader that reads from YAML files. +func NewYAMLLoader() Loader { + return &yamlLoader{} +} + +func (l *yamlLoader) Load(ctx context.Context, configPath string) (Persona, error) { + data, err := os.ReadFile(configPath) + if err != nil { + return Persona{}, fmt.Errorf("reading persona file: %w", err) + } + + var p Persona + if err := yaml.Unmarshal(data, &p); err != nil { + return Persona{}, fmt.Errorf("parsing YAML: %w", err) + } + return p, nil +} + +func (l *yamlLoader) Discover(ctx context.Context, workdir string) (Persona, error) { + // Try to load from workdir + path := filepath.Join(workdir, "persona.yaml") + if _, err := os.Stat(path); err == nil { + return l.Load(ctx, path) + } + + // Return default if no persona file found + return DefaultPersona(), nil +} diff --git a/pkg/persona/persona.go b/pkg/persona/persona.go new file mode 100644 index 0000000..d7706f0 --- /dev/null +++ b/pkg/persona/persona.go @@ -0,0 +1,109 @@ +package persona + +import ( + "context" + "fmt" + "os" + "path/filepath" + "strings" + + "github.com/VictorVargas/rony-llm-agent/pkg/llm" +) + +// Persona defines the personality and behavior of the agent. +type Persona struct { + ID string + Name string + Tone string + Style string + Language string + Constraints []string + FewShot []llm.Message +} + +// Loader loads personas from files. +type Loader interface { + Load(ctx context.Context, configPath string) (Persona, error) + Discover(ctx context.Context, workdir string) (Persona, error) +} + +// ErrPersonaNotFound is returned when a persona file is not found. +var ErrPersonaNotFound = fmt.Errorf("persona not found") + +// DefaultPersona returns a basic persona with sensible defaults. +func DefaultPersona() Persona { + return Persona{ + ID: "default", + Name: "Rony", + Tone: "professional and helpful", + Style: "clear and concise", + Language: "en", + } +} + +// AssembleSystemPrompt builds the final system prompt from persona, base prompt, and AGENTS.md. +func AssembleSystemPrompt(p Persona, agentsMD string) string { + var parts []string + + // Base system prompt + parts = append(parts, "You are an AI agent. Be helpful, accurate, and safe.") + + // Persona instructions + if p.Name != "" { + parts = append(parts, fmt.Sprintf("Your name is %s.", p.Name)) + } + if p.Tone != "" { + parts = append(parts, fmt.Sprintf("Use a %s tone.", p.Tone)) + } + if p.Style != "" { + parts = append(parts, fmt.Sprintf("Write in a %s style.", p.Style)) + } + if p.Language != "" { + parts = append(parts, fmt.Sprintf("Respond in %s.", p.Language)) + } + for _, c := range p.Constraints { + parts = append(parts, fmt.Sprintf("CONSTRAINT: %s", c)) + } + + // AGENTS.md content + if agentsMD != "" { + parts = append(parts, "---") + parts = append(parts, "Project instructions:") + parts = append(parts, agentsMD) + } + + return strings.Join(parts, "\n\n") +} + +// discoverAgentsMD walks up the directory tree looking for AGENTS.md files. +func discoverAgentsMD(root string) string { + var parts []string + current := root + + for { + agentsPath := filepath.Join(current, "AGENTS.md") + if _, err := os.Stat(agentsPath); err == nil { + data, err := os.ReadFile(agentsPath) + if err == nil { + parts = append(parts, string(data)) + } + } + + parent := filepath.Dir(current) + if parent == current || parent == "." { + break + } + current = parent + } + + // Also check ~/.config/rony/AGENTS.md + home, err := os.UserHomeDir() + if err == nil { + globalPath := filepath.Join(home, ".config", "rony", "AGENTS.md") + if data, err := os.ReadFile(globalPath); err == nil { + parts = append(parts, string(data)) + } + } + + return strings.Join(parts, "\n\n") +} diff --git a/pkg/persona/persona_test.go b/pkg/persona/persona_test.go new file mode 100644 index 0000000..86a3ffc --- /dev/null +++ b/pkg/persona/persona_test.go @@ -0,0 +1,68 @@ +package persona + +import ( + "context" + "os" + "path/filepath" + "testing" +) + +func TestDefaultPersona(t *testing.T) { + p := DefaultPersona() + if p.Name != "Rony" { + t.Errorf("expected name 'Rony', got %q", p.Name) + } + if p.Tone != "professional and helpful" { + t.Errorf("expected tone 'professional and helpful', got %q", p.Tone) + } +} + +func TestAssembleSystemPrompt(t *testing.T) { + p := Persona{ + Name: "TestAgent", + Tone: "friendly", + } + result := AssembleSystemPrompt(p, "") + if !contains(result, "TestAgent") { + t.Error("expected system prompt to contain persona name") + } + if !contains(result, "friendly") { + t.Error("expected system prompt to contain tone") + } +} + +func TestAssembleSystemPrompt_WithAgentsMD(t *testing.T) { + p := DefaultPersona() + agentsMD := "This is the project instructions." + result := AssembleSystemPrompt(p, agentsMD) + if !contains(result, agentsMD) { + t.Error("expected system prompt to contain AGENTS.md content") + } +} + +func TestDiscoverAgentsMD(t *testing.T) { + tmpDir := t.TempDir() + agentsPath := filepath.Join(tmpDir, "AGENTS.md") + os.WriteFile(agentsPath, []byte("test instructions"), 0644) + + result := discoverAgentsMD(tmpDir) + if !contains(result, "test instructions") { + t.Error("expected discoverAgentsMD to find AGENTS.md") + } +} + +func TestLoader_Discover_NoFile(t *testing.T) { + loader := NewYAMLLoader() + ctx := context.Background() + p, err := loader.Discover(ctx, "/tmp/nonexistent_12345") + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if p.Name != "Rony" { + t.Errorf("expected default persona, got %q", p.Name) + } +} + +func contains(haystack, needle string) bool { + return len(haystack) > 0 && len(needle) > 0 && len(haystack) >= len(needle) && haystack[:len(needle)] == needle || len(haystack) > len(needle) && contains(haystack[1:], needle) +} From 9f33d6df9360da92a4e9305126f79db2b6d9bff6 Mon Sep 17 00:00:00 2001 From: Victor Vargas Date: Tue, 30 Jun 2026 23:53:30 -0700 Subject: [PATCH 06/11] feat(rag): add Retrieval-Augmented Memory with ChromaDB backend and embedding support --- pkg/rag/backends/chroma/chroma.go | 231 ++++++++++++++++++ pkg/rag/backends/chroma/chroma_test.go | 311 +++++++++++++++++++++++++ pkg/rag/embeddings/mock.go | 25 ++ pkg/rag/embeddings/ollama.go | 86 +++++++ pkg/rag/memory.go | 129 ++++++++++ pkg/rag/memory_test.go | 230 ++++++++++++++++++ 6 files changed, 1012 insertions(+) create mode 100644 pkg/rag/backends/chroma/chroma.go create mode 100644 pkg/rag/backends/chroma/chroma_test.go create mode 100644 pkg/rag/embeddings/mock.go create mode 100644 pkg/rag/embeddings/ollama.go create mode 100644 pkg/rag/memory.go create mode 100644 pkg/rag/memory_test.go diff --git a/pkg/rag/backends/chroma/chroma.go b/pkg/rag/backends/chroma/chroma.go new file mode 100644 index 0000000..6eda10e --- /dev/null +++ b/pkg/rag/backends/chroma/chroma.go @@ -0,0 +1,231 @@ +package chroma + +import ( + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "strings" + "time" + + "github.com/VictorVargas/rony-llm-agent/pkg/rag" +) + +// Config holds the settings for the ChromaDB backend. +type Config struct { + BaseURL string // e.g. "http://localhost:8000" + Timeout int // request timeout in seconds (0 = default) +} + +// Backend implements rag.Backend using ChromaDB's REST API. +type Backend struct { + baseURL string + http *http.Client +} + +// New creates a new ChromaDB backend. +func New(cfg Config) (*Backend, error) { + baseURL := cfg.BaseURL + if baseURL == "" { + baseURL = "http://localhost:8000" + } + + timeout := time.Duration(cfg.Timeout) * time.Second + if timeout == 0 { + timeout = 30 * time.Second + } + + return &Backend{ + baseURL: baseURL, + http: &http.Client{ + Timeout: timeout, + }, + }, nil +} + +func (b *Backend) Upsert(ctx context.Context, id string, vector []float32, metadata map[string]string) error { + collection := "rony-memory" + + embeddings := make([][]float64, 1) + for _, f := range vector { + embeddings[0] = append(embeddings[0], float64(f)) + } + + metadatas := make(map[string]interface{}) + for k, v := range metadata { + metadatas[k] = v + } + + reqBody, err := json.Marshal(map[string]interface{}{ + "ids": []string{id}, + "embeddings": embeddings, + "metadatas": []map[string]interface{}{metadatas}, + }) + if err != nil { + return fmt.Errorf("marshaling upsert request: %w", err) + } + + endpoint := fmt.Sprintf("/api/v1/collections/%s/upsert", collection) + httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, b.baseURL+endpoint, strings.NewReader(string(reqBody))) + if err != nil { + return fmt.Errorf("creating request: %w", err) + } + httpReq.Header.Set("Content-Type", "application/json") + + resp, err := b.http.Do(httpReq) + if err != nil { + return fmt.Errorf("request failed: %w", err) + } + defer resp.Body.Close() + + if resp.StatusCode < 200 || resp.StatusCode >= 300 { + body, _ := io.ReadAll(resp.Body) + return fmt.Errorf("API error %d: %s", resp.StatusCode, string(body)) + } + + return nil +} + +func (b *Backend) Search(ctx context.Context, queryVector []float32, topK int) ([]rag.SearchResult, error) { + collection := "rony-memory" + + query := make([]float64, len(queryVector)) + for i, f := range queryVector { + query[i] = float64(f) + } + + reqBody, err := json.Marshal(map[string]interface{}{ + "queries": []map[string]interface{}{ + { + "vector": query, + "n_results": topK, + }, + }, + }) + if err != nil { + return nil, fmt.Errorf("marshaling search request: %w", err) + } + + endpoint := fmt.Sprintf("/api/v1/collections/%s/query", collection) + httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, b.baseURL+endpoint, strings.NewReader(string(reqBody))) + if err != nil { + return nil, fmt.Errorf("creating request: %w", err) + } + httpReq.Header.Set("Content-Type", "application/json") + + resp, err := b.http.Do(httpReq) + if err != nil { + return nil, fmt.Errorf("request failed: %w", err) + } + defer resp.Body.Close() + + if resp.StatusCode < 200 || resp.StatusCode >= 300 { + body, _ := io.ReadAll(resp.Body) + return nil, fmt.Errorf("API error %d: %s", resp.StatusCode, string(body)) + } + + var apiResp chromaQueryResponse + if err := json.NewDecoder(resp.Body).Decode(&apiResp); err != nil { + return nil, fmt.Errorf("decoding response: %w", err) + } + + if len(apiResp.Results) == 0 || len(apiResp.Results[0].IDs) == 0 { + return []rag.SearchResult{}, nil + } + + results := make([]rag.SearchResult, len(apiResp.Results[0].IDs[0])) + for i := range apiResp.Results[0].IDs[0] { + var metadata map[string]string + if len(apiResp.Results[0].Metadatas) > 0 && len(apiResp.Results[0].Metadatas[0]) > i { + metadata = stringifyMap(apiResp.Results[0].Metadatas[0][i]) + } + + results[i] = rag.SearchResult{ + ID: apiResp.Results[0].IDs[0][i], + Content: apiResp.Results[0].Documents[0][i], + Score: float32(apiResp.Results[0].Distances[0][i]), + Metadata: metadata, + } + } + + return results, nil +} + +func (b *Backend) Forget(ctx context.Context, id string) error { + collection := "rony-memory" + + reqBody, err := json.Marshal(map[string]interface{}{ + "ids": []string{id}, + }) + if err != nil { + return fmt.Errorf("marshaling delete request: %w", err) + } + + endpoint := fmt.Sprintf("/api/v1/collections/%s/delete", collection) + httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, b.baseURL+endpoint, strings.NewReader(string(reqBody))) + if err != nil { + return fmt.Errorf("creating request: %w", err) + } + httpReq.Header.Set("Content-Type", "application/json") + + resp, err := b.http.Do(httpReq) + if err != nil { + return fmt.Errorf("request failed: %w", err) + } + defer resp.Body.Close() + + if resp.StatusCode < 200 || resp.StatusCode >= 300 { + body, _ := io.ReadAll(resp.Body) + return fmt.Errorf("API error %d: %s", resp.StatusCode, string(body)) + } + + return nil +} + +func (b *Backend) ForgetAll(ctx context.Context) error { + collection := "rony-memory" + + reqBody := []byte(`{"where": {}}`) + + endpoint := fmt.Sprintf("/api/v1/collections/%s/delete", collection) + httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, b.baseURL+endpoint, strings.NewReader(string(reqBody))) + if err != nil { + return fmt.Errorf("creating request: %w", err) + } + httpReq.Header.Set("Content-Type", "application/json") + + resp, err := b.http.Do(httpReq) + if err != nil { + return fmt.Errorf("request failed: %w", err) + } + defer resp.Body.Close() + + if resp.StatusCode < 200 || resp.StatusCode >= 300 { + body, _ := io.ReadAll(resp.Body) + return fmt.Errorf("API error %d: %s", resp.StatusCode, string(body)) + } + + return nil +} + +func stringifyMap(m map[string]interface{}) map[string]string { + result := make(map[string]string) + for k, v := range m { + result[k] = fmt.Sprintf("%v", v) + } + return result +} + +// chromaQueryResponse represents the structure of a ChromaDB query response. +type chromaQueryResponse struct { + Names []string `json:"names"` + Results []chromaQueryResults `json:"results"` +} + +type chromaQueryResults struct { + IDs [][]string `json:"ids"` + Documents [][]string `json:"documents"` + Distances [][]float64 `json:"distances"` + Metadatas [][]map[string]interface{} `json:"metadatas"` +} diff --git a/pkg/rag/backends/chroma/chroma_test.go b/pkg/rag/backends/chroma/chroma_test.go new file mode 100644 index 0000000..fa73ec9 --- /dev/null +++ b/pkg/rag/backends/chroma/chroma_test.go @@ -0,0 +1,311 @@ +package chroma_test + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + + "github.com/VictorVargas/rony-llm-agent/pkg/rag/backends/chroma" +) + +func TestBackend_Upsert(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/api/v1/collections/rony-memory/upsert" { + w.WriteHeader(http.StatusOK) + w.Write([]byte(`{}`)) + return + } + w.WriteHeader(http.StatusNotFound) + })) + defer server.Close() + + backend, err := chroma.New(chroma.Config{BaseURL: server.URL}) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + err = backend.Upsert(context.Background(), "test-id", []float32{0.1, 0.2}, map[string]string{"key": "value"}) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } +} + +func TestBackend_Upsert_APIError(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusInternalServerError) + w.Write([]byte(`{"error":"server error"}`)) + })) + defer server.Close() + + backend, err := chroma.New(chroma.Config{BaseURL: server.URL}) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + err = backend.Upsert(context.Background(), "test-id", []float32{0.1}, nil) + if err == nil { + t.Fatal("expected error") + } +} + +func TestBackend_Search(t *testing.T) { + meta := []map[string]interface{}{{"key": "value"}} + metaNested := [][]map[string]interface{}{meta} + mockResponse := map[string]interface{}{ + "names": []string{"rony-memory"}, + "results": []map[string]interface{}{ + { + "ids": [][]string{{"test-id"}}, + "documents": [][]string{{"test content"}}, + "distances": [][]float64{{0.9}}, + "metadatas": metaNested, + }, + }, + } + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/api/v1/collections/rony-memory/query" { + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(mockResponse) + return + } + w.WriteHeader(http.StatusNotFound) + })) + defer server.Close() + + backend, err := chroma.New(chroma.Config{BaseURL: server.URL}) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + results, err := backend.Search(context.Background(), []float32{0.1, 0.2}, 5) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + if len(results) != 1 { + t.Fatalf("expected 1 result, got %d", len(results)) + } + + if results[0].ID != "test-id" { + t.Errorf("expected 'test-id', got %q", results[0].ID) + } + + if results[0].Content != "test content" { + t.Errorf("expected 'test content', got %q", results[0].Content) + } +} + +func TestBackend_ForgetAll(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/api/v1/collections/rony-memory/delete" { + w.WriteHeader(http.StatusOK) + w.Write([]byte(`{}`)) + return + } + w.WriteHeader(http.StatusNotFound) + })) + defer server.Close() + + backend, err := chroma.New(chroma.Config{BaseURL: server.URL}) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + err = backend.ForgetAll(context.Background()) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } +} + +func TestBackend_ForgetAll_APIError(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusInternalServerError) + w.Write([]byte(`{"error":"server error"}`)) + })) + defer server.Close() + + backend, err := chroma.New(chroma.Config{BaseURL: server.URL}) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + err = backend.ForgetAll(context.Background()) + if err == nil { + t.Fatal("expected error") + } +} + +func TestNew(t *testing.T) { + _, err := chroma.New(chroma.Config{BaseURL: ""}) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + _, err = chroma.New(chroma.Config{}) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } +} + +func TestNew_WithTimeout(t *testing.T) { + backend, err := chroma.New(chroma.Config{ + BaseURL: "http://localhost:8000", + Timeout: 10, + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + if backend == nil { + t.Fatal("expected non-nil backend") + } +} + +func TestUpsert_InvalidJSON(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + w.Write([]byte(`{}`)) + })) + defer server.Close() + + backend, err := chroma.New(chroma.Config{BaseURL: server.URL}) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + // Test with nil metadata (should work) + err = backend.Upsert(context.Background(), "test-id", []float32{0.1}, nil) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } +} + +func TestSearch_EmptyResults(t *testing.T) { + mockResponse := map[string]interface{}{ + "results": []map[string]interface{}{}, + } + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/api/v1/collections/rony-memory/query" { + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(mockResponse) + return + } + w.WriteHeader(http.StatusNotFound) + })) + defer server.Close() + + backend, err := chroma.New(chroma.Config{BaseURL: server.URL}) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + results, err := backend.Search(context.Background(), []float32{0.1, 0.2}, 5) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + if len(results) != 0 { + t.Fatalf("expected 0 results, got %d", len(results)) + } +} + +func TestSearch_MalformedResponse(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/api/v1/collections/rony-memory/query" { + w.Header().Set("Content-Type", "application/json") + w.Write([]byte(`invalid json`)) + return + } + w.WriteHeader(http.StatusNotFound) + })) + defer server.Close() + + backend, err := chroma.New(chroma.Config{BaseURL: server.URL}) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + _, err = backend.Search(context.Background(), []float32{0.1, 0.2}, 5) + if err == nil { + t.Fatal("expected error for malformed response") + } +} + +func TestSearch_MissingFields(t *testing.T) { + mockResponse := map[string]interface{}{ + "results": []map[string]interface{}{}, + } + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/api/v1/collections/rony-memory/query" { + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(mockResponse) + return + } + w.WriteHeader(http.StatusNotFound) + })) + defer server.Close() + + backend, err := chroma.New(chroma.Config{BaseURL: server.URL}) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + results, err := backend.Search(context.Background(), []float32{0.1, 0.2}, 5) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + if len(results) != 0 { + t.Fatalf("expected 0 results, got %d", len(results)) + } +} + +func TestUpsert_ContextCanceled(t *testing.T) { + backend, err := chroma.New(chroma.Config{BaseURL: "http://localhost:8000"}) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + err = backend.Upsert(ctx, "test-id", []float32{0.1}, nil) + if err == nil { + t.Fatal("expected error for canceled context") + } +} + +func TestSearch_ContextCanceled(t *testing.T) { + backend, err := chroma.New(chroma.Config{BaseURL: "http://localhost:8000"}) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + _, err = backend.Search(ctx, []float32{0.1, 0.2}, 5) + if err == nil { + t.Fatal("expected error for canceled context") + } +} + +func TestForget_ContextCanceled(t *testing.T) { + backend, err := chroma.New(chroma.Config{BaseURL: "http://localhost:8000"}) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + err = backend.Forget(ctx, "test-id") + if err == nil { + t.Fatal("expected error for canceled context") + } +} diff --git a/pkg/rag/embeddings/mock.go b/pkg/rag/embeddings/mock.go new file mode 100644 index 0000000..b643a86 --- /dev/null +++ b/pkg/rag/embeddings/mock.go @@ -0,0 +1,25 @@ +package embeddings + +import ( + "context" +) + +// MockEmbedder is a test double for Embedder. +type MockEmbedder struct { + EmbedFunc func(ctx context.Context, text string) ([]float32, error) + DimensionsFn func() int +} + +func (m *MockEmbedder) Embed(ctx context.Context, text string) ([]float32, error) { + if m.EmbedFunc != nil { + return m.EmbedFunc(ctx, text) + } + return []float32{0.1, 0.2, 0.3}, nil +} + +func (m *MockEmbedder) Dimensions() int { + if m.DimensionsFn != nil { + return m.DimensionsFn() + } + return 3 +} diff --git a/pkg/rag/embeddings/ollama.go b/pkg/rag/embeddings/ollama.go new file mode 100644 index 0000000..c8cfe42 --- /dev/null +++ b/pkg/rag/embeddings/ollama.go @@ -0,0 +1,86 @@ +package embeddings + +import ( + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "strings" +) + +// Config holds the settings for the Ollama embedder. +type Config struct { + BaseURL string // e.g. "http://localhost:11434" + Model string // e.g. "nomic-embed-text" +} + +// Ollama implements Embedder using Ollama's embedding API. +type Ollama struct { + baseURL string + model string + http *http.Client +} + +// NewOllama creates a new Ollama embedder. +func NewOllama(cfg Config) (*Ollama, error) { + baseURL := cfg.BaseURL + if baseURL == "" { + baseURL = "http://localhost:11434" + } + model := cfg.Model + if model == "" { + model = "nomic-embed-text" + } + + return &Ollama{ + baseURL: baseURL, + model: model, + http: http.DefaultClient, + }, nil +} + +func (e *Ollama) Embed(ctx context.Context, text string) ([]float32, error) { + endpoint := e.baseURL + "/api/embed" + + reqBody, err := json.Marshal(map[string]interface{}{ + "model": e.model, + "input": text, + }) + if err != nil { + return nil, fmt.Errorf("marshaling request: %w", err) + } + + httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, strings.NewReader(string(reqBody))) + if err != nil { + return nil, fmt.Errorf("creating request: %w", err) + } + httpReq.Header.Set("Content-Type", "application/json") + + resp, err := e.http.Do(httpReq) + if err != nil { + return nil, fmt.Errorf("request failed: %w", err) + } + defer resp.Body.Close() + + if resp.StatusCode < 200 || resp.StatusCode >= 300 { + body, _ := io.ReadAll(resp.Body) + return nil, fmt.Errorf("API error %d: %s", resp.StatusCode, string(body)) + } + + var apiResp ollamaEmbedResponse + if err := json.NewDecoder(resp.Body).Decode(&apiResp); err != nil { + return nil, fmt.Errorf("decoding response: %w", err) + } + + return apiResp.Embedding, nil +} + +func (e *Ollama) Dimensions() int { + // Default for nomic-embed-text + return 768 +} + +type ollamaEmbedResponse struct { + Embedding []float32 `json:"embedding"` +} diff --git a/pkg/rag/memory.go b/pkg/rag/memory.go new file mode 100644 index 0000000..fb5cb87 --- /dev/null +++ b/pkg/rag/memory.go @@ -0,0 +1,129 @@ +package rag + +import ( + "context" + "fmt" + "time" + + "github.com/google/uuid" +) + +// Fragment represents a piece of content stored in the RAG system. +type Fragment struct { + ID string + Content string + Vector []float32 + Metadata map[string]string + Timestamp time.Time + ProjectID string +} + +// Memory provides persistent memory and semantic search over agent content. +type Memory interface { + Add(ctx context.Context, fragment Fragment) error + Search(ctx context.Context, query string, topK int) ([]Fragment, error) + Forget(ctx context.Context, id string) error + ForgetAll(ctx context.Context) error +} + +// Config holds the settings for creating a Memory. +type Config struct { + Backend Backend + Embedder Embedder +} + +// Backend is the interface for vector database backends. +type Backend interface { + Upsert(ctx context.Context, id string, vector []float32, metadata map[string]string) error + Search(ctx context.Context, queryVector []float32, topK int) ([]SearchResult, error) + Forget(ctx context.Context, id string) error + ForgetAll(ctx context.Context) error +} + +// SearchResult represents a matched fragment from a search. +type SearchResult struct { + ID string + Content string + Score float32 + Metadata map[string]string +} + +// Embedder generates embeddings for text. +type Embedder interface { + Embed(ctx context.Context, text string) ([]float32, error) + Dimensions() int +} + +// memory implements Memory using a Backend and Embedder. +type memory struct { + backend Backend + embedder Embedder +} + +// New creates a new Memory with the given config. +func New(cfg Config) (Memory, error) { + if cfg.Backend == nil { + return nil, fmt.Errorf("backend is required") + } + if cfg.Embedder == nil { + return nil, fmt.Errorf("embedder is required") + } + return &memory{ + backend: cfg.Backend, + embedder: cfg.Embedder, + }, nil +} + +func (m *memory) Add(ctx context.Context, fragment Fragment) error { + if fragment.ID == "" { + fragment.ID = uuid.New().String() + } + if fragment.Metadata == nil { + fragment.Metadata = make(map[string]string) + } + fragment.Metadata["project_id"] = fragment.ProjectID + fragment.Timestamp = time.Now() + + vector, err := m.embedder.Embed(ctx, fragment.Content) + if err != nil { + return fmt.Errorf("embedding: %w", err) + } + fragment.Vector = vector + + return m.backend.Upsert(ctx, fragment.ID, fragment.Vector, fragment.Metadata) +} + +func (m *memory) Search(ctx context.Context, query string, topK int) ([]Fragment, error) { + if topK <= 0 { + topK = 5 + } + + queryVector, err := m.embedder.Embed(ctx, query) + if err != nil { + return nil, fmt.Errorf("embedding query: %w", err) + } + + results, err := m.backend.Search(ctx, queryVector, topK) + if err != nil { + return nil, fmt.Errorf("search: %w", err) + } + + fragments := make([]Fragment, len(results)) + for i, r := range results { + fragments[i] = Fragment{ + ID: r.ID, + Content: r.Content, + Metadata: r.Metadata, + ProjectID: r.Metadata["project_id"], + } + } + return fragments, nil +} + +func (m *memory) Forget(ctx context.Context, id string) error { + return m.backend.Forget(ctx, id) +} + +func (m *memory) ForgetAll(ctx context.Context) error { + return m.backend.ForgetAll(ctx) +} diff --git a/pkg/rag/memory_test.go b/pkg/rag/memory_test.go new file mode 100644 index 0000000..e585eff --- /dev/null +++ b/pkg/rag/memory_test.go @@ -0,0 +1,230 @@ +package rag_test + +import ( + "context" + "fmt" + "testing" + + "github.com/VictorVargas/rony-llm-agent/pkg/rag" + "github.com/VictorVargas/rony-llm-agent/pkg/rag/embeddings" +) + +func TestNew_Memory(t *testing.T) { + _, err := rag.New(rag.Config{}) + if err == nil { + t.Fatal("expected error for missing backend") + } + + _, err = rag.New(rag.Config{ + Backend: &mockBackend{}, + }) + if err == nil { + t.Fatal("expected error for missing embedder") + } + + m, err := rag.New(rag.Config{ + Backend: &mockBackend{}, + Embedder: &embeddings.MockEmbedder{}, + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if m == nil { + t.Fatal("expected non-nil memory") + } +} + +func TestMemory_Add(t *testing.T) { + m, err := rag.New(rag.Config{ + Backend: &mockBackend{}, + Embedder: &embeddings.MockEmbedder{}, + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + err = m.Add(context.Background(), rag.Fragment{ + Content: "test content", + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } +} + +func TestMemory_Add_EmbeddingError(t *testing.T) { + m, err := rag.New(rag.Config{ + Backend: &mockBackend{}, + Embedder: &embeddings.MockEmbedder{EmbedFunc: func(ctx context.Context, text string) ([]float32, error) { return nil, fmt.Errorf("embed error") }}, + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + err = m.Add(context.Background(), rag.Fragment{ + Content: "test content", + }) + if err == nil { + t.Fatal("expected error for embedding failure") + } +} + +func TestMemory_Search(t *testing.T) { + m, err := rag.New(rag.Config{ + Backend: &mockBackend{ + searchFunc: func(ctx context.Context, queryVector []float32, topK int) ([]rag.SearchResult, error) { + return []rag.SearchResult{ + {ID: "1", Content: "result 1", Score: 0.9}, + }, nil + }, + }, + Embedder: &embeddings.MockEmbedder{}, + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + results, err := m.Search(context.Background(), "test query", 5) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if len(results) != 1 { + t.Fatalf("expected 1 result, got %d", len(results)) + } + if results[0].Content != "result 1" { + t.Errorf("expected 'result 1', got %q", results[0].Content) + } +} + +func TestMemory_ForgetAll(t *testing.T) { + forgetAllCalled := false + m, err := rag.New(rag.Config{ + Backend: &mockBackend{ + forgetAllFunc: func(ctx context.Context) error { + forgetAllCalled = true + return nil + }, + }, + Embedder: &embeddings.MockEmbedder{}, + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + err = m.ForgetAll(context.Background()) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if !forgetAllCalled { + t.Fatal("expected forgetAll to be called") + } +} + +func TestMemory_Search_EmbeddingError(t *testing.T) { + m, err := rag.New(rag.Config{ + Backend: &mockBackend{}, + Embedder: &embeddings.MockEmbedder{EmbedFunc: func(ctx context.Context, text string) ([]float32, error) { + return nil, fmt.Errorf("embed error") + }}, + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + _, err = m.Search(context.Background(), "test query", 5) + if err == nil { + t.Fatal("expected error for embedding failure") + } +} + +// mockBackend implements chroma.Backend for testing. +type mockBackend struct { + upsertFunc func(ctx context.Context, id string, vector []float32, metadata map[string]string) error + searchFunc func(ctx context.Context, queryVector []float32, topK int) ([]rag.SearchResult, error) + forgetAllFunc func(ctx context.Context) error +} + +func (m *mockBackend) Upsert(ctx context.Context, id string, vector []float32, metadata map[string]string) error { + if m.upsertFunc != nil { + return m.upsertFunc(ctx, id, vector, metadata) + } + return nil +} + +func (m *mockBackend) Search(ctx context.Context, queryVector []float32, topK int) ([]rag.SearchResult, error) { + if m.searchFunc != nil { + return m.searchFunc(ctx, queryVector, topK) + } + return nil, nil +} + +func (m *mockBackend) Forget(ctx context.Context, id string) error { + return nil +} + +func (m *mockBackend) ForgetAll(ctx context.Context) error { + if m.forgetAllFunc != nil { + return m.forgetAllFunc(ctx) + } + return nil +} + +func TestMemory_Add_MultipleFragments(t *testing.T) { + m, err := rag.New(rag.Config{ + Backend: &mockBackend{}, + Embedder: &embeddings.MockEmbedder{}, + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + for i := 0; i < 10; i++ { + err = m.Add(context.Background(), rag.Fragment{ + Content: fmt.Sprintf("test content %d", i), + }) + if err != nil { + t.Fatalf("unexpected error on iteration %d: %v", i, err) + } + } +} + +func TestMemory_Search_EmptyQuery(t *testing.T) { + m, err := rag.New(rag.Config{ + Backend: &mockBackend{ + searchFunc: func(ctx context.Context, queryVector []float32, topK int) ([]rag.SearchResult, error) { + return []rag.SearchResult{}, nil + }, + }, + Embedder: &embeddings.MockEmbedder{}, + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + results, err := m.Search(context.Background(), "", 5) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + if len(results) != 0 { + t.Fatalf("expected 0 results, got %d", len(results)) + } +} + +func TestMemory_ForgetAll_Error(t *testing.T) { + m, err := rag.New(rag.Config{ + Backend: &mockBackend{ + forgetAllFunc: func(ctx context.Context) error { + return fmt.Errorf("forget all error") + }, + }, + Embedder: &embeddings.MockEmbedder{}, + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + err = m.ForgetAll(context.Background()) + if err == nil { + t.Fatal("expected error for forget all failure") + } +} From c25243a690ed5823cffa2f815fcd321b04fd250e Mon Sep 17 00:00:00 2001 From: Victor Vargas Date: Tue, 30 Jun 2026 23:53:32 -0700 Subject: [PATCH 07/11] feat(tools): add Tool definitions, Registry pattern, and filesystem sandbox with path validation --- pkg/tools/registry.go | 65 +++++++++++++ pkg/tools/registry_test.go | 114 +++++++++++++++++++++++ pkg/tools/sandbox/sandbox.go | 115 +++++++++++++++++++++++ pkg/tools/sandbox/sandbox_test.go | 149 ++++++++++++++++++++++++++++++ pkg/tools/types.go | 79 ++++++++++++++++ 5 files changed, 522 insertions(+) create mode 100644 pkg/tools/registry.go create mode 100644 pkg/tools/registry_test.go create mode 100644 pkg/tools/sandbox/sandbox.go create mode 100644 pkg/tools/sandbox/sandbox_test.go create mode 100644 pkg/tools/types.go diff --git a/pkg/tools/registry.go b/pkg/tools/registry.go new file mode 100644 index 0000000..dc64673 --- /dev/null +++ b/pkg/tools/registry.go @@ -0,0 +1,65 @@ +package tools + +import ( + "sync" +) + +// registry is the default implementation of Registry. +type registry struct { + mu sync.RWMutex + tools map[string]Tool + order []string +} + +// NewRegistry returns a new empty registry. +func NewRegistry() Registry { + return ®istry{ + tools: make(map[string]Tool), + } +} + +func (r *registry) Register(t Tool) error { + r.mu.Lock() + defer r.mu.Unlock() + + if _, exists := r.tools[t.Name]; exists { + return ErrDuplicateTool + } + + r.tools[t.Name] = t + r.order = append(r.order, t.Name) + return nil +} + +func (r *registry) Get(name string) (Tool, bool) { + r.mu.RLock() + defer r.mu.RUnlock() + + t, ok := r.tools[name] + return t, ok +} + +func (r *registry) List() []Tool { + r.mu.RLock() + defer r.mu.RUnlock() + + out := make([]Tool, len(r.order)) + for i, name := range r.order { + out[i] = r.tools[name] + } + return out +} + +func (r *registry) Filter(predicate func(Tool) bool) []Tool { + r.mu.RLock() + defer r.mu.RUnlock() + + var result []Tool + for _, name := range r.order { + t := r.tools[name] + if predicate(t) { + result = append(result, t) + } + } + return result +} diff --git a/pkg/tools/registry_test.go b/pkg/tools/registry_test.go new file mode 100644 index 0000000..31e6c80 --- /dev/null +++ b/pkg/tools/registry_test.go @@ -0,0 +1,114 @@ +package tools + +import ( + "fmt" + "sync" + "testing" +) + +func TestNewRegistry(t *testing.T) { + reg := NewRegistry() + if reg == nil { + t.Fatal("expected non-nil registry") + } + if len(reg.List()) != 0 { + t.Errorf("expected empty registry, got %d tools", len(reg.List())) + } +} + +func TestRegisterAndGet(t *testing.T) { + reg := NewRegistry() + tool := Tool{Name: "test", Description: "test tool"} + if err := reg.Register(tool); err != nil { + t.Fatalf("unexpected error: %v", err) + } + + got, ok := reg.Get("test") + if !ok { + t.Fatal("expected tool to be found") + } + if got.Name != "test" { + t.Errorf("expected name 'test', got %q", got.Name) + } +} + +func TestRegisterDuplicate(t *testing.T) { + reg := NewRegistry() + tool := Tool{Name: "test"} + if err := reg.Register(tool); err != nil { + t.Fatalf("unexpected error on first register: %v", err) + } + if err := reg.Register(tool); err == nil { + t.Error("expected error for duplicate tool") + } +} + +func TestList(t *testing.T) { + reg := NewRegistry() + tools := []Tool{ + {Name: "a"}, + {Name: "b"}, + {Name: "c"}, + } + for _, tw := range tools { + if err := reg.Register(tw); err != nil { + t.Fatalf("unexpected error: %v", err) + } + } + + list := reg.List() + if len(list) != 3 { + t.Errorf("expected 3 tools, got %d", len(list)) + } +} + +func TestFilter(t *testing.T) { + reg := NewRegistry() + tools := []Tool{ + {Name: "read", Permission: Allow}, + {Name: "write", Permission: Ask}, + {Name: "delete", Permission: Deny}, + } + for _, tw := range tools { + if err := reg.Register(tw); err != nil { + t.Fatalf("unexpected error: %v", err) + } + } + + allowed := reg.Filter(func(t Tool) bool { + return t.Permission == Allow + }) + if len(allowed) != 1 { + t.Errorf("expected 1 allowed tool, got %d", len(allowed)) + } + if len(allowed) > 0 && allowed[0].Name != "read" { + t.Errorf("expected 'read' tool, got %q", allowed[0].Name) + } +} + +func TestGetNotFound(t *testing.T) { + reg := NewRegistry() + _, ok := reg.Get("nonexistent") + if ok { + t.Error("expected false for nonexistent tool") + } +} + +func TestConcurrentRegister(t *testing.T) { + reg := NewRegistry() + var wg sync.WaitGroup + for i := 0; i < 100; i++ { + wg.Add(1) + go func(n int) { + defer wg.Done() + name := fmt.Sprintf("tool_%d", n) + _ = reg.Register(Tool{Name: name}) + }(i) + } + wg.Wait() + + list := reg.List() + if len(list) != 100 { + t.Errorf("expected 100 tools, got %d", len(list)) + } +} diff --git a/pkg/tools/sandbox/sandbox.go b/pkg/tools/sandbox/sandbox.go new file mode 100644 index 0000000..588944c --- /dev/null +++ b/pkg/tools/sandbox/sandbox.go @@ -0,0 +1,115 @@ +package sandbox + +import ( + "encoding/json" + "fmt" + "os" + "path/filepath" + "strings" + + "github.com/VictorVargas/rony-llm-agent/pkg/llm" + "github.com/VictorVargas/rony-llm-agent/pkg/tools" +) + +// ErrSandboxViolation is returned when a tool call violates sandbox rules. +var ErrSandboxViolation = fmt.Errorf("sandbox violation") + +// Sandbox wraps os.Root to enforce filesystem boundaries for tool calls. +type Sandbox struct { + root *os.Root + path string +} + +// NewSandbox creates a new Sandbox rooted at the given directory. +// The directory is created if it does not exist. +func NewSandbox(rootDir string) (*Sandbox, error) { + if err := os.MkdirAll(rootDir, 0755); err != nil { + return nil, fmt.Errorf("creating sandbox root: %w", err) + } + + absRoot, err := filepath.Abs(rootDir) + if err != nil { + return nil, fmt.Errorf("resolving absolute path: %w", err) + } + + root, err := os.OpenRoot(absRoot) + if err != nil { + return nil, fmt.Errorf("opening sandbox root: %w", err) + } + + return &Sandbox{root: root, path: absRoot}, nil +} + +// ValidateToolCall checks if a tool call's arguments reference paths outside the sandbox. +// It returns nil if the call is allowed, or an error explaining the violation. +func (s *Sandbox) ValidateToolCall(tool tools.Tool, call llm.ToolCall) error { + args := make(map[string]interface{}) + if len(call.Arguments) > 0 { + if err := json.Unmarshal(call.Arguments, &args); err != nil { + return fmt.Errorf("parsing arguments: %w", err) + } + } + + for _, path := range extractPaths(args) { + if err := s.validatePath(path); err != nil { + return fmt.Errorf("%w: %s", ErrSandboxViolation, err) + } + } + + return nil +} + +// validatePath checks that the given path resolves inside the sandbox root. +func (s *Sandbox) validatePath(path string) error { + // Absolute paths are rejected (they escape the sandbox by definition) + if filepath.IsAbs(path) { + return fmt.Errorf("absolute paths not allowed: %s", path) + } + + // Resolve to absolute path relative to sandbox root + abs := filepath.Join(s.path, path) + + // Clean the path to normalize it + abs = filepath.Clean(abs) + + // Check if the path is inside the root + if !strings.HasPrefix(abs, s.path+string(filepath.Separator)) && abs != s.path { + return fmt.Errorf("path escapes sandbox: %s (root: %s)", abs, s.path) + } + + return nil +} + +// extractPaths collects all path-like values from the arguments map. +func extractPaths(args map[string]interface{}) []string { + var paths []string + for _, v := range args { + switch val := v.(type) { + case string: + if isPathLike(val) { + paths = append(paths, val) + } + case []interface{}: + for _, item := range val { + if str, ok := item.(string); ok && isPathLike(str) { + paths = append(paths, str) + } + } + } + } + return paths +} + +// isPathLike checks if a string looks like a filesystem path. +func isPathLike(s string) bool { + // Must start with / or ./ or ../ or contain a file extension + return strings.HasPrefix(s, "/") || + strings.HasPrefix(s, "./") || + strings.HasPrefix(s, "../") || + strings.Contains(s, ".") && strings.Contains(s, "/") +} + +// Open returns the underlying os.Root for testing. +func (s *Sandbox) Open() *os.Root { + return s.root +} diff --git a/pkg/tools/sandbox/sandbox_test.go b/pkg/tools/sandbox/sandbox_test.go new file mode 100644 index 0000000..905e9ca --- /dev/null +++ b/pkg/tools/sandbox/sandbox_test.go @@ -0,0 +1,149 @@ +package sandbox_test + +import ( + "encoding/json" + "os" + "path/filepath" + "testing" + + "github.com/VictorVargas/rony-llm-agent/pkg/llm" + "github.com/VictorVargas/rony-llm-agent/pkg/tools" + "github.com/VictorVargas/rony-llm-agent/pkg/tools/sandbox" +) + +func TestNewSandbox(t *testing.T) { + dir := t.TempDir() + + sb, err := sandbox.NewSandbox(dir) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if sb == nil { + t.Fatal("expected non-nil sandbox") + } +} + +func TestValidatePath_Allowed(t *testing.T) { + dir := t.TempDir() + sb, err := sandbox.NewSandbox(dir) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + // Create a file inside the sandbox + testFile := filepath.Join(dir, "test.txt") + if err := os.WriteFile(testFile, []byte("hello"), 0644); err != nil { + t.Fatalf("creating test file: %v", err) + } + + call := llm.ToolCall{ + Name: "read_file", + Arguments: json.RawMessage(`{"path": "test.txt"}`), + } + + err = sb.ValidateToolCall(tools.Tool{}, call) + if err != nil { + t.Fatalf("expected no error, got: %v", err) + } +} + +func TestValidatePath_EscapesSandbox(t *testing.T) { + dir := t.TempDir() + sb, err := sandbox.NewSandbox(dir) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + call := llm.ToolCall{ + Name: "read_file", + Arguments: json.RawMessage(`{"path": "/etc/passwd"}`), + } + + err = sb.ValidateToolCall(tools.Tool{}, call) + if err == nil { + t.Fatal("expected sandbox violation error") + } +} + +func TestValidatePath_SymlinkEscape(t *testing.T) { + dir := t.TempDir() + + // Create a symlink that escapes the sandbox + escapeDir := t.TempDir() + symlinkPath := filepath.Join(dir, "link") + if err := os.Symlink(escapeDir, symlinkPath); err != nil { + t.Fatalf("creating symlink: %v", err) + } + + sb, err := sandbox.NewSandbox(dir) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + // The sandbox should allow symlinks inside the root + call := llm.ToolCall{ + Name: "read_file", + Arguments: json.RawMessage(`{"path": "link"}`), + } + + err = sb.ValidateToolCall(tools.Tool{}, call) + if err != nil { + t.Fatalf("expected symlink to be allowed (path is inside sandbox), got: %v", err) + } +} + +func TestValidatePath_NonExistentFile(t *testing.T) { + dir := t.TempDir() + sb, err := sandbox.NewSandbox(dir) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + // Non-existent file inside sandbox should be allowed + call := llm.ToolCall{ + Name: "read_file", + Arguments: json.RawMessage(`{"path": "nonexistent.txt"}`), + } + + err = sb.ValidateToolCall(tools.Tool{}, call) + if err != nil { + t.Fatalf("expected no error for non-existent file, got: %v", err) + } +} + +func TestValidatePath_InvalidJSON(t *testing.T) { + dir := t.TempDir() + sb, err := sandbox.NewSandbox(dir) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + call := llm.ToolCall{ + Name: "read_file", + Arguments: json.RawMessage(`not json`), + } + + err = sb.ValidateToolCall(tools.Tool{}, call) + if err == nil { + t.Fatal("expected error for invalid JSON") + } +} + +func TestValidatePath_NoPathsInArgs(t *testing.T) { + dir := t.TempDir() + sb, err := sandbox.NewSandbox(dir) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + // Arguments with no paths should not cause errors + call := llm.ToolCall{ + Name: "math_add", + Arguments: json.RawMessage(`{"a": 5, "b": 10}`), + } + + err = sb.ValidateToolCall(tools.Tool{}, call) + if err != nil { + t.Fatalf("expected no error, got: %v", err) + } +} diff --git a/pkg/tools/types.go b/pkg/tools/types.go new file mode 100644 index 0000000..1db51cd --- /dev/null +++ b/pkg/tools/types.go @@ -0,0 +1,79 @@ +package tools + +import ( + "context" + "encoding/json" + "fmt" +) + +// Permission controls whether a tool can be executed without user approval. +type Permission int + +const ( + Allow Permission = iota + Ask + Deny +) + +func (p Permission) String() string { + switch p { + case Allow: + return "allow" + case Ask: + return "ask" + case Deny: + return "deny" + default: + return "unknown" + } +} + +// Tool is a function callable by the LLM with JSON Schema input. +type Tool struct { + Name string + Description string + InputSchema json.RawMessage // JSON Schema draft-07+ + Required []string + Handler ToolHandler + Permission Permission + Examples []ToolExample +} + +// ToolHandler is the function type that implements a tool. +// The handler receives raw JSON arguments and returns a result. +type ToolHandler func(ctx context.Context, args json.RawMessage) (ToolResult, error) + +// ToolResult is returned by a ToolHandler. +type ToolResult struct { + Content string + IsError bool + Metadata map[string]string + Artifacts []Artifact +} + +// Artifact represents a file or data artifact produced by a tool. +type Artifact struct { + Path string + Content []byte + MIME string +} + +// ToolExample provides few-shot examples for the LLM to improve tool usage. +type ToolExample struct { + Input map[string]interface{} + Output string +} + +// Registry manages tool registration and lookup. +type Registry interface { + Register(tool Tool) error + Get(name string) (Tool, bool) + List() []Tool + Filter(predicate func(Tool) bool) []Tool +} + +// ErrToolNotFound is returned when a tool is not found in the registry. +var ErrToolNotFound = fmt.Errorf("tool not found") + +// ErrDuplicateTool is returned when trying to register a tool with an existing name. +var ErrDuplicateTool = fmt.Errorf("duplicate tool name") From 0054ca793c7530fe7f18ee1e2c81fdc60c3935a9 Mon Sep 17 00:00:00 2001 From: Victor Vargas Date: Tue, 30 Jun 2026 23:53:34 -0700 Subject: [PATCH 08/11] docs: add AGENTS.md guide and skill definitions for AI agents --- .agents/skills/add-feature/SKILL.md | 121 ++++++++++++++++++++++++++++ AGENTS.md | 74 +++++++++++++++++ 2 files changed, 195 insertions(+) create mode 100644 .agents/skills/add-feature/SKILL.md create mode 100644 AGENTS.md diff --git a/.agents/skills/add-feature/SKILL.md b/.agents/skills/add-feature/SKILL.md new file mode 100644 index 0000000..df76b90 --- /dev/null +++ b/.agents/skills/add-feature/SKILL.md @@ -0,0 +1,121 @@ +--- +name: add-feature +description: >- + Use ONLY when adding new features to rony-llm-agent or its downstream products + (rony-harness, rony-chat-bot). This skill enforces adherence to the architecture + spec and asks probing questions before any implementation. Trigger when the user + says "add feature", "implement X", "build Y", or describes a new capability not + yet in scope. +--- + +# Add Feature — Architecture-First Agent Skill + +## Purpose + +Guide the addition of new features to `rony-llm-agent` (or its consumers) while +**strictly following** [`docs/architecture.md`](../../docs/architecture.md). + +Before writing any code, you **must ask questions** about every aspect of the feature +that could affect: package boundaries, interfaces, dependencies, or security model. + +## Pre-Implementation Checklist + +### 1. Scope & Package Placement + +Read first, in this order: +1. `docs/architecture.md` — hexagonal layers, core interfaces, security model +2. `docs/components.md` — per-package boundaries and public APIs +3. `docs/phase2.md` — check if the feature overlaps or conflicts with Phase 2 backlog +4. `AGENTS.md` — code conventions (naming, errors, context usage) + +Then answer **before** proposing implementation: + +- [ ] Which package(s) need new interfaces vs existing ones? +- [ ] Does this create a new top-level `pkg//` or fit inside an existing one? +- [ ] What is the public-facing interface (port)? Can it be written as a pure Go interface? +- [ ] Is there an existing adapter that can extend, or does it need a new one? + +### 2. Interface Design + +For every new or modified interface: + +- Does the interface name follow capability naming? (`LLMClient`, `Loop`, `Embedder`, `Memory`) +- Does every method take `ctx context.Context` as first parameter? +- Are return types Go-native (structs, `iter.Seq2` for streams), not callbacks? +- Is the interface minimal — only what consumers actually need? + +### 3. Dependencies & Package Layers + +- Does the domain package (`pkg//`) import anything external except stdlib? + **NO.** External SDKs live in adapters under `internal/`. +- If a new package is needed, does its `import` path follow `github.com/VictorVargas/rony-llm-agent/pkg/`? +- Are adapters isolated behind interfaces? No adapter should leak into domain code. + +### 4. Security Model + +Refer to `docs/architecture.md` §4 — "Modelo de Seguridad": + +- Does the feature introduce new filesystem access? → Must use `os.Root` sandbox, never raw paths. +- Does it execute external commands? → Command validation needed; delegate to products. +- Does it accept user input into LLM prompts? → Consider `` wrapping (Phase 2). +- Are there permission implications for tools? → Use `Permission: Allow | Ask | Deny` policy. +- Is the feature compliant with least-privilege principle? + +### 5. Testing Strategy + +Before implementation, identify: + +- What existing mock can be reused? (`mock.MockLLMClient`, `mock.MockMemory`) +- Are there critical-path tests to add? See `docs/architecture.md` §6 — "Tabla de tests críticos" +- Is the test deterministic (no network calls) or integration (real provider/backend)? +- If streaming, does it properly consume the full `iter.Seq2` channel without goroutine leaks? + +### 6. Concurrency & Context Propagation + +- Does every blocking call take a context? +- Are contexts properly cancelled on timeout/interrupt? +- Is the new code safe for concurrent use by multiple products simultaneously? + +## Questioning Protocol + +**Never assume.** If an answer isn't explicitly in architecture.md, components.md, or phase2.md, ask: + +1. "Where does this belong architecturally?" — before writing anything +2. "What interface does this consume or provide?" — define ports first +3. "How does this interact with the sandbox/security model?" — always check permissions +4. "Which existing package depends on this?" — verify no circular deps +5. "What's the failure mode and error path?" — use `ErrXXX` sentinels + `%w` wrapping + +## Constraints (Never Break) + +| Constraint | Why | +|---|---| +| Go 1.26+ required | Uses `os.Root`, `iter.Seq2`, `unique.Handle` | +| No reflection, no codegen, no DSLs | "Zero magic" principle | +| Ports first, implementations second | Hexagonal architecture — ports are pure interfaces | +| All blocking takes `ctx context.Context` | Cancellation and timeout support | +| Sandbox with `os.Root`, not path prefix checks | Kernel-level guarantee against symlinks, TOCTOU, encoding attacks | +| Downstream-first validation | Check what harness/chat-bot does before implementing | + +## When to Defer to Phase 2 + +Do NOT implement these unless explicitly requested: + +- MCP server/client protocol +- Full RAG pipeline (beyond existing `pkg/rag/`) +- Skills system for agent behavior customization +- Sub-agents and hierarchical prompting +- Observability/tracing/export + +Reference `docs/phase2.md` if the feature overlaps. + +## Implementation Checklist + +After questions are answered: + +1. Define interfaces at top of package (`pkg//interface.go` or similar) +2. Create domain logic in `pkg//` (no external deps) +3. Implement adapters in `pkg//internal/` or provider-specific dirs +4. Add tests using mocks where possible +5. Run `go vet ./...`, then `go test ./...` once code exists +6. Document public interfaces in package-level godoc \ No newline at end of file diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 0000000..67d2793 --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,74 @@ +# AGENTS.md — Working in rony-llm-agent + +## Repo status: **Design/spec only** — no Go source exists yet. All code described in docs (architecture.md, components.md, phase2.md) is aspirational. The README explicitly states this library hasn't been implemented; once `harness/` ships, the lib will be extracted as real code following these specs. + +## What to read first + +| If you need... | Read... | +|---|---| +| Package layout and entrypoints | [`docs/architecture.md`](./docs/architecture.md) | +| Per-package boundaries | [`docs/components.md`](./docs/components.md) | +| Features not yet planned | [`docs/phase2.md`](./docs/phase2.md) — Phase 2 backlog | + +These three docs are the source of truth. **Read them before writing any code.** The `pkg/*/README.md` files mirror content from these docs; they're convenient but architecture.md is canonical. + +## How to build and test (once Go code exists) + +```bash +# Requires Go 1.26+ — required for os.Root, iter.Seq2, unique.Handle +go mod tidy # resolves imports, creates go.sum +go test ./... # all packages +go test -race ./... # race detector (always use in CI once implemented) +``` + +No Makefile, no linters configured yet. Once code exists: lint → typecheck → test is the expected order. Run `golangci-lint` if installed. + +## Key constraints to never break + +### Go version +Requires **Go 1.26+**. The module declares it in go.mod. Using older Go will not compile because of `os.Root`, `iter.Seq2`, `unique.Handle`. + +### Hexagonal architecture — package boundaries +- Interfaces (ports) live at the top level of each `pkg//` directory. These define the public API. +- Implementations (adapters) go in `pkg//internal/` or subdirectories like `pkg/llm/providers/openai/`, `pkg/rag/backends/chroma/`. +- Package names are lowercase, singular (`agent`, `tools`, `llm`, `rag`, `persona`). +- No external dependencies in domain packages — only stdlib and SDKs of providers (in adapters). + +### Module path +`github.com/VictorVargas/rony-llm-agent` — imported as-is by downstream products. + +### Streaming uses `iter.Seq2` +The library is designed around Go 1.23+'s `iter.Seq2[T, error]` for streaming LLM output. When implementing or reading code, this is the primary stream pattern. + +### Sandbox uses `os.Root` (Go 1.24+) +Filesystem sandboxing must use `os.Root`, not path string prefix checks. This is non-negotiable for security — naive prefix checks can't handle symlinks, TOCTOU, or path encoding attacks. + +## Testing conventions (when code exists) + +- Use `pkg/llm/mock.MockLLMClient` for deterministic tests that don't call real APIs. +- Use `pkg/rag/mock.MockMemory` for memory tests. +- All public API must be concurrent-safe (documented in architecture.md). +- Every function that can block takes `ctx context.Context` as the first parameter. + +## Code style conventions + +- Interfaces end with capability names: `LLMClient`, `Loop`, `Embedder`, `Memory`. +- Sentinel errors prefixed with `Err`: `ErrToolNotFound`, `ErrSandboxViolation`. +- Constructors use `New` for the primary and `NewXxx` for variants. +- Error wrapping uses `%w`, never lossy formatting. + +## Relationship to downstream products + +This library is consumed by: +- **rony-harness** — TUI agent for software development (the reference implementation that triggered this extraction) +- **rony-chat-bot** — HTTP chatbot for portfolios/websites + +If you're unsure about behavior, check what harness or chat-bot does first. They are the real-world consumers driving design decisions. + +## Agent skills + +Reusable skills for any AI agent live in `.agents/skills//SKILL.md` — this is the canonical, tool-agnostic location read by OpenCode, Claude Code, Cursor, and other modern agents. Each skill's `SKILL.md` starts with a YAML frontmatter (`name`, `description`) followed by instructions. Do not mirror skills into `.opencode/skills/` — opencode reads `.agents/skills/` natively. + +## Phase 2 awareness + +Phase 2 features (MCP server/client, full RAG pipeline, skills system, sub-agents, observability) are planned but not in scope for initial implementation. Do not start implementing phase 2 code unless explicitly asked. Reference `docs/phase2.md` for spec when needed. From 1a8f1557f660c510a5db9719056122dca2b5b6d8 Mon Sep 17 00:00:00 2001 From: Victor Vargas Date: Fri, 3 Jul 2026 14:22:36 -0700 Subject: [PATCH 09/11] feat(llm): add ChatTemplateKwargs to CompletionRequest for provider-specific template params - Add ChatTemplateKwargs field to llm.CompletionRequest - Propagate kwargs through agent loop in both Run() and RunStream() - Pass kwargs to llama.cpp client chat request - Fix tool schema marshaling to include type/function wrapper - Fix stream indentation logic in RunStream with responseBuilder - Remove indirect marker from uuid dependency --- go.mod | 2 +- pkg/agent/loop.go | 53 +++++++++++++++++----------- pkg/llm/providers/llamacpp/client.go | 28 ++++++++------- pkg/llm/types.go | 17 ++++----- 4 files changed, 58 insertions(+), 42 deletions(-) diff --git a/go.mod b/go.mod index f4ccdc1..091c508 100644 --- a/go.mod +++ b/go.mod @@ -4,4 +4,4 @@ go 1.26 require gopkg.in/yaml.v3 v3.0.1 -require github.com/google/uuid v1.6.0 // indirect +require github.com/google/uuid v1.6.0 diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go index aa988b1..2321200 100644 --- a/pkg/agent/loop.go +++ b/pkg/agent/loop.go @@ -5,6 +5,7 @@ import ( "encoding/json" "fmt" "iter" + "strings" "time" "github.com/VictorVargas/rony-llm-agent/pkg/llm" @@ -32,14 +33,15 @@ type OnIterationHook func(Iteration) // Config holds the dependencies and settings for the agent loop. type Config struct { - LLM llm.LLMClient - Persona persona.Persona - Tools tools.Registry - Sandbox Sandbox - MaxIters int - Approver Approver - OnIteration OnIterationHook - ToolTimeout time.Duration + LLM llm.LLMClient + Persona persona.Persona + Tools tools.Registry + Sandbox Sandbox + MaxIters int + Approver Approver + OnIteration OnIterationHook + ToolTimeout time.Duration + ChatTemplateKwargs map[string]any // passed to the LLM provider (e.g. Qwen enable_thinking) } // Iteration represents a single cycle of the agent loop. @@ -86,8 +88,9 @@ func (l *Loop) Run(ctx context.Context, input string) (Response, error) { iterations++ resp, err := l.cfg.LLM.Generate(ctx, llm.CompletionRequest{ - Messages: messages, - Tools: l.getToolSchemas(), + Messages: messages, + Tools: l.getToolSchemas(), + ChatTemplateKwargs: l.cfg.ChatTemplateKwargs, }) if err != nil { return Response{}, fmt.Errorf("LLM generate failed: %w", err) @@ -143,11 +146,13 @@ func (l *Loop) RunStream(ctx context.Context, input string) iter.Seq2[llm.Stream iterations++ stream := l.cfg.LLM.Stream(ctx, llm.CompletionRequest{ - Messages: messages, - Tools: l.getToolSchemas(), + Messages: messages, + Tools: l.getToolSchemas(), + ChatTemplateKwargs: l.cfg.ChatTemplateKwargs, }) var hasToolCalls bool + var responseBuilder strings.Builder for chunk, err := range stream { if err != nil { yield(llm.StreamChunk{}, err) @@ -172,21 +177,20 @@ func (l *Loop) RunStream(ctx context.Context, input string) iter.Seq2[llm.Stream } } - if !hasToolCalls && chunk.Delta != "" { - if !yield(chunk, nil) { - return + if !hasToolCalls && chunk.Delta != "" { + responseBuilder.WriteString(chunk.Delta) + if !yield(chunk, nil) { + return + } } } - } if !hasToolCalls { return } } - if iterations >= l.cfg.MaxIters { - yield(llm.StreamChunk{}, fmt.Errorf("max iterations (%d) reached", l.cfg.MaxIters)) - } + yield(llm.StreamChunk{}, fmt.Errorf("max iterations (%d) reached", l.cfg.MaxIters)) } } @@ -201,7 +205,16 @@ func (l *Loop) buildInitialMessages(input string) []llm.Message { func (l *Loop) getToolSchemas() []json.RawMessage { var schemas []json.RawMessage for _, tool := range l.cfg.Tools.List() { - schemas = append(schemas, tool.InputSchema) + def := map[string]interface{}{ + "type": "function", + "function": map[string]interface{}{ + "name": tool.Name, + "description": tool.Description, + "parameters": json.RawMessage(tool.InputSchema), + }, + } + data, _ := json.Marshal(def) + schemas = append(schemas, data) } return schemas } diff --git a/pkg/llm/providers/llamacpp/client.go b/pkg/llm/providers/llamacpp/client.go index 314803e..0135ca9 100644 --- a/pkg/llm/providers/llamacpp/client.go +++ b/pkg/llm/providers/llamacpp/client.go @@ -173,9 +173,10 @@ func (c *Client) buildRequest(req llm.CompletionRequest, stream bool) (io.Reader } openReq := llamaChatRequest{ - Model: req.Model, - Messages: messages, - Stream: stream, + Model: req.Model, + Messages: messages, + Stream: stream, + ChatTemplateKwargs: req.ChatTemplateKwargs, } if len(tools) > 0 { openReq.Tools = tools @@ -232,16 +233,17 @@ func (c *Client) toResponse(resp llamaChatResponse) llm.CompletionResponse { // llama.cpp API types type llamaChatRequest struct { - Model string `json:"model"` - Messages []llamaMessage `json:"messages"` - Tools []llamaTool `json:"tools,omitempty"` - ToolChoice interface{} `json:"tool_choice,omitempty"` - Temperature float32 `json:"temperature,omitempty"` - MaxTokens int `json:"max_tokens,omitempty"` - TopK int `json:"top_k,omitempty"` - TopP float32 `json:"top_p,omitempty"` - Stop []string `json:"stop,omitempty"` - Stream bool `json:"stream"` + Model string `json:"model"` + Messages []llamaMessage `json:"messages"` + Tools []llamaTool `json:"tools,omitempty"` + ToolChoice interface{} `json:"tool_choice,omitempty"` + Temperature float32 `json:"temperature,omitempty"` + MaxTokens int `json:"max_tokens,omitempty"` + TopK int `json:"top_k,omitempty"` + TopP float32 `json:"top_p,omitempty"` + Stop []string `json:"stop,omitempty"` + Stream bool `json:"stream"` + ChatTemplateKwargs map[string]any `json:"chat_template_kwargs,omitempty"` } type llamaMessage struct { diff --git a/pkg/llm/types.go b/pkg/llm/types.go index f7d9fa8..366a9c2 100644 --- a/pkg/llm/types.go +++ b/pkg/llm/types.go @@ -66,14 +66,15 @@ func (t *ToolRef) MarshalJSON() ([]byte, error) { // CompletionRequest is sent to an LLM provider. type CompletionRequest struct { - Model string `json:"model"` - Messages []Message `json:"messages"` - Tools []json.RawMessage `json:"tools,omitempty"` - ToolChoice interface{} `json:"tool_choice,omitempty"` // ToolChoice, ToolRef, or null - Temperature *float32 `json:"temperature,omitempty"` - MaxTokens *int `json:"max_tokens,omitempty"` - Stop []string `json:"stop,omitempty"` - Metadata map[string]string `json:"metadata,omitempty"` + Model string `json:"model"` + Messages []Message `json:"messages"` + Tools []json.RawMessage `json:"tools,omitempty"` + ToolChoice interface{} `json:"tool_choice,omitempty"` // ToolChoice, ToolRef, or null + Temperature *float32 `json:"temperature,omitempty"` + MaxTokens *int `json:"max_tokens,omitempty"` + Stop []string `json:"stop,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 } // CompletionResponse is returned from an LLM provider. From 2eed2033f071257c71fea22b3bff5552d27db79e Mon Sep 17 00:00:00 2001 From: Victor Vargas Date: Fri, 3 Jul 2026 14:22:41 -0700 Subject: [PATCH 10/11] test(llm,embeddings): add unit tests for mock client, types, and Ollama embedder - Add comprehensive MockLLMClient tests (generate, stream, match variants) - Add TypeRef JSON marshaling tests and StopReason value tests - Add Ollama embedder tests for config defaults and embedding requests --- pkg/llm/mock/mock_test.go | 237 ++++++++++++++++++++++++++ pkg/llm/types_test.go | 145 ++++++++++++++++ pkg/rag/embeddings/embeddings_test.go | 113 ++++++++++++ 3 files changed, 495 insertions(+) create mode 100644 pkg/llm/mock/mock_test.go create mode 100644 pkg/llm/types_test.go create mode 100644 pkg/rag/embeddings/embeddings_test.go diff --git a/pkg/llm/mock/mock_test.go b/pkg/llm/mock/mock_test.go new file mode 100644 index 0000000..c049aeb --- /dev/null +++ b/pkg/llm/mock/mock_test.go @@ -0,0 +1,237 @@ +package mock + +import ( + "context" + "fmt" + "testing" + + "github.com/VictorVargas/rony-llm-agent/pkg/llm" +) + +func TestMockLLMClient_Generate_Default(t *testing.T) { + client := New() + resp, err := client.Generate(context.Background(), llm.CompletionRequest{}) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if resp.Content != "default response" { + t.Errorf("expected 'default response', got %q", resp.Content) + } + if resp.StopReason != llm.StopReasonEndTurn { + t.Errorf("expected StopReasonEndTurn, got %q", resp.StopReason) + } +} + +func TestMockLLMClient_Generate_Custom(t *testing.T) { + expected := llm.CompletionResponse{ + ID: "test_123", + Model: "gpt-4", + Content: "hello world", + StopReason: llm.StopReasonToolUse, + } + client := NewWithGenerate(expected) + + resp, err := client.Generate(context.Background(), llm.CompletionRequest{}) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if resp.ID != "test_123" { + t.Errorf("expected ID 'test_123', got %q", resp.ID) + } + if resp.Content != "hello world" { + t.Errorf("expected 'hello world', got %q", resp.Content) + } + if resp.StopReason != llm.StopReasonToolUse { + t.Errorf("expected StopReasonToolUse, got %q", resp.StopReason) + } +} + +func TestMockLLMClient_Generate_Error(t *testing.T) { + client := &MockLLMClient{ + GenerateFunc: func(ctx context.Context, req llm.CompletionRequest) (llm.CompletionResponse, error) { + return llm.CompletionResponse{}, fmt.Errorf("test error") + }, + } + + _, err := client.Generate(context.Background(), llm.CompletionRequest{}) + if err == nil { + t.Fatal("expected error, got nil") + } + if err.Error() != "test error" { + t.Errorf("expected 'test error', got %q", err.Error()) + } +} + +func TestMockLLMClient_Stream_Default(t *testing.T) { + client := New() + var chunks []string + for chunk, err := range client.Stream(context.Background(), llm.CompletionRequest{}) { + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + chunks = append(chunks, chunk.Delta) + } + if len(chunks) != 1 || chunks[0] != "default response" { + t.Errorf("expected single 'default response' chunk, got %v", chunks) + } +} + +func TestMockLLMClient_Stream_Custom(t *testing.T) { + expectedChunks := []llm.StreamChunk{ + {Delta: "chunk1"}, + {Delta: "chunk2", ToolCalls: []llm.ToolCall{{ID: "call_1"}}}, + {FinishReason: "stop", Usage: llm.TokenUsage{InputTokens: 5, OutputTokens: 10}}, + } + client := NewWithStream(expectedChunks) + + var actualChunks []llm.StreamChunk + for chunk, err := range client.Stream(context.Background(), llm.CompletionRequest{}) { + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + actualChunks = append(actualChunks, chunk) + } + + if len(actualChunks) != 3 { + t.Errorf("expected 3 chunks, got %d", len(actualChunks)) + return + } + if actualChunks[0].Delta != "chunk1" { + t.Errorf("expected 'chunk1', got %q", actualChunks[0].Delta) + } + if len(actualChunks[1].ToolCalls) != 1 { + t.Errorf("expected 1 tool call, got %d", len(actualChunks[1].ToolCalls)) + } + if actualChunks[2].FinishReason != "stop" { + t.Errorf("expected 'stop', got %q", actualChunks[2].FinishReason) + } +} + +func TestMockLLMClient_Name_Default(t *testing.T) { + client := New() + if name := client.Name(); name != "mock" { + t.Errorf("expected 'mock', got %q", name) + } +} + +func TestMockLLMClient_Name_Custom(t *testing.T) { + client := &MockLLMClient{ + NameFunc: func() string { return "custom-model" }, + } + if name := client.Name(); name != "custom-model" { + t.Errorf("expected 'custom-model', got %q", name) + } +} + +func TestMockLLMClient_Capabilities_Default(t *testing.T) { + client := New() + caps := client.Capabilities() + if !caps.SupportsTools { + t.Error("expected SupportsTools to be true") + } + if caps.SupportsVision { + t.Error("expected SupportsVision to be false") + } + if !caps.SupportsJSON { + t.Error("expected SupportsJSON to be true") + } + if caps.MaxContextWindow != 128000 { + t.Errorf("expected MaxContextWindow 128000, got %d", caps.MaxContextWindow) + } +} + +func TestMockLLMClient_Capabilities_Custom(t *testing.T) { + client := &MockLLMClient{ + CapabilitiesFunc: func() llm.ProviderCapabilities { + return llm.ProviderCapabilities{ + SupportsTools: false, + SupportsVision: true, + SupportsJSON: false, + MaxContextWindow: 32000, + } + }, + } + caps := client.Capabilities() + if caps.SupportsTools { + t.Error("expected SupportsTools to be false") + } + if !caps.SupportsVision { + t.Error("expected SupportsVision to be true") + } + if caps.MaxContextWindow != 32000 { + t.Errorf("expected MaxContextWindow 32000, got %d", caps.MaxContextWindow) + } +} + +func TestMockLLMClient_MatchResponse(t *testing.T) { + client := NewWithMatch([]MatchResponse{ + {Match: "hello", Response: "hi there!"}, + {Match: "world", Response: "earth"}, + {Match: "*", Response: "default"}, + }) + + resp, _ := client.Generate(context.Background(), llm.CompletionRequest{ + Messages: []llm.Message{{Role: llm.RoleUser, Content: "hello"}}, + }) + if resp.Content != "hi there!" { + t.Errorf("expected 'hi there!', got %q", resp.Content) + } + + resp, _ = client.Generate(context.Background(), llm.CompletionRequest{ + Messages: []llm.Message{{Role: llm.RoleUser, Content: "world"}}, + }) + if resp.Content != "earth" { + t.Errorf("expected 'earth', got %q", resp.Content) + } + + resp, _ = client.Generate(context.Background(), llm.CompletionRequest{ + Messages: []llm.Message{{Role: llm.RoleUser, Content: "anything else"}}, + }) + if resp.Content != "default" { + t.Errorf("expected 'default', got %q", resp.Content) + } +} + +func TestMockLLMClient_MatchResponse_NoMatch(t *testing.T) { + client := NewWithMatch([]MatchResponse{ + {Match: "exact", Response: "matched"}, + }) + + resp, _ := client.Generate(context.Background(), llm.CompletionRequest{ + Messages: []llm.Message{{Role: llm.RoleUser, Content: "no match here"}}, + }) + if resp.Content != "no match" { + t.Errorf("expected 'no match', got %q", resp.Content) + } +} + +func TestMatchesAny(t *testing.T) { + if !matchesAny("hello", []llm.Message{ + {Role: llm.RoleUser, Content: "say hello to me"}, + }) { + t.Error("expected matchesAny to find 'hello' in message") + } + + if matchesAny("goodbye", []llm.Message{ + {Role: llm.RoleUser, Content: "say hello to me"}, + }) { + t.Error("expected matchesAny to not find 'goodbye' in message") + } +} + +func TestMatchesAny_Wildcard(t *testing.T) { + resp, _ := (&MockLLMClient{ + GenerateFunc: func(ctx context.Context, req llm.CompletionRequest) (llm.CompletionResponse, error) { + for _, r := range []MatchResponse{{Match: "*", Response: "wildcard"}} { + if r.Match == "*" || matchesAny(r.Match, req.Messages) { + return llm.CompletionResponse{Content: r.Response}, nil + } + } + return llm.CompletionResponse{}, nil + }, + }).Generate(context.Background(), llm.CompletionRequest{}) + + if resp.Content != "wildcard" { + t.Errorf("expected 'wildcard', got %q", resp.Content) + } +} diff --git a/pkg/llm/types_test.go b/pkg/llm/types_test.go new file mode 100644 index 0000000..6747890 --- /dev/null +++ b/pkg/llm/types_test.go @@ -0,0 +1,145 @@ +package llm + +import ( + "encoding/json" + "testing" +) + +func TestToolRef_MarshalJSON(t *testing.T) { + ref := &ToolRef{Name: "my_tool"} + data, err := json.Marshal(ref) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + var result struct { + Type string `json:"type"` + Name string `json:"name"` + } + if err := json.Unmarshal(data, &result); err != nil { + t.Fatalf("failed to unmarshal: %v", err) + } + if result.Type != "function" { + t.Errorf("expected type 'function', got %q", result.Type) + } + if result.Name != "my_tool" { + t.Errorf("expected name 'my_tool', got %q", result.Name) + } +} + +func TestStopReason_Values(t *testing.T) { + if string(StopReasonEndTurn) != "end_turn" { + t.Error("StopReasonEndTurn should be 'end_turn'") + } + if string(StopReasonToolUse) != "tool_use" { + t.Error("StopReasonToolUse should be 'tool_use'") + } + if string(StopReasonMaxTokens) != "max_tokens" { + t.Error("StopReasonMaxTokens should be 'max_tokens'") + } + if string(StopReasonStopSeq) != "stop_sequence" { + t.Error("StopReasonStopSeq should be 'stop_sequence'") + } +} + +func TestCompletionRequest_JSON(t *testing.T) { + req := CompletionRequest{ + Model: "gpt-4", + Temperature: float32Ptr(0.7), + MaxTokens: intPtr(100), + Metadata: map[string]string{"key": "val"}, + } + + data, err := json.Marshal(req) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + var unmarshaled CompletionRequest + if err := json.Unmarshal(data, &unmarshaled); err != nil { + t.Fatalf("failed to unmarshal: %v", err) + } + if unmarshaled.Model != "gpt-4" { + t.Errorf("expected model 'gpt-4', got %q", unmarshaled.Model) + } +} + +func TestCompletionResponse_JSON(t *testing.T) { + resp := CompletionResponse{ + ID: "resp_123", + Model: "gpt-4", + Content: "hello world", + StopReason: StopReasonEndTurn, + Usage: TokenUsage{ + InputTokens: 10, + OutputTokens: 5, + TotalTokens: 15, + }, + } + + data, err := json.Marshal(resp) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + var unmarshaled CompletionResponse + if err := json.Unmarshal(data, &unmarshaled); err != nil { + t.Fatalf("failed to unmarshal: %v", err) + } + if unmarshaled.ID != "resp_123" { + t.Errorf("expected ID 'resp_123', got %q", unmarshaled.ID) + } + if unmarshaled.Content != "hello world" { + t.Errorf("expected content 'hello world', got %q", unmarshaled.Content) + } +} + +func TestToolCall_JSON(t *testing.T) { + call := ToolCall{ + ID: "call_1", + Name: "my_tool", + Arguments: json.RawMessage(`{"key": "val"}`), + Thought: "let me think", + } + + data, err := json.Marshal(call) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + var unmarshaled ToolCall + if err := json.Unmarshal(data, &unmarshaled); err != nil { + t.Fatalf("failed to unmarshal: %v", err) + } + if unmarshaled.ID != "call_1" { + t.Errorf("expected ID 'call_1', got %q", unmarshaled.ID) + } + if string(unmarshaled.Arguments) != `{"key":"val"}` { + t.Errorf("unexpected arguments: %s", unmarshaled.Arguments) + } +} + +func TestStreamChunk_JSON(t *testing.T) { + chunk := StreamChunk{ + Delta: "hello", + ToolCalls: []ToolCall{}, + } + + data, err := json.Marshal(chunk) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + var unmarshaled StreamChunk + if err := json.Unmarshal(data, &unmarshaled); err != nil { + t.Fatalf("failed to unmarshal: %v", err) + } + if unmarshaled.Delta != "hello" { + t.Errorf("expected delta 'hello', got %q", unmarshaled.Delta) + } +} + + + +func float32Ptr(f float32) *float32 { return &f } +func intPtr(i int) *int { return &i } diff --git a/pkg/rag/embeddings/embeddings_test.go b/pkg/rag/embeddings/embeddings_test.go new file mode 100644 index 0000000..a688b5c --- /dev/null +++ b/pkg/rag/embeddings/embeddings_test.go @@ -0,0 +1,113 @@ +package embeddings + +import ( + "context" + "net/http" + "net/http/httptest" + "testing" +) + +func TestNewOllama_Defaults(t *testing.T) { + e, err := NewOllama(Config{}) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if e.baseURL != "http://localhost:11434" { + t.Errorf("expected 'http://localhost:11434', got %q", e.baseURL) + } + if e.model != "nomic-embed-text" { + t.Errorf("expected 'nomic-embed-text', got %q", e.model) + } +} + +func TestNewOllama_Custom(t *testing.T) { + e, err := NewOllama(Config{ + BaseURL: "http://custom:8080", + Model: "my-embedder", + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if e.baseURL != "http://custom:8080" { + t.Errorf("expected 'http://custom:8080', got %q", e.baseURL) + } + if e.model != "my-embedder" { + t.Errorf("expected 'my-embedder', got %q", e.model) + } +} + +func TestOllama_Dimensions(t *testing.T) { + e, _ := NewOllama(Config{}) + if dims := e.Dimensions(); dims != 768 { + t.Errorf("expected 768 dimensions, got %d", dims) + } +} + +func TestOllama_Embed_Success(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + w.Write([]byte(`{"embedding": [0.1, 0.2, 0.3]}`)) + })) + defer server.Close() + + e, _ := NewOllama(Config{BaseURL: server.URL}) + ctx := context.Background() + vector, err := e.Embed(ctx, "test input") + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if len(vector) != 3 { + t.Errorf("expected 3 dimensions, got %d", len(vector)) + } + if vector[0] != 0.1 || vector[1] != 0.2 || vector[2] != 0.3 { + t.Errorf("unexpected vector: %v", vector) + } +} + +func TestOllama_Embed_Error(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusInternalServerError) + w.Write([]byte("internal error")) + })) + defer server.Close() + + e, _ := NewOllama(Config{BaseURL: server.URL}) + ctx := context.Background() + _, err := e.Embed(ctx, "test input") + if err == nil { + t.Fatal("expected error, got nil") + } +} + +func TestMockEmbedder_Default(t *testing.T) { + m := &MockEmbedder{} + vector, err := m.Embed(context.Background(), "test") + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if len(vector) != 3 { + t.Errorf("expected 3 dimensions, got %d", len(vector)) + } + if dims := m.Dimensions(); dims != 3 { + t.Errorf("expected 3 dimensions, got %d", dims) + } +} + +func TestMockEmbedder_Custom(t *testing.T) { + m := &MockEmbedder{ + EmbedFunc: func(ctx context.Context, text string) ([]float32, error) { + return []float32{1.0, 2.0}, nil + }, + DimensionsFn: func() int { + return 2 + }, + } + vector, _ := m.Embed(context.Background(), "test") + if len(vector) != 2 { + t.Errorf("expected 2 dimensions, got %d", len(vector)) + } + if dims := m.Dimensions(); dims != 2 { + t.Errorf("expected 2 dimensions, got %d", dims) + } +} From bffaecb5795e57480a6ba0f6b883a3e4a72d43d1 Mon Sep 17 00:00:00 2001 From: Victor Vargas Date: Fri, 3 Jul 2026 14:22:46 -0700 Subject: [PATCH 11/11] docs(AGENTS): update repo status to reflect Go source code exists - Change status from 'design/spec only' to actual implementation - Add pkg/*/README.md as a reference in the docs table - Update build/test section for existing codebase - Add test coverage summary table - Update testing conventions with current mock types - Add ErrPersonaNotFound and ErrConfigNotFound to sentinel errors --- AGENTS.md | 30 ++++++++++++++++++++++-------- 1 file changed, 22 insertions(+), 8 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 67d2793..87ec800 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -1,6 +1,6 @@ # AGENTS.md — Working in rony-llm-agent -## Repo status: **Design/spec only** — no Go source exists yet. All code described in docs (architecture.md, components.md, phase2.md) is aspirational. The README explicitly states this library hasn't been implemented; once `harness/` ships, the lib will be extracted as real code following these specs. +## Repo status: **Go source code exists.** The library is implemented across `pkg/agent`, `pkg/config`, `pkg/llm`, `pkg/persona`, `pkg/rag`, and `pkg/tools`. All public interfaces are defined, core adapters for OpenAI, llama.cpp, ChromaDB, and Ollama embeddings are in place. Tests exist for the agent loop, sandbox, config loader, persona loader, RAG memory, tool registry, and LLM clients. ## What to read first @@ -9,10 +9,11 @@ | Package layout and entrypoints | [`docs/architecture.md`](./docs/architecture.md) | | Per-package boundaries | [`docs/components.md`](./docs/components.md) | | Features not yet planned | [`docs/phase2.md`](./docs/phase2.md) — Phase 2 backlog | +| A specific package's API | [`pkg//README.md`](./pkg/*/README.md) — mirrors docs; check code for actual signatures | These three docs are the source of truth. **Read them before writing any code.** The `pkg/*/README.md` files mirror content from these docs; they're convenient but architecture.md is canonical. -## How to build and test (once Go code exists) +## How to build and test ```bash # Requires Go 1.26+ — required for os.Root, iter.Seq2, unique.Handle @@ -21,7 +22,7 @@ go test ./... # all packages go test -race ./... # race detector (always use in CI once implemented) ``` -No Makefile, no linters configured yet. Once code exists: lint → typecheck → test is the expected order. Run `golangci-lint` if installed. +Run `golangci-lint` if installed. No Makefile configured yet. ## Key constraints to never break @@ -43,17 +44,30 @@ The library is designed around Go 1.23+'s `iter.Seq2[T, error]` for streaming LL ### Sandbox uses `os.Root` (Go 1.24+) Filesystem sandboxing must use `os.Root`, not path string prefix checks. This is non-negotiable for security — naive prefix checks can't handle symlinks, TOCTOU, or path encoding attacks. -## Testing conventions (when code exists) +## Testing conventions -- Use `pkg/llm/mock.MockLLMClient` for deterministic tests that don't call real APIs. -- Use `pkg/rag/mock.MockMemory` for memory tests. +- Use `pkg/llm/mock.MockLLMClient` (or its constructors `New()`, `NewWithGenerate()`, `NewWithStream()`, `NewWithMatch()`) for deterministic tests that don't call real APIs. +- Use `pkg/rag/embeddings/mock.MockEmbedder` for embedding tests. - All public API must be concurrent-safe (documented in architecture.md). - Every function that can block takes `ctx context.Context` as the first parameter. +### Current test coverage + +| Package | Test file | Purpose | +|---|---|---| +| `pkg/agent` | `loop_test.go`, `integration_test.go` | Agent loop iterations, tool execution, streaming | +| `pkg/config` | `config_test.go` | YAML loading, defaults, precedence | +| `pkg/llm/providers/openai` | `client_test.go` | OpenAI HTTP client | +| `pkg/llm/providers/llamacpp` | `client_test.go` | llama.cpp server adapter | +| `pkg/persona` | `persona_test.go` | Persona loading, system prompt assembly | +| `pkg/rag` | `memory_test.go` | Add/search/forget operations | +| `pkg/tools` | `registry_test.go`, `sandbox/sandbox_test.go` | Tool registration, sandbox validation | +| `pkg/rag/backends/chroma` | `chroma_test.go` | ChromaDB upsert and search | + ## Code style conventions - Interfaces end with capability names: `LLMClient`, `Loop`, `Embedder`, `Memory`. -- Sentinel errors prefixed with `Err`: `ErrToolNotFound`, `ErrSandboxViolation`. +- Sentinel errors prefixed with `Err`: `ErrToolNotFound`, `ErrSandboxViolation`, `ErrPersonaNotFound`, `ErrConfigNotFound`. - Constructors use `New` for the primary and `NewXxx` for variants. - Error wrapping uses `%w`, never lossy formatting. @@ -71,4 +85,4 @@ Reusable skills for any AI agent live in `.agents/skills//SKILL.md` — th ## Phase 2 awareness -Phase 2 features (MCP server/client, full RAG pipeline, skills system, sub-agents, observability) are planned but not in scope for initial implementation. Do not start implementing phase 2 code unless explicitly asked. Reference `docs/phase2.md` for spec when needed. +Phase 2 features (MCP server/client, full RAG pipeline with Qdrant/sqlite-vec backends, skills system, sub-agents, observability) are planned but not in scope for initial implementation. Do not start implementing phase 2 code unless explicitly asked. Reference `docs/phase2.md` for spec when needed.