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.