mirror of
https://github.com/Warky-Devs/vecna.git
synced 2026-05-05 01:26:58 +00:00
103 lines
2.9 KiB
Go
103 lines
2.9 KiB
Go
package server
|
|
|
|
import (
|
|
"encoding/json"
|
|
"fmt"
|
|
"net/http"
|
|
"time"
|
|
|
|
"github.com/uptrace/bunrouter"
|
|
|
|
"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 {
|
|
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.resolveClient(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 := h.adapter.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")
|
|
}
|
|
}
|