vecna/pkg/server/openai.go

package server

import (
	"encoding/json"
	"fmt"
	"net/http"
	"time"

	"github.com/uptrace/bunrouter"

	"github.com/Warky-Devs/vecna.git/pkg/adapter"
	"github.com/Warky-Devs/vecna.git/pkg/embedclient"
)

type openAIEmbedRequest struct {
	Input interface{} `json:"input"` // string or []string
	Model string      `json:"model"`
}

type openAIEmbedResponse struct {
	Object string             `json:"object"`
	Data   []openAIEmbedDatum `json:"data"`
	Model  string             `json:"model"`
	Usage  openAIUsage        `json:"usage"`
}

type openAIEmbedDatum struct {
	Object    string    `json:"object"`
	Embedding []float32 `json:"embedding"`
	Index     int       `json:"index"`
}

type openAIUsage struct {
	PromptTokens int `json:"prompt_tokens"`
	TotalTokens  int `json:"total_tokens"`
}

func (h *handler) openAIEmbeddings(w http.ResponseWriter, req bunrouter.Request) error {
	return h.openAIEmbeddingsWithAdapter(w, req, h.adapter, "")
}

func (h *handler) openAIEmbeddingsMapped(w http.ResponseWriter, req bunrouter.Request) error {
	em, err := h.resolveExtraMap(req.Param("mapping"))
	if err != nil {
		return writeJSON(w, http.StatusNotFound, map[string]string{"error": err.Error()})
	}
	return h.openAIEmbeddingsWithAdapter(w, req, em.Adapter, em.ForwardTarget)
}

func (h *handler) openAIEmbeddingsWithAdapter(w http.ResponseWriter, req bunrouter.Request, adp adapter.Adapter, targetOverride string) error {
	var body openAIEmbedRequest
	if err := json.NewDecoder(req.Body).Decode(&body); err != nil {
		return writeJSON(w, http.StatusBadRequest, map[string]string{"error": "invalid request body"})
	}

	texts, err := toStringSlice(body.Input)
	if err != nil {
		return writeJSON(w, http.StatusBadRequest, map[string]string{"error": err.Error()})
	}

	client, targetName, targetURL := h.resolveClientOverride(targetOverride, body.Model)
	trace := TraceFromContext(req.Context())
	trace.ForwardTarget = targetName
	trace.ForwardURL = targetURL

	t0 := time.Now()
	embedResp, err := client.Embed(req.Context(), embedclient.Request{Texts: texts, Model: body.Model})
	trace.ForwardDuration = time.Since(t0)
	if err != nil {
		return writeJSON(w, http.StatusBadGateway, map[string]string{"error": err.Error()})
	}
	trace.ForwardModel = embedResp.Model
	trace.PromptTokens = embedResp.Usage.PromptTokens
	trace.TotalTokens = embedResp.Usage.TotalTokens

	t1 := time.Now()
	data := make([]openAIEmbedDatum, len(embedResp.Embeddings))
	for i, vec := range embedResp.Embeddings {
		adapted, adaptErr := adp.Adapt(vec)
		if adaptErr != nil {
			return writeJSON(w, http.StatusInternalServerError, map[string]string{"error": adaptErr.Error()})
		}
		data[i] = openAIEmbedDatum{Object: "embedding", Embedding: adapted, Index: i}
	}
	trace.TranslateDuration = time.Since(t1)

	writeTraceHeaders(w, trace)

	return writeJSON(w, http.StatusOK, openAIEmbedResponse{
		Object: "list",
		Data:   data,
		Model:  embedResp.Model,
		Usage:  openAIUsage{PromptTokens: embedResp.Usage.PromptTokens, TotalTokens: embedResp.Usage.TotalTokens},
	})
}

// toStringSlice accepts a JSON string or array of strings.
func toStringSlice(v interface{}) ([]string, error) {
	switch val := v.(type) {
	case string:
		return []string{val}, nil
	case []interface{}:
		out := make([]string, len(val))
		for i, item := range val {
			s, ok := item.(string)
			if !ok {
				return nil, fmt.Errorf("input array element %d is not a string", i)
			}
			out[i] = s
		}
		return out, nil
	default:
		return nil, fmt.Errorf("input must be a string or array of strings")
	}
}