Cyrene_v2/internal/runtime/model/openai.go

package model

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

	"git.yeij.top/AskaEth/Cyrene_v2/internal/config"
)

// Wire constraints for the OpenAI-compatible /chat/completions protocol.
const (
	// sseDataPrefix marks an SSE event line carrying a JSON payload.
	sseDataPrefix = "data: "
	// sseDone is the terminal SSE sentinel.
	sseDone = "[DONE]"
	// sseBufInit is the initial bufio.Scanner read buffer size.
	sseBufInit = 64 * 1024
	// maxSSELineBytes caps a single SSE event line.
	maxSSELineBytes = 1024 * 1024
	// errBodyCap bounds how much of an error body is read for the message.
	errBodyCap = 512
	// defaultHTTPTimeout bounds a single provider call when Options.Timeout is 0.
	defaultHTTPTimeout = 180 * time.Second
)

// openAIChatRequest is the JSON body sent to /chat/completions. Field names
// match the provider protocol so the struct marshals directly.
type openAIChatRequest struct {
	Model         string            `json:"model"`
	Messages      []openAIMessage   `json:"messages"`
	Temperature   *float64          `json:"temperature,omitempty"`
	MaxTokens     int               `json:"max_tokens,omitempty"`
	Stream        bool              `json:"stream,omitempty"`
	StreamOptions *openAIStreamOpts `json:"stream_options,omitempty"`
	Tools         []openAITool      `json:"tools,omitempty"`
}

// openAIStreamOpts asks the provider to include usage in the final stream chunk.
type openAIStreamOpts struct {
	IncludeUsage bool `json:"include_usage"`
}

// openAIMessage is one request message in the wire body. Content is either a
// plain string (text-only) or an array of {type,text|image_url} parts.
type openAIMessage struct {
	Role    string `json:"role"`
	Content any    `json:"content"`
}

// openAITool is the wire form of a function tool schema (M2 placeholder).
type openAITool struct {
	Type     string       `json:"type"`
	Function openAIToolFn `json:"function"`
}

// openAIToolFn carries the function definition.
type openAIToolFn struct {
	Name        string `json:"name"`
	Description string `json:"description,omitempty"`
	Parameters  any    `json:"parameters,omitempty"`
}

// openAIChatResponse is the JSON body returned by /chat/completions, for both
// the non-streaming payload and each SSE stream chunk.
type openAIChatResponse struct {
	ID      string         `json:"id"`
	Model   string         `json:"model"`
	Choices []openAIChoice `json:"choices"`
	Usage   *openAIUsage   `json:"usage,omitempty"`
}

// openAIChoice is one candidate completion.
type openAIChoice struct {
	Index        int                   `json:"index"`
	Message      openAIResponseMessage `json:"message"`
	Delta        *openAIDelta          `json:"delta,omitempty"`
	FinishReason string                `json:"finish_reason,omitempty"`
}

// openAIResponseMessage is the assistant message inside a non-streaming choice.
type openAIResponseMessage struct {
	Role    string `json:"role"`
	Content string `json:"content"`
}

// openAIDelta is the incremental payload inside a stream chunk.
type openAIDelta struct {
	Role    string `json:"role"`
	Content string `json:"content"`
}

// openAIUsage reports token consumption.
type openAIUsage struct {
	PromptTokens     int `json:"prompt_tokens"`
	CompletionTokens int `json:"completion_tokens"`
	TotalTokens      int `json:"total_tokens"`
}

// OpenAIChatAdapter calls an OpenAI-compatible /chat/completions endpoint. It is
// constructed from a config.LLMConfig and implements Adapter for both
// non-streaming and token-level SSE streaming.
type OpenAIChatAdapter struct {
	name     string
	baseURL  string
	apiKey   string
	model    string
	client   *http.Client
	recorder Recorder
}

// OpenAIOption configures an OpenAIChatAdapter.
type OpenAIOption func(*OpenAIChatAdapter)

// WithHTTPClient overrides the default HTTP client (used to inject mocks or a
// tuned transport in tests and production).
func WithHTTPClient(c *http.Client) OpenAIOption {
	return func(a *OpenAIChatAdapter) { a.client = c }
}

// WithRecorder wires a per-call Recorder for CallRecord accounting.
func WithRecorder(r Recorder) OpenAIOption {
	return func(a *OpenAIChatAdapter) { a.recorder = r }
}

// WithModel overrides the adapter's default model name.
func WithModel(m string) OpenAIOption {
	return func(a *OpenAIChatAdapter) { a.model = m }
}

// NewOpenAIChat builds an adapter from a supplier config. baseURL is trimmed of
// any trailing slash before the /chat/completions path is appended.
func NewOpenAIChat(cfg config.LLMConfig, opts ...OpenAIOption) *OpenAIChatAdapter {
	a := &OpenAIChatAdapter{
		name:    "openai",
		baseURL: strings.TrimRight(cfg.BaseURL, "/"),
		apiKey:  cfg.APIKey,
		model:   cfg.Model,
		client: &http.Client{
			Transport: http.DefaultTransport,
			Timeout:   defaultHTTPTimeout,
		},
	}
	for _, o := range opts {
		o(a)
	}
	return a
}

// Name returns the adapter identifier.
func (a *OpenAIChatAdapter) Name() string { return a.name }

// Chat performs a non-streaming completion.
func (a *OpenAIChatAdapter) Chat(ctx context.Context, req Request, opts Options) (Response, error) {
	start := time.Now()
	if a.baseURL == "" {
		return Response{}, ErrUpstream
	}
	ctx, cancel := withCallTimeout(ctx, opts)
	defer cancel()
	model, err := a.resolveModel(req)
	if err != nil {
		return Response{}, err
	}
	req.Model = model
	body, err := a.buildRequest(req, false)
	if err != nil {
		return Response{}, err
	}
	httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, a.endpoint(), bytes.NewReader(body))
	if err != nil {
		return Response{}, fmt.Errorf("model: %w: build request: %v", ErrUpstream, err)
	}
	resp, err := a.doRequest(httpReq, false)
	if err != nil {
		return Response{}, err
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		return Response{}, a.upstreamError(resp)
	}
	var wire openAIChatResponse
	if err := json.NewDecoder(resp.Body).Decode(&wire); err != nil {
		return Response{}, fmt.Errorf("model: %w: decode response: %v", ErrUpstream, err)
	}
	out := responseFromWire(&wire)
	a.record(req, start, out.TotalTokens)
	return out, nil
}

// ChatStream streams token-level deltas via onChunk and returns the aggregated
// response with usage when the provider reports it in the final chunk.
func (a *OpenAIChatAdapter) ChatStream(ctx context.Context, req Request, onChunk func(delta string)) (Response, error) {
	start := time.Now()
	if a.baseURL == "" {
		return Response{}, ErrUpstream
	}
	model, err := a.resolveModel(req)
	if err != nil {
		return Response{}, err
	}
	req.Model = model
	body, err := a.buildRequest(req, true)
	if err != nil {
		return Response{}, err
	}
	httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, a.endpoint(), bytes.NewReader(body))
	if err != nil {
		return Response{}, fmt.Errorf("model: %w: build request: %v", ErrUpstream, err)
	}
	resp, err := a.doRequest(httpReq, true)
	if err != nil {
		return Response{}, err
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		return Response{}, a.upstreamError(resp)
	}
	out, err := scanSSE(resp.Body, onChunk)
	if err != nil {
		return Response{}, err
	}
	out.Model = model
	a.record(req, start, out.TotalTokens)
	return out, nil
}

// withCallTimeout wraps ctx with opts.Timeout when it is positive.
func withCallTimeout(parent context.Context, opts Options) (context.Context, context.CancelFunc) {
	if opts.Timeout > 0 {
		return context.WithTimeout(parent, opts.Timeout)
	}
	return parent, func() {}
}

// resolveModel returns the effective model name.
func (a *OpenAIChatAdapter) resolveModel(req Request) (string, error) {
	if req.Model != "" {
		return req.Model, nil
	}
	if a.model != "" {
		return a.model, nil
	}
	return "", ErrNoModel
}

// endpoint returns the full /chat/completions URL.
func (a *OpenAIChatAdapter) endpoint() string {
	return a.baseURL + "/chat/completions"
}

func (a *OpenAIChatAdapter) buildRequest(req Request, stream bool) ([]byte, error) {
	msgs, err := buildMessages(req.Messages)
	if err != nil {
		return nil, err
	}
	wire := openAIChatRequest{
		Model:       req.Model,
		Messages:    msgs,
		Temperature: req.Temperature,
		MaxTokens:   req.MaxTokens,
		Stream:      stream,
		Tools:       buildTools(req.Tools),
	}
	if stream {
		wire.StreamOptions = &openAIStreamOpts{IncludeUsage: true}
	}
	buf, err := json.Marshal(wire)
	if err != nil {
		return nil, fmt.Errorf("model: %w: marshal request: %v", ErrBadRequest, err)
	}
	return buf, nil
}

func buildMessages(msgs []Message) ([]openAIMessage, error) {
	out := make([]openAIMessage, 0, len(msgs))
	for _, m := range msgs {
		content, err := encodeContent(m.Content)
		if err != nil {
			return nil, err
		}
		out = append(out, openAIMessage{Role: string(m.Role), Content: content})
	}
	return out, nil
}

func encodeContent(parts []Content) (any, error) {
	if len(parts) == 0 {
		return "", nil
	}
	if len(parts) == 1 && parts[0].Type == ContentTypeText {
		return parts[0].Text, nil
	}
	arr := make([]any, 0, len(parts))
	for _, p := range parts {
		switch p.Type {
		case ContentTypeText:
			arr = append(arr, map[string]any{"type": "text", "text": p.Text})
		case ContentTypeImage:
			arr = append(arr, map[string]any{
				"type":      "image_url",
				"image_url": map[string]any{"url": p.ImageURL},
			})
		default:
			return nil, fmt.Errorf("model: %w: unsupported content type %q", ErrBadRequest, p.Type)
		}
	}
	return arr, nil
}

func buildTools(tools []ToolSpec) []openAITool {
	if len(tools) == 0 {
		return nil
	}
	out := make([]openAITool, 0, len(tools))
	for _, t := range tools {
		out = append(out, openAITool{Type: "function", Function: openAIToolFn{
			Name:        t.Name,
			Description: t.Description,
			Parameters:  t.Parameters,
		}})
	}
	return out
}

func (a *OpenAIChatAdapter) doRequest(httpReq *http.Request, stream bool) (*http.Response, error) {
	if a.apiKey != "" {
		httpReq.Header.Set("Authorization", "Bearer "+a.apiKey)
	}
	httpReq.Header.Set("Content-Type", "application/json")
	if stream {
		httpReq.Header.Set("Accept", "text/event-stream")
	} else {
		httpReq.Header.Set("Accept", "application/json")
	}
	resp, err := a.client.Do(httpReq)
	if err != nil {
		return nil, fmt.Errorf("model: %w: %v", ErrUpstream, err)
	}
	return resp, nil
}

func (a *OpenAIChatAdapter) upstreamError(resp *http.Response) error {
	snippet, _ := io.ReadAll(io.LimitReader(resp.Body, errBodyCap))
	return fmt.Errorf("model: %w: status %d: %s", ErrUpstream, resp.StatusCode, strings.TrimSpace(string(snippet)))
}

func (a *OpenAIChatAdapter) record(req Request, start time.Time, tokens int) {
	if a.recorder == nil {
		return
	}
	a.recorder(CallRecord{
		Model:   req.Model,
		Purpose: req.Purpose,
		Tokens:  tokens,
		TookMs:  time.Since(start).Milliseconds(),
	})
}

func responseFromWire(wire *openAIChatResponse) Response {
	out := Response{Model: wire.Model}
	if len(wire.Choices) > 0 {
		text := wire.Choices[0].Message.Content
		out.Content = text
		out.Message = Message{Role: RoleAssistant, Content: []Content{{Type: ContentTypeText, Text: text}}}
	}
	if u := wire.Usage; u != nil {
		out.PromptTokens = u.PromptTokens
		out.CompletionTokens = u.CompletionTokens
		out.TotalTokens = u.TotalTokens
	}
	return out
}

// scanSSE incrementally decodes stream chunks, accumulates the full content
// with a strings.Builder (never reading the whole body at once) and captures
// usage from the final chunk.
func scanSSE(r io.Reader, onChunk func(string)) (Response, error) {
	scanner := bufio.NewScanner(r)
	scanner.Buffer(make([]byte, 0, sseBufInit), maxSSELineBytes)
	agg := Response{}
	var b strings.Builder
	for scanner.Scan() {
		line := strings.TrimSpace(scanner.Text())
		if line == "" || !strings.HasPrefix(line, sseDataPrefix) {
			continue
		}
		data := strings.TrimSpace(strings.TrimPrefix(line, sseDataPrefix))
		if data == sseDone {
			break
		}
		var chunk openAIChatResponse
		if err := json.Unmarshal([]byte(data), &chunk); err != nil {
			return Response{}, fmt.Errorf("model: %w: decode stream chunk: %v", ErrUpstream, err)
		}
		if u := chunk.Usage; u != nil {
			agg.PromptTokens = u.PromptTokens
			agg.CompletionTokens = u.CompletionTokens
			agg.TotalTokens = u.TotalTokens
		}
		if len(chunk.Choices) == 0 || chunk.Choices[0].Delta == nil {
			continue
		}
		delta := chunk.Choices[0].Delta.Content
		if delta == "" {
			continue
		}
		b.WriteString(delta)
		if onChunk != nil {
			onChunk(delta)
		}
	}
	if err := scanner.Err(); err != nil {
		return Response{}, fmt.Errorf("model: %w: read stream: %v", ErrUpstream, err)
	}
	agg.Content = b.String()
	agg.Message = Message{Role: RoleAssistant, Content: []Content{{Type: ContentTypeText, Text: agg.Content}}}
	return agg, nil
}