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 writeError(w, http.StatusNotFound, err.Error())
	}
	return h.openAIEmbeddingsWithAdapter(w, req, em.Adapter, em.ForwardTarget)
}

// openAIEmbeddingsGet is a GET convenience variant (e.g. for browser/curl testing):
// ?input=foo&input=bar&model=name in place of a JSON body.
func (h *handler) openAIEmbeddingsGet(w http.ResponseWriter, req bunrouter.Request) error {
	return h.openAIEmbeddingsGetWithAdapter(w, req, h.adapter, "")
}

func (h *handler) openAIEmbeddingsGetMapped(w http.ResponseWriter, req bunrouter.Request) error {
	em, err := h.resolveExtraMap(req.Param("mapping"))
	if err != nil {
		return writeError(w, http.StatusNotFound, err.Error())
	}
	return h.openAIEmbeddingsGetWithAdapter(w, req, em.Adapter, em.ForwardTarget)
}

func (h *handler) openAIEmbeddingsGetWithAdapter(w http.ResponseWriter, req bunrouter.Request, adp adapter.Adapter, targetOverride string) error {
	q := req.URL.Query()
	texts := q["input"]
	if len(texts) == 0 {
		return writeError(w, http.StatusBadRequest, "query parameter \"input\" is required")
	}
	return h.processEmbeddings(w, req, adp, targetOverride, texts, q.Get("model"))
}

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 writeError(w, http.StatusBadRequest, "invalid request body")
	}

	texts, err := toStringSlice(body.Input)
	if err != nil {
		return writeError(w, http.StatusBadRequest, err.Error())
	}

	return h.processEmbeddings(w, req, adp, targetOverride, texts, body.Model)
}

func (h *handler) processEmbeddings(w http.ResponseWriter, req bunrouter.Request, adp adapter.Adapter, targetOverride string, texts []string, model string) error {
	client, targetName, targetURL := h.resolveClientOverride(targetOverride, 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: model})
	trace.ForwardDuration = time.Since(t0)
	if err != nil {
		return writeError(w, http.StatusBadGateway, 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 writeError(w, http.StatusInternalServerError, 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")
	}
}