116 lines
3.4 KiB
Go
116 lines
3.4 KiB
Go
package vertex
|
|
|
|
import (
|
|
"github.com/maximhq/bifrost/core/schemas"
|
|
)
|
|
|
|
// ToVertexEmbeddingRequest converts a Bifrost embedding request to Vertex AI format
|
|
func ToVertexEmbeddingRequest(bifrostReq *schemas.BifrostEmbeddingRequest) *VertexEmbeddingRequest {
|
|
if bifrostReq == nil || bifrostReq.Input == nil || (bifrostReq.Input.Text == nil && bifrostReq.Input.Texts == nil) {
|
|
return nil
|
|
}
|
|
// Create the request
|
|
vertexReq := &VertexEmbeddingRequest{}
|
|
if bifrostReq.Params != nil {
|
|
vertexReq.ExtraParams = bifrostReq.Params.ExtraParams
|
|
}
|
|
var texts []string
|
|
if bifrostReq.Input.Text != nil {
|
|
texts = []string{*bifrostReq.Input.Text}
|
|
} else {
|
|
texts = bifrostReq.Input.Texts
|
|
}
|
|
|
|
// Create instances for each text
|
|
instances := make([]VertexEmbeddingInstance, 0, len(texts))
|
|
for _, text := range texts {
|
|
instance := VertexEmbeddingInstance{
|
|
Content: text,
|
|
}
|
|
|
|
// Add optional task_type and title from params
|
|
if bifrostReq.Params != nil {
|
|
if taskTypeStr, ok := schemas.SafeExtractStringPointer(bifrostReq.Params.ExtraParams["task_type"]); ok {
|
|
delete(vertexReq.ExtraParams, "task_type")
|
|
instance.TaskType = taskTypeStr
|
|
}
|
|
if title, ok := schemas.SafeExtractStringPointer(bifrostReq.Params.ExtraParams["title"]); ok {
|
|
delete(vertexReq.ExtraParams, "title")
|
|
instance.Title = title
|
|
}
|
|
}
|
|
|
|
instances = append(instances, instance)
|
|
}
|
|
vertexReq.Instances = instances
|
|
// Add parameters if present
|
|
if bifrostReq.Params != nil {
|
|
parameters := &VertexEmbeddingParameters{}
|
|
|
|
// Set autoTruncate (defaults to true)
|
|
autoTruncate := true
|
|
if bifrostReq.Params.ExtraParams != nil {
|
|
if autoTruncateVal, ok := schemas.SafeExtractBool(bifrostReq.Params.ExtraParams["autoTruncate"]); ok {
|
|
delete(vertexReq.ExtraParams, "autoTruncate")
|
|
autoTruncate = autoTruncateVal
|
|
}
|
|
}
|
|
parameters.AutoTruncate = &autoTruncate
|
|
|
|
// Add outputDimensionality if specified
|
|
if bifrostReq.Params.Dimensions != nil {
|
|
delete(vertexReq.ExtraParams, "dimensions")
|
|
parameters.OutputDimensionality = bifrostReq.Params.Dimensions
|
|
}
|
|
|
|
vertexReq.Parameters = parameters
|
|
}
|
|
|
|
return vertexReq
|
|
}
|
|
|
|
// ToBifrostEmbeddingResponse converts a Vertex AI embedding response to Bifrost format
|
|
func (response *VertexEmbeddingResponse) ToBifrostEmbeddingResponse() *schemas.BifrostEmbeddingResponse {
|
|
if response == nil || len(response.Predictions) == 0 {
|
|
return nil
|
|
}
|
|
|
|
// Convert predictions to Bifrost embeddings
|
|
embeddings := make([]schemas.EmbeddingData, 0, len(response.Predictions))
|
|
var usage *schemas.BifrostLLMUsage
|
|
|
|
for i, prediction := range response.Predictions {
|
|
if prediction.Embeddings == nil || len(prediction.Embeddings.Values) == 0 {
|
|
continue
|
|
}
|
|
|
|
// Create embedding object
|
|
embedding := schemas.EmbeddingData{
|
|
Object: "embedding",
|
|
Embedding: schemas.EmbeddingStruct{
|
|
EmbeddingArray: append([]float64(nil), prediction.Embeddings.Values...),
|
|
},
|
|
Index: i,
|
|
}
|
|
|
|
// Extract statistics if available
|
|
if prediction.Embeddings.Statistics != nil {
|
|
if usage == nil {
|
|
usage = &schemas.BifrostLLMUsage{}
|
|
}
|
|
usage.TotalTokens += prediction.Embeddings.Statistics.TokenCount
|
|
usage.PromptTokens += prediction.Embeddings.Statistics.TokenCount
|
|
}
|
|
|
|
embeddings = append(embeddings, embedding)
|
|
}
|
|
|
|
return &schemas.BifrostEmbeddingResponse{
|
|
Object: "list",
|
|
Data: embeddings,
|
|
Usage: usage,
|
|
ExtraFields: schemas.BifrostResponseExtraFields{
|
|
},
|
|
}
|
|
}
|