Phase 4: Embedding Ecosystem

[oss-amikos]

Summary

Phase 4: Embedding Ecosystem
Goal: Expand the embedding ecosystem with sparse/content interfaces, reranking functions, additional dense providers, and an auto-wiring registry.
Status: Verified (UAT 9/9 passed), reviewed (5 rounds of automated review)

This PR adds a complete embedding ecosystem to the Java client: sparse and multimodal embedding interfaces, reranking support, three new dense providers, two sparse providers, and a central registry for auto-wiring providers from server-side collection config.

Changes

Plan 04-01: Embedding Foundation Interfaces

Sparse and multimodal content embedding interfaces with bidirectional adapters.

Key files created:

• SparseEmbeddingFunction — interface for sparse vector providers (BM25, SPLADE)
• ContentEmbeddingFunction — interface for multimodal embedding providers
• Content, Part, BinarySource — value types for multimodal content
• Modality, Intent — enums for content typing and embedding intent
• TextEmbeddingAdapter, ContentToTextAdapter — bidirectional adapters

Plan 04-02: Reranking Interface and Providers

Reranking function interface with Cohere and Jina implementations.

Key files created:

• RerankingFunction — interface for document reranking
• RerankResult — immutable value type for ranked results
• CohereRerankingFunction — Cohere v2/rerank endpoint
• JinaRerankingFunction — Jina v1/rerank endpoint

Plan 04-03: Dense Embedding Providers

Three new dense embedding providers: Gemini, Bedrock, and Voyage.

Key files created:

• GeminiEmbeddingFunction — via Google GenAI SDK (optional dependency)
• BedrockEmbeddingFunction — via AWS SDK BedrockRuntime (optional dependency)
• VoyageEmbeddingFunction — via OkHttp REST calls

Plan 04-04: Sparse Embedding Providers

BM25 local sparse embeddings and ChromaCloud SPLADE remote sparse embeddings.

Key files created:

• BM25EmbeddingFunction — local tokenize → stem → hash → TF-IDF pipeline
• BM25Tokenizer, Murmur3, BM25StopWords — pipeline components
• ChromaCloudSpladeEmbeddingFunction — remote SPLADE via Chroma Cloud API

Plan 04-05: EmbeddingFunctionRegistry

Central registry with three factory maps (dense, sparse, content) and auto-wiring from collection config.

Key files created:

• EmbeddingFunctionRegistry — singleton registry with provider factory maps
• UnsupportedEmbeddingProviderException — typed exception for unknown providers

Key files modified:

• EmbeddingFunctionResolver — now delegates to registry instead of hardcoded dispatch

Requirements Addressed

• EMB-05: Sparse embedding interfaces and providers
• EMB-06: Multimodal content embedding interfaces
• EMB-07: Additional dense embedding providers (Gemini, Bedrock, Voyage)
• EMB-08: Embedding function registry with auto-wiring
• RERANK-01: Reranking interface and providers

Verification

• UAT: 9/9 tests passed
• 5 rounds of automated code review (code quality, error handling, type design, test coverage)
• All 114+ unit tests pass
• Java 8 compatibility maintained
• Thread-safe BM25Tokenizer and EmbeddingFunctionRegistry
• Defensive response DTO parsing with index-context error messages
• Input validation on all public API methods (null query/documents, missing API keys)

Test Plan

• CI passes (build + all tests)
• Verify optional dependencies (Gemini, Bedrock) don't break builds without their SDKs
• Spot-check: BM25 embedding produces deterministic sparse vectors
• Spot-check: Registry resolves built-in providers by name

[tazarov]

@claude PTAL

[claude]

Claude finished @tazarov's task in 2m 6s —— View job

---

PR Review: Phase 4 — Embedding Ecosystem

• Review embedding foundation interfaces (SparseEmbeddingFunction, ContentEmbeddingFunction, adapters, content types)
• Review reranking interface and providers (Cohere, Jina)
• Review dense embedding providers (Gemini, Bedrock, Voyage)
• Review sparse embedding providers (BM25, ChromaCloud SPLADE)
• Review EmbeddingFunctionRegistry and resolver changes
• Review tests
• Review pom.xml dependency changes
• Post final review summary

---

Overall this is a well-structured, comprehensive PR. The interface design is clean, the value types are properly immutable, input validation is thorough with good error messages, and the test coverage is solid with WireMock-based HTTP testing. Here are my findings:

Issues

1. ContentToTextAdapter.embedDocuments(String[]) — null delegates to null list (minor bug)

ContentToTextAdapter.java:64-66: When documents is String[] null, the method calls embedDocuments((List<String>) null), which will throw IllegalArgumentException("documents must not be null"). This works, but the behavior is inconsistent — passing null to the array overload gives a different exception type (IllegalArgumentException) than what the validation in embedDocuments(List) would produce if it were reached directly. This pattern is repeated in GeminiEmbeddingFunction, BedrockEmbeddingFunction, and VoyageEmbeddingFunction, where the list overload throws ChromaException but the null array path also throws IllegalArgumentException. Not a functional bug since null is rejected either way, but the exception type is inconsistent.

2. Cohere/Jina reranking implementations are nearly identical (design)

CohereRerankingFunction.java and JinaRerankingFunction.java are ~95% identical — same structure, same error handling, same sorting. The only differences are the provider name in error messages, the default URL, and model name. Consider extracting a shared AbstractRerankingFunction base class to reduce duplication (~100 lines of shared code).

Extract shared reranking base class →

3. Similarly, RerankRequest and RerankResponse are duplicated

tech.amikos.chromadb.reranking.cohere.RerankRequest and tech.amikos.chromadb.reranking.jina.RerankRequest are byte-for-byte identical (same for RerankResponse). These could be shared in the parent reranking package.

4. GeminiEmbeddingFunction embeds documents one-at-a-time (performance)

GeminiEmbeddingFunction.java:115-121: embedDocuments loops and calls client.models.embedContent() once per document, making N HTTP calls for N documents. If the Gemini SDK supports batch embedding, this would be a significant performance improvement. Worth checking if the SDK has a batch endpoint.

5. BedrockEmbeddingFunction also embeds one-at-a-time

BedrockEmbeddingFunction.java:144-159: Same per-document loop pattern. Titan Embed v2 supports batch via the texts field (list of strings). Using this would reduce API calls from N to 1.

6. BM25Tokenizer creates a new stemmer per tokenize() call

BM25Tokenizer.java:63: A new englishStemmer is allocated on every call to tokenize(). Since the class documents itself as thread-safe, making the stemmer a ThreadLocal<englishStemmer> would avoid per-call allocation while maintaining thread safety.

7. Snowball stemmer pulled in as a compile-scope dependency

pom.xml:155-158: com.github.rholder:snowball-stemmer is added as a compile dependency. Since BM25 is a sparse embedding feature that not all users need, consider making this <optional>true</optional> like Gemini/Bedrock, or at minimum document this in the PR description as a new required transitive dependency.

8. CreateSparseEmbeddingRequest.json() method is misleading

CreateSparseEmbeddingRequest.java:29-31: The json() method returns a MediaType, not JSON. This is confusing — it's never actually used (the class has a separate toJson() method that returns the JSON string). This looks like dead code left over from copy-paste from the Voyage request DTO pattern.

Remove dead json() method →

9. New Gson() allocations in request DTOs

RerankRequest.json() (both Cohere and Jina), CreateSparseEmbeddingRequest.toJson(), and CreateEmbeddingRequest.json() all do new Gson().toJson(this). Gson instances are thread-safe and reusable — a static field would avoid repeated allocation.

10. OkHttpClient instances are not shared

Each provider class (VoyageEmbeddingFunction, CohereRerankingFunction, JinaRerankingFunction, ChromaCloudSpladeEmbeddingFunction) creates its own new OkHttpClient(). OkHttp recommends sharing a single client to share the connection pool and thread pools. Since users may create multiple provider instances, this could lead to resource waste. Consider accepting an OkHttpClient in the constructor or using a shared static instance.

Minor / Nit

• BM25EmbeddingFunction.java:28: Charset.forName("UTF-8") with comment about Java 8 compatibility — StandardCharsets.UTF_8 is available since Java 7, so the comment is misleading. Both work fine on Java 8.
• Modality.java:28: After calling toLowerCase(Locale.ROOT), the regex [a-zA-Z0-9] uppercase check is redundant since input is already lowered. Very minor.
• EmbeddingFunctionRegistry.java:102: The LinkedHashMap diamond operator <> is used in a few places — since this targets Java 8, that's fine, but some earlier code in this same PR uses explicit type parameters (e.g., new LinkedHashMap<String, DenseFactory>()). Minor inconsistency.

What Looks Good

• Interface design: SparseEmbeddingFunction, ContentEmbeddingFunction, and RerankingFunction are clean, minimal interfaces. The bidirectional adapter pattern (TextEmbeddingAdapter / ContentToTextAdapter) is well thought out.
• Content model: Content, Part, BinarySource are properly immutable with defensive copies. Builder validation is solid.
• Registry pattern: EmbeddingFunctionRegistry with three factory maps, case-insensitive name resolution, and graceful handling of optional SDK dependencies via Class.forName checks is well designed.
• Error messages: Consistently include model name and index context, making debugging straightforward.
• Test coverage: Good use of WireMock for HTTP-based providers, thread safety tests for the registry, and thorough edge case coverage (null inputs, empty inputs, malformed responses).
• BM25 implementation: Correct Murmur3, proper cross-client compatibility with Go/Python, sorted sparse vector output.

Summary

The PR is in good shape architecturally. The main actionable items are:

1. Remove the dead CreateSparseEmbeddingRequest.json() method (clear bug/leftover)
2. Consider making snowball-stemmer optional in pom.xml to avoid forcing a new transitive dep on all users
3. Extract shared reranking base to reduce code duplication (optional but recommended)

Everything else is minor or a future optimization opportunity. The interfaces, value types, and registry are well designed.