Vector Search

package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"math"
	"net/http"
	"sort"
)

type EmbedRequest struct {
	Model string `json:"model"`
	Input string `json:"input"`
}

type EmbedResponse struct {
	Embeddings [][]float64 `json:"embeddings"`
}

type Document struct {
	Text      string
	Embedding []float64
}

func embed(text string) ([]float64, error) {
	req := EmbedRequest{Model: "nomic-embed-text", Input: text}
	body, _ := json.Marshal(req)

	resp, err := http.Post("http://localhost:11434/api/embed", "application/json", bytes.NewReader(body))
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	data, _ := io.ReadAll(resp.Body)
	var result EmbedResponse
	json.Unmarshal(data, &result)
	return result.Embeddings[0], nil
}

func cosineSimilarity(a, b []float64) float64 {
	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))
}

func search(docs []Document, query string, topK int) []Document {
	queryVec, _ := embed(query)

	type scored struct {
		doc   Document
		score float64
	}

	var results []scored
	for _, doc := range docs {
		score := cosineSimilarity(queryVec, doc.Embedding)
		results = append(results, scored{doc, score})
	}

	sort.Slice(results, func(i, j int) bool {
		return results[i].score > results[j].score
	})

	top := make([]Document, 0, topK)
	for i := 0; i < topK && i < len(results); i++ {
		top = append(top, results[i].doc)
	}
	return top
}

func main() {
	texts := []string{
		"Go uses goroutines for lightweight concurrency",
		"Error handling in Go uses the error interface",
		"Channels allow goroutines to communicate safely",
		"The defer keyword schedules cleanup functions",
		"Slices are dynamically-sized views into arrays",
	}

	fmt.Println("Embedding documents...")
	var docs []Document
	for _, t := range texts {
		vec, err := embed(t)
		if err != nil {
			fmt.Println("Error embedding:", err)
			return
		}
		docs = append(docs, Document{Text: t, Embedding: vec})
	}
	fmt.Printf("Indexed %d documents\n\n", len(docs))

	queries := []string{
		"How do goroutines talk to each other?",
		"What happens when a function returns an error?",
		"How do I clean up resources?",
	}

	for _, q := range queries {
		fmt.Printf("Q: %s\n", q)
		results := search(docs, q, 2)
		for _, r := range results {
			fmt.Printf("  → %s\n", r.Text)
		}
		fmt.Println()
	}
	// Q: How do goroutines talk to each other?
	//   → Channels allow goroutines to communicate safely
	//   → Go uses goroutines for lightweight concurrency
	//
	// Q: What happens when a function returns an error?
	//   → Error handling in Go uses the error interface
	//   → Channels allow goroutines to communicate safely
	//
	// Q: How do I clean up resources?
	//   → The defer keyword schedules cleanup functions
	//   → Error handling in Go uses the error interface
}