Embedding Model Implementation

public abstract class GeminiEmbeddingModel implements EmbeddingModel {
    /**
     * Creates a Gemini embedding model with the specified configuration.
     *
     * @param modelId The Gemini embedding model identifier (e.g., "text-embedding-004")
     * @param dimension The output embedding dimension (e.g., 768). Null uses model default
     * @param taskType The task type for task-specific embeddings (e.g., "RETRIEVAL_DOCUMENT")
     */
    public GeminiEmbeddingModel(String modelId, Integer dimension, String taskType);

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

    /**
     * Generates embeddings for multiple text segments in a batch.
     *
     * @param textSegments List of text segments to embed
     * @return Response containing list of embedding vectors
     */
    @Override
    public Response<List<Embedding>> embedAll(List<TextSegment> textSegments);

    /**
     * Subclass-specific implementation for single embedding generation.
     * Must be implemented by concrete subclasses to call the Gemini API.
     *
     * @param request The embedding request with text and configuration
     * @return The embedding response from the API
     */
    protected abstract EmbedContentResponse embedContent(EmbedContentRequest request);

    /**
     * Subclass-specific implementation for batch embedding generation.
     * Must be implemented by concrete subclasses to call the batch Gemini API.
     *
     * @param requests The batch embedding requests
     * @return The batch embedding response from the API
     */
    protected abstract EmbedContentResponses batchEmbedContents(EmbedContentRequests requests);
}

public class MyGeminiEmbeddingModel extends GeminiEmbeddingModel {

    private final GeminiRestApi restApi;

    public MyGeminiEmbeddingModel(String apiKey) {
        super(
            "text-embedding-004",  // modelId
            768,                   // dimension
            "RETRIEVAL_DOCUMENT"   // taskType
        );
        this.restApi = createRestApi(apiKey);
    }

    @Override
    protected EmbedContentResponse embedContent(EmbedContentRequest request) {
        return restApi.embedContent(request);
    }

    @Override
    protected EmbedContentResponses batchEmbedContents(EmbedContentRequests requests) {
        return restApi.batchEmbedContents(requests);
    }
}

// Use the model
MyGeminiEmbeddingModel model = new MyGeminiEmbeddingModel(apiKey);
Response<Embedding> response = model.embed("Hello, world!");

float[] vector = response.content().vector();
System.out.println("Embedding dimension: " + vector.length);

MyGeminiEmbeddingModel model = new MyGeminiEmbeddingModel(apiKey);

List<TextSegment> segments = List.of(
    TextSegment.from("First document about machine learning"),
    TextSegment.from("Second document about natural language processing"),
    TextSegment.from("Third document about computer vision")
);

Response<List<Embedding>> response = model.embedAll(segments);

for (int i = 0; i < response.content().size(); i++) {
    Embedding embedding = response.content().get(i);
    System.out.println("Document " + i + " embedding: " + embedding.vector().length + " dimensions");
}

// For search queries
public class QueryEmbeddingModel extends GeminiEmbeddingModel {
    public QueryEmbeddingModel(String apiKey) {
        super(
            "text-embedding-004",
            768,
            "RETRIEVAL_QUERY"      // Optimized for queries
        );
        // ...
    }
}

// For documents
public class DocumentEmbeddingModel extends GeminiEmbeddingModel {
    public DocumentEmbeddingModel(String apiKey) {
        super(
            "text-embedding-004",
            768,
            "RETRIEVAL_DOCUMENT"   // Optimized for documents
        );
        // ...
    }
}

// Use different models for queries vs documents
QueryEmbeddingModel queryModel = new QueryEmbeddingModel(apiKey);
DocumentEmbeddingModel docModel = new DocumentEmbeddingModel(apiKey);

Response<Embedding> queryEmbedding = queryModel.embed("machine learning algorithms");
Response<Embedding> docEmbedding = docModel.embed("A comprehensive guide to ML algorithms...");

// Compare embeddings for semantic search
double similarity = cosineSimilarity(
    queryEmbedding.content().vector(),
    docEmbedding.content().vector()
);

// Use custom dimension (e.g., 256 instead of default 768)
public class CompactEmbeddingModel extends GeminiEmbeddingModel {
    public CompactEmbeddingModel(String apiKey) {
        super(
            "text-embedding-004",
            256,                   // Smaller dimension for efficiency
            "SEMANTIC_SIMILARITY"
        );
        // ...
    }
}

CompactEmbeddingModel model = new CompactEmbeddingModel(apiKey);
Response<Embedding> response = model.embed("Compact embedding");
// Output will be 256-dimensional instead of 768