iasysadmin/internal/llm/client.go

package llm

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

// maxResponseBytes limita el tamaño de una respuesta HTTP del modelo para
// que un servidor defectuoso no pueda agotar la memoria del agente.
const maxResponseBytes = 64 << 20

// OpenAIClient implementa Client sobre la API estilo OpenAI.
//
// Separa el transporte HTTP de la lógica del agente: este tipo solo
// serializa peticiones, gestiona reintentos y decodifica respuestas.
type OpenAIClient struct {
	baseURL    string
	apiKey     string
	model      string
	userAgent  string
	http       *http.Client
	maxRetries int
	backoff    time.Duration
}

// Options configura un OpenAIClient.
type Options struct {
	BaseURL   string
	APIKey    string // se añade solo en la cabecera Authorization, nunca se registra
	Model     string
	UserAgent string
	Timeout   time.Duration
}

// NewOpenAIClient crea un cliente con un único http.Client reutilizable
// (pool de conexiones incluido en http.Transport).
func NewOpenAIClient(o Options) *OpenAIClient {
	transport := http.DefaultTransport.(*http.Transport).Clone()
	transport.MaxIdleConnsPerHost = 4
	return &OpenAIClient{
		baseURL:    strings.TrimRight(o.BaseURL, "/"),
		apiKey:     o.APIKey,
		model:      o.Model,
		userAgent:  o.UserAgent,
		http:       &http.Client{Timeout: o.Timeout, Transport: transport},
		maxRetries: 3,
		backoff:    time.Second,
	}
}

// Formato de cable de la API OpenAI.
type wireRequest struct {
	Model       string    `json:"model"`
	Messages    []Message `json:"messages"`
	Tools       []Tool    `json:"tools,omitempty"`
	MaxTokens   int       `json:"max_tokens,omitempty"`
	Temperature *float64  `json:"temperature,omitempty"`
	Stream      bool      `json:"stream"`
}

type wireResponse struct {
	Choices []struct {
		Message struct {
			Role             string     `json:"role"`
			Content          *string    `json:"content"`
			ReasoningContent string     `json:"reasoning_content"`
			ToolCalls        []ToolCall `json:"tool_calls"`
		} `json:"message"`
		FinishReason string `json:"finish_reason"`
	} `json:"choices"`
	Usage *struct {
		PromptTokens     int `json:"prompt_tokens"`
		CompletionTokens int `json:"completion_tokens"`
		TotalTokens      int `json:"total_tokens"`
	} `json:"usage"`
	Error *struct {
		Message string `json:"message"`
	} `json:"error"`
}

// thinkRe elimina bloques <think>...</think> que algunos modelos (Qwen3,
// DeepSeek-R1) incluyen en el contenido cuando el servidor no los separa.
var thinkRe = regexp.MustCompile(`(?s)^\s*<think>(.*?)</think>\s*`)

// Chat envía la conversación al modelo y devuelve su respuesta.
//
// Reintenta errores transitorios (red, 429, 5xx) con backoff exponencial y
// un número máximo de intentos; nunca reintenta indefinidamente.
func (c *OpenAIClient) Chat(ctx context.Context, req ChatRequest) (ChatResponse, error) {
	body, err := json.Marshal(wireRequest{
		Model:       c.model,
		Messages:    req.Messages,
		Tools:       req.Tools,
		MaxTokens:   req.MaxTokens,
		Temperature: req.Temperature,
	})
	if err != nil {
		return ChatResponse{}, fmt.Errorf("failed to encode chat request: %w", err)
	}

	var lastErr error
	for attempt := 0; attempt <= c.maxRetries; attempt++ {
		if attempt > 0 {
			wait := c.backoff << (attempt - 1)
			select {
			case <-ctx.Done():
				return ChatResponse{}, ctx.Err()
			case <-time.After(wait):
			}
		}
		resp, retry, err := c.doChat(ctx, body)
		if err == nil {
			return resp, nil
		}
		lastErr = err
		if !retry || ctx.Err() != nil {
			break
		}
	}
	return ChatResponse{}, lastErr
}

// doChat realiza un único intento. retry indica si el error es transitorio.
func (c *OpenAIClient) doChat(ctx context.Context, body []byte) (ChatResponse, bool, error) {
	url := c.baseURL + "/chat/completions"
	httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(body))
	if err != nil {
		return ChatResponse{}, false, fmt.Errorf("failed to build request for %s: %w", url, err)
	}
	c.setHeaders(httpReq)
	httpReq.Header.Set("Content-Type", "application/json")

	httpResp, err := c.http.Do(httpReq)
	if err != nil {
		return ChatResponse{}, isTransient(err), fmt.Errorf("POST %s: %w", url, err)
	}
	defer httpResp.Body.Close()

	data, err := io.ReadAll(io.LimitReader(httpResp.Body, maxResponseBytes))
	if err != nil {
		return ChatResponse{}, true, fmt.Errorf("failed to read response from %s: %w", url, err)
	}

	if httpResp.StatusCode != http.StatusOK {
		retry := httpResp.StatusCode == http.StatusTooManyRequests || httpResp.StatusCode >= 500
		return ChatResponse{}, retry, &HTTPError{URL: url, Status: httpResp.StatusCode, Body: snippet(data)}
	}

	var w wireResponse
	if err := json.Unmarshal(data, &w); err != nil {
		return ChatResponse{}, false, fmt.Errorf("invalid JSON from %s: %w (body: %s)", url, err, snippet(data))
	}
	if w.Error != nil && w.Error.Message != "" {
		return ChatResponse{}, false, fmt.Errorf("model server error: %s", w.Error.Message)
	}
	if len(w.Choices) == 0 {
		return ChatResponse{}, false, fmt.Errorf("model server returned no choices (body: %s)", snippet(data))
	}

	ch := w.Choices[0]
	out := ChatResponse{
		FinishReason: ch.FinishReason,
		Reasoning:    ch.Message.ReasoningContent,
		Message: Message{
			Role:      ch.Message.Role,
			ToolCalls: ch.Message.ToolCalls,
		},
	}
	if out.Message.Role == "" {
		out.Message.Role = RoleAssistant
	}
	if ch.Message.Content != nil {
		content := *ch.Message.Content
		if m := thinkRe.FindStringSubmatch(content); m != nil {
			if out.Reasoning == "" {
				out.Reasoning = strings.TrimSpace(m[1])
			}
			content = content[len(m[0]):]
		}
		out.Message.Content = content
	}
	if w.Usage != nil {
		out.Usage = Usage{
			PromptTokens:     w.Usage.PromptTokens,
			CompletionTokens: w.Usage.CompletionTokens,
			TotalTokens:      w.Usage.TotalTokens,
		}
	}
	return out, false, nil
}

// HTTPError es una respuesta no-200 del servidor de inferencia. Permite al
// agente reaccionar a casos concretos (p.ej. servidor sin soporte de tools
// o contexto excedido).
type HTTPError struct {
	URL    string
	Status int
	Body   string
}

func (e *HTTPError) Error() string {
	return fmt.Sprintf("POST %s: HTTP %d: %s", e.URL, e.Status, e.Body)
}

// IsContextOverflow detecta, de forma heurística, errores por exceso de contexto.
func IsContextOverflow(err error) bool {
	var he *HTTPError
	if !errors.As(err, &he) || he.Status < 400 || he.Status > 500 {
		return false
	}
	b := strings.ToLower(he.Body)
	return strings.Contains(b, "context") && (strings.Contains(b, "exceed") || strings.Contains(b, "length") || strings.Contains(b, "too long") || strings.Contains(b, "overflow")) ||
		strings.Contains(b, "too many tokens") || strings.Contains(b, "maximum context")
}

// IsToolsUnsupported detecta, de forma heurística, que el servidor rechaza
// el parámetro "tools".
func IsToolsUnsupported(err error) bool {
	var he *HTTPError
	if !errors.As(err, &he) || he.Status < 400 || he.Status > 501 {
		return false
	}
	b := strings.ToLower(he.Body)
	return strings.Contains(b, "tool") || strings.Contains(b, "function") || he.Status == 501
}

// DiscoverModel intenta averiguar el tamaño de contexto que usa el servidor,
// para cumplir CONTEXT_SIZE=0 sin inventar un número.
//
// Se consultan, en orden, formatos conocidos:
//   - GET {base}/models: campos max_model_len (vLLM), context_length
//     (OpenRouter, LM Studio) o meta.n_ctx_train (llama.cpp).
//   - GET {origin}/props: default_generation_settings.n_ctx (llama.cpp),
//     que es el contexto realmente configurado en el servidor y por tanto
//     tiene prioridad sobre n_ctx_train.
//
// Si nada funciona, devuelve ContextSize 0 (UNKNOWN) sin error: no conocer
// el tamaño no debe impedir el arranque.
func (c *OpenAIClient) DiscoverModel(ctx context.Context) (ModelInfo, error) {
	info := ModelInfo{ID: c.model}

	var models struct {
		Data []struct {
			ID            string `json:"id"`
			MaxModelLen   int    `json:"max_model_len"`
			ContextLength int    `json:"context_length"`
			Meta          struct {
				NCtxTrain int `json:"n_ctx_train"`
			} `json:"meta"`
		} `json:"data"`
	}
	modelsErr := c.getJSON(ctx, c.baseURL+"/models", &models)
	if modelsErr == nil {
		// Se usa la entrada con el ID exacto; si el servidor expone un único
		// modelo (caso típico de llama.cpp), se acepta aunque el ID difiera.
		idx := -1
		for i, m := range models.Data {
			if m.ID == c.model {
				idx = i
				break
			}
		}
		if idx < 0 && len(models.Data) == 1 {
			idx = 0
		}
		if idx >= 0 {
			m := models.Data[idx]
			switch {
			case m.MaxModelLen > 0:
				info.ContextSize, info.Source = m.MaxModelLen, "server /models max_model_len"
			case m.ContextLength > 0:
				info.ContextSize, info.Source = m.ContextLength, "server /models context_length"
			case m.Meta.NCtxTrain > 0:
				info.ContextSize, info.Source = m.Meta.NCtxTrain, "server /models meta.n_ctx_train"
			}
		}
	}

	var props struct {
		DefaultGenerationSettings struct {
			NCtx int `json:"n_ctx"`
		} `json:"default_generation_settings"`
	}
	if err := c.getJSON(ctx, originOf(c.baseURL)+"/props", &props); err == nil && props.DefaultGenerationSettings.NCtx > 0 {
		info.ContextSize, info.Source = props.DefaultGenerationSettings.NCtx, "server /props n_ctx"
	}

	if modelsErr != nil && info.ContextSize == 0 {
		return info, modelsErr
	}
	return info, nil
}

func (c *OpenAIClient) getJSON(ctx context.Context, url string, v any) error {
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
	if err != nil {
		return fmt.Errorf("failed to build request for %s: %w", url, err)
	}
	c.setHeaders(req)
	resp, err := c.http.Do(req)
	if err != nil {
		return fmt.Errorf("GET %s: %w", url, err)
	}
	defer resp.Body.Close()
	data, err := io.ReadAll(io.LimitReader(resp.Body, maxResponseBytes))
	if err != nil {
		return fmt.Errorf("failed to read %s: %w", url, err)
	}
	if resp.StatusCode != http.StatusOK {
		return fmt.Errorf("GET %s: HTTP %d", url, resp.StatusCode)
	}
	if err := json.Unmarshal(data, v); err != nil {
		return fmt.Errorf("invalid JSON from %s: %w", url, err)
	}
	return nil
}

// setHeaders añade autenticación y User-Agent. La API key solo vive aquí.
func (c *OpenAIClient) setHeaders(r *http.Request) {
	if c.apiKey != "" {
		r.Header.Set("Authorization", "Bearer "+c.apiKey)
	}
	if c.userAgent != "" {
		r.Header.Set("User-Agent", c.userAgent)
	}
	r.Header.Set("Accept", "application/json")
}

// isTransient decide si un error de red merece reintento.
func isTransient(err error) bool {
	if errors.Is(err, context.Canceled) {
		return false
	}
	var ne net.Error
	if errors.As(err, &ne) {
		return true
	}
	return errors.Is(err, io.ErrUnexpectedEOF) || errors.Is(err, io.EOF)
}

// originOf convierte "http://host:8080/v1" en "http://host:8080".
func originOf(base string) string {
	if i := strings.Index(base, "://"); i >= 0 {
		if j := strings.Index(base[i+3:], "/"); j >= 0 {
			return base[:i+3+j]
		}
	}
	return base
}

// snippet recorta cuerpos de error para mensajes legibles.
func snippet(b []byte) string {
	s := strings.TrimSpace(string(b))
	if len(s) > 300 {
		s = s[:300] + "..."
	}
	return s
}