byteclaw/internal/llm/openai.go

package llm

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

// OpenAI is a provider for any OpenAI-compatible /chat/completions endpoint
// (OpenAI, OpenRouter, Ollama, LM Studio, vLLM, ...). It always streams.
type OpenAI struct {
	BaseURL     string
	APIKey      string
	Model       string
	Temperature float64
	HTTP        *http.Client
}

func NewOpenAI(baseURL, apiKey, model string, temperature float64, timeout time.Duration) *OpenAI {
	// Default transport: TLS certificate verification stays ON.
	return &OpenAI{
		BaseURL: strings.TrimRight(baseURL, "/"), APIKey: apiKey, Model: model,
		Temperature: temperature, HTTP: &http.Client{Timeout: timeout},
	}
}

type oaTool struct {
	Type     string   `json:"type"`
	Function ToolSpec `json:"function"`
}

type oaFunc struct {
	Name      string `json:"name,omitempty"`
	Arguments string `json:"arguments,omitempty"`
}

type oaToolCall struct {
	Index    *int   `json:"index,omitempty"`
	ID       string `json:"id,omitempty"`
	Type     string `json:"type,omitempty"`
	Function oaFunc `json:"function"`
}

type oaMsg struct {
	Role       string       `json:"role"`
	Content    *string      `json:"content"`
	ToolCalls  []oaToolCall `json:"tool_calls,omitempty"`
	ToolCallID string       `json:"tool_call_id,omitempty"`
	Name       string       `json:"name,omitempty"`
}

type oaReq struct {
	Model       string   `json:"model"`
	Messages    []oaMsg  `json:"messages"`
	Tools       []oaTool `json:"tools,omitempty"`
	Stream      bool     `json:"stream"`
	Temperature float64  `json:"temperature"`
}

type oaChunk struct {
	Choices []struct {
		Delta struct {
			Content   string       `json:"content"`
			ToolCalls []oaToolCall `json:"tool_calls"`
		} `json:"delta"`
		FinishReason *string `json:"finish_reason"`
	} `json:"choices"`
	Error *struct {
		Message string `json:"message"`
	} `json:"error"`
}

func toOA(msgs []Message) []oaMsg {
	out := make([]oaMsg, 0, len(msgs))
	for _, m := range msgs {
		c := m.Content
		om := oaMsg{Role: string(m.Role), Content: &c, ToolCallID: m.ToolCallID, Name: m.Name}
		if m.Role == RoleAssistant && len(m.ToolCalls) > 0 && c == "" {
			om.Content = nil
		}
		for _, tc := range m.ToolCalls {
			om.ToolCalls = append(om.ToolCalls, oaToolCall{ID: tc.ID, Type: "function", Function: oaFunc{Name: tc.Name, Arguments: tc.Arguments}})
		}
		out = append(out, om)
	}
	return out
}

func (p *OpenAI) Chat(ctx context.Context, req Request, onDelta DeltaFunc) (Response, error) {
	body := oaReq{Model: p.Model, Messages: toOA(req.Messages), Stream: true, Temperature: p.Temperature}
	for _, t := range req.Tools {
		body.Tools = append(body.Tools, oaTool{Type: "function", Function: t})
	}
	b, err := json.Marshal(body)
	if err != nil {
		return Response{}, err
	}
	hr, err := http.NewRequestWithContext(ctx, http.MethodPost, p.BaseURL+"/chat/completions", bytes.NewReader(b))
	if err != nil {
		return Response{}, err
	}
	hr.Header.Set("Content-Type", "application/json")
	hr.Header.Set("Accept", "text/event-stream")
	if p.APIKey != "" {
		hr.Header.Set("Authorization", "Bearer "+p.APIKey)
	}
	resp, err := p.HTTP.Do(hr)
	if err != nil {
		return Response{}, fmt.Errorf("llm request: %w", err)
	}
	defer resp.Body.Close()
	if resp.StatusCode/100 != 2 {
		msg, _ := io.ReadAll(io.LimitReader(resp.Body, 4096))
		return Response{}, fmt.Errorf("llm http %d: %s", resp.StatusCode, strings.TrimSpace(string(msg)))
	}
	return parseStream(resp.Body, onDelta)
}

type partialCall struct {
	id, name string
	args     strings.Builder
}

func parseStream(r io.Reader, onDelta DeltaFunc) (Response, error) {
	var out Response
	var text strings.Builder
	calls := map[int]*partialCall{}
	sc := bufio.NewScanner(r)
	sc.Buffer(make([]byte, 64*1024), 4<<20)
	for sc.Scan() {
		line := strings.TrimSpace(sc.Text())
		if !strings.HasPrefix(line, "data:") {
			continue
		}
		data := strings.TrimSpace(strings.TrimPrefix(line, "data:"))
		if data == "[DONE]" {
			break
		}
		var ch oaChunk
		if err := json.Unmarshal([]byte(data), &ch); err != nil {
			return out, fmt.Errorf("decode stream chunk: %w", err)
		}
		if ch.Error != nil {
			return out, fmt.Errorf("llm error: %s", ch.Error.Message)
		}
		for _, c := range ch.Choices {
			if c.Delta.Content != "" {
				text.WriteString(c.Delta.Content)
				if onDelta != nil {
					onDelta(c.Delta.Content)
				}
			}
			for i, tc := range c.Delta.ToolCalls {
				idx := i
				if tc.Index != nil {
					idx = *tc.Index
				}
				pc := calls[idx]
				if pc == nil {
					pc = &partialCall{}
					calls[idx] = pc
				}
				if tc.ID != "" {
					pc.id = tc.ID
				}
				if tc.Function.Name != "" {
					pc.name += tc.Function.Name
				}
				pc.args.WriteString(tc.Function.Arguments)
			}
			if c.FinishReason != nil {
				out.FinishReason = *c.FinishReason
			}
		}
	}
	if err := sc.Err(); err != nil {
		return out, fmt.Errorf("read stream: %w", err)
	}
	out.Content = text.String()
	idxs := make([]int, 0, len(calls))
	for i := range calls {
		idxs = append(idxs, i)
	}
	sort.Ints(idxs)
	for _, i := range idxs {
		pc := calls[i]
		id := pc.id
		if id == "" {
			id = fmt.Sprintf("call_%d", i)
		}
		args := pc.args.String()
		if args == "" {
			args = "{}"
		}
		out.ToolCalls = append(out.ToolCalls, ToolCall{ID: id, Name: pc.name, Arguments: args})
	}
	return out, nil
}