Agent-System/internal/agentx/llm/client.go

package llm

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"strings"
	"time"
)

// Client calls an OpenAI-compatible chat completions endpoint with streaming.
type Client struct {
	BaseURL    string
	APIKey     string
	HTTPClient *http.Client
}

// Stream implements StreamFunc. Transport and provider failures are encoded as
// KindError events; the returned channel is always non-nil and is closed.
func (c *Client) Stream(ctx context.Context, model string, messages []Message, tools []ToolDecl, opts StreamOptions) <-chan Event {
	out := make(chan Event)
	go func() {
		defer close(out)
		c.run(ctx, model, messages, tools, opts, out)
	}()
	return out
}

func (c *Client) run(ctx context.Context, model string, messages []Message, tools []ToolDecl, opts StreamOptions, out chan<- Event) {
	base := strings.TrimRight(c.BaseURL, "/")
	if opts.APIBase != "" {
		base = strings.TrimRight(opts.APIBase, "/")
	}
	key := c.APIKey
	if opts.APIKey != "" {
		key = opts.APIKey
	}
	if base == "" || key == "" || model == "" {
		emitError(out, "", false, "agentx llm client is missing base url, api key, or model")
		return
	}
	body := map[string]any{
		"model":    model,
		"messages": toProviderMessages(messages),
		"stream":   true,
		"stream_options": map[string]any{
			"include_usage": true,
		},
	}
	if opts.MaxTokens > 0 {
		body["max_tokens"] = opts.MaxTokens
	}
	if len(tools) > 0 {
		body["tools"] = toProviderTools(tools)
		body["tool_choice"] = "auto"
	}
	raw, err := json.Marshal(body)
	if err != nil {
		emitError(out, "", false, err.Error())
		return
	}
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, base+"/chat/completions", bytes.NewReader(raw))
	if err != nil {
		emitError(out, "", false, err.Error())
		return
	}
	req.Header.Set("Authorization", "Bearer "+key)
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Accept", "text/event-stream")
	hc := c.HTTPClient
	if hc == nil {
		hc = &http.Client{Timeout: 30 * time.Minute}
	}
	resp, err := hc.Do(req)
	if err != nil {
		if ctx.Err() != nil {
			emitAbort(out, "", ctx.Err())
			return
		}
		emitError(out, "", true, err.Error())
		return
	}
	defer resp.Body.Close()
	if resp.StatusCode >= 400 {
		slurp, _ := io.ReadAll(io.LimitReader(resp.Body, 4096))
		retryable := resp.StatusCode == http.StatusTooManyRequests || resp.StatusCode >= 500
		emitError(out, "", retryable, fmt.Sprintf("provider status %d: %s", resp.StatusCode, strings.TrimSpace(string(slurp))))
		return
	}
	for ev := range ParseSSE(ctx, resp.Body) {
		out <- ev
	}
}

func toProviderMessages(messages []Message) []map[string]any {
	out := make([]map[string]any, 0, len(messages))
	for _, m := range messages {
		switch m.Role {
		case RoleTool:
			out = append(out, map[string]any{
				"role":         "tool",
				"tool_call_id": m.ToolCallID,
				"content":      m.Content,
			})
		case RoleAssistant:
			item := map[string]any{"role": "assistant", "content": m.Content}
			if len(m.ToolCalls) > 0 {
				calls := make([]map[string]any, 0, len(m.ToolCalls))
				for _, tc := range m.ToolCalls {
					calls = append(calls, map[string]any{
						"id":   tc.ID,
						"type": "function",
						"function": map[string]any{
							"name":      tc.Name,
							"arguments": tc.Arguments,
						},
					})
				}
				item["tool_calls"] = calls
			}
			out = append(out, item)
		default:
			out = append(out, map[string]any{
				"role":    string(m.Role),
				"content": m.Content,
			})
		}
	}
	return out
}

func toProviderTools(tools []ToolDecl) []map[string]any {
	out := make([]map[string]any, 0, len(tools))
	for _, t := range tools {
		params := t.Parameters
		if len(params) == 0 {
			params = json.RawMessage(`{"type":"object","properties":{}}`)
		}
		out = append(out, map[string]any{
			"type": "function",
			"function": map[string]any{
				"name":        t.Name,
				"description": t.Description,
				"parameters":  json.RawMessage(params),
			},
		})
	}
	return out
}