Embeddings

/**
 * Text embedding model for Google AI Gemini.
 * Converts text into vector representations for semantic understanding.
 */
public class GoogleAiEmbeddingModel {
    /**
     * Creates a new builder for configuring the embedding model.
     * @return GoogleAiEmbeddingModelBuilder instance
     */
    public static GoogleAiEmbeddingModelBuilder builder();

    /**
     * Embeds a single text segment.
     * @param textSegment The text segment to embed
     * @return Response containing the embedding vector
     */
    public Response<Embedding> embed(TextSegment textSegment);

    /**
     * Embeds a single string.
     * @param text The text string to embed
     * @return Response containing the embedding vector
     */
    public Response<Embedding> embed(String text);

    /**
     * Embeds multiple text segments in a single request.
     * @param textSegments List of text segments to embed
     * @return Response containing list of embedding vectors
     */
    public Response<List<Embedding>> embedAll(List<TextSegment> textSegments);

    /**
     * Returns the model name being used.
     * @return Model name string
     */
    public String modelName();

    /**
     * Returns the known dimension of embeddings produced by this model.
     * @return Embedding dimension (null if not known)
     */
    public Integer knownDimension();
}

/**
 * Builder for GoogleAiEmbeddingModel.
 * Extends base builder with embedding-specific configuration.
 */
public static class GoogleAiEmbeddingModelBuilder
    extends BaseGoogleAiEmbeddingModelBuilder<GoogleAiEmbeddingModelBuilder> {

    /**
     * Builds the GoogleAiEmbeddingModel instance.
     * @return Configured GoogleAiEmbeddingModel
     * @throws IllegalArgumentException if required fields are missing
     */
    public GoogleAiEmbeddingModel build();
}

/**
 * Base builder class providing common configuration for embedding models.
 * @param <B> The builder type for fluent chaining
 */
public abstract class BaseGoogleAiEmbeddingModelBuilder<B> {
    /**
     * Sets the HTTP client builder for customizing requests.
     * @param httpClientBuilder HTTP client builder instance
     * @return Builder instance for chaining
     */
    public B httpClientBuilder(HttpClientBuilder httpClientBuilder);

    /**
     * Sets the model name to use for embeddings.
     * @param modelName Model identifier (e.g., "text-embedding-004")
     * @return Builder instance for chaining
     */
    public B modelName(String modelName);

    /**
     * Sets the API key for authentication (required).
     * @param apiKey Google AI API key
     * @return Builder instance for chaining
     */
    public B apiKey(String apiKey);

    /**
     * Sets the base URL for the API endpoint.
     * @param baseUrl Custom base URL (optional)
     * @return Builder instance for chaining
     */
    public B baseUrl(String baseUrl);

    /**
     * Sets the maximum number of retry attempts for failed requests.
     * @param maxRetries Number of retries (default: 3)
     * @return Builder instance for chaining
     */
    public B maxRetries(Integer maxRetries);

    /**
     * Sets the task type for optimized embeddings.
     * @param taskType Task type enum value
     * @return Builder instance for chaining
     */
    public B taskType(TaskType taskType);

    /**
     * Sets the metadata key for retrieving document titles.
     * @param titleMetadataKey Metadata key string
     * @return Builder instance for chaining
     */
    public B titleMetadataKey(String titleMetadataKey);

    /**
     * Sets the output dimensionality for embeddings.
     * @param outputDimensionality Desired embedding dimension
     * @return Builder instance for chaining
     */
    public B outputDimensionality(Integer outputDimensionality);

    /**
     * Sets the request timeout duration.
     * @param timeout Timeout duration
     * @return Builder instance for chaining
     */
    public B timeout(Duration timeout);

    /**
     * Enables logging of both requests and responses.
     * @param logRequestsAndResponses True to enable full logging
     * @return Builder instance for chaining
     */
    public B logRequestsAndResponses(Boolean logRequestsAndResponses);

    /**
     * Enables logging of requests only.
     * @param logRequests True to enable request logging
     * @return Builder instance for chaining
     */
    public B logRequests(Boolean logRequests);

    /**
     * Enables logging of responses only.
     * @param logResponses True to enable response logging
     * @return Builder instance for chaining
     */
    public B logResponses(Boolean logResponses);

    /**
     * Sets a custom logger for the model.
     * @param logger Logger instance
     * @return Builder instance for chaining
     */
    public B logger(Logger logger);
}

/**
 * Task types for optimizing embedding generation.
 */
public enum TaskType {
    /**
     * Optimize for search queries (query side of retrieval).
     */
    RETRIEVAL_QUERY,

    /**
     * Optimize for documents to be retrieved (document side of retrieval).
     */
    RETRIEVAL_DOCUMENT,

    /**
     * Optimize for measuring semantic similarity between texts.
     */
    SEMANTIC_SIMILARITY,

    /**
     * Optimize for text classification tasks.
     */
    CLASSIFICATION,

    /**
     * Optimize for clustering texts into groups.
     */
    CLUSTERING,

    /**
     * Optimize for question answering tasks.
     */
    QUESTION_ANSWERING,

    /**
     * Optimize for fact verification tasks.
     */
    FACT_VERIFICATION
}

import dev.langchain4j.model.googleai.GoogleAiEmbeddingModel;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.model.output.Response;

// Create embedding model
EmbeddingModel embeddingModel = GoogleAiEmbeddingModel.builder()
    .apiKey(System.getenv("GOOGLE_AI_API_KEY"))
    .modelName("text-embedding-004")
    .build();

// Embed a single text
Response<Embedding> response = embeddingModel.embed("Artificial intelligence is transforming technology");
Embedding embedding = response.content();

// Access vector
float[] vector = embedding.vector();
System.out.println("Embedding dimension: " + vector.length);

import dev.langchain4j.data.segment.TextSegment;
import java.util.List;

GoogleAiEmbeddingModel embeddingModel = GoogleAiEmbeddingModel.builder()
    .apiKey(System.getenv("GOOGLE_AI_API_KEY"))
    .modelName("text-embedding-004")
    .build();

// Embed multiple texts at once
List<TextSegment> segments = List.of(
    TextSegment.from("Machine learning enables computers to learn from data"),
    TextSegment.from("Deep learning uses neural networks with many layers"),
    TextSegment.from("Natural language processing helps computers understand text")
);

Response<List<Embedding>> response = embeddingModel.embedAll(segments);
List<Embedding> embeddings = response.content();

// Process each embedding
for (int i = 0; i < embeddings.size(); i++) {
    float[] vector = embeddings.get(i).vector();
    System.out.println("Text " + (i + 1) + " embedding dimension: " + vector.length);
}

import dev.langchain4j.model.googleai.GoogleAiEmbeddingModel.TaskType;

// Query embedding for search
GoogleAiEmbeddingModel queryModel = GoogleAiEmbeddingModel.builder()
    .apiKey(System.getenv("GOOGLE_AI_API_KEY"))
    .modelName("text-embedding-004")
    .taskType(TaskType.RETRIEVAL_QUERY)
    .build();

Response<Embedding> queryEmbedding = queryModel.embed("What is quantum computing?");

// Document embedding for indexing
GoogleAiEmbeddingModel documentModel = GoogleAiEmbeddingModel.builder()
    .apiKey(System.getenv("GOOGLE_AI_API_KEY"))
    .modelName("text-embedding-004")
    .taskType(TaskType.RETRIEVAL_DOCUMENT)
    .build();

List<TextSegment> documents = List.of(
    TextSegment.from("Quantum computing uses quantum bits or qubits..."),
    TextSegment.from("Classical computers use binary bits for computation..."),
    TextSegment.from("Superposition allows qubits to be in multiple states...")
);

Response<List<Embedding>> docEmbeddings = documentModel.embedAll(documents);

GoogleAiEmbeddingModel similarityModel = GoogleAiEmbeddingModel.builder()
    .apiKey(System.getenv("GOOGLE_AI_API_KEY"))
    .modelName("text-embedding-004")
    .taskType(TaskType.SEMANTIC_SIMILARITY)
    .build();

// Embed two texts
Response<Embedding> embedding1 = similarityModel.embed("I love programming");
Response<Embedding> embedding2 = similarityModel.embed("I enjoy coding");

// Calculate cosine similarity
float[] vec1 = embedding1.content().vector();
float[] vec2 = embedding2.content().vector();

float dotProduct = 0.0f;
float norm1 = 0.0f;
float norm2 = 0.0f;

for (int i = 0; i < vec1.length; i++) {
    dotProduct += vec1[i] * vec2[i];
    norm1 += vec1[i] * vec1[i];
    norm2 += vec2[i] * vec2[i];
}

float similarity = dotProduct / (float)(Math.sqrt(norm1) * Math.sqrt(norm2));
System.out.println("Similarity: " + similarity); // High similarity expected

GoogleAiEmbeddingModel classificationModel = GoogleAiEmbeddingModel.builder()
    .apiKey(System.getenv("GOOGLE_AI_API_KEY"))
    .modelName("text-embedding-004")
    .taskType(TaskType.CLASSIFICATION)
    .build();

// Embed texts for classification
List<TextSegment> textsToClassify = List.of(
    TextSegment.from("This movie was absolutely amazing!"),
    TextSegment.from("Terrible experience, would not recommend"),
    TextSegment.from("Pretty good, but could be better")
);

Response<List<Embedding>> embeddings = classificationModel.embedAll(textsToClassify);
// Use embeddings as features for classifier

GoogleAiEmbeddingModel clusteringModel = GoogleAiEmbeddingModel.builder()
    .apiKey(System.getenv("GOOGLE_AI_API_KEY"))
    .modelName("text-embedding-004")
    .taskType(TaskType.CLUSTERING)
    .build();

List<TextSegment> textsToCluster = List.of(
    TextSegment.from("Apple releases new iPhone"),
    TextSegment.from("Google announces Android update"),
    TextSegment.from("Climate change impacts global weather"),
    TextSegment.from("Samsung unveils latest smartphone"),
    TextSegment.from("Rising temperatures affect ecosystems")
);

Response<List<Embedding>> embeddings = clusteringModel.embedAll(textsToCluster);
// Use embeddings for clustering algorithm

// Reduce embedding dimension for faster processing/storage
GoogleAiEmbeddingModel compactModel = GoogleAiEmbeddingModel.builder()
    .apiKey(System.getenv("GOOGLE_AI_API_KEY"))
    .modelName("text-embedding-004")
    .outputDimensionality(256) // Reduce from default 768
    .build();

Response<Embedding> response = compactModel.embed("Sample text for compact embedding");
float[] vector = response.content().vector();
System.out.println("Compact dimension: " + vector.length); // 256

import dev.langchain4j.data.segment.TextSegment;
import java.util.HashMap;
import java.util.Map;

GoogleAiEmbeddingModel embeddingModel = GoogleAiEmbeddingModel.builder()
    .apiKey(System.getenv("GOOGLE_AI_API_KEY"))
    .modelName("text-embedding-004")
    .taskType(TaskType.RETRIEVAL_DOCUMENT)
    .titleMetadataKey("title") // Use "title" metadata for context
    .build();

// Create segments with title metadata
Map<String, Object> metadata1 = new HashMap<>();
metadata1.put("title", "Introduction to AI");
TextSegment segment1 = TextSegment.from(
    "Artificial intelligence is the simulation of human intelligence...",
    metadata1
);

Map<String, Object> metadata2 = new HashMap<>();
metadata2.put("title", "Machine Learning Basics");
TextSegment segment2 = TextSegment.from(
    "Machine learning is a subset of AI that learns from data...",
    metadata2
);

// Embeddings will use titles for better context
Response<List<Embedding>> response = embeddingModel.embedAll(List.of(segment1, segment2));

import java.time.Duration;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

Logger logger = LoggerFactory.getLogger("EmbeddingModel");

GoogleAiEmbeddingModel embeddingModel = GoogleAiEmbeddingModel.builder()
    .apiKey(System.getenv("GOOGLE_AI_API_KEY"))
    .modelName("text-embedding-004")
    .taskType(TaskType.SEMANTIC_SIMILARITY)
    .outputDimensionality(512)
    .timeout(Duration.ofSeconds(30))
    .maxRetries(5)
    .logRequestsAndResponses(true)
    .logger(logger)
    .build();

// Model configured with custom settings
Response<Embedding> response = embeddingModel.embed("Test embedding with advanced config");

import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;
import dev.langchain4j.data.document.Document;

// Create embedding model
GoogleAiEmbeddingModel embeddingModel = GoogleAiEmbeddingModel.builder()
    .apiKey(System.getenv("GOOGLE_AI_API_KEY"))
    .modelName("text-embedding-004")
    .taskType(TaskType.RETRIEVAL_DOCUMENT)
    .build();

// Create embedding store
EmbeddingStore<TextSegment> embeddingStore = new InMemoryEmbeddingStore<>();

// Index documents
List<TextSegment> documents = List.of(
    TextSegment.from("The sky is blue during the day"),
    TextSegment.from("Grass is green in spring"),
    TextSegment.from("The ocean contains salt water")
);

Response<List<Embedding>> docEmbeddings = embeddingModel.embedAll(documents);
for (int i = 0; i < documents.size(); i++) {
    embeddingStore.add(docEmbeddings.content().get(i), documents.get(i));
}

// Search with query
GoogleAiEmbeddingModel queryModel = GoogleAiEmbeddingModel.builder()
    .apiKey(System.getenv("GOOGLE_AI_API_KEY"))
    .modelName("text-embedding-004")
    .taskType(TaskType.RETRIEVAL_QUERY)
    .build();

Response<Embedding> queryEmbedding = queryModel.embed("What color is the sky?");
List<EmbeddingMatch<TextSegment>> matches = embeddingStore.findRelevant(
    queryEmbedding.content(),
    3
);

matches.forEach(match ->
    System.out.println("Match: " + match.embedded().text() + " (score: " + match.score() + ")")
);

import dev.langchain4j.data.document.loader.FileSystemDocumentLoader;
import dev.langchain4j.data.document.parser.TextDocumentParser;
import dev.langchain4j.data.document.splitter.DocumentSplitters;
import dev.langchain4j.model.embedding.EmbeddingModel;

EmbeddingModel embeddingModel = GoogleAiEmbeddingModel.builder()
    .apiKey(System.getenv("GOOGLE_AI_API_KEY"))
    .modelName("text-embedding-004")
    .taskType(TaskType.RETRIEVAL_DOCUMENT)
    .build();

// Load and split documents
Document document = FileSystemDocumentLoader.loadDocument("/path/to/document.txt");
List<TextSegment> segments = DocumentSplitters.recursive(500, 50).split(document);

// Embed all segments
Response<List<Embedding>> embeddings = embeddingModel.embedAll(segments);