iop/apps/edge/internal/openai/chat_handler.go

package openai

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"strings"
	"time"

	"go.uber.org/zap"

	edgeservice "iop/apps/edge/internal/service"
)

func (s *Server) handleChatCompletions(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		writeError(w, http.StatusMethodNotAllowed, "method_not_allowed", "method not allowed")
		return
	}
	defer r.Body.Close()

	var req chatCompletionRequest
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		writeError(w, http.StatusBadRequest, "invalid_request_error", "invalid JSON request")
		return
	}
	target := s.resolveTarget(req.Model)
	if target == "" {
		writeError(w, http.StatusBadRequest, "invalid_request_error", "model is required")
		return
	}
	basePrompt := promptFromMessages(req.Messages)
	if strings.TrimSpace(basePrompt) == "" {
		writeError(w, http.StatusBadRequest, "invalid_request_error", "messages are required")
		return
	}
	outputPolicy := s.resolveOutputPolicy(basePrompt)
	messages := req.Messages
	if instruction := strictOutputContractInstruction(outputPolicy); instruction != "" {
		messages = prependSystemMessage(messages, instruction)
	}
	prompt := promptFromMessages(messages)
	input := req.runInput(prompt, messages, outputPolicy.Strict)
	s.logger.Info("openai chat completion input",
		zap.String("model", req.Model),
		zap.String("target", target),
		zap.String("adapter", s.resolveAdapter()),
		zap.Bool("strict_output", outputPolicy.Strict),
		zap.Bool("strict_stream_buffer", outputPolicy.StreamBuffer),
		zap.String("xml_completion_tool", outputPolicy.XMLCompletionTool),
		zap.Bool("contract_instruction", outputPolicy.ContractInstruction),
		zap.Bool("stream", req.Stream),
		zap.Int("message_count", len(req.Messages)),
		zap.Int("prompt_len", len(prompt)),
		zap.String("prompt_preview", previewString(prompt, 1000)),
		zap.Any("input_keys", mapKeys(input)),
	)

	handle, err := s.service.SubmitRun(r.Context(), edgeservice.SubmitRunRequest{
		NodeRef:    s.cfg.NodeRef,
		Adapter:    s.resolveAdapter(),
		Target:     target,
		SessionID:  s.resolveSessionID(),
		Prompt:     prompt,
		Input:      input,
		TimeoutSec: s.resolveTimeoutSec(),
		Metadata: map[string]string{
			"source":        "openai",
			"openai_model":  req.Model,
			"openai_stream": fmt.Sprintf("%t", req.Stream),
			"strict_output": fmt.Sprintf("%t", outputPolicy.Strict),
		},
	})
	if err != nil {
		writeError(w, http.StatusBadGateway, "node_dispatch_error", err.Error())
		return
	}
	defer handle.Close()

	if req.Stream {
		s.streamChatCompletion(w, r, req, handle, outputPolicy)
		return
	}
	s.completeChatCompletion(w, r, req, handle, outputPolicy)
}

func (s *Server) completeChatCompletion(w http.ResponseWriter, r *http.Request, req chatCompletionRequest, handle edgeservice.RunResult, outputPolicy strictOutputPolicy) {
	text, reasoning, usage, err := collectRunResult(r.Context(), handle.Stream(), handle.WaitTimeout())
	if err != nil {
		writeError(w, httpStatusForRunError(err), "run_error", err.Error())
		return
	}
	text, reasoning, normalized := normalizeCompletionOutput(outputPolicy, text, reasoning)
	s.logger.Info("openai chat completion output",
		zap.String("run_id", handle.Dispatch().RunID),
		zap.Bool("strict_output", outputPolicy.Strict),
		zap.Bool("strict_stream_buffer", outputPolicy.StreamBuffer),
		zap.String("xml_completion_tool", outputPolicy.XMLCompletionTool),
		zap.Bool("normalized", normalized),
		zap.Int("content_len", len(text)),
		zap.Int("reasoning_len", len(reasoning)),
		zap.String("content_preview", previewString(text, 1000)),
	)
	created := time.Now().Unix()
	writeJSON(w, http.StatusOK, chatCompletionResponse{
		ID:      "chatcmpl-" + handle.Dispatch().RunID,
		Object:  "chat.completion",
		Created: created,
		Model:   responseModel(req.Model, handle.Dispatch().Target),
		Choices: []chatCompletionChoice{{
			Index:        0,
			Message:      chatMessage{Role: "assistant", Content: text, ReasoningContent: reasoning},
			FinishReason: "stop",
		}},
		Usage: usage,
	})
}

func (s *Server) resolveAdapter() string {
	if s.cfg.Adapter != "" {
		return s.cfg.Adapter
	}
	return "ollama"
}

func (s *Server) resolveTarget(model string) string {
	if s.cfg.Target != "" {
		return s.cfg.Target
	}
	return strings.TrimSpace(model)
}

func (s *Server) resolveSessionID() string {
	if s.cfg.SessionID != "" {
		return s.cfg.SessionID
	}
	return edgeservice.DefaultSessionID
}

func (s *Server) resolveTimeoutSec() int {
	if s.cfg.TimeoutSec > 0 {
		return s.cfg.TimeoutSec
	}
	return edgeservice.DefaultTimeoutSec
}

func (s *Server) resolveStrictOutput() bool {
	return s.cfg.StrictOutput
}

func (s *Server) resolveStrictStreamBuffer() bool {
	return s.cfg.StrictStreamBuffer
}

func (s *Server) resolveOutputPolicy(prompt string) strictOutputPolicy {
	policy := strictOutputPolicy{
		Strict:       s.resolveStrictOutput(),
		StreamBuffer: s.resolveStrictStreamBuffer(),
	}
	if !policy.Strict {
		return policy
	}
	policy.XMLCompletionTool, policy.XMLResultTag = inferXMLCompletionContract(prompt)
	policy.ContractInstruction = policy.XMLCompletionTool != ""
	return policy
}

func promptFromMessages(messages []chatMessage) string {
	var b strings.Builder
	for _, msg := range messages {
		content := strings.TrimSpace(msg.Content)
		if content == "" {
			continue
		}
		role := strings.TrimSpace(msg.Role)
		if role == "" {
			role = "user"
		}
		if b.Len() > 0 {
			b.WriteString("\n")
		}
		b.WriteString(role)
		b.WriteString(": ")
		b.WriteString(content)
	}
	return b.String()
}

func responseModel(requestModel, target string) string {
	if requestModel != "" {
		return requestModel
	}
	return target
}

func httpStatusForRunError(err error) int {
	if errors.Is(err, context.Canceled) {
		return http.StatusRequestTimeout
	}
	return http.StatusBadGateway
}

func writeJSON(w http.ResponseWriter, status int, v any) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(status)
	_ = json.NewEncoder(w).Encode(v)
}

func writeError(w http.ResponseWriter, status int, code, message string) {
	writeJSON(w, status, errorResponse{Error: errorBody{Type: code, Message: message}})
}