Embedding - Ollama

  go get github.com/cloudwego/eino-ext/components/embedding/ollama

package main

import (
	"context"
	"log"
	"os"
	"time"

	"github.com/cloudwego/eino-ext/components/embedding/ollama"
)

func main() {
	ctx := context.Background()

	baseURL := os.Getenv("OLLAMA_BASE_URL")
	if baseURL == "" {
		baseURL = "http://localhost:11434" // 默认本地
	}
	model := os.Getenv("OLLAMA_EMBED_MODEL")
	if model == "" {
		model = "nomic-embed-text"
	}

	embedder, err := ollama.NewEmbedder(ctx, &ollama.EmbeddingConfig{
		BaseURL: baseURL,
		Model:   model,
		Timeout: 10 * time.Second,
	})
	if err != nil {
		log.Fatalf("NewEmbedder of ollama error: %v", err)
		return
	}

	log.Printf("===== call Embedder directly =====")

	vectors, err := embedder.EmbedStrings(ctx, []string{"hello", "how are you"})
	if err != nil {
		log.Fatalf("EmbedStrings of Ollama failed, err=%v", err)
	}

	log.Printf("vectors : %v", vectors)

	// you can use WithModel to specify the model
	vectors, err = embedder.EmbedStrings(ctx, []string{"hello", "how are you"}, embedding.WithModel(model))
	if err != nil {
		log.Fatalf("EmbedStrings of Ollama failed, err=%v", err)
	}

	log.Printf("vectors : %v", vectors)
}

type EmbeddingConfig struct {
    // Timeout specifies the maximum duration to wait for API responses
    // If HTTPClient is set, Timeout will not be used.
    // Optional. Default: no timeout
    Timeout time.Duration `json:"timeout"`
    
    // HTTPClient specifies the client to send HTTP requests.
    // If HTTPClient is set, Timeout will not be used.
    // Optional. Default &http.Client{Timeout: Timeout}
    HTTPClient *http.Client `json:"http_client"`
    
    // BaseURL specifies the Ollama service endpoint URL
    // Format: http(s)://host:port
    // Optional. Default: "http://localhost:11434"
    BaseURL string `json:"base_url"`
    
    // Model specifies the ID of the model to use for embedding generation
    // Required. You can also set it when calling the EmbedStrings with `embedding.WithModel(model)`
    Model string `json:"model"`
    
    // Truncate specifies whether to truncate text to model's maximum context length
    // When set to true, if text to embed exceeds the model's maximum context length,
    // a call to EmbedStrings will return an error
    // Optional.
    Truncate *bool `json:"truncate,omitempty"`
    
    // KeepAlive controls how long the model will stay loaded in memory following this request.
    // Optional. Default 5 minutes
    KeepAlive *time.Duration `json:"keep_alive,omitempty"`
    
    // Options lists model-specific options.
    // Optional
    Options map[string]any `json:"options,omitempty"`
}