chore: use new type for embedding

embedder/observer/observer.go

@@ -0,0 +1,49 @@
+package observer
+
+import (
+	"fmt"
+
+	"github.com/henomis/lingoose/embedder"
+	obs "github.com/henomis/lingoose/observer"
+	"github.com/henomis/lingoose/types"
+)
+
+type EmbeddingObserver interface {
+	Embedding(*obs.Embedding) (*obs.Embedding, error)
+	EmbeddingEnd(*obs.Embedding) (*obs.Embedding, error)
+}
+
+func StartObserveEmbedding(
+	o EmbeddingObserver,
+	name string,
+	modelName string,
+	ModelParameters types.M,
+	traceID string,
+	parentID string,
+	texts []string,
+) (*obs.Embedding, error) {
+	embedding, err := o.Embedding(
+		&obs.Embedding{
+			TraceID:         traceID,
+			ParentID:        parentID,
+			Name:            fmt.Sprintf("embedding-%s", name),
+			Model:           modelName,
+			ModelParameters: ModelParameters,
+			Input:           texts,
+		},
+	)
+	if err != nil {
+		return nil, err
+	}
+	return embedding, nil
+}
+
+func StopObserveEmbedding(
+	o EmbeddingObserver,
+	embedding *obs.Embedding,
+	embeddings []embedder.Embedding,
+) error {
+	embedding.Output = embeddings
+	_, err := o.EmbeddingEnd(embedding)
+	return err
+}

embedder/openai/openai.go

@@ -5,6 +5,8 @@ import (
 	"os"
 
 	"github.com/henomis/lingoose/embedder"
+	embobserver "github.com/henomis/lingoose/embedder/observer"
+	"github.com/henomis/lingoose/observer"
 	"github.com/sashabaranov/go-openai"
 )
 
@@ -17,8 +19,11 @@ const (
 )
 
 type OpenAIEmbedder struct {
-	openAIClient *openai.Client
-	model        Model
+	openAIClient    *openai.Client
+	model           Model
+	Name            string
+	observer        embobserver.EmbeddingObserver
+	observerTraceID string
 }
 
 func New(model Model) *OpenAIEmbedder {
@@ -27,6 +32,7 @@ func New(model Model) *OpenAIEmbedder {
 	return &OpenAIEmbedder{
 		openAIClient: openai.NewClient(openAIKey),
 		model:        model,
+		Name:         "openai",
 	}
 }
 
@@ -36,9 +42,49 @@ func (o *OpenAIEmbedder) WithClient(client *openai.Client) *OpenAIEmbedder {
 	return o
 }
 
+func (o *OpenAIEmbedder) WithObserver(observer embobserver.EmbeddingObserver, traceID string) *OpenAIEmbedder {
+	o.observer = observer
+	o.observerTraceID = traceID
+	return o
+}
+
 // Embed returns the embeddings for the given texts
 func (o *OpenAIEmbedder) Embed(ctx context.Context, texts []string) ([]embedder.Embedding, error) {
-	return o.openAICreateEmebeddings(ctx, texts)
+	var observerEmbedding *observer.Embedding
+	var err error
+
+	if o.observer != nil {
+		observerEmbedding, err = embobserver.StartObserveEmbedding(
+			o.observer,
+			o.Name,
+			string(o.model),
+			nil,
+			o.observerTraceID,
+			observer.ContextValueParentID(ctx),
+			texts,
+		)
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	embeddings, err := o.openAICreateEmebeddings(ctx, texts)
+	if err != nil {
+		return nil, err
+	}
+
+	if o.observer != nil {
+		err = embobserver.StopObserveEmbedding(
+			o.observer,
+			observerEmbedding,
+			embeddings,
+		)
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	return embeddings, nil
 }
 
 func (o *OpenAIEmbedder) openAICreateEmebeddings(ctx context.Context, texts []string) ([]embedder.Embedding, error) {

observer/langfuse/formatter.go

@@ -82,6 +82,20 @@ func observerGenerationToLangfuseGeneration(g *observer.Generation) *model.Gener
 	}
 }
 
+func observerEmbeddingToLangfuseGeneration(e *observer.Embedding) *model.Generation {
+	return &model.Generation{
+		ID:                  e.ID,
+		TraceID:             e.TraceID,
+		Name:                e.Name,
+		ParentObservationID: e.ParentID,
+		Model:               e.Model,
+		ModelParameters:     e.ModelParameters,
+		Input:               e.Input,
+		Output:              e.Output,
+		Metadata:            e.Metadata,
+	}
+}
+
 func observerEventToLangfuseEvent(e *observer.Event) *model.Event {
 	return &model.Event{
 		ID:                  e.ID,

observer/langfuse/langfuse.go

@@ -69,6 +69,25 @@ func (l *Langfuse) GenerationEnd(g *observer.Generation) (*observer.Generation,
 	return g, nil
 }
 
+func (l *Langfuse) Embedding(e *observer.Embedding) (*observer.Embedding, error) {
+	langfuseGeneration := observerEmbeddingToLangfuseGeneration(e)
+	langfuseGeneration, err := l.client.Generation(langfuseGeneration, nil)
+	if err != nil {
+		return nil, err
+	}
+	e.ID = langfuseGeneration.ID
+	return e, nil
+}
+
+func (l *Langfuse) EmbeddingEnd(e *observer.Embedding) (*observer.Embedding, error) {
+	langfuseGeneration := observerEmbeddingToLangfuseGeneration(e)
+	_, err := l.client.Generation(langfuseGeneration, nil)
+	if err != nil {
+		return nil, err
+	}
+	return e, nil
+}
+
 func (l *Langfuse) Event(e *observer.Event) (*observer.Event, error) {
 	langfuseEvent := observerEventToLangfuseEvent(e)
 	langfuseEvent, err := l.client.Event(langfuseEvent, nil)

observer/observer.go

@@ -3,6 +3,7 @@ package observer
 import (
 	"context"
 
+	"github.com/henomis/lingoose/embedder"
 	"github.com/henomis/lingoose/thread"
 	"github.com/henomis/lingoose/types"
 )
@@ -38,6 +39,18 @@ type Generation struct {
 	Metadata        types.M
 }
 
+type Embedding struct {
+	ID              string
+	ParentID        string
+	TraceID         string
+	Name            string
+	Model           string
+	ModelParameters types.M
+	Input           []string
+	Output          []embedder.Embedding
+	Metadata        types.M
+}
+
 type Event struct {
 	ID       string
 	ParentID string