From f4254ae1448f8a2c888aeb516e045790e15ca51b Mon Sep 17 00:00:00 2001 From: nfebe Date: Sun, 19 Jul 2026 20:02:46 +0100 Subject: [PATCH] refactor(ai): Adopt whilesmartgo/agents for the model engine The assistant's OpenAI-compatible model client was a bespoke copy of the chat-completions wire format. It now comes from the shared whilesmartgo/agents library, so the same engine is reused rather than maintained in two places. Behaviour is unchanged: the existing provider tests pass as-is. This is the engine-layer step of adopting the extracted agent core; expressing the tools and the tool loop through the library's Registry and Runner follows. --- go.mod | 1 + go.sum | 2 + internal/ai/openai.go | 198 ++++++++++++------------------------------ 3 files changed, 60 insertions(+), 141 deletions(-) diff --git a/go.mod b/go.mod index f35ff3a..92bff46 100644 --- a/go.mod +++ b/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 diff --git a/go.sum b/go.sum index 7a3baaf..da6510b 100644 --- a/go.sum +++ b/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= diff --git a/internal/ai/openai.go b/internal/ai/openai.go index 98558ff..2240046 100644 --- a/internal/ai/openai.go +++ b/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,10 +24,12 @@ 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}), + ), } } @@ -41,144 +37,64 @@ 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 }