refactor(ai): Adopt whilesmartgo/agents for the model engine

go.mod

@@ -29,6 +29,7 @@ require (
 	github.com/spf13/cobra v1.10.2
 	github.com/testcontainers/testcontainers-go v0.41.0
 	github.com/testcontainers/testcontainers-go/modules/compose v0.41.0
+	github.com/whilesmartgo/agents v0.1.0
 	go.opentelemetry.io/otel v1.41.0
 	go.opentelemetry.io/otel/exporters/otlp/otlpmetric/otlpmetricgrpc v1.41.0
 	go.opentelemetry.io/otel/exporters/otlp/otlpmetric/otlpmetrichttp v1.41.0

go.sum

@@ -510,6 +510,8 @@ github.com/ugorji/go/codec v1.3.0 h1:Qd2W2sQawAfG8XSvzwhBeoGq71zXOC/Q1E9y/wUcsUA
 github.com/ugorji/go/codec v1.3.0/go.mod h1:pRBVtBSKl77K30Bv8R2P+cLSGaTtex6fsA2Wjqmfxj4=
 github.com/vbatts/tar-split v0.12.2 h1:w/Y6tjxpeiFMR47yzZPlPj/FcPLpXbTUi/9H7d3CPa4=
 github.com/vbatts/tar-split v0.12.2/go.mod h1:eF6B6i6ftWQcDqEn3/iGFRFRo8cBIMSJVOpnNdfTMFA=
+github.com/whilesmartgo/agents v0.1.0 h1:AX7Q5e6BY2Fb/1J7prtZgXSx9bsD8GRSDsmfTQ8kgxs=
+github.com/whilesmartgo/agents v0.1.0/go.mod h1:MTuXbfen/M5GGphQk3kggnO52i8thPBf3RfywfuxOOM=
 github.com/xhit/go-str2duration/v2 v2.1.0 h1:lxklc02Drh6ynqX+DdPyp5pCKLUQpRT8bp8Ydu2Bstc=
 github.com/xhit/go-str2duration/v2 v2.1.0/go.mod h1:ohY8p+0f07DiV6Em5LKB0s2YpLtXVyJfNt1+BlmyAsU=
 github.com/yuin/goldmark v1.1.27/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=

internal/ai/openai.go

@@ -1,27 +1,21 @@
 package ai
 
 import (
-	"bytes"
 	"context"
-	"encoding/json"
-	"fmt"
-	"io"
 	"net/http"
 	"strings"
 	"time"
 
 	"github.com/flatrun/agent/pkg/config"
+	"github.com/whilesmartgo/agents"
+	"github.com/whilesmartgo/agents/engine/openai"
 )
 
-const maxResponseBytes = 4 << 20
-
-// openAICompatible speaks the OpenAI chat-completions wire format,
-// which OpenAI, Ollama, vLLM, LM Studio and most gateways all accept.
+// openAICompatible adapts the whilesmartgo/agents OpenAI-compatible engine to
+// FlatRun's Provider. The chat-completions wire format lives in the library;
+// this maps FlatRun's request and response types across that boundary.
 type openAICompatible struct {
-	baseURL string
-	apiKey  string
-	model   string
-	client  *http.Client
+	engine *openai.Engine
 }
 
 func newOpenAICompatible(cfg *config.AIConfig) *openAICompatible {
@@ -30,155 +24,77 @@ func newOpenAICompatible(cfg *config.AIConfig) *openAICompatible {
 		timeout = 60 * time.Second
 	}
 	return &openAICompatible{
-		baseURL: strings.TrimRight(cfg.BaseURL, "/"),
-		apiKey:  cfg.APIKey,
-		model:   cfg.Model,
-		client:  &http.Client{Timeout: timeout},
+		engine: openai.New(
+			strings.TrimRight(cfg.BaseURL, "/"),
+			cfg.Model,
+			openai.WithAPIKey(cfg.APIKey),
+			openai.WithHTTPClient(&http.Client{Timeout: timeout}),
+		),
 	}
 }
 
 func (p *openAICompatible) Name() string {
 	return "openai-compatible"
 }
 
-// wireMessages converts internal messages to the OpenAI chat wire
-// format, where assistant tool calls and tool results are nested
-// differently than our flat representation.
-func wireMessages(messages []Message) []map[string]interface{} {
-	out := make([]map[string]interface{}, 0, len(messages))
-	for _, m := range messages {
-		wm := map[string]interface{}{"role": m.Role, "content": m.Content}
-		if len(m.ToolCalls) > 0 {
-			calls := make([]map[string]interface{}, 0, len(m.ToolCalls))
-			for _, tc := range m.ToolCalls {
-				calls = append(calls, map[string]interface{}{
-					"id":   tc.ID,
-					"type": "function",
-					"function": map[string]interface{}{
-						"name":      tc.Name,
-						"arguments": tc.Arguments,
-					},
-				})
-			}
-			wm["tool_calls"] = calls
-		}
-		if m.ToolCallID != "" {
-			wm["tool_call_id"] = m.ToolCallID
-		}
-		if m.Name != "" {
-			wm["name"] = m.Name
-		}
-		out = append(out, wm)
+func (p *openAICompatible) Complete(ctx context.Context, req Request) (*Response, error) {
+	resp, err := p.engine.Complete(ctx, agents.Request{
+		Messages:    toEngineMessages(req.Messages),
+		Tools:       toEngineTools(req.Tools),
+		MaxTokens:   req.MaxTokens,
+		Temperature: req.Temperature,
+	})
+	if err != nil {
+		return nil, err
 	}
-	return out
+	return &Response{
+		Content:   resp.Content,
+		ToolCalls: fromEngineToolCalls(resp.ToolCalls),
+		Model:     resp.Model,
+		Usage:     Usage{PromptTokens: resp.Usage.PromptTokens, CompletionTokens: resp.Usage.CompletionTokens},
+	}, nil
 }
 
-func wireTools(tools []Tool) []map[string]interface{} {
-	out := make([]map[string]interface{}, 0, len(tools))
-	for _, t := range tools {
-		out = append(out, map[string]interface{}{
-			"type": "function",
-			"function": map[string]interface{}{
-				"name":        t.Name,
-				"description": t.Description,
-				"parameters":  t.Parameters,
-			},
+func toEngineMessages(messages []Message) []agents.Message {
+	out := make([]agents.Message, 0, len(messages))
+	for _, m := range messages {
+		out = append(out, agents.Message{
+			Role:       m.Role,
+			Content:    m.Content,
+			ToolCalls:  toEngineToolCalls(m.ToolCalls),
+			ToolCallID: m.ToolCallID,
+			Name:       m.Name,
 		})
 	}
 	return out
 }
 
-func (p *openAICompatible) Complete(ctx context.Context, req Request) (*Response, error) {
-	payload := map[string]interface{}{
-		"model":    p.model,
-		"messages": wireMessages(req.Messages),
-	}
-	if len(req.Tools) > 0 {
-		payload["tools"] = wireTools(req.Tools)
-	}
-	if req.MaxTokens > 0 {
-		payload["max_tokens"] = req.MaxTokens
-	}
-	if req.Temperature > 0 {
-		payload["temperature"] = req.Temperature
-	}
-	body, err := json.Marshal(payload)
-	if err != nil {
-		return nil, err
-	}
-
-	httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, p.baseURL+"/chat/completions", bytes.NewReader(body))
-	if err != nil {
-		return nil, err
-	}
-	httpReq.Header.Set("Content-Type", "application/json")
-	// Local servers (Ollama, LM Studio) typically run keyless.
-	if p.apiKey != "" {
-		httpReq.Header.Set("Authorization", "Bearer "+p.apiKey)
-	}
-
-	resp, err := p.client.Do(httpReq)
-	if err != nil {
-		return nil, fmt.Errorf("ai provider request failed: %w", err)
-	}
-	defer resp.Body.Close()
-
-	raw, err := io.ReadAll(io.LimitReader(resp.Body, maxResponseBytes))
-	if err != nil {
-		return nil, fmt.Errorf("ai provider response read failed: %w", err)
-	}
-
-	if resp.StatusCode != http.StatusOK {
-		var apiErr struct {
-			Error struct {
-				Message string `json:"message"`
-			} `json:"error"`
-		}
-		msg := strings.TrimSpace(string(raw))
-		if json.Unmarshal(raw, &apiErr) == nil && apiErr.Error.Message != "" {
-			msg = apiErr.Error.Message
-		}
-		return nil, fmt.Errorf("ai provider returned %d: %s", resp.StatusCode, msg)
+func toEngineTools(tools []Tool) []agents.ToolSchema {
+	out := make([]agents.ToolSchema, 0, len(tools))
+	for _, t := range tools {
+		out = append(out, agents.ToolSchema{Name: t.Name, Description: t.Description, Parameters: t.Parameters})
 	}
+	return out
+}
 
-	var parsed struct {
-		Model   string `json:"model"`
-		Choices []struct {
-			Message struct {
-				Content   string `json:"content"`
-				ToolCalls []struct {
-					ID       string `json:"id"`
-					Function struct {
-						Name      string `json:"name"`
-						Arguments string `json:"arguments"`
-					} `json:"function"`
-				} `json:"tool_calls"`
-			} `json:"message"`
-		} `json:"choices"`
-		Usage Usage `json:"usage"`
-	}
-	if err := json.Unmarshal(raw, &parsed); err != nil {
-		return nil, fmt.Errorf("ai provider returned invalid JSON: %w", err)
+func toEngineToolCalls(calls []ToolCall) []agents.ToolCall {
+	if len(calls) == 0 {
+		return nil
 	}
-	if len(parsed.Choices) == 0 {
-		return nil, fmt.Errorf("ai provider returned no choices")
+	out := make([]agents.ToolCall, 0, len(calls))
+	for _, c := range calls {
+		out = append(out, agents.ToolCall{ID: c.ID, Name: c.Name, Arguments: c.Arguments})
 	}
+	return out
+}
 
-	model := parsed.Model
-	if model == "" {
-		model = p.model
+func fromEngineToolCalls(calls []agents.ToolCall) []ToolCall {
+	if len(calls) == 0 {
+		return nil
 	}
-
-	choice := parsed.Choices[0].Message
-	var toolCalls []ToolCall
-	for _, tc := range choice.ToolCalls {
-		toolCalls = append(toolCalls, ToolCall{ID: tc.ID, Name: tc.Function.Name, Arguments: tc.Function.Arguments})
+	out := make([]ToolCall, 0, len(calls))
+	for _, c := range calls {
+		out = append(out, ToolCall{ID: c.ID, Name: c.Name, Arguments: c.Arguments})
 	}
-
-	return &Response{
-		Content:   choice.Content,
-		ToolCalls: toolCalls,
-		Model:     model,
-		Usage:     parsed.Usage,
-	}, nil
+	return out
 }