core/server/providers/openai.go

package providers

import (
	"git.gobha.me/xcaliber/chat-switchboard/capabilities"
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"strings"
)

// OpenAIProvider handles any OpenAI-compatible API.
type OpenAIProvider struct{}

func (p *OpenAIProvider) ID() string { return "openai" }

// ── Chat Completion (non-streaming) ─────────

func (p *OpenAIProvider) ChatCompletion(ctx context.Context, cfg ProviderConfig, req CompletionRequest) (*CompletionResponse, error) {
	req.Stream = false

	body, err := p.doRequest(ctx, cfg, req)
	if err != nil {
		return nil, err
	}
	defer body.Close()

	var resp openaiChatResponse
	if err := json.NewDecoder(body).Decode(&resp); err != nil {
		return nil, fmt.Errorf("decode response: %w", err)
	}

	if len(resp.Choices) == 0 {
		return nil, fmt.Errorf("no choices in response")
	}

	result := &CompletionResponse{
		Content:      resp.Choices[0].Message.Content,
		Model:        resp.Model,
		FinishReason: resp.Choices[0].FinishReason,
		InputTokens:  resp.Usage.PromptTokens,
		OutputTokens: resp.Usage.CompletionTokens,
	}

	// Extract tool calls if present
	if len(resp.Choices[0].Message.ToolCalls) > 0 {
		for _, tc := range resp.Choices[0].Message.ToolCalls {
			result.ToolCalls = append(result.ToolCalls, ToolCall{
				ID:   tc.ID,
				Type: tc.Type,
				Function: FunctionCall{
					Name:      tc.Function.Name,
					Arguments: tc.Function.Arguments,
				},
			})
		}
		if result.FinishReason == "" {
			result.FinishReason = "tool_calls"
		}
	}

	return result, nil
}

// ── Stream Completion ───────────────────────

func (p *OpenAIProvider) StreamCompletion(ctx context.Context, cfg ProviderConfig, req CompletionRequest) (<-chan StreamEvent, error) {
	req.Stream = true

	body, err := p.doRequest(ctx, cfg, req)
	if err != nil {
		return nil, err
	}

	ch := make(chan StreamEvent, 64)

	go func() {
		defer close(ch)
		defer body.Close()

		// Accumulate tool calls across stream chunks
		toolCallMap := map[int]*ToolCall{} // index → accumulated call

		scanner := bufio.NewScanner(body)
		for scanner.Scan() {
			line := scanner.Text()

			if !strings.HasPrefix(line, "data: ") {
				continue
			}

			data := strings.TrimPrefix(line, "data: ")
			if data == "[DONE]" {
				ch <- StreamEvent{Done: true}
				return
			}

			var chunk openaiStreamChunk
			if err := json.Unmarshal([]byte(data), &chunk); err != nil {
				continue
			}

			if len(chunk.Choices) == 0 {
				continue
			}

			choice := chunk.Choices[0]

			// Accumulate tool call deltas
			for _, tc := range choice.Delta.ToolCalls {
				idx := 0
				if tc.Index != nil {
					idx = *tc.Index
				}

				if existing, ok := toolCallMap[idx]; ok {
					// Append arguments fragment
					existing.Function.Arguments += tc.Function.Arguments
				} else {
					// First chunk for this tool — has ID and name
					toolCallMap[idx] = &ToolCall{
						ID:   tc.ID,
						Type: "function",
						Function: FunctionCall{
							Name:      tc.Function.Name,
							Arguments: tc.Function.Arguments,
						},
					}
				}
			}

			ev := StreamEvent{
				Delta:        choice.Delta.Content,
				Model:        chunk.Model,
				FinishReason: choice.FinishReason,
			}

			if ev.FinishReason != "" {
				ev.Done = true
				// Attach accumulated tool calls on final event
				if ev.FinishReason == "tool_calls" && len(toolCallMap) > 0 {
					for _, tc := range toolCallMap {
						ev.ToolCalls = append(ev.ToolCalls, *tc)
					}
				}
			}

			ch <- ev
		}

		if err := scanner.Err(); err != nil {
			ch <- StreamEvent{Error: err}
		}
	}()

	return ch, nil
}

// ── List Models ─────────────────────────────

func (p *OpenAIProvider) ListModels(ctx context.Context, cfg ProviderConfig) ([]Model, error) {
	endpoint := strings.TrimSuffix(cfg.Endpoint, "/") + "/models"

	httpReq, err := http.NewRequestWithContext(ctx, "GET", endpoint, nil)
	if err != nil {
		return nil, err
	}
	if cfg.APIKey != "" {
		httpReq.Header.Set("Authorization", "Bearer "+cfg.APIKey)
	}
	for k, v := range cfg.CustomHeaders {
		httpReq.Header.Set(k, v)
	}

	resp, err := http.DefaultClient.Do(httpReq)
	if err != nil {
		return nil, fmt.Errorf("list models: %w", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		b, _ := io.ReadAll(resp.Body)
		return nil, fmt.Errorf("list models: HTTP %d: %s", resp.StatusCode, string(b))
	}

	var result openaiModelsResponse
	if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
		return nil, fmt.Errorf("decode models: %w", err)
	}

	out := make([]Model, 0, len(result.Data))
	for _, m := range result.Data {
		// Try known table first, then heuristic
		caps, found := capabilities.LookupKnownModel(m.ID)
		if !found {
			caps = capabilities.InferCapabilities(m.ID)
		}
		// Use context_length from API if available and we don't have it
		if m.ContextLength > 0 && caps.MaxContext == 0 {
			caps.MaxContext = m.ContextLength
		}

		out = append(out, Model{
			ID:           m.ID,
			OwnedBy:      m.OwnedBy,
			Capabilities: caps,
		})
	}
	return out, nil
}

// ── HTTP Layer ──────────────────────────────

func (p *OpenAIProvider) doRequest(ctx context.Context, cfg ProviderConfig, req CompletionRequest) (io.ReadCloser, error) {
	endpoint := strings.TrimSuffix(cfg.Endpoint, "/") + "/chat/completions"

	// Build OpenAI-format request
	oaiReq := openaiChatRequest{
		Model:    req.Model,
		Stream:   req.Stream,
		Messages: make([]openaiMessage, 0, len(req.Messages)),
	}
	if req.MaxTokens > 0 {
		oaiReq.MaxTokens = req.MaxTokens
	}
	if req.Temperature != nil {
		oaiReq.Temperature = req.Temperature
	}
	if req.TopP != nil {
		oaiReq.TopP = req.TopP
	}

	// Convert tools
	for _, t := range req.Tools {
		oaiReq.Tools = append(oaiReq.Tools, openaiToolDef{
			Type: t.Type,
			Function: openaiToolDefFunction{
				Name:        t.Function.Name,
				Description: t.Function.Description,
				Parameters:  t.Function.Parameters,
			},
		})
	}

	// Convert messages — handle all roles including tool
	for _, m := range req.Messages {
		oaiMsg := openaiMessage{
			Role:    m.Role,
			Content: m.Content,
		}

		// Assistant messages with tool calls
		if len(m.ToolCalls) > 0 {
			for _, tc := range m.ToolCalls {
				oaiMsg.ToolCalls = append(oaiMsg.ToolCalls, openaiToolCall{
					ID:   tc.ID,
					Type: tc.Type,
					Function: openaiToolCallFunction{
						Name:      tc.Function.Name,
						Arguments: tc.Function.Arguments,
					},
				})
			}
		}

		// Tool result messages
		if m.Role == "tool" {
			oaiMsg.ToolCallID = m.ToolCallID
			oaiMsg.Name = m.Name
		}

		oaiReq.Messages = append(oaiReq.Messages, oaiMsg)
	}

	body, err := json.Marshal(oaiReq)
	if err != nil {
		return nil, fmt.Errorf("marshal request: %w", err)
	}

	httpReq, err := http.NewRequestWithContext(ctx, "POST", endpoint, bytes.NewReader(body))
	if err != nil {
		return nil, err
	}
	httpReq.Header.Set("Content-Type", "application/json")
	if cfg.APIKey != "" {
		httpReq.Header.Set("Authorization", "Bearer "+cfg.APIKey)
	}
	// Inject provider-level custom headers (OpenRouter, Venice, etc.)
	for k, v := range cfg.CustomHeaders {
		httpReq.Header.Set(k, v)
	}

	resp, err := http.DefaultClient.Do(httpReq)
	if err != nil {
		return nil, fmt.Errorf("provider request: %w", err)
	}

	if resp.StatusCode != http.StatusOK {
		b, _ := io.ReadAll(resp.Body)
		resp.Body.Close()
		return nil, fmt.Errorf("provider error: HTTP %d: %s", resp.StatusCode, string(b))
	}

	return resp.Body, nil
}

// ── OpenAI Wire Types ───────────────────────

type openaiMessage struct {
	Role       string              `json:"role"`
	Content    string              `json:"content,omitempty"`
	ToolCalls  []openaiToolCall    `json:"tool_calls,omitempty"`   // assistant requesting tools
	ToolCallID string              `json:"tool_call_id,omitempty"` // role=tool result
	Name       string              `json:"name,omitempty"`         // tool name for role=tool
}

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

type openaiToolCall struct {
	Index    *int                   `json:"index,omitempty"` // present in streaming deltas
	ID       string                 `json:"id,omitempty"`
	Type     string                 `json:"type,omitempty"` // "function"
	Function openaiToolCallFunction `json:"function"`
}

type openaiToolDefFunction struct {
	Name        string          `json:"name"`
	Description string          `json:"description"`
	Parameters  json.RawMessage `json:"parameters"`
}

type openaiToolDef struct {
	Type     string               `json:"type"` // "function"
	Function openaiToolDefFunction `json:"function"`
}

type openaiChatRequest struct {
	Model       string          `json:"model"`
	Messages    []openaiMessage `json:"messages"`
	MaxTokens   int             `json:"max_tokens,omitempty"`
	Temperature *float64        `json:"temperature,omitempty"`
	TopP        *float64        `json:"top_p,omitempty"`
	Stream      bool            `json:"stream"`
	Tools       []openaiToolDef `json:"tools,omitempty"`
}

type openaiChatResponse struct {
	Model   string `json:"model"`
	Choices []struct {
		Message      openaiMessage `json:"message"`
		FinishReason string        `json:"finish_reason"`
	} `json:"choices"`
	Usage struct {
		PromptTokens     int `json:"prompt_tokens"`
		CompletionTokens int `json:"completion_tokens"`
	} `json:"usage"`
}

type openaiStreamChunk struct {
	Model   string `json:"model"`
	Choices []struct {
		Delta struct {
			Content   string           `json:"content"`
			ToolCalls []openaiToolCall `json:"tool_calls,omitempty"`
		} `json:"delta"`
		FinishReason string `json:"finish_reason"`
	} `json:"choices"`
}

type openaiModelsResponse struct {
	Data []struct {
		ID            string `json:"id"`
		OwnedBy       string `json:"owned_by"`
		ContextLength int    `json:"context_length,omitempty"` // Ollama, OpenRouter
	} `json:"data"`
}