Merge pull request #8 from VictorVargas/fix/providers-agent-loop
Fix providers (OpenAI/llamacpp/Anthropic) and agent loop correctness
This commit is contained in:
commit
ee9319b823
19 changed files with 410 additions and 254 deletions
|
|
@ -82,17 +82,21 @@ func (l *Loop) Run(ctx context.Context, input string, history ...llm.Message) (R
|
|||
start := time.Now()
|
||||
|
||||
messages := l.buildInitialMessages(input, history)
|
||||
// Tool schemas don't change between iterations, so build the JSON once
|
||||
// per Run instead of re-marshaling every tool on every loop pass.
|
||||
toolSchemas := l.getToolSchemas()
|
||||
var finalContent string
|
||||
var allToolCalls []llm.ToolCall
|
||||
var totalUsage llm.TokenUsage
|
||||
iterations := 0
|
||||
completed := false
|
||||
|
||||
for iterations < l.cfg.MaxIters {
|
||||
iterations++
|
||||
|
||||
resp, err := l.cfg.LLM.Generate(ctx, llm.CompletionRequest{
|
||||
Messages: messages,
|
||||
Tools: l.getToolSchemas(),
|
||||
Tools: toolSchemas,
|
||||
ChatTemplateKwargs: l.cfg.ChatTemplateKwargs,
|
||||
})
|
||||
if err != nil {
|
||||
|
|
@ -105,6 +109,7 @@ func (l *Loop) Run(ctx context.Context, input string, history ...llm.Message) (R
|
|||
|
||||
if len(resp.ToolCalls) == 0 {
|
||||
finalContent = resp.Content
|
||||
completed = true
|
||||
break
|
||||
}
|
||||
|
||||
|
|
@ -138,7 +143,11 @@ func (l *Loop) Run(ctx context.Context, input string, history ...llm.Message) (R
|
|||
}
|
||||
|
||||
duration := time.Since(start)
|
||||
if iterations >= l.cfg.MaxIters {
|
||||
// Only report the max-iterations failure when the loop actually ran out
|
||||
// of budget without producing a final answer — an answer that arrives
|
||||
// exactly on the last allowed iteration is still a success (the old
|
||||
// `iterations >= MaxIters` check threw that valid response away).
|
||||
if !completed {
|
||||
return Response{}, fmt.Errorf("max iterations (%d) reached", l.cfg.MaxIters)
|
||||
}
|
||||
|
||||
|
|
@ -157,6 +166,8 @@ func (l *Loop) Run(ctx context.Context, input string, history ...llm.Message) (R
|
|||
func (l *Loop) RunStream(ctx context.Context, input string, history ...llm.Message) iter.Seq2[llm.StreamChunk, error] {
|
||||
return func(yield func(llm.StreamChunk, error) bool) {
|
||||
messages := l.buildInitialMessages(input, history)
|
||||
// Same as Run: the schemas are identical on every iteration.
|
||||
toolSchemas := l.getToolSchemas()
|
||||
iterations := 0
|
||||
|
||||
for iterations < l.cfg.MaxIters {
|
||||
|
|
@ -164,7 +175,7 @@ func (l *Loop) RunStream(ctx context.Context, input string, history ...llm.Messa
|
|||
|
||||
stream := l.cfg.LLM.Stream(ctx, llm.CompletionRequest{
|
||||
Messages: messages,
|
||||
Tools: l.getToolSchemas(),
|
||||
Tools: toolSchemas,
|
||||
ChatTemplateKwargs: l.cfg.ChatTemplateKwargs,
|
||||
})
|
||||
|
||||
|
|
@ -226,11 +237,23 @@ func (l *Loop) RunStream(ctx context.Context, input string, history ...llm.Messa
|
|||
// real token counts.
|
||||
hasContent := chunk.Delta != "" || chunk.ReasoningDelta != ""
|
||||
hasUsage := chunk.Usage.TotalTokens > 0
|
||||
if !hasToolCalls && (hasContent || hasUsage) {
|
||||
switch {
|
||||
case !hasToolCalls && (hasContent || hasUsage):
|
||||
responseBuilder.WriteString(chunk.Delta)
|
||||
if !yield(chunk, nil) {
|
||||
return
|
||||
}
|
||||
case hasToolCalls && hasUsage:
|
||||
// The content gate above exists to hide raw provider
|
||||
// deltas during a tool-call round, but it also swallowed
|
||||
// that round's token usage — so callers tracking context
|
||||
// occupancy (e.g. a UI's context bar deciding when to
|
||||
// compact) only ever saw the usage of the final,
|
||||
// tool-free round. Forward the usage on its own,
|
||||
// without the content.
|
||||
if !yield(llm.StreamChunk{Usage: chunk.Usage}, nil) {
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -7,8 +7,8 @@ import (
|
|||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"iter"
|
||||
"net/http"
|
||||
"strings"
|
||||
|
||||
"github.com/VictorVargas/rony-llm-agent/pkg/llm"
|
||||
|
|
@ -19,6 +19,13 @@ const (
|
|||
defaultModel = "claude-opus-4-8"
|
||||
defaultMaxTokens = 8192
|
||||
anthropicVersion = "2023-06-01"
|
||||
// defaultContextWindow is what Capabilities() reports when
|
||||
// Config.ContextWindow is unset: 200k tokens, the standard window for
|
||||
// Claude models. Callers use this number to decide when to compact
|
||||
// their conversation, so over-reporting it (the old code assumed 1M
|
||||
// for anything that wasn't haiku) meant compaction fired far too late
|
||||
// and requests started overflowing the real window.
|
||||
defaultContextWindow = 200000
|
||||
)
|
||||
|
||||
// Config holds the settings needed to create an Anthropic client.
|
||||
|
|
@ -27,6 +34,7 @@ type Config struct {
|
|||
Model string // defaults to claude-opus-4-8
|
||||
BaseURL string // defaults to https://api.anthropic.com/v1
|
||||
MaxTokens int // default max_tokens sent on every request (Anthropic requires one); 0 = defaultMaxTokens
|
||||
ContextWindow int // model's context window in tokens (0 = defaultContextWindow); raise it only for models/plans with an extended window
|
||||
Temperature *float32
|
||||
TopP *float32
|
||||
}
|
||||
|
|
@ -37,6 +45,7 @@ type Client struct {
|
|||
baseURL string
|
||||
model string
|
||||
maxTokens int
|
||||
contextWindow int
|
||||
temperature *float32
|
||||
topP *float32
|
||||
http *http.Client
|
||||
|
|
@ -63,11 +72,17 @@ func New(cfg Config) (*Client, error) {
|
|||
maxTokens = defaultMaxTokens
|
||||
}
|
||||
|
||||
contextWindow := cfg.ContextWindow
|
||||
if contextWindow == 0 {
|
||||
contextWindow = defaultContextWindow
|
||||
}
|
||||
|
||||
return &Client{
|
||||
apiKey: cfg.APIKey,
|
||||
baseURL: baseURL,
|
||||
model: model,
|
||||
maxTokens: maxTokens,
|
||||
contextWindow: contextWindow,
|
||||
temperature: cfg.Temperature,
|
||||
topP: cfg.TopP,
|
||||
http: http.DefaultClient,
|
||||
|
|
@ -257,15 +272,11 @@ func (c *Client) Name() string {
|
|||
}
|
||||
|
||||
func (c *Client) Capabilities() llm.ProviderCapabilities {
|
||||
maxContext := 1000000
|
||||
if strings.Contains(c.model, "haiku") {
|
||||
maxContext = 200000
|
||||
}
|
||||
return llm.ProviderCapabilities{
|
||||
SupportsTools: true,
|
||||
SupportsVision: true,
|
||||
SupportsJSON: true,
|
||||
MaxContextWindow: maxContext,
|
||||
MaxContextWindow: c.contextWindow,
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -2,12 +2,13 @@ package llamacpp
|
|||
|
||||
import (
|
||||
"bufio"
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"iter"
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
|
|
@ -123,7 +124,7 @@ func (c *Client) Generate(ctx context.Context, req llm.CompletionRequest) (llm.C
|
|||
return llm.CompletionResponse{}, fmt.Errorf("decoding response: %w", err)
|
||||
}
|
||||
|
||||
return c.toResponse(apiResp), nil
|
||||
return c.toResponse(apiResp)
|
||||
}
|
||||
|
||||
func (c *Client) Stream(ctx context.Context, req llm.CompletionRequest) iter.Seq2[llm.StreamChunk, error] {
|
||||
|
|
@ -191,6 +192,11 @@ func (c *Client) Stream(ctx context.Context, req llm.CompletionRequest) iter.Seq
|
|||
}
|
||||
|
||||
scanner := bufio.NewScanner(resp.Body)
|
||||
// A single SSE line can exceed bufio.Scanner's 64KB default cap
|
||||
// (e.g. a large tool-call arguments delta or a long reasoning
|
||||
// event), which would kill the stream mid-turn with "token too
|
||||
// long" — same headroom the anthropic client already reserves.
|
||||
scanner.Buffer(make([]byte, 0, 64*1024), 4*1024*1024)
|
||||
for scanner.Scan() {
|
||||
line := scanner.Text()
|
||||
if !strings.HasPrefix(line, "data: ") {
|
||||
|
|
@ -370,11 +376,16 @@ func (c *Client) buildRequest(req llm.CompletionRequest, stream bool) (io.Reader
|
|||
if err != nil {
|
||||
return nil, fmt.Errorf("marshaling request: %w", err)
|
||||
}
|
||||
return strings.NewReader(string(data)), nil
|
||||
return bytes.NewReader(data), nil
|
||||
}
|
||||
|
||||
// toResponse converts a llama.cpp API response to our CompletionResponse.
|
||||
func (c *Client) toResponse(resp llamaChatResponse) llm.CompletionResponse {
|
||||
func (c *Client) toResponse(resp llamaChatResponse) (llm.CompletionResponse, error) {
|
||||
// Guard against a 200 response with no choices (e.g. a misbehaving
|
||||
// server or proxy) — indexing Choices[0] blindly panics the whole app.
|
||||
if len(resp.Choices) == 0 {
|
||||
return llm.CompletionResponse{}, fmt.Errorf("llamacpp: response contained no choices")
|
||||
}
|
||||
choice := resp.Choices[0]
|
||||
result := llm.CompletionResponse{
|
||||
ID: resp.ID,
|
||||
|
|
@ -398,7 +409,7 @@ func (c *Client) toResponse(resp llamaChatResponse) llm.CompletionResponse {
|
|||
TotalTokens: resp.Usage.TotalTokens,
|
||||
}
|
||||
|
||||
return result
|
||||
return result, nil
|
||||
}
|
||||
|
||||
// llama.cpp API types
|
||||
|
|
|
|||
|
|
@ -2,12 +2,13 @@ package openai
|
|||
|
||||
import (
|
||||
"bufio"
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"iter"
|
||||
"net/http"
|
||||
"strings"
|
||||
|
||||
"github.com/VictorVargas/rony-llm-agent/pkg/llm"
|
||||
|
|
@ -23,6 +24,7 @@ type Config struct {
|
|||
// Client implements llm.LLMClient for OpenAI.
|
||||
type Client struct {
|
||||
apiKey string
|
||||
model string
|
||||
baseURL string
|
||||
http *http.Client
|
||||
}
|
||||
|
|
@ -40,6 +42,7 @@ func New(cfg Config) (*Client, error) {
|
|||
|
||||
return &Client{
|
||||
apiKey: cfg.APIKey,
|
||||
model: cfg.Model,
|
||||
baseURL: baseURL,
|
||||
http: http.DefaultClient,
|
||||
}, nil
|
||||
|
|
@ -48,7 +51,7 @@ func New(cfg Config) (*Client, error) {
|
|||
func (c *Client) Generate(ctx context.Context, req llm.CompletionRequest) (llm.CompletionResponse, error) {
|
||||
endpoint := c.baseURL + "/chat/completions"
|
||||
|
||||
payload, err := c.buildRequest(req)
|
||||
payload, err := c.buildRequest(req, false)
|
||||
if err != nil {
|
||||
return llm.CompletionResponse{}, fmt.Errorf("building request: %w", err)
|
||||
}
|
||||
|
|
@ -76,14 +79,14 @@ func (c *Client) Generate(ctx context.Context, req llm.CompletionRequest) (llm.C
|
|||
return llm.CompletionResponse{}, fmt.Errorf("decoding response: %w", err)
|
||||
}
|
||||
|
||||
return c.toResponse(apiResp), nil
|
||||
return c.toResponse(apiResp)
|
||||
}
|
||||
|
||||
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)
|
||||
payload, err := c.buildRequest(req, true)
|
||||
if err != nil {
|
||||
yield(llm.StreamChunk{}, fmt.Errorf("building request: %w", err))
|
||||
return
|
||||
|
|
@ -111,7 +114,45 @@ func (c *Client) Stream(ctx context.Context, req llm.CompletionRequest) iter.Seq
|
|||
return
|
||||
}
|
||||
|
||||
// toolCallAccum buffers one tool call's fragments as they stream in:
|
||||
// the SSE format sends the id/name in the first delta for a given
|
||||
// tool-call index and the (potentially large) arguments JSON in
|
||||
// pieces across many subsequent deltas, so it can't be handed to a
|
||||
// tool handler until it's fully assembled. Same accumulation the
|
||||
// llamacpp client does — without it, tool calls made over a stream
|
||||
// were silently dropped and the agent loop never executed them.
|
||||
type toolCallAccum struct {
|
||||
id string
|
||||
name string
|
||||
args strings.Builder
|
||||
}
|
||||
toolCallFrags := map[int]*toolCallAccum{}
|
||||
var toolCallOrder []int
|
||||
|
||||
flushToolCalls := func() []llm.ToolCall {
|
||||
if len(toolCallOrder) == 0 {
|
||||
return nil
|
||||
}
|
||||
calls := make([]llm.ToolCall, 0, len(toolCallOrder))
|
||||
for _, idx := range toolCallOrder {
|
||||
frag := toolCallFrags[idx]
|
||||
calls = append(calls, llm.ToolCall{
|
||||
ID: frag.id,
|
||||
Name: frag.name,
|
||||
Arguments: json.RawMessage(frag.args.String()),
|
||||
})
|
||||
}
|
||||
toolCallFrags = map[int]*toolCallAccum{}
|
||||
toolCallOrder = nil
|
||||
return calls
|
||||
}
|
||||
|
||||
scanner := bufio.NewScanner(resp.Body)
|
||||
// A single SSE line can exceed bufio.Scanner's 64KB default cap
|
||||
// (e.g. a large tool-call arguments delta), which would kill the
|
||||
// stream with "token too long" — same headroom the anthropic
|
||||
// client already reserves.
|
||||
scanner.Buffer(make([]byte, 0, 64*1024), 4*1024*1024)
|
||||
for scanner.Scan() {
|
||||
line := scanner.Text()
|
||||
if !strings.HasPrefix(line, "data: ") {
|
||||
|
|
@ -128,13 +169,61 @@ func (c *Client) Stream(ctx context.Context, req llm.CompletionRequest) iter.Seq
|
|||
return
|
||||
}
|
||||
|
||||
var usage llm.TokenUsage
|
||||
if event.Usage != nil {
|
||||
usage = llm.TokenUsage{
|
||||
InputTokens: event.Usage.PromptTokens,
|
||||
OutputTokens: event.Usage.CompletionTokens,
|
||||
TotalTokens: event.Usage.TotalTokens,
|
||||
}
|
||||
}
|
||||
|
||||
if len(event.Choices) == 0 {
|
||||
// The usage-only event (per stream_options.include_usage)
|
||||
// carries no choices, so it needs its own chunk.
|
||||
if event.Usage != nil {
|
||||
if !yield(llm.StreamChunk{Usage: usage}, nil) {
|
||||
return
|
||||
}
|
||||
}
|
||||
continue
|
||||
}
|
||||
|
||||
for _, choice := range event.Choices {
|
||||
hasFragment := len(choice.Delta.ToolCalls) > 0
|
||||
for _, tc := range choice.Delta.ToolCalls {
|
||||
frag, ok := toolCallFrags[tc.Index]
|
||||
if !ok {
|
||||
frag = &toolCallAccum{}
|
||||
toolCallFrags[tc.Index] = frag
|
||||
toolCallOrder = append(toolCallOrder, tc.Index)
|
||||
}
|
||||
if tc.ID != "" {
|
||||
frag.id = tc.ID
|
||||
}
|
||||
if tc.Function.Name != "" {
|
||||
frag.name = tc.Function.Name
|
||||
}
|
||||
frag.args.WriteString(tc.Function.Arguments)
|
||||
}
|
||||
|
||||
chunk := llm.StreamChunk{
|
||||
Delta: choice.Delta.Content,
|
||||
Usage: usage,
|
||||
}
|
||||
if choice.FinishReason != "" {
|
||||
chunk.FinishReason = choice.FinishReason
|
||||
chunk.ToolCalls = flushToolCalls()
|
||||
}
|
||||
|
||||
// A fragment-only event (a piece of a tool call's streamed
|
||||
// arguments, with nothing else in this delta) has nothing
|
||||
// yet for the agent loop to act on: it was buffered above,
|
||||
// so skip yielding an empty chunk for it.
|
||||
if hasFragment && chunk.Delta == "" && chunk.FinishReason == "" {
|
||||
continue
|
||||
}
|
||||
|
||||
if !yield(chunk, nil) {
|
||||
return
|
||||
}
|
||||
|
|
@ -160,7 +249,7 @@ func (c *Client) Capabilities() llm.ProviderCapabilities {
|
|||
}
|
||||
|
||||
// buildRequest converts an llm.CompletionRequest to the OpenAI API format.
|
||||
func (c *Client) buildRequest(req llm.CompletionRequest) (io.Reader, error) {
|
||||
func (c *Client) buildRequest(req llm.CompletionRequest, stream bool) (io.Reader, error) {
|
||||
// Convert messages to OpenAI format
|
||||
messages := make([]openaiMessage, len(req.Messages))
|
||||
for i, m := range req.Messages {
|
||||
|
|
@ -195,10 +284,25 @@ func (c *Client) buildRequest(req llm.CompletionRequest) (io.Reader, error) {
|
|||
tools[i] = tool
|
||||
}
|
||||
|
||||
// The per-request model wins when set; otherwise fall back to the
|
||||
// client's configured one (Config.Model used to be discarded entirely,
|
||||
// so every request went out with an empty model — a hard API error on
|
||||
// OpenAI, and the agent loop never sets req.Model).
|
||||
model := req.Model
|
||||
if model == "" {
|
||||
model = c.model
|
||||
}
|
||||
|
||||
openaiReq := openaiChatRequest{
|
||||
Model: req.Model,
|
||||
Model: model,
|
||||
Messages: messages,
|
||||
Stream: false,
|
||||
Stream: stream,
|
||||
}
|
||||
if stream {
|
||||
// Ask for a final SSE event carrying token usage (OpenAI-style
|
||||
// streaming omits it otherwise), so callers can track real token
|
||||
// counts per turn instead of always seeing zero.
|
||||
openaiReq.StreamOptions = &openaiStreamOptions{IncludeUsage: true}
|
||||
}
|
||||
if len(tools) > 0 {
|
||||
openaiReq.Tools = tools
|
||||
|
|
@ -222,11 +326,14 @@ func (c *Client) buildRequest(req llm.CompletionRequest) (io.Reader, error) {
|
|||
if err != nil {
|
||||
return nil, fmt.Errorf("marshaling request: %w", err)
|
||||
}
|
||||
return strings.NewReader(string(data)), nil
|
||||
return bytes.NewReader(data), nil
|
||||
}
|
||||
|
||||
// toResponse converts an OpenAI API response to our CompletionResponse.
|
||||
func (c *Client) toResponse(resp openaiChatResponse) llm.CompletionResponse {
|
||||
func (c *Client) toResponse(resp openaiChatResponse) (llm.CompletionResponse, error) {
|
||||
if len(resp.Choices) == 0 {
|
||||
return llm.CompletionResponse{}, fmt.Errorf("openai: response contained no choices")
|
||||
}
|
||||
choice := resp.Choices[0]
|
||||
result := llm.CompletionResponse{
|
||||
ID: resp.ID,
|
||||
|
|
@ -249,7 +356,7 @@ func (c *Client) toResponse(resp openaiChatResponse) llm.CompletionResponse {
|
|||
TotalTokens: resp.Usage.TotalTokens,
|
||||
}
|
||||
|
||||
return result
|
||||
return result, nil
|
||||
}
|
||||
|
||||
// OpenAI API types
|
||||
|
|
@ -263,6 +370,11 @@ type openaiChatRequest struct {
|
|||
MaxTokens *int `json:"max_tokens,omitempty"`
|
||||
Stop []string `json:"stop,omitempty"`
|
||||
Stream bool `json:"stream"`
|
||||
StreamOptions *openaiStreamOptions `json:"stream_options,omitempty"`
|
||||
}
|
||||
|
||||
type openaiStreamOptions struct {
|
||||
IncludeUsage bool `json:"include_usage"`
|
||||
}
|
||||
|
||||
type openaiMessage struct {
|
||||
|
|
@ -319,6 +431,7 @@ type openaiUsage struct {
|
|||
type openaiStreamEvent struct {
|
||||
ID string `json:"id"`
|
||||
Choices []openaiStreamChoice `json:"choices"`
|
||||
Usage *openaiUsage `json:"usage"`
|
||||
}
|
||||
|
||||
type openaiStreamChoice struct {
|
||||
|
|
|
|||
|
|
@ -139,7 +139,5 @@ func TestStreamChunk_JSON(t *testing.T) {
|
|||
}
|
||||
}
|
||||
|
||||
|
||||
|
||||
func float32Ptr(f float32) *float32 { return &f }
|
||||
func intPtr(i int) *int { return &i }
|
||||
|
|
|
|||
Loading…
Reference in a new issue