main.go

// Package main demonstrates generating text embeddings using blmstudio.
//
// Embeddings turn text into dense numeric vectors useful for semantic search,
// clustering, and similarity comparisons.
//
// This example embeds two short phrases and computes their cosine similarity
// so you can see a concrete output immediately.
//
// # Providers
//
//	🏠 LM Studio (default)  — run any embedding model locally
//	☁️  OpenAI              — set LM_BASE_URL=https://api.openai.com and LM_API_KEY=sk-...
//
// # Run
//
//	# LM Studio (local) — load an embedding model first in LM Studio
//	LM_EMBED_MODEL=nomic-ai/nomic-embed-text-v1.5 go run main.go
//
//	# OpenAI
//	LM_BASE_URL=https://api.openai.com LM_API_KEY=sk-... LM_EMBED_MODEL=text-embedding-3-small go run main.go
package main

import (
	"context"
	"fmt"
	"math"
	"os"

	"github.com/bearaujus/blmstudio"
)

func main() {
	// ── Provider configuration ───────────────────────────────────────────────
	baseURL := getenv("LM_BASE_URL", blmstudio.DefaultBaseURL)
	apiKey := os.Getenv("LM_API_KEY")
	model := getenv("LM_EMBED_MODEL", "nomic-ai/nomic-embed-text-v1.5")

	opts := []blmstudio.Option{blmstudio.WithBaseURL(baseURL)}
	if apiKey != "" {
		opts = append(opts, blmstudio.WithAPIToken(apiKey))
	}

	client := blmstudio.New(opts...)

	// ── Texts to embed ───────────────────────────────────────────────────────
	texts := []string{
		"Hello, world!",
		"Hi there, universe!",
	}

	fmt.Printf("→ model:    %s\n", model)
	fmt.Printf("→ base URL: %s\n\n", baseURL)

	// ── Generate embeddings (batch) ──────────────────────────────────────────
	resp, err := client.Embed(context.Background(), &blmstudio.EmbeddingRequest{
		Model: model,
		Input: blmstudio.NewEmbeddingStringsInput(texts),
	})
	if err != nil {
		if ae, ok := blmstudio.IsAPIError(err); ok {
			fmt.Fprintf(os.Stderr, "API error HTTP %d: %s\n", ae.StatusCode, ae.Message)
			os.Exit(1)
		}
		fmt.Fprintf(os.Stderr, "error: %v\n", err)
		os.Exit(1)
	}

	// ── Print results ────────────────────────────────────────────────────────
	for _, obj := range resp.Data {
		fmt.Printf("text[%d] %q → %d-dim vector  (first 4: %v)\n",
			obj.Index,
			texts[obj.Index],
			len(obj.Embedding),
			round4(obj.Embedding[:min(4, len(obj.Embedding))]),
		)
	}

	// ── Compute cosine similarity between the two vectors ────────────────────
	if len(resp.Data) >= 2 {
		// Data is ordered by Index
		vecs := make([][]float64, len(resp.Data))
		for _, obj := range resp.Data {
			vecs[obj.Index] = obj.Embedding
		}
		sim := cosineSimilarity(vecs[0], vecs[1])
		fmt.Printf("\ncosine similarity: %.4f\n", sim)
	}

	fmt.Printf("\ntokens used: %d\n", resp.Usage.TotalTokens)
}

// cosineSimilarity returns the cosine similarity between two vectors.
func cosineSimilarity(a, b []float64) float64 {
	if len(a) != len(b) || len(a) == 0 {
		return 0
	}
	var dot, normA, normB float64
	for i := range a {
		dot += a[i] * b[i]
		normA += a[i] * a[i]
		normB += b[i] * b[i]
	}
	if normA == 0 || normB == 0 {
		return 0
	}
	return dot / (math.Sqrt(normA) * math.Sqrt(normB))
}

// round4 returns floats rounded to 4 decimal places for display.
func round4(vs []float64) []float64 {
	out := make([]float64, len(vs))
	for i, v := range vs {
		out[i] = math.Round(v*10000) / 10000
	}
	return out
}

func min(a, b int) int {
	if a < b {
		return a
	}
	return b
}

func getenv(key, fallback string) string {
	if v := os.Getenv(key); v != "" {
		return v
	}
	return fallback
}