Understanding Embeddings, Vector Search, and RAG with Spring AI

Cosine Similarity = (A · B) / (||A|| × ||B||)

Where:
  A · B   = dot product of vectors
  ||A||   = magnitude of vector A
  ||B||   = magnitude of vector B

Result: Value between -1 and 1
  1   = identical direction (very similar)
  0   = orthogonal (unrelated)
  -1  = opposite direction (opposite meaning)

Query vector:     [0.8, 0.6, 0.0]
Document vector:  [0.9, 0.5, 0.1]

Dot product:  (0.8 × 0.9) + (0.6 × 0.5) + (0.0 × 0.1) = 1.02
Magnitude A:  √(0.8² + 0.6²) = 1.0
Magnitude B:  √(0.9² + 0.5² + 0.1²) = 1.03

Cosine Similarity = 1.02 / (1.0 × 1.03) = 0.99 (99% similar!)

-- Create database and enable extension
CREATE DATABASE myvectordb;
\c myvectordb
CREATE EXTENSION vector;
CREATE SCHEMA embeddings;

CREATE TABLE embeddings.vector_store (
    id UUID PRIMARY KEY,
    content TEXT,
    metadata JSONB,
    embedding VECTOR(1536)
);

-- Create HNSW index for fast similarity search
CREATE INDEX ON embeddings.vector_store
USING hnsw (embedding vector_cosine_ops);

# Database connection
spring.datasource.url=jdbc:postgresql://localhost:5432/myvectordb
spring.datasource.username=myuser
spring.datasource.password=mypassword

# PgVector configuration
spring.ai.vectorstore.pgvector.index-type=HNSW
spring.ai.vectorstore.pgvector.distance-type=COSINE_DISTANCE
spring.ai.vectorstore.pgvector.dimensions=1536

# OpenAI embeddings
spring.ai.openai.api-key=${OPENAI_API_KEY}
spring.ai.openai.embedding.options.model=text-embedding-3-small

@Service
public class EmbeddingService {

    private static final Logger logger = LoggerFactory.getLogger(EmbeddingService.class);
    private final VectorStore vectorStore;

    public EmbeddingService(VectorStore vectorStore) {
        this.vectorStore = vectorStore;
    }

    /**
     * Store an article with its embedding
     * Spring AI automatically generates the embedding via OpenAI
     */
    public void storeArticle(Article article) {
        // Prepare metadata
        Map<String, Object> metadata = new HashMap<>();
        metadata.put("title", article.getTitle());
        metadata.put("category", article.getCategory());
        metadata.put("url", article.getUrl());

        // Truncate content if needed (OpenAI has 8191 token limit)
        String content = article.getContent();
        if (content.length() > 8000) {
            content = content.substring(0, 8000);
        }

        // Create document - Spring AI handles embedding generation
        Document document = new Document(article.getId(), content, metadata);

        // Store in vector database
        vectorStore.add(List.of(document));
        logger.info("Stored article: {}", article.getTitle());
    }

    /**
     * Perform semantic search
     * Returns articles ranked by similarity to the query
     */
    public SearchResult semanticSearch(String query, int topK) {
        long startTime = System.currentTimeMillis();

        // Build search request
        SearchRequest searchRequest = SearchRequest.builder()
            .query(query)           // User's search query
            .topK(topK)             // Number of results to return
            .build();

        // Execute similarity search
        // Spring AI:
        //   1. Generates embedding for query via OpenAI
        //   2. Executes vector similarity search in PgVector
        //   3. Returns ranked results
        List<Document> documents = vectorStore.similaritySearch(searchRequest);

        // Convert to domain objects
        List<Article> results = documents.stream()
            .map(this::documentToArticle)
            .collect(Collectors.toList());

        long executionTime = System.currentTimeMillis() - startTime;

        logger.info("Search for '{}' found {} results in {}ms",
                    query, results.size(), executionTime);

        return new SearchResult(query, results, executionTime);
    }

    /**
     * Search with metadata filtering
     * Example: Find articles in "AI" category similar to query
     */
    public SearchResult searchWithFilter(String query, String category, int topK) {
        SearchRequest.Builder builder = SearchRequest.builder()
            .query(query)
            .topK(topK);

        // Add metadata filter
        if (category != null) {
            String filter = String.format("category == '%s'", category);
            builder.filterExpression(filter);
        }

        List<Document> documents = vectorStore.similaritySearch(builder.build());
        List<Article> results = documents.stream()
            .map(this::documentToArticle)
            .collect(Collectors.toList());

        return new SearchResult(query, results, 0);
    }

    /**
     * Find similar articles to a given article
     * Uses the article's content as the query
     */
    public List<Article> findSimilar(String articleId, int topK) {
        // First, retrieve the target article
        SearchRequest getAllRequest = SearchRequest.builder()
            .query("article")
            .topK(100)
            .build();

        List<Document> allDocs = vectorStore.similaritySearch(getAllRequest);

        Document targetDoc = allDocs.stream()
            .filter(doc -> doc.getId().equals(articleId))
            .findFirst()
            .orElseThrow(() -> new NotFoundException("Article not found"));

        // Use article content as query to find similar ones
        SearchRequest similarRequest = SearchRequest.builder()
            .query(targetDoc.getText())
            .topK(topK + 1)  // +1 to exclude the query article
            .build();

        List<Document> similar = vectorStore.similaritySearch(similarRequest);

        // Exclude the query article itself and return results
        return similar.stream()
            .filter(doc -> !doc.getId().equals(articleId))
            .limit(topK)
            .map(this::documentToArticle)
            .collect(Collectors.toList());
    }

    /**
     * Convert Spring AI Document to domain Article
     * Extract similarity score from metadata
     */
    private Article documentToArticle(Document doc) {
        Article article = new Article();
        article.setId(doc.getId());
        article.setContent(doc.getText());

        Map<String, Object> metadata = doc.getMetadata();
        article.setTitle((String) metadata.get("title"));
        article.setCategory((String) metadata.get("category"));
        article.setUrl((String) metadata.get("url"));

        // Extract similarity score (PgVector returns distance)
        if (metadata.containsKey("distance")) {
            double distance = ((Number) metadata.get("distance")).doubleValue();
            // Convert distance to similarity score (1 = identical, 0 = unrelated)
            article.setScore(1.0 - distance);
        }

        return article;
    }
}

@Service
public class WikipediaService {

    private final EmbeddingService embeddingService;

    public void indexArticle(String title, String category) {
        // Fetch article from Wikipedia API
        String content = fetchFromWikipedia(title);

        // Create article object
        Article article = new Article();
        article.setId(UUID.randomUUID().toString());
        article.setTitle(title);
        article.setCategory(category);
        article.setContent(content);
        article.setUrl("https://en.wikipedia.org/wiki/" + title);
        article.setCreatedAt(LocalDateTime.now());

        // Store with embedding generation
        embeddingService.storeArticle(article);
    }

    public void loadAIArticles() {
        List<String> articles = List.of(
            "Machine learning",
            "Supervised learning",
            "Unsupervised learning",
            "Neural network",
            "Deep learning",
            "Computer vision",
            "Natural language processing"
        );

        articles.forEach(title -> {
            indexArticle(title, "AI");
            // Add delay to respect API rate limits
            Thread.sleep(500);
        });
    }
}

@RestController
@RequestMapping("/api/search")
public class SearchController {

    private final EmbeddingService embeddingService;

    @GetMapping
    public SearchResult search(
            @RequestParam String q,
            @RequestParam(defaultValue = "10") int limit) {

        return embeddingService.semanticSearch(q, limit);
    }

    @GetMapping("/similar/{articleId}")
    public List<Article> findSimilar(
            @PathVariable String articleId,
            @RequestParam(defaultValue = "5") int limit) {

        return embeddingService.findSimilar(articleId, limit);
    }

    @GetMapping("/recommend")
    public List<Article> recommend(
            @RequestParam String q,
            @RequestParam(defaultValue = "5") int limit) {

        return embeddingService.getRecommendations(q, limit);
    }
}

# Index articles
curl -X POST http://localhost:8080/api/data/load

# Semantic search examples
curl "http://localhost:8080/api/search?q=learning+from+examples&limit=5"
# Returns: Supervised Learning (0.89), Machine Learning (0.76), ...

curl "http://localhost:8080/api/search?q=understanding+images"
# Returns: Computer Vision (0.91), CNN (0.85), ...

curl "http://localhost:8080/api/search?q=text+understanding"
# Returns: NLP (0.93), Transformers (0.87), ...

# Find similar articles
curl "http://localhost:8080/api/search/similar/abc-123?limit=5"

@Service
public class RagService {

    private static final Logger logger = LoggerFactory.getLogger(RagService.class);

    private final VectorStore vectorStore;
    private final ChatClient chatClient;

    private static final String SYSTEM_PROMPT_TEMPLATE = """
        You are a helpful AI assistant that answers questions based on the provided context.
        Use the following pieces of context to answer the user's question.
        If you cannot find the answer in the context, say so - do not make up information.

        Context:
        {context}

        Instructions:
        - Answer based only on the provided context
        - Be concise and accurate
        - If the context doesn't contain relevant information, say
          "I don't have enough information to answer that question."
        - Cite which documents you used if relevant
        """;

    public RagService(VectorStore vectorStore, ChatClient.Builder chatClientBuilder) {
        this.vectorStore = vectorStore;
        this.chatClient = chatClientBuilder.build();
    }

    /**
     * Ask a question using RAG
     * Returns an answer generated from retrieved documents
     */
    public RagResponse ask(String question, int topK) {
        long startTime = System.currentTimeMillis();

        logger.info("RAG request: '{}'", question);

        // Step 1: Retrieve relevant documents using semantic search
        SearchRequest searchRequest = SearchRequest.builder()
            .query(question)
            .topK(topK)
            .build();

        List<Document> relevantDocs = vectorStore.similaritySearch(searchRequest);

        if (relevantDocs.isEmpty()) {
            return new RagResponse(
                question,
                "I don't have any relevant documents to answer your question.",
                List.of(),
                0,
                System.currentTimeMillis() - startTime
            );
        }

        // Step 2: Build context from retrieved documents
        String context = buildContext(relevantDocs);

        logger.debug("Retrieved {} documents, context length: {}",
                    relevantDocs.size(), context.length());

        // Step 3: Generate answer using chat model with context
        SystemPromptTemplate systemPromptTemplate =
            new SystemPromptTemplate(SYSTEM_PROMPT_TEMPLATE);
        Message systemMessage = systemPromptTemplate.createMessage(
            Map.of("context", context)
        );
        UserMessage userMessage = new UserMessage(question);

        Prompt prompt = new Prompt(List.of(systemMessage, userMessage));

        // Call LLM
        long generationStart = System.currentTimeMillis();
        String answer = chatClient.prompt(prompt).call().content();
        long generationTime = System.currentTimeMillis() - generationStart;

        long totalTime = System.currentTimeMillis() - startTime;

        logger.info("RAG completed in {}ms (retrieval: {}ms, generation: {}ms)",
                   totalTime, generationStart - startTime, generationTime);

        // Extract source references
        List<DocumentReference> sources = relevantDocs.stream()
            .map(doc -> new DocumentReference(
                (String) doc.getMetadata().get("title"),
                (String) doc.getMetadata().get("category"),
                doc.getId()
            ))
            .collect(Collectors.toList());

        return new RagResponse(question, answer, sources, relevantDocs.size(), totalTime);
    }

    /**
     * Ask a question with category filtering
     * Example: Only search in "AI" category documents
     */
    public RagResponse askWithFilter(String question, String category, int topK) {
        SearchRequest.Builder builder = SearchRequest.builder()
            .query(question)
            .topK(topK);

        if (category != null && !category.trim().isEmpty()) {
            String filter = String.format("category == '%s'", category);
            builder.filterExpression(filter);
        }

        SearchRequest searchRequest = builder.build();
        List<Document> relevantDocs = vectorStore.similaritySearch(searchRequest);

        // Generate answer same as above...
        String context = buildContext(relevantDocs);

        SystemPromptTemplate systemPromptTemplate =
            new SystemPromptTemplate(SYSTEM_PROMPT_TEMPLATE);
        Message systemMessage = systemPromptTemplate.createMessage(
            Map.of("context", context)
        );
        UserMessage userMessage = new UserMessage(question);
        Prompt prompt = new Prompt(List.of(systemMessage, userMessage));

        String answer = chatClient.prompt(prompt).call().content();

        List<DocumentReference> sources = relevantDocs.stream()
            .map(doc -> new DocumentReference(
                (String) doc.getMetadata().get("title"),
                (String) doc.getMetadata().get("category"),
                doc.getId()
            ))
            .collect(Collectors.toList());

        return new RagResponse(question, answer, sources,
                             relevantDocs.size(), 0);
    }

    /**
     * Build context string from multiple documents
     * Formats documents for LLM consumption
     */
    private String buildContext(List<Document> documents) {
        StringBuilder context = new StringBuilder();

        for (int i = 0; i < documents.size(); i++) {
            Document doc = documents.get(i);
            String title = (String) doc.getMetadata().getOrDefault("title", "Untitled");
            String content = doc.getText();

            // Limit content length to avoid token limits
            if (content.length() > 1500) {
                content = content.substring(0, 1500) + "...";
            }

            context.append(String.format("\n[Document %d: %s]\n%s\n",
                                       i + 1, title, content));
        }

        return context.toString();
    }

    // Response classes
    public record RagResponse(
        String question,
        String answer,
        List<DocumentReference> sources,
        int documentCount,
        long executionTimeMs
    ) {}

    public record DocumentReference(
        String title,
        String category,
        String id
    ) {}
}

Model: text-embedding-3-small
Price: $0.02 per 1 million tokens

Indexing 1,000 articles:
  Average 500 words/article = 500,000 words
  ~666,666 tokens
  Cost: $0.013

Search queries (1,000):
  Average 10 words/query = 10,000 words
  ~13,000 tokens
  Cost: $0.0003

Very cost-effective!

Model: gpt-3.5-turbo
Price: $0.50 per 1M input tokens, $1.50 per 1M output tokens

RAG query with 3 documents:
  Input:  ~2,000 tokens (context + question)
  Output: ~200 tokens (answer)
  Cost:   $0.001 + $0.0003 = $0.0013 per query

1,000 RAG queries = ~$1.30