Integrate composite embedder

This commit is contained in:
Louis Dureuil 2025-02-20 11:37:27 +01:00
parent 4a2643daa2
commit 294cf39cad
No known key found for this signature in database

View File

@ -538,6 +538,8 @@ pub enum Embedder {
Ollama(ollama::Embedder), Ollama(ollama::Embedder),
/// An embedder based on making embedding queries against a generic JSON/REST embedding server. /// An embedder based on making embedding queries against a generic JSON/REST embedding server.
Rest(rest::Embedder), Rest(rest::Embedder),
/// An embedder composed of an embedder at search time and an embedder at indexing time.
Composite(composite::Embedder),
} }
/// Configuration for an embedder. /// Configuration for an embedder.
@ -607,6 +609,7 @@ pub enum EmbedderOptions {
Ollama(ollama::EmbedderOptions), Ollama(ollama::EmbedderOptions),
UserProvided(manual::EmbedderOptions), UserProvided(manual::EmbedderOptions),
Rest(rest::EmbedderOptions), Rest(rest::EmbedderOptions),
Composite(composite::EmbedderOptions),
} }
impl Default for EmbedderOptions { impl Default for EmbedderOptions {
@ -648,6 +651,7 @@ impl Embedder {
Embedder::Ollama(embedder) => embedder.embed(&texts, deadline), Embedder::Ollama(embedder) => embedder.embed(&texts, deadline),
Embedder::UserProvided(embedder) => embedder.embed(&texts), Embedder::UserProvided(embedder) => embedder.embed(&texts),
Embedder::Rest(embedder) => embedder.embed(texts, deadline), Embedder::Rest(embedder) => embedder.embed(texts, deadline),
Embedder::Composite(embedder) => embedder.search.embed(texts, deadline),
} }
} }
@ -676,6 +680,7 @@ impl Embedder {
Embedder::Ollama(embedder) => embedder.embed_index(text_chunks, threads), Embedder::Ollama(embedder) => embedder.embed_index(text_chunks, threads),
Embedder::UserProvided(embedder) => embedder.embed_index(text_chunks), Embedder::UserProvided(embedder) => embedder.embed_index(text_chunks),
Embedder::Rest(embedder) => embedder.embed_index(text_chunks, threads), Embedder::Rest(embedder) => embedder.embed_index(text_chunks, threads),
Embedder::Composite(embedder) => embedder.index.embed_index(text_chunks, threads),
} }
} }
@ -691,6 +696,7 @@ impl Embedder {
Embedder::Ollama(embedder) => embedder.embed_index_ref(texts, threads), Embedder::Ollama(embedder) => embedder.embed_index_ref(texts, threads),
Embedder::UserProvided(embedder) => embedder.embed_index_ref(texts), Embedder::UserProvided(embedder) => embedder.embed_index_ref(texts),
Embedder::Rest(embedder) => embedder.embed_index_ref(texts, threads), Embedder::Rest(embedder) => embedder.embed_index_ref(texts, threads),
Embedder::Composite(embedder) => embedder.index.embed_index_ref(texts, threads),
} }
} }
@ -702,6 +708,7 @@ impl Embedder {
Embedder::Ollama(embedder) => embedder.chunk_count_hint(), Embedder::Ollama(embedder) => embedder.chunk_count_hint(),
Embedder::UserProvided(_) => 100, Embedder::UserProvided(_) => 100,
Embedder::Rest(embedder) => embedder.chunk_count_hint(), Embedder::Rest(embedder) => embedder.chunk_count_hint(),
Embedder::Composite(embedder) => embedder.index.chunk_count_hint(),
} }
} }
@ -713,6 +720,7 @@ impl Embedder {
Embedder::Ollama(embedder) => embedder.prompt_count_in_chunk_hint(), Embedder::Ollama(embedder) => embedder.prompt_count_in_chunk_hint(),
Embedder::UserProvided(_) => 1, Embedder::UserProvided(_) => 1,
Embedder::Rest(embedder) => embedder.prompt_count_in_chunk_hint(), Embedder::Rest(embedder) => embedder.prompt_count_in_chunk_hint(),
Embedder::Composite(embedder) => embedder.index.prompt_count_in_chunk_hint(),
} }
} }
@ -724,6 +732,7 @@ impl Embedder {
Embedder::Ollama(embedder) => embedder.dimensions(), Embedder::Ollama(embedder) => embedder.dimensions(),
Embedder::UserProvided(embedder) => embedder.dimensions(), Embedder::UserProvided(embedder) => embedder.dimensions(),
Embedder::Rest(embedder) => embedder.dimensions(), Embedder::Rest(embedder) => embedder.dimensions(),
Embedder::Composite(embedder) => embedder.dimensions(),
} }
} }
@ -735,6 +744,7 @@ impl Embedder {
Embedder::Ollama(embedder) => embedder.distribution(), Embedder::Ollama(embedder) => embedder.distribution(),
Embedder::UserProvided(embedder) => embedder.distribution(), Embedder::UserProvided(embedder) => embedder.distribution(),
Embedder::Rest(embedder) => embedder.distribution(), Embedder::Rest(embedder) => embedder.distribution(),
Embedder::Composite(embedder) => embedder.distribution(),
} }
} }
@ -745,6 +755,7 @@ impl Embedder {
| Embedder::Ollama(_) | Embedder::Ollama(_)
| Embedder::Rest(_) => true, | Embedder::Rest(_) => true,
Embedder::UserProvided(_) => false, Embedder::UserProvided(_) => false,
Embedder::Composite(embedder) => embedder.index.uses_document_template(),
} }
} }
} }