rony-llm-agent/pkg/rag/embeddings/openai_compatible.go

package embeddings

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"sync/atomic"
)

// OpenAICompatibleConfig holds the settings for any embeddings API that
// follows OpenAI's request/response shape: POST {BaseURL}/embeddings with
// {"input": ..., "model": ...}, returning {"data": [{"embedding": [...]}]}.
// This covers llama.cpp (started with --embeddings), real OpenAI, and most
// third-party providers advertised as "OpenAI-compatible".
type OpenAICompatibleConfig struct {
	BaseURL string // e.g. "http://localhost:8080/v1" or "https://api.openai.com/v1"
	APIKey  string // sent as "Authorization: Bearer <key>" when non-empty; local servers like llama.cpp don't need one
	Model   string // e.g. "text-embedding-3-small"; ignored by servers that only have one model loaded
}

// OpenAICompatible implements Embedder against any OpenAI-shaped
// /embeddings endpoint.
type OpenAICompatible struct {
	baseURL string
	apiKey  string
	model   string
	http    *http.Client
	dims    atomic.Int64 // lazily learned from the first successful response
}

// NewOpenAICompatible creates a new OpenAI-shaped embedder.
func NewOpenAICompatible(cfg OpenAICompatibleConfig) (*OpenAICompatible, error) {
	if cfg.BaseURL == "" {
		return nil, fmt.Errorf("base URL is required")
	}
	return &OpenAICompatible{
		baseURL: cfg.BaseURL,
		apiKey:  cfg.APIKey,
		model:   cfg.Model,
		http:    http.DefaultClient,
	}, nil
}

// NewLlamaCpp is a convenience constructor for a local llama.cpp server
// (defaults to http://localhost:8080/v1, no API key). The server must have
// been started with the `--embeddings` flag, otherwise every call fails
// (llama.cpp returns a 501). Quality depends on the loaded model: dedicated
// embedding models (e.g. nomic-embed-text, bge-m3) work best, but a
// chat/instruct model still produces a usable semantic vector via pooling.
func NewLlamaCpp(cfg LlamaCppConfig) (*OpenAICompatible, error) {
	baseURL := cfg.BaseURL
	if baseURL == "" {
		baseURL = "http://localhost:8080/v1"
	}
	return NewOpenAICompatible(OpenAICompatibleConfig{BaseURL: baseURL, Model: cfg.Model})
}

// LlamaCppConfig holds the settings for NewLlamaCpp.
type LlamaCppConfig struct {
	BaseURL string // e.g. "http://localhost:8080/v1"
	Model   string // optional; llama.cpp embeds with whatever model is loaded regardless of this value
}

func (e *OpenAICompatible) Embed(ctx context.Context, text string) ([]float32, error) {
	endpoint := e.baseURL + "/embeddings"

	reqBody, err := json.Marshal(map[string]interface{}{
		"input": text,
		"model": e.model,
	})
	if err != nil {
		return nil, fmt.Errorf("marshaling request: %w", err)
	}

	httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, bytes.NewReader(reqBody))
	if err != nil {
		return nil, fmt.Errorf("creating request: %w", err)
	}
	httpReq.Header.Set("Content-Type", "application/json")
	if e.apiKey != "" {
		httpReq.Header.Set("Authorization", "Bearer "+e.apiKey)
	}

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

	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		body, _ := io.ReadAll(resp.Body)
		return nil, fmt.Errorf("API error %d: %s", resp.StatusCode, string(body))
	}

	var apiResp openAICompatibleEmbedResponse
	if err := json.NewDecoder(resp.Body).Decode(&apiResp); err != nil {
		return nil, fmt.Errorf("decoding response: %w", err)
	}
	if len(apiResp.Data) == 0 || len(apiResp.Data[0].Embedding) == 0 {
		return nil, fmt.Errorf("empty embeddings response")
	}

	vector := apiResp.Data[0].Embedding
	e.dims.Store(int64(len(vector)))
	return vector, nil
}

// Dimensions returns the vector size learned from the last successful Embed
// call, or 0 if none has succeeded yet (it depends on the model/provider
// behind BaseURL, so it can't be known upfront).
func (e *OpenAICompatible) Dimensions() int {
	return int(e.dims.Load())
}

type openAICompatibleEmbedResponse struct {
	Data []struct {
		Embedding []float32 `json:"embedding"`
		Index     int       `json:"index"`
	} `json:"data"`
}