From 88d27949cd0aef343277393391c77d00409d5533 Mon Sep 17 00:00:00 2001 From: Louis Dureuil Date: Tue, 12 Mar 2024 10:56:16 +0100 Subject: [PATCH 01/86] Add documentation for benchmarks --- BENCHMARKS.md | 354 ++++++++++++++++++++++++++++++++++++++++++++++++ CONTRIBUTING.md | 24 ++++ 2 files changed, 378 insertions(+) create mode 100644 BENCHMARKS.md diff --git a/BENCHMARKS.md b/BENCHMARKS.md new file mode 100644 index 000000000..dd69864cc --- /dev/null +++ b/BENCHMARKS.md @@ -0,0 +1,354 @@ +# Benchmarks + +Currently this repository hosts two kinds of benchmarks: + +1. The older "milli benchmarks", that use [criterion](https://github.com/bheisler/criterion.rs) and live in the "benchmarks" directory. +2. The newer "bench" that are workload-based and so split between the [`workloads`](./workloads/) directory and the [`xtask::bench`](./xtask/src/bench/) module. + +This document describes the newer "bench" benchmarks. For more details on the "milli benchmarks", see [benchmarks/README.md](./benchmarks/README.md). + +## Design philosophy for the benchmarks + +The newer "bench" benchmarks are **integration** benchmarks, in the sense that they spawn an actual Meilisearch server and measure its performance end-to-end, including HTTP request overhead. + +Since this is prone to fluctuating, the benchmarks regain a bit of precision by measuring the runtime of the individual spans using the [logging machinery](./CONTRIBUTING.md#logging) of Meilisearch. + +A span roughly translates to a function call. The benchmark runner collects all the spans by name using the [logs route](https://github.com/orgs/meilisearch/discussions/721) and sums their runtime. The processed results are then sent to the [benchmark dashboard](https://bench.meilisearch.dev), which is in charge of storing and presenting the data. + +## Running the benchmarks + +Benchmarks can run locally or in CI. + +### Locally + +#### With a local benchmark dashboard + +The benchmarks dashboard lives in its [own repository](https://github.com/meilisearch/benchboard). We provide binaries for Ubuntu/Debian, but you can build from source for other platforms (MacOS should work as it was developed under that platform). + +Run the `benchboard` binary to create a fresh database of results. By default it will serve the results and the API to gather results on `http://localhost:9001`. + +From the Meilisearch repository, you can then run benchmarks with: + +```sh +cargo xtask bench -- workloads/my_workload_1.json .. +``` + +This command will build and run Meilisearch locally on port 7700, so make sure that this port is available. +To run benchmarks on a different commit, just use the usual git command to get back to the desired commit. + +#### Without a local benchmark dashboard + +To work with the raw results, you can also skip using a local benchmark dashboard. + +Run: + +```sh +cargo xtask bench --no-dashboard -- workloads/my_workload_1.json workloads/my_workload_2.json .. +``` + +For processing the results, look at [Looking at benchmark results/Without dashboard](#without-dashboard). + +### In CI + +We have dedicated runners to run workloads on CI. Currently, there are three ways of running the CI: + +1. Automatically, on every push to `main`. +2. Manually, by clicking the [`Run workflow`](https://github.com/meilisearch/meilisearch/actions/workflows/bench-manual.yml) button and specifying the target reference (tag, commit or branch) as well as one or multiple workloads to run. The workloads must exist in the Meilisearch repository (conventionally, in the [`workloads`](./workloads/) directory) on the target reference. Globbing (e.g., `workloads/*.json`) works. +3. Manually on a PR, by posting a comment containing a `/bench` command, followed by one or multiple workloads to run. Globbing works. The workloads must exist in the Meilisearch repository in the branch of the PR. + ``` + /bench workloads/movies*.json /hackernews_1M.json + ``` + +## Looking at benchmark results + +### On the dashboard + +Results are available on the global dashboard used by CI at or on your [local dashboard](#with-a-local-benchmark-dashboard). + +The dashboard homepage presents three sections: + +1. The latest invocations (a call to `cargo xtask bench`, either local or by CI) with their reason (generally set to some helpful link in CI) and their status. +2. The latest workloads ran on `main`. +3. The latest workloads ran on other references. + +By default, the workload shows the total runtime delta with the latest applicable commit on `main`. The latest applicable commit is the latest commit for workload invocations that do not originate on `main`, and the latest previous commit for workload invocations that originate on `main`. + +You can explicitly request a detailed comparison by span with the `main` branch, the branch or origin, or any previous commit, by clicking the links at the bottom of the workload invocation. + +In the detailed comparison view, the spans are sorted by improvements, regressions, stable (no statistically significant change) and unstable (the span runtime is comparable to its standard deviation). + +You can click on the name of any span to get a box plot comparing the target commit with multiple commits of the selected branch. + +### Without dashboard + +After the workloads are done running, the reports will live in the Meilisearch repository, in the `bench/reports` directory (by default). + +You can then convert these reports into other formats. + +- To [Firefox profiler](https://profiler.firefox.com) format. Run: + ```sh + cd bench/reports + cargo run --release --bin trace-to-firefox -- my_workload_1-0-trace.json + ``` + You can then upload the resulting `firefox-my_workload_1-0-trace.json` file to the online profiler. + + +## Designing benchmark workloads + +Benchmark workloads conventionally live in the `workloads` directory of the Meilisearch repository. + +They are JSON files with the following structure (comments are not actually supported, to make your own, remove them or copy some existing workload file): + +```jsonc +{ + // Name of the workload. Must be unique to the workload, as it will be used to group results on the dashboard. + "name": "hackernews.ndjson_1M,no-threads", + // Number of consecutive runs of the commands that should be performed. + // Each run uses a fresh instance of Meilisearch and a fresh database. + // Each run produces its own report file. + "run_count": 3, + // List of arguments to add to the Meilisearch command line. + "extra_cli_args": ["--max-indexing-threads=1"], + // List of named assets that can be used in the commands. + "assets": { + // name of the asset. + // Must be unique at the workload level. + // For better results, the same asset (same sha256) should have the same name accross workloads. + // Having multiple assets with the same name and distinct hashes is supported accross workloads, + // but will lead to superfluous downloads. + // + // Assets are stored in the `bench/assets/` directory by default. + "hackernews-100_000.ndjson": { + // If the assets exists in the local filesystem (Meilisearch repository or for your local workloads) + // Its file path can be specified here. + // `null` if the asset should be downloaded from a remote location. + "local_location": null, + // URL of the remote location where the asset can be downloaded. + // Use the `--assets-key` of the runner to pass an API key in the `Authorization: Bearer` header of the download requests. + // `null` if the asset should be imported from a local location. + // if both local and remote locations are specified, then the local one is tried first, then the remote one + // if the file is locally missing or its hash differs. + "remote_location": "https://milli-benchmarks.fra1.digitaloceanspaces.com/bench/datasets/hackernews/hackernews-100_000.ndjson", + // SHA256 of the asset. + // Optional, the `sha256` of the asset will be displayed during a run of the workload if it is missing. + // If present, the hash of the asset in the `bench/assets/` directory will be compared against this hash before + // running the workload. If the hashes differ, the asset will be downloaded anew. + "sha256": "60ecd23485d560edbd90d9ca31f0e6dba1455422f2a44e402600fbb5f7f1b213", + // Optional, one of "Auto", "Json", "NdJson" or "Raw". + // If missing, assumed to be "Auto". + // If "Auto", the format will be determined from the extension in the asset name. + "format": "NdJson" + }, + "hackernews-200_000.ndjson": { + "local_location": null, + "remote_location": "https://milli-benchmarks.fra1.digitaloceanspaces.com/bench/datasets/hackernews/hackernews-200_000.ndjson", + "sha256": "785b0271fdb47cba574fab617d5d332276b835c05dd86e4a95251cf7892a1685" + }, + "hackernews-300_000.ndjson": { + "local_location": null, + "remote_location": "https://milli-benchmarks.fra1.digitaloceanspaces.com/bench/datasets/hackernews/hackernews-300_000.ndjson", + "sha256": "de73c7154652eddfaf69cdc3b2f824d5c452f095f40a20a1c97bb1b5c4d80ab2" + }, + "hackernews-400_000.ndjson": { + "local_location": null, + "remote_location": "https://milli-benchmarks.fra1.digitaloceanspaces.com/bench/datasets/hackernews/hackernews-400_000.ndjson", + "sha256": "c1b00a24689110f366447e434c201c086d6f456d54ed1c4995894102794d8fe7" + }, + "hackernews-500_000.ndjson": { + "local_location": null, + "remote_location": "https://milli-benchmarks.fra1.digitaloceanspaces.com/bench/datasets/hackernews/hackernews-500_000.ndjson", + "sha256": "ae98f9dbef8193d750e3e2dbb6a91648941a1edca5f6e82c143e7996f4840083" + }, + "hackernews-600_000.ndjson": { + "local_location": null, + "remote_location": "https://milli-benchmarks.fra1.digitaloceanspaces.com/bench/datasets/hackernews/hackernews-600_000.ndjson", + "sha256": "b495fdc72c4a944801f786400f22076ab99186bee9699f67cbab2f21f5b74dbe" + }, + "hackernews-700_000.ndjson": { + "local_location": null, + "remote_location": "https://milli-benchmarks.fra1.digitaloceanspaces.com/bench/datasets/hackernews/hackernews-700_000.ndjson", + "sha256": "4b2c63974f3dabaa4954e3d4598b48324d03c522321ac05b0d583f36cb78a28b" + }, + "hackernews-800_000.ndjson": { + "local_location": null, + "remote_location": "https://milli-benchmarks.fra1.digitaloceanspaces.com/bench/datasets/hackernews/hackernews-800_000.ndjson", + "sha256": "cb7b6afe0e6caa1be111be256821bc63b0771b2a0e1fad95af7aaeeffd7ba546" + }, + "hackernews-900_000.ndjson": { + "local_location": null, + "remote_location": "https://milli-benchmarks.fra1.digitaloceanspaces.com/bench/datasets/hackernews/hackernews-900_000.ndjson", + "sha256": "e1154ddcd398f1c867758a93db5bcb21a07b9e55530c188a2917fdef332d3ba9" + }, + "hackernews-1_000_000.ndjson": { + "local_location": null, + "remote_location": "https://milli-benchmarks.fra1.digitaloceanspaces.com/bench/datasets/hackernews/hackernews-1_000_000.ndjson", + "sha256": "27e25efd0b68b159b8b21350d9af76938710cb29ce0393fa71b41c4f3c630ffe" + } + }, + // Core of the workload. + // A list of commands to run sequentially. + // A command is a request to the Meilisearch instance that is executed while the profiling runs. + "commands": [ + { + // Meilisearch route to call. `http://localhost:7700/` will be prepended. + "route": "indexes/movies/settings", + // HTTP method to call. + "method": "PATCH", + // If applicable, body of the request. + // Optional, if missing, the body will be empty. + "body": { + // One of "empty", "inline" or "asset". + // If using "empty", you can skip the entire "body" key. + "inline": { + // when "inline" is used, the body is the JSON object that is the value of the `"inline"` key. + "displayedAttributes": [ + "title", + "by", + "score", + "time" + ], + "searchableAttributes": [ + "title" + ], + "filterableAttributes": [ + "by" + ], + "sortableAttributes": [ + "score", + "time" + ] + } + }, + // Whether to wait before running the next request. + // One of: + // - DontWait: run the next command without waiting the response to this one. + // - WaitForResponse: run the next command as soon as the response from the server is received. + // - WaitForTask: run the next command once **all** the Meilisearch tasks created up to now have finished processing. + "synchronous": "DontWait" + }, + { + "route": "indexes/movies/documents", + "method": "POST", + "body": { + // When using "asset", use the name of an asset as value to use the content of that asset as body. + // the content type is derived of the format of the asset: + // "NdJson" => "application/x-ndjson" + // "Json" => "application/json" + // "Raw" => "application/octet-stream" + // See [AssetFormat::to_content_type](https://github.com/meilisearch/meilisearch/blob/7b670a4afadb132ac4a01b6403108700501a391d/xtask/src/bench/assets.rs#L30) + // for details and up-to-date list. + "asset": "hackernews-100_000.ndjson" + }, + "synchronous": "WaitForTask" + }, + { + "route": "indexes/movies/documents", + "method": "POST", + "body": { + "asset": "hackernews-200_000.ndjson" + }, + "synchronous": "WaitForResponse" + }, + { + "route": "indexes/movies/documents", + "method": "POST", + "body": { + "asset": "hackernews-300_000.ndjson" + }, + "synchronous": "WaitForResponse" + }, + { + "route": "indexes/movies/documents", + "method": "POST", + "body": { + "asset": "hackernews-400_000.ndjson" + }, + "synchronous": "WaitForResponse" + }, + { + "route": "indexes/movies/documents", + "method": "POST", + "body": { + "asset": "hackernews-500_000.ndjson" + }, + "synchronous": "WaitForResponse" + }, + { + "route": "indexes/movies/documents", + "method": "POST", + "body": { + "asset": "hackernews-600_000.ndjson" + }, + "synchronous": "WaitForResponse" + }, + { + "route": "indexes/movies/documents", + "method": "POST", + "body": { + "asset": "hackernews-700_000.ndjson" + }, + "synchronous": "WaitForResponse" + }, + { + "route": "indexes/movies/documents", + "method": "POST", + "body": { + "asset": "hackernews-800_000.ndjson" + }, + "synchronous": "WaitForResponse" + }, + { + "route": "indexes/movies/documents", + "method": "POST", + "body": { + "asset": "hackernews-900_000.ndjson" + }, + "synchronous": "WaitForResponse" + }, + { + "route": "indexes/movies/documents", + "method": "POST", + "body": { + "asset": "hackernews-1_000_000.ndjson" + }, + "synchronous": "WaitForTask" + } + ] +} +``` + + +## Upgrading `https://bench.meilisearch.dev` + +The URL of the server is in our password manager (look for "benchboard"). + +1. Make the needed modifications on the [benchboard repository](https://github.com/meilisearch/benchboard) and merge them to main. +2. Publish a new release to produce the Ubuntu/Debian binary. +3. Download the binary locally, send it to the server: + ``` + scp -6 ~/Downloads/benchboard root@\[\]:/bench/new-benchboard + ``` + Note that the ipv6 must be between escaped square brackets for SCP. +4. SSH to the server: + ``` + ssh root@ + ``` + Note the the ipv6 must **NOT** be between escaped square brackets for SSH 🥲 +5. On the server, set the correct permissions for the new binary: + ``` + chown bench:bench /bench/new-benchboard + chmod 700 /bench/new-benchboard + ``` +6. On the server, move the new binary to the location of the running binary (if unsure, start by making a backup of the running binary): + ``` + mv /bench/{new-,}benchboard + ``` +7. Restart the benchboard service. + ``` + systemctl restart benchboard + ``` +8. Check that the service runs correctly. + ``` + systemctl status benchboard + ``` +9. Check the availability of the service by going to on your browser. diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 073da7031..6d6e6076b 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -81,6 +81,30 @@ Meilisearch follows the [cargo xtask](https://github.com/matklad/cargo-xtask) wo Run `cargo xtask --help` from the root of the repository to find out what is available. +### Logging + +Meilisearch uses [`tracing`](https://lib.rs/crates/tracing) for logging purposes. Tracing logs are structured and can be displayed as JSON to the end user, so prefer passing arguments as fields rather than interpolating them in the message. + +Refer to the [documentation](https://docs.rs/tracing/0.1.40/tracing/index.html#using-the-macros) for the syntax of the spans and events. + +Logging spans are used for 3 distinct purposes: + +1. Regular logging +2. Profiling +3. Benchmarking + +As a result, the spans should follow some rules: + +- They should not be put on functions that are called too often. That is because opening and closing a span causes some overhead. For regular logging, avoid putting spans on functions that are taking less than a few hundred nanoseconds. For profiling or benchmarking, avoid putting spans on functions that are taking less than a few microseconds. +- For profiling and benchmarking, use the `TRACE` level. +- For profiling and benchmarking, use the following `target` prefixes: + - `indexing::` for spans meant when profiling the indexing operations. + - `search::` for spans meant when profiling the search operations. + +### Benchmarking + +See [BENCHMARKS.md](./BENCHMARKS.md) + ## Git Guidelines ### Git Branches From 4a467739cdfd6c6eafb73ae57649199729bf8d7b Mon Sep 17 00:00:00 2001 From: Tamo Date: Tue, 5 Mar 2024 11:21:46 +0100 Subject: [PATCH 02/86] implements a first version of the cutoff without settings --- .../src/analytics/segment_analytics.rs | 9 ++ meilisearch/src/search.rs | 39 ++++++--- milli/examples/search.rs | 3 +- milli/src/index.rs | 1 + milli/src/lib.rs | 35 ++++++++ milli/src/search/hybrid.rs | 2 + milli/src/search/mod.rs | 82 +++++++++++-------- milli/src/search/new/bucket_sort.rs | 27 +++++- milli/src/search/new/matches/mod.rs | 3 +- milli/src/search/new/mod.rs | 16 +++- milli/tests/search/mod.rs | 45 +++++++++- 11 files changed, 210 insertions(+), 52 deletions(-) diff --git a/meilisearch/src/analytics/segment_analytics.rs b/meilisearch/src/analytics/segment_analytics.rs index 7dfc52900..99298bd43 100644 --- a/meilisearch/src/analytics/segment_analytics.rs +++ b/meilisearch/src/analytics/segment_analytics.rs @@ -579,6 +579,7 @@ pub struct SearchAggregator { // requests total_received: usize, total_succeeded: usize, + total_degraded: usize, time_spent: BinaryHeap, // sort @@ -758,9 +759,13 @@ impl SearchAggregator { hits_info: _, facet_distribution: _, facet_stats: _, + degraded, } = result; self.total_succeeded = self.total_succeeded.saturating_add(1); + if *degraded { + self.total_degraded = self.total_degraded.saturating_add(1); + } self.time_spent.push(*processing_time_ms as usize); } @@ -802,6 +807,7 @@ impl SearchAggregator { semantic_ratio, embedder, hybrid, + total_degraded, } = other; if self.timestamp.is_none() { @@ -816,6 +822,7 @@ impl SearchAggregator { // request self.total_received = self.total_received.saturating_add(total_received); self.total_succeeded = self.total_succeeded.saturating_add(total_succeeded); + self.total_degraded = self.total_degraded.saturating_add(total_degraded); self.time_spent.append(time_spent); // sort @@ -921,6 +928,7 @@ impl SearchAggregator { semantic_ratio, embedder, hybrid, + total_degraded, } = self; if total_received == 0 { @@ -940,6 +948,7 @@ impl SearchAggregator { "total_succeeded": total_succeeded, "total_failed": total_received.saturating_sub(total_succeeded), // just to be sure we never panics "total_received": total_received, + "total_degraded": total_degraded, }, "sort": { "with_geoPoint": sort_with_geo_point, diff --git a/meilisearch/src/search.rs b/meilisearch/src/search.rs index e65192d16..9bc7b69fc 100644 --- a/meilisearch/src/search.rs +++ b/meilisearch/src/search.rs @@ -1,7 +1,7 @@ use std::cmp::min; use std::collections::{BTreeMap, BTreeSet, HashSet}; use std::str::FromStr; -use std::time::Instant; +use std::time::{Duration, Instant}; use deserr::Deserr; use either::Either; @@ -14,7 +14,7 @@ use meilisearch_types::heed::RoTxn; use meilisearch_types::index_uid::IndexUid; use meilisearch_types::milli::score_details::{self, ScoreDetails, ScoringStrategy}; use meilisearch_types::milli::vector::DistributionShift; -use meilisearch_types::milli::{FacetValueHit, OrderBy, SearchForFacetValues}; +use meilisearch_types::milli::{FacetValueHit, OrderBy, SearchForFacetValues, TimeBudget}; use meilisearch_types::settings::DEFAULT_PAGINATION_MAX_TOTAL_HITS; use meilisearch_types::{milli, Document}; use milli::tokenizer::TokenizerBuilder; @@ -323,6 +323,9 @@ pub struct SearchResult { pub facet_distribution: Option>>, #[serde(skip_serializing_if = "Option::is_none")] pub facet_stats: Option>, + + #[serde(skip_serializing_if = "std::ops::Not::not")] + pub degraded: bool, } #[derive(Serialize, Debug, Clone, PartialEq)] @@ -382,8 +385,10 @@ fn prepare_search<'t>( query: &'t SearchQuery, features: RoFeatures, distribution: Option, + time_budget: TimeBudget, ) -> Result<(milli::Search<'t>, bool, usize, usize), MeilisearchHttpError> { let mut search = index.search(rtxn); + search.time_budget(time_budget); if query.vector.is_some() { features.check_vector("Passing `vector` as a query parameter")?; @@ -491,19 +496,26 @@ pub fn perform_search( distribution: Option, ) -> Result { let before_search = Instant::now(); + let time_budget = TimeBudget::new(Duration::from_millis(150)); let rtxn = index.read_txn()?; let (search, is_finite_pagination, max_total_hits, offset) = - prepare_search(index, &rtxn, &query, features, distribution)?; + prepare_search(index, &rtxn, &query, features, distribution, time_budget)?; - let milli::SearchResult { documents_ids, matching_words, candidates, document_scores, .. } = - match &query.hybrid { - Some(hybrid) => match *hybrid.semantic_ratio { - ratio if ratio == 0.0 || ratio == 1.0 => search.execute()?, - ratio => search.execute_hybrid(ratio)?, - }, - None => search.execute()?, - }; + let milli::SearchResult { + documents_ids, + matching_words, + candidates, + document_scores, + degraded, + .. + } = match &query.hybrid { + Some(hybrid) => match *hybrid.semantic_ratio { + ratio if ratio == 0.0 || ratio == 1.0 => search.execute()?, + ratio => search.execute_hybrid(ratio)?, + }, + None => search.execute()?, + }; let fields_ids_map = index.fields_ids_map(&rtxn).unwrap(); @@ -700,6 +712,7 @@ pub fn perform_search( processing_time_ms: before_search.elapsed().as_millis(), facet_distribution, facet_stats, + degraded, }; Ok(result) } @@ -712,9 +725,11 @@ pub fn perform_facet_search( features: RoFeatures, ) -> Result { let before_search = Instant::now(); + let time_budget = TimeBudget::new(Duration::from_millis(150)); let rtxn = index.read_txn()?; - let (search, _, _, _) = prepare_search(index, &rtxn, &search_query, features, None)?; + let (search, _, _, _) = + prepare_search(index, &rtxn, &search_query, features, None, time_budget)?; let mut facet_search = SearchForFacetValues::new(facet_name, search, search_query.hybrid.is_some()); if let Some(facet_query) = &facet_query { diff --git a/milli/examples/search.rs b/milli/examples/search.rs index a94677771..8640acf42 100644 --- a/milli/examples/search.rs +++ b/milli/examples/search.rs @@ -6,7 +6,7 @@ use std::time::Instant; use heed::EnvOpenOptions; use milli::{ execute_search, filtered_universe, DefaultSearchLogger, GeoSortStrategy, Index, SearchContext, - SearchLogger, TermsMatchingStrategy, + SearchLogger, TermsMatchingStrategy, TimeBudget, }; #[global_allocator] @@ -65,6 +65,7 @@ fn main() -> Result<(), Box> { None, &mut DefaultSearchLogger, logger, + TimeBudget::max(), )?; if let Some((logger, dir)) = detailed_logger { logger.finish(&mut ctx, Path::new(dir))?; diff --git a/milli/src/index.rs b/milli/src/index.rs index 2c3977403..e79c137e7 100644 --- a/milli/src/index.rs +++ b/milli/src/index.rs @@ -2421,6 +2421,7 @@ pub(crate) mod tests { candidates: _, document_scores: _, mut documents_ids, + degraded: _, } = search.execute().unwrap(); let primary_key_id = index.fields_ids_map(&rtxn).unwrap().id("primary_key").unwrap(); documents_ids.sort_unstable(); diff --git a/milli/src/lib.rs b/milli/src/lib.rs index 5effcea3d..eedd25f7e 100644 --- a/milli/src/lib.rs +++ b/milli/src/lib.rs @@ -30,6 +30,7 @@ pub mod snapshot_tests; use std::collections::{BTreeMap, HashMap}; use std::convert::{TryFrom, TryInto}; +use std::fmt; use std::hash::BuildHasherDefault; use charabia::normalizer::{CharNormalizer, CompatibilityDecompositionNormalizer}; @@ -104,6 +105,40 @@ pub const MAX_WORD_LENGTH: usize = MAX_LMDB_KEY_LENGTH / 2; pub const MAX_POSITION_PER_ATTRIBUTE: u32 = u16::MAX as u32 + 1; +#[derive(Clone, Copy)] +pub struct TimeBudget { + started_at: std::time::Instant, + budget: std::time::Duration, +} + +impl fmt::Debug for TimeBudget { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("TimeBudget") + .field("started_at", &self.started_at) + .field("budget", &self.budget) + .field("left", &(self.budget - self.started_at.elapsed())) + .finish() + } +} + +impl TimeBudget { + pub fn new(budget: std::time::Duration) -> Self { + Self { started_at: std::time::Instant::now(), budget } + } + + pub fn max() -> Self { + Self::new(std::time::Duration::from_secs(u64::MAX)) + } + + pub fn exceeded(&self) -> bool { + self.must_stop() + } + + pub fn must_stop(&self) -> bool { + self.started_at.elapsed() > self.budget + } +} + // Convert an absolute word position into a relative position. // Return the field id of the attribute related to the absolute position // and the relative position in the attribute. diff --git a/milli/src/search/hybrid.rs b/milli/src/search/hybrid.rs index b4c79f7f5..9d8b3860d 100644 --- a/milli/src/search/hybrid.rs +++ b/milli/src/search/hybrid.rs @@ -106,6 +106,7 @@ impl ScoreWithRatioResult { candidates: left.candidates | right.candidates, documents_ids, document_scores, + degraded: false, } } } @@ -131,6 +132,7 @@ impl<'a> Search<'a> { index: self.index, distribution_shift: self.distribution_shift, embedder_name: self.embedder_name.clone(), + time_budget: self.time_budget, }; let vector_query = search.vector.take(); diff --git a/milli/src/search/mod.rs b/milli/src/search/mod.rs index dc8354486..b14d88d03 100644 --- a/milli/src/search/mod.rs +++ b/milli/src/search/mod.rs @@ -11,7 +11,7 @@ use crate::score_details::{ScoreDetails, ScoringStrategy}; use crate::vector::DistributionShift; use crate::{ execute_search, filtered_universe, AscDesc, DefaultSearchLogger, DocumentId, Index, Result, - SearchContext, + SearchContext, TimeBudget, }; // Building these factories is not free. @@ -43,6 +43,8 @@ pub struct Search<'a> { index: &'a Index, distribution_shift: Option, embedder_name: Option, + + time_budget: TimeBudget, } impl<'a> Search<'a> { @@ -64,6 +66,7 @@ impl<'a> Search<'a> { index, distribution_shift: None, embedder_name: None, + time_budget: TimeBudget::max(), } } @@ -143,6 +146,11 @@ impl<'a> Search<'a> { self } + pub fn time_budget(&mut self, time_budget: TimeBudget) -> &mut Search<'a> { + self.time_budget = time_budget; + self + } + pub fn execute_for_candidates(&self, has_vector_search: bool) -> Result { if has_vector_search { let ctx = SearchContext::new(self.index, self.rtxn); @@ -169,36 +177,43 @@ impl<'a> Search<'a> { } let universe = filtered_universe(&ctx, &self.filter)?; - let PartialSearchResult { located_query_terms, candidates, documents_ids, document_scores } = - match self.vector.as_ref() { - Some(vector) => execute_vector_search( - &mut ctx, - vector, - self.scoring_strategy, - universe, - &self.sort_criteria, - self.geo_strategy, - self.offset, - self.limit, - self.distribution_shift, - embedder_name, - )?, - None => execute_search( - &mut ctx, - self.query.as_deref(), - self.terms_matching_strategy, - self.scoring_strategy, - self.exhaustive_number_hits, - universe, - &self.sort_criteria, - self.geo_strategy, - self.offset, - self.limit, - Some(self.words_limit), - &mut DefaultSearchLogger, - &mut DefaultSearchLogger, - )?, - }; + let PartialSearchResult { + located_query_terms, + candidates, + documents_ids, + document_scores, + degraded, + } = match self.vector.as_ref() { + Some(vector) => execute_vector_search( + &mut ctx, + vector, + self.scoring_strategy, + universe, + &self.sort_criteria, + self.geo_strategy, + self.offset, + self.limit, + self.distribution_shift, + embedder_name, + self.time_budget, + )?, + None => execute_search( + &mut ctx, + self.query.as_deref(), + self.terms_matching_strategy, + self.scoring_strategy, + self.exhaustive_number_hits, + universe, + &self.sort_criteria, + self.geo_strategy, + self.offset, + self.limit, + Some(self.words_limit), + &mut DefaultSearchLogger, + &mut DefaultSearchLogger, + self.time_budget, + )?, + }; // consume context and located_query_terms to build MatchingWords. let matching_words = match located_query_terms { @@ -206,7 +221,7 @@ impl<'a> Search<'a> { None => MatchingWords::default(), }; - Ok(SearchResult { matching_words, candidates, document_scores, documents_ids }) + Ok(SearchResult { matching_words, candidates, document_scores, documents_ids, degraded }) } } @@ -229,6 +244,7 @@ impl fmt::Debug for Search<'_> { index: _, distribution_shift, embedder_name, + time_budget, } = self; f.debug_struct("Search") .field("query", query) @@ -244,6 +260,7 @@ impl fmt::Debug for Search<'_> { .field("words_limit", words_limit) .field("distribution_shift", distribution_shift) .field("embedder_name", embedder_name) + .field("time_bduget", time_budget) .finish() } } @@ -254,6 +271,7 @@ pub struct SearchResult { pub candidates: RoaringBitmap, pub documents_ids: Vec, pub document_scores: Vec>, + pub degraded: bool, } #[derive(Debug, Clone, Copy, PartialEq, Eq)] diff --git a/milli/src/search/new/bucket_sort.rs b/milli/src/search/new/bucket_sort.rs index 02528e378..7fc830c1f 100644 --- a/milli/src/search/new/bucket_sort.rs +++ b/milli/src/search/new/bucket_sort.rs @@ -5,12 +5,14 @@ use super::ranking_rules::{BoxRankingRule, RankingRuleQueryTrait}; use super::SearchContext; use crate::score_details::{ScoreDetails, ScoringStrategy}; use crate::search::new::distinct::{apply_distinct_rule, distinct_single_docid, DistinctOutput}; -use crate::Result; +use crate::{Result, TimeBudget}; pub struct BucketSortOutput { pub docids: Vec, pub scores: Vec>, pub all_candidates: RoaringBitmap, + + pub degraded: bool, } // TODO: would probably be good to regroup some of these inside of a struct? @@ -25,6 +27,7 @@ pub fn bucket_sort<'ctx, Q: RankingRuleQueryTrait>( length: usize, scoring_strategy: ScoringStrategy, logger: &mut dyn SearchLogger, + time_budget: TimeBudget, ) -> Result { logger.initial_query(query); logger.ranking_rules(&ranking_rules); @@ -41,6 +44,7 @@ pub fn bucket_sort<'ctx, Q: RankingRuleQueryTrait>( docids: vec![], scores: vec![], all_candidates: universe.clone(), + degraded: false, }); } if ranking_rules.is_empty() { @@ -74,6 +78,7 @@ pub fn bucket_sort<'ctx, Q: RankingRuleQueryTrait>( scores: vec![Default::default(); results.len()], docids: results, all_candidates, + degraded: false, }); } else { let docids: Vec = universe.iter().skip(from).take(length).collect(); @@ -81,6 +86,7 @@ pub fn bucket_sort<'ctx, Q: RankingRuleQueryTrait>( scores: vec![Default::default(); docids.len()], docids, all_candidates: universe.clone(), + degraded: false, }); }; } @@ -154,6 +160,18 @@ pub fn bucket_sort<'ctx, Q: RankingRuleQueryTrait>( } while valid_docids.len() < length { + if time_budget.exceeded() { + let bucket = std::mem::take(&mut ranking_rule_universes[cur_ranking_rule_index]); + maybe_add_to_results!(bucket); + + return Ok(BucketSortOutput { + scores: vec![Default::default(); valid_docids.len()], + docids: valid_docids, + all_candidates, + degraded: true, + }); + } + // The universe for this bucket is zero, so we don't need to sort // anything, just go back to the parent ranking rule. if ranking_rule_universes[cur_ranking_rule_index].is_empty() @@ -219,7 +237,12 @@ pub fn bucket_sort<'ctx, Q: RankingRuleQueryTrait>( )?; } - Ok(BucketSortOutput { docids: valid_docids, scores: valid_scores, all_candidates }) + Ok(BucketSortOutput { + docids: valid_docids, + scores: valid_scores, + all_candidates, + degraded: false, + }) } /// Add the candidates to the results. Take `distinct`, `from`, `length`, and `cur_offset` diff --git a/milli/src/search/new/matches/mod.rs b/milli/src/search/new/matches/mod.rs index 8de1d9262..2913f206d 100644 --- a/milli/src/search/new/matches/mod.rs +++ b/milli/src/search/new/matches/mod.rs @@ -502,7 +502,7 @@ mod tests { use super::*; use crate::index::tests::TempIndex; - use crate::{execute_search, filtered_universe, SearchContext}; + use crate::{execute_search, filtered_universe, SearchContext, TimeBudget}; impl<'a> MatcherBuilder<'a> { fn new_test(rtxn: &'a heed::RoTxn, index: &'a TempIndex, query: &str) -> Self { @@ -522,6 +522,7 @@ mod tests { Some(10), &mut crate::DefaultSearchLogger, &mut crate::DefaultSearchLogger, + TimeBudget::max(), ) .unwrap(); diff --git a/milli/src/search/new/mod.rs b/milli/src/search/new/mod.rs index ae661e3f6..ad996f363 100644 --- a/milli/src/search/new/mod.rs +++ b/milli/src/search/new/mod.rs @@ -52,7 +52,8 @@ use crate::score_details::{ScoreDetails, ScoringStrategy}; use crate::search::new::distinct::apply_distinct_rule; use crate::vector::DistributionShift; use crate::{ - AscDesc, DocumentId, FieldId, Filter, Index, Member, Result, TermsMatchingStrategy, UserError, + AscDesc, DocumentId, FieldId, Filter, Index, Member, Result, TermsMatchingStrategy, TimeBudget, + UserError, }; /// A structure used throughout the execution of a search query. @@ -518,6 +519,7 @@ pub fn execute_vector_search( length: usize, distribution_shift: Option, embedder_name: &str, + time_budget: TimeBudget, ) -> Result { check_sort_criteria(ctx, sort_criteria.as_ref())?; @@ -537,7 +539,7 @@ pub fn execute_vector_search( let placeholder_search_logger: &mut dyn SearchLogger = &mut placeholder_search_logger; - let BucketSortOutput { docids, scores, all_candidates } = bucket_sort( + let BucketSortOutput { docids, scores, all_candidates, degraded } = bucket_sort( ctx, ranking_rules, &PlaceholderQuery, @@ -546,6 +548,7 @@ pub fn execute_vector_search( length, scoring_strategy, placeholder_search_logger, + time_budget, )?; Ok(PartialSearchResult { @@ -553,6 +556,7 @@ pub fn execute_vector_search( document_scores: scores, documents_ids: docids, located_query_terms: None, + degraded, }) } @@ -572,6 +576,7 @@ pub fn execute_search( words_limit: Option, placeholder_search_logger: &mut dyn SearchLogger, query_graph_logger: &mut dyn SearchLogger, + time_budget: TimeBudget, ) -> Result { check_sort_criteria(ctx, sort_criteria.as_ref())?; @@ -648,6 +653,7 @@ pub fn execute_search( length, scoring_strategy, query_graph_logger, + time_budget, )? } else { let ranking_rules = @@ -661,10 +667,11 @@ pub fn execute_search( length, scoring_strategy, placeholder_search_logger, + time_budget, )? }; - let BucketSortOutput { docids, scores, mut all_candidates } = bucket_sort_output; + let BucketSortOutput { docids, scores, mut all_candidates, degraded } = bucket_sort_output; let fields_ids_map = ctx.index.fields_ids_map(ctx.txn)?; // The candidates is the universe unless the exhaustive number of hits @@ -682,6 +689,7 @@ pub fn execute_search( document_scores: scores, documents_ids: docids, located_query_terms, + degraded, }) } @@ -742,4 +750,6 @@ pub struct PartialSearchResult { pub candidates: RoaringBitmap, pub documents_ids: Vec, pub document_scores: Vec>, + + pub degraded: bool, } diff --git a/milli/tests/search/mod.rs b/milli/tests/search/mod.rs index 9193ab762..ab6befa60 100644 --- a/milli/tests/search/mod.rs +++ b/milli/tests/search/mod.rs @@ -1,14 +1,19 @@ use std::cmp::Reverse; use std::collections::HashSet; use std::io::Cursor; +use std::time::Duration; use big_s::S; use either::{Either, Left, Right}; use heed::EnvOpenOptions; use maplit::{btreemap, hashset}; +use meili_snap::snapshot; use milli::documents::{DocumentsBatchBuilder, DocumentsBatchReader}; use milli::update::{IndexDocuments, IndexDocumentsConfig, IndexerConfig, Settings}; -use milli::{AscDesc, Criterion, DocumentId, Index, Member, Object, TermsMatchingStrategy}; +use milli::{ + AscDesc, Criterion, DocumentId, Filter, Index, Member, Object, Search, TermsMatchingStrategy, + TimeBudget, +}; use serde::{Deserialize, Deserializer}; use slice_group_by::GroupBy; @@ -349,3 +354,41 @@ where let result = serde_json::Value::deserialize(deserializer)?; Ok(Some(result)) } + +#[test] +fn basic_degraded_search() { + use Criterion::*; + let criteria = vec![Words, Typo, Proximity, Attribute, Exactness]; + let index = setup_search_index_with_criteria(&criteria); + let rtxn = index.read_txn().unwrap(); + + let mut search = Search::new(&rtxn, &index); + search.query(TEST_QUERY); + search.limit(EXTERNAL_DOCUMENTS_IDS.len()); + search.time_budget(TimeBudget::new(Duration::from_millis(0))); + + let result = search.execute().unwrap(); + assert!(result.degraded); +} + +#[test] +fn degraded_search_cannot_skip_filter() { + use Criterion::*; + let criteria = vec![Words, Typo, Proximity, Attribute, Exactness]; + let index = setup_search_index_with_criteria(&criteria); + let rtxn = index.read_txn().unwrap(); + + let mut search = Search::new(&rtxn, &index); + search.query(TEST_QUERY); + search.limit(EXTERNAL_DOCUMENTS_IDS.len()); + search.time_budget(TimeBudget::new(Duration::from_millis(0))); + let filter_condition = Filter::from_str("tag = etiopia").unwrap().unwrap(); + search.filter(filter_condition); + + let result = search.execute().unwrap(); + assert!(result.degraded); + snapshot!(format!("{:?}\n{:?}", result.candidates, result.documents_ids), @r###" + RoaringBitmap<[0, 2, 5, 8, 11, 14]> + [0, 2, 5, 8, 11, 14] + "###); +} From d1db4951195bae1aaa13dd3481bc4ff0003122d7 Mon Sep 17 00:00:00 2001 From: Tamo Date: Mon, 11 Mar 2024 18:24:21 +0100 Subject: [PATCH 03/86] add a settings for the search cutoff --- dump/src/lib.rs | 1 + dump/src/reader/compat/v5_to_v6.rs | 1 + meilisearch-types/src/error.rs | 1 + meilisearch-types/src/settings.rs | 76 +++++++++++++++++----- meilisearch/src/routes/indexes/settings.rs | 22 ++++++- meilisearch/src/search.rs | 5 +- meilisearch/tests/dumps/mod.rs | 39 +++++++---- meilisearch/tests/settings/get_settings.rs | 7 +- milli/src/index.rs | 13 ++++ milli/src/lib.rs | 6 ++ milli/src/update/settings.rs | 33 ++++++++++ 11 files changed, 169 insertions(+), 35 deletions(-) diff --git a/dump/src/lib.rs b/dump/src/lib.rs index be0053a7c..e7cadacbe 100644 --- a/dump/src/lib.rs +++ b/dump/src/lib.rs @@ -277,6 +277,7 @@ pub(crate) mod test { }), pagination: Setting::NotSet, embedders: Setting::NotSet, + search_cutoff: Setting::NotSet, _kind: std::marker::PhantomData, }; settings.check() diff --git a/dump/src/reader/compat/v5_to_v6.rs b/dump/src/reader/compat/v5_to_v6.rs index e00d3a599..2b8997847 100644 --- a/dump/src/reader/compat/v5_to_v6.rs +++ b/dump/src/reader/compat/v5_to_v6.rs @@ -379,6 +379,7 @@ impl From> for v6::Settings { v5::Setting::NotSet => v6::Setting::NotSet, }, embedders: v6::Setting::NotSet, + search_cutoff: v6::Setting::NotSet, _kind: std::marker::PhantomData, } } diff --git a/meilisearch-types/src/error.rs b/meilisearch-types/src/error.rs index 965d2e672..bf9492ff6 100644 --- a/meilisearch-types/src/error.rs +++ b/meilisearch-types/src/error.rs @@ -259,6 +259,7 @@ InvalidSettingsProximityPrecision , InvalidRequest , BAD_REQUEST ; InvalidSettingsFaceting , InvalidRequest , BAD_REQUEST ; InvalidSettingsFilterableAttributes , InvalidRequest , BAD_REQUEST ; InvalidSettingsPagination , InvalidRequest , BAD_REQUEST ; +InvalidSettingsSearchCutoff , InvalidRequest , BAD_REQUEST ; InvalidSettingsEmbedders , InvalidRequest , BAD_REQUEST ; InvalidSettingsRankingRules , InvalidRequest , BAD_REQUEST ; InvalidSettingsSearchableAttributes , InvalidRequest , BAD_REQUEST ; diff --git a/meilisearch-types/src/settings.rs b/meilisearch-types/src/settings.rs index ca46abb0c..d05201943 100644 --- a/meilisearch-types/src/settings.rs +++ b/meilisearch-types/src/settings.rs @@ -202,6 +202,9 @@ pub struct Settings { #[serde(default, skip_serializing_if = "Setting::is_not_set")] #[deserr(default, error = DeserrJsonError)] pub embedders: Setting>>, + #[serde(default, skip_serializing_if = "Setting::is_not_set")] + #[deserr(default, error = DeserrJsonError)] + pub search_cutoff: Setting, #[serde(skip)] #[deserr(skip)] @@ -227,6 +230,7 @@ impl Settings { faceting: Setting::Reset, pagination: Setting::Reset, embedders: Setting::Reset, + search_cutoff: Setting::Reset, _kind: PhantomData, } } @@ -249,6 +253,7 @@ impl Settings { faceting, pagination, embedders, + search_cutoff, .. } = self; @@ -269,6 +274,7 @@ impl Settings { faceting, pagination, embedders, + search_cutoff, _kind: PhantomData, } } @@ -315,6 +321,7 @@ impl Settings { faceting: self.faceting, pagination: self.pagination, embedders: self.embedders, + search_cutoff: self.search_cutoff, _kind: PhantomData, } } @@ -347,19 +354,40 @@ pub fn apply_settings_to_builder( settings: &Settings, builder: &mut milli::update::Settings, ) { - match settings.searchable_attributes { + let Settings { + displayed_attributes, + searchable_attributes, + filterable_attributes, + sortable_attributes, + ranking_rules, + stop_words, + non_separator_tokens, + separator_tokens, + dictionary, + synonyms, + distinct_attribute, + proximity_precision, + typo_tolerance, + faceting, + pagination, + embedders, + search_cutoff, + _kind, + } = settings; + + match searchable_attributes { Setting::Set(ref names) => builder.set_searchable_fields(names.clone()), Setting::Reset => builder.reset_searchable_fields(), Setting::NotSet => (), } - match settings.displayed_attributes { + match displayed_attributes { Setting::Set(ref names) => builder.set_displayed_fields(names.clone()), Setting::Reset => builder.reset_displayed_fields(), Setting::NotSet => (), } - match settings.filterable_attributes { + match filterable_attributes { Setting::Set(ref facets) => { builder.set_filterable_fields(facets.clone().into_iter().collect()) } @@ -367,13 +395,13 @@ pub fn apply_settings_to_builder( Setting::NotSet => (), } - match settings.sortable_attributes { + match sortable_attributes { Setting::Set(ref fields) => builder.set_sortable_fields(fields.iter().cloned().collect()), Setting::Reset => builder.reset_sortable_fields(), Setting::NotSet => (), } - match settings.ranking_rules { + match ranking_rules { Setting::Set(ref criteria) => { builder.set_criteria(criteria.iter().map(|c| c.clone().into()).collect()) } @@ -381,13 +409,13 @@ pub fn apply_settings_to_builder( Setting::NotSet => (), } - match settings.stop_words { + match stop_words { Setting::Set(ref stop_words) => builder.set_stop_words(stop_words.clone()), Setting::Reset => builder.reset_stop_words(), Setting::NotSet => (), } - match settings.non_separator_tokens { + match non_separator_tokens { Setting::Set(ref non_separator_tokens) => { builder.set_non_separator_tokens(non_separator_tokens.clone()) } @@ -395,7 +423,7 @@ pub fn apply_settings_to_builder( Setting::NotSet => (), } - match settings.separator_tokens { + match separator_tokens { Setting::Set(ref separator_tokens) => { builder.set_separator_tokens(separator_tokens.clone()) } @@ -403,31 +431,31 @@ pub fn apply_settings_to_builder( Setting::NotSet => (), } - match settings.dictionary { + match dictionary { Setting::Set(ref dictionary) => builder.set_dictionary(dictionary.clone()), Setting::Reset => builder.reset_dictionary(), Setting::NotSet => (), } - match settings.synonyms { + match synonyms { Setting::Set(ref synonyms) => builder.set_synonyms(synonyms.clone().into_iter().collect()), Setting::Reset => builder.reset_synonyms(), Setting::NotSet => (), } - match settings.distinct_attribute { + match distinct_attribute { Setting::Set(ref attr) => builder.set_distinct_field(attr.clone()), Setting::Reset => builder.reset_distinct_field(), Setting::NotSet => (), } - match settings.proximity_precision { + match proximity_precision { Setting::Set(ref precision) => builder.set_proximity_precision((*precision).into()), Setting::Reset => builder.reset_proximity_precision(), Setting::NotSet => (), } - match settings.typo_tolerance { + match typo_tolerance { Setting::Set(ref value) => { match value.enabled { Setting::Set(val) => builder.set_autorize_typos(val), @@ -482,7 +510,7 @@ pub fn apply_settings_to_builder( Setting::NotSet => (), } - match &settings.faceting { + match faceting { Setting::Set(FacetingSettings { max_values_per_facet, sort_facet_values_by }) => { match max_values_per_facet { Setting::Set(val) => builder.set_max_values_per_facet(*val), @@ -504,7 +532,7 @@ pub fn apply_settings_to_builder( Setting::NotSet => (), } - match settings.pagination { + match pagination { Setting::Set(ref value) => match value.max_total_hits { Setting::Set(val) => builder.set_pagination_max_total_hits(val), Setting::Reset => builder.reset_pagination_max_total_hits(), @@ -514,11 +542,17 @@ pub fn apply_settings_to_builder( Setting::NotSet => (), } - match settings.embedders.clone() { - Setting::Set(value) => builder.set_embedder_settings(value), + match embedders { + Setting::Set(value) => builder.set_embedder_settings(value.clone()), Setting::Reset => builder.reset_embedder_settings(), Setting::NotSet => (), } + + match search_cutoff { + Setting::Set(cutoff) => builder.set_search_cutoff(*cutoff), + Setting::Reset => builder.reset_search_cutoff(), + Setting::NotSet => (), + } } pub fn settings( @@ -607,6 +641,8 @@ pub fn settings( .collect(); let embedders = if embedders.is_empty() { Setting::NotSet } else { Setting::Set(embedders) }; + let search_cutoff = index.search_cutoff(rtxn)?; + Ok(Settings { displayed_attributes: match displayed_attributes { Some(attrs) => Setting::Set(attrs), @@ -633,6 +669,10 @@ pub fn settings( faceting: Setting::Set(faceting), pagination: Setting::Set(pagination), embedders, + search_cutoff: match search_cutoff { + Some(cutoff) => Setting::Set(cutoff), + None => Setting::Reset, + }, _kind: PhantomData, }) } @@ -783,6 +823,7 @@ pub(crate) mod test { faceting: Setting::NotSet, pagination: Setting::NotSet, embedders: Setting::NotSet, + search_cutoff: Setting::NotSet, _kind: PhantomData::, }; @@ -809,6 +850,7 @@ pub(crate) mod test { faceting: Setting::NotSet, pagination: Setting::NotSet, embedders: Setting::NotSet, + search_cutoff: Setting::NotSet, _kind: PhantomData::, }; diff --git a/meilisearch/src/routes/indexes/settings.rs b/meilisearch/src/routes/indexes/settings.rs index c782e78cb..1d03c9a91 100644 --- a/meilisearch/src/routes/indexes/settings.rs +++ b/meilisearch/src/routes/indexes/settings.rs @@ -624,6 +624,25 @@ fn embedder_analytics( ) } +make_setting_route!( + "/search_cutoff", + patch, + u64, + meilisearch_types::deserr::DeserrJsonError< + meilisearch_types::error::deserr_codes::InvalidSettingsSearchCutoff, + >, + search_cutoff, + "search_cutoff", + analytics, + |setting: &Option, req: &HttpRequest| { + analytics.publish( + "Search Cutoff Updated".to_string(), + serde_json::json!({"search_cutoff": setting }), + Some(req), + ); + } +); + macro_rules! generate_configure { ($($mod:ident),*) => { pub fn configure(cfg: &mut web::ServiceConfig) { @@ -765,7 +784,8 @@ pub async fn update_all( "synonyms": { "total": new_settings.synonyms.as_ref().set().map(|synonyms| synonyms.len()), }, - "embedders": crate::routes::indexes::settings::embedder_analytics(new_settings.embedders.as_ref().set()) + "embedders": crate::routes::indexes::settings::embedder_analytics(new_settings.embedders.as_ref().set()), + "search_cutoff": new_settings.search_cutoff.as_ref().set(), }), Some(&req), ); diff --git a/meilisearch/src/search.rs b/meilisearch/src/search.rs index 9bc7b69fc..f83e14187 100644 --- a/meilisearch/src/search.rs +++ b/meilisearch/src/search.rs @@ -496,8 +496,11 @@ pub fn perform_search( distribution: Option, ) -> Result { let before_search = Instant::now(); - let time_budget = TimeBudget::new(Duration::from_millis(150)); let rtxn = index.read_txn()?; + let time_budget = match index.search_cutoff(&rtxn)? { + Some(cutoff) => TimeBudget::new(Duration::from_millis(cutoff)), + None => TimeBudget::default(), + }; let (search, is_finite_pagination, max_total_hits, offset) = prepare_search(index, &rtxn, &query, features, distribution, time_budget)?; diff --git a/meilisearch/tests/dumps/mod.rs b/meilisearch/tests/dumps/mod.rs index e8061ae4a..7bf97f8b2 100644 --- a/meilisearch/tests/dumps/mod.rs +++ b/meilisearch/tests/dumps/mod.rs @@ -77,7 +77,8 @@ async fn import_dump_v1_movie_raw() { }, "pagination": { "maxTotalHits": 1000 - } + }, + "searchCutoff": null } "### ); @@ -238,7 +239,8 @@ async fn import_dump_v1_movie_with_settings() { }, "pagination": { "maxTotalHits": 1000 - } + }, + "searchCutoff": null } "### ); @@ -385,7 +387,8 @@ async fn import_dump_v1_rubygems_with_settings() { }, "pagination": { "maxTotalHits": 1000 - } + }, + "searchCutoff": null } "### ); @@ -518,7 +521,8 @@ async fn import_dump_v2_movie_raw() { }, "pagination": { "maxTotalHits": 1000 - } + }, + "searchCutoff": null } "### ); @@ -663,7 +667,8 @@ async fn import_dump_v2_movie_with_settings() { }, "pagination": { "maxTotalHits": 1000 - } + }, + "searchCutoff": null } "### ); @@ -807,7 +812,8 @@ async fn import_dump_v2_rubygems_with_settings() { }, "pagination": { "maxTotalHits": 1000 - } + }, + "searchCutoff": null } "### ); @@ -940,7 +946,8 @@ async fn import_dump_v3_movie_raw() { }, "pagination": { "maxTotalHits": 1000 - } + }, + "searchCutoff": null } "### ); @@ -1085,7 +1092,8 @@ async fn import_dump_v3_movie_with_settings() { }, "pagination": { "maxTotalHits": 1000 - } + }, + "searchCutoff": null } "### ); @@ -1229,7 +1237,8 @@ async fn import_dump_v3_rubygems_with_settings() { }, "pagination": { "maxTotalHits": 1000 - } + }, + "searchCutoff": null } "### ); @@ -1362,7 +1371,8 @@ async fn import_dump_v4_movie_raw() { }, "pagination": { "maxTotalHits": 1000 - } + }, + "searchCutoff": null } "### ); @@ -1507,7 +1517,8 @@ async fn import_dump_v4_movie_with_settings() { }, "pagination": { "maxTotalHits": 1000 - } + }, + "searchCutoff": null } "### ); @@ -1651,7 +1662,8 @@ async fn import_dump_v4_rubygems_with_settings() { }, "pagination": { "maxTotalHits": 1000 - } + }, + "searchCutoff": null } "### ); @@ -1895,7 +1907,8 @@ async fn import_dump_v6_containing_experimental_features() { }, "pagination": { "maxTotalHits": 1000 - } + }, + "searchCutoff": null } "###); diff --git a/meilisearch/tests/settings/get_settings.rs b/meilisearch/tests/settings/get_settings.rs index 5642e854f..000443f36 100644 --- a/meilisearch/tests/settings/get_settings.rs +++ b/meilisearch/tests/settings/get_settings.rs @@ -49,12 +49,12 @@ async fn get_settings_unexisting_index() { async fn get_settings() { let server = Server::new().await; let index = server.index("test"); - index.create(None).await; - index.wait_task(0).await; + let (response, _code) = index.create(None).await; + index.wait_task(response.uid()).await; let (response, code) = index.settings().await; assert_eq!(code, 200); let settings = response.as_object().unwrap(); - assert_eq!(settings.keys().len(), 15); + assert_eq!(settings.keys().len(), 16); assert_eq!(settings["displayedAttributes"], json!(["*"])); assert_eq!(settings["searchableAttributes"], json!(["*"])); assert_eq!(settings["filterableAttributes"], json!([])); @@ -84,6 +84,7 @@ async fn get_settings() { }) ); assert_eq!(settings["proximityPrecision"], json!("byWord")); + assert_eq!(settings["searchCutoff"], json!(null)); } #[actix_rt::test] diff --git a/milli/src/index.rs b/milli/src/index.rs index e79c137e7..d921de9e4 100644 --- a/milli/src/index.rs +++ b/milli/src/index.rs @@ -67,6 +67,7 @@ pub mod main_key { pub const PAGINATION_MAX_TOTAL_HITS: &str = "pagination-max-total-hits"; pub const PROXIMITY_PRECISION: &str = "proximity-precision"; pub const EMBEDDING_CONFIGS: &str = "embedding_configs"; + pub const SEARCH_CUTOFF: &str = "search_cutoff"; } pub mod db_name { @@ -1505,6 +1506,18 @@ impl Index { _ => "default".to_owned(), }) } + + pub(crate) fn put_search_cutoff(&self, wtxn: &mut RwTxn<'_>, cutoff: u64) -> heed::Result<()> { + self.main.remap_types::().put(wtxn, main_key::SEARCH_CUTOFF, &cutoff) + } + + pub fn search_cutoff(&self, rtxn: &RoTxn<'_>) -> Result> { + Ok(self.main.remap_types::().get(rtxn, main_key::SEARCH_CUTOFF)?) + } + + pub(crate) fn delete_search_cutoff(&self, wtxn: &mut RwTxn<'_>) -> heed::Result { + self.main.remap_key_type::().delete(wtxn, main_key::SEARCH_CUTOFF) + } } #[cfg(test)] diff --git a/milli/src/lib.rs b/milli/src/lib.rs index eedd25f7e..896aadb50 100644 --- a/milli/src/lib.rs +++ b/milli/src/lib.rs @@ -121,6 +121,12 @@ impl fmt::Debug for TimeBudget { } } +impl Default for TimeBudget { + fn default() -> Self { + Self::new(std::time::Duration::from_millis(150)) + } +} + impl TimeBudget { pub fn new(budget: std::time::Duration) -> Self { Self { started_at: std::time::Instant::now(), budget } diff --git a/milli/src/update/settings.rs b/milli/src/update/settings.rs index 63b45e3aa..1e720ba56 100644 --- a/milli/src/update/settings.rs +++ b/milli/src/update/settings.rs @@ -150,6 +150,7 @@ pub struct Settings<'a, 't, 'i> { pagination_max_total_hits: Setting, proximity_precision: Setting, embedder_settings: Setting>>, + search_cutoff: Setting, } impl<'a, 't, 'i> Settings<'a, 't, 'i> { @@ -183,6 +184,7 @@ impl<'a, 't, 'i> Settings<'a, 't, 'i> { pagination_max_total_hits: Setting::NotSet, proximity_precision: Setting::NotSet, embedder_settings: Setting::NotSet, + search_cutoff: Setting::NotSet, indexer_config, } } @@ -373,6 +375,14 @@ impl<'a, 't, 'i> Settings<'a, 't, 'i> { self.embedder_settings = Setting::Reset; } + pub fn set_search_cutoff(&mut self, value: u64) { + self.search_cutoff = Setting::Set(value); + } + + pub fn reset_search_cutoff(&mut self) { + self.search_cutoff = Setting::Reset; + } + #[tracing::instrument( level = "trace" skip(self, progress_callback, should_abort, old_fields_ids_map), @@ -1026,6 +1036,24 @@ impl<'a, 't, 'i> Settings<'a, 't, 'i> { Ok(update) } + fn update_search_cutoff(&mut self) -> Result { + let changed = match self.search_cutoff { + Setting::Set(new) => { + let old = self.index.search_cutoff(self.wtxn)?; + if old == Some(new) { + false + } else { + self.index.put_search_cutoff(self.wtxn, new)?; + true + } + } + Setting::Reset => self.index.delete_search_cutoff(self.wtxn)?, + Setting::NotSet => false, + }; + + Ok(changed) + } + pub fn execute(mut self, progress_callback: FP, should_abort: FA) -> Result<()> where FP: Fn(UpdateIndexingStep) + Sync, @@ -1071,6 +1099,9 @@ impl<'a, 't, 'i> Settings<'a, 't, 'i> { // 3. Keep the old vectors but reattempt indexing on a prompt change: only actually changed prompt will need embedding + storage let embedding_configs_updated = self.update_embedding_configs()?; + // never trigger re-indexing + self.update_search_cutoff()?; + if stop_words_updated || non_separator_tokens_updated || separator_tokens_updated @@ -2027,6 +2058,7 @@ mod tests { pagination_max_total_hits, proximity_precision, embedder_settings, + search_cutoff, } = settings; assert!(matches!(searchable_fields, Setting::NotSet)); assert!(matches!(displayed_fields, Setting::NotSet)); @@ -2050,6 +2082,7 @@ mod tests { assert!(matches!(pagination_max_total_hits, Setting::NotSet)); assert!(matches!(proximity_precision, Setting::NotSet)); assert!(matches!(embedder_settings, Setting::NotSet)); + assert!(matches!(search_cutoff, Setting::NotSet)); }) .unwrap(); } From b72495eb5892cf56d3b30b6d575491b1e80f6889 Mon Sep 17 00:00:00 2001 From: Tamo Date: Tue, 12 Mar 2024 18:19:02 +0100 Subject: [PATCH 04/86] fix the settings tests --- meilisearch/src/routes/indexes/settings.rs | 10 ++++++---- meilisearch/tests/settings/get_settings.rs | 4 +++- 2 files changed, 9 insertions(+), 5 deletions(-) diff --git a/meilisearch/src/routes/indexes/settings.rs b/meilisearch/src/routes/indexes/settings.rs index 1d03c9a91..41fc58a87 100644 --- a/meilisearch/src/routes/indexes/settings.rs +++ b/meilisearch/src/routes/indexes/settings.rs @@ -138,6 +138,7 @@ macro_rules! make_setting_route { debug!(returns = ?settings, "Update settings"); let mut json = serde_json::json!(&settings); + dbg!(&json); let val = json[$camelcase_attr].take(); Ok(HttpResponse::Ok().json(val)) @@ -625,14 +626,14 @@ fn embedder_analytics( } make_setting_route!( - "/search_cutoff", - patch, + "/search-cutoff", + put, u64, meilisearch_types::deserr::DeserrJsonError< meilisearch_types::error::deserr_codes::InvalidSettingsSearchCutoff, >, search_cutoff, - "search_cutoff", + "searchCutoff", analytics, |setting: &Option, req: &HttpRequest| { analytics.publish( @@ -673,7 +674,8 @@ generate_configure!( typo_tolerance, pagination, faceting, - embedders + embedders, + search_cutoff ); pub async fn update_all( diff --git a/meilisearch/tests/settings/get_settings.rs b/meilisearch/tests/settings/get_settings.rs index 000443f36..d573f38e0 100644 --- a/meilisearch/tests/settings/get_settings.rs +++ b/meilisearch/tests/settings/get_settings.rs @@ -35,6 +35,7 @@ static DEFAULT_SETTINGS_VALUES: Lazy> = Lazy::new(| "maxTotalHits": json!(1000), }), ); + map.insert("search_cutoff", json!(null)); map }); @@ -286,7 +287,8 @@ test_setting_routes!( ranking_rules put, synonyms put, pagination patch, - faceting patch + faceting patch, + search_cutoff put ); #[actix_rt::test] From b8cda6c300f8ca351a739319b2fcfadcc80e327b Mon Sep 17 00:00:00 2001 From: Tamo Date: Thu, 14 Mar 2024 17:34:46 +0100 Subject: [PATCH 05/86] fix the search cutoff and add a test --- meilisearch/tests/search/mod.rs | 109 +++++++ milli/src/lib.rs | 37 ++- milli/src/score_details.rs | 12 + milli/src/search/hybrid.rs | 2 +- milli/src/search/mod.rs | 4 +- milli/src/search/new/bucket_sort.rs | 16 +- milli/src/search/new/tests/cutoff.rs | 419 +++++++++++++++++++++++++++ milli/src/search/new/tests/mod.rs | 1 + milli/tests/search/mod.rs | 45 +-- 9 files changed, 590 insertions(+), 55 deletions(-) create mode 100644 milli/src/search/new/tests/cutoff.rs diff --git a/meilisearch/tests/search/mod.rs b/meilisearch/tests/search/mod.rs index 90098c5b6..62dd73c63 100644 --- a/meilisearch/tests/search/mod.rs +++ b/meilisearch/tests/search/mod.rs @@ -834,6 +834,115 @@ async fn test_score_details() { .await; } +#[actix_rt::test] +async fn test_degraded_score_details() { + let server = Server::new().await; + let index = server.index("test"); + + let documents = NESTED_DOCUMENTS.clone(); + + index.add_documents(json!(documents), None).await; + // We can't really use anything else than 0ms here; otherwise, the test will get flaky. + let (res, _code) = index.update_settings(json!({ "searchCutoff": 0 })).await; + index.wait_task(res.uid()).await; + + index + .search( + json!({ + "q": "b", + "showRankingScoreDetails": true, + }), + |response, code| { + meili_snap::snapshot!(code, @"200 OK"); + meili_snap::snapshot!(meili_snap::json_string!(response["hits"]), @r###" + [ + { + "id": 852, + "father": "jean", + "mother": "michelle", + "doggos": [ + { + "name": "bobby", + "age": 2 + }, + { + "name": "buddy", + "age": 4 + } + ], + "cattos": "pésti", + "_vectors": { + "manual": [ + 1, + 2, + 3 + ] + }, + "_rankingScoreDetails": { + "skipped": 0.0 + } + }, + { + "id": 654, + "father": "pierre", + "mother": "sabine", + "doggos": [ + { + "name": "gros bill", + "age": 8 + } + ], + "cattos": [ + "simba", + "pestiféré" + ], + "_vectors": { + "manual": [ + 1, + 2, + 54 + ] + }, + "_rankingScoreDetails": { + "skipped": 0.0 + } + }, + { + "id": 951, + "father": "jean-baptiste", + "mother": "sophie", + "doggos": [ + { + "name": "turbo", + "age": 5 + }, + { + "name": "fast", + "age": 6 + } + ], + "cattos": [ + "moumoute", + "gomez" + ], + "_vectors": { + "manual": [ + 10, + 23, + 32 + ] + }, + "_rankingScoreDetails": { + "skipped": 0.0 + } + } + ] + "###); + }, + ) + .await; +} + #[actix_rt::test] async fn experimental_feature_vector_store() { let server = Server::new().await; diff --git a/milli/src/lib.rs b/milli/src/lib.rs index 896aadb50..df44ca127 100644 --- a/milli/src/lib.rs +++ b/milli/src/lib.rs @@ -105,10 +105,15 @@ pub const MAX_WORD_LENGTH: usize = MAX_LMDB_KEY_LENGTH / 2; pub const MAX_POSITION_PER_ATTRIBUTE: u32 = u16::MAX as u32 + 1; -#[derive(Clone, Copy)] +#[derive(Clone)] pub struct TimeBudget { started_at: std::time::Instant, budget: std::time::Duration, + + /// When testing the time budget, ensuring we did more than iteration of the bucket sort can be useful. + /// But to avoid being flaky, the only option is to add the ability to stop after a specific number of calls instead of a `Duration`. + #[cfg(test)] + stop_after: Option<(std::sync::Arc, usize)>, } impl fmt::Debug for TimeBudget { @@ -129,18 +134,40 @@ impl Default for TimeBudget { impl TimeBudget { pub fn new(budget: std::time::Duration) -> Self { - Self { started_at: std::time::Instant::now(), budget } + Self { + started_at: std::time::Instant::now(), + budget, + + #[cfg(test)] + stop_after: None, + } } pub fn max() -> Self { Self::new(std::time::Duration::from_secs(u64::MAX)) } - pub fn exceeded(&self) -> bool { - self.must_stop() + #[cfg(test)] + pub fn with_stop_after(mut self, stop_after: usize) -> Self { + use std::sync::atomic::AtomicUsize; + use std::sync::Arc; + + self.stop_after = Some((Arc::new(AtomicUsize::new(0)), stop_after)); + self } - pub fn must_stop(&self) -> bool { + pub fn exceeded(&self) -> bool { + #[cfg(test)] + if let Some((current, stop_after)) = &self.stop_after { + let current = current.fetch_add(1, std::sync::atomic::Ordering::Relaxed); + if current >= *stop_after { + return true; + } else { + // if a number has been specified then we ignore entirely the time budget + return false; + } + } + self.started_at.elapsed() > self.budget } } diff --git a/milli/src/score_details.rs b/milli/src/score_details.rs index f6b9db58c..f2c6fb58a 100644 --- a/milli/src/score_details.rs +++ b/milli/src/score_details.rs @@ -17,6 +17,9 @@ pub enum ScoreDetails { Sort(Sort), Vector(Vector), GeoSort(GeoSort), + + /// Returned when we don't have the time to finish applying all the subsequent ranking-rules + Skipped, } #[derive(Clone, Copy)] @@ -50,6 +53,7 @@ impl ScoreDetails { ScoreDetails::Sort(_) => None, ScoreDetails::GeoSort(_) => None, ScoreDetails::Vector(_) => None, + ScoreDetails::Skipped => Some(Rank { rank: 0, max_rank: 1 }), } } @@ -97,6 +101,7 @@ impl ScoreDetails { ScoreDetails::Vector(vector) => RankOrValue::Score( vector.value_similarity.as_ref().map(|(_, s)| *s as f64).unwrap_or(0.0f64), ), + ScoreDetails::Skipped => RankOrValue::Score(0.), } } @@ -256,6 +261,13 @@ impl ScoreDetails { details_map.insert(vector, details); order += 1; } + ScoreDetails::Skipped => { + details_map.insert( + "skipped".to_string(), + serde_json::Number::from_f64(0.).unwrap().into(), + ); + order += 1; + } } } details_map diff --git a/milli/src/search/hybrid.rs b/milli/src/search/hybrid.rs index 9d8b3860d..1e14f4430 100644 --- a/milli/src/search/hybrid.rs +++ b/milli/src/search/hybrid.rs @@ -132,7 +132,7 @@ impl<'a> Search<'a> { index: self.index, distribution_shift: self.distribution_shift, embedder_name: self.embedder_name.clone(), - time_budget: self.time_budget, + time_budget: self.time_budget.clone(), }; let vector_query = search.vector.take(); diff --git a/milli/src/search/mod.rs b/milli/src/search/mod.rs index b14d88d03..f6ab8a7de 100644 --- a/milli/src/search/mod.rs +++ b/milli/src/search/mod.rs @@ -195,7 +195,7 @@ impl<'a> Search<'a> { self.limit, self.distribution_shift, embedder_name, - self.time_budget, + self.time_budget.clone(), )?, None => execute_search( &mut ctx, @@ -211,7 +211,7 @@ impl<'a> Search<'a> { Some(self.words_limit), &mut DefaultSearchLogger, &mut DefaultSearchLogger, - self.time_budget, + self.time_budget.clone(), )?, }; diff --git a/milli/src/search/new/bucket_sort.rs b/milli/src/search/new/bucket_sort.rs index 7fc830c1f..521fcb983 100644 --- a/milli/src/search/new/bucket_sort.rs +++ b/milli/src/search/new/bucket_sort.rs @@ -161,11 +161,21 @@ pub fn bucket_sort<'ctx, Q: RankingRuleQueryTrait>( while valid_docids.len() < length { if time_budget.exceeded() { - let bucket = std::mem::take(&mut ranking_rule_universes[cur_ranking_rule_index]); - maybe_add_to_results!(bucket); + loop { + let bucket = std::mem::take(&mut ranking_rule_universes[cur_ranking_rule_index]); + ranking_rule_scores.push(ScoreDetails::Skipped); + maybe_add_to_results!(bucket); + ranking_rule_scores.pop(); + + if cur_ranking_rule_index == 0 { + break; + } + + back!(); + } return Ok(BucketSortOutput { - scores: vec![Default::default(); valid_docids.len()], + scores: valid_scores, docids: valid_docids, all_candidates, degraded: true, diff --git a/milli/src/search/new/tests/cutoff.rs b/milli/src/search/new/tests/cutoff.rs new file mode 100644 index 000000000..4256abc2b --- /dev/null +++ b/milli/src/search/new/tests/cutoff.rs @@ -0,0 +1,419 @@ +//! This module test the search cutoff and ensure a few things: +//! 1. A basic test works and mark the search as degraded +//! 2. A test that ensure the filters are affectively applied even with a cutoff of 0 +//! 3. A test that ensure the cutoff works well with the ranking scores + +use std::time::Duration; + +use big_s::S; +use maplit::hashset; +use meili_snap::snapshot; + +use crate::index::tests::TempIndex; +use crate::{Criterion, Filter, Search, TimeBudget}; + +fn create_index() -> TempIndex { + let index = TempIndex::new(); + + index + .update_settings(|s| { + s.set_primary_key("id".to_owned()); + s.set_searchable_fields(vec!["text".to_owned()]); + s.set_filterable_fields(hashset! { S("id") }); + s.set_criteria(vec![Criterion::Words, Criterion::Typo]); + }) + .unwrap(); + + // reverse the ID / insertion order so we see better what was sorted from what got the insertion order ordering + index + .add_documents(documents!([ + { + "id": 4, + "text": "hella puppo kefir", + }, + { + "id": 3, + "text": "hella puppy kefir", + }, + { + "id": 2, + "text": "hello", + }, + { + "id": 1, + "text": "hello puppy", + }, + { + "id": 0, + "text": "hello puppy kefir", + }, + ])) + .unwrap(); + index +} + +#[test] +fn basic_degraded_search() { + let index = create_index(); + let rtxn = index.read_txn().unwrap(); + + let mut search = Search::new(&rtxn, &index); + search.query("hello puppy kefir"); + search.limit(3); + search.time_budget(TimeBudget::new(Duration::from_millis(0))); + + let result = search.execute().unwrap(); + assert!(result.degraded); +} + +#[test] +fn degraded_search_cannot_skip_filter() { + let index = create_index(); + let rtxn = index.read_txn().unwrap(); + + let mut search = Search::new(&rtxn, &index); + search.query("hello puppy kefir"); + search.limit(100); + search.time_budget(TimeBudget::new(Duration::from_millis(0))); + let filter_condition = Filter::from_str("id > 2").unwrap().unwrap(); + search.filter(filter_condition); + + let result = search.execute().unwrap(); + assert!(result.degraded); + snapshot!(format!("{:?}\n{:?}", result.candidates, result.documents_ids), @r###" + RoaringBitmap<[0, 1]> + [0, 1] + "###); +} + +#[test] +fn degraded_search_and_score_details() { + let index = create_index(); + let rtxn = index.read_txn().unwrap(); + + let mut search = Search::new(&rtxn, &index); + search.query("hello puppy kefir"); + search.limit(4); + search.time_budget(TimeBudget::max()); + + let result = search.execute().unwrap(); + snapshot!(format!("{:#?}\n{:#?}", result.documents_ids, result.document_scores), @r###" + [ + 4, + 1, + 0, + 3, + ] + [ + [ + Words( + Words { + matching_words: 3, + max_matching_words: 3, + }, + ), + Typo( + Typo { + typo_count: 0, + max_typo_count: 3, + }, + ), + ], + [ + Words( + Words { + matching_words: 3, + max_matching_words: 3, + }, + ), + Typo( + Typo { + typo_count: 1, + max_typo_count: 3, + }, + ), + ], + [ + Words( + Words { + matching_words: 3, + max_matching_words: 3, + }, + ), + ], + [ + Words( + Words { + matching_words: 2, + max_matching_words: 3, + }, + ), + ], + ] + "###); + + // Do ONE loop iteration. Not much can be deduced, almost everyone matched the words first bucket. + search.time_budget(TimeBudget::max().with_stop_after(1)); + + let result = search.execute().unwrap(); + snapshot!(format!("{:#?}\n{:#?}", result.documents_ids, result.document_scores), @r###" + [ + 0, + 1, + 4, + 2, + ] + [ + [ + Words( + Words { + matching_words: 3, + max_matching_words: 3, + }, + ), + Skipped, + ], + [ + Words( + Words { + matching_words: 3, + max_matching_words: 3, + }, + ), + Skipped, + ], + [ + Words( + Words { + matching_words: 3, + max_matching_words: 3, + }, + ), + Skipped, + ], + [ + Skipped, + ], + ] + "###); + + // Do TWO loop iterations. The first document should be entirely sorted + search.time_budget(TimeBudget::max().with_stop_after(2)); + + let result = search.execute().unwrap(); + snapshot!(format!("{:#?}\n{:#?}", result.documents_ids, result.document_scores), @r###" + [ + 4, + 0, + 1, + 2, + ] + [ + [ + Words( + Words { + matching_words: 3, + max_matching_words: 3, + }, + ), + Typo( + Typo { + typo_count: 0, + max_typo_count: 3, + }, + ), + ], + [ + Words( + Words { + matching_words: 3, + max_matching_words: 3, + }, + ), + Skipped, + ], + [ + Words( + Words { + matching_words: 3, + max_matching_words: 3, + }, + ), + Skipped, + ], + [ + Skipped, + ], + ] + "###); + + // Do THREE loop iterations. The second document should be entirely sorted as well + search.time_budget(TimeBudget::max().with_stop_after(3)); + + let result = search.execute().unwrap(); + snapshot!(format!("{:#?}\n{:#?}", result.documents_ids, result.document_scores), @r###" + [ + 4, + 1, + 0, + 2, + ] + [ + [ + Words( + Words { + matching_words: 3, + max_matching_words: 3, + }, + ), + Typo( + Typo { + typo_count: 0, + max_typo_count: 3, + }, + ), + ], + [ + Words( + Words { + matching_words: 3, + max_matching_words: 3, + }, + ), + Typo( + Typo { + typo_count: 1, + max_typo_count: 3, + }, + ), + ], + [ + Words( + Words { + matching_words: 3, + max_matching_words: 3, + }, + ), + Skipped, + ], + [ + Skipped, + ], + ] + "###); + + // Do FOUR loop iterations. The third document should be entirely sorted as well + // The words bucket have still not progressed thus the last document doesn't have any info yet. + search.time_budget(TimeBudget::max().with_stop_after(4)); + + let result = search.execute().unwrap(); + snapshot!(format!("{:#?}\n{:#?}", result.documents_ids, result.document_scores), @r###" + [ + 4, + 1, + 0, + 2, + ] + [ + [ + Words( + Words { + matching_words: 3, + max_matching_words: 3, + }, + ), + Typo( + Typo { + typo_count: 0, + max_typo_count: 3, + }, + ), + ], + [ + Words( + Words { + matching_words: 3, + max_matching_words: 3, + }, + ), + Typo( + Typo { + typo_count: 1, + max_typo_count: 3, + }, + ), + ], + [ + Words( + Words { + matching_words: 3, + max_matching_words: 3, + }, + ), + ], + [ + Skipped, + ], + ] + "###); + + // After FIVE loop iteration. The words ranking rule gave us a new bucket. + // Since we reached the limit we were able to early exit without checking the typo ranking rule. + search.time_budget(TimeBudget::max().with_stop_after(5)); + + let result = search.execute().unwrap(); + snapshot!(format!("{:#?}\n{:#?}", result.documents_ids, result.document_scores), @r###" + [ + 4, + 1, + 0, + 3, + ] + [ + [ + Words( + Words { + matching_words: 3, + max_matching_words: 3, + }, + ), + Typo( + Typo { + typo_count: 0, + max_typo_count: 3, + }, + ), + ], + [ + Words( + Words { + matching_words: 3, + max_matching_words: 3, + }, + ), + Typo( + Typo { + typo_count: 1, + max_typo_count: 3, + }, + ), + ], + [ + Words( + Words { + matching_words: 3, + max_matching_words: 3, + }, + ), + ], + [ + Words( + Words { + matching_words: 2, + max_matching_words: 3, + }, + ), + ], + ] + "###); +} diff --git a/milli/src/search/new/tests/mod.rs b/milli/src/search/new/tests/mod.rs index e500d16fb..26199b79b 100644 --- a/milli/src/search/new/tests/mod.rs +++ b/milli/src/search/new/tests/mod.rs @@ -1,5 +1,6 @@ pub mod attribute_fid; pub mod attribute_position; +pub mod cutoff; pub mod distinct; pub mod exactness; pub mod geo_sort; diff --git a/milli/tests/search/mod.rs b/milli/tests/search/mod.rs index ab6befa60..9193ab762 100644 --- a/milli/tests/search/mod.rs +++ b/milli/tests/search/mod.rs @@ -1,19 +1,14 @@ use std::cmp::Reverse; use std::collections::HashSet; use std::io::Cursor; -use std::time::Duration; use big_s::S; use either::{Either, Left, Right}; use heed::EnvOpenOptions; use maplit::{btreemap, hashset}; -use meili_snap::snapshot; use milli::documents::{DocumentsBatchBuilder, DocumentsBatchReader}; use milli::update::{IndexDocuments, IndexDocumentsConfig, IndexerConfig, Settings}; -use milli::{ - AscDesc, Criterion, DocumentId, Filter, Index, Member, Object, Search, TermsMatchingStrategy, - TimeBudget, -}; +use milli::{AscDesc, Criterion, DocumentId, Index, Member, Object, TermsMatchingStrategy}; use serde::{Deserialize, Deserializer}; use slice_group_by::GroupBy; @@ -354,41 +349,3 @@ where let result = serde_json::Value::deserialize(deserializer)?; Ok(Some(result)) } - -#[test] -fn basic_degraded_search() { - use Criterion::*; - let criteria = vec![Words, Typo, Proximity, Attribute, Exactness]; - let index = setup_search_index_with_criteria(&criteria); - let rtxn = index.read_txn().unwrap(); - - let mut search = Search::new(&rtxn, &index); - search.query(TEST_QUERY); - search.limit(EXTERNAL_DOCUMENTS_IDS.len()); - search.time_budget(TimeBudget::new(Duration::from_millis(0))); - - let result = search.execute().unwrap(); - assert!(result.degraded); -} - -#[test] -fn degraded_search_cannot_skip_filter() { - use Criterion::*; - let criteria = vec![Words, Typo, Proximity, Attribute, Exactness]; - let index = setup_search_index_with_criteria(&criteria); - let rtxn = index.read_txn().unwrap(); - - let mut search = Search::new(&rtxn, &index); - search.query(TEST_QUERY); - search.limit(EXTERNAL_DOCUMENTS_IDS.len()); - search.time_budget(TimeBudget::new(Duration::from_millis(0))); - let filter_condition = Filter::from_str("tag = etiopia").unwrap().unwrap(); - search.filter(filter_condition); - - let result = search.execute().unwrap(); - assert!(result.degraded); - snapshot!(format!("{:?}\n{:?}", result.candidates, result.documents_ids), @r###" - RoaringBitmap<[0, 2, 5, 8, 11, 14]> - [0, 2, 5, 8, 11, 14] - "###); -} From ad9192fbbf38a29413b20cdf7678a522673b8ad7 Mon Sep 17 00:00:00 2001 From: Tamo Date: Thu, 14 Mar 2024 17:42:33 +0100 Subject: [PATCH 06/86] reduce the size of an integration test --- meilisearch/tests/search/mod.rs | 46 +++++---------------------------- 1 file changed, 6 insertions(+), 40 deletions(-) diff --git a/meilisearch/tests/search/mod.rs b/meilisearch/tests/search/mod.rs index 62dd73c63..8c947a329 100644 --- a/meilisearch/tests/search/mod.rs +++ b/meilisearch/tests/search/mod.rs @@ -850,6 +850,7 @@ async fn test_degraded_score_details() { .search( json!({ "q": "b", + "attributesToRetrieve": ["doggos.name", "cattos"], "showRankingScoreDetails": true, }), |response, code| { @@ -857,81 +858,46 @@ async fn test_degraded_score_details() { meili_snap::snapshot!(meili_snap::json_string!(response["hits"]), @r###" [ { - "id": 852, - "father": "jean", - "mother": "michelle", "doggos": [ { - "name": "bobby", - "age": 2 + "name": "bobby" }, { - "name": "buddy", - "age": 4 + "name": "buddy" } ], "cattos": "pésti", - "_vectors": { - "manual": [ - 1, - 2, - 3 - ] - }, "_rankingScoreDetails": { "skipped": 0.0 } }, { - "id": 654, - "father": "pierre", - "mother": "sabine", "doggos": [ { - "name": "gros bill", - "age": 8 + "name": "gros bill" } ], "cattos": [ "simba", "pestiféré" ], - "_vectors": { - "manual": [ - 1, - 2, - 54 - ] - }, "_rankingScoreDetails": { "skipped": 0.0 } }, { - "id": 951, - "father": "jean-baptiste", - "mother": "sophie", "doggos": [ { - "name": "turbo", - "age": 5 + "name": "turbo" }, { - "name": "fast", - "age": 6 + "name": "fast" } ], "cattos": [ "moumoute", "gomez" ], - "_vectors": { - "manual": [ - 10, - 23, - 32 - ] - }, "_rankingScoreDetails": { "skipped": 0.0 } From 038c26c118ef041e5843f29b5f4862c879e39979 Mon Sep 17 00:00:00 2001 From: Tamo Date: Thu, 14 Mar 2024 17:52:08 +0100 Subject: [PATCH 07/86] stop returning the degraded boolean when a search was cutoff --- meilisearch/src/search.rs | 3 +- meilisearch/tests/search/mod.rs | 95 ++++++++++++++++++--------------- 2 files changed, 53 insertions(+), 45 deletions(-) diff --git a/meilisearch/src/search.rs b/meilisearch/src/search.rs index f83e14187..0333eb0d5 100644 --- a/meilisearch/src/search.rs +++ b/meilisearch/src/search.rs @@ -324,7 +324,8 @@ pub struct SearchResult { #[serde(skip_serializing_if = "Option::is_none")] pub facet_stats: Option>, - #[serde(skip_serializing_if = "std::ops::Not::not")] + // This information is only used for analytics purposes + #[serde(skip)] pub degraded: bool, } diff --git a/meilisearch/tests/search/mod.rs b/meilisearch/tests/search/mod.rs index 8c947a329..3e5c4278a 100644 --- a/meilisearch/tests/search/mod.rs +++ b/meilisearch/tests/search/mod.rs @@ -855,54 +855,61 @@ async fn test_degraded_score_details() { }), |response, code| { meili_snap::snapshot!(code, @"200 OK"); - meili_snap::snapshot!(meili_snap::json_string!(response["hits"]), @r###" - [ - { - "doggos": [ - { - "name": "bobby" - }, - { - "name": "buddy" + meili_snap::snapshot!(meili_snap::json_string!(response), @r###" + { + "hits": [ + { + "doggos": [ + { + "name": "bobby" + }, + { + "name": "buddy" + } + ], + "cattos": "pésti", + "_rankingScoreDetails": { + "skipped": 0.0 } - ], - "cattos": "pésti", - "_rankingScoreDetails": { - "skipped": 0.0 - } - }, - { - "doggos": [ - { - "name": "gros bill" + }, + { + "doggos": [ + { + "name": "gros bill" + } + ], + "cattos": [ + "simba", + "pestiféré" + ], + "_rankingScoreDetails": { + "skipped": 0.0 } - ], - "cattos": [ - "simba", - "pestiféré" - ], - "_rankingScoreDetails": { - "skipped": 0.0 - } - }, - { - "doggos": [ - { - "name": "turbo" - }, - { - "name": "fast" + }, + { + "doggos": [ + { + "name": "turbo" + }, + { + "name": "fast" + } + ], + "cattos": [ + "moumoute", + "gomez" + ], + "_rankingScoreDetails": { + "skipped": 0.0 } - ], - "cattos": [ - "moumoute", - "gomez" - ], - "_rankingScoreDetails": { - "skipped": 0.0 } - } - ] + ], + "query": "b", + "processingTimeMs": 0, + "limit": 20, + "offset": 0, + "estimatedTotalHits": 3 + } "###); }, ) From 6a0c399c2f827a46600deda0b0d3695d1ed6af19 Mon Sep 17 00:00:00 2001 From: Tamo Date: Mon, 18 Mar 2024 12:06:00 +0100 Subject: [PATCH 08/86] rename the search_cutoff parameter to search_cutoff_ms --- dump/src/lib.rs | 2 +- dump/src/reader/compat/v5_to_v6.rs | 2 +- meilisearch-types/src/settings.rs | 22 +++++++++--------- meilisearch/src/routes/indexes/settings.rs | 12 +++++----- meilisearch/tests/common/mod.rs | 1 + meilisearch/tests/dumps/mod.rs | 26 +++++++++++----------- meilisearch/tests/search/mod.rs | 6 ++--- meilisearch/tests/settings/get_settings.rs | 6 ++--- 8 files changed, 39 insertions(+), 38 deletions(-) diff --git a/dump/src/lib.rs b/dump/src/lib.rs index e7cadacbe..a7af2d5d0 100644 --- a/dump/src/lib.rs +++ b/dump/src/lib.rs @@ -277,7 +277,7 @@ pub(crate) mod test { }), pagination: Setting::NotSet, embedders: Setting::NotSet, - search_cutoff: Setting::NotSet, + search_cutoff_ms: Setting::NotSet, _kind: std::marker::PhantomData, }; settings.check() diff --git a/dump/src/reader/compat/v5_to_v6.rs b/dump/src/reader/compat/v5_to_v6.rs index 2b8997847..a883f0ba0 100644 --- a/dump/src/reader/compat/v5_to_v6.rs +++ b/dump/src/reader/compat/v5_to_v6.rs @@ -379,7 +379,7 @@ impl From> for v6::Settings { v5::Setting::NotSet => v6::Setting::NotSet, }, embedders: v6::Setting::NotSet, - search_cutoff: v6::Setting::NotSet, + search_cutoff_ms: v6::Setting::NotSet, _kind: std::marker::PhantomData, } } diff --git a/meilisearch-types/src/settings.rs b/meilisearch-types/src/settings.rs index d05201943..23fe98347 100644 --- a/meilisearch-types/src/settings.rs +++ b/meilisearch-types/src/settings.rs @@ -204,7 +204,7 @@ pub struct Settings { pub embedders: Setting>>, #[serde(default, skip_serializing_if = "Setting::is_not_set")] #[deserr(default, error = DeserrJsonError)] - pub search_cutoff: Setting, + pub search_cutoff_ms: Setting, #[serde(skip)] #[deserr(skip)] @@ -230,7 +230,7 @@ impl Settings { faceting: Setting::Reset, pagination: Setting::Reset, embedders: Setting::Reset, - search_cutoff: Setting::Reset, + search_cutoff_ms: Setting::Reset, _kind: PhantomData, } } @@ -253,7 +253,7 @@ impl Settings { faceting, pagination, embedders, - search_cutoff, + search_cutoff_ms, .. } = self; @@ -274,7 +274,7 @@ impl Settings { faceting, pagination, embedders, - search_cutoff, + search_cutoff_ms, _kind: PhantomData, } } @@ -321,7 +321,7 @@ impl Settings { faceting: self.faceting, pagination: self.pagination, embedders: self.embedders, - search_cutoff: self.search_cutoff, + search_cutoff_ms: self.search_cutoff_ms, _kind: PhantomData, } } @@ -371,7 +371,7 @@ pub fn apply_settings_to_builder( faceting, pagination, embedders, - search_cutoff, + search_cutoff_ms, _kind, } = settings; @@ -548,7 +548,7 @@ pub fn apply_settings_to_builder( Setting::NotSet => (), } - match search_cutoff { + match search_cutoff_ms { Setting::Set(cutoff) => builder.set_search_cutoff(*cutoff), Setting::Reset => builder.reset_search_cutoff(), Setting::NotSet => (), @@ -641,7 +641,7 @@ pub fn settings( .collect(); let embedders = if embedders.is_empty() { Setting::NotSet } else { Setting::Set(embedders) }; - let search_cutoff = index.search_cutoff(rtxn)?; + let search_cutoff_ms = index.search_cutoff(rtxn)?; Ok(Settings { displayed_attributes: match displayed_attributes { @@ -669,7 +669,7 @@ pub fn settings( faceting: Setting::Set(faceting), pagination: Setting::Set(pagination), embedders, - search_cutoff: match search_cutoff { + search_cutoff_ms: match search_cutoff_ms { Some(cutoff) => Setting::Set(cutoff), None => Setting::Reset, }, @@ -823,7 +823,7 @@ pub(crate) mod test { faceting: Setting::NotSet, pagination: Setting::NotSet, embedders: Setting::NotSet, - search_cutoff: Setting::NotSet, + search_cutoff_ms: Setting::NotSet, _kind: PhantomData::, }; @@ -850,7 +850,7 @@ pub(crate) mod test { faceting: Setting::NotSet, pagination: Setting::NotSet, embedders: Setting::NotSet, - search_cutoff: Setting::NotSet, + search_cutoff_ms: Setting::NotSet, _kind: PhantomData::, }; diff --git a/meilisearch/src/routes/indexes/settings.rs b/meilisearch/src/routes/indexes/settings.rs index 41fc58a87..4c03eb1a1 100644 --- a/meilisearch/src/routes/indexes/settings.rs +++ b/meilisearch/src/routes/indexes/settings.rs @@ -626,19 +626,19 @@ fn embedder_analytics( } make_setting_route!( - "/search-cutoff", + "/search-cutoff-ms", put, u64, meilisearch_types::deserr::DeserrJsonError< meilisearch_types::error::deserr_codes::InvalidSettingsSearchCutoff, >, - search_cutoff, - "searchCutoff", + search_cutoff_ms, + "searchCutoffMs", analytics, |setting: &Option, req: &HttpRequest| { analytics.publish( "Search Cutoff Updated".to_string(), - serde_json::json!({"search_cutoff": setting }), + serde_json::json!({"search_cutoff_ms": setting }), Some(req), ); } @@ -675,7 +675,7 @@ generate_configure!( pagination, faceting, embedders, - search_cutoff + search_cutoff_ms ); pub async fn update_all( @@ -787,7 +787,7 @@ pub async fn update_all( "total": new_settings.synonyms.as_ref().set().map(|synonyms| synonyms.len()), }, "embedders": crate::routes::indexes::settings::embedder_analytics(new_settings.embedders.as_ref().set()), - "search_cutoff": new_settings.search_cutoff.as_ref().set(), + "search_cutoff_ms": new_settings.search_cutoff_ms.as_ref().set(), }), Some(&req), ); diff --git a/meilisearch/tests/common/mod.rs b/meilisearch/tests/common/mod.rs index 2b9e5e1d7..3117dd185 100644 --- a/meilisearch/tests/common/mod.rs +++ b/meilisearch/tests/common/mod.rs @@ -16,6 +16,7 @@ pub use server::{default_settings, Server}; pub struct Value(pub serde_json::Value); impl Value { + #[track_caller] pub fn uid(&self) -> u64 { if let Some(uid) = self["uid"].as_u64() { uid diff --git a/meilisearch/tests/dumps/mod.rs b/meilisearch/tests/dumps/mod.rs index 7bf97f8b2..1a31437f8 100644 --- a/meilisearch/tests/dumps/mod.rs +++ b/meilisearch/tests/dumps/mod.rs @@ -78,7 +78,7 @@ async fn import_dump_v1_movie_raw() { "pagination": { "maxTotalHits": 1000 }, - "searchCutoff": null + "searchCutoffMs": null } "### ); @@ -240,7 +240,7 @@ async fn import_dump_v1_movie_with_settings() { "pagination": { "maxTotalHits": 1000 }, - "searchCutoff": null + "searchCutoffMs": null } "### ); @@ -388,7 +388,7 @@ async fn import_dump_v1_rubygems_with_settings() { "pagination": { "maxTotalHits": 1000 }, - "searchCutoff": null + "searchCutoffMs": null } "### ); @@ -522,7 +522,7 @@ async fn import_dump_v2_movie_raw() { "pagination": { "maxTotalHits": 1000 }, - "searchCutoff": null + "searchCutoffMs": null } "### ); @@ -668,7 +668,7 @@ async fn import_dump_v2_movie_with_settings() { "pagination": { "maxTotalHits": 1000 }, - "searchCutoff": null + "searchCutoffMs": null } "### ); @@ -813,7 +813,7 @@ async fn import_dump_v2_rubygems_with_settings() { "pagination": { "maxTotalHits": 1000 }, - "searchCutoff": null + "searchCutoffMs": null } "### ); @@ -947,7 +947,7 @@ async fn import_dump_v3_movie_raw() { "pagination": { "maxTotalHits": 1000 }, - "searchCutoff": null + "searchCutoffMs": null } "### ); @@ -1093,7 +1093,7 @@ async fn import_dump_v3_movie_with_settings() { "pagination": { "maxTotalHits": 1000 }, - "searchCutoff": null + "searchCutoffMs": null } "### ); @@ -1238,7 +1238,7 @@ async fn import_dump_v3_rubygems_with_settings() { "pagination": { "maxTotalHits": 1000 }, - "searchCutoff": null + "searchCutoffMs": null } "### ); @@ -1372,7 +1372,7 @@ async fn import_dump_v4_movie_raw() { "pagination": { "maxTotalHits": 1000 }, - "searchCutoff": null + "searchCutoffMs": null } "### ); @@ -1518,7 +1518,7 @@ async fn import_dump_v4_movie_with_settings() { "pagination": { "maxTotalHits": 1000 }, - "searchCutoff": null + "searchCutoffMs": null } "### ); @@ -1663,7 +1663,7 @@ async fn import_dump_v4_rubygems_with_settings() { "pagination": { "maxTotalHits": 1000 }, - "searchCutoff": null + "searchCutoffMs": null } "### ); @@ -1908,7 +1908,7 @@ async fn import_dump_v6_containing_experimental_features() { "pagination": { "maxTotalHits": 1000 }, - "searchCutoff": null + "searchCutoffMs": null } "###); diff --git a/meilisearch/tests/search/mod.rs b/meilisearch/tests/search/mod.rs index 3e5c4278a..971539a31 100644 --- a/meilisearch/tests/search/mod.rs +++ b/meilisearch/tests/search/mod.rs @@ -843,7 +843,7 @@ async fn test_degraded_score_details() { index.add_documents(json!(documents), None).await; // We can't really use anything else than 0ms here; otherwise, the test will get flaky. - let (res, _code) = index.update_settings(json!({ "searchCutoff": 0 })).await; + let (res, _code) = index.update_settings(json!({ "searchCutoffMs": 0 })).await; index.wait_task(res.uid()).await; index @@ -855,7 +855,7 @@ async fn test_degraded_score_details() { }), |response, code| { meili_snap::snapshot!(code, @"200 OK"); - meili_snap::snapshot!(meili_snap::json_string!(response), @r###" + meili_snap::snapshot!(meili_snap::json_string!(response, { ".processingTimeMs" => "[duration]" }), @r###" { "hits": [ { @@ -905,7 +905,7 @@ async fn test_degraded_score_details() { } ], "query": "b", - "processingTimeMs": 0, + "processingTimeMs": "[duration]", "limit": 20, "offset": 0, "estimatedTotalHits": 3 diff --git a/meilisearch/tests/settings/get_settings.rs b/meilisearch/tests/settings/get_settings.rs index d573f38e0..09e38e55a 100644 --- a/meilisearch/tests/settings/get_settings.rs +++ b/meilisearch/tests/settings/get_settings.rs @@ -35,7 +35,7 @@ static DEFAULT_SETTINGS_VALUES: Lazy> = Lazy::new(| "maxTotalHits": json!(1000), }), ); - map.insert("search_cutoff", json!(null)); + map.insert("search_cutoff_ms", json!(null)); map }); @@ -85,7 +85,7 @@ async fn get_settings() { }) ); assert_eq!(settings["proximityPrecision"], json!("byWord")); - assert_eq!(settings["searchCutoff"], json!(null)); + assert_eq!(settings["searchCutoffMs"], json!(null)); } #[actix_rt::test] @@ -288,7 +288,7 @@ test_setting_routes!( synonyms put, pagination patch, faceting patch, - search_cutoff put + search_cutoff_ms put ); #[actix_rt::test] From 7bd881b9bcdeb2e84b9ec4870584d11efa580897 Mon Sep 17 00:00:00 2001 From: Tamo Date: Mon, 18 Mar 2024 18:39:05 +0100 Subject: [PATCH 09/86] adds the degraded searches to the prometheus dashboard --- assets/grafana-dashboard.json | 64 ++++++++++++++++++++++++ meilisearch/src/metrics.rs | 5 ++ meilisearch/src/routes/indexes/search.rs | 4 ++ 3 files changed, 73 insertions(+) diff --git a/assets/grafana-dashboard.json b/assets/grafana-dashboard.json index 37f7b1ca2..74a456b97 100644 --- a/assets/grafana-dashboard.json +++ b/assets/grafana-dashboard.json @@ -238,6 +238,70 @@ "title": "Total Searches (1h)", "type": "gauge" }, + { + "datasource": { + "type": "prometheus" + }, + "fieldConfig": { + "defaults": { + "color": { + "mode": "thresholds" + }, + "mappings": [], + "thresholds": { + "mode": "absolute", + "steps": [ + { + "color": "green", + "value": null + }, + { + "color": "red", + "value": 80 + } + ] + } + }, + "overrides": [] + }, + "gridPos": { + "h": 6, + "w": 4, + "x": 8, + "y": 1 + }, + "id": 26, + "options": { + "orientation": "auto", + "reduceOptions": { + "calcs": [ + "lastNotNull" + ], + "fields": "", + "values": false + }, + "showThresholdLabels": false, + "showThresholdMarkers": true, + "text": {} + }, + "pluginVersion": "9.5.2", + "targets": [ + { + "datasource": { + "type": "prometheus" + }, + "editorMode": "builder", + "exemplar": true, + "expr": "round(increase(meilisearch_degraded_search_requests{job=\"$job\"}[1h]))", + "interval": "", + "legendFormat": "", + "range": true, + "refId": "A" + } + ], + "title": "Total Degraded Searches (1h)", + "type": "gauge" + }, { "datasource": { "type": "prometheus" diff --git a/meilisearch/src/metrics.rs b/meilisearch/src/metrics.rs index bfe704979..652e6c227 100644 --- a/meilisearch/src/metrics.rs +++ b/meilisearch/src/metrics.rs @@ -22,6 +22,11 @@ lazy_static! { &["method", "path"] ) .expect("Can't create a metric"); + pub static ref MEILISEARCH_DEGRADED_SEARCH_REQUESTS: IntGauge = register_int_gauge!(opts!( + "meilisearch_degraded_search_requests", + "Meilisearch number of degraded search requests" + )) + .expect("Can't create a metric"); pub static ref MEILISEARCH_DB_SIZE_BYTES: IntGauge = register_int_gauge!(opts!("meilisearch_db_size_bytes", "Meilisearch DB Size In Bytes")) .expect("Can't create a metric"); diff --git a/meilisearch/src/routes/indexes/search.rs b/meilisearch/src/routes/indexes/search.rs index 3adfce970..6a430b6a3 100644 --- a/meilisearch/src/routes/indexes/search.rs +++ b/meilisearch/src/routes/indexes/search.rs @@ -17,6 +17,7 @@ use crate::analytics::{Analytics, SearchAggregator}; use crate::extractors::authentication::policies::*; use crate::extractors::authentication::GuardedData; use crate::extractors::sequential_extractor::SeqHandler; +use crate::metrics::MEILISEARCH_DEGRADED_SEARCH_REQUESTS; use crate::search::{ add_search_rules, perform_search, HybridQuery, MatchingStrategy, SearchQuery, SemanticRatio, DEFAULT_CROP_LENGTH, DEFAULT_CROP_MARKER, DEFAULT_HIGHLIGHT_POST_TAG, @@ -247,6 +248,9 @@ pub async fn search_with_post( .await?; if let Ok(ref search_result) = search_result { aggregate.succeed(search_result); + if search_result.degraded { + MEILISEARCH_DEGRADED_SEARCH_REQUESTS.inc(); + } } analytics.post_search(aggregate); From 4369e9e97c47401066f4a2c076ebd717f0fccf5b Mon Sep 17 00:00:00 2001 From: Tamo Date: Tue, 19 Mar 2024 11:14:28 +0100 Subject: [PATCH 10/86] add an error code test on the setting --- meilisearch-types/src/error.rs | 2 +- meilisearch-types/src/settings.rs | 2 +- meilisearch/src/routes/indexes/settings.rs | 3 +-- meilisearch/tests/common/index.rs | 5 ++++ meilisearch/tests/settings/errors.rs | 28 ++++++++++++++++++++++ 5 files changed, 36 insertions(+), 4 deletions(-) diff --git a/meilisearch-types/src/error.rs b/meilisearch-types/src/error.rs index bf9492ff6..aed77411a 100644 --- a/meilisearch-types/src/error.rs +++ b/meilisearch-types/src/error.rs @@ -259,7 +259,7 @@ InvalidSettingsProximityPrecision , InvalidRequest , BAD_REQUEST ; InvalidSettingsFaceting , InvalidRequest , BAD_REQUEST ; InvalidSettingsFilterableAttributes , InvalidRequest , BAD_REQUEST ; InvalidSettingsPagination , InvalidRequest , BAD_REQUEST ; -InvalidSettingsSearchCutoff , InvalidRequest , BAD_REQUEST ; +InvalidSettingsSearchCutoffMs , InvalidRequest , BAD_REQUEST ; InvalidSettingsEmbedders , InvalidRequest , BAD_REQUEST ; InvalidSettingsRankingRules , InvalidRequest , BAD_REQUEST ; InvalidSettingsSearchableAttributes , InvalidRequest , BAD_REQUEST ; diff --git a/meilisearch-types/src/settings.rs b/meilisearch-types/src/settings.rs index 23fe98347..5480e72c6 100644 --- a/meilisearch-types/src/settings.rs +++ b/meilisearch-types/src/settings.rs @@ -203,7 +203,7 @@ pub struct Settings { #[deserr(default, error = DeserrJsonError)] pub embedders: Setting>>, #[serde(default, skip_serializing_if = "Setting::is_not_set")] - #[deserr(default, error = DeserrJsonError)] + #[deserr(default, error = DeserrJsonError)] pub search_cutoff_ms: Setting, #[serde(skip)] diff --git a/meilisearch/src/routes/indexes/settings.rs b/meilisearch/src/routes/indexes/settings.rs index 4c03eb1a1..5dabd7b0d 100644 --- a/meilisearch/src/routes/indexes/settings.rs +++ b/meilisearch/src/routes/indexes/settings.rs @@ -138,7 +138,6 @@ macro_rules! make_setting_route { debug!(returns = ?settings, "Update settings"); let mut json = serde_json::json!(&settings); - dbg!(&json); let val = json[$camelcase_attr].take(); Ok(HttpResponse::Ok().json(val)) @@ -630,7 +629,7 @@ make_setting_route!( put, u64, meilisearch_types::deserr::DeserrJsonError< - meilisearch_types::error::deserr_codes::InvalidSettingsSearchCutoff, + meilisearch_types::error::deserr_codes::InvalidSettingsSearchCutoffMs, >, search_cutoff_ms, "searchCutoffMs", diff --git a/meilisearch/tests/common/index.rs b/meilisearch/tests/common/index.rs index 16fc10e98..9ed6a6077 100644 --- a/meilisearch/tests/common/index.rs +++ b/meilisearch/tests/common/index.rs @@ -328,6 +328,11 @@ impl Index<'_> { self.service.patch_encoded(url, settings, self.encoder).await } + pub async fn update_settings_search_cutoff_ms(&self, settings: Value) -> (Value, StatusCode) { + let url = format!("/indexes/{}/settings/search-cutoff-ms", urlencode(self.uid.as_ref())); + self.service.put_encoded(url, settings, self.encoder).await + } + pub async fn delete_settings(&self) -> (Value, StatusCode) { let url = format!("/indexes/{}/settings", urlencode(self.uid.as_ref())); self.service.delete(url).await diff --git a/meilisearch/tests/settings/errors.rs b/meilisearch/tests/settings/errors.rs index 687cef1f8..2bd17d649 100644 --- a/meilisearch/tests/settings/errors.rs +++ b/meilisearch/tests/settings/errors.rs @@ -337,3 +337,31 @@ async fn settings_bad_pagination() { } "###); } + +#[actix_rt::test] +async fn settings_bad_search_cutoff_ms() { + let server = Server::new().await; + let index = server.index("test"); + + let (response, code) = index.update_settings(json!({ "searchCutoffMs": "doggo" })).await; + snapshot!(code, @"400 Bad Request"); + snapshot!(json_string!(response), @r###" + { + "message": "Invalid value type at `.searchCutoffMs`: expected a positive integer, but found a string: `\"doggo\"`", + "code": "invalid_settings_search_cutoff_ms", + "type": "invalid_request", + "link": "https://docs.meilisearch.com/errors#invalid_settings_search_cutoff_ms" + } + "###); + + let (response, code) = index.update_settings_search_cutoff_ms(json!("doggo")).await; + snapshot!(code, @"400 Bad Request"); + snapshot!(json_string!(response), @r###" + { + "message": "Invalid value type: expected a positive integer, but found a string: `\"doggo\"`", + "code": "invalid_settings_search_cutoff_ms", + "type": "invalid_request", + "link": "https://docs.meilisearch.com/errors#invalid_settings_search_cutoff_ms" + } + "###); +} From 2a92c041006630e0ef573b159acfdd7bd6cfceac Mon Sep 17 00:00:00 2001 From: Louis Dureuil Date: Tue, 19 Mar 2024 11:31:32 +0100 Subject: [PATCH 11/86] Adding new assets --- BENCHMARKS.md | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/BENCHMARKS.md b/BENCHMARKS.md index dd69864cc..b3d311c45 100644 --- a/BENCHMARKS.md +++ b/BENCHMARKS.md @@ -317,6 +317,14 @@ They are JSON files with the following structure (comments are not actually supp } ``` +### Adding new assets + +Assets reside in our DigitalOcean S3 space. Assuming you have team access to the DigitalOcean S3 space: + +1. go to +2. upload your dataset: + 1. if your dataset is a single file, upload that single file using the "upload" button, + 2. otherwise, create a folder using the "create folder" button, then inside that folder upload your individual files. ## Upgrading `https://bench.meilisearch.dev` From bfec9468d47414e7f260261504ed46a83c291e65 Mon Sep 17 00:00:00 2001 From: Tamo Date: Tue, 19 Mar 2024 14:49:15 +0100 Subject: [PATCH 12/86] Update milli/src/search/mod.rs Co-authored-by: Louis Dureuil --- milli/src/search/mod.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/milli/src/search/mod.rs b/milli/src/search/mod.rs index f6ab8a7de..b3dd0c091 100644 --- a/milli/src/search/mod.rs +++ b/milli/src/search/mod.rs @@ -260,7 +260,7 @@ impl fmt::Debug for Search<'_> { .field("words_limit", words_limit) .field("distribution_shift", distribution_shift) .field("embedder_name", embedder_name) - .field("time_bduget", time_budget) + .field("time_budget", time_budget) .finish() } } From 0ae39644f7d97e83c8edfb19f344cbb2eb24fc40 Mon Sep 17 00:00:00 2001 From: Tamo Date: Tue, 19 Mar 2024 15:07:06 +0100 Subject: [PATCH 13/86] fix the facet search --- meilisearch/src/search.rs | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/meilisearch/src/search.rs b/meilisearch/src/search.rs index 0333eb0d5..3c00ca802 100644 --- a/meilisearch/src/search.rs +++ b/meilisearch/src/search.rs @@ -729,8 +729,11 @@ pub fn perform_facet_search( features: RoFeatures, ) -> Result { let before_search = Instant::now(); - let time_budget = TimeBudget::new(Duration::from_millis(150)); let rtxn = index.read_txn()?; + let time_budget = match index.search_cutoff(&rtxn)? { + Some(cutoff) => TimeBudget::new(Duration::from_millis(cutoff)), + None => TimeBudget::default(), + }; let (search, _, _, _) = prepare_search(index, &rtxn, &search_query, features, None, time_budget)?; From 7b9e0d29442df352a99aee7eab839464cd5764c1 Mon Sep 17 00:00:00 2001 From: Tamo Date: Tue, 19 Mar 2024 15:11:21 +0100 Subject: [PATCH 14/86] forward the degraded parameter to the hybrid search --- milli/src/search/hybrid.rs | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/milli/src/search/hybrid.rs b/milli/src/search/hybrid.rs index 1e14f4430..47ac5f46b 100644 --- a/milli/src/search/hybrid.rs +++ b/milli/src/search/hybrid.rs @@ -10,6 +10,7 @@ struct ScoreWithRatioResult { matching_words: MatchingWords, candidates: RoaringBitmap, document_scores: Vec<(u32, ScoreWithRatio)>, + degraded: bool, } type ScoreWithRatio = (Vec, f32); @@ -72,6 +73,7 @@ impl ScoreWithRatioResult { matching_words: results.matching_words, candidates: results.candidates, document_scores, + degraded: results.degraded, } } @@ -106,7 +108,7 @@ impl ScoreWithRatioResult { candidates: left.candidates | right.candidates, documents_ids, document_scores, - degraded: false, + degraded: left.degraded | right.degraded, } } } From d8fe4fe49d12b36bd9b82963c7e2fcc278d2f894 Mon Sep 17 00:00:00 2001 From: Tamo Date: Tue, 19 Mar 2024 15:45:04 +0100 Subject: [PATCH 15/86] return the order in the score details --- meilisearch/tests/search/mod.rs | 12 +++++++++--- milli/src/score_details.rs | 8 +++----- 2 files changed, 12 insertions(+), 8 deletions(-) diff --git a/meilisearch/tests/search/mod.rs b/meilisearch/tests/search/mod.rs index 971539a31..88470187a 100644 --- a/meilisearch/tests/search/mod.rs +++ b/meilisearch/tests/search/mod.rs @@ -869,7 +869,9 @@ async fn test_degraded_score_details() { ], "cattos": "pésti", "_rankingScoreDetails": { - "skipped": 0.0 + "skipped": { + "order": 0 + } } }, { @@ -883,7 +885,9 @@ async fn test_degraded_score_details() { "pestiféré" ], "_rankingScoreDetails": { - "skipped": 0.0 + "skipped": { + "order": 0 + } } }, { @@ -900,7 +904,9 @@ async fn test_degraded_score_details() { "gomez" ], "_rankingScoreDetails": { - "skipped": 0.0 + "skipped": { + "order": 0 + } } } ], diff --git a/milli/src/score_details.rs b/milli/src/score_details.rs index f2c6fb58a..08dfcdbb6 100644 --- a/milli/src/score_details.rs +++ b/milli/src/score_details.rs @@ -101,7 +101,7 @@ impl ScoreDetails { ScoreDetails::Vector(vector) => RankOrValue::Score( vector.value_similarity.as_ref().map(|(_, s)| *s as f64).unwrap_or(0.0f64), ), - ScoreDetails::Skipped => RankOrValue::Score(0.), + ScoreDetails::Skipped => RankOrValue::Rank(Rank { rank: 0, max_rank: 1 }), } } @@ -262,10 +262,8 @@ impl ScoreDetails { order += 1; } ScoreDetails::Skipped => { - details_map.insert( - "skipped".to_string(), - serde_json::Number::from_f64(0.).unwrap().into(), - ); + details_map + .insert("skipped".to_string(), serde_json::json!({ "order": order })); order += 1; } } From 098ab594eb156f5ba34ee4db81893f4e2c146b1f Mon Sep 17 00:00:00 2001 From: Louis Dureuil Date: Tue, 19 Mar 2024 17:32:32 +0100 Subject: [PATCH 16/86] A score of 0.0 is now lesser than a sort result MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit handles the niche case 🐩 in the hybrid search where: 1. a sort ranking rule is the first rule. 2. the keyword search is skipped at the first rule. 3. the semantic search is not skipped at the first rule. Previously, we would have the skipped search winning, whereas we want the non skipped one winning. --- milli/src/search/hybrid.rs | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/milli/src/search/hybrid.rs b/milli/src/search/hybrid.rs index 47ac5f46b..a8b7f0fcf 100644 --- a/milli/src/search/hybrid.rs +++ b/milli/src/search/hybrid.rs @@ -50,8 +50,12 @@ fn compare_scores( order => return order, } } - (Some(ScoreValue::Score(_)), Some(_)) => return Ordering::Greater, - (Some(_), Some(ScoreValue::Score(_))) => return Ordering::Less, + (Some(ScoreValue::Score(x)), Some(_)) => { + return if x == 0. { Ordering::Less } else { Ordering::Greater } + } + (Some(_), Some(ScoreValue::Score(x))) => { + return if x == 0. { Ordering::Greater } else { Ordering::Less } + } // if we have this, we're bad (Some(ScoreValue::GeoSort(_)), Some(ScoreValue::Sort(_))) | (Some(ScoreValue::Sort(_)), Some(ScoreValue::GeoSort(_))) => { From 2c3af8e51379b698276a23864c229cafc3984d77 Mon Sep 17 00:00:00 2001 From: Tamo Date: Tue, 19 Mar 2024 18:07:11 +0100 Subject: [PATCH 17/86] query the detailed score detail in the test --- milli/src/search/new/tests/cutoff.rs | 31 ++++++++++++++++++++++++++-- 1 file changed, 29 insertions(+), 2 deletions(-) diff --git a/milli/src/search/new/tests/cutoff.rs b/milli/src/search/new/tests/cutoff.rs index 4256abc2b..664b139f3 100644 --- a/milli/src/search/new/tests/cutoff.rs +++ b/milli/src/search/new/tests/cutoff.rs @@ -10,6 +10,7 @@ use maplit::hashset; use meili_snap::snapshot; use crate::index::tests::TempIndex; +use crate::score_details::ScoringStrategy; use crate::{Criterion, Filter, Search, TimeBudget}; fn create_index() -> TempIndex { @@ -94,6 +95,7 @@ fn degraded_search_and_score_details() { let mut search = Search::new(&rtxn, &index); search.query("hello puppy kefir"); search.limit(4); + search.scoring_strategy(ScoringStrategy::Detailed); search.time_budget(TimeBudget::max()); let result = search.execute().unwrap(); @@ -140,6 +142,12 @@ fn degraded_search_and_score_details() { max_matching_words: 3, }, ), + Typo( + Typo { + typo_count: 2, + max_typo_count: 3, + }, + ), ], [ Words( @@ -148,6 +156,12 @@ fn degraded_search_and_score_details() { max_matching_words: 3, }, ), + Typo( + Typo { + typo_count: 0, + max_typo_count: 2, + }, + ), ], ] "###); @@ -350,6 +364,12 @@ fn degraded_search_and_score_details() { max_matching_words: 3, }, ), + Typo( + Typo { + typo_count: 2, + max_typo_count: 3, + }, + ), ], [ Skipped, @@ -357,9 +377,9 @@ fn degraded_search_and_score_details() { ] "###); - // After FIVE loop iteration. The words ranking rule gave us a new bucket. + // After SIX loop iteration. The words ranking rule gave us a new bucket. // Since we reached the limit we were able to early exit without checking the typo ranking rule. - search.time_budget(TimeBudget::max().with_stop_after(5)); + search.time_budget(TimeBudget::max().with_stop_after(6)); let result = search.execute().unwrap(); snapshot!(format!("{:#?}\n{:#?}", result.documents_ids, result.document_scores), @r###" @@ -405,6 +425,12 @@ fn degraded_search_and_score_details() { max_matching_words: 3, }, ), + Typo( + Typo { + typo_count: 2, + max_typo_count: 3, + }, + ), ], [ Words( @@ -413,6 +439,7 @@ fn degraded_search_and_score_details() { max_matching_words: 3, }, ), + Skipped, ], ] "###); From 6079141ea6d77ac08a4b2c44ecc5f7fb07feb57f Mon Sep 17 00:00:00 2001 From: Tamo Date: Tue, 19 Mar 2024 18:30:14 +0100 Subject: [PATCH 18/86] snapshot the scores side by side with the score details --- milli/src/search/new/tests/cutoff.rs | 69 +++++++++++----------------- 1 file changed, 26 insertions(+), 43 deletions(-) diff --git a/milli/src/search/new/tests/cutoff.rs b/milli/src/search/new/tests/cutoff.rs index 664b139f3..63b67f2e7 100644 --- a/milli/src/search/new/tests/cutoff.rs +++ b/milli/src/search/new/tests/cutoff.rs @@ -10,7 +10,7 @@ use maplit::hashset; use meili_snap::snapshot; use crate::index::tests::TempIndex; -use crate::score_details::ScoringStrategy; +use crate::score_details::{ScoreDetails, ScoringStrategy}; use crate::{Criterion, Filter, Search, TimeBudget}; fn create_index() -> TempIndex { @@ -88,6 +88,7 @@ fn degraded_search_cannot_skip_filter() { } #[test] +#[allow(clippy::format_collect)] // the test is already quite big fn degraded_search_and_score_details() { let index = create_index(); let rtxn = index.read_txn().unwrap(); @@ -99,13 +100,10 @@ fn degraded_search_and_score_details() { search.time_budget(TimeBudget::max()); let result = search.execute().unwrap(); - snapshot!(format!("{:#?}\n{:#?}", result.documents_ids, result.document_scores), @r###" - [ - 4, - 1, - 0, - 3, - ] + snapshot!(format!("IDs: {:?}\nScores: {}\nScore Details:\n{:#?}", result.documents_ids, result.document_scores.iter().map(|scores| format!("{:.4} ", ScoreDetails::global_score(scores.iter()))).collect::(), result.document_scores), @r###" + IDs: [4, 1, 0, 3] + Scores: 1.0000 0.9167 0.8333 0.6667 + Score Details: [ [ Words( @@ -170,13 +168,10 @@ fn degraded_search_and_score_details() { search.time_budget(TimeBudget::max().with_stop_after(1)); let result = search.execute().unwrap(); - snapshot!(format!("{:#?}\n{:#?}", result.documents_ids, result.document_scores), @r###" - [ - 0, - 1, - 4, - 2, - ] + snapshot!(format!("IDs: {:?}\nScores: {}\nScore Details:\n{:#?}", result.documents_ids, result.document_scores.iter().map(|scores| format!("{:.4} ", ScoreDetails::global_score(scores.iter()))).collect::(), result.document_scores), @r###" + IDs: [0, 1, 4, 2] + Scores: 0.6667 0.6667 0.6667 0.0000 + Score Details: [ [ Words( @@ -215,13 +210,10 @@ fn degraded_search_and_score_details() { search.time_budget(TimeBudget::max().with_stop_after(2)); let result = search.execute().unwrap(); - snapshot!(format!("{:#?}\n{:#?}", result.documents_ids, result.document_scores), @r###" - [ - 4, - 0, - 1, - 2, - ] + snapshot!(format!("IDs: {:?}\nScores: {}\nScore Details:\n{:#?}", result.documents_ids, result.document_scores.iter().map(|scores| format!("{:.4} ", ScoreDetails::global_score(scores.iter()))).collect::(), result.document_scores), @r###" + IDs: [4, 0, 1, 2] + Scores: 1.0000 0.6667 0.6667 0.0000 + Score Details: [ [ Words( @@ -265,13 +257,10 @@ fn degraded_search_and_score_details() { search.time_budget(TimeBudget::max().with_stop_after(3)); let result = search.execute().unwrap(); - snapshot!(format!("{:#?}\n{:#?}", result.documents_ids, result.document_scores), @r###" - [ - 4, - 1, - 0, - 2, - ] + snapshot!(format!("IDs: {:?}\nScores: {}\nScore Details:\n{:#?}", result.documents_ids, result.document_scores.iter().map(|scores| format!("{:.4} ", ScoreDetails::global_score(scores.iter()))).collect::(), result.document_scores), @r###" + IDs: [4, 1, 0, 2] + Scores: 1.0000 0.9167 0.6667 0.0000 + Score Details: [ [ Words( @@ -321,13 +310,10 @@ fn degraded_search_and_score_details() { search.time_budget(TimeBudget::max().with_stop_after(4)); let result = search.execute().unwrap(); - snapshot!(format!("{:#?}\n{:#?}", result.documents_ids, result.document_scores), @r###" - [ - 4, - 1, - 0, - 2, - ] + snapshot!(format!("IDs: {:?}\nScores: {}\nScore Details:\n{:#?}", result.documents_ids, result.document_scores.iter().map(|scores| format!("{:.4} ", ScoreDetails::global_score(scores.iter()))).collect::(), result.document_scores), @r###" + IDs: [4, 1, 0, 2] + Scores: 1.0000 0.9167 0.8333 0.0000 + Score Details: [ [ Words( @@ -382,13 +368,10 @@ fn degraded_search_and_score_details() { search.time_budget(TimeBudget::max().with_stop_after(6)); let result = search.execute().unwrap(); - snapshot!(format!("{:#?}\n{:#?}", result.documents_ids, result.document_scores), @r###" - [ - 4, - 1, - 0, - 3, - ] + snapshot!(format!("IDs: {:?}\nScores: {}\nScore Details:\n{:#?}", result.documents_ids, result.document_scores.iter().map(|scores| format!("{:.4} ", ScoreDetails::global_score(scores.iter()))).collect::(), result.document_scores), @r###" + IDs: [4, 1, 0, 3] + Scores: 1.0000 0.9167 0.8333 0.3333 + Score Details: [ [ Words( From c5322df519ba9bb7c1010e5b8cf14edef5b8d168 Mon Sep 17 00:00:00 2001 From: Tamo Date: Wed, 20 Mar 2024 10:08:28 +0100 Subject: [PATCH 19/86] Revert "Revert "Merge remote-tracking branch 'origin/main' into release-v1.7.1"" --- .github/workflows/bench-pr.yml | 2 +- .github/workflows/milestone-workflow.yml | 19 + CONTRIBUTING.md | 2 +- Cargo.lock | 195 ++++------- Cargo.toml | 2 +- meilisearch-types/Cargo.toml | 2 +- meilisearch/Cargo.toml | 12 +- meilisearch/src/main.rs | 2 +- meilisearch/src/option.rs | 4 +- meilisearch/src/routes/indexes/settings.rs | 1 + meilisearch/src/search.rs | 21 +- meilisearch/tests/documents/add_documents.rs | 239 ++++++++++++- meilisearch/tests/search/facet_search.rs | 43 +++ milli/src/index.rs | 14 +- milli/src/lib.rs | 7 +- milli/src/order_by_map.rs | 57 +++ milli/src/search/facet/facet_range_search.rs | 4 +- milli/src/search/facet/mod.rs | 3 + milli/src/search/facet/search.rs | 326 ++++++++++++++++++ milli/src/search/mod.rs | 249 +------------ milli/src/search/new/tests/typo_proximity.rs | 2 +- milli/src/update/settings.rs | 14 +- milli/src/vector/error.rs | 39 +++ milli/src/vector/mod.rs | 18 + milli/src/vector/ollama.rs | 307 +++++++++++++++++ milli/src/vector/openai.rs | 20 +- milli/src/vector/settings.rs | 29 +- workloads/settings-add-remove-filters.json | 94 +++++ workloads/settings-proximity-precision.json | 86 +++++ .../settings-remove-add-swap-searchable.json | 114 ++++++ workloads/settings-typo.json | 115 ++++++ xtask/src/bench/dashboard.rs | 312 +++++++++-------- xtask/src/bench/mod.rs | 24 +- xtask/src/bench/workload.rs | 16 +- 34 files changed, 1784 insertions(+), 610 deletions(-) create mode 100644 milli/src/order_by_map.rs create mode 100644 milli/src/search/facet/search.rs create mode 100644 milli/src/vector/ollama.rs create mode 100644 workloads/settings-add-remove-filters.json create mode 100644 workloads/settings-proximity-precision.json create mode 100644 workloads/settings-remove-add-swap-searchable.json create mode 100644 workloads/settings-typo.json diff --git a/.github/workflows/bench-pr.yml b/.github/workflows/bench-pr.yml index 6f4956542..418a23717 100644 --- a/.github/workflows/bench-pr.yml +++ b/.github/workflows/bench-pr.yml @@ -43,4 +43,4 @@ jobs: - name: Run benchmarks on PR ${{ github.event.issue.id }} run: | - cargo xtask bench --api-key "${{ secrets.BENCHMARK_API_KEY }}" --dashboard-url "${{ vars.BENCHMARK_DASHBOARD_URL }}" --reason "[Comment](${{ github.event.comment.url }}) on [#${{github.event.issue.id}}](${{ github.event.issue.url }})" -- ${{ steps.command.outputs.command-arguments }} \ No newline at end of file + cargo xtask bench --api-key "${{ secrets.BENCHMARK_API_KEY }}" --dashboard-url "${{ vars.BENCHMARK_DASHBOARD_URL }}" --reason "[Comment](${{ github.event.comment.html_url }}) on [#${{ github.event.issue.number }}](${{ github.event.issue.html_url }})" -- ${{ steps.command.outputs.command-arguments }} \ No newline at end of file diff --git a/.github/workflows/milestone-workflow.yml b/.github/workflows/milestone-workflow.yml index 2b8b7bf62..2ede3dc21 100644 --- a/.github/workflows/milestone-workflow.yml +++ b/.github/workflows/milestone-workflow.yml @@ -110,6 +110,25 @@ jobs: --milestone $MILESTONE_VERSION \ --assignee curquiza + create-update-version-issue: + needs: get-release-version + # Create the changelog issue if the release is not only a patch release + if: github.event.action == 'created' + runs-on: ubuntu-latest + env: + ISSUE_TEMPLATE: issue-template.md + steps: + - uses: actions/checkout@v3 + - name: Download the issue template + run: curl -s https://raw.githubusercontent.com/meilisearch/engine-team/main/issue-templates/update-version-issue.md > $ISSUE_TEMPLATE + - name: Create the issue + run: | + gh issue create \ + --title "Update version in Cargo.toml for $MILESTONE_VERSION" \ + --label 'maintenance' \ + --body-file $ISSUE_TEMPLATE \ + --milestone $MILESTONE_VERSION + # ---------------- # MILESTONE CLOSED # ---------------- diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 6d6e6076b..f33416820 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -4,7 +4,7 @@ First, thank you for contributing to Meilisearch! The goal of this document is t Remember that there are many ways to contribute other than writing code: writing [tutorials or blog posts](https://github.com/meilisearch/awesome-meilisearch), improving [the documentation](https://github.com/meilisearch/documentation), submitting [bug reports](https://github.com/meilisearch/meilisearch/issues/new?assignees=&labels=&template=bug_report.md&title=) and [feature requests](https://github.com/meilisearch/product/discussions/categories/feedback-feature-proposal)... -The code in this repository is only concerned with managing multiple indexes, handling the update store, and exposing an HTTP API. Search and indexation are the domain of our core engine, [`milli`](https://github.com/meilisearch/milli), while tokenization is handled by [our `charabia` library](https://github.com/meilisearch/charabia/). +Meilisearch can manage multiple indexes, handle the update store, and expose an HTTP API. Search and indexation are the domain of our core engine, [`milli`](https://github.com/meilisearch/meilisearch/tree/main/milli), while tokenization is handled by [our `charabia` library](https://github.com/meilisearch/charabia/). If Meilisearch does not offer optimized support for your language, please consider contributing to `charabia` by following the [CONTRIBUTING.md file](https://github.com/meilisearch/charabia/blob/main/CONTRIBUTING.md) and integrating your intended normalizer/segmenter. diff --git a/Cargo.lock b/Cargo.lock index a1527c31c..bdca7e24c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -36,9 +36,9 @@ dependencies = [ [[package]] name = "actix-http" -version = "3.5.1" +version = "3.6.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "129d4c88e98860e1758c5de288d1632b07970a16d59bdf7b8d66053d582bb71f" +checksum = "d223b13fd481fc0d1f83bb12659ae774d9e3601814c68a0bc539731698cca743" dependencies = [ "actix-codec", "actix-rt", @@ -138,9 +138,9 @@ dependencies = [ [[package]] name = "actix-tls" -version = "3.1.1" +version = "3.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "72616e7fbec0aa99c6f3164677fa48ff5a60036d0799c98cab894a44f3e0efc3" +checksum = "d4cce60a2f2b477bc72e5cde0af1812a6e82d8fd85b5570a5dcf2a5bf2c5be5f" dependencies = [ "actix-rt", "actix-service", @@ -148,13 +148,11 @@ dependencies = [ "futures-core", "impl-more", "pin-project-lite", - "rustls 0.21.6", - "rustls-webpki", "tokio", - "tokio-rustls 0.23.4", + "tokio-rustls", "tokio-util", "tracing", - "webpki-roots 0.22.6", + "webpki-roots", ] [[package]] @@ -169,9 +167,9 @@ dependencies = [ [[package]] name = "actix-web" -version = "4.4.1" +version = "4.5.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e43428f3bf11dee6d166b00ec2df4e3aa8cc1606aaa0b7433c146852e2f4e03b" +checksum = "43a6556ddebb638c2358714d853257ed226ece6023ef9364f23f0c70737ea984" dependencies = [ "actix-codec", "actix-http", @@ -259,9 +257,9 @@ dependencies = [ [[package]] name = "ahash" -version = "0.8.8" +version = "0.8.11" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "42cd52102d3df161c77a887b608d7a4897d7cc112886a9537b738a887a03aaff" +checksum = "e89da841a80418a9b391ebaea17f5c112ffaaa96f621d2c285b5174da76b9011" dependencies = [ "cfg-if", "getrandom", @@ -496,7 +494,7 @@ checksum = "8c3c1a368f70d6cf7302d78f8f7093da241fb8e8807c05cc9e51a125895a6d5b" [[package]] name = "benchmarks" -version = "1.7.1" +version = "1.8.0" dependencies = [ "anyhow", "bytes", @@ -630,7 +628,7 @@ dependencies = [ [[package]] name = "build-info" -version = "1.7.1" +version = "1.8.0" dependencies = [ "anyhow", "time", @@ -835,9 +833,9 @@ checksum = "37b2a672a2cb129a2e41c10b1224bb368f9f37a2b16b612598138befd7b37eb5" [[package]] name = "cc" -version = "1.0.82" +version = "1.0.83" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "305fe645edc1442a0fa8b6726ba61d422798d37a52e12eaecf4b022ebbb88f01" +checksum = "f1174fb0b6ec23863f8b971027804a42614e347eafb0a95bf0b12cdae21fc4d0" dependencies = [ "jobserver", "libc", @@ -1531,7 +1529,7 @@ dependencies = [ [[package]] name = "dump" -version = "1.7.1" +version = "1.8.0" dependencies = [ "anyhow", "big_s", @@ -1769,7 +1767,7 @@ dependencies = [ [[package]] name = "file-store" -version = "1.7.1" +version = "1.8.0" dependencies = [ "faux", "tempfile", @@ -1792,7 +1790,7 @@ dependencies = [ [[package]] name = "filter-parser" -version = "1.7.1" +version = "1.8.0" dependencies = [ "insta", "nom", @@ -1812,7 +1810,7 @@ dependencies = [ [[package]] name = "flatten-serde-json" -version = "1.7.1" +version = "1.8.0" dependencies = [ "criterion", "serde_json", @@ -1930,7 +1928,7 @@ dependencies = [ [[package]] name = "fuzzers" -version = "1.7.1" +version = "1.8.0" dependencies = [ "arbitrary", "clap", @@ -2104,8 +2102,10 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "be4136b2a15dd319360be1c07d9933517ccf0be8f16bf62a3bee4f0d618df427" dependencies = [ "cfg-if", + "js-sys", "libc", "wasi", + "wasm-bindgen", ] [[package]] @@ -2224,7 +2224,7 @@ dependencies = [ "atomic-polyfill", "hash32", "rustc_version", - "spin 0.9.8", + "spin", "stable_deref_trait", ] @@ -2393,9 +2393,9 @@ dependencies = [ "futures-util", "http 0.2.11", "hyper", - "rustls 0.21.6", + "rustls", "tokio", - "tokio-rustls 0.24.1", + "tokio-rustls", ] [[package]] @@ -2422,7 +2422,7 @@ checksum = "206ca75c9c03ba3d4ace2460e57b189f39f43de612c2f85836e65c929701bb2d" [[package]] name = "index-scheduler" -version = "1.7.1" +version = "1.8.0" dependencies = [ "anyhow", "big_s", @@ -2609,7 +2609,7 @@ dependencies = [ [[package]] name = "json-depth-checker" -version = "1.7.1" +version = "1.8.0" dependencies = [ "criterion", "serde_json", @@ -2617,13 +2617,14 @@ dependencies = [ [[package]] name = "jsonwebtoken" -version = "8.3.0" +version = "9.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6971da4d9c3aa03c3d8f3ff0f4155b534aad021292003895a469716b2a230378" +checksum = "5c7ea04a7c5c055c175f189b6dc6ba036fd62306b58c66c9f6389036c503a3f4" dependencies = [ "base64 0.21.7", + "js-sys", "pem", - "ring 0.16.20", + "ring", "serde", "serde_json", "simple_asn1", @@ -3117,7 +3118,7 @@ checksum = "490cc448043f947bae3cbee9c203358d62dbee0db12107a74be5c30ccfd09771" [[package]] name = "meili-snap" -version = "1.7.1" +version = "1.8.0" dependencies = [ "insta", "md5", @@ -3126,7 +3127,7 @@ dependencies = [ [[package]] name = "meilisearch" -version = "1.7.1" +version = "1.8.0" dependencies = [ "actix-cors", "actix-http", @@ -3184,7 +3185,7 @@ dependencies = [ "rayon", "regex", "reqwest", - "rustls 0.20.9", + "rustls", "rustls-pemfile", "segment", "serde", @@ -3219,7 +3220,7 @@ dependencies = [ [[package]] name = "meilisearch-auth" -version = "1.7.1" +version = "1.8.0" dependencies = [ "base64 0.21.7", "enum-iterator", @@ -3238,7 +3239,7 @@ dependencies = [ [[package]] name = "meilisearch-types" -version = "1.7.1" +version = "1.8.0" dependencies = [ "actix-web", "anyhow", @@ -3268,7 +3269,7 @@ dependencies = [ [[package]] name = "meilitool" -version = "1.7.1" +version = "1.8.0" dependencies = [ "anyhow", "clap", @@ -3307,7 +3308,7 @@ dependencies = [ [[package]] name = "milli" -version = "1.7.1" +version = "1.8.0" dependencies = [ "arroy", "big_s", @@ -3413,9 +3414,9 @@ dependencies = [ [[package]] name = "mio" -version = "0.8.9" +version = "0.8.11" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3dce281c5e46beae905d4de1870d8b1509a9142b62eedf18b443b011ca8343d0" +checksum = "a4a650543ca06a924e8b371db273b2756685faae30f8487da1b56505a8f78b0c" dependencies = [ "libc", "log", @@ -3733,11 +3734,12 @@ checksum = "19b17cddbe7ec3f8bc800887bab5e717348c95ea2ca0b1bf0837fb964dc67099" [[package]] name = "pem" -version = "1.1.1" +version = "3.0.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a8835c273a76a90455d7344889b0964598e3316e2a79ede8e36f16bdcf2228b8" +checksum = "1b8fcc794035347fb64beda2d3b462595dd2753e3f268d89c5aae77e8cf2c310" dependencies = [ - "base64 0.13.1", + "base64 0.21.7", + "serde", ] [[package]] @@ -3748,7 +3750,7 @@ checksum = "e3148f5046208a5d56bcfc03053e3ca6334e51da8dfb19b6cdc8b306fae3283e" [[package]] name = "permissive-json-pointer" -version = "1.7.1" +version = "1.8.0" dependencies = [ "big_s", "serde_json", @@ -4239,14 +4241,14 @@ dependencies = [ "once_cell", "percent-encoding", "pin-project-lite", - "rustls 0.21.6", + "rustls", "rustls-pemfile", "serde", "serde_json", "serde_urlencoded", "system-configuration", "tokio", - "tokio-rustls 0.24.1", + "tokio-rustls", "tokio-util", "tower-service", "url", @@ -4254,7 +4256,7 @@ dependencies = [ "wasm-bindgen-futures", "wasm-streams", "web-sys", - "webpki-roots 0.25.3", + "webpki-roots", "winreg", ] @@ -4272,30 +4274,15 @@ checksum = "b9b1a3d5f46d53f4a3478e2be4a5a5ce5108ea58b100dcd139830eae7f79a3a1" [[package]] name = "ring" -version = "0.16.20" +version = "0.17.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3053cf52e236a3ed746dfc745aa9cacf1b791d846bdaf412f60a8d7d6e17c8fc" -dependencies = [ - "cc", - "libc", - "once_cell", - "spin 0.5.2", - "untrusted 0.7.1", - "web-sys", - "winapi", -] - -[[package]] -name = "ring" -version = "0.17.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9babe80d5c16becf6594aa32ad2be8fe08498e7ae60b77de8df700e67f191d7e" +checksum = "688c63d65483050968b2a8937f7995f443e27041a0f7700aa59b0822aedebb74" dependencies = [ "cc", "getrandom", "libc", - "spin 0.9.8", - "untrusted 0.9.0", + "spin", + "untrusted", "windows-sys 0.48.0", ] @@ -4373,24 +4360,12 @@ dependencies = [ [[package]] name = "rustls" -version = "0.20.9" +version = "0.21.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1b80e3dec595989ea8510028f30c408a4630db12c9cbb8de34203b89d6577e99" +checksum = "f9d5a6813c0759e4609cd494e8e725babae6a2ca7b62a5536a13daaec6fcb7ba" dependencies = [ "log", - "ring 0.16.20", - "sct", - "webpki", -] - -[[package]] -name = "rustls" -version = "0.21.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1d1feddffcfcc0b33f5c6ce9a29e341e4cd59c3f78e7ee45f4a40c038b1d6cbb" -dependencies = [ - "log", - "ring 0.16.20", + "ring", "rustls-webpki", "sct", ] @@ -4410,8 +4385,8 @@ version = "0.101.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8b6275d1ee7a1cd780b64aca7726599a1dbc893b1e64144529e55c3c2f745765" dependencies = [ - "ring 0.17.3", - "untrusted 0.9.0", + "ring", + "untrusted", ] [[package]] @@ -4453,12 +4428,12 @@ checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" [[package]] name = "sct" -version = "0.7.0" +version = "0.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d53dcdb7c9f8158937a7981b48accfd39a43af418591a5d008c7b22b5e1b7ca4" +checksum = "da046153aa2352493d6cb7da4b6e5c0c057d8a1d0a9aa8560baffdd945acd414" dependencies = [ - "ring 0.16.20", - "untrusted 0.7.1", + "ring", + "untrusted", ] [[package]] @@ -4721,12 +4696,6 @@ dependencies = [ "winapi", ] -[[package]] -name = "spin" -version = "0.5.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6e63cff320ae2c57904679ba7cb63280a3dc4613885beafb148ee7bf9aa9042d" - [[package]] name = "spin" version = "0.9.8" @@ -5080,24 +5049,13 @@ dependencies = [ "syn 2.0.48", ] -[[package]] -name = "tokio-rustls" -version = "0.23.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c43ee83903113e03984cb9e5cebe6c04a5116269e900e3ddba8f068a62adda59" -dependencies = [ - "rustls 0.20.9", - "tokio", - "webpki", -] - [[package]] name = "tokio-rustls" version = "0.24.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c28327cf380ac148141087fbfb9de9d7bd4e84ab5d2c28fbc911d753de8a7081" dependencies = [ - "rustls 0.21.6", + "rustls", "tokio", ] @@ -5366,12 +5324,6 @@ version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39ec24b3121d976906ece63c9daad25b85969647682eee313cb5779fdd69e14e" -[[package]] -name = "untrusted" -version = "0.7.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a156c684c91ea7d62626509bce3cb4e1d9ed5c4d978f7b4352658f96a4c26b4a" - [[package]] name = "untrusted" version = "0.9.0" @@ -5388,13 +5340,13 @@ dependencies = [ "flate2", "log", "once_cell", - "rustls 0.21.6", + "rustls", "rustls-webpki", "serde", "serde_json", "socks", "url", - "webpki-roots 0.25.3", + "webpki-roots", ] [[package]] @@ -5630,25 +5582,6 @@ dependencies = [ "wasm-bindgen", ] -[[package]] -name = "webpki" -version = "0.22.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "07ecc0cd7cac091bf682ec5efa18b1cff79d617b84181f38b3951dbe135f607f" -dependencies = [ - "ring 0.16.20", - "untrusted 0.7.1", -] - -[[package]] -name = "webpki-roots" -version = "0.22.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b6c71e40d7d2c34a5106301fb632274ca37242cd0c9d3e64dbece371a40a2d87" -dependencies = [ - "webpki", -] - [[package]] name = "webpki-roots" version = "0.25.3" @@ -5943,7 +5876,7 @@ dependencies = [ [[package]] name = "xtask" -version = "1.7.1" +version = "1.8.0" dependencies = [ "anyhow", "build-info", diff --git a/Cargo.toml b/Cargo.toml index 5337ec5c3..1d0e0ca0d 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -21,7 +21,7 @@ members = [ ] [workspace.package] -version = "1.7.1" +version = "1.8.0" authors = [ "Quentin de Quelen ", "Clément Renault ", diff --git a/meilisearch-types/Cargo.toml b/meilisearch-types/Cargo.toml index b5460fb56..7709d33d7 100644 --- a/meilisearch-types/Cargo.toml +++ b/meilisearch-types/Cargo.toml @@ -11,7 +11,7 @@ edition.workspace = true license.workspace = true [dependencies] -actix-web = { version = "4.4.1", default-features = false } +actix-web = { version = "4.5.1", default-features = false } anyhow = "1.0.79" convert_case = "0.6.0" csv = "1.3.0" diff --git a/meilisearch/Cargo.toml b/meilisearch/Cargo.toml index b65c466ca..04b919904 100644 --- a/meilisearch/Cargo.toml +++ b/meilisearch/Cargo.toml @@ -14,18 +14,18 @@ default-run = "meilisearch" [dependencies] actix-cors = "0.7.0" -actix-http = { version = "3.5.1", default-features = false, features = [ +actix-http = { version = "3.6.0", default-features = false, features = [ "compress-brotli", "compress-gzip", - "rustls", + "rustls-0_21", ] } actix-utils = "3.0.1" -actix-web = { version = "4.4.1", default-features = false, features = [ +actix-web = { version = "4.5.1", default-features = false, features = [ "macros", "compress-brotli", "compress-gzip", "cookies", - "rustls", + "rustls-0_21", ] } actix-web-static-files = { git = "https://github.com/kilork/actix-web-static-files.git", rev = "2d3b6160", optional = true } anyhow = { version = "1.0.79", features = ["backtrace"] } @@ -52,7 +52,7 @@ index-scheduler = { path = "../index-scheduler" } indexmap = { version = "2.1.0", features = ["serde"] } is-terminal = "0.4.10" itertools = "0.11.0" -jsonwebtoken = "8.3.0" +jsonwebtoken = "9.2.0" lazy_static = "1.4.0" meilisearch-auth = { path = "../meilisearch-auth" } meilisearch-types = { path = "../meilisearch-types" } @@ -75,7 +75,7 @@ reqwest = { version = "0.11.23", features = [ "rustls-tls", "json", ], default-features = false } -rustls = "0.20.8" +rustls = "0.21.6" rustls-pemfile = "1.0.2" segment = { version = "0.2.3", optional = true } serde = { version = "1.0.195", features = ["derive"] } diff --git a/meilisearch/src/main.rs b/meilisearch/src/main.rs index 3451325b2..af02f58e1 100644 --- a/meilisearch/src/main.rs +++ b/meilisearch/src/main.rs @@ -151,7 +151,7 @@ async fn run_http( .keep_alive(KeepAlive::Os); if let Some(config) = opt_clone.get_ssl_config()? { - http_server.bind_rustls(opt_clone.http_addr, config)?.run().await?; + http_server.bind_rustls_021(opt_clone.http_addr, config)?.run().await?; } else { http_server.bind(&opt_clone.http_addr)?.run().await?; } diff --git a/meilisearch/src/option.rs b/meilisearch/src/option.rs index 92d53fd32..43bf2c62c 100644 --- a/meilisearch/src/option.rs +++ b/meilisearch/src/option.rs @@ -564,11 +564,11 @@ impl Opt { } if self.ssl_require_auth { let verifier = AllowAnyAuthenticatedClient::new(client_auth_roots); - config.with_client_cert_verifier(verifier) + config.with_client_cert_verifier(Arc::from(verifier)) } else { let verifier = AllowAnyAnonymousOrAuthenticatedClient::new(client_auth_roots); - config.with_client_cert_verifier(verifier) + config.with_client_cert_verifier(Arc::from(verifier)) } } None => config.with_no_client_auth(), diff --git a/meilisearch/src/routes/indexes/settings.rs b/meilisearch/src/routes/indexes/settings.rs index c71d83279..c782e78cb 100644 --- a/meilisearch/src/routes/indexes/settings.rs +++ b/meilisearch/src/routes/indexes/settings.rs @@ -604,6 +604,7 @@ fn embedder_analytics( EmbedderSource::OpenAi => sources.insert("openAi"), EmbedderSource::HuggingFace => sources.insert("huggingFace"), EmbedderSource::UserProvided => sources.insert("userProvided"), + EmbedderSource::Ollama => sources.insert("ollama"), }; } }; diff --git a/meilisearch/src/search.rs b/meilisearch/src/search.rs index 27de36c6d..e65192d16 100644 --- a/meilisearch/src/search.rs +++ b/meilisearch/src/search.rs @@ -530,7 +530,7 @@ pub fn perform_search( // The attributes to retrieve are the ones explicitly marked as to retrieve (all by default), // but these attributes must be also be present // - in the fields_ids_map - // - in the the displayed attributes + // - in the displayed attributes let to_retrieve_ids: BTreeSet<_> = query .attributes_to_retrieve .as_ref() @@ -671,27 +671,16 @@ pub fn perform_search( let sort_facet_values_by = index.sort_facet_values_by(&rtxn).map_err(milli::Error::from)?; - let default_sort_facet_values_by = - sort_facet_values_by.get("*").copied().unwrap_or_default(); if fields.iter().all(|f| f != "*") { - let fields: Vec<_> = fields - .iter() - .map(|n| { - ( - n, - sort_facet_values_by - .get(n) - .copied() - .unwrap_or(default_sort_facet_values_by), - ) - }) - .collect(); + let fields: Vec<_> = + fields.iter().map(|n| (n, sort_facet_values_by.get(n))).collect(); facet_distribution.facets(fields); } + let distribution = facet_distribution .candidates(candidates) - .default_order_by(default_sort_facet_values_by) + .default_order_by(sort_facet_values_by.get("*")) .execute()?; let stats = facet_distribution.compute_stats()?; (Some(distribution), Some(stats)) diff --git a/meilisearch/tests/documents/add_documents.rs b/meilisearch/tests/documents/add_documents.rs index e6af85229..b1262fa2d 100644 --- a/meilisearch/tests/documents/add_documents.rs +++ b/meilisearch/tests/documents/add_documents.rs @@ -1237,8 +1237,8 @@ async fn error_add_documents_missing_document_id() { } #[actix_rt::test] -#[ignore] // // TODO: Fix in an other PR: this does not provoke any error. -async fn error_document_field_limit_reached() { +#[should_panic] +async fn error_document_field_limit_reached_in_one_document() { let server = Server::new().await; let index = server.index("test"); @@ -1246,22 +1246,241 @@ async fn error_document_field_limit_reached() { let mut big_object = std::collections::HashMap::new(); big_object.insert("id".to_owned(), "wow"); - for i in 0..65535 { + for i in 0..(u16::MAX as usize + 1) { let key = i.to_string(); big_object.insert(key, "I am a text!"); } let documents = json!([big_object]); - let (_response, code) = index.update_documents(documents, Some("id")).await; - snapshot!(code, @"202"); + let (response, code) = index.update_documents(documents, Some("id")).await; + snapshot!(code, @"500 Internal Server Error"); - index.wait_task(0).await; - let (response, code) = index.get_task(0).await; - snapshot!(code, @"200"); + let response = index.wait_task(response.uid()).await; + snapshot!(code, @"202 Accepted"); // Documents without a primary key are not accepted. - snapshot!(json_string!(response, { ".duration" => "[duration]", ".enqueuedAt" => "[date]", ".startedAt" => "[date]", ".finishedAt" => "[date]" }), - @""); + snapshot!(response, + @r###" + { + "uid": 1, + "indexUid": "test", + "status": "succeeded", + "type": "documentAdditionOrUpdate", + "canceledBy": null, + "details": { + "receivedDocuments": 1, + "indexedDocuments": 1 + }, + "error": null, + "duration": "[duration]", + "enqueuedAt": "[date]", + "startedAt": "[date]", + "finishedAt": "[date]" + } + "###); +} + +#[actix_rt::test] +async fn error_document_field_limit_reached_over_multiple_documents() { + let server = Server::new().await; + let index = server.index("test"); + + index.create(Some("id")).await; + + let mut big_object = std::collections::HashMap::new(); + big_object.insert("id".to_owned(), "wow"); + for i in 0..(u16::MAX / 2) { + let key = i.to_string(); + big_object.insert(key, "I am a text!"); + } + + let documents = json!([big_object]); + + let (response, code) = index.update_documents(documents, Some("id")).await; + snapshot!(code, @"202 Accepted"); + + let response = index.wait_task(response.uid()).await; + snapshot!(code, @"202 Accepted"); + snapshot!(response, + @r###" + { + "uid": 1, + "indexUid": "test", + "status": "succeeded", + "type": "documentAdditionOrUpdate", + "canceledBy": null, + "details": { + "receivedDocuments": 1, + "indexedDocuments": 1 + }, + "error": null, + "duration": "[duration]", + "enqueuedAt": "[date]", + "startedAt": "[date]", + "finishedAt": "[date]" + } + "###); + + let mut big_object = std::collections::HashMap::new(); + big_object.insert("id".to_owned(), "waw"); + for i in (u16::MAX as usize / 2)..(u16::MAX as usize + 1) { + let key = i.to_string(); + big_object.insert(key, "I am a text!"); + } + + let documents = json!([big_object]); + + let (response, code) = index.update_documents(documents, Some("id")).await; + snapshot!(code, @"202 Accepted"); + + let response = index.wait_task(response.uid()).await; + snapshot!(code, @"202 Accepted"); + snapshot!(response, + @r###" + { + "uid": 2, + "indexUid": "test", + "status": "failed", + "type": "documentAdditionOrUpdate", + "canceledBy": null, + "details": { + "receivedDocuments": 1, + "indexedDocuments": 0 + }, + "error": { + "message": "A document cannot contain more than 65,535 fields.", + "code": "max_fields_limit_exceeded", + "type": "invalid_request", + "link": "https://docs.meilisearch.com/errors#max_fields_limit_exceeded" + }, + "duration": "[duration]", + "enqueuedAt": "[date]", + "startedAt": "[date]", + "finishedAt": "[date]" + } + "###); +} + +#[actix_rt::test] +async fn error_document_field_limit_reached_in_one_nested_document() { + let server = Server::new().await; + let index = server.index("test"); + + index.create(Some("id")).await; + + let mut nested = std::collections::HashMap::new(); + for i in 0..(u16::MAX as usize + 1) { + let key = i.to_string(); + nested.insert(key, "I am a text!"); + } + let mut big_object = std::collections::HashMap::new(); + big_object.insert("id".to_owned(), "wow"); + + let documents = json!([big_object]); + + let (response, code) = index.update_documents(documents, Some("id")).await; + snapshot!(code, @"202 Accepted"); + + let response = index.wait_task(response.uid()).await; + snapshot!(code, @"202 Accepted"); + // Documents without a primary key are not accepted. + snapshot!(response, + @r###" + { + "uid": 1, + "indexUid": "test", + "status": "succeeded", + "type": "documentAdditionOrUpdate", + "canceledBy": null, + "details": { + "receivedDocuments": 1, + "indexedDocuments": 1 + }, + "error": null, + "duration": "[duration]", + "enqueuedAt": "[date]", + "startedAt": "[date]", + "finishedAt": "[date]" + } + "###); +} + +#[actix_rt::test] +async fn error_document_field_limit_reached_over_multiple_documents_with_nested_fields() { + let server = Server::new().await; + let index = server.index("test"); + + index.create(Some("id")).await; + + let mut nested = std::collections::HashMap::new(); + for i in 0..(u16::MAX / 2) { + let key = i.to_string(); + nested.insert(key, "I am a text!"); + } + let mut big_object = std::collections::HashMap::new(); + big_object.insert("id".to_owned(), "wow"); + + let documents = json!([big_object]); + + let (response, code) = index.update_documents(documents, Some("id")).await; + snapshot!(code, @"202 Accepted"); + + let response = index.wait_task(response.uid()).await; + snapshot!(code, @"202 Accepted"); + snapshot!(response, + @r###" + { + "uid": 1, + "indexUid": "test", + "status": "succeeded", + "type": "documentAdditionOrUpdate", + "canceledBy": null, + "details": { + "receivedDocuments": 1, + "indexedDocuments": 1 + }, + "error": null, + "duration": "[duration]", + "enqueuedAt": "[date]", + "startedAt": "[date]", + "finishedAt": "[date]" + } + "###); + + let mut nested = std::collections::HashMap::new(); + for i in 0..(u16::MAX / 2) { + let key = i.to_string(); + nested.insert(key, "I am a text!"); + } + let mut big_object = std::collections::HashMap::new(); + big_object.insert("id".to_owned(), "wow"); + + let documents = json!([big_object]); + + let (response, code) = index.update_documents(documents, Some("id")).await; + snapshot!(code, @"202 Accepted"); + + let response = index.wait_task(response.uid()).await; + snapshot!(code, @"202 Accepted"); + snapshot!(response, + @r###" + { + "uid": 2, + "indexUid": "test", + "status": "succeeded", + "type": "documentAdditionOrUpdate", + "canceledBy": null, + "details": { + "receivedDocuments": 1, + "indexedDocuments": 1 + }, + "error": null, + "duration": "[duration]", + "enqueuedAt": "[date]", + "startedAt": "[date]", + "finishedAt": "[date]" + } + "###); } #[actix_rt::test] diff --git a/meilisearch/tests/search/facet_search.rs b/meilisearch/tests/search/facet_search.rs index 5f9f631f9..12d2226a9 100644 --- a/meilisearch/tests/search/facet_search.rs +++ b/meilisearch/tests/search/facet_search.rs @@ -123,6 +123,28 @@ async fn simple_facet_search_with_max_values() { assert_eq!(dbg!(response)["facetHits"].as_array().unwrap().len(), 1); } +#[actix_rt::test] +async fn simple_facet_search_by_count_with_max_values() { + let server = Server::new().await; + let index = server.index("test"); + + let documents = DOCUMENTS.clone(); + index + .update_settings_faceting( + json!({ "maxValuesPerFacet": 1, "sortFacetValuesBy": { "*": "count" } }), + ) + .await; + index.update_settings_filterable_attributes(json!(["genres"])).await; + index.add_documents(documents, None).await; + index.wait_task(2).await; + + let (response, code) = + index.facet_search(json!({"facetName": "genres", "facetQuery": "a"})).await; + + assert_eq!(code, 200, "{}", response); + assert_eq!(dbg!(response)["facetHits"].as_array().unwrap().len(), 1); +} + #[actix_rt::test] async fn non_filterable_facet_search_error() { let server = Server::new().await; @@ -157,3 +179,24 @@ async fn facet_search_dont_support_words() { assert_eq!(code, 200, "{}", response); assert_eq!(response["facetHits"].as_array().unwrap().len(), 0); } + +#[actix_rt::test] +async fn simple_facet_search_with_sort_by_count() { + let server = Server::new().await; + let index = server.index("test"); + + let documents = DOCUMENTS.clone(); + index.update_settings_faceting(json!({ "sortFacetValuesBy": { "*": "count" } })).await; + index.update_settings_filterable_attributes(json!(["genres"])).await; + index.add_documents(documents, None).await; + index.wait_task(2).await; + + let (response, code) = + index.facet_search(json!({"facetName": "genres", "facetQuery": "a"})).await; + + assert_eq!(code, 200, "{}", response); + let hits = response["facetHits"].as_array().unwrap(); + assert_eq!(hits.len(), 2); + assert_eq!(hits[0], json!({ "value": "Action", "count": 3 })); + assert_eq!(hits[1], json!({ "value": "Adventure", "count": 2 })); +} diff --git a/milli/src/index.rs b/milli/src/index.rs index 6ad39dcb1..2c3977403 100644 --- a/milli/src/index.rs +++ b/milli/src/index.rs @@ -20,13 +20,13 @@ use crate::heed_codec::facet::{ use crate::heed_codec::{ BEU16StrCodec, FstSetCodec, ScriptLanguageCodec, StrBEU16Codec, StrRefCodec, }; +use crate::order_by_map::OrderByMap; use crate::proximity::ProximityPrecision; use crate::vector::EmbeddingConfig; use crate::{ default_criteria, CboRoaringBitmapCodec, Criterion, DocumentId, ExternalDocumentsIds, FacetDistribution, FieldDistribution, FieldId, FieldIdWordCountCodec, GeoPoint, ObkvCodec, - OrderBy, Result, RoaringBitmapCodec, RoaringBitmapLenCodec, Search, U8StrStrCodec, BEU16, - BEU32, BEU64, + Result, RoaringBitmapCodec, RoaringBitmapLenCodec, Search, U8StrStrCodec, BEU16, BEU32, BEU64, }; pub const DEFAULT_MIN_WORD_LEN_ONE_TYPO: u8 = 5; @@ -1373,21 +1373,19 @@ impl Index { self.main.remap_key_type::().delete(txn, main_key::MAX_VALUES_PER_FACET) } - pub fn sort_facet_values_by(&self, txn: &RoTxn) -> heed::Result> { - let mut orders = self + pub fn sort_facet_values_by(&self, txn: &RoTxn) -> heed::Result { + let orders = self .main - .remap_types::>>() + .remap_types::>() .get(txn, main_key::SORT_FACET_VALUES_BY)? .unwrap_or_default(); - // Insert the default ordering if it is not already overwritten by the user. - orders.entry("*".to_string()).or_insert(OrderBy::Lexicographic); Ok(orders) } pub(crate) fn put_sort_facet_values_by( &self, txn: &mut RwTxn, - val: &HashMap, + val: &OrderByMap, ) -> heed::Result<()> { self.main.remap_types::>().put(txn, main_key::SORT_FACET_VALUES_BY, &val) } diff --git a/milli/src/lib.rs b/milli/src/lib.rs index f6b398304..5effcea3d 100644 --- a/milli/src/lib.rs +++ b/milli/src/lib.rs @@ -16,6 +16,7 @@ pub mod facet; mod fields_ids_map; pub mod heed_codec; pub mod index; +pub mod order_by_map; pub mod prompt; pub mod proximity; pub mod score_details; @@ -56,10 +57,10 @@ pub use self::heed_codec::{ UncheckedU8StrStrCodec, }; pub use self::index::Index; +pub use self::search::facet::{FacetValueHit, SearchForFacetValues}; pub use self::search::{ - FacetDistribution, FacetValueHit, Filter, FormatOptions, MatchBounds, MatcherBuilder, - MatchingWords, OrderBy, Search, SearchForFacetValues, SearchResult, TermsMatchingStrategy, - DEFAULT_VALUES_PER_FACET, + FacetDistribution, Filter, FormatOptions, MatchBounds, MatcherBuilder, MatchingWords, OrderBy, + Search, SearchResult, TermsMatchingStrategy, DEFAULT_VALUES_PER_FACET, }; pub type Result = std::result::Result; diff --git a/milli/src/order_by_map.rs b/milli/src/order_by_map.rs new file mode 100644 index 000000000..287e62c3a --- /dev/null +++ b/milli/src/order_by_map.rs @@ -0,0 +1,57 @@ +use std::collections::{hash_map, HashMap}; +use std::iter::FromIterator; + +use serde::{Deserialize, Deserializer, Serialize}; + +use crate::OrderBy; + +#[derive(Serialize)] +pub struct OrderByMap(HashMap); + +impl OrderByMap { + pub fn get(&self, key: impl AsRef) -> OrderBy { + self.0 + .get(key.as_ref()) + .copied() + .unwrap_or_else(|| self.0.get("*").copied().unwrap_or_default()) + } + + pub fn insert(&mut self, key: String, value: OrderBy) -> Option { + self.0.insert(key, value) + } +} + +impl Default for OrderByMap { + fn default() -> Self { + let mut map = HashMap::new(); + map.insert("*".to_string(), OrderBy::Lexicographic); + OrderByMap(map) + } +} + +impl FromIterator<(String, OrderBy)> for OrderByMap { + fn from_iter>(iter: T) -> Self { + OrderByMap(iter.into_iter().collect()) + } +} + +impl IntoIterator for OrderByMap { + type Item = (String, OrderBy); + type IntoIter = hash_map::IntoIter; + + fn into_iter(self) -> Self::IntoIter { + self.0.into_iter() + } +} + +impl<'de> Deserialize<'de> for OrderByMap { + fn deserialize(deserializer: D) -> Result + where + D: Deserializer<'de>, + { + let mut map = Deserialize::deserialize(deserializer).map(OrderByMap)?; + // Insert the default ordering if it is not already overwritten by the user. + map.0.entry("*".to_string()).or_insert(OrderBy::default()); + Ok(map) + } +} diff --git a/milli/src/search/facet/facet_range_search.rs b/milli/src/search/facet/facet_range_search.rs index f1a26ded5..e340fbac5 100644 --- a/milli/src/search/facet/facet_range_search.rs +++ b/milli/src/search/facet/facet_range_search.rs @@ -168,7 +168,7 @@ impl<'t, 'b, 'bitmap> FacetRangeSearch<'t, 'b, 'bitmap> { } // should we stop? - // We should if the the search range doesn't include any + // We should if the search range doesn't include any // element from the previous key or its successors let should_stop = { match self.right { @@ -232,7 +232,7 @@ impl<'t, 'b, 'bitmap> FacetRangeSearch<'t, 'b, 'bitmap> { } // should we stop? - // We should if the the search range doesn't include any + // We should if the search range doesn't include any // element from the previous key or its successors let should_stop = { match self.right { diff --git a/milli/src/search/facet/mod.rs b/milli/src/search/facet/mod.rs index f676ee109..34a9cdcb8 100644 --- a/milli/src/search/facet/mod.rs +++ b/milli/src/search/facet/mod.rs @@ -6,15 +6,18 @@ use roaring::RoaringBitmap; pub use self::facet_distribution::{FacetDistribution, OrderBy, DEFAULT_VALUES_PER_FACET}; pub use self::filter::{BadGeoError, Filter}; +pub use self::search::{FacetValueHit, SearchForFacetValues}; use crate::heed_codec::facet::{FacetGroupKeyCodec, FacetGroupValueCodec, OrderedF64Codec}; use crate::heed_codec::BytesRefCodec; use crate::{Index, Result}; + mod facet_distribution; mod facet_distribution_iter; mod facet_range_search; mod facet_sort_ascending; mod facet_sort_descending; mod filter; +mod search; fn facet_extreme_value<'t>( mut extreme_it: impl Iterator> + 't, diff --git a/milli/src/search/facet/search.rs b/milli/src/search/facet/search.rs new file mode 100644 index 000000000..0251d6b8d --- /dev/null +++ b/milli/src/search/facet/search.rs @@ -0,0 +1,326 @@ +use std::cmp::{Ordering, Reverse}; +use std::collections::BinaryHeap; +use std::ops::ControlFlow; + +use charabia::normalizer::NormalizerOption; +use charabia::Normalize; +use fst::automaton::{Automaton, Str}; +use fst::{IntoStreamer, Streamer}; +use roaring::RoaringBitmap; +use tracing::error; + +use crate::error::UserError; +use crate::heed_codec::facet::{FacetGroupKey, FacetGroupValue}; +use crate::search::build_dfa; +use crate::{DocumentId, FieldId, OrderBy, Result, Search}; + +/// The maximum number of values per facet returned by the facet search route. +const DEFAULT_MAX_NUMBER_OF_VALUES_PER_FACET: usize = 100; + +pub struct SearchForFacetValues<'a> { + query: Option, + facet: String, + search_query: Search<'a>, + max_values: usize, + is_hybrid: bool, +} + +impl<'a> SearchForFacetValues<'a> { + pub fn new( + facet: String, + search_query: Search<'a>, + is_hybrid: bool, + ) -> SearchForFacetValues<'a> { + SearchForFacetValues { + query: None, + facet, + search_query, + max_values: DEFAULT_MAX_NUMBER_OF_VALUES_PER_FACET, + is_hybrid, + } + } + + pub fn query(&mut self, query: impl Into) -> &mut Self { + self.query = Some(query.into()); + self + } + + pub fn max_values(&mut self, max: usize) -> &mut Self { + self.max_values = max; + self + } + + fn one_original_value_of( + &self, + field_id: FieldId, + facet_str: &str, + any_docid: DocumentId, + ) -> Result> { + let index = self.search_query.index; + let rtxn = self.search_query.rtxn; + let key: (FieldId, _, &str) = (field_id, any_docid, facet_str); + Ok(index.field_id_docid_facet_strings.get(rtxn, &key)?.map(|v| v.to_owned())) + } + + pub fn execute(&self) -> Result> { + let index = self.search_query.index; + let rtxn = self.search_query.rtxn; + + let filterable_fields = index.filterable_fields(rtxn)?; + if !filterable_fields.contains(&self.facet) { + let (valid_fields, hidden_fields) = + index.remove_hidden_fields(rtxn, filterable_fields)?; + + return Err(UserError::InvalidFacetSearchFacetName { + field: self.facet.clone(), + valid_fields, + hidden_fields, + } + .into()); + } + + let fields_ids_map = index.fields_ids_map(rtxn)?; + let fid = match fields_ids_map.id(&self.facet) { + Some(fid) => fid, + // we return an empty list of results when the attribute has been + // set as filterable but no document contains this field (yet). + None => return Ok(Vec::new()), + }; + + let fst = match self.search_query.index.facet_id_string_fst.get(rtxn, &fid)? { + Some(fst) => fst, + None => return Ok(Vec::new()), + }; + + let search_candidates = self + .search_query + .execute_for_candidates(self.is_hybrid || self.search_query.vector.is_some())?; + + let mut results = match index.sort_facet_values_by(rtxn)?.get(&self.facet) { + OrderBy::Lexicographic => ValuesCollection::by_lexicographic(self.max_values), + OrderBy::Count => ValuesCollection::by_count(self.max_values), + }; + + match self.query.as_ref() { + Some(query) => { + let options = NormalizerOption { lossy: true, ..Default::default() }; + let query = query.normalize(&options); + let query = query.as_ref(); + + let authorize_typos = self.search_query.index.authorize_typos(rtxn)?; + let field_authorizes_typos = + !self.search_query.index.exact_attributes_ids(rtxn)?.contains(&fid); + + if authorize_typos && field_authorizes_typos { + let exact_words_fst = self.search_query.index.exact_words(rtxn)?; + if exact_words_fst.map_or(false, |fst| fst.contains(query)) { + if fst.contains(query) { + self.fetch_original_facets_using_normalized( + fid, + query, + query, + &search_candidates, + &mut results, + )?; + } + } else { + let one_typo = self.search_query.index.min_word_len_one_typo(rtxn)?; + let two_typos = self.search_query.index.min_word_len_two_typos(rtxn)?; + + let is_prefix = true; + let automaton = if query.len() < one_typo as usize { + build_dfa(query, 0, is_prefix) + } else if query.len() < two_typos as usize { + build_dfa(query, 1, is_prefix) + } else { + build_dfa(query, 2, is_prefix) + }; + + let mut stream = fst.search(automaton).into_stream(); + while let Some(facet_value) = stream.next() { + let value = std::str::from_utf8(facet_value)?; + if self + .fetch_original_facets_using_normalized( + fid, + value, + query, + &search_candidates, + &mut results, + )? + .is_break() + { + break; + } + } + } + } else { + let automaton = Str::new(query).starts_with(); + let mut stream = fst.search(automaton).into_stream(); + while let Some(facet_value) = stream.next() { + let value = std::str::from_utf8(facet_value)?; + if self + .fetch_original_facets_using_normalized( + fid, + value, + query, + &search_candidates, + &mut results, + )? + .is_break() + { + break; + } + } + } + } + None => { + let prefix = FacetGroupKey { field_id: fid, level: 0, left_bound: "" }; + for result in index.facet_id_string_docids.prefix_iter(rtxn, &prefix)? { + let (FacetGroupKey { left_bound, .. }, FacetGroupValue { bitmap, .. }) = + result?; + let count = search_candidates.intersection_len(&bitmap); + if count != 0 { + let value = self + .one_original_value_of(fid, left_bound, bitmap.min().unwrap())? + .unwrap_or_else(|| left_bound.to_string()); + if results.insert(FacetValueHit { value, count }).is_break() { + break; + } + } + } + } + } + + Ok(results.into_sorted_vec()) + } + + fn fetch_original_facets_using_normalized( + &self, + fid: FieldId, + value: &str, + query: &str, + search_candidates: &RoaringBitmap, + results: &mut ValuesCollection, + ) -> Result> { + let index = self.search_query.index; + let rtxn = self.search_query.rtxn; + + let database = index.facet_id_normalized_string_strings; + let key = (fid, value); + let original_strings = match database.get(rtxn, &key)? { + Some(original_strings) => original_strings, + None => { + error!("the facet value is missing from the facet database: {key:?}"); + return Ok(ControlFlow::Continue(())); + } + }; + for original in original_strings { + let key = FacetGroupKey { field_id: fid, level: 0, left_bound: original.as_str() }; + let docids = match index.facet_id_string_docids.get(rtxn, &key)? { + Some(FacetGroupValue { bitmap, .. }) => bitmap, + None => { + error!("the facet value is missing from the facet database: {key:?}"); + return Ok(ControlFlow::Continue(())); + } + }; + let count = search_candidates.intersection_len(&docids); + if count != 0 { + let value = self + .one_original_value_of(fid, &original, docids.min().unwrap())? + .unwrap_or_else(|| query.to_string()); + if results.insert(FacetValueHit { value, count }).is_break() { + break; + } + } + } + + Ok(ControlFlow::Continue(())) + } +} + +#[derive(Debug, Clone, serde::Serialize, PartialEq)] +pub struct FacetValueHit { + /// The original facet value + pub value: String, + /// The number of documents associated to this facet + pub count: u64, +} + +impl PartialOrd for FacetValueHit { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } +} + +impl Ord for FacetValueHit { + fn cmp(&self, other: &Self) -> Ordering { + self.count.cmp(&other.count).then_with(|| self.value.cmp(&other.value)) + } +} + +impl Eq for FacetValueHit {} + +/// A wrapper type that collects the best facet values by +/// lexicographic or number of associated values. +enum ValuesCollection { + /// Keeps the top values according to the lexicographic order. + Lexicographic { max: usize, content: Vec }, + /// Keeps the top values according to the number of values associated to them. + /// + /// Note that it is a max heap and we need to move the smallest counts + /// at the top to be able to pop them when we reach the max_values limit. + Count { max: usize, content: BinaryHeap> }, +} + +impl ValuesCollection { + pub fn by_lexicographic(max: usize) -> Self { + ValuesCollection::Lexicographic { max, content: Vec::new() } + } + + pub fn by_count(max: usize) -> Self { + ValuesCollection::Count { max, content: BinaryHeap::new() } + } + + pub fn insert(&mut self, value: FacetValueHit) -> ControlFlow<()> { + match self { + ValuesCollection::Lexicographic { max, content } => { + if content.len() < *max { + content.push(value); + if content.len() < *max { + return ControlFlow::Continue(()); + } + } + ControlFlow::Break(()) + } + ValuesCollection::Count { max, content } => { + if content.len() == *max { + // Peeking gives us the worst value in the list as + // this is a max-heap and we reversed it. + let Some(mut peek) = content.peek_mut() else { return ControlFlow::Break(()) }; + if peek.0.count <= value.count { + // Replace the current worst value in the heap + // with the new one we received that is better. + *peek = Reverse(value); + } + } else { + content.push(Reverse(value)); + } + ControlFlow::Continue(()) + } + } + } + + /// Returns the list of facet values in descending order of, either, + /// count or lexicographic order of the value depending on the type. + pub fn into_sorted_vec(self) -> Vec { + match self { + ValuesCollection::Lexicographic { content, .. } => content.into_iter().collect(), + ValuesCollection::Count { content, .. } => { + // Convert the heap into a vec of hits by removing the Reverse wrapper. + // Hits are already in the right order as they were reversed and there + // are output in ascending order. + content.into_sorted_vec().into_iter().map(|Reverse(hit)| hit).collect() + } + } + } +} diff --git a/milli/src/search/mod.rs b/milli/src/search/mod.rs index e411bd032..dc8354486 100644 --- a/milli/src/search/mod.rs +++ b/milli/src/search/mod.rs @@ -1,25 +1,17 @@ use std::fmt; -use std::ops::ControlFlow; -use charabia::normalizer::NormalizerOption; -use charabia::Normalize; -use fst::automaton::{Automaton, Str}; -use fst::{IntoStreamer, Streamer}; use levenshtein_automata::{LevenshteinAutomatonBuilder as LevBuilder, DFA}; use once_cell::sync::Lazy; use roaring::bitmap::RoaringBitmap; -use tracing::error; pub use self::facet::{FacetDistribution, Filter, OrderBy, DEFAULT_VALUES_PER_FACET}; pub use self::new::matches::{FormatOptions, MatchBounds, MatcherBuilder, MatchingWords}; use self::new::{execute_vector_search, PartialSearchResult}; -use crate::error::UserError; -use crate::heed_codec::facet::{FacetGroupKey, FacetGroupValue}; use crate::score_details::{ScoreDetails, ScoringStrategy}; use crate::vector::DistributionShift; use crate::{ - execute_search, filtered_universe, AscDesc, DefaultSearchLogger, DocumentId, FieldId, Index, - Result, SearchContext, + execute_search, filtered_universe, AscDesc, DefaultSearchLogger, DocumentId, Index, Result, + SearchContext, }; // Building these factories is not free. @@ -27,9 +19,6 @@ static LEVDIST0: Lazy = Lazy::new(|| LevBuilder::new(0, true)); static LEVDIST1: Lazy = Lazy::new(|| LevBuilder::new(1, true)); static LEVDIST2: Lazy = Lazy::new(|| LevBuilder::new(2, true)); -/// The maximum number of values per facet returned by the facet search route. -const DEFAULT_MAX_NUMBER_OF_VALUES_PER_FACET: usize = 100; - pub mod facet; mod fst_utils; pub mod hybrid; @@ -302,240 +291,6 @@ pub fn build_dfa(word: &str, typos: u8, is_prefix: bool) -> DFA { } } -pub struct SearchForFacetValues<'a> { - query: Option, - facet: String, - search_query: Search<'a>, - max_values: usize, - is_hybrid: bool, -} - -impl<'a> SearchForFacetValues<'a> { - pub fn new( - facet: String, - search_query: Search<'a>, - is_hybrid: bool, - ) -> SearchForFacetValues<'a> { - SearchForFacetValues { - query: None, - facet, - search_query, - max_values: DEFAULT_MAX_NUMBER_OF_VALUES_PER_FACET, - is_hybrid, - } - } - - pub fn query(&mut self, query: impl Into) -> &mut Self { - self.query = Some(query.into()); - self - } - - pub fn max_values(&mut self, max: usize) -> &mut Self { - self.max_values = max; - self - } - - fn one_original_value_of( - &self, - field_id: FieldId, - facet_str: &str, - any_docid: DocumentId, - ) -> Result> { - let index = self.search_query.index; - let rtxn = self.search_query.rtxn; - let key: (FieldId, _, &str) = (field_id, any_docid, facet_str); - Ok(index.field_id_docid_facet_strings.get(rtxn, &key)?.map(|v| v.to_owned())) - } - - pub fn execute(&self) -> Result> { - let index = self.search_query.index; - let rtxn = self.search_query.rtxn; - - let filterable_fields = index.filterable_fields(rtxn)?; - if !filterable_fields.contains(&self.facet) { - let (valid_fields, hidden_fields) = - index.remove_hidden_fields(rtxn, filterable_fields)?; - - return Err(UserError::InvalidFacetSearchFacetName { - field: self.facet.clone(), - valid_fields, - hidden_fields, - } - .into()); - } - - let fields_ids_map = index.fields_ids_map(rtxn)?; - let fid = match fields_ids_map.id(&self.facet) { - Some(fid) => fid, - // we return an empty list of results when the attribute has been - // set as filterable but no document contains this field (yet). - None => return Ok(Vec::new()), - }; - - let fst = match self.search_query.index.facet_id_string_fst.get(rtxn, &fid)? { - Some(fst) => fst, - None => return Ok(vec![]), - }; - - let search_candidates = self - .search_query - .execute_for_candidates(self.is_hybrid || self.search_query.vector.is_some())?; - - match self.query.as_ref() { - Some(query) => { - let options = NormalizerOption { lossy: true, ..Default::default() }; - let query = query.normalize(&options); - let query = query.as_ref(); - - let authorize_typos = self.search_query.index.authorize_typos(rtxn)?; - let field_authorizes_typos = - !self.search_query.index.exact_attributes_ids(rtxn)?.contains(&fid); - - if authorize_typos && field_authorizes_typos { - let exact_words_fst = self.search_query.index.exact_words(rtxn)?; - if exact_words_fst.map_or(false, |fst| fst.contains(query)) { - let mut results = vec![]; - if fst.contains(query) { - self.fetch_original_facets_using_normalized( - fid, - query, - query, - &search_candidates, - &mut results, - )?; - } - Ok(results) - } else { - let one_typo = self.search_query.index.min_word_len_one_typo(rtxn)?; - let two_typos = self.search_query.index.min_word_len_two_typos(rtxn)?; - - let is_prefix = true; - let automaton = if query.len() < one_typo as usize { - build_dfa(query, 0, is_prefix) - } else if query.len() < two_typos as usize { - build_dfa(query, 1, is_prefix) - } else { - build_dfa(query, 2, is_prefix) - }; - - let mut stream = fst.search(automaton).into_stream(); - let mut results = vec![]; - while let Some(facet_value) = stream.next() { - let value = std::str::from_utf8(facet_value)?; - if self - .fetch_original_facets_using_normalized( - fid, - value, - query, - &search_candidates, - &mut results, - )? - .is_break() - { - break; - } - } - - Ok(results) - } - } else { - let automaton = Str::new(query).starts_with(); - let mut stream = fst.search(automaton).into_stream(); - let mut results = vec![]; - while let Some(facet_value) = stream.next() { - let value = std::str::from_utf8(facet_value)?; - if self - .fetch_original_facets_using_normalized( - fid, - value, - query, - &search_candidates, - &mut results, - )? - .is_break() - { - break; - } - } - - Ok(results) - } - } - None => { - let mut results = vec![]; - let prefix = FacetGroupKey { field_id: fid, level: 0, left_bound: "" }; - for result in index.facet_id_string_docids.prefix_iter(rtxn, &prefix)? { - let (FacetGroupKey { left_bound, .. }, FacetGroupValue { bitmap, .. }) = - result?; - let count = search_candidates.intersection_len(&bitmap); - if count != 0 { - let value = self - .one_original_value_of(fid, left_bound, bitmap.min().unwrap())? - .unwrap_or_else(|| left_bound.to_string()); - results.push(FacetValueHit { value, count }); - } - if results.len() >= self.max_values { - break; - } - } - Ok(results) - } - } - } - - fn fetch_original_facets_using_normalized( - &self, - fid: FieldId, - value: &str, - query: &str, - search_candidates: &RoaringBitmap, - results: &mut Vec, - ) -> Result> { - let index = self.search_query.index; - let rtxn = self.search_query.rtxn; - - let database = index.facet_id_normalized_string_strings; - let key = (fid, value); - let original_strings = match database.get(rtxn, &key)? { - Some(original_strings) => original_strings, - None => { - error!("the facet value is missing from the facet database: {key:?}"); - return Ok(ControlFlow::Continue(())); - } - }; - for original in original_strings { - let key = FacetGroupKey { field_id: fid, level: 0, left_bound: original.as_str() }; - let docids = match index.facet_id_string_docids.get(rtxn, &key)? { - Some(FacetGroupValue { bitmap, .. }) => bitmap, - None => { - error!("the facet value is missing from the facet database: {key:?}"); - return Ok(ControlFlow::Continue(())); - } - }; - let count = search_candidates.intersection_len(&docids); - if count != 0 { - let value = self - .one_original_value_of(fid, &original, docids.min().unwrap())? - .unwrap_or_else(|| query.to_string()); - results.push(FacetValueHit { value, count }); - } - if results.len() >= self.max_values { - return Ok(ControlFlow::Break(())); - } - } - - Ok(ControlFlow::Continue(())) - } -} - -#[derive(Debug, Clone, serde::Serialize, PartialEq)] -pub struct FacetValueHit { - /// The original facet value - pub value: String, - /// The number of documents associated to this facet - pub count: u64, -} - #[cfg(test)] mod test { #[allow(unused_imports)] diff --git a/milli/src/search/new/tests/typo_proximity.rs b/milli/src/search/new/tests/typo_proximity.rs index 8dd110704..e71d32331 100644 --- a/milli/src/search/new/tests/typo_proximity.rs +++ b/milli/src/search/new/tests/typo_proximity.rs @@ -5,7 +5,7 @@ The typo ranking rule should transform the query graph such that it only contain the combinations of word derivations that it used to compute its bucket. The proximity ranking rule should then look for proximities only between those specific derivations. -For example, given the the search query `beautiful summer` and the dataset: +For example, given the search query `beautiful summer` and the dataset: ```text { "id": 0, "text": "beautigul summer...... beautiful day in the summer" } { "id": 1, "text": "beautiful summer" } diff --git a/milli/src/update/settings.rs b/milli/src/update/settings.rs index 2f53718ac..46014202b 100644 --- a/milli/src/update/settings.rs +++ b/milli/src/update/settings.rs @@ -14,12 +14,13 @@ use super::IndexerConfig; use crate::criterion::Criterion; use crate::error::UserError; use crate::index::{DEFAULT_MIN_WORD_LEN_ONE_TYPO, DEFAULT_MIN_WORD_LEN_TWO_TYPOS}; +use crate::order_by_map::OrderByMap; use crate::proximity::ProximityPrecision; use crate::update::index_documents::IndexDocumentsMethod; use crate::update::{IndexDocuments, UpdateIndexingStep}; use crate::vector::settings::{check_set, check_unset, EmbedderSource, EmbeddingSettings}; use crate::vector::{Embedder, EmbeddingConfig, EmbeddingConfigs}; -use crate::{FieldsIdsMap, Index, OrderBy, Result}; +use crate::{FieldsIdsMap, Index, Result}; #[derive(Debug, Clone, PartialEq, Eq, Copy)] pub enum Setting { @@ -145,7 +146,7 @@ pub struct Settings<'a, 't, 'i> { /// Attributes on which typo tolerance is disabled. exact_attributes: Setting>, max_values_per_facet: Setting, - sort_facet_values_by: Setting>, + sort_facet_values_by: Setting, pagination_max_total_hits: Setting, proximity_precision: Setting, embedder_settings: Setting>>, @@ -340,7 +341,7 @@ impl<'a, 't, 'i> Settings<'a, 't, 'i> { self.max_values_per_facet = Setting::Reset; } - pub fn set_sort_facet_values_by(&mut self, value: HashMap) { + pub fn set_sort_facet_values_by(&mut self, value: OrderByMap) { self.sort_facet_values_by = Setting::Set(value); } @@ -1186,6 +1187,13 @@ pub fn validate_embedding_settings( } } } + EmbedderSource::Ollama => { + // Dimensions get inferred, only model name is required + check_unset(&dimensions, "dimensions", inferred_source, name)?; + check_set(&model, "model", inferred_source, name)?; + check_unset(&api_key, "apiKey", inferred_source, name)?; + check_unset(&revision, "revision", inferred_source, name)?; + } EmbedderSource::HuggingFace => { check_unset(&api_key, "apiKey", inferred_source, name)?; check_unset(&dimensions, "dimensions", inferred_source, name)?; diff --git a/milli/src/vector/error.rs b/milli/src/vector/error.rs index fbe4ee878..9bbdeaa90 100644 --- a/milli/src/vector/error.rs +++ b/milli/src/vector/error.rs @@ -2,6 +2,7 @@ use std::path::PathBuf; use hf_hub::api::sync::ApiError; +use super::ollama::OllamaError; use crate::error::FaultSource; use crate::vector::openai::OpenAiError; @@ -71,6 +72,17 @@ pub enum EmbedErrorKind { OpenAiRuntimeInit(std::io::Error), #[error("initializing web client for sending embedding requests failed: {0}")] InitWebClient(reqwest::Error), + // Dedicated Ollama error kinds, might have to merge them into one cohesive error type for all backends. + #[error("unexpected response from Ollama: {0}")] + OllamaUnexpected(reqwest::Error), + #[error("sent too many requests to Ollama: {0}")] + OllamaTooManyRequests(OllamaError), + #[error("received internal error from Ollama: {0}")] + OllamaInternalServerError(OllamaError), + #[error("model not found. Meilisearch will not automatically download models from the Ollama library, please pull the model manually: {0}")] + OllamaModelNotFoundError(OllamaError), + #[error("received unhandled HTTP status code {0} from Ollama")] + OllamaUnhandledStatusCode(u16), } impl EmbedError { @@ -129,6 +141,26 @@ impl EmbedError { pub fn openai_initialize_web_client(inner: reqwest::Error) -> Self { Self { kind: EmbedErrorKind::InitWebClient(inner), fault: FaultSource::Runtime } } + + pub(crate) fn ollama_unexpected(inner: reqwest::Error) -> EmbedError { + Self { kind: EmbedErrorKind::OllamaUnexpected(inner), fault: FaultSource::Bug } + } + + pub(crate) fn ollama_model_not_found(inner: OllamaError) -> EmbedError { + Self { kind: EmbedErrorKind::OllamaModelNotFoundError(inner), fault: FaultSource::User } + } + + pub(crate) fn ollama_too_many_requests(inner: OllamaError) -> EmbedError { + Self { kind: EmbedErrorKind::OllamaTooManyRequests(inner), fault: FaultSource::Runtime } + } + + pub(crate) fn ollama_internal_server_error(inner: OllamaError) -> EmbedError { + Self { kind: EmbedErrorKind::OllamaInternalServerError(inner), fault: FaultSource::Runtime } + } + + pub(crate) fn ollama_unhandled_status_code(code: u16) -> EmbedError { + Self { kind: EmbedErrorKind::OllamaUnhandledStatusCode(code), fault: FaultSource::Bug } + } } #[derive(Debug, thiserror::Error)] @@ -195,6 +227,13 @@ impl NewEmbedderError { } } + pub fn ollama_could_not_determine_dimension(inner: EmbedError) -> NewEmbedderError { + Self { + kind: NewEmbedderErrorKind::CouldNotDetermineDimension(inner), + fault: FaultSource::User, + } + } + pub fn openai_invalid_api_key_format(inner: reqwest::header::InvalidHeaderValue) -> Self { Self { kind: NewEmbedderErrorKind::InvalidApiKeyFormat(inner), fault: FaultSource::User } } diff --git a/milli/src/vector/mod.rs b/milli/src/vector/mod.rs index 6aa324da9..035ac555e 100644 --- a/milli/src/vector/mod.rs +++ b/milli/src/vector/mod.rs @@ -10,6 +10,8 @@ pub mod manual; pub mod openai; pub mod settings; +pub mod ollama; + pub use self::error::Error; pub type Embedding = Vec; @@ -76,6 +78,7 @@ pub enum Embedder { HuggingFace(hf::Embedder), OpenAi(openai::Embedder), UserProvided(manual::Embedder), + Ollama(ollama::Embedder), } #[derive(Debug, Clone, Default, serde::Deserialize, serde::Serialize)] @@ -127,6 +130,7 @@ impl IntoIterator for EmbeddingConfigs { pub enum EmbedderOptions { HuggingFace(hf::EmbedderOptions), OpenAi(openai::EmbedderOptions), + Ollama(ollama::EmbedderOptions), UserProvided(manual::EmbedderOptions), } @@ -144,6 +148,10 @@ impl EmbedderOptions { pub fn openai(api_key: Option) -> Self { Self::OpenAi(openai::EmbedderOptions::with_default_model(api_key)) } + + pub fn ollama() -> Self { + Self::Ollama(ollama::EmbedderOptions::with_default_model()) + } } impl Embedder { @@ -151,6 +159,7 @@ impl Embedder { Ok(match options { EmbedderOptions::HuggingFace(options) => Self::HuggingFace(hf::Embedder::new(options)?), EmbedderOptions::OpenAi(options) => Self::OpenAi(openai::Embedder::new(options)?), + EmbedderOptions::Ollama(options) => Self::Ollama(ollama::Embedder::new(options)?), EmbedderOptions::UserProvided(options) => { Self::UserProvided(manual::Embedder::new(options)) } @@ -167,6 +176,10 @@ impl Embedder { let client = embedder.new_client()?; embedder.embed(texts, &client).await } + Embedder::Ollama(embedder) => { + let client = embedder.new_client()?; + embedder.embed(texts, &client).await + } Embedder::UserProvided(embedder) => embedder.embed(texts), } } @@ -181,6 +194,7 @@ impl Embedder { match self { Embedder::HuggingFace(embedder) => embedder.embed_chunks(text_chunks), Embedder::OpenAi(embedder) => embedder.embed_chunks(text_chunks), + Embedder::Ollama(embedder) => embedder.embed_chunks(text_chunks), Embedder::UserProvided(embedder) => embedder.embed_chunks(text_chunks), } } @@ -189,6 +203,7 @@ impl Embedder { match self { Embedder::HuggingFace(embedder) => embedder.chunk_count_hint(), Embedder::OpenAi(embedder) => embedder.chunk_count_hint(), + Embedder::Ollama(embedder) => embedder.chunk_count_hint(), Embedder::UserProvided(_) => 1, } } @@ -197,6 +212,7 @@ impl Embedder { match self { Embedder::HuggingFace(embedder) => embedder.prompt_count_in_chunk_hint(), Embedder::OpenAi(embedder) => embedder.prompt_count_in_chunk_hint(), + Embedder::Ollama(embedder) => embedder.prompt_count_in_chunk_hint(), Embedder::UserProvided(_) => 1, } } @@ -205,6 +221,7 @@ impl Embedder { match self { Embedder::HuggingFace(embedder) => embedder.dimensions(), Embedder::OpenAi(embedder) => embedder.dimensions(), + Embedder::Ollama(embedder) => embedder.dimensions(), Embedder::UserProvided(embedder) => embedder.dimensions(), } } @@ -213,6 +230,7 @@ impl Embedder { match self { Embedder::HuggingFace(embedder) => embedder.distribution(), Embedder::OpenAi(embedder) => embedder.distribution(), + Embedder::Ollama(embedder) => embedder.distribution(), Embedder::UserProvided(_embedder) => None, } } diff --git a/milli/src/vector/ollama.rs b/milli/src/vector/ollama.rs new file mode 100644 index 000000000..76988f70b --- /dev/null +++ b/milli/src/vector/ollama.rs @@ -0,0 +1,307 @@ +// Copied from "openai.rs" with the sections I actually understand changed for Ollama. +// The common components of the Ollama and OpenAI interfaces might need to be extracted. + +use std::fmt::Display; + +use reqwest::StatusCode; + +use super::error::{EmbedError, NewEmbedderError}; +use super::openai::Retry; +use super::{DistributionShift, Embedding, Embeddings}; + +#[derive(Debug)] +pub struct Embedder { + headers: reqwest::header::HeaderMap, + options: EmbedderOptions, +} + +#[derive(Debug, Clone, Hash, PartialEq, Eq, serde::Deserialize, serde::Serialize)] +pub struct EmbedderOptions { + pub embedding_model: EmbeddingModel, +} + +#[derive( + Debug, Clone, Hash, PartialEq, Eq, serde::Serialize, serde::Deserialize, deserr::Deserr, +)] +#[deserr(deny_unknown_fields)] +pub struct EmbeddingModel { + name: String, + dimensions: usize, +} + +#[derive(Debug, serde::Serialize)] +struct OllamaRequest<'a> { + model: &'a str, + prompt: &'a str, +} + +#[derive(Debug, serde::Deserialize)] +struct OllamaResponse { + embedding: Embedding, +} + +#[derive(Debug, serde::Deserialize)] +pub struct OllamaError { + error: String, +} + +impl EmbeddingModel { + pub fn max_token(&self) -> usize { + // this might not be the same for all models + 8192 + } + + pub fn default_dimensions(&self) -> usize { + // Dimensions for nomic-embed-text + 768 + } + + pub fn name(&self) -> String { + self.name.clone() + } + + pub fn from_name(name: &str) -> Self { + Self { name: name.to_string(), dimensions: 0 } + } + + pub fn supports_overriding_dimensions(&self) -> bool { + false + } +} + +impl Default for EmbeddingModel { + fn default() -> Self { + Self { name: "nomic-embed-text".to_string(), dimensions: 0 } + } +} + +impl EmbedderOptions { + pub fn with_default_model() -> Self { + Self { embedding_model: Default::default() } + } + + pub fn with_embedding_model(embedding_model: EmbeddingModel) -> Self { + Self { embedding_model } + } +} + +impl Embedder { + pub fn new_client(&self) -> Result { + reqwest::ClientBuilder::new() + .default_headers(self.headers.clone()) + .build() + .map_err(EmbedError::openai_initialize_web_client) + } + + pub fn new(options: EmbedderOptions) -> Result { + let mut headers = reqwest::header::HeaderMap::new(); + headers.insert( + reqwest::header::CONTENT_TYPE, + reqwest::header::HeaderValue::from_static("application/json"), + ); + + let mut embedder = Self { options, headers }; + + let rt = tokio::runtime::Builder::new_current_thread() + .enable_io() + .enable_time() + .build() + .map_err(EmbedError::openai_runtime_init) + .map_err(NewEmbedderError::ollama_could_not_determine_dimension)?; + + // Get dimensions from Ollama + let request = + OllamaRequest { model: &embedder.options.embedding_model.name(), prompt: "test" }; + // TODO: Refactor into shared error type + let client = embedder + .new_client() + .map_err(NewEmbedderError::ollama_could_not_determine_dimension)?; + + rt.block_on(async move { + let response = client + .post(get_ollama_path()) + .json(&request) + .send() + .await + .map_err(EmbedError::ollama_unexpected) + .map_err(NewEmbedderError::ollama_could_not_determine_dimension)?; + + // Process error in case model not found + let response = Self::check_response(response).await.map_err(|_err| { + let e = EmbedError::ollama_model_not_found(OllamaError { + error: format!("model: {}", embedder.options.embedding_model.name()), + }); + NewEmbedderError::ollama_could_not_determine_dimension(e) + })?; + + let response: OllamaResponse = response + .json() + .await + .map_err(EmbedError::ollama_unexpected) + .map_err(NewEmbedderError::ollama_could_not_determine_dimension)?; + + let embedding = Embeddings::from_single_embedding(response.embedding); + + embedder.options.embedding_model.dimensions = embedding.dimension(); + + tracing::info!( + "ollama model {} with dimensionality {} added", + embedder.options.embedding_model.name(), + embedding.dimension() + ); + + Ok(embedder) + }) + } + + async fn check_response(response: reqwest::Response) -> Result { + if !response.status().is_success() { + // Not the same number of possible error cases covered as with OpenAI. + match response.status() { + StatusCode::TOO_MANY_REQUESTS => { + let error_response: OllamaError = response + .json() + .await + .map_err(EmbedError::ollama_unexpected) + .map_err(Retry::retry_later)?; + + return Err(Retry::rate_limited(EmbedError::ollama_too_many_requests( + OllamaError { error: error_response.error }, + ))); + } + StatusCode::SERVICE_UNAVAILABLE => { + let error_response: OllamaError = response + .json() + .await + .map_err(EmbedError::ollama_unexpected) + .map_err(Retry::retry_later)?; + return Err(Retry::retry_later(EmbedError::ollama_internal_server_error( + OllamaError { error: error_response.error }, + ))); + } + StatusCode::NOT_FOUND => { + let error_response: OllamaError = response + .json() + .await + .map_err(EmbedError::ollama_unexpected) + .map_err(Retry::give_up)?; + + return Err(Retry::give_up(EmbedError::ollama_model_not_found(OllamaError { + error: error_response.error, + }))); + } + code => { + return Err(Retry::give_up(EmbedError::ollama_unhandled_status_code( + code.as_u16(), + ))); + } + } + } + Ok(response) + } + + pub async fn embed( + &self, + texts: Vec, + client: &reqwest::Client, + ) -> Result>, EmbedError> { + // Ollama only embedds one document at a time. + let mut results = Vec::with_capacity(texts.len()); + + // The retry loop is inside the texts loop, might have to switch that around + for text in texts { + // Retries copied from openai.rs + for attempt in 0..7 { + let retry_duration = match self.try_embed(&text, client).await { + Ok(result) => { + results.push(result); + break; + } + Err(retry) => { + tracing::warn!("Failed: {}", retry.error); + retry.into_duration(attempt) + } + }?; + tracing::warn!( + "Attempt #{}, retrying after {}ms.", + attempt, + retry_duration.as_millis() + ); + tokio::time::sleep(retry_duration).await; + } + } + + Ok(results) + } + + async fn try_embed( + &self, + text: &str, + client: &reqwest::Client, + ) -> Result, Retry> { + let request = OllamaRequest { model: &self.options.embedding_model.name(), prompt: text }; + let response = client + .post(get_ollama_path()) + .json(&request) + .send() + .await + .map_err(EmbedError::openai_network) + .map_err(Retry::retry_later)?; + + let response = Self::check_response(response).await?; + + let response: OllamaResponse = response + .json() + .await + .map_err(EmbedError::openai_unexpected) + .map_err(Retry::retry_later)?; + + tracing::trace!("response: {:?}", response.embedding); + + let embedding = Embeddings::from_single_embedding(response.embedding); + Ok(embedding) + } + + pub fn embed_chunks( + &self, + text_chunks: Vec>, + ) -> Result>>, EmbedError> { + let rt = tokio::runtime::Builder::new_current_thread() + .enable_io() + .enable_time() + .build() + .map_err(EmbedError::openai_runtime_init)?; + let client = self.new_client()?; + rt.block_on(futures::future::try_join_all( + text_chunks.into_iter().map(|prompts| self.embed(prompts, &client)), + )) + } + + // Defaults copied from openai.rs + pub fn chunk_count_hint(&self) -> usize { + 10 + } + + pub fn prompt_count_in_chunk_hint(&self) -> usize { + 10 + } + + pub fn dimensions(&self) -> usize { + self.options.embedding_model.dimensions + } + + pub fn distribution(&self) -> Option { + None + } +} + +impl Display for OllamaError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{}", self.error) + } +} + +fn get_ollama_path() -> String { + // Important: Hostname not enough, has to be entire path to embeddings endpoint + std::env::var("MEILI_OLLAMA_URL").unwrap_or("http://localhost:11434/api/embeddings".to_string()) +} diff --git a/milli/src/vector/openai.rs b/milli/src/vector/openai.rs index 33442dda4..dcf3f4c89 100644 --- a/milli/src/vector/openai.rs +++ b/milli/src/vector/openai.rs @@ -419,12 +419,12 @@ impl Embedder { // retrying in case of failure -struct Retry { - error: EmbedError, +pub struct Retry { + pub error: EmbedError, strategy: RetryStrategy, } -enum RetryStrategy { +pub enum RetryStrategy { GiveUp, Retry, RetryTokenized, @@ -432,23 +432,23 @@ enum RetryStrategy { } impl Retry { - fn give_up(error: EmbedError) -> Self { + pub fn give_up(error: EmbedError) -> Self { Self { error, strategy: RetryStrategy::GiveUp } } - fn retry_later(error: EmbedError) -> Self { + pub fn retry_later(error: EmbedError) -> Self { Self { error, strategy: RetryStrategy::Retry } } - fn retry_tokenized(error: EmbedError) -> Self { + pub fn retry_tokenized(error: EmbedError) -> Self { Self { error, strategy: RetryStrategy::RetryTokenized } } - fn rate_limited(error: EmbedError) -> Self { + pub fn rate_limited(error: EmbedError) -> Self { Self { error, strategy: RetryStrategy::RetryAfterRateLimit } } - fn into_duration(self, attempt: u32) -> Result { + pub fn into_duration(self, attempt: u32) -> Result { match self.strategy { RetryStrategy::GiveUp => Err(self.error), RetryStrategy::Retry => Ok(tokio::time::Duration::from_millis((10u64).pow(attempt))), @@ -459,11 +459,11 @@ impl Retry { } } - fn must_tokenize(&self) -> bool { + pub fn must_tokenize(&self) -> bool { matches!(self.strategy, RetryStrategy::RetryTokenized) } - fn into_error(self) -> EmbedError { + pub fn into_error(self) -> EmbedError { self.error } } diff --git a/milli/src/vector/settings.rs b/milli/src/vector/settings.rs index 834a1c81d..89571e98a 100644 --- a/milli/src/vector/settings.rs +++ b/milli/src/vector/settings.rs @@ -1,7 +1,7 @@ use deserr::Deserr; use serde::{Deserialize, Serialize}; -use super::openai; +use super::{ollama, openai}; use crate::prompt::PromptData; use crate::update::Setting; use crate::vector::EmbeddingConfig; @@ -80,11 +80,15 @@ impl EmbeddingSettings { Self::SOURCE => { &[EmbedderSource::HuggingFace, EmbedderSource::OpenAi, EmbedderSource::UserProvided] } - Self::MODEL => &[EmbedderSource::HuggingFace, EmbedderSource::OpenAi], + Self::MODEL => { + &[EmbedderSource::HuggingFace, EmbedderSource::OpenAi, EmbedderSource::Ollama] + } Self::REVISION => &[EmbedderSource::HuggingFace], Self::API_KEY => &[EmbedderSource::OpenAi], Self::DIMENSIONS => &[EmbedderSource::OpenAi, EmbedderSource::UserProvided], - Self::DOCUMENT_TEMPLATE => &[EmbedderSource::HuggingFace, EmbedderSource::OpenAi], + Self::DOCUMENT_TEMPLATE => { + &[EmbedderSource::HuggingFace, EmbedderSource::OpenAi, EmbedderSource::Ollama] + } _other => unreachable!("unknown field"), } } @@ -101,6 +105,7 @@ impl EmbeddingSettings { EmbedderSource::HuggingFace => { &[Self::SOURCE, Self::MODEL, Self::REVISION, Self::DOCUMENT_TEMPLATE] } + EmbedderSource::Ollama => &[Self::SOURCE, Self::MODEL, Self::DOCUMENT_TEMPLATE], EmbedderSource::UserProvided => &[Self::SOURCE, Self::DIMENSIONS], } } @@ -134,6 +139,7 @@ pub enum EmbedderSource { #[default] OpenAi, HuggingFace, + Ollama, UserProvided, } @@ -143,6 +149,7 @@ impl std::fmt::Display for EmbedderSource { EmbedderSource::OpenAi => "openAi", EmbedderSource::HuggingFace => "huggingFace", EmbedderSource::UserProvided => "userProvided", + EmbedderSource::Ollama => "ollama", }; f.write_str(s) } @@ -195,6 +202,14 @@ impl From for EmbeddingSettings { dimensions: options.dimensions.map(Setting::Set).unwrap_or_default(), document_template: Setting::Set(prompt.template), }, + super::EmbedderOptions::Ollama(options) => Self { + source: Setting::Set(EmbedderSource::Ollama), + model: Setting::Set(options.embedding_model.name().to_owned()), + revision: Setting::NotSet, + api_key: Setting::NotSet, + dimensions: Setting::NotSet, + document_template: Setting::Set(prompt.template), + }, super::EmbedderOptions::UserProvided(options) => Self { source: Setting::Set(EmbedderSource::UserProvided), model: Setting::NotSet, @@ -229,6 +244,14 @@ impl From for EmbeddingConfig { } this.embedder_options = super::EmbedderOptions::OpenAi(options); } + EmbedderSource::Ollama => { + let mut options: ollama::EmbedderOptions = + super::ollama::EmbedderOptions::with_default_model(); + if let Some(model) = model.set() { + options.embedding_model = super::ollama::EmbeddingModel::from_name(&model); + } + this.embedder_options = super::EmbedderOptions::Ollama(options); + } EmbedderSource::HuggingFace => { let mut options = super::hf::EmbedderOptions::default(); if let Some(model) = model.set() { diff --git a/workloads/settings-add-remove-filters.json b/workloads/settings-add-remove-filters.json new file mode 100644 index 000000000..04a57c707 --- /dev/null +++ b/workloads/settings-add-remove-filters.json @@ -0,0 +1,94 @@ +{ + "name": "settings-add-remove-filters.json", + "run_count": 2, + "extra_cli_args": [ + "--max-indexing-threads=4" + ], + "assets": { + "150k-people.json": { + "local_location": null, + "remote_location": "https://milli-benchmarks.fra1.digitaloceanspaces.com/bench/datasets/150k-people.json", + "sha256": "28c359a0956958af0ba204ec11bad3045a0864a10b4838914fea25a01724f84b" + } + }, + "commands": [ + { + "route": "indexes/peoples/settings", + "method": "PATCH", + "body": { + "inline": { + "searchableAttributes": [ + "last_name", + "first_name", + "featured_job_organization_name", + "facebook_url", + "twitter_url", + "linkedin_url" + ], + "filterableAttributes": [ + "city", + "region", + "country_code" + ], + "dictionary": [ + "https://", + "http://", + "www.", + "crunchbase.com", + "facebook.com", + "twitter.com", + "linkedin.com" + ], + "stopWords": [ + "https://", + "http://", + "www.", + "crunchbase.com", + "facebook.com", + "twitter.com", + "linkedin.com" + ] + } + }, + "synchronous": "DontWait" + }, + { + "route": "indexes/peoples/documents", + "method": "POST", + "body": { + "asset": "150k-people.json" + }, + "synchronous": "WaitForTask" + }, + { + "route": "indexes/peoples/settings", + "method": "PATCH", + "body": { + "inline": { + "filterableAttributes": [ + "city", + "region", + "country_code", + "featured_job_title", + "featured_job_organization_name" + ] + } + }, + "synchronous": "WaitForTask" + }, + { + "route": "indexes/peoples/settings", + "method": "PATCH", + "body": { + "inline": { + "filterableAttributes": [ + "city", + "region", + "country_code" + ] + } + }, + "synchronous": "WaitForTask" + } + ] +} \ No newline at end of file diff --git a/workloads/settings-proximity-precision.json b/workloads/settings-proximity-precision.json new file mode 100644 index 000000000..48cfad49d --- /dev/null +++ b/workloads/settings-proximity-precision.json @@ -0,0 +1,86 @@ +{ + "name": "settings-proximity-precision.json", + "run_count": 2, + "extra_cli_args": [ + "--max-indexing-threads=4" + ], + "assets": { + "150k-people.json": { + "local_location": null, + "remote_location": "https://milli-benchmarks.fra1.digitaloceanspaces.com/bench/datasets/150k-people.json", + "sha256": "28c359a0956958af0ba204ec11bad3045a0864a10b4838914fea25a01724f84b" + } + }, + "commands": [ + { + "route": "indexes/peoples/settings", + "method": "PATCH", + "body": { + "inline": { + "searchableAttributes": [ + "last_name", + "first_name", + "featured_job_organization_name", + "facebook_url", + "twitter_url", + "linkedin_url" + ], + "filterableAttributes": [ + "city", + "region", + "country_code", + "featured_job_title", + "featured_job_organization_name" + ], + "dictionary": [ + "https://", + "http://", + "www.", + "crunchbase.com", + "facebook.com", + "twitter.com", + "linkedin.com" + ], + "stopWords": [ + "https://", + "http://", + "www.", + "crunchbase.com", + "facebook.com", + "twitter.com", + "linkedin.com" + ] + } + }, + "synchronous": "DontWait" + }, + { + "route": "indexes/peoples/documents", + "method": "POST", + "body": { + "asset": "150k-people.json" + }, + "synchronous": "WaitForTask" + }, + { + "route": "indexes/peoples/settings", + "method": "PATCH", + "body": { + "inline": { + "proximityPrecision": "byAttribute" + } + }, + "synchronous": "WaitForTask" + }, + { + "route": "indexes/peoples/settings", + "method": "PATCH", + "body": { + "inline": { + "proximityPrecision": "byWord" + } + }, + "synchronous": "WaitForTask" + } + ] +} \ No newline at end of file diff --git a/workloads/settings-remove-add-swap-searchable.json b/workloads/settings-remove-add-swap-searchable.json new file mode 100644 index 000000000..ba315680f --- /dev/null +++ b/workloads/settings-remove-add-swap-searchable.json @@ -0,0 +1,114 @@ +{ + "name": "settings-remove-add-swap-searchable.json", + "run_count": 2, + "extra_cli_args": [ + "--max-indexing-threads=4" + ], + "assets": { + "150k-people.json": { + "local_location": null, + "remote_location": "https://milli-benchmarks.fra1.digitaloceanspaces.com/bench/datasets/150k-people.json", + "sha256": "28c359a0956958af0ba204ec11bad3045a0864a10b4838914fea25a01724f84b" + } + }, + "commands": [ + { + "route": "indexes/peoples/settings", + "method": "PATCH", + "body": { + "inline": { + "searchableAttributes": [ + "last_name", + "first_name", + "featured_job_organization_name", + "facebook_url", + "twitter_url", + "linkedin_url" + ], + "filterableAttributes": [ + "city", + "region", + "country_code", + "featured_job_title", + "featured_job_organization_name" + ], + "dictionary": [ + "https://", + "http://", + "www.", + "crunchbase.com", + "facebook.com", + "twitter.com", + "linkedin.com" + ], + "stopWords": [ + "https://", + "http://", + "www.", + "crunchbase.com", + "facebook.com", + "twitter.com", + "linkedin.com" + ] + } + }, + "synchronous": "DontWait" + }, + { + "route": "indexes/peoples/documents", + "method": "POST", + "body": { + "asset": "150k-people.json" + }, + "synchronous": "WaitForTask" + }, + { + "route": "indexes/peoples/settings", + "method": "PATCH", + "body": { + "inline": { + "searchableAttributes": [ + "last_name", + "first_name", + "featured_job_organization_name" + ] + } + }, + "synchronous": "WaitForTask" + }, + { + "route": "indexes/peoples/settings", + "method": "PATCH", + "body": { + "inline": { + "searchableAttributes": [ + "last_name", + "first_name", + "featured_job_organization_name", + "facebook_url", + "twitter_url", + "linkedin_url" + ] + } + }, + "synchronous": "WaitForTask" + }, + { + "route": "indexes/peoples/settings", + "method": "PATCH", + "body": { + "inline": { + "searchableAttributes": [ + "first_name", + "last_name", + "featured_job_organization_name", + "facebook_url", + "twitter_url", + "linkedin_url" + ] + } + }, + "synchronous": "WaitForTask" + } + ] +} \ No newline at end of file diff --git a/workloads/settings-typo.json b/workloads/settings-typo.json new file mode 100644 index 000000000..a272e6d1f --- /dev/null +++ b/workloads/settings-typo.json @@ -0,0 +1,115 @@ +{ + "name": "settings-typo.json", + "run_count": 2, + "extra_cli_args": [ + "--max-indexing-threads=4" + ], + "assets": { + "150k-people.json": { + "local_location": null, + "remote_location": "https://milli-benchmarks.fra1.digitaloceanspaces.com/bench/datasets/150k-people.json", + "sha256": "28c359a0956958af0ba204ec11bad3045a0864a10b4838914fea25a01724f84b" + } + }, + "commands": [ + { + "route": "indexes/peoples/settings", + "method": "PATCH", + "body": { + "inline": { + "searchableAttributes": [ + "last_name", + "first_name", + "featured_job_title", + "featured_job_organization_name", + "facebook_url", + "twitter_url", + "linkedin_url" + ], + "filterableAttributes": [ + "city", + "region", + "country_code", + "featured_job_title", + "featured_job_organization_name" + ], + "dictionary": [ + "https://", + "http://", + "www.", + "crunchbase.com", + "facebook.com", + "twitter.com", + "linkedin.com" + ], + "stopWords": [ + "https://", + "http://", + "www.", + "crunchbase.com", + "facebook.com", + "twitter.com", + "linkedin.com" + ] + } + }, + "synchronous": "DontWait" + }, + { + "route": "indexes/peoples/documents", + "method": "POST", + "body": { + "asset": "150k-people.json" + }, + "synchronous": "WaitForTask" + }, + { + "route": "indexes/peoples/settings", + "method": "PATCH", + "body": { + "inline": { + "typoTolerance": { + "disableOnAttributes": ["featured_job_organization_name"] + } + } + }, + "synchronous": "WaitForTask" + }, + { + "route": "indexes/peoples/settings", + "method": "PATCH", + "body": { + "inline": { + "typoTolerance": { + "disableOnAttributes": [] + } + } + }, + "synchronous": "WaitForTask" + }, + { + "route": "indexes/peoples/settings", + "method": "PATCH", + "body": { + "inline": { + "typoTolerance": { + "disableOnWords": ["Ben","Elowitz","Kevin","Flaherty", "Ron", "Dustin", "Owen", "Chris", "Mark", "Matt", "Peter", "Van", "Head", "of"] + } + } + }, + "synchronous": "WaitForTask" + }, + { + "route": "indexes/peoples/settings", + "method": "PATCH", + "body": { + "inline": { + "typoTolerance": { + "disableOnWords": [] + } + } + }, + "synchronous": "WaitForTask" + } + ] +} \ No newline at end of file diff --git a/xtask/src/bench/dashboard.rs b/xtask/src/bench/dashboard.rs index 833426207..3ba0ca58b 100644 --- a/xtask/src/bench/dashboard.rs +++ b/xtask/src/bench/dashboard.rs @@ -11,157 +11,179 @@ use super::client::Client; use super::env_info; use super::workload::Workload; -pub async fn cancel_on_ctrl_c( - invocation_uuid: Uuid, - dashboard_client: Client, - abort_handle: AbortHandle, -) { - tracing::info!("press Ctrl-C to cancel the invocation"); - match ctrl_c().await { - Ok(()) => { - tracing::info!(%invocation_uuid, "received Ctrl-C, cancelling invocation"); - mark_as_failed(dashboard_client, invocation_uuid, None).await; - abort_handle.abort(); +#[derive(Debug, Clone)] +pub enum DashboardClient { + Client(Client), + Dry, +} + +impl DashboardClient { + pub fn new(dashboard_url: &str, api_key: Option<&str>) -> anyhow::Result { + let dashboard_client = Client::new( + Some(format!("{}/api/v1", dashboard_url)), + api_key, + Some(std::time::Duration::from_secs(60)), + )?; + + Ok(Self::Client(dashboard_client)) + } + + pub fn new_dry() -> Self { + Self::Dry + } + + pub async fn send_machine_info(&self, env: &env_info::Environment) -> anyhow::Result<()> { + let Self::Client(dashboard_client) = self else { return Ok(()) }; + + let response = dashboard_client + .put("machine") + .json(&json!({"hostname": env.hostname})) + .send() + .await + .context("sending machine information")?; + if !response.status().is_success() { + bail!( + "could not send machine information: {} {}", + response.status(), + response.text().await.unwrap_or_else(|_| "unknown".into()) + ); } - Err(error) => tracing::warn!( - error = &error as &dyn std::error::Error, - "failed to listen to Ctrl-C signal, invocation won't be canceled on Ctrl-C" - ), + Ok(()) } -} -pub async fn mark_as_failed( - dashboard_client: Client, - invocation_uuid: Uuid, - failure_reason: Option, -) { - let response = dashboard_client - .post("cancel-invocation") - .json(&json!({ - "invocation_uuid": invocation_uuid, - "failure_reason": failure_reason, - })) - .send() - .await; - let response = match response { - Ok(response) => response, - Err(response_error) => { - tracing::error!(error = &response_error as &dyn std::error::Error, %invocation_uuid, "could not mark invocation as failed"); - return; + pub async fn create_invocation( + &self, + build_info: build_info::BuildInfo, + commit_message: &str, + env: env_info::Environment, + max_workloads: usize, + reason: Option<&str>, + ) -> anyhow::Result { + let Self::Client(dashboard_client) = self else { return Ok(Uuid::now_v7()) }; + + let response = dashboard_client + .put("invocation") + .json(&json!({ + "commit": { + "sha1": build_info.commit_sha1, + "message": commit_message, + "commit_date": build_info.commit_timestamp, + "branch": build_info.branch, + "tag": build_info.describe.and_then(|describe| describe.as_tag()), + }, + "machine_hostname": env.hostname, + "max_workloads": max_workloads, + "reason": reason + })) + .send() + .await + .context("sending invocation")?; + if !response.status().is_success() { + bail!( + "could not send new invocation: {}", + response.text().await.unwrap_or_else(|_| "unknown".into()) + ); } - }; - - if !response.status().is_success() { - tracing::error!( - %invocation_uuid, - "could not mark invocation as failed: {}", - response.text().await.unwrap() - ); - return; - } - tracing::warn!(%invocation_uuid, "marked invocation as failed or canceled"); -} - -pub async fn send_machine_info( - dashboard_client: &Client, - env: &env_info::Environment, -) -> anyhow::Result<()> { - let response = dashboard_client - .put("machine") - .json(&json!({"hostname": env.hostname})) - .send() - .await - .context("sending machine information")?; - if !response.status().is_success() { - bail!( - "could not send machine information: {} {}", - response.status(), - response.text().await.unwrap_or_else(|_| "unknown".into()) - ); - } - Ok(()) -} - -pub async fn create_invocation( - dashboard_client: &Client, - build_info: build_info::BuildInfo, - commit_message: &str, - env: env_info::Environment, - max_workloads: usize, - reason: Option<&str>, -) -> anyhow::Result { - let response = dashboard_client - .put("invocation") - .json(&json!({ - "commit": { - "sha1": build_info.commit_sha1, - "message": commit_message, - "commit_date": build_info.commit_timestamp, - "branch": build_info.branch, - "tag": build_info.describe.and_then(|describe| describe.as_tag()), - }, - "machine_hostname": env.hostname, - "max_workloads": max_workloads, - "reason": reason - })) - .send() - .await - .context("sending invocation")?; - if !response.status().is_success() { - bail!( - "could not send new invocation: {}", - response.text().await.unwrap_or_else(|_| "unknown".into()) - ); - } - let invocation_uuid: Uuid = - response.json().await.context("could not deserialize invocation response as JSON")?; - Ok(invocation_uuid) -} - -pub async fn create_workload( - dashboard_client: &Client, - invocation_uuid: Uuid, - workload: &Workload, -) -> anyhow::Result { - let response = dashboard_client - .put("workload") - .json(&json!({ - "invocation_uuid": invocation_uuid, - "name": &workload.name, - "max_runs": workload.run_count, - })) - .send() - .await - .context("could not create new workload")?; - - if !response.status().is_success() { - bail!("creating new workload failed: {}", response.text().await.unwrap()) + let invocation_uuid: Uuid = + response.json().await.context("could not deserialize invocation response as JSON")?; + Ok(invocation_uuid) } - let workload_uuid: Uuid = - response.json().await.context("could not deserialize JSON as UUID")?; - Ok(workload_uuid) -} + pub async fn create_workload( + &self, + invocation_uuid: Uuid, + workload: &Workload, + ) -> anyhow::Result { + let Self::Client(dashboard_client) = self else { return Ok(Uuid::now_v7()) }; -pub async fn create_run( - dashboard_client: Client, - workload_uuid: Uuid, - report: &BTreeMap, -) -> anyhow::Result<()> { - let response = dashboard_client - .put("run") - .json(&json!({ - "workload_uuid": workload_uuid, - "data": report - })) - .send() - .await - .context("sending new run")?; - if !response.status().is_success() { - bail!( - "sending new run failed: {}", - response.text().await.unwrap_or_else(|_| "unknown".into()) - ) + let response = dashboard_client + .put("workload") + .json(&json!({ + "invocation_uuid": invocation_uuid, + "name": &workload.name, + "max_runs": workload.run_count, + })) + .send() + .await + .context("could not create new workload")?; + + if !response.status().is_success() { + bail!("creating new workload failed: {}", response.text().await.unwrap()) + } + + let workload_uuid: Uuid = + response.json().await.context("could not deserialize JSON as UUID")?; + Ok(workload_uuid) + } + + pub async fn create_run( + &self, + workload_uuid: Uuid, + report: &BTreeMap, + ) -> anyhow::Result<()> { + let Self::Client(dashboard_client) = self else { return Ok(()) }; + + let response = dashboard_client + .put("run") + .json(&json!({ + "workload_uuid": workload_uuid, + "data": report + })) + .send() + .await + .context("sending new run")?; + if !response.status().is_success() { + bail!( + "sending new run failed: {}", + response.text().await.unwrap_or_else(|_| "unknown".into()) + ) + } + Ok(()) + } + + pub async fn cancel_on_ctrl_c(self, invocation_uuid: Uuid, abort_handle: AbortHandle) { + tracing::info!("press Ctrl-C to cancel the invocation"); + match ctrl_c().await { + Ok(()) => { + tracing::info!(%invocation_uuid, "received Ctrl-C, cancelling invocation"); + self.mark_as_failed(invocation_uuid, None).await; + abort_handle.abort(); + } + Err(error) => tracing::warn!( + error = &error as &dyn std::error::Error, + "failed to listen to Ctrl-C signal, invocation won't be canceled on Ctrl-C" + ), + } + } + + pub async fn mark_as_failed(&self, invocation_uuid: Uuid, failure_reason: Option) { + if let DashboardClient::Client(client) = self { + let response = client + .post("cancel-invocation") + .json(&json!({ + "invocation_uuid": invocation_uuid, + "failure_reason": failure_reason, + })) + .send() + .await; + let response = match response { + Ok(response) => response, + Err(response_error) => { + tracing::error!(error = &response_error as &dyn std::error::Error, %invocation_uuid, "could not mark invocation as failed"); + return; + } + }; + + if !response.status().is_success() { + tracing::error!( + %invocation_uuid, + "could not mark invocation as failed: {}", + response.text().await.unwrap() + ); + return; + } + } + + tracing::warn!(%invocation_uuid, "marked invocation as failed or canceled"); } - Ok(()) } diff --git a/xtask/src/bench/mod.rs b/xtask/src/bench/mod.rs index 62c11b604..844b64f63 100644 --- a/xtask/src/bench/mod.rs +++ b/xtask/src/bench/mod.rs @@ -50,6 +50,10 @@ pub struct BenchDeriveArgs { #[arg(long, default_value_t = default_dashboard_url())] dashboard_url: String, + /// Don't actually send results to the dashboard + #[arg(long)] + no_dashboard: bool, + /// Directory to output reports. #[arg(long, default_value_t = default_report_folder())] report_folder: String, @@ -103,11 +107,11 @@ pub fn run(args: BenchDeriveArgs) -> anyhow::Result<()> { let assets_client = Client::new(None, args.assets_key.as_deref(), Some(std::time::Duration::from_secs(3600)))?; // 1h - let dashboard_client = Client::new( - Some(format!("{}/api/v1", args.dashboard_url)), - args.api_key.as_deref(), - Some(std::time::Duration::from_secs(60)), - )?; + let dashboard_client = if args.no_dashboard { + dashboard::DashboardClient::new_dry() + } else { + dashboard::DashboardClient::new(&args.dashboard_url, args.api_key.as_deref())? + }; // reporting uses its own client because keeping the stream open to wait for entries // blocks any other requests @@ -127,12 +131,12 @@ pub fn run(args: BenchDeriveArgs) -> anyhow::Result<()> { // enter runtime rt.block_on(async { - dashboard::send_machine_info(&dashboard_client, &env).await?; + dashboard_client.send_machine_info(&env).await?; let commit_message = build_info.commit_msg.context("missing commit message")?.split('\n').next().unwrap(); let max_workloads = args.workload_file.len(); let reason: Option<&str> = args.reason.as_deref(); - let invocation_uuid = dashboard::create_invocation(&dashboard_client, build_info, commit_message, env, max_workloads, reason).await?; + let invocation_uuid = dashboard_client.create_invocation( build_info, commit_message, env, max_workloads, reason).await?; tracing::info!(workload_count = args.workload_file.len(), "handling workload files"); @@ -167,7 +171,7 @@ pub fn run(args: BenchDeriveArgs) -> anyhow::Result<()> { let abort_handle = workload_runs.abort_handle(); tokio::spawn({ let dashboard_client = dashboard_client.clone(); - dashboard::cancel_on_ctrl_c(invocation_uuid, dashboard_client, abort_handle) + dashboard_client.cancel_on_ctrl_c(invocation_uuid, abort_handle) }); // wait for the end of the main task, handle result @@ -178,7 +182,7 @@ pub fn run(args: BenchDeriveArgs) -> anyhow::Result<()> { } Ok(Err(error)) => { tracing::error!(%invocation_uuid, error = %error, "invocation failed, attempting to report the failure to dashboard"); - dashboard::mark_as_failed(dashboard_client, invocation_uuid, Some(error.to_string())).await; + dashboard_client.mark_as_failed(invocation_uuid, Some(error.to_string())).await; tracing::warn!(%invocation_uuid, "invocation marked as failed following error"); Err(error) }, @@ -186,7 +190,7 @@ pub fn run(args: BenchDeriveArgs) -> anyhow::Result<()> { match join_error.try_into_panic() { Ok(panic) => { tracing::error!("invocation panicked, attempting to report the failure to dashboard"); - dashboard::mark_as_failed(dashboard_client, invocation_uuid, Some("Panicked".into())).await; + dashboard_client.mark_as_failed( invocation_uuid, Some("Panicked".into())).await; std::panic::resume_unwind(panic) } Err(_) => { diff --git a/xtask/src/bench/workload.rs b/xtask/src/bench/workload.rs index b3e952f29..d82c5ad19 100644 --- a/xtask/src/bench/workload.rs +++ b/xtask/src/bench/workload.rs @@ -12,8 +12,9 @@ use uuid::Uuid; use super::assets::Asset; use super::client::Client; use super::command::SyncMode; +use super::dashboard::DashboardClient; use super::BenchDeriveArgs; -use crate::bench::{assets, dashboard, meili_process}; +use crate::bench::{assets, meili_process}; #[derive(Deserialize)] pub struct Workload { @@ -25,7 +26,7 @@ pub struct Workload { } async fn run_commands( - dashboard_client: &Client, + dashboard_client: &DashboardClient, logs_client: &Client, meili_client: &Client, workload_uuid: Uuid, @@ -64,7 +65,7 @@ async fn run_commands( #[tracing::instrument(skip(assets_client, dashboard_client, logs_client, meili_client, workload, master_key, args), fields(workload = workload.name))] pub async fn execute( assets_client: &Client, - dashboard_client: &Client, + dashboard_client: &DashboardClient, logs_client: &Client, meili_client: &Client, invocation_uuid: Uuid, @@ -74,8 +75,7 @@ pub async fn execute( ) -> anyhow::Result<()> { assets::fetch_assets(assets_client, &workload.assets, &args.asset_folder).await?; - let workload_uuid = - dashboard::create_workload(dashboard_client, invocation_uuid, &workload).await?; + let workload_uuid = dashboard_client.create_workload(invocation_uuid, &workload).await?; let mut tasks = Vec::new(); @@ -113,7 +113,7 @@ pub async fn execute( #[allow(clippy::too_many_arguments)] // not best code quality, but this is a benchmark runner #[tracing::instrument(skip(dashboard_client, logs_client, meili_client, workload, master_key, args), fields(workload = %workload.name))] async fn execute_run( - dashboard_client: &Client, + dashboard_client: &DashboardClient, logs_client: &Client, meili_client: &Client, workload_uuid: Uuid, @@ -202,7 +202,7 @@ async fn start_report( } async fn stop_report( - dashboard_client: &Client, + dashboard_client: &DashboardClient, logs_client: &Client, workload_uuid: Uuid, filename: String, @@ -232,7 +232,7 @@ async fn stop_report( .context("could not convert trace to report")?; let context = || format!("writing report to {filename}"); - dashboard::create_run(dashboard_client, workload_uuid, &report).await?; + dashboard_client.create_run(workload_uuid, &report).await?; let mut output_file = std::io::BufWriter::new( std::fs::File::options() From 4628b7b7bd8a995073494bc0e25f1a642de695a2 Mon Sep 17 00:00:00 2001 From: 6543 <6543@obermui.de> Date: Wed, 20 Mar 2024 13:39:00 +0100 Subject: [PATCH 20/86] bump charabia to 0.8.8 and update lock file --- Cargo.lock | 182 +++++++++++++++++++++++++++++++---------------- milli/Cargo.toml | 2 +- 2 files changed, 120 insertions(+), 64 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index bdca7e24c..9f1ebb60b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -152,7 +152,7 @@ dependencies = [ "tokio-rustls", "tokio-util", "tracing", - "webpki-roots", + "webpki-roots 0.25.3", ] [[package]] @@ -306,9 +306,9 @@ checksum = "4b46cbb362ab8752921c97e041f5e366ee6297bd428a31275b9fcf1e380f7299" [[package]] name = "anstream" -version = "0.6.7" +version = "0.6.13" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4cd2405b3ac1faab2990b74d728624cd9fd115651fcecc7c2d8daf01376275ba" +checksum = "d96bd03f33fe50a863e394ee9718a706f988b9079b20c3784fb726e7678b62fb" dependencies = [ "anstyle", "anstyle-parse", @@ -320,9 +320,9 @@ dependencies = [ [[package]] name = "anstyle" -version = "1.0.1" +version = "1.0.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3a30da5c5f2d5e72842e00bcb57657162cdabef0931f40e2deb9b4140440cecd" +checksum = "8901269c6307e8d93993578286ac0edf7f195079ffff5ebdeea6a59ffb7e36bc" [[package]] name = "anstyle-parse" @@ -877,9 +877,9 @@ dependencies = [ [[package]] name = "charabia" -version = "0.8.7" +version = "0.8.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3a9071b1586dd067b5fdfd2069fab932c047ca5bbce4bd2bdee8af0f4b155053" +checksum = "60dc1a562fc8cb53d552d371758a4ecd76d15cc7489d2b968529cd9cadcbd854" dependencies = [ "aho-corasick", "cow-utils", @@ -1643,9 +1643,9 @@ checksum = "a246d82be1c9d791c5dfde9a2bd045fc3cbba3fa2b11ad558f27d01712f00569" [[package]] name = "encoding_rs" -version = "0.8.32" +version = "0.8.33" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "071a31f4ee85403370b58aca746f01041ede6f0da2730960ad001edc2b71b394" +checksum = "7268b386296a025e474d5140678f75d6de9493ae55a5d709eeb9dd08149945e1" dependencies = [ "cfg-if", ] @@ -1692,16 +1692,26 @@ dependencies = [ ] [[package]] -name = "env_logger" -version = "0.10.1" +name = "env_filter" +version = "0.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "95b3f3e67048839cb0d0781f445682a35113da7121f7c949db0e2be96a4fbece" +checksum = "a009aa4810eb158359dda09d0c87378e4bbb89b5a801f016885a4707ba24f7ea" dependencies = [ - "humantime", - "is-terminal", "log", "regex", - "termcolor", +] + +[[package]] +name = "env_logger" +version = "0.11.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6c012a26a7f605efc424dd53697843a72be7dc86ad2d01f7814337794a12231d" +dependencies = [ + "anstream", + "anstyle", + "env_filter", + "humantime", + "log", ] [[package]] @@ -2393,7 +2403,7 @@ dependencies = [ "futures-util", "http 0.2.11", "hyper", - "rustls", + "rustls 0.21.10", "tokio", "tokio-rustls", ] @@ -2455,9 +2465,9 @@ dependencies = [ [[package]] name = "indexmap" -version = "2.1.0" +version = "2.2.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d530e1a18b1cb4c484e6e34556a0d948706958449fca0cab753d649f2bce3d1f" +checksum = "7b0b929d511467233429c45a44ac1dcaa21ba0f5ba11e4879e6ed28ddb4f9df4" dependencies = [ "equivalent", "hashbrown", @@ -2736,9 +2746,9 @@ dependencies = [ [[package]] name = "lindera-cc-cedict-builder" -version = "0.27.2" +version = "0.28.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a90d23f7cef31c6ab7ac0d4f3b23940754207f7b5a80b080c39193caffe99ac2" +checksum = "ca21f2ee3ca40e7f3ebbd568d041be1531c2c28dbf540e737aeba934ab53f330" dependencies = [ "anyhow", "bincode", @@ -2755,9 +2765,9 @@ dependencies = [ [[package]] name = "lindera-compress" -version = "0.27.2" +version = "0.28.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1927b7d2bd4ffc19e07691bf8609722663c341f80260a1c636cee8f1ec420dce" +checksum = "34da125091f3b3a49351f418484a16cb2a23f6888cd53fe219edad19d263da5d" dependencies = [ "anyhow", "flate2", @@ -2766,9 +2776,9 @@ dependencies = [ [[package]] name = "lindera-core" -version = "0.27.2" +version = "0.28.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3299caa2b81c9a076535a4651a83bf7d624c15f2349f243187fffc64b5a78251" +checksum = "09d4b717a8a31b73a3cbd3552e0abda14e0c85d97dc8b911035342533defdbad" dependencies = [ "anyhow", "bincode", @@ -2783,9 +2793,9 @@ dependencies = [ [[package]] name = "lindera-decompress" -version = "0.27.2" +version = "0.28.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7b82b8d2323a67dc8ff0c40751d199b7ba94cd5e3c13a5b31622d318acc79e5b" +checksum = "98f4476c99cb4ffa54fbfc42953adf69ada7276cfbb594bce9829547de012058" dependencies = [ "anyhow", "flate2", @@ -2794,9 +2804,9 @@ dependencies = [ [[package]] name = "lindera-dictionary" -version = "0.27.2" +version = "0.28.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cddf783b459d54b130d956889bec052c25fcb478a304e03fa9b2289387572bc5" +checksum = "a45b92f0ce331c2202c6cec3135e4bfce29525ab3bb97a613c27c8e0a29fa967" dependencies = [ "anyhow", "bincode", @@ -2814,9 +2824,9 @@ dependencies = [ [[package]] name = "lindera-ipadic-builder" -version = "0.27.2" +version = "0.28.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "27c708f08f14b0806f6c4cce5324b4bcba27209463026b78c31f399f8be9d30d" +checksum = "642dee52201852df209cb43423ff1ca4d161a329f5cdba049a7b5820118345f2" dependencies = [ "anyhow", "bincode", @@ -2835,9 +2845,9 @@ dependencies = [ [[package]] name = "lindera-ipadic-neologd-builder" -version = "0.27.2" +version = "0.28.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e5e67eb91652203d202f7d27ead220d1d8c9099552709b8429eae9c70f2312fb" +checksum = "325144b154e68159373e944d1cd7f67c6ff9965a2af41240a8e41732b3fdb3af" dependencies = [ "anyhow", "bincode", @@ -2856,9 +2866,9 @@ dependencies = [ [[package]] name = "lindera-ko-dic" -version = "0.27.2" +version = "0.28.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d45da8d9a5888f4d4e78bb29fc82ff9ae519962efb0d2d92343b6cf8e373952f" +checksum = "b484a2f9964e7424264fda304beb6ff6ad883c347accfe1115e777dedef3661d" dependencies = [ "bincode", "byteorder", @@ -2873,9 +2883,9 @@ dependencies = [ [[package]] name = "lindera-ko-dic-builder" -version = "0.27.2" +version = "0.28.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "41c0933295dc945178bbc08f34111dc3ef22bfee38820f78453c8f8d4f3463d1" +checksum = "b9413d4d9bf7af921f5ac64414a290c7ba81695e8ba08dd2f6c950b57c281a69" dependencies = [ "anyhow", "bincode", @@ -2893,12 +2903,11 @@ dependencies = [ [[package]] name = "lindera-tokenizer" -version = "0.27.2" +version = "0.28.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "348ce9bb3f2e5edc577420b98cca05b2177f3af50ef5ae278a1d8a1351d56197" +checksum = "9987c818462d51ca67e131e40f0386e25e8c557e195059b1257f95731561185d" dependencies = [ "bincode", - "byteorder", "lindera-core", "lindera-dictionary", "once_cell", @@ -2908,9 +2917,9 @@ dependencies = [ [[package]] name = "lindera-unidic" -version = "0.27.2" +version = "0.28.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "74022a57c395ed7e213a9cd5833207e3c583145078ee9a164aeaec68b30c9d8e" +checksum = "0c379cf436b2627cd7d3498642e491eadbff9b3e01231c516ce9f9b1893ab7c3" dependencies = [ "bincode", "byteorder", @@ -2925,9 +2934,9 @@ dependencies = [ [[package]] name = "lindera-unidic-builder" -version = "0.27.2" +version = "0.28.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a34e5564ee81af82603cd6a03c3abe6e17cc0ae598bfa5078809f06e59e96e08" +checksum = "601ec33b5174141396a7a4ca066278863840221fec32d0be19091e7fae91ed94" dependencies = [ "anyhow", "bincode", @@ -3185,7 +3194,7 @@ dependencies = [ "rayon", "regex", "reqwest", - "rustls", + "rustls 0.21.10", "rustls-pemfile", "segment", "serde", @@ -4241,7 +4250,7 @@ dependencies = [ "once_cell", "percent-encoding", "pin-project-lite", - "rustls", + "rustls 0.21.10", "rustls-pemfile", "serde", "serde_json", @@ -4256,7 +4265,7 @@ dependencies = [ "wasm-bindgen-futures", "wasm-streams", "web-sys", - "webpki-roots", + "webpki-roots 0.25.3", "winreg", ] @@ -4366,10 +4375,24 @@ checksum = "f9d5a6813c0759e4609cd494e8e725babae6a2ca7b62a5536a13daaec6fcb7ba" dependencies = [ "log", "ring", - "rustls-webpki", + "rustls-webpki 0.101.7", "sct", ] +[[package]] +name = "rustls" +version = "0.22.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e87c9956bd9807afa1f77e0f7594af32566e830e088a5576d27c5b6f30f49d41" +dependencies = [ + "log", + "ring", + "rustls-pki-types", + "rustls-webpki 0.102.2", + "subtle", + "zeroize", +] + [[package]] name = "rustls-pemfile" version = "1.0.4" @@ -4379,6 +4402,12 @@ dependencies = [ "base64 0.21.7", ] +[[package]] +name = "rustls-pki-types" +version = "1.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5ede67b28608b4c60685c7d54122d4400d90f62b40caee7700e700380a390fa8" + [[package]] name = "rustls-webpki" version = "0.101.7" @@ -4389,6 +4418,17 @@ dependencies = [ "untrusted", ] +[[package]] +name = "rustls-webpki" +version = "0.102.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "faaa0a62740bedb9b2ef5afa303da42764c012f743917351dc9a237ea1663610" +dependencies = [ + "ring", + "rustls-pki-types", + "untrusted", +] + [[package]] name = "rustversion" version = "1.0.14" @@ -4467,9 +4507,9 @@ checksum = "a3f0bf26fd526d2a95683cd0f87bf103b8539e2ca1ef48ce002d67aad59aa0b4" [[package]] name = "serde" -version = "1.0.195" +version = "1.0.197" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "63261df402c67811e9ac6def069e4786148c4563f4b50fd4bf30aa370d626b02" +checksum = "3fb1c873e1b9b056a4dc4c0c198b24c3ffa059243875552b2bd0933b1aee4ce2" dependencies = [ "serde_derive", ] @@ -4485,9 +4525,9 @@ dependencies = [ [[package]] name = "serde_derive" -version = "1.0.195" +version = "1.0.197" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "46fe8f8603d81ba86327b23a2e9cdf49e1255fb94a4c5f297f6ee0547178ea2c" +checksum = "7eb0b34b42edc17f6b7cac84a52a1c5f0e1bb2227e997ca9011ea3dd34e8610b" dependencies = [ "proc-macro2", "quote", @@ -4496,9 +4536,9 @@ dependencies = [ [[package]] name = "serde_json" -version = "1.0.111" +version = "1.0.114" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "176e46fa42316f18edd598015a5166857fc835ec732f5215eac6b7bdbf0a84f4" +checksum = "c5f09b1bd632ef549eaa9f60a1f8de742bdbc698e6cee2095fc84dde5f549ae0" dependencies = [ "indexmap", "itoa", @@ -4888,18 +4928,18 @@ dependencies = [ [[package]] name = "thiserror" -version = "1.0.56" +version = "1.0.58" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d54378c645627613241d077a3a79db965db602882668f9136ac42af9ecb730ad" +checksum = "03468839009160513471e86a034bb2c5c0e4baae3b43f79ffc55c4a5427b3297" dependencies = [ "thiserror-impl", ] [[package]] name = "thiserror-impl" -version = "1.0.56" +version = "1.0.58" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fa0faa943b50f3db30a20aa7e265dbc66076993efed8463e8de414e5d06d3471" +checksum = "c61f3ba182994efc43764a46c018c347bc492c79f024e705f46567b418f6d4f7" dependencies = [ "proc-macro2", "quote", @@ -5055,7 +5095,7 @@ version = "0.24.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c28327cf380ac148141087fbfb9de9d7bd4e84ab5d2c28fbc911d753de8a7081" dependencies = [ - "rustls", + "rustls 0.21.10", "tokio", ] @@ -5332,21 +5372,22 @@ checksum = "8ecb6da28b8a351d773b68d5825ac39017e680750f980f3a1a85cd8dd28a47c1" [[package]] name = "ureq" -version = "2.9.1" +version = "2.9.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f8cdd25c339e200129fe4de81451814e5228c9b771d57378817d6117cc2b3f97" +checksum = "11f214ce18d8b2cbe84ed3aa6486ed3f5b285cf8d8fbdbce9f3f767a724adc35" dependencies = [ "base64 0.21.7", "flate2", "log", "once_cell", - "rustls", - "rustls-webpki", + "rustls 0.22.2", + "rustls-pki-types", + "rustls-webpki 0.102.2", "serde", "serde_json", "socks", "url", - "webpki-roots", + "webpki-roots 0.26.1", ] [[package]] @@ -5588,6 +5629,15 @@ version = "0.25.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1778a42e8b3b90bff8d0f5032bf22250792889a5cdc752aa0020c84abe3aaf10" +[[package]] +name = "webpki-roots" +version = "0.26.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b3de34ae270483955a94f4b21bdaaeb83d508bb84a01435f393818edb0012009" +dependencies = [ + "rustls-pki-types", +] + [[package]] name = "whatlang" version = "0.16.4" @@ -5987,6 +6037,12 @@ dependencies = [ "synstructure", ] +[[package]] +name = "zeroize" +version = "1.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "525b4ec142c6b68a2d10f01f7bbf6755599ca3f81ea53b8431b7dd348f5fdb2d" + [[package]] name = "zerovec" version = "0.10.1" diff --git a/milli/Cargo.toml b/milli/Cargo.toml index 1dfa495ea..fa4215404 100644 --- a/milli/Cargo.toml +++ b/milli/Cargo.toml @@ -17,7 +17,7 @@ bincode = "1.3.3" bstr = "1.9.0" bytemuck = { version = "1.14.0", features = ["extern_crate_alloc"] } byteorder = "1.5.0" -charabia = { version = "0.8.7", default-features = false } +charabia = { version = "0.8.8", default-features = false } concat-arrays = "0.1.2" crossbeam-channel = "0.5.11" deserr = "0.6.1" From c67f04c74624300df85a8185d911498b983f991b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Cl=C3=A9mentine=20U=2E=20-=20curqui?= Date: Wed, 20 Mar 2024 18:45:56 +0100 Subject: [PATCH 21/86] Update sprint_issue.md --- .github/ISSUE_TEMPLATE/sprint_issue.md | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/.github/ISSUE_TEMPLATE/sprint_issue.md b/.github/ISSUE_TEMPLATE/sprint_issue.md index 1e90f5946..0a3eb6843 100644 --- a/.github/ISSUE_TEMPLATE/sprint_issue.md +++ b/.github/ISSUE_TEMPLATE/sprint_issue.md @@ -2,7 +2,7 @@ name: New sprint issue about: ⚠️ Should only be used by the engine team ⚠️ title: '' -labels: '' +labels: 'missing usage in PRD, impacts docs' assignees: '' --- @@ -21,11 +21,7 @@ Related spec: WIP ## TODO - - -- [ ] Release a prototype -- [ ] If prototype validated, merge changes into `main` -- [ ] Update the spec + ### Reminders when modifying the Setting API From 8394be948426c48cd5ed2702f9b15de59efd6d7f Mon Sep 17 00:00:00 2001 From: curquiza Date: Thu, 21 Mar 2024 15:49:25 +0100 Subject: [PATCH 22/86] Add automation to create openAPI issue --- .github/workflows/milestone-workflow.yml | 21 ++++++++++++++++++++- 1 file changed, 20 insertions(+), 1 deletion(-) diff --git a/.github/workflows/milestone-workflow.yml b/.github/workflows/milestone-workflow.yml index 2ede3dc21..c15684661 100644 --- a/.github/workflows/milestone-workflow.yml +++ b/.github/workflows/milestone-workflow.yml @@ -112,7 +112,7 @@ jobs: create-update-version-issue: needs: get-release-version - # Create the changelog issue if the release is not only a patch release + # Create the update-version issue even if the release is a patch release if: github.event.action == 'created' runs-on: ubuntu-latest env: @@ -129,6 +129,25 @@ jobs: --body-file $ISSUE_TEMPLATE \ --milestone $MILESTONE_VERSION + create-update-openapi-issue: + needs: get-release-version + # Create the openAPI issue if the release is not only a patch release + if: github.event.action == 'created' && needs.get-release-version.outputs.is-patch == 'false' + runs-on: ubuntu-latest + env: + ISSUE_TEMPLATE: issue-template.md + steps: + - uses: actions/checkout@v3 + - name: Download the issue template + run: curl -s https://raw.githubusercontent.com/meilisearch/engine-team/main/issue-templates/update-openapi-issue.md > $ISSUE_TEMPLATE + - name: Create the issue + run: | + gh issue create \ + --title "Update Open API file for $MILESTONE_VERSION" \ + --label 'maintenance' \ + --body-file $ISSUE_TEMPLATE \ + --milestone $MILESTONE_VERSION + # ---------------- # MILESTONE CLOSED # ---------------- From 9865c5804664d46bccbf072635d846ced52472ab Mon Sep 17 00:00:00 2001 From: availhang Date: Fri, 22 Mar 2024 15:23:13 +0800 Subject: [PATCH 23/86] chore: remove repetitive words Signed-off-by: availhang --- BENCHMARKS.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/BENCHMARKS.md b/BENCHMARKS.md index b3d311c45..e588b1b5b 100644 --- a/BENCHMARKS.md +++ b/BENCHMARKS.md @@ -341,7 +341,7 @@ The URL of the server is in our password manager (look for "benchboard"). ``` ssh root@ ``` - Note the the ipv6 must **NOT** be between escaped square brackets for SSH 🥲 + Note the ipv6 must **NOT** be between escaped square brackets for SSH 🥲 5. On the server, set the correct permissions for the new binary: ``` chown bench:bench /bench/new-benchboard From 58330703585af8899b903aaee6a40e5d14e96668 Mon Sep 17 00:00:00 2001 From: Rohan Kumar Date: Tue, 30 Jan 2024 11:26:01 +0530 Subject: [PATCH 24/86] feat: add status code label to prometheus http request counter --- meilisearch/src/metrics.rs | 2 +- meilisearch/src/middleware.rs | 11 ++++++++--- 2 files changed, 9 insertions(+), 4 deletions(-) diff --git a/meilisearch/src/metrics.rs b/meilisearch/src/metrics.rs index 652e6c227..3be12c7ce 100644 --- a/meilisearch/src/metrics.rs +++ b/meilisearch/src/metrics.rs @@ -19,7 +19,7 @@ lazy_static! { pub static ref MEILISEARCH_HTTP_RESPONSE_TIME_CUSTOM_BUCKETS: [f64; 29] = create_buckets(); pub static ref MEILISEARCH_HTTP_REQUESTS_TOTAL: IntCounterVec = register_int_counter_vec!( opts!("meilisearch_http_requests_total", "Meilisearch HTTP requests total"), - &["method", "path"] + &["method", "path", "status"] ) .expect("Can't create a metric"); pub static ref MEILISEARCH_DEGRADED_SEARCH_REQUESTS: IntGauge = register_int_gauge!(opts!( diff --git a/meilisearch/src/middleware.rs b/meilisearch/src/middleware.rs index 5b87dee34..6707bb6d5 100644 --- a/meilisearch/src/middleware.rs +++ b/meilisearch/src/middleware.rs @@ -65,9 +65,6 @@ where .with_label_values(&[&request_method, request_path]) .start_timer(), ); - crate::metrics::MEILISEARCH_HTTP_REQUESTS_TOTAL - .with_label_values(&[&request_method, request_path]) - .inc(); } }; @@ -76,6 +73,14 @@ where Box::pin(async move { let res = fut.await?; + crate::metrics::MEILISEARCH_HTTP_REQUESTS_TOTAL + .with_label_values(&[ + res.request().method().as_str(), + res.request().path(), + res.status().as_str(), + ]) + .inc(); + if let Some(histogram_timer) = histogram_timer { histogram_timer.observe_duration(); }; From 325435ad4340af19c0f8ceb152a2478bd8a0b9cb Mon Sep 17 00:00:00 2001 From: Rohan Kumar Date: Sun, 24 Mar 2024 21:22:02 +0530 Subject: [PATCH 25/86] feat: add request rate and error rate panels to grafana dashboard --- assets/grafana-dashboard.json | 974 ++++++++++++++++++++++------------ 1 file changed, 631 insertions(+), 343 deletions(-) diff --git a/assets/grafana-dashboard.json b/assets/grafana-dashboard.json index 74a456b97..2cfa85a46 100644 --- a/assets/grafana-dashboard.json +++ b/assets/grafana-dashboard.json @@ -1,4 +1,47 @@ { + "__inputs": [ + { + "name": "DS_PROMETHEUS", + "label": "prometheus", + "description": "", + "type": "datasource", + "pluginId": "prometheus", + "pluginName": "Prometheus" + } + ], + "__elements": {}, + "__requires": [ + { + "type": "panel", + "id": "gauge", + "name": "Gauge", + "version": "" + }, + { + "type": "grafana", + "id": "grafana", + "name": "Grafana", + "version": "10.4.1" + }, + { + "type": "panel", + "id": "heatmap", + "name": "Heatmap", + "version": "" + }, + { + "type": "datasource", + "id": "prometheus", + "name": "Prometheus", + "version": "1.0.0" + }, + { + "type": "panel", + "id": "timeseries", + "name": "Time series", + "version": "" + } + ], "annotations": { "list": [ { @@ -24,7 +67,7 @@ "editable": true, "fiscalYearStartMonth": 0, "graphTooltip": 0, - "id": 2, + "id": null, "links": [], "liveNow": false, "panels": [ @@ -54,7 +97,8 @@ }, { "datasource": { - "type": "prometheus" + "type": "prometheus", + "uid": "${DS_PROMETHEUS}" }, "fieldConfig": { "defaults": { @@ -86,6 +130,8 @@ "id": 22, "interval": "5s", "options": { + "minVizHeight": 75, + "minVizWidth": 75, "orientation": "auto", "reduceOptions": { "calcs": [ @@ -96,13 +142,15 @@ }, "showThresholdLabels": false, "showThresholdMarkers": true, + "sizing": "auto", "text": {} }, - "pluginVersion": "10.0.1", + "pluginVersion": "10.4.1", "targets": [ { "datasource": { - "type": "prometheus" + "type": "prometheus", + "uid": "${DS_PROMETHEUS}" }, "editorMode": "builder", "exemplar": true, @@ -118,7 +166,8 @@ }, { "datasource": { - "type": "prometheus" + "type": "prometheus", + "uid": "${DS_PROMETHEUS}" }, "fieldConfig": { "defaults": { @@ -146,6 +195,8 @@ }, "id": 18, "options": { + "minVizHeight": 75, + "minVizWidth": 75, "orientation": "auto", "reduceOptions": { "calcs": [ @@ -156,13 +207,15 @@ }, "showThresholdLabels": false, "showThresholdMarkers": true, + "sizing": "auto", "text": {} }, - "pluginVersion": "10.0.1", + "pluginVersion": "10.4.1", "targets": [ { "datasource": { - "type": "prometheus" + "type": "prometheus", + "uid": "${DS_PROMETHEUS}" }, "editorMode": "builder", "expr": "meilisearch_index_docs_count{job=\"$job\", index=\"$Index\", instance=\"$instance\"}", @@ -176,71 +229,8 @@ }, { "datasource": { - "type": "prometheus" - }, - "fieldConfig": { - "defaults": { - "color": { - "mode": "thresholds" - }, - "mappings": [], - "thresholds": { - "mode": "absolute", - "steps": [ - { - "color": "green", - "value": null - }, - { - "color": "red", - "value": 80 - } - ] - } - }, - "overrides": [] - }, - "gridPos": { - "h": 6, - "w": 4, - "x": 12, - "y": 1 - }, - "id": 19, - "options": { - "orientation": "auto", - "reduceOptions": { - "calcs": [ - "lastNotNull" - ], - "fields": "", - "values": false - }, - "showThresholdLabels": false, - "showThresholdMarkers": true, - "text": {} - }, - "pluginVersion": "10.0.1", - "targets": [ - { - "datasource": { - "type": "prometheus" - }, - "editorMode": "builder", - "exemplar": true, - "expr": "round(increase(meilisearch_http_requests_total{method=\"POST\", path=\"/indexes/$Index/search\", job=\"$job\"}[1h]))", - "interval": "", - "legendFormat": "", - "range": true, - "refId": "A" - } - ], - "title": "Total Searches (1h)", - "type": "gauge" - }, - { - "datasource": { - "type": "prometheus" + "type": "prometheus", + "uid": "${DS_PROMETHEUS}" }, "fieldConfig": { "defaults": { @@ -272,6 +262,8 @@ }, "id": 26, "options": { + "minVizHeight": 75, + "minVizWidth": 75, "orientation": "auto", "reduceOptions": { "calcs": [ @@ -282,13 +274,15 @@ }, "showThresholdLabels": false, "showThresholdMarkers": true, + "sizing": "auto", "text": {} }, - "pluginVersion": "9.5.2", + "pluginVersion": "10.4.1", "targets": [ { "datasource": { - "type": "prometheus" + "type": "prometheus", + "uid": "${DS_PROMETHEUS}" }, "editorMode": "builder", "exemplar": true, @@ -304,7 +298,77 @@ }, { "datasource": { - "type": "prometheus" + "type": "prometheus", + "uid": "${DS_PROMETHEUS}" + }, + "fieldConfig": { + "defaults": { + "color": { + "mode": "thresholds" + }, + "mappings": [], + "thresholds": { + "mode": "absolute", + "steps": [ + { + "color": "green", + "value": null + }, + { + "color": "red", + "value": 80 + } + ] + } + }, + "overrides": [] + }, + "gridPos": { + "h": 6, + "w": 4, + "x": 12, + "y": 1 + }, + "id": 19, + "options": { + "minVizHeight": 75, + "minVizWidth": 75, + "orientation": "auto", + "reduceOptions": { + "calcs": [ + "lastNotNull" + ], + "fields": "", + "values": false + }, + "showThresholdLabels": false, + "showThresholdMarkers": true, + "sizing": "auto", + "text": {} + }, + "pluginVersion": "10.4.1", + "targets": [ + { + "datasource": { + "type": "prometheus", + "uid": "${DS_PROMETHEUS}" + }, + "editorMode": "builder", + "exemplar": true, + "expr": "round(increase(meilisearch_http_requests_total{method=\"POST\", path=\"/indexes/$Index/search\", job=\"$job\"}[1h]))", + "interval": "", + "legendFormat": "", + "range": true, + "refId": "A" + } + ], + "title": "Total Searches (1h)", + "type": "gauge" + }, + { + "datasource": { + "type": "prometheus", + "uid": "${DS_PROMETHEUS}" }, "fieldConfig": { "defaults": { @@ -332,6 +396,8 @@ }, "id": 20, "options": { + "minVizHeight": 75, + "minVizWidth": 75, "orientation": "auto", "reduceOptions": { "calcs": [ @@ -342,13 +408,15 @@ }, "showThresholdLabels": false, "showThresholdMarkers": true, + "sizing": "auto", "text": {} }, - "pluginVersion": "10.0.1", + "pluginVersion": "10.4.1", "targets": [ { "datasource": { - "type": "prometheus" + "type": "prometheus", + "uid": "${DS_PROMETHEUS}" }, "editorMode": "builder", "exemplar": true, @@ -364,7 +432,8 @@ }, { "datasource": { - "type": "prometheus" + "type": "prometheus", + "uid": "${DS_PROMETHEUS}" }, "fieldConfig": { "defaults": { @@ -392,6 +461,8 @@ }, "id": 21, "options": { + "minVizHeight": 75, + "minVizWidth": 75, "orientation": "auto", "reduceOptions": { "calcs": [ @@ -402,13 +473,15 @@ }, "showThresholdLabels": false, "showThresholdMarkers": true, + "sizing": "auto", "text": {} }, - "pluginVersion": "10.0.1", + "pluginVersion": "10.4.1", "targets": [ { "datasource": { - "type": "prometheus" + "type": "prometheus", + "uid": "${DS_PROMETHEUS}" }, "editorMode": "builder", "exemplar": true, @@ -424,7 +497,8 @@ }, { "datasource": { - "type": "prometheus" + "type": "prometheus", + "uid": "${DS_PROMETHEUS}" }, "description": "", "fieldConfig": { @@ -433,6 +507,7 @@ "mode": "palette-classic" }, "custom": { + "axisBorderShow": false, "axisCenteredZero": false, "axisColorMode": "text", "axisLabel": "", @@ -446,6 +521,7 @@ "tooltip": false, "viz": false }, + "insertNulls": false, "lineInterpolation": "linear", "lineStyle": { "fill": "solid" @@ -471,7 +547,8 @@ "mode": "absolute", "steps": [ { - "color": "green" + "color": "green", + "value": null }, { "color": "red", @@ -507,7 +584,8 @@ "targets": [ { "datasource": { - "type": "prometheus" + "type": "prometheus", + "uid": "${DS_PROMETHEUS}" }, "editorMode": "builder", "exemplar": true, @@ -519,7 +597,8 @@ }, { "datasource": { - "type": "prometheus" + "type": "prometheus", + "uid": "${DS_PROMETHEUS}" }, "editorMode": "builder", "expr": "meilisearch_used_db_size_bytes{job=\"$job\", instance=\"$instance\"}", @@ -534,7 +613,8 @@ }, { "datasource": { - "type": "prometheus" + "type": "prometheus", + "uid": "${DS_PROMETHEUS}" }, "fieldConfig": { "defaults": { @@ -542,6 +622,7 @@ "mode": "continuous-YlBl" }, "custom": { + "axisBorderShow": false, "axisCenteredZero": false, "axisColorMode": "text", "axisLabel": "", @@ -555,6 +636,7 @@ "tooltip": false, "viz": false }, + "insertNulls": false, "lineInterpolation": "linear", "lineWidth": 1, "pointSize": 5, @@ -613,7 +695,8 @@ "targets": [ { "datasource": { - "type": "prometheus" + "type": "prometheus", + "uid": "${DS_PROMETHEUS}" }, "editorMode": "builder", "exemplar": true, @@ -629,7 +712,192 @@ }, { "datasource": { - "type": "prometheus" + "type": "prometheus", + "uid": "${DS_PROMETHEUS}" + }, + "description": "number of requests per second, faceted by its HTTP status code", + "fieldConfig": { + "defaults": { + "color": { + "mode": "palette-classic" + }, + "custom": { + "axisBorderShow": false, + "axisCenteredZero": false, + "axisColorMode": "text", + "axisLabel": "", + "axisPlacement": "auto", + "barAlignment": 0, + "drawStyle": "line", + "fillOpacity": 0, + "gradientMode": "none", + "hideFrom": { + "legend": false, + "tooltip": false, + "viz": false + }, + "insertNulls": false, + "lineInterpolation": "linear", + "lineWidth": 1, + "pointSize": 5, + "scaleDistribution": { + "type": "linear" + }, + "showPoints": "auto", + "spanNulls": false, + "stacking": { + "group": "A", + "mode": "none" + }, + "thresholdsStyle": { + "mode": "off" + } + }, + "mappings": [], + "thresholds": { + "mode": "absolute", + "steps": [ + { + "color": "green", + "value": null + } + ] + } + }, + "overrides": [] + }, + "gridPos": { + "h": 11, + "w": 12, + "x": 0, + "y": 18 + }, + "id": 27, + "options": { + "legend": { + "calcs": [], + "displayMode": "list", + "placement": "right", + "showLegend": true + }, + "tooltip": { + "mode": "single", + "sort": "none" + } + }, + "targets": [ + { + "datasource": { + "type": "prometheus", + "uid": "${DS_PROMETHEUS}" + }, + "editorMode": "code", + "expr": "sum(rate(meilisearch_http_requests_total{instance=~\"$instance\",job=~\"$job\"} [$__rate_interval])) by (status)", + "instant": false, + "legendFormat": "HTTP Server - {{status}}", + "range": true, + "refId": "A" + } + ], + "title": "Request Rate", + "type": "timeseries" + }, + { + "datasource": { + "type": "prometheus", + "uid": "${DS_PROMETHEUS}" + }, + "description": "percentage of 4xx and 5xx HTTP responses over the total of the requests", + "fieldConfig": { + "defaults": { + "color": { + "mode": "palette-classic" + }, + "custom": { + "axisBorderShow": false, + "axisCenteredZero": false, + "axisColorMode": "text", + "axisLabel": "", + "axisPlacement": "auto", + "barAlignment": 0, + "drawStyle": "line", + "fillOpacity": 0, + "gradientMode": "none", + "hideFrom": { + "legend": false, + "tooltip": false, + "viz": false + }, + "insertNulls": false, + "lineInterpolation": "linear", + "lineWidth": 1, + "pointSize": 5, + "scaleDistribution": { + "type": "linear" + }, + "showPoints": "auto", + "spanNulls": false, + "stacking": { + "group": "A", + "mode": "none" + }, + "thresholdsStyle": { + "mode": "off" + } + }, + "mappings": [], + "thresholds": { + "mode": "absolute", + "steps": [ + { + "color": "green", + "value": null + } + ] + } + }, + "overrides": [] + }, + "gridPos": { + "h": 11, + "w": 12, + "x": 12, + "y": 18 + }, + "id": 28, + "options": { + "legend": { + "calcs": [], + "displayMode": "list", + "placement": "right", + "showLegend": true + }, + "tooltip": { + "mode": "single", + "sort": "none" + } + }, + "targets": [ + { + "datasource": { + "type": "prometheus", + "uid": "${DS_PROMETHEUS}" + }, + "editorMode": "code", + "expr": "sum by (status) (rate(meilisearch_http_requests_total{status=~\"(4|5).*\",instance=~\"$instance\",job=~\"$job\"}[$__rate_interval])) / ignoring(status) group_left sum(rate(meilisearch_http_requests_total{instance=~\"$instance\",job=~\"$job\"}[$__rate_interval]))", + "instant": false, + "legendFormat": "HTTP Server - {{status}}", + "range": true, + "refId": "A" + } + ], + "title": "Error Rate", + "type": "timeseries" + }, + { + "datasource": { + "type": "prometheus", + "uid": "${DS_PROMETHEUS}" }, "fieldConfig": { "defaults": { @@ -637,6 +905,7 @@ "mode": "continuous-YlBl" }, "custom": { + "axisBorderShow": false, "axisCenteredZero": false, "axisColorMode": "text", "axisLabel": "", @@ -650,6 +919,7 @@ "tooltip": false, "viz": false }, + "insertNulls": false, "lineInterpolation": "linear", "lineWidth": 1, "pointSize": 5, @@ -687,7 +957,7 @@ "h": 11, "w": 12, "x": 0, - "y": 18 + "y": 29 }, "id": 1, "interval": "5s", @@ -706,7 +976,8 @@ "targets": [ { "datasource": { - "type": "prometheus" + "type": "prometheus", + "uid": "${DS_PROMETHEUS}" }, "editorMode": "builder", "exemplar": true, @@ -731,7 +1002,8 @@ }, "dataFormat": "tsbuckets", "datasource": { - "type": "prometheus" + "type": "prometheus", + "uid": "${DS_PROMETHEUS}" }, "fieldConfig": { "defaults": { @@ -752,7 +1024,7 @@ "h": 11, "w": 12, "x": 12, - "y": 18 + "y": 29 }, "heatmap": {}, "hideZeroBuckets": false, @@ -789,7 +1061,8 @@ }, "showValue": "never", "tooltip": { - "show": true, + "mode": "single", + "showColorScale": false, "yHistogram": false }, "yAxis": { @@ -799,12 +1072,13 @@ "unit": "s" } }, - "pluginVersion": "10.0.1", + "pluginVersion": "10.4.1", "reverseYBuckets": false, "targets": [ { "datasource": { - "type": "prometheus" + "type": "prometheus", + "uid": "${DS_PROMETHEUS}" }, "editorMode": "builder", "exemplar": true, @@ -836,13 +1110,17 @@ "yBucketNumber": 10 }, { - "datasource": {}, + "datasource": { + "type": "prometheus", + "uid": "${DS_PROMETHEUS}" + }, "fieldConfig": { "defaults": { "color": { "mode": "palette-classic" }, "custom": { + "axisBorderShow": false, "axisCenteredZero": false, "axisColorMode": "text", "axisLabel": "", @@ -856,6 +1134,7 @@ "tooltip": false, "viz": false }, + "insertNulls": false, "lineInterpolation": "linear", "lineWidth": 1, "pointSize": 5, @@ -877,8 +1156,7 @@ "mode": "absolute", "steps": [ { - "color": "green", - "value": null + "color": "green" }, { "color": "red", @@ -894,7 +1172,7 @@ "h": 11, "w": 12, "x": 0, - "y": 29 + "y": 40 }, "id": 23, "interval": "5s", @@ -914,7 +1192,8 @@ "targets": [ { "datasource": { - "type": "prometheus" + "type": "prometheus", + "uid": "${DS_PROMETHEUS}" }, "editorMode": "builder", "exemplar": true, @@ -929,13 +1208,17 @@ "type": "timeseries" }, { - "datasource": {}, + "datasource": { + "type": "prometheus", + "uid": "${DS_PROMETHEUS}" + }, "fieldConfig": { "defaults": { "color": { "mode": "palette-classic" }, "custom": { + "axisBorderShow": false, "axisCenteredZero": false, "axisColorMode": "text", "axisLabel": "", @@ -949,6 +1232,7 @@ "tooltip": false, "viz": false }, + "insertNulls": false, "lineInterpolation": "linear", "lineWidth": 1, "pointSize": 5, @@ -970,8 +1254,7 @@ "mode": "absolute", "steps": [ { - "color": "green", - "value": null + "color": "green" }, { "color": "red", @@ -987,7 +1270,7 @@ "h": 11, "w": 12, "x": 12, - "y": 29 + "y": 40 }, "id": 24, "interval": "5s", @@ -1007,7 +1290,8 @@ "targets": [ { "datasource": { - "type": "prometheus" + "type": "prometheus", + "uid": "${DS_PROMETHEUS}" }, "editorMode": "builder", "exemplar": true, @@ -1022,13 +1306,17 @@ "type": "timeseries" }, { - "datasource": {}, + "datasource": { + "type": "prometheus", + "uid": "${DS_PROMETHEUS}" + }, "fieldConfig": { "defaults": { "color": { "mode": "palette-classic" }, "custom": { + "axisBorderShow": false, "axisCenteredZero": false, "axisColorMode": "text", "axisLabel": "", @@ -1042,6 +1330,7 @@ "tooltip": false, "viz": false }, + "insertNulls": false, "lineInterpolation": "linear", "lineWidth": 1, "pointSize": 5, @@ -1063,8 +1352,7 @@ "mode": "absolute", "steps": [ { - "color": "green", - "value": null + "color": "green" }, { "color": "red", @@ -1080,7 +1368,7 @@ "h": 11, "w": 12, "x": 0, - "y": 40 + "y": 51 }, "id": 25, "interval": "5s", @@ -1100,7 +1388,8 @@ "targets": [ { "datasource": { - "type": "prometheus" + "type": "prometheus", + "uid": "${DS_PROMETHEUS}" }, "editorMode": "builder", "exemplar": true, @@ -1115,7 +1404,7 @@ "type": "timeseries" }, { - "collapsed": false, + "collapsed": true, "datasource": { "type": "prometheus", "uid": "i51CxikVz" @@ -1124,10 +1413,227 @@ "h": 1, "w": 24, "x": 0, - "y": 51 + "y": 62 }, "id": 12, - "panels": [], + "panels": [ + { + "datasource": { + "type": "prometheus", + "uid": "${DS_PROMETHEUS}" + }, + "fieldConfig": { + "defaults": { + "color": { + "mode": "continuous-YlBl" + }, + "custom": { + "axisBorderShow": false, + "axisCenteredZero": false, + "axisColorMode": "text", + "axisLabel": "", + "axisPlacement": "auto", + "barAlignment": 0, + "drawStyle": "line", + "fillOpacity": 15, + "gradientMode": "none", + "hideFrom": { + "legend": false, + "tooltip": false, + "viz": false + }, + "insertNulls": false, + "lineInterpolation": "linear", + "lineWidth": 1, + "pointSize": 5, + "scaleDistribution": { + "type": "linear" + }, + "showPoints": "never", + "spanNulls": false, + "stacking": { + "group": "A", + "mode": "none" + }, + "thresholdsStyle": { + "mode": "off" + } + }, + "decimals": 2, + "mappings": [], + "thresholds": { + "mode": "absolute", + "steps": [ + { + "color": "green" + }, + { + "color": "red", + "value": 80 + } + ] + }, + "unit": "percentunit" + }, + "overrides": [] + }, + "gridPos": { + "h": 11, + "w": 12, + "x": 0, + "y": 60 + }, + "id": 4, + "interval": "5s", + "options": { + "legend": { + "calcs": [], + "displayMode": "list", + "placement": "bottom", + "showLegend": true + }, + "tooltip": { + "mode": "single", + "sort": "none" + } + }, + "targets": [ + { + "datasource": { + "type": "prometheus", + "uid": "${DS_PROMETHEUS}" + }, + "editorMode": "builder", + "exemplar": true, + "expr": "rate(process_cpu_seconds_total{job=\"$job\", instance=\"$instance\"}[1m])", + "interval": "", + "legendFormat": "process", + "range": true, + "refId": "A" + }, + { + "datasource": { + "type": "prometheus", + "uid": "${DS_PROMETHEUS}" + }, + "exemplar": true, + "expr": "sum(rate(container_cpu_usage_seconds_total{name='mongodb-redis'}[1m])) by (name)", + "interval": "", + "legendFormat": "container", + "refId": "B" + } + ], + "title": "CPU usage", + "type": "timeseries" + }, + { + "datasource": { + "type": "prometheus", + "uid": "${DS_PROMETHEUS}" + }, + "fieldConfig": { + "defaults": { + "color": { + "mode": "continuous-YlBl" + }, + "custom": { + "axisBorderShow": false, + "axisCenteredZero": false, + "axisColorMode": "text", + "axisLabel": "MiB", + "axisPlacement": "left", + "barAlignment": 0, + "drawStyle": "line", + "fillOpacity": 15, + "gradientMode": "none", + "hideFrom": { + "legend": false, + "tooltip": false, + "viz": false + }, + "insertNulls": false, + "lineInterpolation": "linear", + "lineWidth": 1, + "pointSize": 5, + "scaleDistribution": { + "type": "linear" + }, + "showPoints": "never", + "spanNulls": false, + "stacking": { + "group": "A", + "mode": "none" + }, + "thresholdsStyle": { + "mode": "off" + } + }, + "mappings": [], + "thresholds": { + "mode": "absolute", + "steps": [ + { + "color": "green" + }, + { + "color": "red", + "value": 80 + } + ] + } + }, + "overrides": [] + }, + "gridPos": { + "h": 11, + "w": 12, + "x": 12, + "y": 60 + }, + "id": 5, + "interval": "5s", + "options": { + "legend": { + "calcs": [], + "displayMode": "list", + "placement": "bottom", + "showLegend": true + }, + "tooltip": { + "mode": "single", + "sort": "none" + } + }, + "targets": [ + { + "datasource": { + "type": "prometheus", + "uid": "${DS_PROMETHEUS}" + }, + "editorMode": "builder", + "exemplar": true, + "expr": "process_resident_memory_bytes{job=\"$job\", instance=\"$instance\"} / 1024 / 1024", + "interval": "", + "legendFormat": "process", + "range": true, + "refId": "A" + }, + { + "datasource": { + "type": "prometheus", + "uid": "${DS_PROMETHEUS}" + }, + "exemplar": true, + "expr": "container_memory_usage_bytes{name=\"mongodb-redis\"} / 1024 / 1024", + "interval": "", + "legendFormat": "container", + "refId": "B" + } + ], + "title": "Memory usage", + "type": "timeseries" + } + ], "targets": [ { "datasource": { @@ -1139,230 +1645,18 @@ ], "title": "System metrics", "type": "row" - }, - { - "datasource": { - "type": "prometheus" - }, - "fieldConfig": { - "defaults": { - "color": { - "mode": "continuous-YlBl" - }, - "custom": { - "axisCenteredZero": false, - "axisColorMode": "text", - "axisLabel": "", - "axisPlacement": "auto", - "barAlignment": 0, - "drawStyle": "line", - "fillOpacity": 15, - "gradientMode": "none", - "hideFrom": { - "legend": false, - "tooltip": false, - "viz": false - }, - "lineInterpolation": "linear", - "lineWidth": 1, - "pointSize": 5, - "scaleDistribution": { - "type": "linear" - }, - "showPoints": "never", - "spanNulls": false, - "stacking": { - "group": "A", - "mode": "none" - }, - "thresholdsStyle": { - "mode": "off" - } - }, - "decimals": 2, - "mappings": [], - "thresholds": { - "mode": "absolute", - "steps": [ - { - "color": "green", - "value": null - }, - { - "color": "red", - "value": 80 - } - ] - }, - "unit": "percentunit" - }, - "overrides": [] - }, - "gridPos": { - "h": 11, - "w": 12, - "x": 0, - "y": 52 - }, - "id": 4, - "interval": "5s", - "options": { - "legend": { - "calcs": [], - "displayMode": "list", - "placement": "bottom", - "showLegend": true - }, - "tooltip": { - "mode": "single", - "sort": "none" - } - }, - "targets": [ - { - "datasource": { - "type": "prometheus" - }, - "editorMode": "builder", - "exemplar": true, - "expr": "rate(process_cpu_seconds_total{job=\"$job\", instance=\"$instance\"}[1m])", - "interval": "", - "legendFormat": "process", - "range": true, - "refId": "A" - }, - { - "datasource": { - "type": "prometheus" - }, - "exemplar": true, - "expr": "sum(rate(container_cpu_usage_seconds_total{name='mongodb-redis'}[1m])) by (name)", - "interval": "", - "legendFormat": "container", - "refId": "B" - } - ], - "title": "CPU usage", - "type": "timeseries" - }, - { - "datasource": { - "type": "prometheus" - }, - "fieldConfig": { - "defaults": { - "color": { - "mode": "continuous-YlBl" - }, - "custom": { - "axisCenteredZero": false, - "axisColorMode": "text", - "axisLabel": "MiB", - "axisPlacement": "left", - "barAlignment": 0, - "drawStyle": "line", - "fillOpacity": 15, - "gradientMode": "none", - "hideFrom": { - "legend": false, - "tooltip": false, - "viz": false - }, - "lineInterpolation": "linear", - "lineWidth": 1, - "pointSize": 5, - "scaleDistribution": { - "type": "linear" - }, - "showPoints": "never", - "spanNulls": false, - "stacking": { - "group": "A", - "mode": "none" - }, - "thresholdsStyle": { - "mode": "off" - } - }, - "mappings": [], - "thresholds": { - "mode": "absolute", - "steps": [ - { - "color": "green", - "value": null - }, - { - "color": "red", - "value": 80 - } - ] - } - }, - "overrides": [] - }, - "gridPos": { - "h": 11, - "w": 12, - "x": 12, - "y": 52 - }, - "id": 5, - "interval": "5s", - "options": { - "legend": { - "calcs": [], - "displayMode": "list", - "placement": "bottom", - "showLegend": true - }, - "tooltip": { - "mode": "single", - "sort": "none" - } - }, - "targets": [ - { - "datasource": { - "type": "prometheus" - }, - "editorMode": "builder", - "exemplar": true, - "expr": "process_resident_memory_bytes{job=\"$job\", instance=\"$instance\"} / 1024 / 1024", - "interval": "", - "legendFormat": "process", - "range": true, - "refId": "A" - }, - { - "datasource": { - "type": "prometheus" - }, - "exemplar": true, - "expr": "container_memory_usage_bytes{name=\"mongodb-redis\"} / 1024 / 1024", - "interval": "", - "legendFormat": "container", - "refId": "B" - } - ], - "title": "Memory usage", - "type": "timeseries" } ], - "refresh": "5s", - "schemaVersion": 38, - "style": "dark", + "refresh": "10s", + "schemaVersion": 39, "tags": [], "templating": { "list": [ { - "current": { - "selected": false, - "text": "localhost:7700", - "value": "localhost:7700" - }, + "current": {}, "datasource": { - "type": "prometheus" + "type": "prometheus", + "uid": "${DS_PROMETHEUS}" }, "definition": "label_values(instance)", "hide": 0, @@ -1382,13 +1676,10 @@ "type": "query" }, { - "current": { - "selected": false, - "text": "index-word-count-10-count", - "value": "index-word-count-10-count" - }, + "current": {}, "datasource": { - "type": "prometheus" + "type": "prometheus", + "uid": "${DS_PROMETHEUS}" }, "definition": "label_values(index)", "hide": 0, @@ -1408,13 +1699,10 @@ "type": "query" }, { - "current": { - "selected": true, - "text": "meilisearch", - "value": "meilisearch" - }, + "current": {}, "datasource": { - "type": "prometheus" + "type": "prometheus", + "uid": "${DS_PROMETHEUS}" }, "definition": "label_values(job)", "description": "Prometheus job_name from scrape config (default is meilisearch)", @@ -1452,6 +1740,6 @@ "timezone": "", "title": "Meilisearch", "uid": "7wcZ94dnz", - "version": 5, + "version": 2, "weekStart": "" -} +} \ No newline at end of file From 13a84ae557b151ca163071fadb90e6c91cb97052 Mon Sep 17 00:00:00 2001 From: Rohan Kumar Date: Mon, 25 Mar 2024 11:07:07 +0530 Subject: [PATCH 26/86] fix: set the histogram bucket boundaries to follow the otel spec --- meilisearch/src/metrics.rs | 14 +------------- 1 file changed, 1 insertion(+), 13 deletions(-) diff --git a/meilisearch/src/metrics.rs b/meilisearch/src/metrics.rs index 652e6c227..5f73ccf33 100644 --- a/meilisearch/src/metrics.rs +++ b/meilisearch/src/metrics.rs @@ -4,19 +4,7 @@ use prometheus::{ register_int_gauge_vec, HistogramVec, IntCounterVec, IntGauge, IntGaugeVec, }; -/// Create evenly distributed buckets -fn create_buckets() -> [f64; 29] { - (0..10) - .chain((10..100).step_by(10)) - .chain((100..=1000).step_by(100)) - .map(|i| i as f64 / 1000.) - .collect::>() - .try_into() - .unwrap() -} - lazy_static! { - pub static ref MEILISEARCH_HTTP_RESPONSE_TIME_CUSTOM_BUCKETS: [f64; 29] = create_buckets(); pub static ref MEILISEARCH_HTTP_REQUESTS_TOTAL: IntCounterVec = register_int_counter_vec!( opts!("meilisearch_http_requests_total", "Meilisearch HTTP requests total"), &["method", "path"] @@ -47,7 +35,7 @@ lazy_static! { "meilisearch_http_response_time_seconds", "Meilisearch HTTP response times", &["method", "path"], - MEILISEARCH_HTTP_RESPONSE_TIME_CUSTOM_BUCKETS.to_vec() + vec![0.005, 0.01, 0.025, 0.05, 0.075, 0.1, 0.25, 0.5, 0.75, 1.0, 2.5, 5.0, 7.5, 10.0] ) .expect("Can't create a metric"); pub static ref MEILISEARCH_NB_TASKS: IntGaugeVec = register_int_gauge_vec!( From bc58e8a310aa0774265c1e51ec38162cff526da2 Mon Sep 17 00:00:00 2001 From: Louis Dureuil Date: Tue, 12 Mar 2024 15:00:26 +0100 Subject: [PATCH 27/86] Documentation for the vector module --- milli/src/vector/mod.rs | 66 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 66 insertions(+) diff --git a/milli/src/vector/mod.rs b/milli/src/vector/mod.rs index 035ac555e..aeb0be1ca 100644 --- a/milli/src/vector/mod.rs +++ b/milli/src/vector/mod.rs @@ -16,46 +16,62 @@ pub use self::error::Error; pub type Embedding = Vec; +/// One or multiple embeddings stored consecutively in a flat vector. pub struct Embeddings { data: Vec, dimension: usize, } impl Embeddings { + /// Declares an empty vector of embeddings of the specified dimensions. pub fn new(dimension: usize) -> Self { Self { data: Default::default(), dimension } } + /// Declares a vector of embeddings containing a single element. + /// + /// The dimension is inferred from the length of the passed embedding. pub fn from_single_embedding(embedding: Vec) -> Self { Self { dimension: embedding.len(), data: embedding } } + /// Declares a vector of embeddings from its components. + /// + /// `data.len()` must be a multiple of `dimension`, otherwise an error is returned. pub fn from_inner(data: Vec, dimension: usize) -> Result> { let mut this = Self::new(dimension); this.append(data)?; Ok(this) } + /// Returns the number of embeddings in this vector of embeddings. pub fn embedding_count(&self) -> usize { self.data.len() / self.dimension } + /// Dimension of a single embedding. pub fn dimension(&self) -> usize { self.dimension } + /// Deconstructs self into the inner flat vector. pub fn into_inner(self) -> Vec { self.data } + /// A reference to the inner flat vector. pub fn as_inner(&self) -> &[F] { &self.data } + /// Iterates over the embeddings contained in the flat vector. pub fn iter(&self) -> impl Iterator + '_ { self.data.as_slice().chunks_exact(self.dimension) } + /// Push an embedding at the end of the embeddings. + /// + /// If `embedding.len() != self.dimension`, then the push operation fails. pub fn push(&mut self, mut embedding: Vec) -> Result<(), Vec> { if embedding.len() != self.dimension { return Err(embedding); @@ -64,6 +80,9 @@ impl Embeddings { Ok(()) } + /// Append a flat vector of embeddings a the end of the embeddings. + /// + /// If `embeddings.len() % self.dimension != 0`, then the append operation fails. pub fn append(&mut self, mut embeddings: Vec) -> Result<(), Vec> { if embeddings.len() % self.dimension != 0 { return Err(embeddings); @@ -73,37 +92,57 @@ impl Embeddings { } } +/// An embedder can be used to transform text into embeddings. #[derive(Debug)] pub enum Embedder { + /// An embedder based on running local models, fetched from the Hugging Face Hub. HuggingFace(hf::Embedder), + /// An embedder based on making embedding queries against the OpenAI API. OpenAi(openai::Embedder), + /// An embedder based on the user providing the embeddings in the documents and queries. UserProvided(manual::Embedder), Ollama(ollama::Embedder), } +/// Configuration for an embedder. #[derive(Debug, Clone, Default, serde::Deserialize, serde::Serialize)] pub struct EmbeddingConfig { + /// Options of the embedder, specific to each kind of embedder pub embedder_options: EmbedderOptions, + /// Document template pub prompt: PromptData, // TODO: add metrics and anything needed } +/// Map of embedder configurations. +/// +/// Each configuration is mapped to a name. #[derive(Clone, Default)] pub struct EmbeddingConfigs(HashMap, Arc)>); impl EmbeddingConfigs { + /// Create the map from its internal component.s pub fn new(data: HashMap, Arc)>) -> Self { Self(data) } + /// Get an embedder configuration and template from its name. pub fn get(&self, name: &str) -> Option<(Arc, Arc)> { self.0.get(name).cloned() } + /// Get the default embedder configuration, if any. pub fn get_default(&self) -> Option<(Arc, Arc)> { self.get_default_embedder_name().and_then(|default| self.get(&default)) } + /// Get the name of the default embedder configuration. + /// + /// The default embedder is determined as follows: + /// + /// - If there is only one embedder, it is always the default. + /// - If there are multiple embedders and one of them is called `default`, then that one is the default embedder. + /// - In all other cases, there is no default embedder. pub fn get_default_embedder_name(&self) -> Option { let mut it = self.0.keys(); let first_name = it.next(); @@ -126,6 +165,7 @@ impl IntoIterator for EmbeddingConfigs { } } +/// Options of an embedder, specific to each kind of embedder. #[derive(Debug, Clone, Hash, PartialEq, Eq, serde::Deserialize, serde::Serialize)] pub enum EmbedderOptions { HuggingFace(hf::EmbedderOptions), @@ -141,10 +181,12 @@ impl Default for EmbedderOptions { } impl EmbedderOptions { + /// Default options for the Hugging Face embedder pub fn huggingface() -> Self { Self::HuggingFace(hf::EmbedderOptions::new()) } + /// Default options for the OpenAI embedder pub fn openai(api_key: Option) -> Self { Self::OpenAi(openai::EmbedderOptions::with_default_model(api_key)) } @@ -155,6 +197,7 @@ impl EmbedderOptions { } impl Embedder { + /// Spawns a new embedder built from its options. pub fn new(options: EmbedderOptions) -> std::result::Result { Ok(match options { EmbedderOptions::HuggingFace(options) => Self::HuggingFace(hf::Embedder::new(options)?), @@ -166,6 +209,9 @@ impl Embedder { }) } + /// Embed one or multiple texts. + /// + /// Each text can be embedded as one or multiple embeddings. pub async fn embed( &self, texts: Vec, @@ -184,6 +230,10 @@ impl Embedder { } } + /// Embed multiple chunks of texts. + /// + /// Each chunk is composed of one or multiple texts. + /// /// # Panics /// /// - if called from an asynchronous context @@ -199,6 +249,7 @@ impl Embedder { } } + /// Indicates the preferred number of chunks to pass to [`Self::embed_chunks`] pub fn chunk_count_hint(&self) -> usize { match self { Embedder::HuggingFace(embedder) => embedder.chunk_count_hint(), @@ -208,6 +259,7 @@ impl Embedder { } } + /// Indicates the preferred number of texts in a single chunk passed to [`Self::embed`] pub fn prompt_count_in_chunk_hint(&self) -> usize { match self { Embedder::HuggingFace(embedder) => embedder.prompt_count_in_chunk_hint(), @@ -217,6 +269,7 @@ impl Embedder { } } + /// Indicates the dimensions of a single embedding produced by the embedder. pub fn dimensions(&self) -> usize { match self { Embedder::HuggingFace(embedder) => embedder.dimensions(), @@ -226,6 +279,7 @@ impl Embedder { } } + /// An optional distribution used to apply an affine transformation to the similarity score of a document. pub fn distribution(&self) -> Option { match self { Embedder::HuggingFace(embedder) => embedder.distribution(), @@ -236,9 +290,20 @@ impl Embedder { } } +/// Describes the mean and sigma of distribution of embedding similarity in the embedding space. +/// +/// The intended use is to make the similarity score more comparable to the regular ranking score. +/// This allows to correct effects where results are too "packed" around a certain value. #[derive(Debug, Clone, Copy)] pub struct DistributionShift { + /// Value where the results are "packed". + /// + /// Similarity scores are translated so that they are packed around 0.5 instead pub current_mean: f32, + + /// standard deviation of a similarity score. + /// + /// Set below 0.4 to make the results less packed around the mean, and above 0.4 to make them more packed. pub current_sigma: f32, } @@ -280,6 +345,7 @@ impl DistributionShift { } } +/// Whether CUDA is supported in this version of Meilisearch. pub const fn is_cuda_enabled() -> bool { cfg!(feature = "cuda") } From c3d02f092dddf5f3e0f336f7774cca72eb8ed0bb Mon Sep 17 00:00:00 2001 From: Louis Dureuil Date: Thu, 14 Mar 2024 11:14:31 +0100 Subject: [PATCH 28/86] OpenAI sync --- Cargo.lock | 1 + milli/Cargo.toml | 1 + milli/src/vector/error.rs | 28 +- milli/src/vector/hf.rs | 2 +- milli/src/vector/mod.rs | 9 +- milli/src/vector/openai.rs | 554 +++++++++++++++++-------------------- 6 files changed, 274 insertions(+), 321 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index b44b151d1..60d0e4c0e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3378,6 +3378,7 @@ dependencies = [ "tokenizers", "tokio", "tracing", + "ureq", "uuid", ] diff --git a/milli/Cargo.toml b/milli/Cargo.toml index fa4215404..59b3699cc 100644 --- a/milli/Cargo.toml +++ b/milli/Cargo.toml @@ -91,6 +91,7 @@ liquid = "0.26.4" arroy = "0.2.0" rand = "0.8.5" tracing = "0.1.40" +ureq = { version = "2.9.6", features = ["json"] } [dev-dependencies] mimalloc = { version = "0.1.39", default-features = false } diff --git a/milli/src/vector/error.rs b/milli/src/vector/error.rs index 9bbdeaa90..1def4f7a9 100644 --- a/milli/src/vector/error.rs +++ b/milli/src/vector/error.rs @@ -53,17 +53,17 @@ pub enum EmbedErrorKind { #[error("could not run model: {0}")] ModelForward(candle_core::Error), #[error("could not reach OpenAI: {0}")] - OpenAiNetwork(reqwest::Error), + OpenAiNetwork(ureq::Transport), #[error("unexpected response from OpenAI: {0}")] - OpenAiUnexpected(reqwest::Error), - #[error("could not authenticate against OpenAI: {0}")] - OpenAiAuth(OpenAiError), - #[error("sent too many requests to OpenAI: {0}")] - OpenAiTooManyRequests(OpenAiError), + OpenAiUnexpected(ureq::Error), + #[error("could not authenticate against OpenAI: {0:?}")] + OpenAiAuth(Option), + #[error("sent too many requests to OpenAI: {0:?}")] + OpenAiTooManyRequests(Option), #[error("received internal error from OpenAI: {0:?}")] OpenAiInternalServerError(Option), - #[error("sent too many tokens in a request to OpenAI: {0}")] - OpenAiTooManyTokens(OpenAiError), + #[error("sent too many tokens in a request to OpenAI: {0:?}")] + OpenAiTooManyTokens(Option), #[error("received unhandled HTTP status code {0} from OpenAI")] OpenAiUnhandledStatusCode(u16), #[error("attempt to embed the following text in a configuration where embeddings must be user provided: {0:?}")] @@ -102,19 +102,19 @@ impl EmbedError { Self { kind: EmbedErrorKind::ModelForward(inner), fault: FaultSource::Runtime } } - pub fn openai_network(inner: reqwest::Error) -> Self { + pub fn openai_network(inner: ureq::Transport) -> Self { Self { kind: EmbedErrorKind::OpenAiNetwork(inner), fault: FaultSource::Runtime } } - pub fn openai_unexpected(inner: reqwest::Error) -> EmbedError { + pub fn openai_unexpected(inner: ureq::Error) -> EmbedError { Self { kind: EmbedErrorKind::OpenAiUnexpected(inner), fault: FaultSource::Bug } } - pub(crate) fn openai_auth_error(inner: OpenAiError) -> EmbedError { + pub(crate) fn openai_auth_error(inner: Option) -> EmbedError { Self { kind: EmbedErrorKind::OpenAiAuth(inner), fault: FaultSource::User } } - pub(crate) fn openai_too_many_requests(inner: OpenAiError) -> EmbedError { + pub(crate) fn openai_too_many_requests(inner: Option) -> EmbedError { Self { kind: EmbedErrorKind::OpenAiTooManyRequests(inner), fault: FaultSource::Runtime } } @@ -122,7 +122,7 @@ impl EmbedError { Self { kind: EmbedErrorKind::OpenAiInternalServerError(inner), fault: FaultSource::Runtime } } - pub(crate) fn openai_too_many_tokens(inner: OpenAiError) -> EmbedError { + pub(crate) fn openai_too_many_tokens(inner: Option) -> EmbedError { Self { kind: EmbedErrorKind::OpenAiTooManyTokens(inner), fault: FaultSource::Bug } } @@ -220,7 +220,7 @@ impl NewEmbedderError { Self { kind: NewEmbedderErrorKind::LoadModel(inner), fault: FaultSource::Runtime } } - pub fn hf_could_not_determine_dimension(inner: EmbedError) -> NewEmbedderError { + pub fn could_not_determine_dimension(inner: EmbedError) -> NewEmbedderError { Self { kind: NewEmbedderErrorKind::CouldNotDetermineDimension(inner), fault: FaultSource::Runtime, diff --git a/milli/src/vector/hf.rs b/milli/src/vector/hf.rs index 04e169c71..939b6210a 100644 --- a/milli/src/vector/hf.rs +++ b/milli/src/vector/hf.rs @@ -131,7 +131,7 @@ impl Embedder { let embeddings = this .embed(vec!["test".into()]) - .map_err(NewEmbedderError::hf_could_not_determine_dimension)?; + .map_err(NewEmbedderError::could_not_determine_dimension)?; this.dimensions = embeddings.first().unwrap().dimension(); Ok(this) diff --git a/milli/src/vector/mod.rs b/milli/src/vector/mod.rs index aeb0be1ca..86dde8ad4 100644 --- a/milli/src/vector/mod.rs +++ b/milli/src/vector/mod.rs @@ -98,7 +98,7 @@ pub enum Embedder { /// An embedder based on running local models, fetched from the Hugging Face Hub. HuggingFace(hf::Embedder), /// An embedder based on making embedding queries against the OpenAI API. - OpenAi(openai::Embedder), + OpenAi(openai::sync::Embedder), /// An embedder based on the user providing the embeddings in the documents and queries. UserProvided(manual::Embedder), Ollama(ollama::Embedder), @@ -201,7 +201,7 @@ impl Embedder { pub fn new(options: EmbedderOptions) -> std::result::Result { Ok(match options { EmbedderOptions::HuggingFace(options) => Self::HuggingFace(hf::Embedder::new(options)?), - EmbedderOptions::OpenAi(options) => Self::OpenAi(openai::Embedder::new(options)?), + EmbedderOptions::OpenAi(options) => Self::OpenAi(openai::sync::Embedder::new(options)?), EmbedderOptions::Ollama(options) => Self::Ollama(ollama::Embedder::new(options)?), EmbedderOptions::UserProvided(options) => { Self::UserProvided(manual::Embedder::new(options)) @@ -218,10 +218,7 @@ impl Embedder { ) -> std::result::Result>, EmbedError> { match self { Embedder::HuggingFace(embedder) => embedder.embed(texts), - Embedder::OpenAi(embedder) => { - let client = embedder.new_client()?; - embedder.embed(texts, &client).await - } + Embedder::OpenAi(embedder) => embedder.embed(texts), Embedder::Ollama(embedder) => { let client = embedder.new_client()?; embedder.embed(texts, &client).await diff --git a/milli/src/vector/openai.rs b/milli/src/vector/openai.rs index dcf3f4c89..5d13d5ee2 100644 --- a/milli/src/vector/openai.rs +++ b/milli/src/vector/openai.rs @@ -1,18 +1,10 @@ use std::fmt::Display; -use reqwest::StatusCode; use serde::{Deserialize, Serialize}; use super::error::{EmbedError, NewEmbedderError}; use super::{DistributionShift, Embedding, Embeddings}; -#[derive(Debug)] -pub struct Embedder { - headers: reqwest::header::HeaderMap, - tokenizer: tiktoken_rs::CoreBPE, - options: EmbedderOptions, -} - #[derive(Debug, Clone, Hash, PartialEq, Eq, serde::Deserialize, serde::Serialize)] pub struct EmbedderOptions { pub api_key: Option, @@ -125,298 +117,6 @@ impl EmbedderOptions { } } -impl Embedder { - pub fn new_client(&self) -> Result { - reqwest::ClientBuilder::new() - .default_headers(self.headers.clone()) - .build() - .map_err(EmbedError::openai_initialize_web_client) - } - - pub fn new(options: EmbedderOptions) -> Result { - let mut headers = reqwest::header::HeaderMap::new(); - let mut inferred_api_key = Default::default(); - let api_key = options.api_key.as_ref().unwrap_or_else(|| { - inferred_api_key = infer_api_key(); - &inferred_api_key - }); - headers.insert( - reqwest::header::AUTHORIZATION, - reqwest::header::HeaderValue::from_str(&format!("Bearer {}", api_key)) - .map_err(NewEmbedderError::openai_invalid_api_key_format)?, - ); - headers.insert( - reqwest::header::CONTENT_TYPE, - reqwest::header::HeaderValue::from_static("application/json"), - ); - - // looking at the code it is very unclear that this can actually fail. - let tokenizer = tiktoken_rs::cl100k_base().unwrap(); - - Ok(Self { options, headers, tokenizer }) - } - - pub async fn embed( - &self, - texts: Vec, - client: &reqwest::Client, - ) -> Result>, EmbedError> { - let mut tokenized = false; - - for attempt in 0..7 { - let result = if tokenized { - self.try_embed_tokenized(&texts, client).await - } else { - self.try_embed(&texts, client).await - }; - - let retry_duration = match result { - Ok(embeddings) => return Ok(embeddings), - Err(retry) => { - tracing::warn!("Failed: {}", retry.error); - tokenized |= retry.must_tokenize(); - retry.into_duration(attempt) - } - }?; - - let retry_duration = retry_duration.min(std::time::Duration::from_secs(60)); // don't wait more than a minute - tracing::warn!( - "Attempt #{}, retrying after {}ms.", - attempt, - retry_duration.as_millis() - ); - tokio::time::sleep(retry_duration).await; - } - - let result = if tokenized { - self.try_embed_tokenized(&texts, client).await - } else { - self.try_embed(&texts, client).await - }; - - result.map_err(Retry::into_error) - } - - async fn check_response(response: reqwest::Response) -> Result { - if !response.status().is_success() { - match response.status() { - StatusCode::UNAUTHORIZED => { - let error_response: OpenAiErrorResponse = response - .json() - .await - .map_err(EmbedError::openai_unexpected) - .map_err(Retry::retry_later)?; - - return Err(Retry::give_up(EmbedError::openai_auth_error( - error_response.error, - ))); - } - StatusCode::TOO_MANY_REQUESTS => { - let error_response: OpenAiErrorResponse = response - .json() - .await - .map_err(EmbedError::openai_unexpected) - .map_err(Retry::retry_later)?; - - return Err(Retry::rate_limited(EmbedError::openai_too_many_requests( - error_response.error, - ))); - } - StatusCode::INTERNAL_SERVER_ERROR - | StatusCode::BAD_GATEWAY - | StatusCode::SERVICE_UNAVAILABLE => { - let error_response: Result = response.json().await; - return Err(Retry::retry_later(EmbedError::openai_internal_server_error( - error_response.ok().map(|error_response| error_response.error), - ))); - } - StatusCode::BAD_REQUEST => { - // Most probably, one text contained too many tokens - let error_response: OpenAiErrorResponse = response - .json() - .await - .map_err(EmbedError::openai_unexpected) - .map_err(Retry::retry_later)?; - - tracing::warn!("OpenAI: received `BAD_REQUEST`. Input was maybe too long, retrying on tokenized version. For best performance, limit the size of your prompt."); - - return Err(Retry::retry_tokenized(EmbedError::openai_too_many_tokens( - error_response.error, - ))); - } - code => { - return Err(Retry::retry_later(EmbedError::openai_unhandled_status_code( - code.as_u16(), - ))); - } - } - } - Ok(response) - } - - async fn try_embed + serde::Serialize>( - &self, - texts: &[S], - client: &reqwest::Client, - ) -> Result>, Retry> { - for text in texts { - tracing::trace!("Received prompt: {}", text.as_ref()) - } - let request = OpenAiRequest { - model: self.options.embedding_model.name(), - input: texts, - dimensions: self.overriden_dimensions(), - }; - let response = client - .post(OPENAI_EMBEDDINGS_URL) - .json(&request) - .send() - .await - .map_err(EmbedError::openai_network) - .map_err(Retry::retry_later)?; - - let response = Self::check_response(response).await?; - - let response: OpenAiResponse = response - .json() - .await - .map_err(EmbedError::openai_unexpected) - .map_err(Retry::retry_later)?; - - tracing::trace!("response: {:?}", response.data); - - Ok(response - .data - .into_iter() - .map(|data| Embeddings::from_single_embedding(data.embedding)) - .collect()) - } - - async fn try_embed_tokenized( - &self, - text: &[String], - client: &reqwest::Client, - ) -> Result>, Retry> { - pub const OVERLAP_SIZE: usize = 200; - let mut all_embeddings = Vec::with_capacity(text.len()); - for text in text { - let max_token_count = self.options.embedding_model.max_token(); - let encoded = self.tokenizer.encode_ordinary(text.as_str()); - let len = encoded.len(); - if len < max_token_count { - all_embeddings.append(&mut self.try_embed(&[text], client).await?); - continue; - } - - let mut tokens = encoded.as_slice(); - let mut embeddings_for_prompt = Embeddings::new(self.dimensions()); - while tokens.len() > max_token_count { - let window = &tokens[..max_token_count]; - embeddings_for_prompt.push(self.embed_tokens(window, client).await?).unwrap(); - - tokens = &tokens[max_token_count - OVERLAP_SIZE..]; - } - - // end of text - embeddings_for_prompt.push(self.embed_tokens(tokens, client).await?).unwrap(); - - all_embeddings.push(embeddings_for_prompt); - } - Ok(all_embeddings) - } - - async fn embed_tokens( - &self, - tokens: &[usize], - client: &reqwest::Client, - ) -> Result { - for attempt in 0..9 { - let duration = match self.try_embed_tokens(tokens, client).await { - Ok(embedding) => return Ok(embedding), - Err(retry) => retry.into_duration(attempt), - } - .map_err(Retry::retry_later)?; - - tokio::time::sleep(duration).await; - } - - self.try_embed_tokens(tokens, client) - .await - .map_err(|retry| Retry::give_up(retry.into_error())) - } - - async fn try_embed_tokens( - &self, - tokens: &[usize], - client: &reqwest::Client, - ) -> Result { - let request = OpenAiTokensRequest { - model: self.options.embedding_model.name(), - input: tokens, - dimensions: self.overriden_dimensions(), - }; - let response = client - .post(OPENAI_EMBEDDINGS_URL) - .json(&request) - .send() - .await - .map_err(EmbedError::openai_network) - .map_err(Retry::retry_later)?; - - let response = Self::check_response(response).await?; - - let mut response: OpenAiResponse = response - .json() - .await - .map_err(EmbedError::openai_unexpected) - .map_err(Retry::retry_later)?; - Ok(response.data.pop().map(|data| data.embedding).unwrap_or_default()) - } - - pub fn embed_chunks( - &self, - text_chunks: Vec>, - ) -> Result>>, EmbedError> { - let rt = tokio::runtime::Builder::new_current_thread() - .enable_io() - .enable_time() - .build() - .map_err(EmbedError::openai_runtime_init)?; - let client = self.new_client()?; - rt.block_on(futures::future::try_join_all( - text_chunks.into_iter().map(|prompts| self.embed(prompts, &client)), - )) - } - - pub fn chunk_count_hint(&self) -> usize { - 10 - } - - pub fn prompt_count_in_chunk_hint(&self) -> usize { - 10 - } - - pub fn dimensions(&self) -> usize { - if self.options.embedding_model.supports_overriding_dimensions() { - self.options.dimensions.unwrap_or(self.options.embedding_model.default_dimensions()) - } else { - self.options.embedding_model.default_dimensions() - } - } - - pub fn distribution(&self) -> Option { - self.options.embedding_model.distribution() - } - - fn overriden_dimensions(&self) -> Option { - if self.options.embedding_model.supports_overriding_dimensions() { - self.options.dimensions - } else { - None - } - } -} - // retrying in case of failure pub struct Retry { @@ -524,3 +224,257 @@ fn infer_api_key() -> String { .or_else(|_| std::env::var("OPENAI_API_KEY")) .unwrap_or_default() } + +pub mod sync { + use rayon::iter::{IntoParallelIterator, ParallelIterator as _}; + + use super::{ + EmbedError, Embedding, Embeddings, NewEmbedderError, OpenAiErrorResponse, OpenAiRequest, + OpenAiResponse, OpenAiTokensRequest, Retry, OPENAI_EMBEDDINGS_URL, + }; + use crate::vector::DistributionShift; + + const REQUEST_PARALLELISM: usize = 10; + + #[derive(Debug)] + pub struct Embedder { + tokenizer: tiktoken_rs::CoreBPE, + options: super::EmbedderOptions, + bearer: String, + threads: rayon::ThreadPool, + } + + impl Embedder { + pub fn new(options: super::EmbedderOptions) -> Result { + let mut inferred_api_key = Default::default(); + let api_key = options.api_key.as_ref().unwrap_or_else(|| { + inferred_api_key = super::infer_api_key(); + &inferred_api_key + }); + let bearer = format!("Bearer {api_key}"); + + // looking at the code it is very unclear that this can actually fail. + let tokenizer = tiktoken_rs::cl100k_base().unwrap(); + + // FIXME: unwrap + let threads = rayon::ThreadPoolBuilder::new() + .num_threads(REQUEST_PARALLELISM) + .thread_name(|index| format!("embedder-chunk-{index}")) + .build() + .unwrap(); + + Ok(Self { options, bearer, tokenizer, threads }) + } + + pub fn embed(&self, texts: Vec) -> Result>, EmbedError> { + let mut tokenized = false; + + let client = ureq::agent(); + + for attempt in 0..7 { + let result = if tokenized { + self.try_embed_tokenized(&texts, &client) + } else { + self.try_embed(&texts, &client) + }; + + let retry_duration = match result { + Ok(embeddings) => return Ok(embeddings), + Err(retry) => { + tracing::warn!("Failed: {}", retry.error); + tokenized |= retry.must_tokenize(); + retry.into_duration(attempt) + } + }?; + + let retry_duration = retry_duration.min(std::time::Duration::from_secs(60)); // don't wait more than a minute + tracing::warn!( + "Attempt #{}, retrying after {}ms.", + attempt, + retry_duration.as_millis() + ); + std::thread::sleep(retry_duration); + } + + let result = if tokenized { + self.try_embed_tokenized(&texts, &client) + } else { + self.try_embed(&texts, &client) + }; + + result.map_err(Retry::into_error) + } + + fn check_response( + response: Result, + ) -> Result { + match response { + Ok(response) => Ok(response), + Err(ureq::Error::Status(code, response)) => { + let error_response: Option = response.into_json().ok(); + let error = error_response.map(|response| response.error); + Err(match code { + 401 => Retry::give_up(EmbedError::openai_auth_error(error)), + 429 => Retry::rate_limited(EmbedError::openai_too_many_requests(error)), + 400 => { + tracing::warn!("OpenAI: received `BAD_REQUEST`. Input was maybe too long, retrying on tokenized version. For best performance, limit the size of your document template."); + + Retry::retry_tokenized(EmbedError::openai_too_many_tokens(error)) + } + 500..=599 => { + Retry::retry_later(EmbedError::openai_internal_server_error(error)) + } + x => Retry::retry_later(EmbedError::openai_unhandled_status_code(code)), + }) + } + Err(ureq::Error::Transport(transport)) => { + Err(Retry::retry_later(EmbedError::openai_network(transport))) + } + } + } + + fn try_embed + serde::Serialize>( + &self, + texts: &[S], + client: &ureq::Agent, + ) -> Result>, Retry> { + for text in texts { + tracing::trace!("Received prompt: {}", text.as_ref()) + } + let request = OpenAiRequest { + model: self.options.embedding_model.name(), + input: texts, + dimensions: self.overriden_dimensions(), + }; + let response = client + .post(OPENAI_EMBEDDINGS_URL) + .set("Authorization", &self.bearer) + .send_json(&request); + + let response = Self::check_response(response)?; + + let response: OpenAiResponse = response + .into_json() + .map_err(EmbedError::openai_unexpected) + .map_err(Retry::retry_later)?; + + tracing::trace!("response: {:?}", response.data); + + Ok(response + .data + .into_iter() + .map(|data| Embeddings::from_single_embedding(data.embedding)) + .collect()) + } + + fn try_embed_tokenized( + &self, + text: &[String], + client: &ureq::Agent, + ) -> Result>, Retry> { + pub const OVERLAP_SIZE: usize = 200; + let mut all_embeddings = Vec::with_capacity(text.len()); + for text in text { + let max_token_count = self.options.embedding_model.max_token(); + let encoded = self.tokenizer.encode_ordinary(text.as_str()); + let len = encoded.len(); + if len < max_token_count { + all_embeddings.append(&mut self.try_embed(&[text], client)?); + continue; + } + + let mut tokens = encoded.as_slice(); + let mut embeddings_for_prompt = Embeddings::new(self.dimensions()); + while tokens.len() > max_token_count { + let window = &tokens[..max_token_count]; + embeddings_for_prompt.push(self.embed_tokens(window, client)?).unwrap(); + + tokens = &tokens[max_token_count - OVERLAP_SIZE..]; + } + + // end of text + embeddings_for_prompt.push(self.embed_tokens(tokens, client)?).unwrap(); + + all_embeddings.push(embeddings_for_prompt); + } + Ok(all_embeddings) + } + + fn embed_tokens(&self, tokens: &[usize], client: &ureq::Agent) -> Result { + for attempt in 0..9 { + let duration = match self.try_embed_tokens(tokens, client) { + Ok(embedding) => return Ok(embedding), + Err(retry) => retry.into_duration(attempt), + } + .map_err(Retry::retry_later)?; + + std::thread::sleep(duration); + } + + self.try_embed_tokens(tokens, client) + .map_err(|retry| Retry::give_up(retry.into_error())) + } + + fn try_embed_tokens( + &self, + tokens: &[usize], + client: &ureq::Agent, + ) -> Result { + let request = OpenAiTokensRequest { + model: self.options.embedding_model.name(), + input: tokens, + dimensions: self.overriden_dimensions(), + }; + let response = client + .post(OPENAI_EMBEDDINGS_URL) + .set("Authorization", &self.bearer) + .send_json(&request); + + let response = Self::check_response(response)?; + + let mut response: OpenAiResponse = response + .into_json() + .map_err(EmbedError::openai_unexpected) + .map_err(Retry::retry_later)?; + + Ok(response.data.pop().map(|data| data.embedding).unwrap_or_default()) + } + + pub fn embed_chunks( + &self, + text_chunks: Vec>, + ) -> Result>>, EmbedError> { + self.threads + .install(move || text_chunks.into_par_iter().map(|chunk| self.embed(chunk))) + .collect() + } + + pub fn chunk_count_hint(&self) -> usize { + 10 + } + + pub fn prompt_count_in_chunk_hint(&self) -> usize { + 10 + } + + pub fn dimensions(&self) -> usize { + if self.options.embedding_model.supports_overriding_dimensions() { + self.options.dimensions.unwrap_or(self.options.embedding_model.default_dimensions()) + } else { + self.options.embedding_model.default_dimensions() + } + } + + pub fn distribution(&self) -> Option { + self.options.embedding_model.distribution() + } + + fn overriden_dimensions(&self) -> Option { + if self.options.embedding_model.supports_overriding_dimensions() { + self.options.dimensions + } else { + None + } + } + } +} From 8708cbef2538d28c65b7511e9706b9c1a093762a Mon Sep 17 00:00:00 2001 From: Louis Dureuil Date: Thu, 14 Mar 2024 14:44:43 +0100 Subject: [PATCH 29/86] Add RestEmbedder --- milli/src/vector/error.rs | 109 ++++++++++++++++++++++ milli/src/vector/mod.rs | 1 + milli/src/vector/rest.rs | 185 ++++++++++++++++++++++++++++++++++++++ 3 files changed, 295 insertions(+) create mode 100644 milli/src/vector/rest.rs diff --git a/milli/src/vector/error.rs b/milli/src/vector/error.rs index 1def4f7a9..b2eb37e81 100644 --- a/milli/src/vector/error.rs +++ b/milli/src/vector/error.rs @@ -83,6 +83,32 @@ pub enum EmbedErrorKind { OllamaModelNotFoundError(OllamaError), #[error("received unhandled HTTP status code {0} from Ollama")] OllamaUnhandledStatusCode(u16), + #[error("error serializing template context: {0}")] + RestTemplateContextSerialization(liquid::Error), + #[error( + "error rendering request template: {0}. Hint: available variable in the context: {{{{input}}}}'" + )] + RestTemplateError(liquid::Error), + #[error("error deserialization the response body as JSON: {0}")] + RestResponseDeserialization(std::io::Error), + #[error("component `{0}` not found in path `{1}` in response: `{2}`")] + RestResponseMissingEmbeddings(String, String, String), + #[error("expected a response parseable as a vector or an array of vectors: {0}")] + RestResponseFormat(serde_json::Error), + #[error("expected a response containing {0} embeddings, got only {1}")] + RestResponseEmbeddingCount(usize, usize), + #[error("could not authenticate against embedding server: {0:?}")] + RestUnauthorized(Option), + #[error("sent too many requests to embedding server: {0:?}")] + RestTooManyRequests(Option), + #[error("sent a bad request to embedding server: {0:?}")] + RestBadRequest(Option), + #[error("received internal error from embedding server: {0:?}")] + RestInternalServerError(u16, Option), + #[error("received HTTP {0} from embedding server: {0:?}")] + RestOtherStatusCode(u16, Option), + #[error("could not reach embedding server: {0}")] + RestNetwork(ureq::Transport), } impl EmbedError { @@ -161,6 +187,89 @@ impl EmbedError { pub(crate) fn ollama_unhandled_status_code(code: u16) -> EmbedError { Self { kind: EmbedErrorKind::OllamaUnhandledStatusCode(code), fault: FaultSource::Bug } } + + pub(crate) fn rest_template_context_serialization(error: liquid::Error) -> EmbedError { + Self { + kind: EmbedErrorKind::RestTemplateContextSerialization(error), + fault: FaultSource::Bug, + } + } + + pub(crate) fn rest_template_render(error: liquid::Error) -> EmbedError { + Self { kind: EmbedErrorKind::RestTemplateError(error), fault: FaultSource::User } + } + + pub(crate) fn rest_response_deserialization(error: std::io::Error) -> EmbedError { + Self { + kind: EmbedErrorKind::RestResponseDeserialization(error), + fault: FaultSource::Runtime, + } + } + + pub(crate) fn rest_response_missing_embeddings>( + response: serde_json::Value, + component: &str, + response_field: &[S], + ) -> EmbedError { + let response_field: Vec<&str> = response_field.iter().map(AsRef::as_ref).collect(); + let response_field = response_field.join("."); + + Self { + kind: EmbedErrorKind::RestResponseMissingEmbeddings( + component.to_owned(), + response_field, + serde_json::to_string_pretty(&response).unwrap_or_default(), + ), + fault: FaultSource::Undecided, + } + } + + pub(crate) fn rest_response_format(error: serde_json::Error) -> EmbedError { + Self { kind: EmbedErrorKind::RestResponseFormat(error), fault: FaultSource::Undecided } + } + + pub(crate) fn rest_response_embedding_count(expected: usize, got: usize) -> EmbedError { + Self { + kind: EmbedErrorKind::RestResponseEmbeddingCount(expected, got), + fault: FaultSource::Runtime, + } + } + + pub(crate) fn rest_unauthorized(error_response: Option) -> EmbedError { + Self { kind: EmbedErrorKind::RestUnauthorized(error_response), fault: FaultSource::User } + } + + pub(crate) fn rest_too_many_requests(error_response: Option) -> EmbedError { + Self { + kind: EmbedErrorKind::RestTooManyRequests(error_response), + fault: FaultSource::Runtime, + } + } + + pub(crate) fn rest_bad_request(error_response: Option) -> EmbedError { + Self { kind: EmbedErrorKind::RestBadRequest(error_response), fault: FaultSource::User } + } + + pub(crate) fn rest_internal_server_error( + code: u16, + error_response: Option, + ) -> EmbedError { + Self { + kind: EmbedErrorKind::RestInternalServerError(code, error_response), + fault: FaultSource::Runtime, + } + } + + pub(crate) fn rest_other_status_code(code: u16, error_response: Option) -> EmbedError { + Self { + kind: EmbedErrorKind::RestOtherStatusCode(code, error_response), + fault: FaultSource::Undecided, + } + } + + pub(crate) fn rest_network(transport: ureq::Transport) -> EmbedError { + Self { kind: EmbedErrorKind::RestNetwork(transport), fault: FaultSource::Runtime } + } } #[derive(Debug, thiserror::Error)] diff --git a/milli/src/vector/mod.rs b/milli/src/vector/mod.rs index 86dde8ad4..7eef3d442 100644 --- a/milli/src/vector/mod.rs +++ b/milli/src/vector/mod.rs @@ -11,6 +11,7 @@ pub mod openai; pub mod settings; pub mod ollama; +pub mod rest; pub use self::error::Error; diff --git a/milli/src/vector/rest.rs b/milli/src/vector/rest.rs new file mode 100644 index 000000000..975bd3790 --- /dev/null +++ b/milli/src/vector/rest.rs @@ -0,0 +1,185 @@ +use rayon::iter::{IntoParallelIterator as _, ParallelIterator as _}; + +use super::openai::Retry; +use super::{DistributionShift, EmbedError, Embeddings, NewEmbedderError}; +use crate::VectorOrArrayOfVectors; + +pub struct Embedder { + client: ureq::Agent, + options: EmbedderOptions, + bearer: Option, + dimensions: usize, +} + +pub struct EmbedderOptions { + api_key: Option, + distribution: Option, + dimensions: Option, + url: String, + query: liquid::Template, + response_field: Vec, +} + +impl Embedder { + pub fn new(options: EmbedderOptions) -> Result { + let bearer = options.api_key.as_deref().map(|api_key| format!("Bearer: {api_key}")); + + let client = ureq::agent(); + + let dimensions = if let Some(dimensions) = options.dimensions { + dimensions + } else { + infer_dimensions(&client, &options, bearer.as_deref())? + }; + + Ok(Self { client, dimensions, options, bearer }) + } + + pub fn embed(&self, texts: Vec) -> Result>, EmbedError> { + embed(&self.client, &self.options, self.bearer.as_deref(), texts.as_slice()) + } + + pub fn embed_chunks( + &self, + text_chunks: Vec>, + threads: &rayon::ThreadPool, + ) -> Result>>, EmbedError> { + threads + .install(move || text_chunks.into_par_iter().map(|chunk| self.embed(chunk))) + .collect() + } + + pub fn chunk_count_hint(&self) -> usize { + 10 + } + + pub fn prompt_count_in_chunk_hint(&self) -> usize { + 10 + } + + pub fn dimensions(&self) -> usize { + self.dimensions + } + + pub fn distribution(&self) -> Option { + self.options.distribution + } +} + +fn infer_dimensions( + client: &ureq::Agent, + options: &EmbedderOptions, + bearer: Option<&str>, +) -> Result { + let v = embed(client, options, bearer, ["test"].as_slice()) + .map_err(NewEmbedderError::could_not_determine_dimension)?; + // unwrap: guaranteed that v.len() == ["test"].len() == 1, otherwise the previous line terminated in error + Ok(v.first().unwrap().dimension()) +} + +fn embed( + client: &ureq::Agent, + options: &EmbedderOptions, + bearer: Option<&str>, + inputs: &[S], +) -> Result>, EmbedError> +where + S: serde::Serialize, +{ + let request = client.post(&options.url); + let request = + if let Some(bearer) = bearer { request.set("Authorization", bearer) } else { request }; + let request = request.set("Content-Type", "application/json"); + + let body = options + .query + .render( + &liquid::to_object(&serde_json::json!({ + "input": inputs, + })) + .map_err(EmbedError::rest_template_context_serialization)?, + ) + .map_err(EmbedError::rest_template_render)?; + + for attempt in 0..7 { + let response = request.send_string(&body); + let result = check_response(response); + + let retry_duration = match result { + Ok(response) => { + return response_to_embedding(response, &options.response_field, inputs.len()) + } + Err(retry) => { + tracing::warn!("Failed: {}", retry.error); + retry.into_duration(attempt) + } + }?; + + let retry_duration = retry_duration.min(std::time::Duration::from_secs(60)); // don't wait more than a minute + tracing::warn!("Attempt #{}, retrying after {}ms.", attempt, retry_duration.as_millis()); + std::thread::sleep(retry_duration); + } + + let response = request.send_string(&body); + let result = check_response(response); + result + .map_err(Retry::into_error) + .and_then(|response| response_to_embedding(response, &options.response_field, inputs.len())) +} + +fn check_response(response: Result) -> Result { + match response { + Ok(response) => Ok(response), + Err(ureq::Error::Status(code, response)) => { + let error_response: Option = response.into_string().ok(); + Err(match code { + 401 => Retry::give_up(EmbedError::rest_unauthorized(error_response)), + 429 => Retry::rate_limited(EmbedError::rest_too_many_requests(error_response)), + 400 => Retry::give_up(EmbedError::rest_bad_request(error_response)), + 500..=599 => { + Retry::retry_later(EmbedError::rest_internal_server_error(code, error_response)) + } + x => Retry::retry_later(EmbedError::rest_other_status_code(code, error_response)), + }) + } + Err(ureq::Error::Transport(transport)) => { + Err(Retry::retry_later(EmbedError::rest_network(transport))) + } + } +} + +fn response_to_embedding>( + response: ureq::Response, + response_field: &[S], + expected_count: usize, +) -> Result>, EmbedError> { + let response: serde_json::Value = + response.into_json().map_err(EmbedError::rest_response_deserialization)?; + + let mut current_value = &response; + for component in response_field { + let component = component.as_ref(); + let current_value = current_value.get(component).ok_or_else(|| { + EmbedError::rest_response_missing_embeddings(response, component, response_field) + })?; + } + + let embeddings = current_value.to_owned(); + + let embeddings: VectorOrArrayOfVectors = + serde_json::from_value(embeddings).map_err(EmbedError::rest_response_format)?; + + let embeddings = embeddings.into_array_of_vectors(); + + let embeddings: Vec> = embeddings + .into_iter() + .flatten() + .map(|embedding| Embeddings::from_single_embedding(embedding)) + .collect(); + + if embeddings.len() != expected_count { + return Err(EmbedError::rest_response_embedding_count(expected_count, embeddings.len())); + } + + Ok(embeddings) +} From ac52c857e8f5ecf85a42e32abe7a14450fdfdd66 Mon Sep 17 00:00:00 2001 From: Louis Dureuil Date: Tue, 19 Mar 2024 15:41:37 +0100 Subject: [PATCH 30/86] Update ollama and openai impls to use the rest embedder internally --- .../extract/extract_vector_points.rs | 10 +- .../src/update/index_documents/extract/mod.rs | 15 +- milli/src/vector/error.rs | 116 +---- milli/src/vector/mod.rs | 22 +- milli/src/vector/ollama.rs | 307 ++---------- milli/src/vector/openai.rs | 452 +++++------------- milli/src/vector/rest.rs | 247 ++++++++-- milli/src/vector/settings.rs | 4 +- 8 files changed, 394 insertions(+), 779 deletions(-) diff --git a/milli/src/update/index_documents/extract/extract_vector_points.rs b/milli/src/update/index_documents/extract/extract_vector_points.rs index ece841659..40b32bf9c 100644 --- a/milli/src/update/index_documents/extract/extract_vector_points.rs +++ b/milli/src/update/index_documents/extract/extract_vector_points.rs @@ -339,6 +339,7 @@ pub fn extract_embeddings( prompt_reader: grenad::Reader, indexer: GrenadParameters, embedder: Arc, + request_threads: &rayon::ThreadPool, ) -> Result>> { puffin::profile_function!(); let n_chunks = embedder.chunk_count_hint(); // chunk level parallelism @@ -376,7 +377,10 @@ pub fn extract_embeddings( if chunks.len() == chunks.capacity() { let chunked_embeds = embedder - .embed_chunks(std::mem::replace(&mut chunks, Vec::with_capacity(n_chunks))) + .embed_chunks( + std::mem::replace(&mut chunks, Vec::with_capacity(n_chunks)), + request_threads, + ) .map_err(crate::vector::Error::from) .map_err(crate::Error::from)?; @@ -394,7 +398,7 @@ pub fn extract_embeddings( // send last chunk if !chunks.is_empty() { let chunked_embeds = embedder - .embed_chunks(std::mem::take(&mut chunks)) + .embed_chunks(std::mem::take(&mut chunks), request_threads) .map_err(crate::vector::Error::from) .map_err(crate::Error::from)?; for (docid, embeddings) in chunks_ids @@ -408,7 +412,7 @@ pub fn extract_embeddings( if !current_chunk.is_empty() { let embeds = embedder - .embed_chunks(vec![std::mem::take(&mut current_chunk)]) + .embed_chunks(vec![std::mem::take(&mut current_chunk)], request_threads) .map_err(crate::vector::Error::from) .map_err(crate::Error::from)?; diff --git a/milli/src/update/index_documents/extract/mod.rs b/milli/src/update/index_documents/extract/mod.rs index 43f3f4947..5689bb04f 100644 --- a/milli/src/update/index_documents/extract/mod.rs +++ b/milli/src/update/index_documents/extract/mod.rs @@ -238,7 +238,15 @@ fn send_original_documents_data( let documents_chunk_cloned = original_documents_chunk.clone(); let lmdb_writer_sx_cloned = lmdb_writer_sx.clone(); + + let request_threads = rayon::ThreadPoolBuilder::new() + .num_threads(crate::vector::REQUEST_PARALLELISM) + .thread_name(|index| format!("embedding-request-{index}")) + .build() + .unwrap(); + rayon::spawn(move || { + /// FIXME: unwrap for (name, (embedder, prompt)) in embedders { let result = extract_vector_points( documents_chunk_cloned.clone(), @@ -249,7 +257,12 @@ fn send_original_documents_data( ); match result { Ok(ExtractedVectorPoints { manual_vectors, remove_vectors, prompts }) => { - let embeddings = match extract_embeddings(prompts, indexer, embedder.clone()) { + let embeddings = match extract_embeddings( + prompts, + indexer, + embedder.clone(), + &request_threads, + ) { Ok(results) => Some(results), Err(error) => { let _ = lmdb_writer_sx_cloned.send(Err(error)); diff --git a/milli/src/vector/error.rs b/milli/src/vector/error.rs index b2eb37e81..92f077924 100644 --- a/milli/src/vector/error.rs +++ b/milli/src/vector/error.rs @@ -2,9 +2,7 @@ use std::path::PathBuf; use hf_hub::api::sync::ApiError; -use super::ollama::OllamaError; use crate::error::FaultSource; -use crate::vector::openai::OpenAiError; #[derive(Debug, thiserror::Error)] #[error("Error while generating embeddings: {inner}")] @@ -52,43 +50,12 @@ pub enum EmbedErrorKind { TensorValue(candle_core::Error), #[error("could not run model: {0}")] ModelForward(candle_core::Error), - #[error("could not reach OpenAI: {0}")] - OpenAiNetwork(ureq::Transport), - #[error("unexpected response from OpenAI: {0}")] - OpenAiUnexpected(ureq::Error), - #[error("could not authenticate against OpenAI: {0:?}")] - OpenAiAuth(Option), - #[error("sent too many requests to OpenAI: {0:?}")] - OpenAiTooManyRequests(Option), - #[error("received internal error from OpenAI: {0:?}")] - OpenAiInternalServerError(Option), - #[error("sent too many tokens in a request to OpenAI: {0:?}")] - OpenAiTooManyTokens(Option), - #[error("received unhandled HTTP status code {0} from OpenAI")] - OpenAiUnhandledStatusCode(u16), #[error("attempt to embed the following text in a configuration where embeddings must be user provided: {0:?}")] ManualEmbed(String), #[error("could not initialize asynchronous runtime: {0}")] OpenAiRuntimeInit(std::io::Error), - #[error("initializing web client for sending embedding requests failed: {0}")] - InitWebClient(reqwest::Error), - // Dedicated Ollama error kinds, might have to merge them into one cohesive error type for all backends. - #[error("unexpected response from Ollama: {0}")] - OllamaUnexpected(reqwest::Error), - #[error("sent too many requests to Ollama: {0}")] - OllamaTooManyRequests(OllamaError), - #[error("received internal error from Ollama: {0}")] - OllamaInternalServerError(OllamaError), - #[error("model not found. Meilisearch will not automatically download models from the Ollama library, please pull the model manually: {0}")] - OllamaModelNotFoundError(OllamaError), - #[error("received unhandled HTTP status code {0} from Ollama")] - OllamaUnhandledStatusCode(u16), - #[error("error serializing template context: {0}")] - RestTemplateContextSerialization(liquid::Error), - #[error( - "error rendering request template: {0}. Hint: available variable in the context: {{{{input}}}}'" - )] - RestTemplateError(liquid::Error), + #[error("model not found. Meilisearch will not automatically download models from the Ollama library, please pull the model manually: {0:?}")] + OllamaModelNotFoundError(Option), #[error("error deserialization the response body as JSON: {0}")] RestResponseDeserialization(std::io::Error), #[error("component `{0}` not found in path `{1}` in response: `{2}`")] @@ -128,77 +95,14 @@ impl EmbedError { Self { kind: EmbedErrorKind::ModelForward(inner), fault: FaultSource::Runtime } } - pub fn openai_network(inner: ureq::Transport) -> Self { - Self { kind: EmbedErrorKind::OpenAiNetwork(inner), fault: FaultSource::Runtime } - } - - pub fn openai_unexpected(inner: ureq::Error) -> EmbedError { - Self { kind: EmbedErrorKind::OpenAiUnexpected(inner), fault: FaultSource::Bug } - } - - pub(crate) fn openai_auth_error(inner: Option) -> EmbedError { - Self { kind: EmbedErrorKind::OpenAiAuth(inner), fault: FaultSource::User } - } - - pub(crate) fn openai_too_many_requests(inner: Option) -> EmbedError { - Self { kind: EmbedErrorKind::OpenAiTooManyRequests(inner), fault: FaultSource::Runtime } - } - - pub(crate) fn openai_internal_server_error(inner: Option) -> EmbedError { - Self { kind: EmbedErrorKind::OpenAiInternalServerError(inner), fault: FaultSource::Runtime } - } - - pub(crate) fn openai_too_many_tokens(inner: Option) -> EmbedError { - Self { kind: EmbedErrorKind::OpenAiTooManyTokens(inner), fault: FaultSource::Bug } - } - - pub(crate) fn openai_unhandled_status_code(code: u16) -> EmbedError { - Self { kind: EmbedErrorKind::OpenAiUnhandledStatusCode(code), fault: FaultSource::Bug } - } - pub(crate) fn embed_on_manual_embedder(texts: String) -> EmbedError { Self { kind: EmbedErrorKind::ManualEmbed(texts), fault: FaultSource::User } } - pub(crate) fn openai_runtime_init(inner: std::io::Error) -> EmbedError { - Self { kind: EmbedErrorKind::OpenAiRuntimeInit(inner), fault: FaultSource::Runtime } - } - - pub fn openai_initialize_web_client(inner: reqwest::Error) -> Self { - Self { kind: EmbedErrorKind::InitWebClient(inner), fault: FaultSource::Runtime } - } - - pub(crate) fn ollama_unexpected(inner: reqwest::Error) -> EmbedError { - Self { kind: EmbedErrorKind::OllamaUnexpected(inner), fault: FaultSource::Bug } - } - - pub(crate) fn ollama_model_not_found(inner: OllamaError) -> EmbedError { + pub(crate) fn ollama_model_not_found(inner: Option) -> EmbedError { Self { kind: EmbedErrorKind::OllamaModelNotFoundError(inner), fault: FaultSource::User } } - pub(crate) fn ollama_too_many_requests(inner: OllamaError) -> EmbedError { - Self { kind: EmbedErrorKind::OllamaTooManyRequests(inner), fault: FaultSource::Runtime } - } - - pub(crate) fn ollama_internal_server_error(inner: OllamaError) -> EmbedError { - Self { kind: EmbedErrorKind::OllamaInternalServerError(inner), fault: FaultSource::Runtime } - } - - pub(crate) fn ollama_unhandled_status_code(code: u16) -> EmbedError { - Self { kind: EmbedErrorKind::OllamaUnhandledStatusCode(code), fault: FaultSource::Bug } - } - - pub(crate) fn rest_template_context_serialization(error: liquid::Error) -> EmbedError { - Self { - kind: EmbedErrorKind::RestTemplateContextSerialization(error), - fault: FaultSource::Bug, - } - } - - pub(crate) fn rest_template_render(error: liquid::Error) -> EmbedError { - Self { kind: EmbedErrorKind::RestTemplateError(error), fault: FaultSource::User } - } - pub(crate) fn rest_response_deserialization(error: std::io::Error) -> EmbedError { Self { kind: EmbedErrorKind::RestResponseDeserialization(error), @@ -335,17 +239,6 @@ impl NewEmbedderError { fault: FaultSource::Runtime, } } - - pub fn ollama_could_not_determine_dimension(inner: EmbedError) -> NewEmbedderError { - Self { - kind: NewEmbedderErrorKind::CouldNotDetermineDimension(inner), - fault: FaultSource::User, - } - } - - pub fn openai_invalid_api_key_format(inner: reqwest::header::InvalidHeaderValue) -> Self { - Self { kind: NewEmbedderErrorKind::InvalidApiKeyFormat(inner), fault: FaultSource::User } - } } #[derive(Debug, thiserror::Error)] @@ -392,7 +285,4 @@ pub enum NewEmbedderErrorKind { CouldNotDetermineDimension(EmbedError), #[error("loading model failed: {0}")] LoadModel(candle_core::Error), - // openai - #[error("The API key passed to Authorization error was in an invalid format: {0}")] - InvalidApiKeyFormat(reqwest::header::InvalidHeaderValue), } diff --git a/milli/src/vector/mod.rs b/milli/src/vector/mod.rs index 7eef3d442..39232e387 100644 --- a/milli/src/vector/mod.rs +++ b/milli/src/vector/mod.rs @@ -17,6 +17,8 @@ pub use self::error::Error; pub type Embedding = Vec; +pub const REQUEST_PARALLELISM: usize = 40; + /// One or multiple embeddings stored consecutively in a flat vector. pub struct Embeddings { data: Vec, @@ -99,7 +101,7 @@ pub enum Embedder { /// An embedder based on running local models, fetched from the Hugging Face Hub. HuggingFace(hf::Embedder), /// An embedder based on making embedding queries against the OpenAI API. - OpenAi(openai::sync::Embedder), + OpenAi(openai::Embedder), /// An embedder based on the user providing the embeddings in the documents and queries. UserProvided(manual::Embedder), Ollama(ollama::Embedder), @@ -202,7 +204,7 @@ impl Embedder { pub fn new(options: EmbedderOptions) -> std::result::Result { Ok(match options { EmbedderOptions::HuggingFace(options) => Self::HuggingFace(hf::Embedder::new(options)?), - EmbedderOptions::OpenAi(options) => Self::OpenAi(openai::sync::Embedder::new(options)?), + EmbedderOptions::OpenAi(options) => Self::OpenAi(openai::Embedder::new(options)?), EmbedderOptions::Ollama(options) => Self::Ollama(ollama::Embedder::new(options)?), EmbedderOptions::UserProvided(options) => { Self::UserProvided(manual::Embedder::new(options)) @@ -213,17 +215,14 @@ impl Embedder { /// Embed one or multiple texts. /// /// Each text can be embedded as one or multiple embeddings. - pub async fn embed( + pub fn embed( &self, texts: Vec, ) -> std::result::Result>, EmbedError> { match self { Embedder::HuggingFace(embedder) => embedder.embed(texts), Embedder::OpenAi(embedder) => embedder.embed(texts), - Embedder::Ollama(embedder) => { - let client = embedder.new_client()?; - embedder.embed(texts, &client).await - } + Embedder::Ollama(embedder) => embedder.embed(texts), Embedder::UserProvided(embedder) => embedder.embed(texts), } } @@ -231,18 +230,15 @@ impl Embedder { /// Embed multiple chunks of texts. /// /// Each chunk is composed of one or multiple texts. - /// - /// # Panics - /// - /// - if called from an asynchronous context pub fn embed_chunks( &self, text_chunks: Vec>, + threads: &rayon::ThreadPool, ) -> std::result::Result>>, EmbedError> { match self { Embedder::HuggingFace(embedder) => embedder.embed_chunks(text_chunks), - Embedder::OpenAi(embedder) => embedder.embed_chunks(text_chunks), - Embedder::Ollama(embedder) => embedder.embed_chunks(text_chunks), + Embedder::OpenAi(embedder) => embedder.embed_chunks(text_chunks, threads), + Embedder::Ollama(embedder) => embedder.embed_chunks(text_chunks, threads), Embedder::UserProvided(embedder) => embedder.embed_chunks(text_chunks), } } diff --git a/milli/src/vector/ollama.rs b/milli/src/vector/ollama.rs index 76988f70b..9c44e8052 100644 --- a/milli/src/vector/ollama.rs +++ b/milli/src/vector/ollama.rs @@ -1,293 +1,94 @@ -// Copied from "openai.rs" with the sections I actually understand changed for Ollama. -// The common components of the Ollama and OpenAI interfaces might need to be extracted. +use rayon::iter::{IntoParallelIterator as _, ParallelIterator as _}; -use std::fmt::Display; - -use reqwest::StatusCode; - -use super::error::{EmbedError, NewEmbedderError}; -use super::openai::Retry; -use super::{DistributionShift, Embedding, Embeddings}; +use super::error::{EmbedError, EmbedErrorKind, NewEmbedderError, NewEmbedderErrorKind}; +use super::rest::{Embedder as RestEmbedder, EmbedderOptions as RestEmbedderOptions}; +use super::{DistributionShift, Embeddings}; #[derive(Debug)] pub struct Embedder { - headers: reqwest::header::HeaderMap, - options: EmbedderOptions, + rest_embedder: RestEmbedder, } #[derive(Debug, Clone, Hash, PartialEq, Eq, serde::Deserialize, serde::Serialize)] pub struct EmbedderOptions { - pub embedding_model: EmbeddingModel, -} - -#[derive( - Debug, Clone, Hash, PartialEq, Eq, serde::Serialize, serde::Deserialize, deserr::Deserr, -)] -#[deserr(deny_unknown_fields)] -pub struct EmbeddingModel { - name: String, - dimensions: usize, -} - -#[derive(Debug, serde::Serialize)] -struct OllamaRequest<'a> { - model: &'a str, - prompt: &'a str, -} - -#[derive(Debug, serde::Deserialize)] -struct OllamaResponse { - embedding: Embedding, -} - -#[derive(Debug, serde::Deserialize)] -pub struct OllamaError { - error: String, -} - -impl EmbeddingModel { - pub fn max_token(&self) -> usize { - // this might not be the same for all models - 8192 - } - - pub fn default_dimensions(&self) -> usize { - // Dimensions for nomic-embed-text - 768 - } - - pub fn name(&self) -> String { - self.name.clone() - } - - pub fn from_name(name: &str) -> Self { - Self { name: name.to_string(), dimensions: 0 } - } - - pub fn supports_overriding_dimensions(&self) -> bool { - false - } -} - -impl Default for EmbeddingModel { - fn default() -> Self { - Self { name: "nomic-embed-text".to_string(), dimensions: 0 } - } + pub embedding_model: String, } impl EmbedderOptions { pub fn with_default_model() -> Self { - Self { embedding_model: Default::default() } + Self { embedding_model: "nomic-embed-text".into() } } - pub fn with_embedding_model(embedding_model: EmbeddingModel) -> Self { + pub fn with_embedding_model(embedding_model: String) -> Self { Self { embedding_model } } } impl Embedder { - pub fn new_client(&self) -> Result { - reqwest::ClientBuilder::new() - .default_headers(self.headers.clone()) - .build() - .map_err(EmbedError::openai_initialize_web_client) - } - pub fn new(options: EmbedderOptions) -> Result { - let mut headers = reqwest::header::HeaderMap::new(); - headers.insert( - reqwest::header::CONTENT_TYPE, - reqwest::header::HeaderValue::from_static("application/json"), - ); - - let mut embedder = Self { options, headers }; - - let rt = tokio::runtime::Builder::new_current_thread() - .enable_io() - .enable_time() - .build() - .map_err(EmbedError::openai_runtime_init) - .map_err(NewEmbedderError::ollama_could_not_determine_dimension)?; - - // Get dimensions from Ollama - let request = - OllamaRequest { model: &embedder.options.embedding_model.name(), prompt: "test" }; - // TODO: Refactor into shared error type - let client = embedder - .new_client() - .map_err(NewEmbedderError::ollama_could_not_determine_dimension)?; - - rt.block_on(async move { - let response = client - .post(get_ollama_path()) - .json(&request) - .send() - .await - .map_err(EmbedError::ollama_unexpected) - .map_err(NewEmbedderError::ollama_could_not_determine_dimension)?; - - // Process error in case model not found - let response = Self::check_response(response).await.map_err(|_err| { - let e = EmbedError::ollama_model_not_found(OllamaError { - error: format!("model: {}", embedder.options.embedding_model.name()), - }); - NewEmbedderError::ollama_could_not_determine_dimension(e) - })?; - - let response: OllamaResponse = response - .json() - .await - .map_err(EmbedError::ollama_unexpected) - .map_err(NewEmbedderError::ollama_could_not_determine_dimension)?; - - let embedding = Embeddings::from_single_embedding(response.embedding); - - embedder.options.embedding_model.dimensions = embedding.dimension(); - - tracing::info!( - "ollama model {} with dimensionality {} added", - embedder.options.embedding_model.name(), - embedding.dimension() - ); - - Ok(embedder) - }) - } - - async fn check_response(response: reqwest::Response) -> Result { - if !response.status().is_success() { - // Not the same number of possible error cases covered as with OpenAI. - match response.status() { - StatusCode::TOO_MANY_REQUESTS => { - let error_response: OllamaError = response - .json() - .await - .map_err(EmbedError::ollama_unexpected) - .map_err(Retry::retry_later)?; - - return Err(Retry::rate_limited(EmbedError::ollama_too_many_requests( - OllamaError { error: error_response.error }, - ))); - } - StatusCode::SERVICE_UNAVAILABLE => { - let error_response: OllamaError = response - .json() - .await - .map_err(EmbedError::ollama_unexpected) - .map_err(Retry::retry_later)?; - return Err(Retry::retry_later(EmbedError::ollama_internal_server_error( - OllamaError { error: error_response.error }, - ))); - } - StatusCode::NOT_FOUND => { - let error_response: OllamaError = response - .json() - .await - .map_err(EmbedError::ollama_unexpected) - .map_err(Retry::give_up)?; - - return Err(Retry::give_up(EmbedError::ollama_model_not_found(OllamaError { - error: error_response.error, - }))); - } - code => { - return Err(Retry::give_up(EmbedError::ollama_unhandled_status_code( - code.as_u16(), - ))); - } + let model = options.embedding_model.as_str(); + let rest_embedder = match RestEmbedder::new(RestEmbedderOptions { + api_key: None, + distribution: None, + dimensions: None, + url: get_ollama_path(), + query: serde_json::json!({ + "model": model, + }), + input_field: vec!["prompt".to_owned()], + path_to_embeddings: Default::default(), + embedding_object: vec!["embedding".to_owned()], + input_type: super::rest::InputType::Text, + }) { + Ok(embedder) => embedder, + Err(NewEmbedderError { + kind: + NewEmbedderErrorKind::CouldNotDetermineDimension(EmbedError { + kind: super::error::EmbedErrorKind::RestOtherStatusCode(404, error), + fault: _, + }), + fault: _, + }) => { + return Err(NewEmbedderError::could_not_determine_dimension( + EmbedError::ollama_model_not_found(error), + )) } - } - Ok(response) + Err(error) => return Err(error), + }; + + Ok(Self { rest_embedder }) } - pub async fn embed( - &self, - texts: Vec, - client: &reqwest::Client, - ) -> Result>, EmbedError> { - // Ollama only embedds one document at a time. - let mut results = Vec::with_capacity(texts.len()); - - // The retry loop is inside the texts loop, might have to switch that around - for text in texts { - // Retries copied from openai.rs - for attempt in 0..7 { - let retry_duration = match self.try_embed(&text, client).await { - Ok(result) => { - results.push(result); - break; - } - Err(retry) => { - tracing::warn!("Failed: {}", retry.error); - retry.into_duration(attempt) - } - }?; - tracing::warn!( - "Attempt #{}, retrying after {}ms.", - attempt, - retry_duration.as_millis() - ); - tokio::time::sleep(retry_duration).await; + pub fn embed(&self, texts: Vec) -> Result>, EmbedError> { + match self.rest_embedder.embed(texts) { + Ok(embeddings) => Ok(embeddings), + Err(EmbedError { kind: EmbedErrorKind::RestOtherStatusCode(404, error), fault: _ }) => { + Err(EmbedError::ollama_model_not_found(error)) } + Err(error) => Err(error), } - - Ok(results) - } - - async fn try_embed( - &self, - text: &str, - client: &reqwest::Client, - ) -> Result, Retry> { - let request = OllamaRequest { model: &self.options.embedding_model.name(), prompt: text }; - let response = client - .post(get_ollama_path()) - .json(&request) - .send() - .await - .map_err(EmbedError::openai_network) - .map_err(Retry::retry_later)?; - - let response = Self::check_response(response).await?; - - let response: OllamaResponse = response - .json() - .await - .map_err(EmbedError::openai_unexpected) - .map_err(Retry::retry_later)?; - - tracing::trace!("response: {:?}", response.embedding); - - let embedding = Embeddings::from_single_embedding(response.embedding); - Ok(embedding) } pub fn embed_chunks( &self, text_chunks: Vec>, + threads: &rayon::ThreadPool, ) -> Result>>, EmbedError> { - let rt = tokio::runtime::Builder::new_current_thread() - .enable_io() - .enable_time() - .build() - .map_err(EmbedError::openai_runtime_init)?; - let client = self.new_client()?; - rt.block_on(futures::future::try_join_all( - text_chunks.into_iter().map(|prompts| self.embed(prompts, &client)), - )) + threads.install(move || { + text_chunks.into_par_iter().map(move |chunk| self.embed(chunk)).collect() + }) } - // Defaults copied from openai.rs pub fn chunk_count_hint(&self) -> usize { - 10 + self.rest_embedder.chunk_count_hint() } pub fn prompt_count_in_chunk_hint(&self) -> usize { - 10 + self.rest_embedder.prompt_count_in_chunk_hint() } pub fn dimensions(&self) -> usize { - self.options.embedding_model.dimensions + self.rest_embedder.dimensions() } pub fn distribution(&self) -> Option { @@ -295,12 +96,6 @@ impl Embedder { } } -impl Display for OllamaError { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!(f, "{}", self.error) - } -} - fn get_ollama_path() -> String { // Important: Hostname not enough, has to be entire path to embeddings endpoint std::env::var("MEILI_OLLAMA_URL").unwrap_or("http://localhost:11434/api/embeddings".to_string()) diff --git a/milli/src/vector/openai.rs b/milli/src/vector/openai.rs index 5d13d5ee2..b2638966e 100644 --- a/milli/src/vector/openai.rs +++ b/milli/src/vector/openai.rs @@ -1,9 +1,9 @@ -use std::fmt::Display; - -use serde::{Deserialize, Serialize}; +use rayon::iter::{IntoParallelIterator, ParallelIterator as _}; use super::error::{EmbedError, NewEmbedderError}; -use super::{DistributionShift, Embedding, Embeddings}; +use super::rest::{Embedder as RestEmbedder, EmbedderOptions as RestEmbedderOptions}; +use super::{DistributionShift, Embeddings}; +use crate::vector::error::EmbedErrorKind; #[derive(Debug, Clone, Hash, PartialEq, Eq, serde::Deserialize, serde::Serialize)] pub struct EmbedderOptions { @@ -12,6 +12,32 @@ pub struct EmbedderOptions { pub dimensions: Option, } +impl EmbedderOptions { + pub fn dimensions(&self) -> usize { + if self.embedding_model.supports_overriding_dimensions() { + self.dimensions.unwrap_or(self.embedding_model.default_dimensions()) + } else { + self.embedding_model.default_dimensions() + } + } + + pub fn query(&self) -> serde_json::Value { + let model = self.embedding_model.name(); + + let mut query = serde_json::json!({ + "model": model, + }); + + if self.embedding_model.supports_overriding_dimensions() { + if let Some(dimensions) = self.dimensions { + query["dimensions"] = dimensions.into(); + } + } + + query + } +} + #[derive( Debug, Clone, @@ -117,364 +143,112 @@ impl EmbedderOptions { } } -// retrying in case of failure - -pub struct Retry { - pub error: EmbedError, - strategy: RetryStrategy, -} - -pub enum RetryStrategy { - GiveUp, - Retry, - RetryTokenized, - RetryAfterRateLimit, -} - -impl Retry { - pub fn give_up(error: EmbedError) -> Self { - Self { error, strategy: RetryStrategy::GiveUp } - } - - pub fn retry_later(error: EmbedError) -> Self { - Self { error, strategy: RetryStrategy::Retry } - } - - pub fn retry_tokenized(error: EmbedError) -> Self { - Self { error, strategy: RetryStrategy::RetryTokenized } - } - - pub fn rate_limited(error: EmbedError) -> Self { - Self { error, strategy: RetryStrategy::RetryAfterRateLimit } - } - - pub fn into_duration(self, attempt: u32) -> Result { - match self.strategy { - RetryStrategy::GiveUp => Err(self.error), - RetryStrategy::Retry => Ok(tokio::time::Duration::from_millis((10u64).pow(attempt))), - RetryStrategy::RetryTokenized => Ok(tokio::time::Duration::from_millis(1)), - RetryStrategy::RetryAfterRateLimit => { - Ok(tokio::time::Duration::from_millis(100 + 10u64.pow(attempt))) - } - } - } - - pub fn must_tokenize(&self) -> bool { - matches!(self.strategy, RetryStrategy::RetryTokenized) - } - - pub fn into_error(self) -> EmbedError { - self.error - } -} - -// openai api structs - -#[derive(Debug, Serialize)] -struct OpenAiRequest<'a, S: AsRef + serde::Serialize> { - model: &'a str, - input: &'a [S], - #[serde(skip_serializing_if = "Option::is_none")] - dimensions: Option, -} - -#[derive(Debug, Serialize)] -struct OpenAiTokensRequest<'a> { - model: &'a str, - input: &'a [usize], - #[serde(skip_serializing_if = "Option::is_none")] - dimensions: Option, -} - -#[derive(Debug, Deserialize)] -struct OpenAiResponse { - data: Vec, -} - -#[derive(Debug, Deserialize)] -struct OpenAiErrorResponse { - error: OpenAiError, -} - -#[derive(Debug, Deserialize)] -pub struct OpenAiError { - message: String, - // type: String, - code: Option, -} - -impl Display for OpenAiError { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - match &self.code { - Some(code) => write!(f, "{} ({})", self.message, code), - None => write!(f, "{}", self.message), - } - } -} - -#[derive(Debug, Deserialize)] -struct OpenAiEmbedding { - embedding: Embedding, - // object: String, - // index: usize, -} - fn infer_api_key() -> String { std::env::var("MEILI_OPENAI_API_KEY") .or_else(|_| std::env::var("OPENAI_API_KEY")) .unwrap_or_default() } -pub mod sync { - use rayon::iter::{IntoParallelIterator, ParallelIterator as _}; +#[derive(Debug)] +pub struct Embedder { + tokenizer: tiktoken_rs::CoreBPE, + rest_embedder: RestEmbedder, + options: EmbedderOptions, +} - use super::{ - EmbedError, Embedding, Embeddings, NewEmbedderError, OpenAiErrorResponse, OpenAiRequest, - OpenAiResponse, OpenAiTokensRequest, Retry, OPENAI_EMBEDDINGS_URL, - }; - use crate::vector::DistributionShift; +impl Embedder { + pub fn new(options: EmbedderOptions) -> Result { + let mut inferred_api_key = Default::default(); + let api_key = options.api_key.as_ref().unwrap_or_else(|| { + inferred_api_key = infer_api_key(); + &inferred_api_key + }); - const REQUEST_PARALLELISM: usize = 10; + let rest_embedder = RestEmbedder::new(RestEmbedderOptions { + api_key: Some(api_key.clone()), + distribution: options.embedding_model.distribution(), + dimensions: Some(options.dimensions()), + url: OPENAI_EMBEDDINGS_URL.to_owned(), + query: options.query(), + input_field: vec!["input".to_owned()], + input_type: crate::vector::rest::InputType::TextArray, + path_to_embeddings: vec!["data".to_owned()], + embedding_object: vec!["embedding".to_owned()], + })?; - #[derive(Debug)] - pub struct Embedder { - tokenizer: tiktoken_rs::CoreBPE, - options: super::EmbedderOptions, - bearer: String, - threads: rayon::ThreadPool, + // looking at the code it is very unclear that this can actually fail. + let tokenizer = tiktoken_rs::cl100k_base().unwrap(); + + Ok(Self { options, rest_embedder, tokenizer }) } - impl Embedder { - pub fn new(options: super::EmbedderOptions) -> Result { - let mut inferred_api_key = Default::default(); - let api_key = options.api_key.as_ref().unwrap_or_else(|| { - inferred_api_key = super::infer_api_key(); - &inferred_api_key - }); - let bearer = format!("Bearer {api_key}"); - - // looking at the code it is very unclear that this can actually fail. - let tokenizer = tiktoken_rs::cl100k_base().unwrap(); - - // FIXME: unwrap - let threads = rayon::ThreadPoolBuilder::new() - .num_threads(REQUEST_PARALLELISM) - .thread_name(|index| format!("embedder-chunk-{index}")) - .build() - .unwrap(); - - Ok(Self { options, bearer, tokenizer, threads }) + pub fn embed(&self, texts: Vec) -> Result>, EmbedError> { + match self.rest_embedder.embed_ref(&texts) { + Ok(embeddings) => Ok(embeddings), + Err(EmbedError { kind: EmbedErrorKind::RestBadRequest(error), fault: _ }) => { + tracing::warn!(error=?error, "OpenAI: received `BAD_REQUEST`. Input was maybe too long, retrying on tokenized version. For best performance, limit the size of your document template."); + self.try_embed_tokenized(&texts) + } + Err(error) => Err(error), } + } - pub fn embed(&self, texts: Vec) -> Result>, EmbedError> { - let mut tokenized = false; - - let client = ureq::agent(); - - for attempt in 0..7 { - let result = if tokenized { - self.try_embed_tokenized(&texts, &client) - } else { - self.try_embed(&texts, &client) - }; - - let retry_duration = match result { - Ok(embeddings) => return Ok(embeddings), - Err(retry) => { - tracing::warn!("Failed: {}", retry.error); - tokenized |= retry.must_tokenize(); - retry.into_duration(attempt) - } - }?; - - let retry_duration = retry_duration.min(std::time::Duration::from_secs(60)); // don't wait more than a minute - tracing::warn!( - "Attempt #{}, retrying after {}ms.", - attempt, - retry_duration.as_millis() - ); - std::thread::sleep(retry_duration); + fn try_embed_tokenized(&self, text: &[String]) -> Result>, EmbedError> { + pub const OVERLAP_SIZE: usize = 200; + let mut all_embeddings = Vec::with_capacity(text.len()); + for text in text { + let max_token_count = self.options.embedding_model.max_token(); + let encoded = self.tokenizer.encode_ordinary(text.as_str()); + let len = encoded.len(); + if len < max_token_count { + all_embeddings.append(&mut self.rest_embedder.embed_ref(&[text])?); + continue; } - let result = if tokenized { - self.try_embed_tokenized(&texts, &client) - } else { - self.try_embed(&texts, &client) - }; + let mut tokens = encoded.as_slice(); + let mut embeddings_for_prompt = Embeddings::new(self.dimensions()); + while tokens.len() > max_token_count { + let window = &tokens[..max_token_count]; + let embedding = self.rest_embedder.embed_tokens(window)?; + /// FIXME: unwrap + embeddings_for_prompt.append(embedding.into_inner()).unwrap(); - result.map_err(Retry::into_error) - } - - fn check_response( - response: Result, - ) -> Result { - match response { - Ok(response) => Ok(response), - Err(ureq::Error::Status(code, response)) => { - let error_response: Option = response.into_json().ok(); - let error = error_response.map(|response| response.error); - Err(match code { - 401 => Retry::give_up(EmbedError::openai_auth_error(error)), - 429 => Retry::rate_limited(EmbedError::openai_too_many_requests(error)), - 400 => { - tracing::warn!("OpenAI: received `BAD_REQUEST`. Input was maybe too long, retrying on tokenized version. For best performance, limit the size of your document template."); - - Retry::retry_tokenized(EmbedError::openai_too_many_tokens(error)) - } - 500..=599 => { - Retry::retry_later(EmbedError::openai_internal_server_error(error)) - } - x => Retry::retry_later(EmbedError::openai_unhandled_status_code(code)), - }) - } - Err(ureq::Error::Transport(transport)) => { - Err(Retry::retry_later(EmbedError::openai_network(transport))) - } - } - } - - fn try_embed + serde::Serialize>( - &self, - texts: &[S], - client: &ureq::Agent, - ) -> Result>, Retry> { - for text in texts { - tracing::trace!("Received prompt: {}", text.as_ref()) - } - let request = OpenAiRequest { - model: self.options.embedding_model.name(), - input: texts, - dimensions: self.overriden_dimensions(), - }; - let response = client - .post(OPENAI_EMBEDDINGS_URL) - .set("Authorization", &self.bearer) - .send_json(&request); - - let response = Self::check_response(response)?; - - let response: OpenAiResponse = response - .into_json() - .map_err(EmbedError::openai_unexpected) - .map_err(Retry::retry_later)?; - - tracing::trace!("response: {:?}", response.data); - - Ok(response - .data - .into_iter() - .map(|data| Embeddings::from_single_embedding(data.embedding)) - .collect()) - } - - fn try_embed_tokenized( - &self, - text: &[String], - client: &ureq::Agent, - ) -> Result>, Retry> { - pub const OVERLAP_SIZE: usize = 200; - let mut all_embeddings = Vec::with_capacity(text.len()); - for text in text { - let max_token_count = self.options.embedding_model.max_token(); - let encoded = self.tokenizer.encode_ordinary(text.as_str()); - let len = encoded.len(); - if len < max_token_count { - all_embeddings.append(&mut self.try_embed(&[text], client)?); - continue; - } - - let mut tokens = encoded.as_slice(); - let mut embeddings_for_prompt = Embeddings::new(self.dimensions()); - while tokens.len() > max_token_count { - let window = &tokens[..max_token_count]; - embeddings_for_prompt.push(self.embed_tokens(window, client)?).unwrap(); - - tokens = &tokens[max_token_count - OVERLAP_SIZE..]; - } - - // end of text - embeddings_for_prompt.push(self.embed_tokens(tokens, client)?).unwrap(); - - all_embeddings.push(embeddings_for_prompt); - } - Ok(all_embeddings) - } - - fn embed_tokens(&self, tokens: &[usize], client: &ureq::Agent) -> Result { - for attempt in 0..9 { - let duration = match self.try_embed_tokens(tokens, client) { - Ok(embedding) => return Ok(embedding), - Err(retry) => retry.into_duration(attempt), - } - .map_err(Retry::retry_later)?; - - std::thread::sleep(duration); + tokens = &tokens[max_token_count - OVERLAP_SIZE..]; } - self.try_embed_tokens(tokens, client) - .map_err(|retry| Retry::give_up(retry.into_error())) + // end of text + let embedding = self.rest_embedder.embed_tokens(tokens)?; + /// FIXME: unwrap + embeddings_for_prompt.append(embedding.into_inner()).unwrap(); + + all_embeddings.push(embeddings_for_prompt); } + Ok(all_embeddings) + } - fn try_embed_tokens( - &self, - tokens: &[usize], - client: &ureq::Agent, - ) -> Result { - let request = OpenAiTokensRequest { - model: self.options.embedding_model.name(), - input: tokens, - dimensions: self.overriden_dimensions(), - }; - let response = client - .post(OPENAI_EMBEDDINGS_URL) - .set("Authorization", &self.bearer) - .send_json(&request); + pub fn embed_chunks( + &self, + text_chunks: Vec>, + threads: &rayon::ThreadPool, + ) -> Result>>, EmbedError> { + threads.install(move || { + text_chunks.into_par_iter().map(move |chunk| self.embed(chunk)).collect() + }) + } - let response = Self::check_response(response)?; + pub fn chunk_count_hint(&self) -> usize { + self.rest_embedder.chunk_count_hint() + } - let mut response: OpenAiResponse = response - .into_json() - .map_err(EmbedError::openai_unexpected) - .map_err(Retry::retry_later)?; + pub fn prompt_count_in_chunk_hint(&self) -> usize { + self.rest_embedder.prompt_count_in_chunk_hint() + } - Ok(response.data.pop().map(|data| data.embedding).unwrap_or_default()) - } + pub fn dimensions(&self) -> usize { + self.options.dimensions() + } - pub fn embed_chunks( - &self, - text_chunks: Vec>, - ) -> Result>>, EmbedError> { - self.threads - .install(move || text_chunks.into_par_iter().map(|chunk| self.embed(chunk))) - .collect() - } - - pub fn chunk_count_hint(&self) -> usize { - 10 - } - - pub fn prompt_count_in_chunk_hint(&self) -> usize { - 10 - } - - pub fn dimensions(&self) -> usize { - if self.options.embedding_model.supports_overriding_dimensions() { - self.options.dimensions.unwrap_or(self.options.embedding_model.default_dimensions()) - } else { - self.options.embedding_model.default_dimensions() - } - } - - pub fn distribution(&self) -> Option { - self.options.embedding_model.distribution() - } - - fn overriden_dimensions(&self) -> Option { - if self.options.embedding_model.supports_overriding_dimensions() { - self.options.dimensions - } else { - None - } - } + pub fn distribution(&self) -> Option { + self.options.embedding_model.distribution() } } diff --git a/milli/src/vector/rest.rs b/milli/src/vector/rest.rs index 975bd3790..6fd47d882 100644 --- a/milli/src/vector/rest.rs +++ b/milli/src/vector/rest.rs @@ -1,9 +1,62 @@ use rayon::iter::{IntoParallelIterator as _, ParallelIterator as _}; +use serde::Serialize; -use super::openai::Retry; -use super::{DistributionShift, EmbedError, Embeddings, NewEmbedderError}; -use crate::VectorOrArrayOfVectors; +use super::{ + DistributionShift, EmbedError, Embedding, Embeddings, NewEmbedderError, REQUEST_PARALLELISM, +}; +// retrying in case of failure + +pub struct Retry { + pub error: EmbedError, + strategy: RetryStrategy, +} + +pub enum RetryStrategy { + GiveUp, + Retry, + RetryTokenized, + RetryAfterRateLimit, +} + +impl Retry { + pub fn give_up(error: EmbedError) -> Self { + Self { error, strategy: RetryStrategy::GiveUp } + } + + pub fn retry_later(error: EmbedError) -> Self { + Self { error, strategy: RetryStrategy::Retry } + } + + pub fn retry_tokenized(error: EmbedError) -> Self { + Self { error, strategy: RetryStrategy::RetryTokenized } + } + + pub fn rate_limited(error: EmbedError) -> Self { + Self { error, strategy: RetryStrategy::RetryAfterRateLimit } + } + + pub fn into_duration(self, attempt: u32) -> Result { + match self.strategy { + RetryStrategy::GiveUp => Err(self.error), + RetryStrategy::Retry => Ok(std::time::Duration::from_millis((10u64).pow(attempt))), + RetryStrategy::RetryTokenized => Ok(std::time::Duration::from_millis(1)), + RetryStrategy::RetryAfterRateLimit => { + Ok(std::time::Duration::from_millis(100 + 10u64.pow(attempt))) + } + } + } + + pub fn must_tokenize(&self) -> bool { + matches!(self.strategy, RetryStrategy::RetryTokenized) + } + + pub fn into_error(self) -> EmbedError { + self.error + } +} + +#[derive(Debug)] pub struct Embedder { client: ureq::Agent, options: EmbedderOptions, @@ -11,20 +64,35 @@ pub struct Embedder { dimensions: usize, } +#[derive(Debug)] pub struct EmbedderOptions { - api_key: Option, - distribution: Option, - dimensions: Option, - url: String, - query: liquid::Template, - response_field: Vec, + pub api_key: Option, + pub distribution: Option, + pub dimensions: Option, + pub url: String, + pub query: serde_json::Value, + pub input_field: Vec, + // path to the array of embeddings + pub path_to_embeddings: Vec, + // shape of a single embedding + pub embedding_object: Vec, + pub input_type: InputType, +} + +#[derive(Debug)] +pub enum InputType { + Text, + TextArray, } impl Embedder { pub fn new(options: EmbedderOptions) -> Result { - let bearer = options.api_key.as_deref().map(|api_key| format!("Bearer: {api_key}")); + let bearer = options.api_key.as_deref().map(|api_key| format!("Bearer {api_key}")); - let client = ureq::agent(); + let client = ureq::AgentBuilder::new() + .max_idle_connections(REQUEST_PARALLELISM * 2) + .max_idle_connections_per_host(REQUEST_PARALLELISM * 2) + .build(); let dimensions = if let Some(dimensions) = options.dimensions { dimensions @@ -36,7 +104,20 @@ impl Embedder { } pub fn embed(&self, texts: Vec) -> Result>, EmbedError> { - embed(&self.client, &self.options, self.bearer.as_deref(), texts.as_slice()) + embed(&self.client, &self.options, self.bearer.as_deref(), texts.as_slice(), texts.len()) + } + + pub fn embed_ref(&self, texts: &[S]) -> Result>, EmbedError> + where + S: AsRef + Serialize, + { + embed(&self.client, &self.options, self.bearer.as_deref(), texts, texts.len()) + } + + pub fn embed_tokens(&self, tokens: &[usize]) -> Result, EmbedError> { + let mut embeddings = embed(&self.client, &self.options, self.bearer.as_deref(), tokens, 1)?; + // unwrap: guaranteed that embeddings.len() == 1, otherwise the previous line terminated in error + Ok(embeddings.pop().unwrap()) } pub fn embed_chunks( @@ -44,17 +125,20 @@ impl Embedder { text_chunks: Vec>, threads: &rayon::ThreadPool, ) -> Result>>, EmbedError> { - threads - .install(move || text_chunks.into_par_iter().map(|chunk| self.embed(chunk))) - .collect() + threads.install(move || { + text_chunks.into_par_iter().map(move |chunk| self.embed(chunk)).collect() + }) } pub fn chunk_count_hint(&self) -> usize { - 10 + super::REQUEST_PARALLELISM } pub fn prompt_count_in_chunk_hint(&self) -> usize { - 10 + match self.options.input_type { + InputType::Text => 1, + InputType::TextArray => 10, + } } pub fn dimensions(&self) -> usize { @@ -71,9 +155,9 @@ fn infer_dimensions( options: &EmbedderOptions, bearer: Option<&str>, ) -> Result { - let v = embed(client, options, bearer, ["test"].as_slice()) + let v = embed(client, options, bearer, ["test"].as_slice(), 1) .map_err(NewEmbedderError::could_not_determine_dimension)?; - // unwrap: guaranteed that v.len() == ["test"].len() == 1, otherwise the previous line terminated in error + // unwrap: guaranteed that v.len() == 1, otherwise the previous line terminated in error Ok(v.first().unwrap().dimension()) } @@ -82,33 +166,57 @@ fn embed( options: &EmbedderOptions, bearer: Option<&str>, inputs: &[S], + expected_count: usize, ) -> Result>, EmbedError> where - S: serde::Serialize, + S: Serialize, { let request = client.post(&options.url); let request = if let Some(bearer) = bearer { request.set("Authorization", bearer) } else { request }; let request = request.set("Content-Type", "application/json"); - let body = options - .query - .render( - &liquid::to_object(&serde_json::json!({ - "input": inputs, - })) - .map_err(EmbedError::rest_template_context_serialization)?, - ) - .map_err(EmbedError::rest_template_render)?; + let input_value = match options.input_type { + InputType::Text => serde_json::json!(inputs.first()), + InputType::TextArray => serde_json::json!(inputs), + }; + + let body = match options.input_field.as_slice() { + [] => { + // inject input in body + input_value + } + [input] => { + let mut body = options.query.clone(); + + /// FIXME unwrap + body.as_object_mut().unwrap().insert(input.clone(), input_value); + body + } + [path @ .., input] => { + let mut body = options.query.clone(); + + /// FIXME unwrap + let mut current_value = &mut body; + for component in path { + current_value = current_value + .as_object_mut() + .unwrap() + .entry(component.clone()) + .or_insert(serde_json::json!({})); + } + + current_value.as_object_mut().unwrap().insert(input.clone(), input_value); + body + } + }; for attempt in 0..7 { - let response = request.send_string(&body); + let response = request.clone().send_json(&body); let result = check_response(response); let retry_duration = match result { - Ok(response) => { - return response_to_embedding(response, &options.response_field, inputs.len()) - } + Ok(response) => return response_to_embedding(response, options, expected_count), Err(retry) => { tracing::warn!("Failed: {}", retry.error); retry.into_duration(attempt) @@ -120,11 +228,11 @@ where std::thread::sleep(retry_duration); } - let response = request.send_string(&body); + let response = request.send_json(&body); let result = check_response(response); result .map_err(Retry::into_error) - .and_then(|response| response_to_embedding(response, &options.response_field, inputs.len())) + .and_then(|response| response_to_embedding(response, options, expected_count)) } fn check_response(response: Result) -> Result { @@ -139,7 +247,10 @@ fn check_response(response: Result) -> Result { Retry::retry_later(EmbedError::rest_internal_server_error(code, error_response)) } - x => Retry::retry_later(EmbedError::rest_other_status_code(code, error_response)), + 402..=499 => { + Retry::give_up(EmbedError::rest_other_status_code(code, error_response)) + } + _ => Retry::retry_later(EmbedError::rest_other_status_code(code, error_response)), }) } Err(ureq::Error::Transport(transport)) => { @@ -148,34 +259,66 @@ fn check_response(response: Result) -> Result>( +fn response_to_embedding( response: ureq::Response, - response_field: &[S], + options: &EmbedderOptions, expected_count: usize, ) -> Result>, EmbedError> { let response: serde_json::Value = response.into_json().map_err(EmbedError::rest_response_deserialization)?; let mut current_value = &response; - for component in response_field { + for component in &options.path_to_embeddings { let component = component.as_ref(); - let current_value = current_value.get(component).ok_or_else(|| { - EmbedError::rest_response_missing_embeddings(response, component, response_field) + current_value = current_value.get(component).ok_or_else(|| { + EmbedError::rest_response_missing_embeddings( + response.clone(), + component, + &options.path_to_embeddings, + ) })?; } - let embeddings = current_value.to_owned(); + let embeddings = match options.input_type { + InputType::Text => { + for component in &options.embedding_object { + current_value = current_value.get(component).ok_or_else(|| { + EmbedError::rest_response_missing_embeddings( + response.clone(), + component, + &options.embedding_object, + ) + })?; + } + let embeddings = current_value.to_owned(); + let embeddings: Embedding = + serde_json::from_value(embeddings).map_err(EmbedError::rest_response_format)?; - let embeddings: VectorOrArrayOfVectors = - serde_json::from_value(embeddings).map_err(EmbedError::rest_response_format)?; - - let embeddings = embeddings.into_array_of_vectors(); - - let embeddings: Vec> = embeddings - .into_iter() - .flatten() - .map(|embedding| Embeddings::from_single_embedding(embedding)) - .collect(); + vec![Embeddings::from_single_embedding(embeddings)] + } + InputType::TextArray => { + let empty = vec![]; + let values = current_value.as_array().unwrap_or(&empty); + let mut embeddings: Vec> = Vec::with_capacity(expected_count); + for value in values { + let mut current_value = value; + for component in &options.embedding_object { + current_value = current_value.get(component).ok_or_else(|| { + EmbedError::rest_response_missing_embeddings( + response.clone(), + component, + &options.embedding_object, + ) + })?; + } + let embedding = current_value.to_owned(); + let embedding: Embedding = + serde_json::from_value(embedding).map_err(EmbedError::rest_response_format)?; + embeddings.push(Embeddings::from_single_embedding(embedding)); + } + embeddings + } + }; if embeddings.len() != expected_count { return Err(EmbedError::rest_response_embedding_count(expected_count, embeddings.len())); diff --git a/milli/src/vector/settings.rs b/milli/src/vector/settings.rs index 89571e98a..540693d44 100644 --- a/milli/src/vector/settings.rs +++ b/milli/src/vector/settings.rs @@ -204,7 +204,7 @@ impl From for EmbeddingSettings { }, super::EmbedderOptions::Ollama(options) => Self { source: Setting::Set(EmbedderSource::Ollama), - model: Setting::Set(options.embedding_model.name().to_owned()), + model: Setting::Set(options.embedding_model.to_owned()), revision: Setting::NotSet, api_key: Setting::NotSet, dimensions: Setting::NotSet, @@ -248,7 +248,7 @@ impl From for EmbeddingConfig { let mut options: ollama::EmbedderOptions = super::ollama::EmbedderOptions::with_default_model(); if let Some(model) = model.set() { - options.embedding_model = super::ollama::EmbeddingModel::from_name(&model); + options.embedding_model = model; } this.embedder_options = super::EmbedderOptions::Ollama(options); } From f649f58013c969908a2a2753ab91e1689e766b2d Mon Sep 17 00:00:00 2001 From: Louis Dureuil Date: Tue, 19 Mar 2024 15:42:53 +0100 Subject: [PATCH 31/86] embed no longer async --- meilisearch/src/routes/indexes/search.rs | 7 +++---- meilisearch/src/routes/multi_search.rs | 5 ++--- 2 files changed, 5 insertions(+), 7 deletions(-) diff --git a/meilisearch/src/routes/indexes/search.rs b/meilisearch/src/routes/indexes/search.rs index 6a430b6a3..8de2be13f 100644 --- a/meilisearch/src/routes/indexes/search.rs +++ b/meilisearch/src/routes/indexes/search.rs @@ -202,7 +202,7 @@ pub async fn search_with_url_query( let index = index_scheduler.index(&index_uid)?; let features = index_scheduler.features(); - let distribution = embed(&mut query, index_scheduler.get_ref(), &index).await?; + let distribution = embed(&mut query, index_scheduler.get_ref(), &index)?; let search_result = tokio::task::spawn_blocking(move || perform_search(&index, query, features, distribution)) @@ -241,7 +241,7 @@ pub async fn search_with_post( let features = index_scheduler.features(); - let distribution = embed(&mut query, index_scheduler.get_ref(), &index).await?; + let distribution = embed(&mut query, index_scheduler.get_ref(), &index)?; let search_result = tokio::task::spawn_blocking(move || perform_search(&index, query, features, distribution)) @@ -260,7 +260,7 @@ pub async fn search_with_post( Ok(HttpResponse::Ok().json(search_result)) } -pub async fn embed( +pub fn embed( query: &mut SearchQuery, index_scheduler: &IndexScheduler, index: &milli::Index, @@ -287,7 +287,6 @@ pub async fn embed( let embeddings = embedder .embed(vec![q.to_owned()]) - .await .map_err(milli::vector::Error::from) .map_err(milli::Error::from)? .pop() diff --git a/meilisearch/src/routes/multi_search.rs b/meilisearch/src/routes/multi_search.rs index 86aa58e70..f54b8ae8f 100644 --- a/meilisearch/src/routes/multi_search.rs +++ b/meilisearch/src/routes/multi_search.rs @@ -75,9 +75,8 @@ pub async fn multi_search_with_post( }) .with_index(query_index)?; - let distribution = embed(&mut query, index_scheduler.get_ref(), &index) - .await - .with_index(query_index)?; + let distribution = + embed(&mut query, index_scheduler.get_ref(), &index).with_index(query_index)?; let search_result = tokio::task::spawn_blocking(move || { perform_search(&index, query, features, distribution) From b6b4b6bab7563180475d0f306886870086804a32 Mon Sep 17 00:00:00 2001 From: Louis Dureuil Date: Tue, 19 Mar 2024 15:43:12 +0100 Subject: [PATCH 32/86] Remove the tokio and the reqwests --- Cargo.lock | 3 --- milli/Cargo.toml | 6 ------ 2 files changed, 9 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 60d0e4c0e..6a8c20f12 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3338,7 +3338,6 @@ dependencies = [ "filter-parser", "flatten-serde-json", "fst", - "futures", "fxhash", "geoutils", "grenad", @@ -3362,7 +3361,6 @@ dependencies = [ "rand", "rand_pcg", "rayon", - "reqwest", "roaring", "rstar", "serde", @@ -3376,7 +3374,6 @@ dependencies = [ "tiktoken-rs", "time", "tokenizers", - "tokio", "tracing", "ureq", "uuid", diff --git a/milli/Cargo.toml b/milli/Cargo.toml index 59b3699cc..4833ad00b 100644 --- a/milli/Cargo.toml +++ b/milli/Cargo.toml @@ -80,12 +80,6 @@ tokenizers = { git = "https://github.com/huggingface/tokenizers.git", tag = "v0. hf-hub = { git = "https://github.com/dureuill/hf-hub.git", branch = "rust_tls", default_features = false, features = [ "online", ] } -tokio = { version = "1.35.1", features = ["rt"] } -futures = "0.3.30" -reqwest = { version = "0.11.23", features = [ - "rustls-tls", - "json", -], default-features = false } tiktoken-rs = "0.5.8" liquid = "0.26.4" arroy = "0.2.0" From f87747f4d32af96a152776d021075782ce058c0f Mon Sep 17 00:00:00 2001 From: Louis Dureuil Date: Wed, 20 Mar 2024 13:25:10 +0100 Subject: [PATCH 33/86] Remove unwraps --- .../src/update/index_documents/extract/mod.rs | 4 +--- milli/src/vector/error.rs | 20 +++++++++++++++++-- milli/src/vector/openai.rs | 11 ++++++---- milli/src/vector/rest.rs | 18 +++++++++++++---- 4 files changed, 40 insertions(+), 13 deletions(-) diff --git a/milli/src/update/index_documents/extract/mod.rs b/milli/src/update/index_documents/extract/mod.rs index 5689bb04f..82486f3a8 100644 --- a/milli/src/update/index_documents/extract/mod.rs +++ b/milli/src/update/index_documents/extract/mod.rs @@ -242,11 +242,9 @@ fn send_original_documents_data( let request_threads = rayon::ThreadPoolBuilder::new() .num_threads(crate::vector::REQUEST_PARALLELISM) .thread_name(|index| format!("embedding-request-{index}")) - .build() - .unwrap(); + .build()?; rayon::spawn(move || { - /// FIXME: unwrap for (name, (embedder, prompt)) in embedders { let result = extract_vector_points( documents_chunk_cloned.clone(), diff --git a/milli/src/vector/error.rs b/milli/src/vector/error.rs index 92f077924..1e0bcc7fb 100644 --- a/milli/src/vector/error.rs +++ b/milli/src/vector/error.rs @@ -52,8 +52,6 @@ pub enum EmbedErrorKind { ModelForward(candle_core::Error), #[error("attempt to embed the following text in a configuration where embeddings must be user provided: {0:?}")] ManualEmbed(String), - #[error("could not initialize asynchronous runtime: {0}")] - OpenAiRuntimeInit(std::io::Error), #[error("model not found. Meilisearch will not automatically download models from the Ollama library, please pull the model manually: {0:?}")] OllamaModelNotFoundError(Option), #[error("error deserialization the response body as JSON: {0}")] @@ -76,6 +74,10 @@ pub enum EmbedErrorKind { RestOtherStatusCode(u16, Option), #[error("could not reach embedding server: {0}")] RestNetwork(ureq::Transport), + #[error("was expected '{}' to be an object in query '{0}'", .1.join("."))] + RestNotAnObject(serde_json::Value, Vec), + #[error("while embedding tokenized, was expecting embeddings of dimension `{0}`, got embeddings of dimensions `{1}`")] + OpenAiUnexpectedDimension(usize, usize), } impl EmbedError { @@ -174,6 +176,20 @@ impl EmbedError { pub(crate) fn rest_network(transport: ureq::Transport) -> EmbedError { Self { kind: EmbedErrorKind::RestNetwork(transport), fault: FaultSource::Runtime } } + + pub(crate) fn rest_not_an_object( + query: serde_json::Value, + input_path: Vec, + ) -> EmbedError { + Self { kind: EmbedErrorKind::RestNotAnObject(query, input_path), fault: FaultSource::User } + } + + pub(crate) fn openai_unexpected_dimension(expected: usize, got: usize) -> EmbedError { + Self { + kind: EmbedErrorKind::OpenAiUnexpectedDimension(expected, got), + fault: FaultSource::Runtime, + } + } } #[derive(Debug, thiserror::Error)] diff --git a/milli/src/vector/openai.rs b/milli/src/vector/openai.rs index b2638966e..737878a1a 100644 --- a/milli/src/vector/openai.rs +++ b/milli/src/vector/openai.rs @@ -210,16 +210,19 @@ impl Embedder { while tokens.len() > max_token_count { let window = &tokens[..max_token_count]; let embedding = self.rest_embedder.embed_tokens(window)?; - /// FIXME: unwrap - embeddings_for_prompt.append(embedding.into_inner()).unwrap(); + embeddings_for_prompt.append(embedding.into_inner()).map_err(|got| { + EmbedError::openai_unexpected_dimension(self.dimensions(), got.len()) + })?; tokens = &tokens[max_token_count - OVERLAP_SIZE..]; } // end of text let embedding = self.rest_embedder.embed_tokens(tokens)?; - /// FIXME: unwrap - embeddings_for_prompt.append(embedding.into_inner()).unwrap(); + + embeddings_for_prompt.append(embedding.into_inner()).map_err(|got| { + EmbedError::openai_unexpected_dimension(self.dimensions(), got.len()) + })?; all_embeddings.push(embeddings_for_prompt); } diff --git a/milli/src/vector/rest.rs b/milli/src/vector/rest.rs index 6fd47d882..8650bb68d 100644 --- a/milli/src/vector/rest.rs +++ b/milli/src/vector/rest.rs @@ -189,19 +189,29 @@ where [input] => { let mut body = options.query.clone(); - /// FIXME unwrap - body.as_object_mut().unwrap().insert(input.clone(), input_value); + body.as_object_mut() + .ok_or_else(|| { + EmbedError::rest_not_an_object( + options.query.clone(), + options.input_field.clone(), + ) + })? + .insert(input.clone(), input_value); body } [path @ .., input] => { let mut body = options.query.clone(); - /// FIXME unwrap let mut current_value = &mut body; for component in path { current_value = current_value .as_object_mut() - .unwrap() + .ok_or_else(|| { + EmbedError::rest_not_an_object( + options.query.clone(), + options.input_field.clone(), + ) + })? .entry(component.clone()) .or_insert(serde_json::json!({})); } From a1db342f01076a6ccebb2e76a8fe43fbf2ac22a0 Mon Sep 17 00:00:00 2001 From: Louis Dureuil Date: Mon, 25 Mar 2024 10:05:38 +0100 Subject: [PATCH 34/86] Expose REST embedder to the API --- milli/src/update/index_documents/mod.rs | 6 + milli/src/update/settings.rs | 74 +++++++++- milli/src/vector/hf.rs | 5 +- milli/src/vector/mod.rs | 58 +++++++- milli/src/vector/openai.rs | 22 +-- milli/src/vector/rest.rs | 41 +++++- milli/src/vector/settings.rs | 183 ++++++++++++++++++++++-- 7 files changed, 357 insertions(+), 32 deletions(-) diff --git a/milli/src/update/index_documents/mod.rs b/milli/src/update/index_documents/mod.rs index 7499b68e5..913fbc881 100644 --- a/milli/src/update/index_documents/mod.rs +++ b/milli/src/update/index_documents/mod.rs @@ -2646,6 +2646,12 @@ mod tests { api_key: Setting::NotSet, dimensions: Setting::Set(3), document_template: Setting::NotSet, + url: Setting::NotSet, + query: Setting::NotSet, + input_field: Setting::NotSet, + path_to_embeddings: Setting::NotSet, + embedding_object: Setting::NotSet, + input_type: Setting::NotSet, }), ); settings.set_embedder_settings(embedders); diff --git a/milli/src/update/settings.rs b/milli/src/update/settings.rs index f54f45e1e..4c7289eb7 100644 --- a/milli/src/update/settings.rs +++ b/milli/src/update/settings.rs @@ -1140,6 +1140,12 @@ fn validate_prompt( api_key, dimensions, document_template: Setting::Set(template), + url, + query, + input_field, + path_to_embeddings, + embedding_object, + input_type, }) => { // validate let template = crate::prompt::Prompt::new(template) @@ -1153,6 +1159,12 @@ fn validate_prompt( api_key, dimensions, document_template: Setting::Set(template), + url, + query, + input_field, + path_to_embeddings, + embedding_object, + input_type, })) } new => Ok(new), @@ -1165,8 +1177,20 @@ pub fn validate_embedding_settings( ) -> Result> { let settings = validate_prompt(name, settings)?; let Setting::Set(settings) = settings else { return Ok(settings) }; - let EmbeddingSettings { source, model, revision, api_key, dimensions, document_template } = - settings; + let EmbeddingSettings { + source, + model, + revision, + api_key, + dimensions, + document_template, + url, + query, + input_field, + path_to_embeddings, + embedding_object, + input_type, + } = settings; if let Some(0) = dimensions.set() { return Err(crate::error::UserError::InvalidSettingsDimensions { @@ -1183,11 +1207,25 @@ pub fn validate_embedding_settings( api_key, dimensions, document_template, + url, + query, + input_field, + path_to_embeddings, + embedding_object, + input_type, })); }; match inferred_source { EmbedderSource::OpenAi => { check_unset(&revision, "revision", inferred_source, name)?; + + check_unset(&url, "url", inferred_source, name)?; + check_unset(&query, "query", inferred_source, name)?; + check_unset(&input_field, "inputField", inferred_source, name)?; + check_unset(&path_to_embeddings, "pathToEmbeddings", inferred_source, name)?; + check_unset(&embedding_object, "embeddingObject", inferred_source, name)?; + check_unset(&input_type, "inputType", inferred_source, name)?; + if let Setting::Set(model) = &model { let model = crate::vector::openai::EmbeddingModel::from_name(model.as_str()) .ok_or(crate::error::UserError::InvalidOpenAiModel { @@ -1224,10 +1262,24 @@ pub fn validate_embedding_settings( check_set(&model, "model", inferred_source, name)?; check_unset(&api_key, "apiKey", inferred_source, name)?; check_unset(&revision, "revision", inferred_source, name)?; + + check_unset(&url, "url", inferred_source, name)?; + check_unset(&query, "query", inferred_source, name)?; + check_unset(&input_field, "inputField", inferred_source, name)?; + check_unset(&path_to_embeddings, "pathToEmbeddings", inferred_source, name)?; + check_unset(&embedding_object, "embeddingObject", inferred_source, name)?; + check_unset(&input_type, "inputType", inferred_source, name)?; } EmbedderSource::HuggingFace => { check_unset(&api_key, "apiKey", inferred_source, name)?; check_unset(&dimensions, "dimensions", inferred_source, name)?; + + check_unset(&url, "url", inferred_source, name)?; + check_unset(&query, "query", inferred_source, name)?; + check_unset(&input_field, "inputField", inferred_source, name)?; + check_unset(&path_to_embeddings, "pathToEmbeddings", inferred_source, name)?; + check_unset(&embedding_object, "embeddingObject", inferred_source, name)?; + check_unset(&input_type, "inputType", inferred_source, name)?; } EmbedderSource::UserProvided => { check_unset(&model, "model", inferred_source, name)?; @@ -1235,6 +1287,18 @@ pub fn validate_embedding_settings( check_unset(&api_key, "apiKey", inferred_source, name)?; check_unset(&document_template, "documentTemplate", inferred_source, name)?; check_set(&dimensions, "dimensions", inferred_source, name)?; + + check_unset(&url, "url", inferred_source, name)?; + check_unset(&query, "query", inferred_source, name)?; + check_unset(&input_field, "inputField", inferred_source, name)?; + check_unset(&path_to_embeddings, "pathToEmbeddings", inferred_source, name)?; + check_unset(&embedding_object, "embeddingObject", inferred_source, name)?; + check_unset(&input_type, "inputType", inferred_source, name)?; + } + EmbedderSource::Rest => { + check_unset(&model, "model", inferred_source, name)?; + check_unset(&revision, "revision", inferred_source, name)?; + check_set(&url, "url", inferred_source, name)?; } } Ok(Setting::Set(EmbeddingSettings { @@ -1244,6 +1308,12 @@ pub fn validate_embedding_settings( api_key, dimensions, document_template, + url, + query, + input_field, + path_to_embeddings, + embedding_object, + input_type, })) } diff --git a/milli/src/vector/hf.rs b/milli/src/vector/hf.rs index 939b6210a..e341a553e 100644 --- a/milli/src/vector/hf.rs +++ b/milli/src/vector/hf.rs @@ -194,7 +194,10 @@ impl Embedder { pub fn distribution(&self) -> Option { if self.options.model == "BAAI/bge-base-en-v1.5" { - Some(DistributionShift { current_mean: 0.85, current_sigma: 0.1 }) + Some(DistributionShift { + current_mean: ordered_float::OrderedFloat(0.85), + current_sigma: ordered_float::OrderedFloat(0.1), + }) } else { None } diff --git a/milli/src/vector/mod.rs b/milli/src/vector/mod.rs index 39232e387..65654af4a 100644 --- a/milli/src/vector/mod.rs +++ b/milli/src/vector/mod.rs @@ -1,6 +1,9 @@ use std::collections::HashMap; use std::sync::Arc; +use ordered_float::OrderedFloat; +use serde::{Deserialize, Serialize}; + use self::error::{EmbedError, NewEmbedderError}; use crate::prompt::{Prompt, PromptData}; @@ -104,7 +107,10 @@ pub enum Embedder { OpenAi(openai::Embedder), /// An embedder based on the user providing the embeddings in the documents and queries. UserProvided(manual::Embedder), + /// An embedder based on making embedding queries against an embedding server. Ollama(ollama::Embedder), + /// An embedder based on making embedding queries against a generic JSON/REST embedding server. + Rest(rest::Embedder), } /// Configuration for an embedder. @@ -175,6 +181,7 @@ pub enum EmbedderOptions { OpenAi(openai::EmbedderOptions), Ollama(ollama::EmbedderOptions), UserProvided(manual::EmbedderOptions), + Rest(rest::EmbedderOptions), } impl Default for EmbedderOptions { @@ -209,6 +216,7 @@ impl Embedder { EmbedderOptions::UserProvided(options) => { Self::UserProvided(manual::Embedder::new(options)) } + EmbedderOptions::Rest(options) => Self::Rest(rest::Embedder::new(options)?), }) } @@ -224,6 +232,7 @@ impl Embedder { Embedder::OpenAi(embedder) => embedder.embed(texts), Embedder::Ollama(embedder) => embedder.embed(texts), Embedder::UserProvided(embedder) => embedder.embed(texts), + Embedder::Rest(embedder) => embedder.embed(texts), } } @@ -240,6 +249,7 @@ impl Embedder { Embedder::OpenAi(embedder) => embedder.embed_chunks(text_chunks, threads), Embedder::Ollama(embedder) => embedder.embed_chunks(text_chunks, threads), Embedder::UserProvided(embedder) => embedder.embed_chunks(text_chunks), + Embedder::Rest(embedder) => embedder.embed_chunks(text_chunks, threads), } } @@ -250,6 +260,7 @@ impl Embedder { Embedder::OpenAi(embedder) => embedder.chunk_count_hint(), Embedder::Ollama(embedder) => embedder.chunk_count_hint(), Embedder::UserProvided(_) => 1, + Embedder::Rest(embedder) => embedder.chunk_count_hint(), } } @@ -260,6 +271,7 @@ impl Embedder { Embedder::OpenAi(embedder) => embedder.prompt_count_in_chunk_hint(), Embedder::Ollama(embedder) => embedder.prompt_count_in_chunk_hint(), Embedder::UserProvided(_) => 1, + Embedder::Rest(embedder) => embedder.prompt_count_in_chunk_hint(), } } @@ -270,6 +282,7 @@ impl Embedder { Embedder::OpenAi(embedder) => embedder.dimensions(), Embedder::Ollama(embedder) => embedder.dimensions(), Embedder::UserProvided(embedder) => embedder.dimensions(), + Embedder::Rest(embedder) => embedder.dimensions(), } } @@ -280,6 +293,7 @@ impl Embedder { Embedder::OpenAi(embedder) => embedder.distribution(), Embedder::Ollama(embedder) => embedder.distribution(), Embedder::UserProvided(_embedder) => None, + Embedder::Rest(embedder) => embedder.distribution(), } } } @@ -288,17 +302,47 @@ impl Embedder { /// /// The intended use is to make the similarity score more comparable to the regular ranking score. /// This allows to correct effects where results are too "packed" around a certain value. -#[derive(Debug, Clone, Copy)] +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Deserialize, Serialize)] +#[serde(from = "DistributionShiftSerializable")] +#[serde(into = "DistributionShiftSerializable")] pub struct DistributionShift { /// Value where the results are "packed". /// /// Similarity scores are translated so that they are packed around 0.5 instead - pub current_mean: f32, + pub current_mean: OrderedFloat, /// standard deviation of a similarity score. /// /// Set below 0.4 to make the results less packed around the mean, and above 0.4 to make them more packed. - pub current_sigma: f32, + pub current_sigma: OrderedFloat, +} + +#[derive(Serialize, Deserialize)] +struct DistributionShiftSerializable { + current_mean: f32, + current_sigma: f32, +} + +impl From for DistributionShiftSerializable { + fn from( + DistributionShift { + current_mean: OrderedFloat(current_mean), + current_sigma: OrderedFloat(current_sigma), + }: DistributionShift, + ) -> Self { + Self { current_mean, current_sigma } + } +} + +impl From for DistributionShift { + fn from( + DistributionShiftSerializable { current_mean, current_sigma }: DistributionShiftSerializable, + ) -> Self { + Self { + current_mean: OrderedFloat(current_mean), + current_sigma: OrderedFloat(current_sigma), + } + } } impl DistributionShift { @@ -307,11 +351,13 @@ impl DistributionShift { if sigma <= 0.0 { None } else { - Some(Self { current_mean: mean, current_sigma: sigma }) + Some(Self { current_mean: OrderedFloat(mean), current_sigma: OrderedFloat(sigma) }) } } pub fn shift(&self, score: f32) -> f32 { + let current_mean = self.current_mean.0; + let current_sigma = self.current_sigma.0; // // We're somewhat abusively mapping the distribution of distances to a gaussian. // The parameters we're given is the mean and sigma of the native result distribution. @@ -321,9 +367,9 @@ impl DistributionShift { let target_sigma = 0.4; // a^2 sig1^2 = sig2^2 => a^2 = sig2^2 / sig1^2 => a = sig2 / sig1, assuming a, sig1, and sig2 positive. - let factor = target_sigma / self.current_sigma; + let factor = target_sigma / current_sigma; // a*mu1 + b = mu2 => b = mu2 - a*mu1 - let offset = target_mean - (factor * self.current_mean); + let offset = target_mean - (factor * current_mean); let mut score = factor * score + offset; diff --git a/milli/src/vector/openai.rs b/milli/src/vector/openai.rs index 737878a1a..24e94a9f7 100644 --- a/milli/src/vector/openai.rs +++ b/milli/src/vector/openai.rs @@ -1,3 +1,4 @@ +use ordered_float::OrderedFloat; use rayon::iter::{IntoParallelIterator, ParallelIterator as _}; use super::error::{EmbedError, NewEmbedderError}; @@ -110,15 +111,18 @@ impl EmbeddingModel { fn distribution(&self) -> Option { match self { - EmbeddingModel::TextEmbeddingAda002 => { - Some(DistributionShift { current_mean: 0.90, current_sigma: 0.08 }) - } - EmbeddingModel::TextEmbedding3Large => { - Some(DistributionShift { current_mean: 0.70, current_sigma: 0.1 }) - } - EmbeddingModel::TextEmbedding3Small => { - Some(DistributionShift { current_mean: 0.75, current_sigma: 0.1 }) - } + EmbeddingModel::TextEmbeddingAda002 => Some(DistributionShift { + current_mean: OrderedFloat(0.90), + current_sigma: OrderedFloat(0.08), + }), + EmbeddingModel::TextEmbedding3Large => Some(DistributionShift { + current_mean: OrderedFloat(0.70), + current_sigma: OrderedFloat(0.1), + }), + EmbeddingModel::TextEmbedding3Small => Some(DistributionShift { + current_mean: OrderedFloat(0.75), + current_sigma: OrderedFloat(0.1), + }), } } diff --git a/milli/src/vector/rest.rs b/milli/src/vector/rest.rs index 8650bb68d..b0ea07f82 100644 --- a/milli/src/vector/rest.rs +++ b/milli/src/vector/rest.rs @@ -1,5 +1,6 @@ +use deserr::Deserr; use rayon::iter::{IntoParallelIterator as _, ParallelIterator as _}; -use serde::Serialize; +use serde::{Deserialize, Serialize}; use super::{ DistributionShift, EmbedError, Embedding, Embeddings, NewEmbedderError, REQUEST_PARALLELISM, @@ -64,7 +65,7 @@ pub struct Embedder { dimensions: usize, } -#[derive(Debug)] +#[derive(Debug, Clone, PartialEq, Eq, Deserialize, Serialize)] pub struct EmbedderOptions { pub api_key: Option, pub distribution: Option, @@ -79,7 +80,41 @@ pub struct EmbedderOptions { pub input_type: InputType, } -#[derive(Debug)] +impl Default for EmbedderOptions { + fn default() -> Self { + Self { + url: Default::default(), + query: Default::default(), + input_field: vec!["input".into()], + path_to_embeddings: vec!["data".into()], + embedding_object: vec!["embedding".into()], + input_type: InputType::Text, + api_key: None, + distribution: None, + dimensions: None, + } + } +} + +impl std::hash::Hash for EmbedderOptions { + fn hash(&self, state: &mut H) { + self.api_key.hash(state); + self.distribution.hash(state); + self.dimensions.hash(state); + self.url.hash(state); + // skip hashing the query + // collisions in regular usage should be minimal, + // and the list is limited to 256 values anyway + self.input_field.hash(state); + self.path_to_embeddings.hash(state); + self.embedding_object.hash(state); + self.input_type.hash(state); + } +} + +#[derive(Debug, Clone, Copy, Deserialize, Serialize, PartialEq, Eq, Hash, Deserr)] +#[serde(rename_all = "camelCase")] +#[deserr(rename_all = camelCase, deny_unknown_fields)] pub enum InputType { Text, TextArray, diff --git a/milli/src/vector/settings.rs b/milli/src/vector/settings.rs index 540693d44..c5b0d0326 100644 --- a/milli/src/vector/settings.rs +++ b/milli/src/vector/settings.rs @@ -1,6 +1,7 @@ use deserr::Deserr; use serde::{Deserialize, Serialize}; +use super::rest::InputType; use super::{ollama, openai}; use crate::prompt::PromptData; use crate::update::Setting; @@ -29,6 +30,24 @@ pub struct EmbeddingSettings { #[serde(default, skip_serializing_if = "Setting::is_not_set")] #[deserr(default)] pub document_template: Setting, + #[serde(default, skip_serializing_if = "Setting::is_not_set")] + #[deserr(default)] + pub url: Setting, + #[serde(default, skip_serializing_if = "Setting::is_not_set")] + #[deserr(default)] + pub query: Setting, + #[serde(default, skip_serializing_if = "Setting::is_not_set")] + #[deserr(default)] + pub input_field: Setting>, + #[serde(default, skip_serializing_if = "Setting::is_not_set")] + #[deserr(default)] + pub path_to_embeddings: Setting>, + #[serde(default, skip_serializing_if = "Setting::is_not_set")] + #[deserr(default)] + pub embedding_object: Setting>, + #[serde(default, skip_serializing_if = "Setting::is_not_set")] + #[deserr(default)] + pub input_type: Setting, } pub fn check_unset( @@ -75,20 +94,42 @@ impl EmbeddingSettings { pub const DIMENSIONS: &'static str = "dimensions"; pub const DOCUMENT_TEMPLATE: &'static str = "documentTemplate"; + pub const URL: &'static str = "url"; + pub const QUERY: &'static str = "query"; + pub const INPUT_FIELD: &'static str = "inputField"; + pub const PATH_TO_EMBEDDINGS: &'static str = "pathToEmbeddings"; + pub const EMBEDDING_OBJECT: &'static str = "embeddingObject"; + pub const INPUT_TYPE: &'static str = "inputType"; + pub fn allowed_sources_for_field(field: &'static str) -> &'static [EmbedderSource] { match field { - Self::SOURCE => { - &[EmbedderSource::HuggingFace, EmbedderSource::OpenAi, EmbedderSource::UserProvided] - } + Self::SOURCE => &[ + EmbedderSource::HuggingFace, + EmbedderSource::OpenAi, + EmbedderSource::UserProvided, + EmbedderSource::Rest, + EmbedderSource::Ollama, + ], Self::MODEL => { &[EmbedderSource::HuggingFace, EmbedderSource::OpenAi, EmbedderSource::Ollama] } Self::REVISION => &[EmbedderSource::HuggingFace], - Self::API_KEY => &[EmbedderSource::OpenAi], - Self::DIMENSIONS => &[EmbedderSource::OpenAi, EmbedderSource::UserProvided], - Self::DOCUMENT_TEMPLATE => { - &[EmbedderSource::HuggingFace, EmbedderSource::OpenAi, EmbedderSource::Ollama] + Self::API_KEY => &[EmbedderSource::OpenAi, EmbedderSource::Rest], + Self::DIMENSIONS => { + &[EmbedderSource::OpenAi, EmbedderSource::UserProvided, EmbedderSource::Rest] } + Self::DOCUMENT_TEMPLATE => &[ + EmbedderSource::HuggingFace, + EmbedderSource::OpenAi, + EmbedderSource::Ollama, + EmbedderSource::Rest, + ], + Self::URL => &[EmbedderSource::Rest], + Self::QUERY => &[EmbedderSource::Rest], + Self::INPUT_FIELD => &[EmbedderSource::Rest], + Self::PATH_TO_EMBEDDINGS => &[EmbedderSource::Rest], + Self::EMBEDDING_OBJECT => &[EmbedderSource::Rest], + Self::INPUT_TYPE => &[EmbedderSource::Rest], _other => unreachable!("unknown field"), } } @@ -107,6 +148,18 @@ impl EmbeddingSettings { } EmbedderSource::Ollama => &[Self::SOURCE, Self::MODEL, Self::DOCUMENT_TEMPLATE], EmbedderSource::UserProvided => &[Self::SOURCE, Self::DIMENSIONS], + EmbedderSource::Rest => &[ + Self::SOURCE, + Self::API_KEY, + Self::DIMENSIONS, + Self::DOCUMENT_TEMPLATE, + Self::URL, + Self::QUERY, + Self::INPUT_FIELD, + Self::PATH_TO_EMBEDDINGS, + Self::EMBEDDING_OBJECT, + Self::INPUT_TYPE, + ], } } @@ -141,6 +194,7 @@ pub enum EmbedderSource { HuggingFace, Ollama, UserProvided, + Rest, } impl std::fmt::Display for EmbedderSource { @@ -150,6 +204,7 @@ impl std::fmt::Display for EmbedderSource { EmbedderSource::HuggingFace => "huggingFace", EmbedderSource::UserProvided => "userProvided", EmbedderSource::Ollama => "ollama", + EmbedderSource::Rest => "rest", }; f.write_str(s) } @@ -157,8 +212,20 @@ impl std::fmt::Display for EmbedderSource { impl EmbeddingSettings { pub fn apply(&mut self, new: Self) { - let EmbeddingSettings { source, model, revision, api_key, dimensions, document_template } = - new; + let EmbeddingSettings { + source, + model, + revision, + api_key, + dimensions, + document_template, + url, + query, + input_field, + path_to_embeddings, + embedding_object, + input_type, + } = new; let old_source = self.source; self.source.apply(source); // Reinitialize the whole setting object on a source change @@ -170,6 +237,12 @@ impl EmbeddingSettings { api_key, dimensions, document_template, + url, + query, + input_field, + path_to_embeddings, + embedding_object, + input_type, }; return; } @@ -179,6 +252,13 @@ impl EmbeddingSettings { self.api_key.apply(api_key); self.dimensions.apply(dimensions); self.document_template.apply(document_template); + + self.url.apply(url); + self.query.apply(query); + self.input_field.apply(input_field); + self.path_to_embeddings.apply(path_to_embeddings); + self.embedding_object.apply(embedding_object); + self.input_type.apply(input_type); } } @@ -193,6 +273,12 @@ impl From for EmbeddingSettings { api_key: Setting::NotSet, dimensions: Setting::NotSet, document_template: Setting::Set(prompt.template), + url: Setting::NotSet, + query: Setting::NotSet, + input_field: Setting::NotSet, + path_to_embeddings: Setting::NotSet, + embedding_object: Setting::NotSet, + input_type: Setting::NotSet, }, super::EmbedderOptions::OpenAi(options) => Self { source: Setting::Set(EmbedderSource::OpenAi), @@ -201,6 +287,12 @@ impl From for EmbeddingSettings { api_key: options.api_key.map(Setting::Set).unwrap_or_default(), dimensions: options.dimensions.map(Setting::Set).unwrap_or_default(), document_template: Setting::Set(prompt.template), + url: Setting::NotSet, + query: Setting::NotSet, + input_field: Setting::NotSet, + path_to_embeddings: Setting::NotSet, + embedding_object: Setting::NotSet, + input_type: Setting::NotSet, }, super::EmbedderOptions::Ollama(options) => Self { source: Setting::Set(EmbedderSource::Ollama), @@ -209,6 +301,12 @@ impl From for EmbeddingSettings { api_key: Setting::NotSet, dimensions: Setting::NotSet, document_template: Setting::Set(prompt.template), + url: Setting::NotSet, + query: Setting::NotSet, + input_field: Setting::NotSet, + path_to_embeddings: Setting::NotSet, + embedding_object: Setting::NotSet, + input_type: Setting::NotSet, }, super::EmbedderOptions::UserProvided(options) => Self { source: Setting::Set(EmbedderSource::UserProvided), @@ -217,6 +315,37 @@ impl From for EmbeddingSettings { api_key: Setting::NotSet, dimensions: Setting::Set(options.dimensions), document_template: Setting::NotSet, + url: Setting::NotSet, + query: Setting::NotSet, + input_field: Setting::NotSet, + path_to_embeddings: Setting::NotSet, + embedding_object: Setting::NotSet, + input_type: Setting::NotSet, + }, + super::EmbedderOptions::Rest(super::rest::EmbedderOptions { + api_key, + // TODO: support distribution + distribution: _, + dimensions, + url, + query, + input_field, + path_to_embeddings, + embedding_object, + input_type, + }) => Self { + source: Setting::Set(EmbedderSource::Rest), + model: Setting::NotSet, + revision: Setting::NotSet, + api_key: api_key.map(Setting::Set).unwrap_or_default(), + dimensions: dimensions.map(Setting::Set).unwrap_or_default(), + document_template: Setting::Set(prompt.template), + url: Setting::Set(url), + query: Setting::Set(query), + input_field: Setting::Set(input_field), + path_to_embeddings: Setting::Set(path_to_embeddings), + embedding_object: Setting::Set(embedding_object), + input_type: Setting::Set(input_type), }, } } @@ -225,8 +354,20 @@ impl From for EmbeddingSettings { impl From for EmbeddingConfig { fn from(value: EmbeddingSettings) -> Self { let mut this = Self::default(); - let EmbeddingSettings { source, model, revision, api_key, dimensions, document_template } = - value; + let EmbeddingSettings { + source, + model, + revision, + api_key, + dimensions, + document_template, + url, + query, + input_field, + path_to_embeddings, + embedding_object, + input_type, + } = value; if let Some(source) = source.set() { match source { EmbedderSource::OpenAi => { @@ -274,6 +415,26 @@ impl From for EmbeddingConfig { dimensions: dimensions.set().unwrap(), }); } + EmbedderSource::Rest => { + let embedder_options = super::rest::EmbedderOptions::default(); + + this.embedder_options = + super::EmbedderOptions::Rest(super::rest::EmbedderOptions { + api_key: api_key.set(), + distribution: None, + dimensions: dimensions.set(), + url: url.set().unwrap(), + query: query.set().unwrap_or(embedder_options.query), + input_field: input_field.set().unwrap_or(embedder_options.input_field), + path_to_embeddings: path_to_embeddings + .set() + .unwrap_or(embedder_options.path_to_embeddings), + embedding_object: embedding_object + .set() + .unwrap_or(embedder_options.embedding_object), + input_type: input_type.set().unwrap_or(embedder_options.input_type), + }) + } } } From dfa5e41ea6fa37f1d5df6710da543b00050417d9 Mon Sep 17 00:00:00 2001 From: Louis Dureuil Date: Mon, 25 Mar 2024 10:05:58 +0100 Subject: [PATCH 35/86] Check validity of the URL setting --- Cargo.lock | 1 + meilisearch-types/src/error.rs | 1 + meilisearch/src/routes/indexes/settings.rs | 1 + milli/Cargo.toml | 1 + milli/src/error.rs | 2 ++ milli/src/update/settings.rs | 8 ++++++++ 6 files changed, 14 insertions(+) diff --git a/Cargo.lock b/Cargo.lock index 6a8c20f12..214ba368f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3376,6 +3376,7 @@ dependencies = [ "tokenizers", "tracing", "ureq", + "url", "uuid", ] diff --git a/meilisearch-types/src/error.rs b/meilisearch-types/src/error.rs index aed77411a..1b94201f2 100644 --- a/meilisearch-types/src/error.rs +++ b/meilisearch-types/src/error.rs @@ -353,6 +353,7 @@ impl ErrorCode for milli::Error { | UserError::InvalidOpenAiModelDimensions { .. } | UserError::InvalidOpenAiModelDimensionsMax { .. } | UserError::InvalidSettingsDimensions { .. } + | UserError::InvalidUrl { .. } | UserError::InvalidPrompt(_) => Code::InvalidSettingsEmbedders, UserError::TooManyEmbedders(_) => Code::InvalidSettingsEmbedders, UserError::InvalidPromptForEmbeddings(..) => Code::InvalidSettingsEmbedders, diff --git a/meilisearch/src/routes/indexes/settings.rs b/meilisearch/src/routes/indexes/settings.rs index 5dabd7b0d..99c3d0fbb 100644 --- a/meilisearch/src/routes/indexes/settings.rs +++ b/meilisearch/src/routes/indexes/settings.rs @@ -605,6 +605,7 @@ fn embedder_analytics( EmbedderSource::HuggingFace => sources.insert("huggingFace"), EmbedderSource::UserProvided => sources.insert("userProvided"), EmbedderSource::Ollama => sources.insert("ollama"), + EmbedderSource::Rest => sources.insert("rest"), }; } }; diff --git a/milli/Cargo.toml b/milli/Cargo.toml index 4833ad00b..9f5803f4e 100644 --- a/milli/Cargo.toml +++ b/milli/Cargo.toml @@ -86,6 +86,7 @@ arroy = "0.2.0" rand = "0.8.5" tracing = "0.1.40" ureq = { version = "2.9.6", features = ["json"] } +url = "2.5.0" [dev-dependencies] mimalloc = { version = "0.1.39", default-features = false } diff --git a/milli/src/error.rs b/milli/src/error.rs index 1147085dd..aba80b475 100644 --- a/milli/src/error.rs +++ b/milli/src/error.rs @@ -243,6 +243,8 @@ only composed of alphanumeric characters (a-z A-Z 0-9), hyphens (-) and undersco }, #[error("`.embedders.{embedder_name}.dimensions`: `dimensions` cannot be zero")] InvalidSettingsDimensions { embedder_name: String }, + #[error("`.embedders.{embedder_name}.url`: could not parse `{url}`: {inner_error}")] + InvalidUrl { embedder_name: String, inner_error: url::ParseError, url: String }, } impl From for Error { diff --git a/milli/src/update/settings.rs b/milli/src/update/settings.rs index 4c7289eb7..e902badc0 100644 --- a/milli/src/update/settings.rs +++ b/milli/src/update/settings.rs @@ -1199,6 +1199,14 @@ pub fn validate_embedding_settings( .into()); } + if let Some(url) = url.as_ref().set() { + url::Url::parse(url).map_err(|error| crate::error::UserError::InvalidUrl { + embedder_name: name.to_owned(), + inner_error: error, + url: url.to_owned(), + })?; + } + let Some(inferred_source) = source.set() else { return Ok(Setting::Set(EmbeddingSettings { source, From 58972f35cb0d755fc9fc45de110beedc78565602 Mon Sep 17 00:00:00 2001 From: Louis Dureuil Date: Mon, 25 Mar 2024 11:13:21 +0100 Subject: [PATCH 36/86] Allow `url` parameter for ollama embedder --- milli/src/update/settings.rs | 1 - milli/src/vector/mod.rs | 4 ++-- milli/src/vector/ollama.rs | 11 ++++------- milli/src/vector/settings.rs | 13 ++++++++++--- 4 files changed, 16 insertions(+), 13 deletions(-) diff --git a/milli/src/update/settings.rs b/milli/src/update/settings.rs index e902badc0..62e6b20b7 100644 --- a/milli/src/update/settings.rs +++ b/milli/src/update/settings.rs @@ -1271,7 +1271,6 @@ pub fn validate_embedding_settings( check_unset(&api_key, "apiKey", inferred_source, name)?; check_unset(&revision, "revision", inferred_source, name)?; - check_unset(&url, "url", inferred_source, name)?; check_unset(&query, "query", inferred_source, name)?; check_unset(&input_field, "inputField", inferred_source, name)?; check_unset(&path_to_embeddings, "pathToEmbeddings", inferred_source, name)?; diff --git a/milli/src/vector/mod.rs b/milli/src/vector/mod.rs index 65654af4a..8186e4409 100644 --- a/milli/src/vector/mod.rs +++ b/milli/src/vector/mod.rs @@ -201,8 +201,8 @@ impl EmbedderOptions { Self::OpenAi(openai::EmbedderOptions::with_default_model(api_key)) } - pub fn ollama() -> Self { - Self::Ollama(ollama::EmbedderOptions::with_default_model()) + pub fn ollama(url: Option) -> Self { + Self::Ollama(ollama::EmbedderOptions::with_default_model(url)) } } diff --git a/milli/src/vector/ollama.rs b/milli/src/vector/ollama.rs index 9c44e8052..7edfd13b5 100644 --- a/milli/src/vector/ollama.rs +++ b/milli/src/vector/ollama.rs @@ -12,15 +12,12 @@ pub struct Embedder { #[derive(Debug, Clone, Hash, PartialEq, Eq, serde::Deserialize, serde::Serialize)] pub struct EmbedderOptions { pub embedding_model: String, + pub url: Option, } impl EmbedderOptions { - pub fn with_default_model() -> Self { - Self { embedding_model: "nomic-embed-text".into() } - } - - pub fn with_embedding_model(embedding_model: String) -> Self { - Self { embedding_model } + pub fn with_default_model(url: Option) -> Self { + Self { embedding_model: "nomic-embed-text".into(), url } } } @@ -31,7 +28,7 @@ impl Embedder { api_key: None, distribution: None, dimensions: None, - url: get_ollama_path(), + url: options.url.unwrap_or_else(get_ollama_path), query: serde_json::json!({ "model": model, }), diff --git a/milli/src/vector/settings.rs b/milli/src/vector/settings.rs index c5b0d0326..7760573f6 100644 --- a/milli/src/vector/settings.rs +++ b/milli/src/vector/settings.rs @@ -124,7 +124,7 @@ impl EmbeddingSettings { EmbedderSource::Ollama, EmbedderSource::Rest, ], - Self::URL => &[EmbedderSource::Rest], + Self::URL => &[EmbedderSource::Ollama, EmbedderSource::Rest], Self::QUERY => &[EmbedderSource::Rest], Self::INPUT_FIELD => &[EmbedderSource::Rest], Self::PATH_TO_EMBEDDINGS => &[EmbedderSource::Rest], @@ -146,7 +146,9 @@ impl EmbeddingSettings { EmbedderSource::HuggingFace => { &[Self::SOURCE, Self::MODEL, Self::REVISION, Self::DOCUMENT_TEMPLATE] } - EmbedderSource::Ollama => &[Self::SOURCE, Self::MODEL, Self::DOCUMENT_TEMPLATE], + EmbedderSource::Ollama => { + &[Self::SOURCE, Self::MODEL, Self::DOCUMENT_TEMPLATE, Self::URL] + } EmbedderSource::UserProvided => &[Self::SOURCE, Self::DIMENSIONS], EmbedderSource::Rest => &[ Self::SOURCE, @@ -387,10 +389,15 @@ impl From for EmbeddingConfig { } EmbedderSource::Ollama => { let mut options: ollama::EmbedderOptions = - super::ollama::EmbedderOptions::with_default_model(); + super::ollama::EmbedderOptions::with_default_model(None); if let Some(model) = model.set() { options.embedding_model = model; } + + if let Some(url) = url.set() { + options.url = Some(url) + } + this.embedder_options = super::EmbedderOptions::Ollama(options); } EmbedderSource::HuggingFace => { From 4136630ea5c62ea2f0dc6d9faa4a3171b24da721 Mon Sep 17 00:00:00 2001 From: Louis Dureuil Date: Mon, 25 Mar 2024 11:39:33 +0100 Subject: [PATCH 37/86] Use constants instead of raw strings in set_*set() --- milli/src/update/settings.rs | 121 ++++++++++++++++++++++++----------- 1 file changed, 83 insertions(+), 38 deletions(-) diff --git a/milli/src/update/settings.rs b/milli/src/update/settings.rs index 62e6b20b7..e0c559b85 100644 --- a/milli/src/update/settings.rs +++ b/milli/src/update/settings.rs @@ -1225,14 +1225,24 @@ pub fn validate_embedding_settings( }; match inferred_source { EmbedderSource::OpenAi => { - check_unset(&revision, "revision", inferred_source, name)?; + check_unset(&revision, EmbeddingSettings::REVISION, inferred_source, name)?; - check_unset(&url, "url", inferred_source, name)?; - check_unset(&query, "query", inferred_source, name)?; - check_unset(&input_field, "inputField", inferred_source, name)?; - check_unset(&path_to_embeddings, "pathToEmbeddings", inferred_source, name)?; - check_unset(&embedding_object, "embeddingObject", inferred_source, name)?; - check_unset(&input_type, "inputType", inferred_source, name)?; + check_unset(&url, EmbeddingSettings::URL, inferred_source, name)?; + check_unset(&query, EmbeddingSettings::QUERY, inferred_source, name)?; + check_unset(&input_field, EmbeddingSettings::INPUT_FIELD, inferred_source, name)?; + check_unset( + &path_to_embeddings, + EmbeddingSettings::PATH_TO_EMBEDDINGS, + inferred_source, + name, + )?; + check_unset( + &embedding_object, + EmbeddingSettings::EMBEDDING_OBJECT, + inferred_source, + name, + )?; + check_unset(&input_type, EmbeddingSettings::INPUT_TYPE, inferred_source, name)?; if let Setting::Set(model) = &model { let model = crate::vector::openai::EmbeddingModel::from_name(model.as_str()) @@ -1266,46 +1276,81 @@ pub fn validate_embedding_settings( } EmbedderSource::Ollama => { // Dimensions get inferred, only model name is required - check_unset(&dimensions, "dimensions", inferred_source, name)?; - check_set(&model, "model", inferred_source, name)?; - check_unset(&api_key, "apiKey", inferred_source, name)?; - check_unset(&revision, "revision", inferred_source, name)?; + check_unset(&dimensions, EmbeddingSettings::DIMENSIONS, inferred_source, name)?; + check_set(&model, EmbeddingSettings::MODEL, inferred_source, name)?; + check_unset(&api_key, EmbeddingSettings::API_KEY, inferred_source, name)?; + check_unset(&revision, EmbeddingSettings::REVISION, inferred_source, name)?; - check_unset(&query, "query", inferred_source, name)?; - check_unset(&input_field, "inputField", inferred_source, name)?; - check_unset(&path_to_embeddings, "pathToEmbeddings", inferred_source, name)?; - check_unset(&embedding_object, "embeddingObject", inferred_source, name)?; - check_unset(&input_type, "inputType", inferred_source, name)?; + check_unset(&query, EmbeddingSettings::QUERY, inferred_source, name)?; + check_unset(&input_field, EmbeddingSettings::INPUT_FIELD, inferred_source, name)?; + check_unset( + &path_to_embeddings, + EmbeddingSettings::PATH_TO_EMBEDDINGS, + inferred_source, + name, + )?; + check_unset( + &embedding_object, + EmbeddingSettings::EMBEDDING_OBJECT, + inferred_source, + name, + )?; + check_unset(&input_type, EmbeddingSettings::INPUT_TYPE, inferred_source, name)?; } EmbedderSource::HuggingFace => { - check_unset(&api_key, "apiKey", inferred_source, name)?; - check_unset(&dimensions, "dimensions", inferred_source, name)?; + check_unset(&api_key, EmbeddingSettings::API_KEY, inferred_source, name)?; + check_unset(&dimensions, EmbeddingSettings::DIMENSIONS, inferred_source, name)?; - check_unset(&url, "url", inferred_source, name)?; - check_unset(&query, "query", inferred_source, name)?; - check_unset(&input_field, "inputField", inferred_source, name)?; - check_unset(&path_to_embeddings, "pathToEmbeddings", inferred_source, name)?; - check_unset(&embedding_object, "embeddingObject", inferred_source, name)?; - check_unset(&input_type, "inputType", inferred_source, name)?; + check_unset(&url, EmbeddingSettings::URL, inferred_source, name)?; + check_unset(&query, EmbeddingSettings::QUERY, inferred_source, name)?; + check_unset(&input_field, EmbeddingSettings::INPUT_FIELD, inferred_source, name)?; + check_unset( + &path_to_embeddings, + EmbeddingSettings::PATH_TO_EMBEDDINGS, + inferred_source, + name, + )?; + check_unset( + &embedding_object, + EmbeddingSettings::EMBEDDING_OBJECT, + inferred_source, + name, + )?; + check_unset(&input_type, EmbeddingSettings::INPUT_TYPE, inferred_source, name)?; } EmbedderSource::UserProvided => { - check_unset(&model, "model", inferred_source, name)?; - check_unset(&revision, "revision", inferred_source, name)?; - check_unset(&api_key, "apiKey", inferred_source, name)?; - check_unset(&document_template, "documentTemplate", inferred_source, name)?; - check_set(&dimensions, "dimensions", inferred_source, name)?; + check_unset(&model, EmbeddingSettings::MODEL, inferred_source, name)?; + check_unset(&revision, EmbeddingSettings::REVISION, inferred_source, name)?; + check_unset(&api_key, EmbeddingSettings::API_KEY, inferred_source, name)?; + check_unset( + &document_template, + EmbeddingSettings::DOCUMENT_TEMPLATE, + inferred_source, + name, + )?; + check_set(&dimensions, EmbeddingSettings::DIMENSIONS, inferred_source, name)?; - check_unset(&url, "url", inferred_source, name)?; - check_unset(&query, "query", inferred_source, name)?; - check_unset(&input_field, "inputField", inferred_source, name)?; - check_unset(&path_to_embeddings, "pathToEmbeddings", inferred_source, name)?; - check_unset(&embedding_object, "embeddingObject", inferred_source, name)?; - check_unset(&input_type, "inputType", inferred_source, name)?; + check_unset(&url, EmbeddingSettings::URL, inferred_source, name)?; + check_unset(&query, EmbeddingSettings::QUERY, inferred_source, name)?; + check_unset(&input_field, EmbeddingSettings::INPUT_FIELD, inferred_source, name)?; + check_unset( + &path_to_embeddings, + EmbeddingSettings::PATH_TO_EMBEDDINGS, + inferred_source, + name, + )?; + check_unset( + &embedding_object, + EmbeddingSettings::EMBEDDING_OBJECT, + inferred_source, + name, + )?; + check_unset(&input_type, EmbeddingSettings::INPUT_TYPE, inferred_source, name)?; } EmbedderSource::Rest => { - check_unset(&model, "model", inferred_source, name)?; - check_unset(&revision, "revision", inferred_source, name)?; - check_set(&url, "url", inferred_source, name)?; + check_unset(&model, EmbeddingSettings::MODEL, inferred_source, name)?; + check_unset(&revision, EmbeddingSettings::REVISION, inferred_source, name)?; + check_set(&url, EmbeddingSettings::URL, inferred_source, name)?; } } Ok(Setting::Set(EmbeddingSettings { From 817ccc089a1defa4e88db4494e5d9bdfb38d13e8 Mon Sep 17 00:00:00 2001 From: Louis Dureuil Date: Mon, 25 Mar 2024 11:50:00 +0100 Subject: [PATCH 38/86] also allow `api_key` --- milli/src/update/settings.rs | 1 - milli/src/vector/mod.rs | 4 ++-- milli/src/vector/ollama.rs | 7 ++++--- milli/src/vector/settings.rs | 15 ++++++++------- 4 files changed, 14 insertions(+), 13 deletions(-) diff --git a/milli/src/update/settings.rs b/milli/src/update/settings.rs index e0c559b85..2b1be9453 100644 --- a/milli/src/update/settings.rs +++ b/milli/src/update/settings.rs @@ -1278,7 +1278,6 @@ pub fn validate_embedding_settings( // Dimensions get inferred, only model name is required check_unset(&dimensions, EmbeddingSettings::DIMENSIONS, inferred_source, name)?; check_set(&model, EmbeddingSettings::MODEL, inferred_source, name)?; - check_unset(&api_key, EmbeddingSettings::API_KEY, inferred_source, name)?; check_unset(&revision, EmbeddingSettings::REVISION, inferred_source, name)?; check_unset(&query, EmbeddingSettings::QUERY, inferred_source, name)?; diff --git a/milli/src/vector/mod.rs b/milli/src/vector/mod.rs index 8186e4409..8b25de56d 100644 --- a/milli/src/vector/mod.rs +++ b/milli/src/vector/mod.rs @@ -201,8 +201,8 @@ impl EmbedderOptions { Self::OpenAi(openai::EmbedderOptions::with_default_model(api_key)) } - pub fn ollama(url: Option) -> Self { - Self::Ollama(ollama::EmbedderOptions::with_default_model(url)) + pub fn ollama(api_key: Option, url: Option) -> Self { + Self::Ollama(ollama::EmbedderOptions::with_default_model(api_key, url)) } } diff --git a/milli/src/vector/ollama.rs b/milli/src/vector/ollama.rs index 7edfd13b5..578b6c8e2 100644 --- a/milli/src/vector/ollama.rs +++ b/milli/src/vector/ollama.rs @@ -13,11 +13,12 @@ pub struct Embedder { pub struct EmbedderOptions { pub embedding_model: String, pub url: Option, + pub api_key: Option, } impl EmbedderOptions { - pub fn with_default_model(url: Option) -> Self { - Self { embedding_model: "nomic-embed-text".into(), url } + pub fn with_default_model(api_key: Option, url: Option) -> Self { + Self { embedding_model: "nomic-embed-text".into(), api_key, url } } } @@ -25,7 +26,7 @@ impl Embedder { pub fn new(options: EmbedderOptions) -> Result { let model = options.embedding_model.as_str(); let rest_embedder = match RestEmbedder::new(RestEmbedderOptions { - api_key: None, + api_key: options.api_key, distribution: None, dimensions: None, url: options.url.unwrap_or_else(get_ollama_path), diff --git a/milli/src/vector/settings.rs b/milli/src/vector/settings.rs index 7760573f6..c277dd0cf 100644 --- a/milli/src/vector/settings.rs +++ b/milli/src/vector/settings.rs @@ -114,7 +114,9 @@ impl EmbeddingSettings { &[EmbedderSource::HuggingFace, EmbedderSource::OpenAi, EmbedderSource::Ollama] } Self::REVISION => &[EmbedderSource::HuggingFace], - Self::API_KEY => &[EmbedderSource::OpenAi, EmbedderSource::Rest], + Self::API_KEY => { + &[EmbedderSource::OpenAi, EmbedderSource::Ollama, EmbedderSource::Rest] + } Self::DIMENSIONS => { &[EmbedderSource::OpenAi, EmbedderSource::UserProvided, EmbedderSource::Rest] } @@ -147,7 +149,7 @@ impl EmbeddingSettings { &[Self::SOURCE, Self::MODEL, Self::REVISION, Self::DOCUMENT_TEMPLATE] } EmbedderSource::Ollama => { - &[Self::SOURCE, Self::MODEL, Self::DOCUMENT_TEMPLATE, Self::URL] + &[Self::SOURCE, Self::MODEL, Self::DOCUMENT_TEMPLATE, Self::URL, Self::API_KEY] } EmbedderSource::UserProvided => &[Self::SOURCE, Self::DIMENSIONS], EmbedderSource::Rest => &[ @@ -389,15 +391,14 @@ impl From for EmbeddingConfig { } EmbedderSource::Ollama => { let mut options: ollama::EmbedderOptions = - super::ollama::EmbedderOptions::with_default_model(None); + super::ollama::EmbedderOptions::with_default_model( + api_key.set(), + url.set(), + ); if let Some(model) = model.set() { options.embedding_model = model; } - if let Some(url) = url.set() { - options.url = Some(url) - } - this.embedder_options = super::EmbedderOptions::Ollama(options); } EmbedderSource::HuggingFace => { From f82d05607258cff6ed716585d09024ef294c5fb7 Mon Sep 17 00:00:00 2001 From: Louis Dureuil Date: Tue, 26 Mar 2024 10:36:24 +0100 Subject: [PATCH 39/86] Hide secrets in settings and task queue --- index-scheduler/src/batch.rs | 6 ++- meilisearch-types/src/settings.rs | 53 +++++++++++++++++++++- meilisearch-types/src/task_view.rs | 3 +- meilisearch/src/routes/indexes/settings.rs | 6 +-- meilitool/src/main.rs | 6 ++- 5 files changed, 66 insertions(+), 8 deletions(-) diff --git a/index-scheduler/src/batch.rs b/index-scheduler/src/batch.rs index b7e31c136..3161dc499 100644 --- a/index-scheduler/src/batch.rs +++ b/index-scheduler/src/batch.rs @@ -920,7 +920,11 @@ impl IndexScheduler { } // 3.2. Dump the settings - let settings = meilisearch_types::settings::settings(index, &rtxn)?; + let settings = meilisearch_types::settings::settings( + index, + &rtxn, + meilisearch_types::settings::SecretPolicy::RevealSecrets, + )?; index_dumper.settings(&settings)?; Ok(()) })?; diff --git a/meilisearch-types/src/settings.rs b/meilisearch-types/src/settings.rs index 5480e72c6..ce3a74d69 100644 --- a/meilisearch-types/src/settings.rs +++ b/meilisearch-types/src/settings.rs @@ -211,6 +211,43 @@ pub struct Settings { pub _kind: PhantomData, } +impl Settings { + pub fn hide_secrets(&mut self) { + let Setting::Set(embedders) = &mut self.embedders else { + return; + }; + + for mut embedder in embedders.values_mut() { + let Setting::Set(embedder) = &mut embedder else { + continue; + }; + + let Setting::Set(api_key) = &mut embedder.api_key else { + continue; + }; + + Self::hide_secret(api_key); + } + } + + fn hide_secret(secret: &mut String) { + match secret.len() { + x if x < 10 => { + secret.replace_range(.., "XXX..."); + } + x if x < 20 => { + secret.replace_range(2.., "XXXX..."); + } + x if x < 30 => { + secret.replace_range(3.., "XXXXX..."); + } + _x => { + secret.replace_range(5.., "XXXXXX..."); + } + } + } +} + impl Settings { pub fn cleared() -> Settings { Settings { @@ -555,9 +592,15 @@ pub fn apply_settings_to_builder( } } +pub enum SecretPolicy { + RevealSecrets, + HideSecrets, +} + pub fn settings( index: &Index, rtxn: &crate::heed::RoTxn, + secret_policy: SecretPolicy, ) -> Result, milli::Error> { let displayed_attributes = index.displayed_fields(rtxn)?.map(|fields| fields.into_iter().map(String::from).collect()); @@ -643,7 +686,7 @@ pub fn settings( let search_cutoff_ms = index.search_cutoff(rtxn)?; - Ok(Settings { + let mut settings = Settings { displayed_attributes: match displayed_attributes { Some(attrs) => Setting::Set(attrs), None => Setting::Reset, @@ -674,7 +717,13 @@ pub fn settings( None => Setting::Reset, }, _kind: PhantomData, - }) + }; + + if let SecretPolicy::HideSecrets = secret_policy { + settings.hide_secrets() + } + + Ok(settings) } #[derive(Debug, Clone, PartialEq, Eq, Deserr)] diff --git a/meilisearch-types/src/task_view.rs b/meilisearch-types/src/task_view.rs index 02be91a88..659427c9d 100644 --- a/meilisearch-types/src/task_view.rs +++ b/meilisearch-types/src/task_view.rs @@ -86,7 +86,8 @@ impl From
for DetailsView { ..DetailsView::default() } } - Details::SettingsUpdate { settings } => { + Details::SettingsUpdate { mut settings } => { + settings.hide_secrets(); DetailsView { settings: Some(settings), ..DetailsView::default() } } Details::IndexInfo { primary_key } => { diff --git a/meilisearch/src/routes/indexes/settings.rs b/meilisearch/src/routes/indexes/settings.rs index 99c3d0fbb..0918444ef 100644 --- a/meilisearch/src/routes/indexes/settings.rs +++ b/meilisearch/src/routes/indexes/settings.rs @@ -7,7 +7,7 @@ use meilisearch_types::error::ResponseError; use meilisearch_types::facet_values_sort::FacetValuesSort; use meilisearch_types::index_uid::IndexUid; use meilisearch_types::milli::update::Setting; -use meilisearch_types::settings::{settings, RankingRuleView, Settings, Unchecked}; +use meilisearch_types::settings::{settings, RankingRuleView, SecretPolicy, Settings, Unchecked}; use meilisearch_types::tasks::KindWithContent; use serde_json::json; use tracing::debug; @@ -134,7 +134,7 @@ macro_rules! make_setting_route { let index = index_scheduler.index(&index_uid)?; let rtxn = index.read_txn()?; - let settings = settings(&index, &rtxn)?; + let settings = settings(&index, &rtxn, meilisearch_types::settings::SecretPolicy::HideSecrets)?; debug!(returns = ?settings, "Update settings"); let mut json = serde_json::json!(&settings); @@ -819,7 +819,7 @@ pub async fn get_all( let index = index_scheduler.index(&index_uid)?; let rtxn = index.read_txn()?; - let new_settings = settings(&index, &rtxn)?; + let new_settings = settings(&index, &rtxn, SecretPolicy::HideSecrets)?; debug!(returns = ?new_settings, "Get all settings"); Ok(HttpResponse::Ok().json(new_settings)) } diff --git a/meilitool/src/main.rs b/meilitool/src/main.rs index f199df216..bace7d16b 100644 --- a/meilitool/src/main.rs +++ b/meilitool/src/main.rs @@ -291,7 +291,11 @@ fn export_a_dump( } // 4.2. Dump the settings - let settings = meilisearch_types::settings::settings(&index, &rtxn)?; + let settings = meilisearch_types::settings::settings( + &index, + &rtxn, + meilisearch_types::settings::SecretPolicy::RevealSecrets, + )?; index_dumper.settings(&settings)?; count += 1; } From 9a95ed619dc41f502c11e509646adbe9872acbee Mon Sep 17 00:00:00 2001 From: Louis Dureuil Date: Tue, 26 Mar 2024 10:36:56 +0100 Subject: [PATCH 40/86] Add tests --- index-scheduler/src/lib.rs | 60 ++++++++ ...x_scheduler__tests__settings_update-2.snap | 13 ++ ...x_scheduler__tests__settings_update-3.snap | 23 ++++ ...dex_scheduler__tests__settings_update.snap | 13 ++ .../after_registering_settings_task.snap | 36 +++++ .../settings_update_processed.snap | 40 ++++++ meilisearch/tests/settings/get_settings.rs | 130 ++++++++++++++++++ 7 files changed, 315 insertions(+) create mode 100644 index-scheduler/src/snapshots/index_scheduler__tests__settings_update-2.snap create mode 100644 index-scheduler/src/snapshots/index_scheduler__tests__settings_update-3.snap create mode 100644 index-scheduler/src/snapshots/index_scheduler__tests__settings_update.snap create mode 100644 index-scheduler/src/snapshots/lib.rs/test_settings_update/after_registering_settings_task.snap create mode 100644 index-scheduler/src/snapshots/lib.rs/test_settings_update/settings_update_processed.snap diff --git a/index-scheduler/src/lib.rs b/index-scheduler/src/lib.rs index 46ed76649..84416e869 100644 --- a/index-scheduler/src/lib.rs +++ b/index-scheduler/src/lib.rs @@ -3028,6 +3028,66 @@ mod tests { snapshot!(serde_json::to_string_pretty(&documents).unwrap(), name: "documents"); } + #[test] + fn test_settings_update() { + use meilisearch_types::settings::{Settings, Unchecked}; + use milli::update::Setting; + + let (index_scheduler, mut handle) = IndexScheduler::test(true, vec![]); + + let mut new_settings: Box> = Box::default(); + let mut embedders = BTreeMap::default(); + let embedding_settings = milli::vector::settings::EmbeddingSettings { + source: Setting::Set(milli::vector::settings::EmbedderSource::Rest), + api_key: Setting::Set(S("My super secret")), + url: Setting::Set(S("http://localhost:7777")), + ..Default::default() + }; + embedders.insert(S("default"), Setting::Set(embedding_settings)); + new_settings.embedders = Setting::Set(embedders); + + index_scheduler + .register( + KindWithContent::SettingsUpdate { + index_uid: S("doggos"), + new_settings, + is_deletion: false, + allow_index_creation: true, + }, + None, + false, + ) + .unwrap(); + index_scheduler.assert_internally_consistent(); + + snapshot!(snapshot_index_scheduler(&index_scheduler), name: "after_registering_settings_task"); + + { + let rtxn = index_scheduler.read_txn().unwrap(); + let task = index_scheduler.get_task(&rtxn, 0).unwrap().unwrap(); + let task = meilisearch_types::task_view::TaskView::from_task(&task); + insta::assert_json_snapshot!(task.details); + } + + handle.advance_n_successful_batches(1); + snapshot!(snapshot_index_scheduler(&index_scheduler), name: "settings_update_processed"); + + { + let rtxn = index_scheduler.read_txn().unwrap(); + let task = index_scheduler.get_task(&rtxn, 0).unwrap().unwrap(); + let task = meilisearch_types::task_view::TaskView::from_task(&task); + insta::assert_json_snapshot!(task.details); + } + + // has everything being pushed successfully in milli? + let index = index_scheduler.index("doggos").unwrap(); + let rtxn = index.read_txn().unwrap(); + + let configs = index.embedding_configs(&rtxn).unwrap(); + let (_, embedding_config) = configs.first().unwrap(); + insta::assert_json_snapshot!(embedding_config.embedder_options); + } + #[test] fn test_document_replace_without_autobatching() { let (index_scheduler, mut handle) = IndexScheduler::test(false, vec![]); diff --git a/index-scheduler/src/snapshots/index_scheduler__tests__settings_update-2.snap b/index-scheduler/src/snapshots/index_scheduler__tests__settings_update-2.snap new file mode 100644 index 000000000..85f0926b9 --- /dev/null +++ b/index-scheduler/src/snapshots/index_scheduler__tests__settings_update-2.snap @@ -0,0 +1,13 @@ +--- +source: index-scheduler/src/lib.rs +expression: task.details +--- +{ + "embedders": { + "default": { + "source": "rest", + "apiKey": "MyXXXX...", + "url": "http://localhost:7777" + } + } +} diff --git a/index-scheduler/src/snapshots/index_scheduler__tests__settings_update-3.snap b/index-scheduler/src/snapshots/index_scheduler__tests__settings_update-3.snap new file mode 100644 index 000000000..50a42d678 --- /dev/null +++ b/index-scheduler/src/snapshots/index_scheduler__tests__settings_update-3.snap @@ -0,0 +1,23 @@ +--- +source: index-scheduler/src/lib.rs +expression: embedding_config.embedder_options +--- +{ + "Rest": { + "api_key": "My super secret", + "distribution": null, + "dimensions": null, + "url": "http://localhost:7777", + "query": null, + "input_field": [ + "input" + ], + "path_to_embeddings": [ + "data" + ], + "embedding_object": [ + "embedding" + ], + "input_type": "text" + } +} diff --git a/index-scheduler/src/snapshots/index_scheduler__tests__settings_update.snap b/index-scheduler/src/snapshots/index_scheduler__tests__settings_update.snap new file mode 100644 index 000000000..85f0926b9 --- /dev/null +++ b/index-scheduler/src/snapshots/index_scheduler__tests__settings_update.snap @@ -0,0 +1,13 @@ +--- +source: index-scheduler/src/lib.rs +expression: task.details +--- +{ + "embedders": { + "default": { + "source": "rest", + "apiKey": "MyXXXX...", + "url": "http://localhost:7777" + } + } +} diff --git a/index-scheduler/src/snapshots/lib.rs/test_settings_update/after_registering_settings_task.snap b/index-scheduler/src/snapshots/lib.rs/test_settings_update/after_registering_settings_task.snap new file mode 100644 index 000000000..01bb73993 --- /dev/null +++ b/index-scheduler/src/snapshots/lib.rs/test_settings_update/after_registering_settings_task.snap @@ -0,0 +1,36 @@ +--- +source: index-scheduler/src/lib.rs +--- +### Autobatching Enabled = true +### Processing Tasks: +[] +---------------------------------------------------------------------- +### All Tasks: +0 {uid: 0, status: enqueued, details: { settings: Settings { displayed_attributes: NotSet, searchable_attributes: NotSet, filterable_attributes: NotSet, sortable_attributes: NotSet, ranking_rules: NotSet, stop_words: NotSet, non_separator_tokens: NotSet, separator_tokens: NotSet, dictionary: NotSet, synonyms: NotSet, distinct_attribute: NotSet, proximity_precision: NotSet, typo_tolerance: NotSet, faceting: NotSet, pagination: NotSet, embedders: Set({"default": Set(EmbeddingSettings { source: Set(Rest), model: NotSet, revision: NotSet, api_key: Set("My super secret"), dimensions: NotSet, document_template: NotSet, url: Set("http://localhost:7777"), query: NotSet, input_field: NotSet, path_to_embeddings: NotSet, embedding_object: NotSet, input_type: NotSet })}), search_cutoff_ms: NotSet, _kind: PhantomData } }, kind: SettingsUpdate { index_uid: "doggos", new_settings: Settings { displayed_attributes: NotSet, searchable_attributes: NotSet, filterable_attributes: NotSet, sortable_attributes: NotSet, ranking_rules: NotSet, stop_words: NotSet, non_separator_tokens: NotSet, separator_tokens: NotSet, dictionary: NotSet, synonyms: NotSet, distinct_attribute: NotSet, proximity_precision: NotSet, typo_tolerance: NotSet, faceting: NotSet, pagination: NotSet, embedders: Set({"default": Set(EmbeddingSettings { source: Set(Rest), model: NotSet, revision: NotSet, api_key: Set("My super secret"), dimensions: NotSet, document_template: NotSet, url: Set("http://localhost:7777"), query: NotSet, input_field: NotSet, path_to_embeddings: NotSet, embedding_object: NotSet, input_type: NotSet })}), search_cutoff_ms: NotSet, _kind: PhantomData }, is_deletion: false, allow_index_creation: true }} +---------------------------------------------------------------------- +### Status: +enqueued [0,] +---------------------------------------------------------------------- +### Kind: +"settingsUpdate" [0,] +---------------------------------------------------------------------- +### Index Tasks: +doggos [0,] +---------------------------------------------------------------------- +### Index Mapper: + +---------------------------------------------------------------------- +### Canceled By: + +---------------------------------------------------------------------- +### Enqueued At: +[timestamp] [0,] +---------------------------------------------------------------------- +### Started At: +---------------------------------------------------------------------- +### Finished At: +---------------------------------------------------------------------- +### File Store: + +---------------------------------------------------------------------- + diff --git a/index-scheduler/src/snapshots/lib.rs/test_settings_update/settings_update_processed.snap b/index-scheduler/src/snapshots/lib.rs/test_settings_update/settings_update_processed.snap new file mode 100644 index 000000000..d1d219da1 --- /dev/null +++ b/index-scheduler/src/snapshots/lib.rs/test_settings_update/settings_update_processed.snap @@ -0,0 +1,40 @@ +--- +source: index-scheduler/src/lib.rs +--- +### Autobatching Enabled = true +### Processing Tasks: +[] +---------------------------------------------------------------------- +### All Tasks: +0 {uid: 0, status: succeeded, details: { settings: Settings { displayed_attributes: NotSet, searchable_attributes: NotSet, filterable_attributes: NotSet, sortable_attributes: NotSet, ranking_rules: NotSet, stop_words: NotSet, non_separator_tokens: NotSet, separator_tokens: NotSet, dictionary: NotSet, synonyms: NotSet, distinct_attribute: NotSet, proximity_precision: NotSet, typo_tolerance: NotSet, faceting: NotSet, pagination: NotSet, embedders: Set({"default": Set(EmbeddingSettings { source: Set(Rest), model: NotSet, revision: NotSet, api_key: Set("My super secret"), dimensions: NotSet, document_template: NotSet, url: Set("http://localhost:7777"), query: NotSet, input_field: NotSet, path_to_embeddings: NotSet, embedding_object: NotSet, input_type: NotSet })}), search_cutoff_ms: NotSet, _kind: PhantomData } }, kind: SettingsUpdate { index_uid: "doggos", new_settings: Settings { displayed_attributes: NotSet, searchable_attributes: NotSet, filterable_attributes: NotSet, sortable_attributes: NotSet, ranking_rules: NotSet, stop_words: NotSet, non_separator_tokens: NotSet, separator_tokens: NotSet, dictionary: NotSet, synonyms: NotSet, distinct_attribute: NotSet, proximity_precision: NotSet, typo_tolerance: NotSet, faceting: NotSet, pagination: NotSet, embedders: Set({"default": Set(EmbeddingSettings { source: Set(Rest), model: NotSet, revision: NotSet, api_key: Set("My super secret"), dimensions: NotSet, document_template: NotSet, url: Set("http://localhost:7777"), query: NotSet, input_field: NotSet, path_to_embeddings: NotSet, embedding_object: NotSet, input_type: NotSet })}), search_cutoff_ms: NotSet, _kind: PhantomData }, is_deletion: false, allow_index_creation: true }} +---------------------------------------------------------------------- +### Status: +enqueued [] +succeeded [0,] +---------------------------------------------------------------------- +### Kind: +"settingsUpdate" [0,] +---------------------------------------------------------------------- +### Index Tasks: +doggos [0,] +---------------------------------------------------------------------- +### Index Mapper: +doggos: { number_of_documents: 0, field_distribution: {} } + +---------------------------------------------------------------------- +### Canceled By: + +---------------------------------------------------------------------- +### Enqueued At: +[timestamp] [0,] +---------------------------------------------------------------------- +### Started At: +[timestamp] [0,] +---------------------------------------------------------------------- +### Finished At: +[timestamp] [0,] +---------------------------------------------------------------------- +### File Store: + +---------------------------------------------------------------------- + diff --git a/meilisearch/tests/settings/get_settings.rs b/meilisearch/tests/settings/get_settings.rs index 09e38e55a..980ef3064 100644 --- a/meilisearch/tests/settings/get_settings.rs +++ b/meilisearch/tests/settings/get_settings.rs @@ -88,6 +88,136 @@ async fn get_settings() { assert_eq!(settings["searchCutoffMs"], json!(null)); } +#[actix_rt::test] +async fn secrets_are_hidden_in_settings() { + let server = Server::new().await; + let (response, code) = server.set_features(json!({"vectorStore": true})).await; + + meili_snap::snapshot!(code, @"200 OK"); + meili_snap::snapshot!(meili_snap::json_string!(response), @r###" + { + "vectorStore": true, + "metrics": false, + "logsRoute": false, + "exportPuffinReports": false + } + "###); + + let index = server.index("test"); + let (response, _code) = index.create(None).await; + index.wait_task(response.uid()).await; + + let (response, code) = index + .update_settings(json!({ + "embedders": { + "default": { + "source": "rest", + "url": "https://localhost:7777", + "apiKey": "My super secret value you will never guess" + } + } + })) + .await; + meili_snap::snapshot!(code, @"202 Accepted"); + + meili_snap::snapshot!(meili_snap::json_string!(response, { ".duration" => "[duration]", ".enqueuedAt" => "[date]", ".startedAt" => "[date]", ".finishedAt" => "[date]" }), + @r###" + { + "taskUid": 1, + "indexUid": "test", + "status": "enqueued", + "type": "settingsUpdate", + "enqueuedAt": "[date]" + } + "###); + + let settings_update_uid = response.uid(); + + index.wait_task(settings_update_uid).await; + + let (response, code) = index.settings().await; + meili_snap::snapshot!(code, @"200 OK"); + meili_snap::snapshot!(meili_snap::json_string!(response), @r###" + { + "displayedAttributes": [ + "*" + ], + "searchableAttributes": [ + "*" + ], + "filterableAttributes": [], + "sortableAttributes": [], + "rankingRules": [ + "words", + "typo", + "proximity", + "attribute", + "sort", + "exactness" + ], + "stopWords": [], + "nonSeparatorTokens": [], + "separatorTokens": [], + "dictionary": [], + "synonyms": {}, + "distinctAttribute": null, + "proximityPrecision": "byWord", + "typoTolerance": { + "enabled": true, + "minWordSizeForTypos": { + "oneTypo": 5, + "twoTypos": 9 + }, + "disableOnWords": [], + "disableOnAttributes": [] + }, + "faceting": { + "maxValuesPerFacet": 100, + "sortFacetValuesBy": { + "*": "alpha" + } + }, + "pagination": { + "maxTotalHits": 1000 + }, + "embedders": { + "default": { + "source": "rest", + "apiKey": "My suXXXXXX...", + "documentTemplate": "{% for field in fields %} {{ field.name }}: {{ field.value }}\n{% endfor %}", + "url": "https://localhost:7777", + "query": null, + "inputField": [ + "input" + ], + "pathToEmbeddings": [ + "data" + ], + "embeddingObject": [ + "embedding" + ], + "inputType": "text" + } + }, + "searchCutoffMs": null + } + "###); + + let (response, code) = server.get_task(settings_update_uid).await; + meili_snap::snapshot!(code, @"200 OK"); + meili_snap::snapshot!(meili_snap::json_string!(response["details"]), @r###" + { + "embedders": { + "default": { + "source": "rest", + "apiKey": "My suXXXXXX...", + "url": "https://localhost:7777" + } + } + } + "###); +} + #[actix_rt::test] async fn error_update_settings_unknown_field() { let server = Server::new().await; From c41e1274dc7aa8360d7ed0a64e5e04c2c73c95fd Mon Sep 17 00:00:00 2001 From: Tamo Date: Tue, 26 Mar 2024 15:56:43 +0100 Subject: [PATCH 41/86] push and test the search queue datastructure --- meilisearch-types/src/error.rs | 1 + meilisearch/src/error.rs | 6 ++ meilisearch/src/lib.rs | 1 + meilisearch/src/search_queue.rs | 88 +++++++++++++++++ meilisearch/tests/search/mod.rs | 1 + meilisearch/tests/search/search_queue.rs | 120 +++++++++++++++++++++++ 6 files changed, 217 insertions(+) create mode 100644 meilisearch/src/search_queue.rs create mode 100644 meilisearch/tests/search/search_queue.rs diff --git a/meilisearch-types/src/error.rs b/meilisearch-types/src/error.rs index aed77411a..fe0d75dae 100644 --- a/meilisearch-types/src/error.rs +++ b/meilisearch-types/src/error.rs @@ -305,6 +305,7 @@ MissingSwapIndexes , InvalidRequest , BAD_REQUEST ; MissingTaskFilters , InvalidRequest , BAD_REQUEST ; NoSpaceLeftOnDevice , System , UNPROCESSABLE_ENTITY; PayloadTooLarge , InvalidRequest , PAYLOAD_TOO_LARGE ; +TooManySearchRequests , System , SERVICE_UNAVAILABLE ; TaskNotFound , InvalidRequest , NOT_FOUND ; TooManyOpenFiles , System , UNPROCESSABLE_ENTITY ; TooManyVectors , InvalidRequest , BAD_REQUEST ; diff --git a/meilisearch/src/error.rs b/meilisearch/src/error.rs index a8351fd1f..48c44c12d 100644 --- a/meilisearch/src/error.rs +++ b/meilisearch/src/error.rs @@ -29,6 +29,10 @@ pub enum MeilisearchHttpError { InvalidExpression(&'static [&'static str], Value), #[error("A {0} payload is missing.")] MissingPayload(PayloadType), + #[error("Too many search requests running at the same time: {0}. Retry after {1:?}.")] + TooManySearchRequests(usize, std::time::Duration), + #[error("Internal error: Search limiter is down")] + SearchLimiterIsDown, #[error("The provided payload reached the size limit. The maximum accepted payload size is {}.", Byte::from_bytes(*.0 as u64).get_appropriate_unit(true))] PayloadTooLarge(usize), #[error("Two indexes must be given for each swap. The list `[{}]` contains {} indexes.", @@ -69,6 +73,8 @@ impl ErrorCode for MeilisearchHttpError { MeilisearchHttpError::EmptyFilter => Code::InvalidDocumentFilter, MeilisearchHttpError::InvalidExpression(_, _) => Code::InvalidSearchFilter, MeilisearchHttpError::PayloadTooLarge(_) => Code::PayloadTooLarge, + MeilisearchHttpError::TooManySearchRequests(_, _) => Code::TooManySearchRequests, + MeilisearchHttpError::SearchLimiterIsDown => Code::Internal, MeilisearchHttpError::SwapIndexPayloadWrongLength(_) => Code::InvalidSwapIndexes, MeilisearchHttpError::IndexUid(e) => e.error_code(), MeilisearchHttpError::SerdeJson(_) => Code::Internal, diff --git a/meilisearch/src/lib.rs b/meilisearch/src/lib.rs index 820f1ae42..bfe25b5a7 100644 --- a/meilisearch/src/lib.rs +++ b/meilisearch/src/lib.rs @@ -9,6 +9,7 @@ pub mod middleware; pub mod option; pub mod routes; pub mod search; +pub mod search_queue; use std::fs::File; use std::io::{BufReader, BufWriter}; diff --git a/meilisearch/src/search_queue.rs b/meilisearch/src/search_queue.rs new file mode 100644 index 000000000..1f3cda1a2 --- /dev/null +++ b/meilisearch/src/search_queue.rs @@ -0,0 +1,88 @@ +use std::time::Duration; + +use rand::{rngs::StdRng, Rng, SeedableRng}; +use tokio::sync::{mpsc, oneshot}; + +use crate::error::MeilisearchHttpError; + +#[derive(Debug)] +pub struct SearchQueue { + sender: mpsc::Sender>, + capacity: usize, +} + +#[derive(Debug)] +pub struct Permit { + sender: mpsc::Sender<()>, +} + +impl Drop for Permit { + fn drop(&mut self) { + // if the channel is closed then the whole instance is down + let _ = futures::executor::block_on(self.sender.send(())); + } +} + +impl SearchQueue { + pub fn new(capacity: usize, paralellism: usize) -> Self { + // We can make the search requests wait until we're available. + // they're going to wait anyway right after, so let's not allocate any + // RAM by keeping a capacity of 1. + let (sender, receiver) = mpsc::channel(1); + tokio::task::spawn(Self::run(capacity, paralellism, receiver)); + Self { sender, capacity } + } + + async fn run( + capacity: usize, + parallelism: usize, + mut receive_new_searches: mpsc::Receiver>, + ) { + let mut queue: Vec> = Default::default(); + let mut rng: StdRng = StdRng::from_entropy(); + let mut searches_running = 0; + // by having a capacity of parallelism we ensures that every time a search finish it can release its RAM asap + let (sender, mut search_finished) = mpsc::channel(parallelism); + + loop { + tokio::select! { + search_request = receive_new_searches.recv() => { + let search_request = search_request.unwrap(); + println!("queue contains {} elements and already running {}", queue.len(), searches_running); + if searches_running < parallelism && queue.is_empty() { + println!("We can process the search straight away"); + searches_running += 1; + // if the search requests die it's not a hard error on our side + let _ = search_request.send(Permit { sender: sender.clone() }); + continue; + } + if queue.len() >= capacity { + println!("we're above capacity, dropping a random request"); + let remove = rng.gen_range(0..queue.len()); + let thing = queue.swap_remove(remove); // this will drop the channel and notify the search that it won't be processed + drop(thing); + } + println!("pushed a new search request to the queue {}", queue.len()); + queue.push(search_request); + }, + _ = search_finished.recv() => { + searches_running = searches_running.saturating_sub(1); + if !queue.is_empty() { + println!("processed an element in the queue"); + let remove = rng.gen_range(0..queue.len()); + let channel = queue.swap_remove(remove); + let _ = channel.send(Permit { sender: sender.clone() }); + } + }, + } + } + } + + pub async fn register_search(&self) -> Result { + let (sender, receiver) = oneshot::channel(); + self.sender.send(sender).await.map_err(|_| MeilisearchHttpError::SearchLimiterIsDown)?; + receiver.await.map_err(|_| { + MeilisearchHttpError::TooManySearchRequests(self.capacity, Duration::from_secs(10)) + }) + } +} diff --git a/meilisearch/tests/search/mod.rs b/meilisearch/tests/search/mod.rs index 88470187a..e5925d77e 100644 --- a/meilisearch/tests/search/mod.rs +++ b/meilisearch/tests/search/mod.rs @@ -10,6 +10,7 @@ mod hybrid; mod multi; mod pagination; mod restrict_searchable; +mod search_queue; use once_cell::sync::Lazy; diff --git a/meilisearch/tests/search/search_queue.rs b/meilisearch/tests/search/search_queue.rs new file mode 100644 index 000000000..078e8c152 --- /dev/null +++ b/meilisearch/tests/search/search_queue.rs @@ -0,0 +1,120 @@ +use std::{sync::Arc, time::Duration}; + +use meili_snap::snapshot; +use meilisearch::search_queue::SearchQueue; + +#[actix_rt::test] +async fn search_queue_register() { + let queue = SearchQueue::new(4, 2); + + // First, use all the cores + let permit1 = tokio::time::timeout(Duration::from_secs(1), queue.register_search()) + .await + .expect("I should get a permit straight away") + .unwrap(); + let _permit2 = tokio::time::timeout(Duration::from_secs(1), queue.register_search()) + .await + .expect("I should get a permit straight away") + .unwrap(); + + // If we free one spot we should be able to register one new search + drop(permit1); + + let permit3 = tokio::time::timeout(Duration::from_secs(1), queue.register_search()) + .await + .expect("I should get a permit straight away") + .unwrap(); + + // And again + drop(permit3); + + let _permit4 = tokio::time::timeout(Duration::from_secs(1), queue.register_search()) + .await + .expect("I should get a permit straight away") + .unwrap(); +} + +#[actix_rt::test] +async fn search_queue_wait_till_cores_available() { + let queue = Arc::new(SearchQueue::new(4, 1)); + + // First, use all the cores + let permit1 = tokio::time::timeout(Duration::from_secs(1), queue.register_search()) + .await + .expect("I should get a permit straight away") + .unwrap(); + + let ret = tokio::time::timeout(Duration::from_secs(1), queue.register_search()).await; + assert!(ret.is_err(), "The capacity is full, we should not get a permit"); + + let q = queue.clone(); + let task = tokio::task::spawn(async move { q.register_search().await }); + + // after dropping a permit the previous task should be able to finish + drop(permit1); + let _permit2 = tokio::time::timeout(Duration::from_secs(1), task) + .await + .expect("I should get a permit straight away") + .unwrap(); +} + +#[actix_rt::test] +async fn search_queue_refuse_search_requests() { + let queue = Arc::new(SearchQueue::new(1, 1)); + + // First, use the whole capacity of the + let _permit1 = tokio::time::timeout(Duration::from_secs(1), queue.register_search()) + .await + .expect("I should get a permit straight away") + .unwrap(); + + let q = queue.clone(); + let permit2 = tokio::task::spawn(async move { q.register_search().await }); + + // Here the queue is full. By registering two new search requests the permit 2 and 3 should be thrown out + let q = queue.clone(); + let _permit3 = tokio::task::spawn(async move { q.register_search().await }); + + let permit2 = tokio::time::timeout(Duration::from_secs(1), permit2) + .await + .expect("I should get a result straight away") + .unwrap(); // task should end successfully + + snapshot!(permit2.unwrap_err(), @"Too many search requests running at the same time: 1. Retry after 10s."); +} + +#[actix_rt::test] +async fn search_request_crashes_while_holding_permits() { + let queue = Arc::new(SearchQueue::new(1, 1)); + + let (send, recv) = tokio::sync::oneshot::channel(); + + // This first request take a cpu + let q = queue.clone(); + tokio::task::spawn(async move { + let _permit = q.register_search().await.unwrap(); + recv.await.unwrap(); + panic!("oops an unexpected crash happened") + }); + + // This second request waits in the queue till the first request finishes + let q = queue.clone(); + let task = tokio::task::spawn(async move { + let _permit = q.register_search().await.unwrap(); + }); + + // By sending something in the channel the request holding a CPU will panic and should lose its permit + send.send(()).unwrap(); + + // Then the second request should be able to process and finishes correctly without panic + tokio::time::timeout(Duration::from_secs(1), task) + .await + .expect("I should get a permit straight away") + .unwrap(); + + // I should even be able to take second permit here + let _permit1 = tokio::time::timeout(Duration::from_secs(1), queue.register_search()) + .await + .expect("I should get a permit straight away") + .unwrap(); +} From 3f23fbb46d410eaad6109aee4d849dc7f3234f11 Mon Sep 17 00:00:00 2001 From: Tamo Date: Tue, 26 Mar 2024 16:43:40 +0100 Subject: [PATCH 42/86] create the experimental CLI argument --- meilisearch/src/analytics/segment_analytics.rs | 3 +++ meilisearch/src/option.rs | 16 ++++++++++++++++ 2 files changed, 19 insertions(+) diff --git a/meilisearch/src/analytics/segment_analytics.rs b/meilisearch/src/analytics/segment_analytics.rs index 99298bd43..b334f651d 100644 --- a/meilisearch/src/analytics/segment_analytics.rs +++ b/meilisearch/src/analytics/segment_analytics.rs @@ -252,6 +252,7 @@ impl super::Analytics for SegmentAnalytics { struct Infos { env: String, experimental_enable_metrics: bool, + experimental_search_queue_size: usize, experimental_logs_mode: LogMode, experimental_replication_parameters: bool, experimental_enable_logs_route: bool, @@ -293,6 +294,7 @@ impl From for Infos { let Opt { db_path, experimental_enable_metrics, + experimental_search_queue_size, experimental_logs_mode, experimental_replication_parameters, experimental_enable_logs_route, @@ -342,6 +344,7 @@ impl From for Infos { Self { env, experimental_enable_metrics, + experimental_search_queue_size, experimental_logs_mode, experimental_replication_parameters, experimental_enable_logs_route, diff --git a/meilisearch/src/option.rs b/meilisearch/src/option.rs index 43bf2c62c..a9b8578bb 100644 --- a/meilisearch/src/option.rs +++ b/meilisearch/src/option.rs @@ -54,6 +54,7 @@ const MEILI_EXPERIMENTAL_LOGS_MODE: &str = "MEILI_EXPERIMENTAL_LOGS_MODE"; const MEILI_EXPERIMENTAL_REPLICATION_PARAMETERS: &str = "MEILI_EXPERIMENTAL_REPLICATION_PARAMETERS"; const MEILI_EXPERIMENTAL_ENABLE_LOGS_ROUTE: &str = "MEILI_EXPERIMENTAL_ENABLE_LOGS_ROUTE"; const MEILI_EXPERIMENTAL_ENABLE_METRICS: &str = "MEILI_EXPERIMENTAL_ENABLE_METRICS"; +const MEILI_EXPERIMENTAL_SEARCH_QUEUE_SIZE: &str = "MEILI_EXPERIMENTAL_SEARCH_QUEUE_SIZE"; const MEILI_EXPERIMENTAL_REDUCE_INDEXING_MEMORY_USAGE: &str = "MEILI_EXPERIMENTAL_REDUCE_INDEXING_MEMORY_USAGE"; const MEILI_EXPERIMENTAL_MAX_NUMBER_OF_BATCHED_TASKS: &str = @@ -344,6 +345,16 @@ pub struct Opt { #[serde(default)] pub experimental_enable_metrics: bool, + /// TODO: Update the discussion link + /// Experimental search queue size. For more information, see: + /// + /// Lets you customize the size of the search queue. Meilisearch processes your search requests as fast as possible but once the + /// queue is full it starts returning HTTP 503, Service Unavailable. + /// The default value is 1000. + #[clap(long, env = MEILI_EXPERIMENTAL_SEARCH_QUEUE_SIZE, default_value_t = 1000)] + #[serde(default)] + pub experimental_search_queue_size: usize, + /// Experimental logs mode feature. For more information, see: /// /// Change the mode of the logs on the console. @@ -473,6 +484,7 @@ impl Opt { #[cfg(feature = "analytics")] no_analytics, experimental_enable_metrics, + experimental_search_queue_size, experimental_logs_mode, experimental_enable_logs_route, experimental_replication_parameters, @@ -532,6 +544,10 @@ impl Opt { MEILI_EXPERIMENTAL_ENABLE_METRICS, experimental_enable_metrics.to_string(), ); + export_to_env_if_not_present( + MEILI_EXPERIMENTAL_SEARCH_QUEUE_SIZE, + experimental_search_queue_size.to_string(), + ); export_to_env_if_not_present( MEILI_EXPERIMENTAL_LOGS_MODE, experimental_logs_mode.to_string(), From e433fd53e6fe1a5acbfa002bcc33b7c2e9468d1f Mon Sep 17 00:00:00 2001 From: Tamo Date: Tue, 26 Mar 2024 17:28:03 +0100 Subject: [PATCH 43/86] rename the method to get a permit and use it in all search requests --- meilisearch/src/lib.rs | 9 ++++- .../src/routes/indexes/facet_search.rs | 3 ++ meilisearch/src/routes/indexes/search.rs | 5 +++ meilisearch/src/routes/multi_search.rs | 6 ++++ meilisearch/src/search_queue.rs | 14 ++++---- meilisearch/tests/search/search_queue.rs | 36 +++++++++---------- 6 files changed, 47 insertions(+), 26 deletions(-) diff --git a/meilisearch/src/lib.rs b/meilisearch/src/lib.rs index bfe25b5a7..bb7562c85 100644 --- a/meilisearch/src/lib.rs +++ b/meilisearch/src/lib.rs @@ -13,9 +13,10 @@ pub mod search_queue; use std::fs::File; use std::io::{BufReader, BufWriter}; +use std::num::NonZeroUsize; use std::path::Path; use std::sync::Arc; -use std::thread; +use std::thread::{self, available_parallelism}; use std::time::Duration; use actix_cors::Cors; @@ -39,6 +40,7 @@ use meilisearch_types::versioning::{check_version_file, create_version_file}; use meilisearch_types::{compression, milli, VERSION_FILE_NAME}; pub use option::Opt; use option::ScheduleSnapshot; +use search_queue::SearchQueue; use tracing::{error, info_span}; use tracing_subscriber::filter::Targets; @@ -470,10 +472,15 @@ pub fn configure_data( (logs_route, logs_stderr): (LogRouteHandle, LogStderrHandle), analytics: Arc, ) { + let search_queue = SearchQueue::new( + opt.experimental_search_queue_size, + available_parallelism().unwrap_or(NonZeroUsize::new(2).unwrap()), + ); let http_payload_size_limit = opt.http_payload_size_limit.get_bytes() as usize; config .app_data(index_scheduler) .app_data(auth) + .app_data(web::Data::new(search_queue)) .app_data(web::Data::from(analytics)) .app_data(web::Data::new(logs_route)) .app_data(web::Data::new(logs_stderr)) diff --git a/meilisearch/src/routes/indexes/facet_search.rs b/meilisearch/src/routes/indexes/facet_search.rs index a980fb278..272b8156f 100644 --- a/meilisearch/src/routes/indexes/facet_search.rs +++ b/meilisearch/src/routes/indexes/facet_search.rs @@ -17,6 +17,7 @@ use crate::search::{ DEFAULT_CROP_LENGTH, DEFAULT_CROP_MARKER, DEFAULT_HIGHLIGHT_POST_TAG, DEFAULT_HIGHLIGHT_PRE_TAG, DEFAULT_SEARCH_LIMIT, DEFAULT_SEARCH_OFFSET, }; +use crate::search_queue::SearchQueue; pub fn configure(cfg: &mut web::ServiceConfig) { cfg.service(web::resource("").route(web::post().to(search))); @@ -48,6 +49,7 @@ pub struct FacetSearchQuery { pub async fn search( index_scheduler: GuardedData, Data>, + search_queue: Data, index_uid: web::Path, params: AwebJson, req: HttpRequest, @@ -71,6 +73,7 @@ pub async fn search( let index = index_scheduler.index(&index_uid)?; let features = index_scheduler.features(); + let _permit = search_queue.try_get_search_permit().await?; let search_result = tokio::task::spawn_blocking(move || { perform_facet_search(&index, search_query, facet_query, facet_name, features) }) diff --git a/meilisearch/src/routes/indexes/search.rs b/meilisearch/src/routes/indexes/search.rs index 6a430b6a3..880786138 100644 --- a/meilisearch/src/routes/indexes/search.rs +++ b/meilisearch/src/routes/indexes/search.rs @@ -23,6 +23,7 @@ use crate::search::{ DEFAULT_CROP_LENGTH, DEFAULT_CROP_MARKER, DEFAULT_HIGHLIGHT_POST_TAG, DEFAULT_HIGHLIGHT_PRE_TAG, DEFAULT_SEARCH_LIMIT, DEFAULT_SEARCH_OFFSET, DEFAULT_SEMANTIC_RATIO, }; +use crate::search_queue::SearchQueue; pub fn configure(cfg: &mut web::ServiceConfig) { cfg.service( @@ -182,6 +183,7 @@ fn fix_sort_query_parameters(sort_query: &str) -> Vec { pub async fn search_with_url_query( index_scheduler: GuardedData, Data>, + search_queue: web::Data, index_uid: web::Path, params: AwebQueryParameter, req: HttpRequest, @@ -204,6 +206,7 @@ pub async fn search_with_url_query( let distribution = embed(&mut query, index_scheduler.get_ref(), &index).await?; + let _permit = search_queue.try_get_search_permit().await?; let search_result = tokio::task::spawn_blocking(move || perform_search(&index, query, features, distribution)) .await?; @@ -220,6 +223,7 @@ pub async fn search_with_url_query( pub async fn search_with_post( index_scheduler: GuardedData, Data>, + search_queue: web::Data, index_uid: web::Path, params: AwebJson, req: HttpRequest, @@ -243,6 +247,7 @@ pub async fn search_with_post( let distribution = embed(&mut query, index_scheduler.get_ref(), &index).await?; + let _permit = search_queue.try_get_search_permit().await?; let search_result = tokio::task::spawn_blocking(move || perform_search(&index, query, features, distribution)) .await?; diff --git a/meilisearch/src/routes/multi_search.rs b/meilisearch/src/routes/multi_search.rs index 86aa58e70..55f633e97 100644 --- a/meilisearch/src/routes/multi_search.rs +++ b/meilisearch/src/routes/multi_search.rs @@ -17,6 +17,7 @@ use crate::routes::indexes::search::embed; use crate::search::{ add_search_rules, perform_search, SearchQueryWithIndex, SearchResultWithIndex, }; +use crate::search_queue::SearchQueue; pub fn configure(cfg: &mut web::ServiceConfig) { cfg.service(web::resource("").route(web::post().to(SeqHandler(multi_search_with_post)))); @@ -35,6 +36,7 @@ pub struct SearchQueries { pub async fn multi_search_with_post( index_scheduler: GuardedData, Data>, + search_queue: Data, params: AwebJson, req: HttpRequest, analytics: web::Data, @@ -44,6 +46,10 @@ pub async fn multi_search_with_post( let mut multi_aggregate = MultiSearchAggregator::from_queries(&queries, &req); let features = index_scheduler.features(); + // Since we don't want to process half of the search requests and then get a permit refused + // we're going to get one permit for the whole duration of the multi-search request. + let _permit = search_queue.try_get_search_permit().await?; + // Explicitly expect a `(ResponseError, usize)` for the error type rather than `ResponseError` only, // so that `?` doesn't work if it doesn't use `with_index`, ensuring that it is not forgotten in case of code // changes. diff --git a/meilisearch/src/search_queue.rs b/meilisearch/src/search_queue.rs index 1f3cda1a2..570394e34 100644 --- a/meilisearch/src/search_queue.rs +++ b/meilisearch/src/search_queue.rs @@ -1,4 +1,4 @@ -use std::time::Duration; +use std::{num::NonZeroUsize, time::Duration}; use rand::{rngs::StdRng, Rng, SeedableRng}; use tokio::sync::{mpsc, oneshot}; @@ -24,7 +24,7 @@ impl Drop for Permit { } impl SearchQueue { - pub fn new(capacity: usize, paralellism: usize) -> Self { + pub fn new(capacity: usize, paralellism: NonZeroUsize) -> Self { // We can make the search requests wait until we're available. // they're going to wait anyway right after, so let's not allocate any // RAM by keeping a capacity of 1. @@ -35,21 +35,21 @@ impl SearchQueue { async fn run( capacity: usize, - parallelism: usize, + parallelism: NonZeroUsize, mut receive_new_searches: mpsc::Receiver>, ) { let mut queue: Vec> = Default::default(); let mut rng: StdRng = StdRng::from_entropy(); - let mut searches_running = 0; + let mut searches_running: usize = 0; // by having a capacity of parallelism we ensures that every time a search finish it can release its RAM asap - let (sender, mut search_finished) = mpsc::channel(parallelism); + let (sender, mut search_finished) = mpsc::channel(parallelism.into()); loop { tokio::select! { search_request = receive_new_searches.recv() => { let search_request = search_request.unwrap(); println!("queue contains {} elements and already running {}", queue.len(), searches_running); - if searches_running < parallelism && queue.is_empty() { + if searches_running < usize::from(parallelism) && queue.is_empty() { println!("We can process the search straight away"); searches_running += 1; // if the search requests die it's not a hard error on our side @@ -78,7 +78,7 @@ impl SearchQueue { } } - pub async fn register_search(&self) -> Result { + pub async fn try_get_search_permit(&self) -> Result { let (sender, receiver) = oneshot::channel(); self.sender.send(sender).await.map_err(|_| MeilisearchHttpError::SearchLimiterIsDown)?; receiver.await.map_err(|_| { diff --git a/meilisearch/tests/search/search_queue.rs b/meilisearch/tests/search/search_queue.rs index 078e8c152..15e62ab6d 100644 --- a/meilisearch/tests/search/search_queue.rs +++ b/meilisearch/tests/search/search_queue.rs @@ -1,18 +1,18 @@ -use std::{sync::Arc, time::Duration}; +use std::{num::NonZeroUsize, sync::Arc, time::Duration}; use meili_snap::snapshot; use meilisearch::search_queue::SearchQueue; #[actix_rt::test] async fn search_queue_register() { - let queue = SearchQueue::new(4, 2); + let queue = SearchQueue::new(4, NonZeroUsize::new(2).unwrap()); // First, use all the cores - let permit1 = tokio::time::timeout(Duration::from_secs(1), queue.register_search()) + let permit1 = tokio::time::timeout(Duration::from_secs(1), queue.try_get_search_permit()) .await .expect("I should get a permit straight away") .unwrap(); - let _permit2 = tokio::time::timeout(Duration::from_secs(1), queue.register_search()) + let _permit2 = tokio::time::timeout(Duration::from_secs(1), queue.try_get_search_permit()) .await .expect("I should get a permit straight away") .unwrap(); @@ -20,7 +20,7 @@ async fn search_queue_register() { // If we free one spot we should be able to register one new search drop(permit1); - let permit3 = tokio::time::timeout(Duration::from_secs(1), queue.register_search()) + let permit3 = tokio::time::timeout(Duration::from_secs(1), queue.try_get_search_permit()) .await .expect("I should get a permit straight away") .unwrap(); @@ -28,7 +28,7 @@ async fn search_queue_register() { // And again drop(permit3); - let _permit4 = tokio::time::timeout(Duration::from_secs(1), queue.register_search()) + let _permit4 = tokio::time::timeout(Duration::from_secs(1), queue.try_get_search_permit()) .await .expect("I should get a permit straight away") .unwrap(); @@ -36,19 +36,19 @@ async fn search_queue_register() { #[actix_rt::test] async fn search_queue_wait_till_cores_available() { - let queue = Arc::new(SearchQueue::new(4, 1)); + let queue = Arc::new(SearchQueue::new(4, NonZeroUsize::new(1).unwrap())); // First, use all the cores - let permit1 = tokio::time::timeout(Duration::from_secs(1), queue.register_search()) + let permit1 = tokio::time::timeout(Duration::from_secs(1), queue.try_get_search_permit()) .await .expect("I should get a permit straight away") .unwrap(); - let ret = tokio::time::timeout(Duration::from_secs(1), queue.register_search()).await; + let ret = tokio::time::timeout(Duration::from_secs(1), queue.try_get_search_permit()).await; assert!(ret.is_err(), "The capacity is full, we should not get a permit"); let q = queue.clone(); - let task = tokio::task::spawn(async move { q.register_search().await }); + let task = tokio::task::spawn(async move { q.try_get_search_permit().await }); // after dropping a permit the previous task should be able to finish drop(permit1); @@ -60,20 +60,20 @@ async fn search_queue_wait_till_cores_available() { #[actix_rt::test] async fn search_queue_refuse_search_requests() { - let queue = Arc::new(SearchQueue::new(1, 1)); + let queue = Arc::new(SearchQueue::new(1, NonZeroUsize::new(1).unwrap())); // First, use the whole capacity of the - let _permit1 = tokio::time::timeout(Duration::from_secs(1), queue.register_search()) + let _permit1 = tokio::time::timeout(Duration::from_secs(1), queue.try_get_search_permit()) .await .expect("I should get a permit straight away") .unwrap(); let q = queue.clone(); - let permit2 = tokio::task::spawn(async move { q.register_search().await }); + let permit2 = tokio::task::spawn(async move { q.try_get_search_permit().await }); // Here the queue is full. By registering two new search requests the permit 2 and 3 should be thrown out let q = queue.clone(); - let _permit3 = tokio::task::spawn(async move { q.register_search().await }); + let _permit3 = tokio::task::spawn(async move { q.try_get_search_permit().await }); let permit2 = tokio::time::timeout(Duration::from_secs(1), permit2) .await @@ -85,14 +85,14 @@ async fn search_queue_refuse_search_requests() { #[actix_rt::test] async fn search_request_crashes_while_holding_permits() { - let queue = Arc::new(SearchQueue::new(1, 1)); + let queue = Arc::new(SearchQueue::new(1, NonZeroUsize::new(1).unwrap())); let (send, recv) = tokio::sync::oneshot::channel(); // This first request take a cpu let q = queue.clone(); tokio::task::spawn(async move { - let _permit = q.register_search().await.unwrap(); + let _permit = q.try_get_search_permit().await.unwrap(); recv.await.unwrap(); panic!("oops an unexpected crash happened") }); @@ -100,7 +100,7 @@ async fn search_request_crashes_while_holding_permits() { // This second request waits in the queue till the first request finishes let q = queue.clone(); let task = tokio::task::spawn(async move { - let _permit = q.register_search().await.unwrap(); + let _permit = q.try_get_search_permit().await.unwrap(); }); // By sending something in the channel the request holding a CPU will panic and should lose its permit @@ -113,7 +113,7 @@ async fn search_request_crashes_while_holding_permits() { .unwrap(); // I should even be able to take second permit here - let _permit1 = tokio::time::timeout(Duration::from_secs(1), queue.register_search()) + let _permit1 = tokio::time::timeout(Duration::from_secs(1), queue.try_get_search_permit()) .await .expect("I should get a permit straight away") .unwrap(); From e4a3e603b36808e10ad732ed1c0ae2acf9cd49c5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Cl=C3=A9ment=20Renault?= Date: Tue, 26 Mar 2024 17:31:56 +0100 Subject: [PATCH 44/86] Expose a first working version of the negative keyword --- milli/src/search/new/mod.rs | 22 ++++++++++++++++++- .../src/search/new/query_term/parse_query.rs | 21 ++++++++++++++---- 2 files changed, 38 insertions(+), 5 deletions(-) diff --git a/milli/src/search/new/mod.rs b/milli/src/search/new/mod.rs index ad996f363..ec83b84d1 100644 --- a/milli/src/search/new/mod.rs +++ b/milli/src/search/new/mod.rs @@ -209,6 +209,20 @@ fn resolve_universe( ) } +#[tracing::instrument(level = "trace", skip_all, target = "search")] +fn resolve_negative_words( + ctx: &mut SearchContext, + negative_words: &[Word], +) -> Result { + let mut negative_bitmap = RoaringBitmap::new(); + for &word in negative_words { + if let Some(bitmap) = ctx.word_docids(word)? { + negative_bitmap |= bitmap; + } + } + Ok(negative_bitmap) +} + /// Return the list of initialised ranking rules to be used for a placeholder search. fn get_ranking_rules_for_placeholder_search<'ctx>( ctx: &SearchContext<'ctx>, @@ -620,7 +634,12 @@ pub fn execute_search( let tokens = tokenizer.tokenize(query); drop(entered); - let query_terms = located_query_terms_from_tokens(ctx, tokens, words_limit)?; + let (query_terms, negative_words) = + located_query_terms_from_tokens(ctx, tokens, words_limit)?; + + let ignored_documents = resolve_negative_words(ctx, &negative_words)?; + universe -= ignored_documents; + if query_terms.is_empty() { // Do a placeholder search instead None @@ -630,6 +649,7 @@ pub fn execute_search( } else { None }; + let bucket_sort_output = if let Some(query_terms) = query_terms { let (graph, new_located_query_terms) = QueryGraph::from_query(ctx, &query_terms)?; located_query_terms = Some(new_located_query_terms); diff --git a/milli/src/search/new/query_term/parse_query.rs b/milli/src/search/new/query_term/parse_query.rs index ea997a41a..b23cb2426 100644 --- a/milli/src/search/new/query_term/parse_query.rs +++ b/milli/src/search/new/query_term/parse_query.rs @@ -6,6 +6,7 @@ use charabia::{SeparatorKind, TokenKind}; use super::compute_derivations::partially_initialized_term_from_word; use super::{LocatedQueryTerm, ZeroTypoTerm}; use crate::search::new::query_term::{Lazy, Phrase, QueryTerm}; +use crate::search::new::Word; use crate::{Result, SearchContext, MAX_WORD_LENGTH}; /// Convert the tokenised search query into a list of located query terms. @@ -14,12 +15,14 @@ pub fn located_query_terms_from_tokens( ctx: &mut SearchContext, query: NormalizedTokenIter, words_limit: Option, -) -> Result> { +) -> Result<(Vec, Vec)> { let nbr_typos = number_of_typos_allowed(ctx)?; let mut located_terms = Vec::new(); let mut phrase: Option = None; + let mut negative_next_token = false; + let mut negative_words = Vec::new(); let parts_limit = words_limit.unwrap_or(usize::MAX); @@ -33,7 +36,7 @@ pub fn located_query_terms_from_tokens( } // early return if word limit is exceeded if located_terms.len() >= parts_limit { - return Ok(located_terms); + return Ok((located_terms, negative_words)); } match token.kind { @@ -46,6 +49,11 @@ pub fn located_query_terms_from_tokens( // 3. if the word is the last token of the query we push it as a prefix word. if let Some(phrase) = &mut phrase { phrase.push_word(ctx, &token, position) + } else if negative_next_token { + let word = token.lemma().to_string(); + let word = Word::Original(ctx.word_interner.insert(word)); + negative_words.push(word); + negative_next_token = false; } else if peekable.peek().is_some() { match token.kind { TokenKind::Word => { @@ -63,7 +71,7 @@ pub fn located_query_terms_from_tokens( }; located_terms.push(located_term); } - TokenKind::StopWord | TokenKind::Separator(_) | TokenKind::Unknown => {} + TokenKind::StopWord | TokenKind::Separator(_) | TokenKind::Unknown => (), } } else { let word = token.lemma(); @@ -122,6 +130,10 @@ pub fn located_query_terms_from_tokens( // Start new phrase if the token ends with an opening quote (quote_count % 2 == 1).then_some(PhraseBuilder::empty()) }; + + if phrase.is_none() && token.lemma() == "-" { + negative_next_token = true; + } } _ => (), } @@ -134,7 +146,7 @@ pub fn located_query_terms_from_tokens( } } - Ok(located_terms) + Ok((located_terms, negative_words)) } pub fn number_of_typos_allowed<'ctx>( @@ -317,6 +329,7 @@ mod tests { // panics with `attempt to add with overflow` before let located_query_terms = located_query_terms_from_tokens(&mut ctx, tokens, None)?; assert!(located_query_terms.is_empty()); + Ok(()) } } From 1da9e0f246c9dc0b7f12cc745dfb9711a9a039bc Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Cl=C3=A9ment=20Renault?= Date: Tue, 26 Mar 2024 17:42:09 +0100 Subject: [PATCH 45/86] Better support space around the negative operator (-) --- milli/src/search/new/query_term/parse_query.rs | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/milli/src/search/new/query_term/parse_query.rs b/milli/src/search/new/query_term/parse_query.rs index b23cb2426..e510595ee 100644 --- a/milli/src/search/new/query_term/parse_query.rs +++ b/milli/src/search/new/query_term/parse_query.rs @@ -21,6 +21,7 @@ pub fn located_query_terms_from_tokens( let mut located_terms = Vec::new(); let mut phrase: Option = None; + let mut encountered_whitespace = true; let mut negative_next_token = false; let mut negative_words = Vec::new(); @@ -34,6 +35,7 @@ pub fn located_query_terms_from_tokens( if token.lemma().is_empty() { continue; } + // early return if word limit is exceeded if located_terms.len() >= parts_limit { return Ok((located_terms, negative_words)); @@ -131,12 +133,14 @@ pub fn located_query_terms_from_tokens( (quote_count % 2 == 1).then_some(PhraseBuilder::empty()) }; - if phrase.is_none() && token.lemma() == "-" { - negative_next_token = true; - } + negative_next_token = + phrase.is_none() && token.lemma() == "-" && encountered_whitespace; } _ => (), } + + encountered_whitespace = + token.lemma().chars().last().filter(|c| c.is_whitespace()).is_some(); } // If a quote is never closed, we consider all of the end of the query as a phrase. From e2a1bbae378182897ecb03320e602b6d31e72bbe Mon Sep 17 00:00:00 2001 From: Tamo Date: Tue, 26 Mar 2024 17:53:37 +0100 Subject: [PATCH 46/86] simplify and improve the http error --- meilisearch-types/src/error.rs | 10 +++++++++- meilisearch/src/error.rs | 6 +++--- meilisearch/src/search_queue.rs | 17 +++++------------ 3 files changed, 17 insertions(+), 16 deletions(-) diff --git a/meilisearch-types/src/error.rs b/meilisearch-types/src/error.rs index fe0d75dae..f16097da1 100644 --- a/meilisearch-types/src/error.rs +++ b/meilisearch-types/src/error.rs @@ -2,6 +2,7 @@ use std::{fmt, io}; use actix_web::http::StatusCode; use actix_web::{self as aweb, HttpResponseBuilder}; +use aweb::http::header; use aweb::rt::task::JoinError; use convert_case::Casing; use milli::heed::{Error as HeedError, MdbError}; @@ -56,7 +57,14 @@ where impl aweb::error::ResponseError for ResponseError { fn error_response(&self) -> aweb::HttpResponse { let json = serde_json::to_vec(self).unwrap(); - HttpResponseBuilder::new(self.status_code()).content_type("application/json").body(json) + let mut builder = HttpResponseBuilder::new(self.status_code()); + builder.content_type("application/json"); + + if self.code == StatusCode::SERVICE_UNAVAILABLE { + builder.insert_header((header::RETRY_AFTER, "10")); + } + + builder.body(json) } fn status_code(&self) -> StatusCode { diff --git a/meilisearch/src/error.rs b/meilisearch/src/error.rs index 48c44c12d..bb1156997 100644 --- a/meilisearch/src/error.rs +++ b/meilisearch/src/error.rs @@ -29,8 +29,8 @@ pub enum MeilisearchHttpError { InvalidExpression(&'static [&'static str], Value), #[error("A {0} payload is missing.")] MissingPayload(PayloadType), - #[error("Too many search requests running at the same time: {0}. Retry after {1:?}.")] - TooManySearchRequests(usize, std::time::Duration), + #[error("Too many search requests running at the same time: {0}. Retry after 10s.")] + TooManySearchRequests(usize), #[error("Internal error: Search limiter is down")] SearchLimiterIsDown, #[error("The provided payload reached the size limit. The maximum accepted payload size is {}.", Byte::from_bytes(*.0 as u64).get_appropriate_unit(true))] @@ -73,7 +73,7 @@ impl ErrorCode for MeilisearchHttpError { MeilisearchHttpError::EmptyFilter => Code::InvalidDocumentFilter, MeilisearchHttpError::InvalidExpression(_, _) => Code::InvalidSearchFilter, MeilisearchHttpError::PayloadTooLarge(_) => Code::PayloadTooLarge, - MeilisearchHttpError::TooManySearchRequests(_, _) => Code::TooManySearchRequests, + MeilisearchHttpError::TooManySearchRequests(_) => Code::TooManySearchRequests, MeilisearchHttpError::SearchLimiterIsDown => Code::Internal, MeilisearchHttpError::SwapIndexPayloadWrongLength(_) => Code::InvalidSwapIndexes, MeilisearchHttpError::IndexUid(e) => e.error_code(), diff --git a/meilisearch/src/search_queue.rs b/meilisearch/src/search_queue.rs index 570394e34..b677f81a4 100644 --- a/meilisearch/src/search_queue.rs +++ b/meilisearch/src/search_queue.rs @@ -1,4 +1,4 @@ -use std::{num::NonZeroUsize, time::Duration}; +use std::num::NonZeroUsize; use rand::{rngs::StdRng, Rng, SeedableRng}; use tokio::sync::{mpsc, oneshot}; @@ -25,10 +25,10 @@ impl Drop for Permit { impl SearchQueue { pub fn new(capacity: usize, paralellism: NonZeroUsize) -> Self { - // We can make the search requests wait until we're available. - // they're going to wait anyway right after, so let's not allocate any - // RAM by keeping a capacity of 1. + // Search requests are going to wait until we're available anyway, + // so let's not allocate any RAM and keep a capacity of 1. let (sender, receiver) = mpsc::channel(1); + tokio::task::spawn(Self::run(capacity, paralellism, receiver)); Self { sender, capacity } } @@ -48,27 +48,22 @@ impl SearchQueue { tokio::select! { search_request = receive_new_searches.recv() => { let search_request = search_request.unwrap(); - println!("queue contains {} elements and already running {}", queue.len(), searches_running); if searches_running < usize::from(parallelism) && queue.is_empty() { - println!("We can process the search straight away"); searches_running += 1; // if the search requests die it's not a hard error on our side let _ = search_request.send(Permit { sender: sender.clone() }); continue; } if queue.len() >= capacity { - println!("we're above capacity, dropping a random request"); let remove = rng.gen_range(0..queue.len()); let thing = queue.swap_remove(remove); // this will drop the channel and notify the search that it won't be processed drop(thing); } - println!("pushed a new search request to the queue {}", queue.len()); queue.push(search_request); }, _ = search_finished.recv() => { searches_running = searches_running.saturating_sub(1); if !queue.is_empty() { - println!("processed an element in the queue"); let remove = rng.gen_range(0..queue.len()); let channel = queue.swap_remove(remove); let _ = channel.send(Permit { sender: sender.clone() }); @@ -81,8 +76,6 @@ impl SearchQueue { pub async fn try_get_search_permit(&self) -> Result { let (sender, receiver) = oneshot::channel(); self.sender.send(sender).await.map_err(|_| MeilisearchHttpError::SearchLimiterIsDown)?; - receiver.await.map_err(|_| { - MeilisearchHttpError::TooManySearchRequests(self.capacity, Duration::from_secs(10)) - }) + receiver.await.map_err(|_| MeilisearchHttpError::TooManySearchRequests(self.capacity)) } } From 34262c7a0d9baba51e2ebbc6a1b865c5ab5d2d00 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Cl=C3=A9ment=20Renault?= Date: Tue, 26 Mar 2024 18:01:27 +0100 Subject: [PATCH 47/86] Add analytics for the negative operator --- meilisearch/src/analytics/segment_analytics.rs | 10 ++++++++++ meilisearch/src/search.rs | 6 +++++- milli/src/search/hybrid.rs | 3 +++ milli/src/search/mod.rs | 11 ++++++++++- milli/src/search/new/mod.rs | 5 +++++ 5 files changed, 33 insertions(+), 2 deletions(-) diff --git a/meilisearch/src/analytics/segment_analytics.rs b/meilisearch/src/analytics/segment_analytics.rs index 99298bd43..f20cff5d9 100644 --- a/meilisearch/src/analytics/segment_analytics.rs +++ b/meilisearch/src/analytics/segment_analytics.rs @@ -580,6 +580,7 @@ pub struct SearchAggregator { total_received: usize, total_succeeded: usize, total_degraded: usize, + total_used_negative_operator: usize, time_spent: BinaryHeap, // sort @@ -760,12 +761,16 @@ impl SearchAggregator { facet_distribution: _, facet_stats: _, degraded, + used_negative_operator, } = result; self.total_succeeded = self.total_succeeded.saturating_add(1); if *degraded { self.total_degraded = self.total_degraded.saturating_add(1); } + if *used_negative_operator { + self.total_used_negative_operator = self.total_used_negative_operator.saturating_add(1); + } self.time_spent.push(*processing_time_ms as usize); } @@ -808,6 +813,7 @@ impl SearchAggregator { embedder, hybrid, total_degraded, + total_used_negative_operator, } = other; if self.timestamp.is_none() { @@ -823,6 +829,8 @@ impl SearchAggregator { self.total_received = self.total_received.saturating_add(total_received); self.total_succeeded = self.total_succeeded.saturating_add(total_succeeded); self.total_degraded = self.total_degraded.saturating_add(total_degraded); + self.total_used_negative_operator = + self.total_used_negative_operator.saturating_add(total_used_negative_operator); self.time_spent.append(time_spent); // sort @@ -929,6 +937,7 @@ impl SearchAggregator { embedder, hybrid, total_degraded, + total_used_negative_operator, } = self; if total_received == 0 { @@ -949,6 +958,7 @@ impl SearchAggregator { "total_failed": total_received.saturating_sub(total_succeeded), // just to be sure we never panics "total_received": total_received, "total_degraded": total_degraded, + "total_used_negative_operator": total_used_negative_operator, }, "sort": { "with_geoPoint": sort_with_geo_point, diff --git a/meilisearch/src/search.rs b/meilisearch/src/search.rs index 3c00ca802..db58c6102 100644 --- a/meilisearch/src/search.rs +++ b/meilisearch/src/search.rs @@ -324,9 +324,11 @@ pub struct SearchResult { #[serde(skip_serializing_if = "Option::is_none")] pub facet_stats: Option>, - // This information is only used for analytics purposes + // These fields are only used for analytics purposes #[serde(skip)] pub degraded: bool, + #[serde(skip)] + pub used_negative_operator: bool, } #[derive(Serialize, Debug, Clone, PartialEq)] @@ -512,6 +514,7 @@ pub fn perform_search( candidates, document_scores, degraded, + used_negative_operator, .. } = match &query.hybrid { Some(hybrid) => match *hybrid.semantic_ratio { @@ -717,6 +720,7 @@ pub fn perform_search( facet_distribution, facet_stats, degraded, + used_negative_operator, }; Ok(result) } diff --git a/milli/src/search/hybrid.rs b/milli/src/search/hybrid.rs index a8b7f0fcf..2a6d9f7a5 100644 --- a/milli/src/search/hybrid.rs +++ b/milli/src/search/hybrid.rs @@ -11,6 +11,7 @@ struct ScoreWithRatioResult { candidates: RoaringBitmap, document_scores: Vec<(u32, ScoreWithRatio)>, degraded: bool, + used_negative_operator: bool, } type ScoreWithRatio = (Vec, f32); @@ -78,6 +79,7 @@ impl ScoreWithRatioResult { candidates: results.candidates, document_scores, degraded: results.degraded, + used_negative_operator: results.used_negative_operator, } } @@ -113,6 +115,7 @@ impl ScoreWithRatioResult { documents_ids, document_scores, degraded: left.degraded | right.degraded, + used_negative_operator: left.used_negative_operator | right.used_negative_operator, } } } diff --git a/milli/src/search/mod.rs b/milli/src/search/mod.rs index b3dd0c091..3c709a647 100644 --- a/milli/src/search/mod.rs +++ b/milli/src/search/mod.rs @@ -183,6 +183,7 @@ impl<'a> Search<'a> { documents_ids, document_scores, degraded, + used_negative_operator, } = match self.vector.as_ref() { Some(vector) => execute_vector_search( &mut ctx, @@ -221,7 +222,14 @@ impl<'a> Search<'a> { None => MatchingWords::default(), }; - Ok(SearchResult { matching_words, candidates, document_scores, documents_ids, degraded }) + Ok(SearchResult { + matching_words, + candidates, + document_scores, + documents_ids, + degraded, + used_negative_operator, + }) } } @@ -272,6 +280,7 @@ pub struct SearchResult { pub documents_ids: Vec, pub document_scores: Vec>, pub degraded: bool, + pub used_negative_operator: bool, } #[derive(Debug, Clone, Copy, PartialEq, Eq)] diff --git a/milli/src/search/new/mod.rs b/milli/src/search/new/mod.rs index ec83b84d1..99f4562d6 100644 --- a/milli/src/search/new/mod.rs +++ b/milli/src/search/new/mod.rs @@ -571,6 +571,7 @@ pub fn execute_vector_search( documents_ids: docids, located_query_terms: None, degraded, + used_negative_operator: false, }) } @@ -594,6 +595,7 @@ pub fn execute_search( ) -> Result { check_sort_criteria(ctx, sort_criteria.as_ref())?; + let mut used_negative_operator = false; let mut located_query_terms = None; let query_terms = if let Some(query) = query { let span = tracing::trace_span!(target: "search::tokens", "tokenizer_builder"); @@ -636,6 +638,7 @@ pub fn execute_search( let (query_terms, negative_words) = located_query_terms_from_tokens(ctx, tokens, words_limit)?; + used_negative_operator = !negative_words.is_empty(); let ignored_documents = resolve_negative_words(ctx, &negative_words)?; universe -= ignored_documents; @@ -710,6 +713,7 @@ pub fn execute_search( documents_ids: docids, located_query_terms, degraded, + used_negative_operator, }) } @@ -772,4 +776,5 @@ pub struct PartialSearchResult { pub document_scores: Vec>, pub degraded: bool, + pub used_negative_operator: bool, } From e7704f1fc127268a3dfaea6ffff40633a47cf273 Mon Sep 17 00:00:00 2001 From: Tamo Date: Tue, 26 Mar 2024 18:08:59 +0100 Subject: [PATCH 48/86] add a test to ensure we effectively returns a retry-after when the search queue is full --- meilisearch/tests/search/search_queue.rs | 19 ++++++++++++++++--- 1 file changed, 16 insertions(+), 3 deletions(-) diff --git a/meilisearch/tests/search/search_queue.rs b/meilisearch/tests/search/search_queue.rs index 15e62ab6d..8dca2ab31 100644 --- a/meilisearch/tests/search/search_queue.rs +++ b/meilisearch/tests/search/search_queue.rs @@ -1,5 +1,6 @@ use std::{num::NonZeroUsize, sync::Arc, time::Duration}; +use actix_web::ResponseError; use meili_snap::snapshot; use meilisearch::search_queue::SearchQueue; @@ -35,7 +36,7 @@ async fn search_queue_register() { } #[actix_rt::test] -async fn search_queue_wait_till_cores_available() { +async fn wait_till_cores_are_available() { let queue = Arc::new(SearchQueue::new(4, NonZeroUsize::new(1).unwrap())); // First, use all the cores @@ -59,7 +60,7 @@ async fn search_queue_wait_till_cores_available() { } #[actix_rt::test] -async fn search_queue_refuse_search_requests() { +async fn refuse_search_requests_when_queue_is_full() { let queue = Arc::new(SearchQueue::new(1, NonZeroUsize::new(1).unwrap())); // First, use the whole capacity of the @@ -80,7 +81,19 @@ async fn search_queue_refuse_search_requests() { .expect("I should get a result straight away") .unwrap(); // task should end successfully - snapshot!(permit2.unwrap_err(), @"Too many search requests running at the same time: 1. Retry after 10s."); + let err = meilisearch_types::error::ResponseError::from(permit2.unwrap_err()); + let http_response = err.error_response(); + snapshot!(format!("{:?}", http_response.headers()), @r###"HeaderMap { inner: {"retry-after": Value { inner: ["10"] }, "content-type": Value { inner: ["application/json"] }} }"###); + + let err = serde_json::to_string_pretty(&err).unwrap(); + snapshot!(err, @r###" + { + "message": "Too many search requests running at the same time: 1. Retry after 10s.", + "code": "too_many_search_requests", + "type": "system", + "link": "https://docs.meilisearch.com/errors#too_many_search_requests" + } + "###); } #[actix_rt::test] From 8127c9a115a0091654ecab1513f66a151f8819ba Mon Sep 17 00:00:00 2001 From: Tamo Date: Tue, 26 Mar 2024 19:04:39 +0100 Subject: [PATCH 49/86] handle the case of a queue of zero elements --- meilisearch/src/search_queue.rs | 11 +++++-- meilisearch/tests/search/search_queue.rs | 37 ++++++++++++++++++++++++ 2 files changed, 46 insertions(+), 2 deletions(-) diff --git a/meilisearch/src/search_queue.rs b/meilisearch/src/search_queue.rs index b677f81a4..0dd2abf17 100644 --- a/meilisearch/src/search_queue.rs +++ b/meilisearch/src/search_queue.rs @@ -47,14 +47,21 @@ impl SearchQueue { loop { tokio::select! { search_request = receive_new_searches.recv() => { + // this unwrap is safe because we're sure the `SearchQueue` still lives somewhere in actix-web let search_request = search_request.unwrap(); if searches_running < usize::from(parallelism) && queue.is_empty() { searches_running += 1; // if the search requests die it's not a hard error on our side let _ = search_request.send(Permit { sender: sender.clone() }); continue; - } - if queue.len() >= capacity { + } else if capacity == 0 { + // in the very specific case where we have a capacity of zero + // we must refuse the request straight away without going through + // the queue stuff. + drop(search_request); + continue; + + } else if queue.len() >= capacity { let remove = rng.gen_range(0..queue.len()); let thing = queue.swap_remove(remove); // this will drop the channel and notify the search that it won't be processed drop(thing); diff --git a/meilisearch/tests/search/search_queue.rs b/meilisearch/tests/search/search_queue.rs index 8dca2ab31..6bb8d2169 100644 --- a/meilisearch/tests/search/search_queue.rs +++ b/meilisearch/tests/search/search_queue.rs @@ -131,3 +131,40 @@ async fn search_request_crashes_while_holding_permits() { .expect("I should get a permit straight away") .unwrap(); } + +#[actix_rt::test] +async fn works_with_capacity_of_zero() { + let queue = Arc::new(SearchQueue::new(0, NonZeroUsize::new(1).unwrap())); + + // First, use the whole capacity of the + let permit1 = tokio::time::timeout(Duration::from_secs(1), queue.try_get_search_permit()) + .await + .expect("I should get a permit straight away") + .unwrap(); + + // then we should get an error if we try to register a second search request. + let permit2 = tokio::time::timeout(Duration::from_secs(1), queue.try_get_search_permit()) + .await + .expect("I should get a result straight away"); + + let err = meilisearch_types::error::ResponseError::from(permit2.unwrap_err()); + let http_response = err.error_response(); + snapshot!(format!("{:?}", http_response.headers()), @r###"HeaderMap { inner: {"content-type": Value { inner: ["application/json"] }, "retry-after": Value { inner: ["10"] }} }"###); + + let err = serde_json::to_string_pretty(&err).unwrap(); + snapshot!(err, @r###" + { + "message": "Too many search requests running at the same time: 0. Retry after 10s.", + "code": "too_many_search_requests", + "type": "system", + "link": "https://docs.meilisearch.com/errors#too_many_search_requests" + } + "###); + + drop(permit1); + // After dropping the first permit we should be able to get a new permit + let _permit3 = tokio::time::timeout(Duration::from_secs(1), queue.try_get_search_permit()) + .await + .expect("I should get a permit straight away") + .unwrap(); +} From 8f5d9f501ac265ffdd6b17b477843a78f85d0187 Mon Sep 17 00:00:00 2001 From: Tamo Date: Tue, 26 Mar 2024 19:18:32 +0100 Subject: [PATCH 50/86] update the discussion link --- meilisearch/src/option.rs | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/meilisearch/src/option.rs b/meilisearch/src/option.rs index a9b8578bb..651af7336 100644 --- a/meilisearch/src/option.rs +++ b/meilisearch/src/option.rs @@ -345,8 +345,7 @@ pub struct Opt { #[serde(default)] pub experimental_enable_metrics: bool, - /// TODO: Update the discussion link - /// Experimental search queue size. For more information, see: + /// Experimental search queue size. For more information, see: /// /// Lets you customize the size of the search queue. Meilisearch processes your search requests as fast as possible but once the /// queue is full it starts returning HTTP 503, Service Unavailable. From 2e36f069c20fdc8da55894a87d2f4c0237308cf2 Mon Sep 17 00:00:00 2001 From: Tamo Date: Tue, 26 Mar 2024 19:23:55 +0100 Subject: [PATCH 51/86] fmt imports --- meilisearch/src/search_queue.rs | 3 ++- meilisearch/tests/search/search_queue.rs | 4 +++- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/meilisearch/src/search_queue.rs b/meilisearch/src/search_queue.rs index 0dd2abf17..9c0a6704a 100644 --- a/meilisearch/src/search_queue.rs +++ b/meilisearch/src/search_queue.rs @@ -1,6 +1,7 @@ use std::num::NonZeroUsize; -use rand::{rngs::StdRng, Rng, SeedableRng}; +use rand::rngs::StdRng; +use rand::{Rng, SeedableRng}; use tokio::sync::{mpsc, oneshot}; use crate::error::MeilisearchHttpError; diff --git a/meilisearch/tests/search/search_queue.rs b/meilisearch/tests/search/search_queue.rs index 6bb8d2169..717becc3b 100644 --- a/meilisearch/tests/search/search_queue.rs +++ b/meilisearch/tests/search/search_queue.rs @@ -1,4 +1,6 @@ -use std::{num::NonZeroUsize, sync::Arc, time::Duration}; +use std::num::NonZeroUsize; +use std::sync::Arc; +use std::time::Duration; use actix_web::ResponseError; use meili_snap::snapshot; From 55df9daaa0a4ea515e802f3985f162e2b2b54565 Mon Sep 17 00:00:00 2001 From: Tamo Date: Tue, 26 Mar 2024 19:34:55 +0100 Subject: [PATCH 52/86] adds a comment about the safety of an operation --- meilisearch/src/search_queue.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/meilisearch/src/search_queue.rs b/meilisearch/src/search_queue.rs index 9c0a6704a..b661fefb0 100644 --- a/meilisearch/src/search_queue.rs +++ b/meilisearch/src/search_queue.rs @@ -72,6 +72,7 @@ impl SearchQueue { _ = search_finished.recv() => { searches_running = searches_running.saturating_sub(1); if !queue.is_empty() { + // Can't panic: the queue wasn't empty thus the range isn't empty. let remove = rng.gen_range(0..queue.len()); let channel = queue.swap_remove(remove); let _ = channel.send(Permit { sender: sender.clone() }); From 3a1f458139ee4313d4d015635a0c24874dab3097 Mon Sep 17 00:00:00 2001 From: Tamo Date: Tue, 26 Mar 2024 19:42:10 +0100 Subject: [PATCH 53/86] fix a flaky test --- meilisearch/tests/search/search_queue.rs | 16 ++++++++++++++-- 1 file changed, 14 insertions(+), 2 deletions(-) diff --git a/meilisearch/tests/search/search_queue.rs b/meilisearch/tests/search/search_queue.rs index 717becc3b..3b4fbf252 100644 --- a/meilisearch/tests/search/search_queue.rs +++ b/meilisearch/tests/search/search_queue.rs @@ -85,7 +85,13 @@ async fn refuse_search_requests_when_queue_is_full() { let err = meilisearch_types::error::ResponseError::from(permit2.unwrap_err()); let http_response = err.error_response(); - snapshot!(format!("{:?}", http_response.headers()), @r###"HeaderMap { inner: {"retry-after": Value { inner: ["10"] }, "content-type": Value { inner: ["application/json"] }} }"###); + let mut headers: Vec<_> = http_response + .headers() + .iter() + .map(|(name, value)| (name.to_string(), value.to_str().unwrap().to_string())) + .collect(); + headers.sort(); + snapshot!(format!("{headers:?}"), @r###"[("content-type", "application/json"), ("retry-after", "10")]"###); let err = serde_json::to_string_pretty(&err).unwrap(); snapshot!(err, @r###" @@ -151,7 +157,13 @@ async fn works_with_capacity_of_zero() { let err = meilisearch_types::error::ResponseError::from(permit2.unwrap_err()); let http_response = err.error_response(); - snapshot!(format!("{:?}", http_response.headers()), @r###"HeaderMap { inner: {"content-type": Value { inner: ["application/json"] }, "retry-after": Value { inner: ["10"] }} }"###); + let mut headers: Vec<_> = http_response + .headers() + .iter() + .map(|(name, value)| (name.to_string(), value.to_str().unwrap().to_string())) + .collect(); + headers.sort(); + snapshot!(format!("{headers:?}"), @r###"[("content-type", "application/json"), ("retry-after", "10")]"###); let err = serde_json::to_string_pretty(&err).unwrap(); snapshot!(err, @r###" From 087a96d22e2b8d5e4506acc1f3c605abf4f978cd Mon Sep 17 00:00:00 2001 From: Tamo Date: Wed, 27 Mar 2024 11:05:37 +0100 Subject: [PATCH 54/86] fix flaky test --- meilisearch/src/search_queue.rs | 21 ++++++++++++--------- 1 file changed, 12 insertions(+), 9 deletions(-) diff --git a/meilisearch/src/search_queue.rs b/meilisearch/src/search_queue.rs index b661fefb0..ad222cdf8 100644 --- a/meilisearch/src/search_queue.rs +++ b/meilisearch/src/search_queue.rs @@ -47,6 +47,18 @@ impl SearchQueue { loop { tokio::select! { + // biased select because we wants to free up space before trying to register new tasks + biased; + _ = search_finished.recv() => { + searches_running = searches_running.saturating_sub(1); + if !queue.is_empty() { + // Can't panic: the queue wasn't empty thus the range isn't empty. + let remove = rng.gen_range(0..queue.len()); + let channel = queue.swap_remove(remove); + let _ = channel.send(Permit { sender: sender.clone() }); + } + }, + search_request = receive_new_searches.recv() => { // this unwrap is safe because we're sure the `SearchQueue` still lives somewhere in actix-web let search_request = search_request.unwrap(); @@ -69,15 +81,6 @@ impl SearchQueue { } queue.push(search_request); }, - _ = search_finished.recv() => { - searches_running = searches_running.saturating_sub(1); - if !queue.is_empty() { - // Can't panic: the queue wasn't empty thus the range isn't empty. - let remove = rng.gen_range(0..queue.len()); - let channel = queue.swap_remove(remove); - let _ = channel.send(Permit { sender: sender.clone() }); - } - }, } } } From afd1da5642e564b7a896c8529dcfa2d6e161578f Mon Sep 17 00:00:00 2001 From: Louis Dureuil Date: Wed, 27 Mar 2024 11:50:22 +0100 Subject: [PATCH 55/86] Add distribution to all embedders --- milli/src/update/index_documents/mod.rs | 1 + milli/src/update/settings.rs | 5 +++++ milli/src/vector/hf.rs | 20 ++++++++++++-------- milli/src/vector/manual.rs | 10 ++++++++-- milli/src/vector/mod.rs | 3 ++- milli/src/vector/ollama.rs | 7 ++++--- milli/src/vector/openai.rs | 13 +++++++++---- 7 files changed, 41 insertions(+), 18 deletions(-) diff --git a/milli/src/update/index_documents/mod.rs b/milli/src/update/index_documents/mod.rs index 913fbc881..dbacb4002 100644 --- a/milli/src/update/index_documents/mod.rs +++ b/milli/src/update/index_documents/mod.rs @@ -2652,6 +2652,7 @@ mod tests { path_to_embeddings: Setting::NotSet, embedding_object: Setting::NotSet, input_type: Setting::NotSet, + distribution: Setting::NotSet, }), ); settings.set_embedder_settings(embedders); diff --git a/milli/src/update/settings.rs b/milli/src/update/settings.rs index 2b1be9453..9f47768c1 100644 --- a/milli/src/update/settings.rs +++ b/milli/src/update/settings.rs @@ -1146,6 +1146,7 @@ fn validate_prompt( path_to_embeddings, embedding_object, input_type, + distribution, }) => { // validate let template = crate::prompt::Prompt::new(template) @@ -1165,6 +1166,7 @@ fn validate_prompt( path_to_embeddings, embedding_object, input_type, + distribution, })) } new => Ok(new), @@ -1190,6 +1192,7 @@ pub fn validate_embedding_settings( path_to_embeddings, embedding_object, input_type, + distribution, } = settings; if let Some(0) = dimensions.set() { @@ -1221,6 +1224,7 @@ pub fn validate_embedding_settings( path_to_embeddings, embedding_object, input_type, + distribution, })); }; match inferred_source { @@ -1365,6 +1369,7 @@ pub fn validate_embedding_settings( path_to_embeddings, embedding_object, input_type, + distribution, })) } diff --git a/milli/src/vector/hf.rs b/milli/src/vector/hf.rs index e341a553e..725d702ec 100644 --- a/milli/src/vector/hf.rs +++ b/milli/src/vector/hf.rs @@ -33,6 +33,7 @@ enum WeightSource { pub struct EmbedderOptions { pub model: String, pub revision: Option, + pub distribution: Option, } impl EmbedderOptions { @@ -40,6 +41,7 @@ impl EmbedderOptions { Self { model: "BAAI/bge-base-en-v1.5".to_string(), revision: Some("617ca489d9e86b49b8167676d8220688b99db36e".into()), + distribution: None, } } } @@ -193,13 +195,15 @@ impl Embedder { } pub fn distribution(&self) -> Option { - if self.options.model == "BAAI/bge-base-en-v1.5" { - Some(DistributionShift { - current_mean: ordered_float::OrderedFloat(0.85), - current_sigma: ordered_float::OrderedFloat(0.1), - }) - } else { - None - } + self.options.distribution.or_else(|| { + if self.options.model == "BAAI/bge-base-en-v1.5" { + Some(DistributionShift { + current_mean: ordered_float::OrderedFloat(0.85), + current_sigma: ordered_float::OrderedFloat(0.1), + }) + } else { + None + } + }) } } diff --git a/milli/src/vector/manual.rs b/milli/src/vector/manual.rs index 7ed48a251..e5d3689c0 100644 --- a/milli/src/vector/manual.rs +++ b/milli/src/vector/manual.rs @@ -1,19 +1,21 @@ use super::error::EmbedError; -use super::Embeddings; +use super::{DistributionShift, Embeddings}; #[derive(Debug, Clone, Copy)] pub struct Embedder { dimensions: usize, + distribution: Option, } #[derive(Debug, Clone, Hash, PartialEq, Eq, serde::Deserialize, serde::Serialize)] pub struct EmbedderOptions { pub dimensions: usize, + pub distribution: Option, } impl Embedder { pub fn new(options: EmbedderOptions) -> Self { - Self { dimensions: options.dimensions } + Self { dimensions: options.dimensions, distribution: options.distribution } } pub fn embed(&self, mut texts: Vec) -> Result>, EmbedError> { @@ -31,4 +33,8 @@ impl Embedder { ) -> Result>>, EmbedError> { text_chunks.into_iter().map(|prompts| self.embed(prompts)).collect() } + + pub fn distribution(&self) -> Option { + self.distribution + } } diff --git a/milli/src/vector/mod.rs b/milli/src/vector/mod.rs index 8b25de56d..4a3a9920e 100644 --- a/milli/src/vector/mod.rs +++ b/milli/src/vector/mod.rs @@ -1,6 +1,7 @@ use std::collections::HashMap; use std::sync::Arc; +use deserr::{DeserializeError, Deserr}; use ordered_float::OrderedFloat; use serde::{Deserialize, Serialize}; @@ -292,7 +293,7 @@ impl Embedder { Embedder::HuggingFace(embedder) => embedder.distribution(), Embedder::OpenAi(embedder) => embedder.distribution(), Embedder::Ollama(embedder) => embedder.distribution(), - Embedder::UserProvided(_embedder) => None, + Embedder::UserProvided(embedder) => embedder.distribution(), Embedder::Rest(embedder) => embedder.distribution(), } } diff --git a/milli/src/vector/ollama.rs b/milli/src/vector/ollama.rs index 578b6c8e2..cf5030fb4 100644 --- a/milli/src/vector/ollama.rs +++ b/milli/src/vector/ollama.rs @@ -14,11 +14,12 @@ pub struct EmbedderOptions { pub embedding_model: String, pub url: Option, pub api_key: Option, + pub distribution: Option, } impl EmbedderOptions { pub fn with_default_model(api_key: Option, url: Option) -> Self { - Self { embedding_model: "nomic-embed-text".into(), api_key, url } + Self { embedding_model: "nomic-embed-text".into(), api_key, url, distribution: None } } } @@ -27,8 +28,8 @@ impl Embedder { let model = options.embedding_model.as_str(); let rest_embedder = match RestEmbedder::new(RestEmbedderOptions { api_key: options.api_key, - distribution: None, dimensions: None, + distribution: options.distribution, url: options.url.unwrap_or_else(get_ollama_path), query: serde_json::json!({ "model": model, @@ -90,7 +91,7 @@ impl Embedder { } pub fn distribution(&self) -> Option { - None + self.rest_embedder.distribution() } } diff --git a/milli/src/vector/openai.rs b/milli/src/vector/openai.rs index 24e94a9f7..141de486b 100644 --- a/milli/src/vector/openai.rs +++ b/milli/src/vector/openai.rs @@ -11,6 +11,7 @@ pub struct EmbedderOptions { pub api_key: Option, pub embedding_model: EmbeddingModel, pub dimensions: Option, + pub distribution: Option, } impl EmbedderOptions { @@ -37,6 +38,10 @@ impl EmbedderOptions { query } + + pub fn distribution(&self) -> Option { + self.distribution.or(self.embedding_model.distribution()) + } } #[derive( @@ -139,11 +144,11 @@ pub const OPENAI_EMBEDDINGS_URL: &str = "https://api.openai.com/v1/embeddings"; impl EmbedderOptions { pub fn with_default_model(api_key: Option) -> Self { - Self { api_key, embedding_model: Default::default(), dimensions: None } + Self { api_key, embedding_model: Default::default(), dimensions: None, distribution: None } } pub fn with_embedding_model(api_key: Option, embedding_model: EmbeddingModel) -> Self { - Self { api_key, embedding_model, dimensions: None } + Self { api_key, embedding_model, dimensions: None, distribution: None } } } @@ -170,7 +175,7 @@ impl Embedder { let rest_embedder = RestEmbedder::new(RestEmbedderOptions { api_key: Some(api_key.clone()), - distribution: options.embedding_model.distribution(), + distribution: None, dimensions: Some(options.dimensions()), url: OPENAI_EMBEDDINGS_URL.to_owned(), query: options.query(), @@ -256,6 +261,6 @@ impl Embedder { } pub fn distribution(&self) -> Option { - self.options.embedding_model.distribution() + self.options.distribution() } } From 168ded3b9d17f290190c326ed9574c2937b216f9 Mon Sep 17 00:00:00 2001 From: Louis Dureuil Date: Wed, 27 Mar 2024 11:50:33 +0100 Subject: [PATCH 56/86] Deserr for distribution --- milli/src/vector/mod.rs | 57 +++++++++++++++++++++++++++++++++-------- 1 file changed, 46 insertions(+), 11 deletions(-) diff --git a/milli/src/vector/mod.rs b/milli/src/vector/mod.rs index 4a3a9920e..1cb0a18f7 100644 --- a/milli/src/vector/mod.rs +++ b/milli/src/vector/mod.rs @@ -318,10 +318,50 @@ pub struct DistributionShift { pub current_sigma: OrderedFloat, } -#[derive(Serialize, Deserialize)] +impl Deserr for DistributionShift +where + E: DeserializeError, +{ + fn deserialize_from_value( + value: deserr::Value, + location: deserr::ValuePointerRef, + ) -> Result { + let value = DistributionShiftSerializable::deserialize_from_value(value, location)?; + if value.mean < 0. || value.mean > 1. { + return Err(deserr::take_cf_content(E::error::( + None, + deserr::ErrorKind::Unexpected { + msg: format!( + "the distribution mean must be in the range [0, 1], got {}", + value.mean + ), + }, + location, + ))); + } + if value.sigma <= 0. || value.sigma > 1. { + return Err(deserr::take_cf_content(E::error::( + None, + deserr::ErrorKind::Unexpected { + msg: format!( + "the distribution sigma must be in the range ]0, 1], got {}", + value.sigma + ), + }, + location, + ))); + } + + Ok(value.into()) + } +} + +#[derive(Serialize, Deserialize, Deserr)] +#[serde(deny_unknown_fields)] +#[deserr(deny_unknown_fields)] struct DistributionShiftSerializable { - current_mean: f32, - current_sigma: f32, + mean: f32, + sigma: f32, } impl From for DistributionShiftSerializable { @@ -331,18 +371,13 @@ impl From for DistributionShiftSerializable { current_sigma: OrderedFloat(current_sigma), }: DistributionShift, ) -> Self { - Self { current_mean, current_sigma } + Self { mean: current_mean, sigma: current_sigma } } } impl From for DistributionShift { - fn from( - DistributionShiftSerializable { current_mean, current_sigma }: DistributionShiftSerializable, - ) -> Self { - Self { - current_mean: OrderedFloat(current_mean), - current_sigma: OrderedFloat(current_sigma), - } + fn from(DistributionShiftSerializable { mean, sigma }: DistributionShiftSerializable) -> Self { + Self { current_mean: OrderedFloat(mean), current_sigma: OrderedFloat(sigma) } } } From a25456120d352fd10729ae2751f4774f61840f9d Mon Sep 17 00:00:00 2001 From: Louis Dureuil Date: Wed, 27 Mar 2024 11:51:04 +0100 Subject: [PATCH 57/86] Expose distribution in settings --- milli/src/vector/settings.rs | 55 ++++++++++++++++++++++++++++-------- 1 file changed, 44 insertions(+), 11 deletions(-) diff --git a/milli/src/vector/settings.rs b/milli/src/vector/settings.rs index c277dd0cf..b13b84178 100644 --- a/milli/src/vector/settings.rs +++ b/milli/src/vector/settings.rs @@ -2,7 +2,7 @@ use deserr::Deserr; use serde::{Deserialize, Serialize}; use super::rest::InputType; -use super::{ollama, openai}; +use super::{ollama, openai, DistributionShift}; use crate::prompt::PromptData; use crate::update::Setting; use crate::vector::EmbeddingConfig; @@ -48,6 +48,9 @@ pub struct EmbeddingSettings { #[serde(default, skip_serializing_if = "Setting::is_not_set")] #[deserr(default)] pub input_type: Setting, + #[serde(default, skip_serializing_if = "Setting::is_not_set")] + #[deserr(default)] + pub distribution: Setting, } pub fn check_unset( @@ -101,6 +104,8 @@ impl EmbeddingSettings { pub const EMBEDDING_OBJECT: &'static str = "embeddingObject"; pub const INPUT_TYPE: &'static str = "inputType"; + pub const DISTRIBUTION: &'static str = "distribution"; + pub fn allowed_sources_for_field(field: &'static str) -> &'static [EmbedderSource] { match field { Self::SOURCE => &[ @@ -132,6 +137,13 @@ impl EmbeddingSettings { Self::PATH_TO_EMBEDDINGS => &[EmbedderSource::Rest], Self::EMBEDDING_OBJECT => &[EmbedderSource::Rest], Self::INPUT_TYPE => &[EmbedderSource::Rest], + Self::DISTRIBUTION => &[ + EmbedderSource::HuggingFace, + EmbedderSource::Ollama, + EmbedderSource::OpenAi, + EmbedderSource::Rest, + EmbedderSource::UserProvided, + ], _other => unreachable!("unknown field"), } } @@ -144,14 +156,24 @@ impl EmbeddingSettings { Self::API_KEY, Self::DOCUMENT_TEMPLATE, Self::DIMENSIONS, + Self::DISTRIBUTION, ], - EmbedderSource::HuggingFace => { - &[Self::SOURCE, Self::MODEL, Self::REVISION, Self::DOCUMENT_TEMPLATE] - } - EmbedderSource::Ollama => { - &[Self::SOURCE, Self::MODEL, Self::DOCUMENT_TEMPLATE, Self::URL, Self::API_KEY] - } - EmbedderSource::UserProvided => &[Self::SOURCE, Self::DIMENSIONS], + EmbedderSource::HuggingFace => &[ + Self::SOURCE, + Self::MODEL, + Self::REVISION, + Self::DOCUMENT_TEMPLATE, + Self::DISTRIBUTION, + ], + EmbedderSource::Ollama => &[ + Self::SOURCE, + Self::MODEL, + Self::DOCUMENT_TEMPLATE, + Self::URL, + Self::API_KEY, + Self::DISTRIBUTION, + ], + EmbedderSource::UserProvided => &[Self::SOURCE, Self::DIMENSIONS, Self::DISTRIBUTION], EmbedderSource::Rest => &[ Self::SOURCE, Self::API_KEY, @@ -163,6 +185,7 @@ impl EmbeddingSettings { Self::PATH_TO_EMBEDDINGS, Self::EMBEDDING_OBJECT, Self::INPUT_TYPE, + Self::DISTRIBUTION, ], } } @@ -283,6 +306,7 @@ impl From for EmbeddingSettings { path_to_embeddings: Setting::NotSet, embedding_object: Setting::NotSet, input_type: Setting::NotSet, + distribution: options.distribution.map(Setting::Set).unwrap_or_default(), }, super::EmbedderOptions::OpenAi(options) => Self { source: Setting::Set(EmbedderSource::OpenAi), @@ -297,6 +321,7 @@ impl From for EmbeddingSettings { path_to_embeddings: Setting::NotSet, embedding_object: Setting::NotSet, input_type: Setting::NotSet, + distribution: options.distribution.map(Setting::Set).unwrap_or_default(), }, super::EmbedderOptions::Ollama(options) => Self { source: Setting::Set(EmbedderSource::Ollama), @@ -311,6 +336,7 @@ impl From for EmbeddingSettings { path_to_embeddings: Setting::NotSet, embedding_object: Setting::NotSet, input_type: Setting::NotSet, + distribution: options.distribution.map(Setting::Set).unwrap_or_default(), }, super::EmbedderOptions::UserProvided(options) => Self { source: Setting::Set(EmbedderSource::UserProvided), @@ -325,11 +351,10 @@ impl From for EmbeddingSettings { path_to_embeddings: Setting::NotSet, embedding_object: Setting::NotSet, input_type: Setting::NotSet, + distribution: options.distribution.map(Setting::Set).unwrap_or_default(), }, super::EmbedderOptions::Rest(super::rest::EmbedderOptions { api_key, - // TODO: support distribution - distribution: _, dimensions, url, query, @@ -337,6 +362,7 @@ impl From for EmbeddingSettings { path_to_embeddings, embedding_object, input_type, + distribution, }) => Self { source: Setting::Set(EmbedderSource::Rest), model: Setting::NotSet, @@ -350,6 +376,7 @@ impl From for EmbeddingSettings { path_to_embeddings: Setting::Set(path_to_embeddings), embedding_object: Setting::Set(embedding_object), input_type: Setting::Set(input_type), + distribution: distribution.map(Setting::Set).unwrap_or_default(), }, } } @@ -371,7 +398,9 @@ impl From for EmbeddingConfig { path_to_embeddings, embedding_object, input_type, + distribution, } = value; + if let Some(source) = source.set() { match source { EmbedderSource::OpenAi => { @@ -387,6 +416,7 @@ impl From for EmbeddingConfig { if let Some(dimensions) = dimensions.set() { options.dimensions = Some(dimensions); } + options.distribution = distribution.set(); this.embedder_options = super::EmbedderOptions::OpenAi(options); } EmbedderSource::Ollama => { @@ -399,6 +429,7 @@ impl From for EmbeddingConfig { options.embedding_model = model; } + options.distribution = distribution.set(); this.embedder_options = super::EmbedderOptions::Ollama(options); } EmbedderSource::HuggingFace => { @@ -415,12 +446,14 @@ impl From for EmbeddingConfig { if let Some(revision) = revision.set() { options.revision = Some(revision); } + options.distribution = distribution.set(); this.embedder_options = super::EmbedderOptions::HuggingFace(options); } EmbedderSource::UserProvided => { this.embedder_options = super::EmbedderOptions::UserProvided(super::manual::EmbedderOptions { dimensions: dimensions.set().unwrap(), + distribution: distribution.set(), }); } EmbedderSource::Rest => { @@ -429,7 +462,6 @@ impl From for EmbeddingConfig { this.embedder_options = super::EmbedderOptions::Rest(super::rest::EmbedderOptions { api_key: api_key.set(), - distribution: None, dimensions: dimensions.set(), url: url.set().unwrap(), query: query.set().unwrap_or(embedder_options.query), @@ -441,6 +473,7 @@ impl From for EmbeddingConfig { .set() .unwrap_or(embedder_options.embedding_object), input_type: input_type.set().unwrap_or(embedder_options.input_type), + distribution: distribution.set(), }) } } From 4ff02557837ef1e6a57ebb8306a9e81f24796c42 Mon Sep 17 00:00:00 2001 From: Louis Dureuil Date: Wed, 27 Mar 2024 11:51:14 +0100 Subject: [PATCH 58/86] remove unused function --- milli/src/vector/settings.rs | 52 ------------------------------------ 1 file changed, 52 deletions(-) diff --git a/milli/src/vector/settings.rs b/milli/src/vector/settings.rs index b13b84178..78f83cbea 100644 --- a/milli/src/vector/settings.rs +++ b/milli/src/vector/settings.rs @@ -237,58 +237,6 @@ impl std::fmt::Display for EmbedderSource { } } -impl EmbeddingSettings { - pub fn apply(&mut self, new: Self) { - let EmbeddingSettings { - source, - model, - revision, - api_key, - dimensions, - document_template, - url, - query, - input_field, - path_to_embeddings, - embedding_object, - input_type, - } = new; - let old_source = self.source; - self.source.apply(source); - // Reinitialize the whole setting object on a source change - if old_source != self.source { - *self = EmbeddingSettings { - source, - model, - revision, - api_key, - dimensions, - document_template, - url, - query, - input_field, - path_to_embeddings, - embedding_object, - input_type, - }; - return; - } - - self.model.apply(model); - self.revision.apply(revision); - self.api_key.apply(api_key); - self.dimensions.apply(dimensions); - self.document_template.apply(document_template); - - self.url.apply(url); - self.query.apply(query); - self.input_field.apply(input_field); - self.path_to_embeddings.apply(path_to_embeddings); - self.embedding_object.apply(embedding_object); - self.input_type.apply(input_type); - } -} - impl From for EmbeddingSettings { fn from(value: EmbeddingConfig) -> Self { let EmbeddingConfig { embedder_options, prompt } = value; From 572fb3a51d76d0e92d7e0841e770758757509106 Mon Sep 17 00:00:00 2001 From: Louis Dureuil Date: Wed, 27 Mar 2024 11:48:00 +0100 Subject: [PATCH 59/86] Finer granularity for embedder needs reindex --- milli/src/update/settings.rs | 7 ++++- milli/src/vector/settings.rs | 60 ++++++++++++++++++++++++++++++++++++ 2 files changed, 66 insertions(+), 1 deletion(-) diff --git a/milli/src/update/settings.rs b/milli/src/update/settings.rs index 9f47768c1..b784b3f92 100644 --- a/milli/src/update/settings.rs +++ b/milli/src/update/settings.rs @@ -976,7 +976,12 @@ impl<'a, 't, 'i> Settings<'a, 't, 'i> { match joined { // updated config EitherOrBoth::Both((name, mut old), (_, new)) => { - changed |= old.apply(new); + changed |= EmbeddingSettings::apply_and_need_reindex(&mut old, new); + if changed { + tracing::debug!(embedder = name, "need reindex"); + } else { + tracing::debug!(embedder = name, "skip reindex"); + } let new = validate_embedding_settings(old, &name)?; new_configs.insert(name, new); } diff --git a/milli/src/vector/settings.rs b/milli/src/vector/settings.rs index 78f83cbea..18a86368f 100644 --- a/milli/src/vector/settings.rs +++ b/milli/src/vector/settings.rs @@ -210,6 +210,66 @@ impl EmbeddingSettings { *model = Setting::Set(openai::EmbeddingModel::default().name().to_owned()) } } + + pub(crate) fn apply_and_need_reindex( + old: &mut Setting, + new: Setting, + ) -> bool { + match (old, new) { + ( + Setting::Set(EmbeddingSettings { + source: old_source, + model: old_model, + revision: old_revision, + api_key: old_api_key, + dimensions: old_dimensions, + document_template: old_document_template, + url: old_url, + query: old_query, + input_field: old_input_field, + path_to_embeddings: old_path_to_embeddings, + embedding_object: old_embedding_object, + input_type: old_input_type, + distribution: old_distribution, + }), + Setting::Set(EmbeddingSettings { + source: new_source, + model: new_model, + revision: new_revision, + api_key: new_api_key, + dimensions: new_dimensions, + document_template: new_document_template, + url: new_url, + query: new_query, + input_field: new_input_field, + path_to_embeddings: new_path_to_embeddings, + embedding_object: new_embedding_object, + input_type: new_input_type, + distribution: new_distribution, + }), + ) => { + let mut needs_reindex = false; + + needs_reindex |= old_source.apply(new_source); + needs_reindex |= old_model.apply(new_model); + needs_reindex |= old_revision.apply(new_revision); + needs_reindex |= old_dimensions.apply(new_dimensions); + needs_reindex |= old_document_template.apply(new_document_template); + needs_reindex |= old_url.apply(new_url); + needs_reindex |= old_query.apply(new_query); + needs_reindex |= old_input_field.apply(new_input_field); + needs_reindex |= old_path_to_embeddings.apply(new_path_to_embeddings); + needs_reindex |= old_embedding_object.apply(new_embedding_object); + needs_reindex |= old_input_type.apply(new_input_type); + + old_distribution.apply(new_distribution); + old_api_key.apply(new_api_key); + needs_reindex + } + (Setting::Reset, Setting::Reset) | (_, Setting::NotSet) => false, + _ => true, + } + } } #[derive(Debug, Clone, Copy, Default, Serialize, Deserialize, PartialEq, Eq, Deserr)] From 92224f109a9ef91d69cac1a721b48906619b9eb0 Mon Sep 17 00:00:00 2001 From: Louis Dureuil Date: Wed, 27 Mar 2024 12:19:10 +0100 Subject: [PATCH 60/86] Fix tests --- .../test_settings_update/after_registering_settings_task.snap | 2 +- .../lib.rs/test_settings_update/settings_update_processed.snap | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/index-scheduler/src/snapshots/lib.rs/test_settings_update/after_registering_settings_task.snap b/index-scheduler/src/snapshots/lib.rs/test_settings_update/after_registering_settings_task.snap index 01bb73993..8c081b84b 100644 --- a/index-scheduler/src/snapshots/lib.rs/test_settings_update/after_registering_settings_task.snap +++ b/index-scheduler/src/snapshots/lib.rs/test_settings_update/after_registering_settings_task.snap @@ -6,7 +6,7 @@ source: index-scheduler/src/lib.rs [] ---------------------------------------------------------------------- ### All Tasks: -0 {uid: 0, status: enqueued, details: { settings: Settings { displayed_attributes: NotSet, searchable_attributes: NotSet, filterable_attributes: NotSet, sortable_attributes: NotSet, ranking_rules: NotSet, stop_words: NotSet, non_separator_tokens: NotSet, separator_tokens: NotSet, dictionary: NotSet, synonyms: NotSet, distinct_attribute: NotSet, proximity_precision: NotSet, typo_tolerance: NotSet, faceting: NotSet, pagination: NotSet, embedders: Set({"default": Set(EmbeddingSettings { source: Set(Rest), model: NotSet, revision: NotSet, api_key: Set("My super secret"), dimensions: NotSet, document_template: NotSet, url: Set("http://localhost:7777"), query: NotSet, input_field: NotSet, path_to_embeddings: NotSet, embedding_object: NotSet, input_type: NotSet })}), search_cutoff_ms: NotSet, _kind: PhantomData } }, kind: SettingsUpdate { index_uid: "doggos", new_settings: Settings { displayed_attributes: NotSet, searchable_attributes: NotSet, filterable_attributes: NotSet, sortable_attributes: NotSet, ranking_rules: NotSet, stop_words: NotSet, non_separator_tokens: NotSet, separator_tokens: NotSet, dictionary: NotSet, synonyms: NotSet, distinct_attribute: NotSet, proximity_precision: NotSet, typo_tolerance: NotSet, faceting: NotSet, pagination: NotSet, embedders: Set({"default": Set(EmbeddingSettings { source: Set(Rest), model: NotSet, revision: NotSet, api_key: Set("My super secret"), dimensions: NotSet, document_template: NotSet, url: Set("http://localhost:7777"), query: NotSet, input_field: NotSet, path_to_embeddings: NotSet, embedding_object: NotSet, input_type: NotSet })}), search_cutoff_ms: NotSet, _kind: PhantomData }, is_deletion: false, allow_index_creation: true }} +0 {uid: 0, status: enqueued, details: { settings: Settings { displayed_attributes: NotSet, searchable_attributes: NotSet, filterable_attributes: NotSet, sortable_attributes: NotSet, ranking_rules: NotSet, stop_words: NotSet, non_separator_tokens: NotSet, separator_tokens: NotSet, dictionary: NotSet, synonyms: NotSet, distinct_attribute: NotSet, proximity_precision: NotSet, typo_tolerance: NotSet, faceting: NotSet, pagination: NotSet, embedders: Set({"default": Set(EmbeddingSettings { source: Set(Rest), model: NotSet, revision: NotSet, api_key: Set("My super secret"), dimensions: NotSet, document_template: NotSet, url: Set("http://localhost:7777"), query: NotSet, input_field: NotSet, path_to_embeddings: NotSet, embedding_object: NotSet, input_type: NotSet, distribution: NotSet })}), search_cutoff_ms: NotSet, _kind: PhantomData } }, kind: SettingsUpdate { index_uid: "doggos", new_settings: Settings { displayed_attributes: NotSet, searchable_attributes: NotSet, filterable_attributes: NotSet, sortable_attributes: NotSet, ranking_rules: NotSet, stop_words: NotSet, non_separator_tokens: NotSet, separator_tokens: NotSet, dictionary: NotSet, synonyms: NotSet, distinct_attribute: NotSet, proximity_precision: NotSet, typo_tolerance: NotSet, faceting: NotSet, pagination: NotSet, embedders: Set({"default": Set(EmbeddingSettings { source: Set(Rest), model: NotSet, revision: NotSet, api_key: Set("My super secret"), dimensions: NotSet, document_template: NotSet, url: Set("http://localhost:7777"), query: NotSet, input_field: NotSet, path_to_embeddings: NotSet, embedding_object: NotSet, input_type: NotSet, distribution: NotSet })}), search_cutoff_ms: NotSet, _kind: PhantomData }, is_deletion: false, allow_index_creation: true }} ---------------------------------------------------------------------- ### Status: enqueued [0,] diff --git a/index-scheduler/src/snapshots/lib.rs/test_settings_update/settings_update_processed.snap b/index-scheduler/src/snapshots/lib.rs/test_settings_update/settings_update_processed.snap index d1d219da1..f6fb6a186 100644 --- a/index-scheduler/src/snapshots/lib.rs/test_settings_update/settings_update_processed.snap +++ b/index-scheduler/src/snapshots/lib.rs/test_settings_update/settings_update_processed.snap @@ -6,7 +6,7 @@ source: index-scheduler/src/lib.rs [] ---------------------------------------------------------------------- ### All Tasks: -0 {uid: 0, status: succeeded, details: { settings: Settings { displayed_attributes: NotSet, searchable_attributes: NotSet, filterable_attributes: NotSet, sortable_attributes: NotSet, ranking_rules: NotSet, stop_words: NotSet, non_separator_tokens: NotSet, separator_tokens: NotSet, dictionary: NotSet, synonyms: NotSet, distinct_attribute: NotSet, proximity_precision: NotSet, typo_tolerance: NotSet, faceting: NotSet, pagination: NotSet, embedders: Set({"default": Set(EmbeddingSettings { source: Set(Rest), model: NotSet, revision: NotSet, api_key: Set("My super secret"), dimensions: NotSet, document_template: NotSet, url: Set("http://localhost:7777"), query: NotSet, input_field: NotSet, path_to_embeddings: NotSet, embedding_object: NotSet, input_type: NotSet })}), search_cutoff_ms: NotSet, _kind: PhantomData } }, kind: SettingsUpdate { index_uid: "doggos", new_settings: Settings { displayed_attributes: NotSet, searchable_attributes: NotSet, filterable_attributes: NotSet, sortable_attributes: NotSet, ranking_rules: NotSet, stop_words: NotSet, non_separator_tokens: NotSet, separator_tokens: NotSet, dictionary: NotSet, synonyms: NotSet, distinct_attribute: NotSet, proximity_precision: NotSet, typo_tolerance: NotSet, faceting: NotSet, pagination: NotSet, embedders: Set({"default": Set(EmbeddingSettings { source: Set(Rest), model: NotSet, revision: NotSet, api_key: Set("My super secret"), dimensions: NotSet, document_template: NotSet, url: Set("http://localhost:7777"), query: NotSet, input_field: NotSet, path_to_embeddings: NotSet, embedding_object: NotSet, input_type: NotSet })}), search_cutoff_ms: NotSet, _kind: PhantomData }, is_deletion: false, allow_index_creation: true }} +0 {uid: 0, status: succeeded, details: { settings: Settings { displayed_attributes: NotSet, searchable_attributes: NotSet, filterable_attributes: NotSet, sortable_attributes: NotSet, ranking_rules: NotSet, stop_words: NotSet, non_separator_tokens: NotSet, separator_tokens: NotSet, dictionary: NotSet, synonyms: NotSet, distinct_attribute: NotSet, proximity_precision: NotSet, typo_tolerance: NotSet, faceting: NotSet, pagination: NotSet, embedders: Set({"default": Set(EmbeddingSettings { source: Set(Rest), model: NotSet, revision: NotSet, api_key: Set("My super secret"), dimensions: NotSet, document_template: NotSet, url: Set("http://localhost:7777"), query: NotSet, input_field: NotSet, path_to_embeddings: NotSet, embedding_object: NotSet, input_type: NotSet, distribution: NotSet })}), search_cutoff_ms: NotSet, _kind: PhantomData } }, kind: SettingsUpdate { index_uid: "doggos", new_settings: Settings { displayed_attributes: NotSet, searchable_attributes: NotSet, filterable_attributes: NotSet, sortable_attributes: NotSet, ranking_rules: NotSet, stop_words: NotSet, non_separator_tokens: NotSet, separator_tokens: NotSet, dictionary: NotSet, synonyms: NotSet, distinct_attribute: NotSet, proximity_precision: NotSet, typo_tolerance: NotSet, faceting: NotSet, pagination: NotSet, embedders: Set({"default": Set(EmbeddingSettings { source: Set(Rest), model: NotSet, revision: NotSet, api_key: Set("My super secret"), dimensions: NotSet, document_template: NotSet, url: Set("http://localhost:7777"), query: NotSet, input_field: NotSet, path_to_embeddings: NotSet, embedding_object: NotSet, input_type: NotSet, distribution: NotSet })}), search_cutoff_ms: NotSet, _kind: PhantomData }, is_deletion: false, allow_index_creation: true }} ---------------------------------------------------------------------- ### Status: enqueued [] From cde7ce4f44372eaeb66d4135f8e2498d5c2f4f2f Mon Sep 17 00:00:00 2001 From: Louis Dureuil Date: Wed, 27 Mar 2024 14:02:09 +0100 Subject: [PATCH 61/86] Add test --- meilisearch/tests/search/hybrid.rs | 32 ++++++++++++++++++++++++++++++ 1 file changed, 32 insertions(+) diff --git a/meilisearch/tests/search/hybrid.rs b/meilisearch/tests/search/hybrid.rs index 85bc96d86..8decb7ded 100644 --- a/meilisearch/tests/search/hybrid.rs +++ b/meilisearch/tests/search/hybrid.rs @@ -87,6 +87,38 @@ async fn simple_search() { snapshot!(response["hits"], @r###"[{"title":"Captain Marvel","desc":"a Shazam ersatz","id":"3","_vectors":{"default":[2.0,3.0]},"_semanticScore":0.99029034},{"title":"Captain Planet","desc":"He's not part of the Marvel Cinematic Universe","id":"2","_vectors":{"default":[1.0,2.0]},"_semanticScore":0.97434163},{"title":"Shazam!","desc":"a Captain Marvel ersatz","id":"1","_vectors":{"default":[1.0,3.0]},"_semanticScore":0.9472136}]"###); } +#[actix_rt::test] +async fn distribution_shift() { + let server = Server::new().await; + let index = index_with_documents(&server, &SIMPLE_SEARCH_DOCUMENTS).await; + + let search = json!({"q": "Captain", "vector": [1.0, 1.0], "showRankingScore": true, "hybrid": {"semanticRatio": 1.0}}); + let (response, code) = index.search_post(search.clone()).await; + snapshot!(code, @"200 OK"); + snapshot!(response["hits"], @r###"[{"title":"Captain Marvel","desc":"a Shazam ersatz","id":"3","_vectors":{"default":[2.0,3.0]},"_rankingScore":0.990290343761444,"_semanticScore":0.99029034},{"title":"Captain Planet","desc":"He's not part of the Marvel Cinematic Universe","id":"2","_vectors":{"default":[1.0,2.0]},"_rankingScore":0.974341630935669,"_semanticScore":0.97434163},{"title":"Shazam!","desc":"a Captain Marvel ersatz","id":"1","_vectors":{"default":[1.0,3.0]},"_rankingScore":0.9472135901451112,"_semanticScore":0.9472136}]"###); + + let (response, code) = index + .update_settings(json!({ + "embedders": { + "default": { + "distribution": { + "mean": 0.998, + "sigma": 0.01 + } + } + } + })) + .await; + + snapshot!(code, @"202 Accepted"); + let response = server.wait_task(response.uid()).await; + snapshot!(response["details"], @r###"{"embedders":{"default":{"distribution":{"mean":0.998,"sigma":0.01}}}}"###); + + let (response, code) = index.search_post(search).await; + snapshot!(code, @"200 OK"); + snapshot!(response["hits"], @r###"[{"title":"Captain Marvel","desc":"a Shazam ersatz","id":"3","_vectors":{"default":[2.0,3.0]},"_rankingScore":0.19161224365234375,"_semanticScore":0.19161224},{"title":"Captain Planet","desc":"He's not part of the Marvel Cinematic Universe","id":"2","_vectors":{"default":[1.0,2.0]},"_rankingScore":1.1920928955078125e-7,"_semanticScore":1.1920929e-7},{"title":"Shazam!","desc":"a Captain Marvel ersatz","id":"1","_vectors":{"default":[1.0,3.0]},"_rankingScore":1.1920928955078125e-7,"_semanticScore":1.1920929e-7}]"###); +} + #[actix_rt::test] async fn highlighter() { let server = Server::new().await; From 03c886ac1b58179da966a9d30bc06481552e63ee Mon Sep 17 00:00:00 2001 From: Tamo Date: Wed, 27 Mar 2024 15:38:36 +0100 Subject: [PATCH 62/86] adds a bit of documentation --- meilisearch/src/search_queue.rs | 29 ++++++++++++++++++++++++++++- 1 file changed, 28 insertions(+), 1 deletion(-) diff --git a/meilisearch/src/search_queue.rs b/meilisearch/src/search_queue.rs index ad222cdf8..c0d7c7706 100644 --- a/meilisearch/src/search_queue.rs +++ b/meilisearch/src/search_queue.rs @@ -1,3 +1,22 @@ +//! This file implements a queue of searches to process and the ability to control how many searches can be run in parallel. +//! We need this because we don't want to process more search requests than we have cores. +//! That slows down everything and consumes RAM for no reason. +//! The steps to do a search are to get the `SearchQueue` data structure and try to get a search permit. +//! This can fail if the queue is full, and we need to drop your search request to register a new one. +//! +//! ### How to do a search request +//! +//! In order to do a search request you should try to get a search permit. +//! Retrieve the `SearchQueue` structure from actix-web (`search_queue: Data`) +//! and right before processing the search, calls the `SearchQueue::try_get_search_permit` method: `search_queue.try_get_search_permit().await?;` +//! +//! What is going to happen at this point is that you're going to send a oneshot::Sender over an async mpsc channel. +//! Then, the queue/scheduler is going to either: +//! - Drop your oneshot channel => that means there are too many searches going on, and yours won't be executed. +//! You should exit and free all the RAM you use ASAP. +//! - Sends you a Permit => that will unlock the method, and you will be able to process your search. +//! And should drop the Permit only once you have freed all the RAM consumed by the method. + use std::num::NonZeroUsize; use rand::rngs::StdRng; @@ -12,6 +31,8 @@ pub struct SearchQueue { capacity: usize, } +/// You should only run search requests while holding this permit. +/// Once it's dropped, a new search request will be able to process. #[derive(Debug)] pub struct Permit { sender: mpsc::Sender<()>, @@ -34,6 +55,10 @@ impl SearchQueue { Self { sender, capacity } } + /// This function is the main loop, it's in charge on scheduling which search request should execute first and + /// how many should executes at the same time. + /// + /// It **must never** panic or exit. async fn run( capacity: usize, parallelism: NonZeroUsize, @@ -42,7 +67,7 @@ impl SearchQueue { let mut queue: Vec> = Default::default(); let mut rng: StdRng = StdRng::from_entropy(); let mut searches_running: usize = 0; - // by having a capacity of parallelism we ensures that every time a search finish it can release its RAM asap + // By having a capacity of parallelism we ensures that every time a search finish it can release its RAM asap let (sender, mut search_finished) = mpsc::channel(parallelism.into()); loop { @@ -85,6 +110,8 @@ impl SearchQueue { } } + /// Returns a search `Permit`. + /// It should be dropped as soon as you've freed all the RAM associated with the search request being processed. pub async fn try_get_search_permit(&self) -> Result { let (sender, receiver) = oneshot::channel(); self.sender.send(sender).await.map_err(|_| MeilisearchHttpError::SearchLimiterIsDown)?; From b7c582e4f3a44bf96e1b34b46ab6731cb5d54665 Mon Sep 17 00:00:00 2001 From: Tamo Date: Wed, 27 Mar 2024 15:49:43 +0100 Subject: [PATCH 63/86] connect the search queue with the health route --- meilisearch/src/routes/mod.rs | 3 +++ meilisearch/src/search_queue.rs | 10 ++++++++++ 2 files changed, 13 insertions(+) diff --git a/meilisearch/src/routes/mod.rs b/meilisearch/src/routes/mod.rs index 1c1465582..7cf886017 100644 --- a/meilisearch/src/routes/mod.rs +++ b/meilisearch/src/routes/mod.rs @@ -15,6 +15,7 @@ use tracing::debug; use crate::analytics::Analytics; use crate::extractors::authentication::policies::*; use crate::extractors::authentication::GuardedData; +use crate::search_queue::SearchQueue; use crate::Opt; const PAGINATION_DEFAULT_LIMIT: usize = 20; @@ -385,10 +386,12 @@ pub async fn get_health( req: HttpRequest, index_scheduler: Data, auth_controller: Data, + search_queue: Data, analytics: web::Data, ) -> Result { analytics.health_seen(&req); + search_queue.health().unwrap(); index_scheduler.health().unwrap(); auth_controller.health().unwrap(); diff --git a/meilisearch/src/search_queue.rs b/meilisearch/src/search_queue.rs index c0d7c7706..6d5044d20 100644 --- a/meilisearch/src/search_queue.rs +++ b/meilisearch/src/search_queue.rs @@ -117,4 +117,14 @@ impl SearchQueue { self.sender.send(sender).await.map_err(|_| MeilisearchHttpError::SearchLimiterIsDown)?; receiver.await.map_err(|_| MeilisearchHttpError::TooManySearchRequests(self.capacity)) } + + /// Returns `Ok(())` if everything seems normal. + /// Returns `Err(MeilisearchHttpError::SearchLimiterIsDown)` if the search limiter seems down. + pub fn health(&self) -> Result<(), MeilisearchHttpError> { + if self.sender.is_closed() { + Err(MeilisearchHttpError::SearchLimiterIsDown) + } else { + Ok(()) + } + } } From 06a11b5b216692e3a03c2b47aa49f6e531a8d19e Mon Sep 17 00:00:00 2001 From: Tamo Date: Wed, 27 Mar 2024 17:34:49 +0100 Subject: [PATCH 64/86] Improve error message --- meilisearch/src/error.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/meilisearch/src/error.rs b/meilisearch/src/error.rs index bb1156997..5a0b04020 100644 --- a/meilisearch/src/error.rs +++ b/meilisearch/src/error.rs @@ -31,7 +31,7 @@ pub enum MeilisearchHttpError { MissingPayload(PayloadType), #[error("Too many search requests running at the same time: {0}. Retry after 10s.")] TooManySearchRequests(usize), - #[error("Internal error: Search limiter is down")] + #[error("Internal error: Search limiter is down.")] SearchLimiterIsDown, #[error("The provided payload reached the size limit. The maximum accepted payload size is {}.", Byte::from_bytes(*.0 as u64).get_appropriate_unit(true))] PayloadTooLarge(usize), From 8f2606d79d8aa5bb79df036d0e8d673327010fb7 Mon Sep 17 00:00:00 2001 From: Bruno Casali Date: Wed, 27 Mar 2024 14:26:47 -0300 Subject: [PATCH 65/86] fixes typos --- index-scheduler/src/lib.rs | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/index-scheduler/src/lib.rs b/index-scheduler/src/lib.rs index 46ed76649..59bfa1ecf 100644 --- a/index-scheduler/src/lib.rs +++ b/index-scheduler/src/lib.rs @@ -1301,8 +1301,8 @@ impl IndexScheduler { wtxn.commit().map_err(Error::HeedTransaction)?; - // Once the tasks are commited, we should delete all the update files associated ASAP to avoid leaking files in case of a restart - tracing::debug!("Deleting the upadate files"); + // Once the tasks are committed, we should delete all the update files associated ASAP to avoid leaking files in case of a restart + tracing::debug!("Deleting the update files"); //We take one read transaction **per thread**. Then, every thread is going to pull out new IDs from the roaring bitmap with the help of an atomic shared index into the bitmap let idx = AtomicU32::new(0); @@ -1332,7 +1332,7 @@ impl IndexScheduler { Ok(TickOutcome::TickAgain(processed_tasks)) } - /// Once the tasks changes have been commited we must send all the tasks that were updated to our webhook if there is one. + /// Once the tasks changes have been committed we must send all the tasks that were updated to our webhook if there is one. fn notify_webhook(&self, updated: &RoaringBitmap) -> Result<()> { if let Some(ref url) = self.webhook_url { struct TaskReader<'a, 'b> { From 69f8b2730d9128816a18a5c100379f96e815ddc2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Cl=C3=A9ment=20Renault?= Date: Tue, 26 Mar 2024 18:07:43 +0100 Subject: [PATCH 66/86] Fix the tests --- milli/src/index.rs | 1 + milli/src/search/new/matches/matching_words.rs | 2 +- milli/src/search/new/query_term/parse_query.rs | 2 +- 3 files changed, 3 insertions(+), 2 deletions(-) diff --git a/milli/src/index.rs b/milli/src/index.rs index d921de9e4..b2ea105da 100644 --- a/milli/src/index.rs +++ b/milli/src/index.rs @@ -2435,6 +2435,7 @@ pub(crate) mod tests { document_scores: _, mut documents_ids, degraded: _, + used_negative_operator: _, } = search.execute().unwrap(); let primary_key_id = index.fields_ids_map(&rtxn).unwrap().id("primary_key").unwrap(); documents_ids.sort_unstable(); diff --git a/milli/src/search/new/matches/matching_words.rs b/milli/src/search/new/matches/matching_words.rs index 9e01f6dcf..acc0de6b0 100644 --- a/milli/src/search/new/matches/matching_words.rs +++ b/milli/src/search/new/matches/matching_words.rs @@ -261,7 +261,7 @@ pub(crate) mod tests { let mut builder = TokenizerBuilder::default(); let tokenizer = builder.build(); let tokens = tokenizer.tokenize("split this world"); - let query_terms = located_query_terms_from_tokens(&mut ctx, tokens, None).unwrap(); + let (query_terms, _) = located_query_terms_from_tokens(&mut ctx, tokens, None).unwrap(); let matching_words = MatchingWords::new(ctx, query_terms); assert_eq!( diff --git a/milli/src/search/new/query_term/parse_query.rs b/milli/src/search/new/query_term/parse_query.rs index e510595ee..bdb937384 100644 --- a/milli/src/search/new/query_term/parse_query.rs +++ b/milli/src/search/new/query_term/parse_query.rs @@ -331,7 +331,7 @@ mod tests { let rtxn = index.read_txn()?; let mut ctx = SearchContext::new(&index, &rtxn); // panics with `attempt to add with overflow` before - let located_query_terms = located_query_terms_from_tokens(&mut ctx, tokens, None)?; + let (located_query_terms, _) = located_query_terms_from_tokens(&mut ctx, tokens, None)?; assert!(located_query_terms.is_empty()); Ok(()) From 877f4b1045ae0d0eb9c72a5c36c0d939a736d906 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Cl=C3=A9ment=20Renault?= Date: Thu, 28 Mar 2024 15:51:43 +0100 Subject: [PATCH 67/86] Support negative phrases --- .../src/search/new/matches/matching_words.rs | 4 +- milli/src/search/new/mod.rs | 26 +++++++- milli/src/search/new/query_term/mod.rs | 9 ++- .../src/search/new/query_term/parse_query.rs | 64 +++++++++++++++---- 4 files changed, 85 insertions(+), 18 deletions(-) diff --git a/milli/src/search/new/matches/matching_words.rs b/milli/src/search/new/matches/matching_words.rs index acc0de6b0..56bf6c169 100644 --- a/milli/src/search/new/matches/matching_words.rs +++ b/milli/src/search/new/matches/matching_words.rs @@ -240,6 +240,7 @@ pub(crate) mod tests { use super::super::super::located_query_terms_from_tokens; use super::*; use crate::index::tests::TempIndex; + use crate::search::new::query_term::ExtractedTokens; pub(crate) fn temp_index_with_documents() -> TempIndex { let temp_index = TempIndex::new(); @@ -261,7 +262,8 @@ pub(crate) mod tests { let mut builder = TokenizerBuilder::default(); let tokenizer = builder.build(); let tokens = tokenizer.tokenize("split this world"); - let (query_terms, _) = located_query_terms_from_tokens(&mut ctx, tokens, None).unwrap(); + let ExtractedTokens { query_terms, .. } = + located_query_terms_from_tokens(&mut ctx, tokens, None).unwrap(); let matching_words = MatchingWords::new(ctx, query_terms); assert_eq!( diff --git a/milli/src/search/new/mod.rs b/milli/src/search/new/mod.rs index 99f4562d6..1f0ae7b29 100644 --- a/milli/src/search/new/mod.rs +++ b/milli/src/search/new/mod.rs @@ -33,7 +33,9 @@ use interner::{DedupInterner, Interner}; pub use logger::visual::VisualSearchLogger; pub use logger::{DefaultSearchLogger, SearchLogger}; use query_graph::{QueryGraph, QueryNode}; -use query_term::{located_query_terms_from_tokens, LocatedQueryTerm, Phrase, QueryTerm}; +use query_term::{ + located_query_terms_from_tokens, ExtractedTokens, LocatedQueryTerm, Phrase, QueryTerm, +}; use ranking_rules::{ BoxRankingRule, PlaceholderQuery, RankingRule, RankingRuleOutput, RankingRuleQueryTrait, }; @@ -223,6 +225,21 @@ fn resolve_negative_words( Ok(negative_bitmap) } +#[tracing::instrument(level = "trace", skip_all, target = "search")] +fn resolve_negative_phrases( + ctx: &mut SearchContext, + negative_phrases: &[LocatedQueryTerm], +) -> Result { + let mut negative_bitmap = RoaringBitmap::new(); + for term in negative_phrases { + let query_term = ctx.term_interner.get(term.value); + if let Some(phrase) = query_term.original_phrase() { + negative_bitmap |= ctx.get_phrase_docids(phrase)?; + } + } + Ok(negative_bitmap) +} + /// Return the list of initialised ranking rules to be used for a placeholder search. fn get_ranking_rules_for_placeholder_search<'ctx>( ctx: &SearchContext<'ctx>, @@ -636,12 +653,15 @@ pub fn execute_search( let tokens = tokenizer.tokenize(query); drop(entered); - let (query_terms, negative_words) = + let ExtractedTokens { query_terms, negative_words, negative_phrases } = located_query_terms_from_tokens(ctx, tokens, words_limit)?; - used_negative_operator = !negative_words.is_empty(); + used_negative_operator = !negative_words.is_empty() || !negative_phrases.is_empty(); let ignored_documents = resolve_negative_words(ctx, &negative_words)?; + let ignored_phrases = resolve_negative_phrases(ctx, &negative_phrases)?; + universe -= ignored_documents; + universe -= ignored_phrases; if query_terms.is_empty() { // Do a placeholder search instead diff --git a/milli/src/search/new/query_term/mod.rs b/milli/src/search/new/query_term/mod.rs index a37e60ed0..70f1d3c4f 100644 --- a/milli/src/search/new/query_term/mod.rs +++ b/milli/src/search/new/query_term/mod.rs @@ -9,7 +9,9 @@ use std::ops::RangeInclusive; use either::Either; pub use ntypo_subset::NTypoTermSubset; -pub use parse_query::{located_query_terms_from_tokens, make_ngram, number_of_typos_allowed}; +pub use parse_query::{ + located_query_terms_from_tokens, make_ngram, number_of_typos_allowed, ExtractedTokens, +}; pub use phrase::Phrase; use super::interner::{DedupInterner, Interned}; @@ -478,6 +480,11 @@ impl QueryTerm { pub fn original_word(&self, ctx: &SearchContext) -> String { ctx.word_interner.get(self.original).clone() } + + pub fn original_phrase(&self) -> Option> { + self.zero_typo.phrase + } + pub fn all_computed_derivations(&self) -> (Vec>, Vec>) { let mut words = BTreeSet::new(); let mut phrases = BTreeSet::new(); diff --git a/milli/src/search/new/query_term/parse_query.rs b/milli/src/search/new/query_term/parse_query.rs index bdb937384..86be7da77 100644 --- a/milli/src/search/new/query_term/parse_query.rs +++ b/milli/src/search/new/query_term/parse_query.rs @@ -9,21 +9,34 @@ use crate::search::new::query_term::{Lazy, Phrase, QueryTerm}; use crate::search::new::Word; use crate::{Result, SearchContext, MAX_WORD_LENGTH}; +#[derive(Clone)] +/// Extraction of the content of a query. +pub struct ExtractedTokens { + /// The terms to search for in the database. + pub query_terms: Vec, + /// The words that must not appear in the results. + pub negative_words: Vec, + /// The phrases that must not appear in the results. + pub negative_phrases: Vec, +} + /// Convert the tokenised search query into a list of located query terms. #[tracing::instrument(level = "trace", skip_all, target = "search::query")] pub fn located_query_terms_from_tokens( ctx: &mut SearchContext, query: NormalizedTokenIter, words_limit: Option, -) -> Result<(Vec, Vec)> { +) -> Result { let nbr_typos = number_of_typos_allowed(ctx)?; - let mut located_terms = Vec::new(); + let mut query_terms = Vec::new(); + let mut negative_phrase = false; let mut phrase: Option = None; let mut encountered_whitespace = true; let mut negative_next_token = false; let mut negative_words = Vec::new(); + let mut negative_phrases = Vec::new(); let parts_limit = words_limit.unwrap_or(usize::MAX); @@ -37,8 +50,8 @@ pub fn located_query_terms_from_tokens( } // early return if word limit is exceeded - if located_terms.len() >= parts_limit { - return Ok((located_terms, negative_words)); + if query_terms.len() >= parts_limit { + return Ok(ExtractedTokens { query_terms, negative_words, negative_phrases }); } match token.kind { @@ -71,7 +84,7 @@ pub fn located_query_terms_from_tokens( value: ctx.term_interner.push(term), positions: position..=position, }; - located_terms.push(located_term); + query_terms.push(located_term); } TokenKind::StopWord | TokenKind::Separator(_) | TokenKind::Unknown => (), } @@ -88,7 +101,7 @@ pub fn located_query_terms_from_tokens( value: ctx.term_interner.push(term), positions: position..=position, }; - located_terms.push(located_term); + query_terms.push(located_term); } } TokenKind::Separator(separator_kind) => { @@ -104,7 +117,14 @@ pub fn located_query_terms_from_tokens( let phrase = if separator_kind == SeparatorKind::Hard { if let Some(phrase) = phrase { if let Some(located_query_term) = phrase.build(ctx) { - located_terms.push(located_query_term) + // as we are evaluating a negative operator we put the phrase + // in the negative one *but* we don't reset the negative operator + // as we are immediatly starting a new negative phrase. + if negative_phrase { + negative_phrases.push(located_query_term); + } else { + query_terms.push(located_query_term); + } } Some(PhraseBuilder::empty()) } else { @@ -125,12 +145,24 @@ pub fn located_query_terms_from_tokens( // Per the check above, quote_count > 0 quote_count -= 1; if let Some(located_query_term) = phrase.build(ctx) { - located_terms.push(located_query_term) + // we were evaluating a negative operator so we + // put the phrase in the negative phrases + if negative_phrase { + negative_phrases.push(located_query_term); + negative_phrase = false; + } else { + query_terms.push(located_query_term); + } } } // Start new phrase if the token ends with an opening quote - (quote_count % 2 == 1).then_some(PhraseBuilder::empty()) + if quote_count % 2 == 1 { + negative_phrase = negative_next_token; + Some(PhraseBuilder::empty()) + } else { + None + } }; negative_next_token = @@ -146,11 +178,16 @@ pub fn located_query_terms_from_tokens( // If a quote is never closed, we consider all of the end of the query as a phrase. if let Some(phrase) = phrase.take() { if let Some(located_query_term) = phrase.build(ctx) { - located_terms.push(located_query_term); + // put the phrase in the negative set if we are evaluating a negative operator. + if negative_phrase { + negative_phrases.push(located_query_term); + } else { + query_terms.push(located_query_term); + } } } - Ok((located_terms, negative_words)) + Ok(ExtractedTokens { query_terms, negative_words, negative_phrases }) } pub fn number_of_typos_allowed<'ctx>( @@ -331,8 +368,9 @@ mod tests { let rtxn = index.read_txn()?; let mut ctx = SearchContext::new(&index, &rtxn); // panics with `attempt to add with overflow` before - let (located_query_terms, _) = located_query_terms_from_tokens(&mut ctx, tokens, None)?; - assert!(located_query_terms.is_empty()); + let ExtractedTokens { query_terms, .. } = + located_query_terms_from_tokens(&mut ctx, tokens, None)?; + assert!(query_terms.is_empty()); Ok(()) } From 182cb42953acbd61bd21d4f9733b5b0bfb27f5f7 Mon Sep 17 00:00:00 2001 From: redistay Date: Tue, 2 Apr 2024 19:37:55 +0800 Subject: [PATCH 68/86] chore: fix some typos in conments Signed-off-by: redistay --- dump/src/reader/v2/updates.rs | 2 +- dump/src/writer.rs | 2 +- filter-parser/src/value.rs | 2 +- index-scheduler/src/autobatcher.rs | 2 +- milli/src/index.rs | 4 ++-- milli/src/update/words_prefixes_fst.rs | 2 +- 6 files changed, 7 insertions(+), 7 deletions(-) diff --git a/dump/src/reader/v2/updates.rs b/dump/src/reader/v2/updates.rs index bf1227c7a..82bde6bf1 100644 --- a/dump/src/reader/v2/updates.rs +++ b/dump/src/reader/v2/updates.rs @@ -61,7 +61,7 @@ pub enum IndexDocumentsMethod { #[cfg_attr(test, derive(serde::Serialize))] #[non_exhaustive] pub enum UpdateFormat { - /// The given update is a real **comma seperated** CSV with headers on the first line. + /// The given update is a real **comma separated** CSV with headers on the first line. Csv, /// The given update is a JSON array with documents inside. Json, diff --git a/dump/src/writer.rs b/dump/src/writer.rs index 3c8126876..ae89e7005 100644 --- a/dump/src/writer.rs +++ b/dump/src/writer.rs @@ -219,7 +219,7 @@ pub(crate) mod test { fn _create_directory_hierarchy(dir: &Path, depth: usize) -> String { let mut ret = String::new(); - // the entries are not guarenteed to be returned in the same order thus we need to sort them. + // the entries are not guaranteed to be returned in the same order thus we need to sort them. let mut entries = fs::read_dir(dir).unwrap().collect::, _>>().unwrap(); diff --git a/filter-parser/src/value.rs b/filter-parser/src/value.rs index 1d70cb025..b1ecac55b 100644 --- a/filter-parser/src/value.rs +++ b/filter-parser/src/value.rs @@ -42,7 +42,7 @@ fn quoted_by(quote: char, input: Span) -> IResult { ))); } } - // if it was preceeded by a `\` or if it was anything else we can continue to advance + // if it was preceded by a `\` or if it was anything else we can continue to advance } Ok(( diff --git a/index-scheduler/src/autobatcher.rs b/index-scheduler/src/autobatcher.rs index 8f3ffa146..dc184947c 100644 --- a/index-scheduler/src/autobatcher.rs +++ b/index-scheduler/src/autobatcher.rs @@ -870,7 +870,7 @@ mod tests { debug_snapshot!(autobatch_from(false,None, [doc_imp(UpdateDocuments, false, None), settings(false), idx_del()]), @"Some((IndexDeletion { ids: [0, 2, 1] }, false))"); debug_snapshot!(autobatch_from(false,None, [doc_imp(ReplaceDocuments,false, None), settings(false), doc_clr(), idx_del()]), @"Some((IndexDeletion { ids: [1, 3, 0, 2] }, false))"); debug_snapshot!(autobatch_from(false,None, [doc_imp(UpdateDocuments, false, None), settings(false), doc_clr(), idx_del()]), @"Some((IndexDeletion { ids: [1, 3, 0, 2] }, false))"); - // The third and final case is when the first task doesn't create an index but is directly followed by a task creating an index. In this case we can't batch whith what + // The third and final case is when the first task doesn't create an index but is directly followed by a task creating an index. In this case we can't batch whit what // follows because we first need to process the erronous batch. debug_snapshot!(autobatch_from(false,None, [doc_imp(ReplaceDocuments,false, None), settings(true), idx_del()]), @"Some((DocumentOperation { method: ReplaceDocuments, allow_index_creation: false, primary_key: None, operation_ids: [0] }, false))"); debug_snapshot!(autobatch_from(false,None, [doc_imp(UpdateDocuments, false, None), settings(true), idx_del()]), @"Some((DocumentOperation { method: UpdateDocuments, allow_index_creation: false, primary_key: None, operation_ids: [0] }, false))"); diff --git a/milli/src/index.rs b/milli/src/index.rs index d921de9e4..74f71ff29 100644 --- a/milli/src/index.rs +++ b/milli/src/index.rs @@ -1116,7 +1116,7 @@ impl Index { /* words prefixes fst */ - /// Writes the FST which is the words prefixes dictionnary of the engine. + /// Writes the FST which is the words prefixes dictionary of the engine. pub(crate) fn put_words_prefixes_fst>( &self, wtxn: &mut RwTxn, @@ -1129,7 +1129,7 @@ impl Index { ) } - /// Returns the FST which is the words prefixes dictionnary of the engine. + /// Returns the FST which is the words prefixes dictionary of the engine. pub fn words_prefixes_fst<'t>(&self, rtxn: &'t RoTxn) -> Result>> { match self.main.remap_types::().get(rtxn, main_key::WORDS_PREFIXES_FST_KEY)? { Some(bytes) => Ok(fst::Set::new(bytes)?.map_data(Cow::Borrowed)?), diff --git a/milli/src/update/words_prefixes_fst.rs b/milli/src/update/words_prefixes_fst.rs index bb1830727..8b438cef3 100644 --- a/milli/src/update/words_prefixes_fst.rs +++ b/milli/src/update/words_prefixes_fst.rs @@ -20,7 +20,7 @@ impl<'t, 'i> WordsPrefixesFst<'t, 'i> { /// Set the number of words required to make a prefix be part of the words prefixes /// database. If a word prefix is supposed to match more than this number of words in the - /// dictionnary, therefore this prefix is added to the words prefixes datastructures. + /// dictionary, therefore this prefix is added to the words prefixes datastructures. /// /// Default value is 100. This value must be higher than 50 and will be clamped /// to this bound otherwise. From d55d4962509d882a135dd62d62ad2596c3779519 Mon Sep 17 00:00:00 2001 From: Thomas Gauges Date: Tue, 2 Apr 2024 15:03:10 +0200 Subject: [PATCH 69/86] Fix milli/Cargo.toml for usage as dependency via git --- Cargo.lock | 25 +++++++++++++++++-------- milli/Cargo.toml | 6 +++--- 2 files changed, 20 insertions(+), 11 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index aba12d3b8..d358bf10b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -555,6 +555,17 @@ dependencies = [ "syn 2.0.48", ] +[[package]] +name = "bindgen_cuda" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1f8489af5b7d17a81bffe37e0f4d6e1e4de87c87329d05447f22c35d95a1227d" +dependencies = [ + "glob", + "num_cpus", + "rayon", +] + [[package]] name = "bit-set" version = "0.5.3" @@ -731,7 +742,7 @@ dependencies = [ [[package]] name = "candle-core" version = "0.3.3" -source = "git+https://github.com/huggingface/candle.git#5270224f407502b82fe90bc2622894ce3871b002" +source = "git+https://github.com/huggingface/candle.git?rev=5270224f407502b82fe90bc2622894ce3871b002#5270224f407502b82fe90bc2622894ce3871b002" dependencies = [ "byteorder", "candle-kernels", @@ -752,18 +763,16 @@ dependencies = [ [[package]] name = "candle-kernels" -version = "0.3.1" -source = "git+https://github.com/huggingface/candle.git#f4fcf6090045ac44122fd5f0a7e46db6e3e16528" +version = "0.3.3" +source = "git+https://github.com/huggingface/candle.git?rev=5270224f407502b82fe90bc2622894ce3871b002#5270224f407502b82fe90bc2622894ce3871b002" dependencies = [ - "anyhow", - "glob", - "rayon", + "bindgen_cuda", ] [[package]] name = "candle-nn" version = "0.3.3" -source = "git+https://github.com/huggingface/candle.git#5270224f407502b82fe90bc2622894ce3871b002" +source = "git+https://github.com/huggingface/candle.git?rev=5270224f407502b82fe90bc2622894ce3871b002#5270224f407502b82fe90bc2622894ce3871b002" dependencies = [ "candle-core", "half 2.3.1", @@ -777,7 +786,7 @@ dependencies = [ [[package]] name = "candle-transformers" version = "0.3.3" -source = "git+https://github.com/huggingface/candle.git#5270224f407502b82fe90bc2622894ce3871b002" +source = "git+https://github.com/huggingface/candle.git?rev=5270224f407502b82fe90bc2622894ce3871b002#5270224f407502b82fe90bc2622894ce3871b002" dependencies = [ "byteorder", "candle-core", diff --git a/milli/Cargo.toml b/milli/Cargo.toml index 1dfa495ea..ce6392b59 100644 --- a/milli/Cargo.toml +++ b/milli/Cargo.toml @@ -71,9 +71,9 @@ itertools = "0.11.0" puffin = "0.16.0" csv = "1.3.0" -candle-core = { git = "https://github.com/huggingface/candle.git", version = "0.3.1" } -candle-transformers = { git = "https://github.com/huggingface/candle.git", version = "0.3.1" } -candle-nn = { git = "https://github.com/huggingface/candle.git", version = "0.3.1" } +candle-core = { git = "https://github.com/huggingface/candle.git", rev="5270224f407502b82fe90bc2622894ce3871b002", version = "0.3.3" } +candle-transformers = { git = "https://github.com/huggingface/candle.git", rev="5270224f407502b82fe90bc2622894ce3871b002", version = "0.3.3" } +candle-nn = { git = "https://github.com/huggingface/candle.git", rev="5270224f407502b82fe90bc2622894ce3871b002", version = "0.3.3" } tokenizers = { git = "https://github.com/huggingface/tokenizers.git", tag = "v0.14.1", version = "0.14.1", default_features = false, features = [ "onig", ] } From a1eccc762a69b2b4feea75e8c51e11e24b0b4273 Mon Sep 17 00:00:00 2001 From: Louis Dureuil Date: Wed, 3 Apr 2024 11:05:59 +0200 Subject: [PATCH 70/86] Prefer safetensors to pytorch when both are available --- milli/src/vector/hf.rs | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/milli/src/vector/hf.rs b/milli/src/vector/hf.rs index e341a553e..4c3c0670c 100644 --- a/milli/src/vector/hf.rs +++ b/milli/src/vector/hf.rs @@ -87,11 +87,11 @@ impl Embedder { let config = api.get("config.json").map_err(NewEmbedderError::api_get)?; let tokenizer = api.get("tokenizer.json").map_err(NewEmbedderError::api_get)?; let (weights, source) = { - api.get("pytorch_model.bin") - .map(|filename| (filename, WeightSource::Pytorch)) + api.get("model.safetensors") + .map(|filename| (filename, WeightSource::Safetensors)) .or_else(|_| { - api.get("model.safetensors") - .map(|filename| (filename, WeightSource::Safetensors)) + api.get("pytorch_model.bin") + .map(|filename| (filename, WeightSource::Pytorch)) }) .map_err(NewEmbedderError::api_get)? }; From 58cafcc8240c74ff14d067705fd8476c9ef9daf2 Mon Sep 17 00:00:00 2001 From: Louis Dureuil Date: Wed, 3 Apr 2024 13:11:56 +0200 Subject: [PATCH 71/86] Update candle --- Cargo.lock | 373 ++++++++++++++++++++++------------------------- milli/Cargo.toml | 8 +- 2 files changed, 182 insertions(+), 199 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 3a5e8b987..a11312639 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -47,7 +47,7 @@ dependencies = [ "actix-utils", "ahash", "base64 0.21.7", - "bitflags 2.4.1", + "bitflags 2.5.0", "brotli", "bytes", "bytestring", @@ -80,7 +80,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e01ed3140b2f8d422c68afa1ed2e85d996ea619c988ac834d255db32138655cb" dependencies = [ "quote", - "syn 2.0.48", + "syn 2.0.58", ] [[package]] @@ -246,9 +246,9 @@ checksum = "f26201604c87b1e01bd3d98f8d5d9a8fcbb815e8cedb41ffccbeb4bf593a35fe" [[package]] name = "aes" -version = "0.8.3" +version = "0.8.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ac1f845298e95f983ff1944b728ae08b8cebab80d684f0a832ed0fc74dfa27e2" +checksum = "b169f7a6d4742236a0a00c541b845991d0ac43e546831af1249753ab4c3aa3a0" dependencies = [ "cfg-if", "cipher", @@ -386,7 +386,7 @@ dependencies = [ "byteorder", "heed", "log", - "memmap2 0.9.3", + "memmap2 0.9.4", "ordered-float", "rand", "rayon", @@ -424,7 +424,7 @@ checksum = "16e62a023e7c117e27523144c5d2459f4397fcc3cab0085af8e2224f643a0193" dependencies = [ "proc-macro2", "quote", - "syn 2.0.48", + "syn 2.0.58", ] [[package]] @@ -435,7 +435,7 @@ checksum = "c980ee35e870bd1a4d2c8294d4c04d0499e67bca1e4b5cefcc693c2fa00caea9" dependencies = [ "proc-macro2", "quote", - "syn 2.0.48", + "syn 2.0.58", ] [[package]] @@ -455,9 +455,9 @@ dependencies = [ [[package]] name = "autocfg" -version = "1.1.0" +version = "1.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d468802bab17cbc0cc575e9b053f41e72aa36bfa6b7f55e3529ffa43161b97fa" +checksum = "f1fdabc7756949593fe60f30ec81974b613357de856987752631dea1e3394c80" [[package]] name = "backtrace" @@ -541,7 +541,7 @@ version = "0.68.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "726e4313eb6ec35d2730258ad4e15b547ee75d6afaa1361a922e78e59b7d8078" dependencies = [ - "bitflags 2.4.1", + "bitflags 2.5.0", "cexpr", "clang-sys", "lazy_static", @@ -552,7 +552,7 @@ dependencies = [ "regex", "rustc-hash", "shlex", - "syn 2.0.48", + "syn 2.0.58", ] [[package]] @@ -589,9 +589,9 @@ checksum = "bef38d45163c2f1dde094a7dfd33ccf595c92905c8f8f4fdc18d06fb1037718a" [[package]] name = "bitflags" -version = "2.4.1" +version = "2.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "327762f6e5a765692301e5bb513e0d9fef63be86bbc14528052b1cd3e6f03e07" +checksum = "cf4b9d6a944f767f8e5e0db018570623c85f3d925ac718db4e06d0187adb21c1" dependencies = [ "serde", ] @@ -648,9 +648,9 @@ dependencies = [ [[package]] name = "bumpalo" -version = "3.13.0" +version = "3.15.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a3e2c3daef883ecc1b5d58c15adae93470a91d425f3532ba1695849656af3fc1" +checksum = "7ff69b9dd49fd426c69a0db9fc04dd934cdb6645ff000864d98f7e2af8830eaa" [[package]] name = "byte-unit" @@ -670,22 +670,22 @@ checksum = "2c676a478f63e9fa2dd5368a42f28bba0d6c560b775f38583c8bbaa7fcd67c9c" [[package]] name = "bytemuck" -version = "1.14.0" +version = "1.15.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "374d28ec25809ee0e23827c2ab573d729e293f281dfe393500e7ad618baa61c6" +checksum = "5d6d68c57235a3a081186990eca2867354726650f42f7516ca50c28d6281fd15" dependencies = [ "bytemuck_derive", ] [[package]] name = "bytemuck_derive" -version = "1.4.1" +version = "1.6.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fdde5c9cd29ebd706ce1b35600920a33550e402fc998a2e53ad3b42c3c47a192" +checksum = "4da9a32f3fed317401fa3c862968128267c3106685286e15d5aaa3d7389c2f60" dependencies = [ "proc-macro2", "quote", - "syn 2.0.48", + "syn 2.0.58", ] [[package]] @@ -741,15 +741,16 @@ dependencies = [ [[package]] name = "candle-core" -version = "0.3.3" -source = "git+https://github.com/huggingface/candle.git?rev=5270224f407502b82fe90bc2622894ce3871b002#5270224f407502b82fe90bc2622894ce3871b002" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6f1b20174c1707e20f4cb364a355b449803c03e9b0c9193324623cf9787a4e00" dependencies = [ "byteorder", "candle-kernels", "cudarc", "gemm", - "half 2.3.1", - "memmap2 0.9.3", + "half 2.4.0", + "memmap2 0.9.4", "num-traits", "num_cpus", "rand", @@ -763,19 +764,21 @@ dependencies = [ [[package]] name = "candle-kernels" -version = "0.3.3" -source = "git+https://github.com/huggingface/candle.git?rev=5270224f407502b82fe90bc2622894ce3871b002#5270224f407502b82fe90bc2622894ce3871b002" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5845911a44164ebb73b56a0e23793ba1b583bad102af7400fe4768babc5815b2" dependencies = [ "bindgen_cuda", ] [[package]] name = "candle-nn" -version = "0.3.3" -source = "git+https://github.com/huggingface/candle.git?rev=5270224f407502b82fe90bc2622894ce3871b002#5270224f407502b82fe90bc2622894ce3871b002" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "66a27533c8edfc915a6459f9850641ef523a829fa1a181c670766c1f752d873a" dependencies = [ "candle-core", - "half 2.3.1", + "half 2.4.0", "num-traits", "rayon", "safetensors", @@ -785,8 +788,9 @@ dependencies = [ [[package]] name = "candle-transformers" -version = "0.3.3" -source = "git+https://github.com/huggingface/candle.git?rev=5270224f407502b82fe90bc2622894ce3871b002#5270224f407502b82fe90bc2622894ce3871b002" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b5847699f0643da05e57fc473672566e93dc36d82c1b7eeb970c6154d3434fe1" dependencies = [ "byteorder", "candle-core", @@ -798,7 +802,6 @@ dependencies = [ "serde_json", "serde_plain", "tracing", - "wav", ] [[package]] @@ -842,9 +845,9 @@ checksum = "37b2a672a2cb129a2e41c10b1224bb368f9f37a2b16b612598138befd7b37eb5" [[package]] name = "cc" -version = "1.0.83" +version = "1.0.90" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f1174fb0b6ec23863f8b971027804a42614e347eafb0a95bf0b12cdae21fc4d0" +checksum = "8cd6604a82acf3039f1144f54b8eb34e91ffba622051189e71b781822d5ee1f5" dependencies = [ "jobserver", "libc", @@ -991,7 +994,7 @@ dependencies = [ "heck", "proc-macro2", "quote", - "syn 2.0.48", + "syn 2.0.58", ] [[package]] @@ -1098,18 +1101,18 @@ checksum = "79bb3adfaf5f75d24b01aee375f7555907840fa2800e5ec8fa3b9e2031830173" [[package]] name = "cpufeatures" -version = "0.2.9" +version = "0.2.12" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a17b76ff3a4162b0b27f354a0c87015ddad39d35f9c0c36607a3bdd175dde1f1" +checksum = "53fe5e26ff1b7aef8bca9c6080520cfb8d9333c7568e1829cef191a9723e5504" dependencies = [ "libc", ] [[package]] name = "crc32fast" -version = "1.3.2" +version = "1.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b540bd8bc810d3885c6ea91e2018302f68baba2129ab3e88f32389ee9370880d" +checksum = "b3855a8a784b474f333699ef2bbca9db2c4a1f6d9088a90a2d25b1eb53111eaa" dependencies = [ "cfg-if", ] @@ -1255,7 +1258,7 @@ version = "0.10.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9395df0cab995685664e79cc35ad6302bf08fb9c5d82301875a183affe1278b1" dependencies = [ - "half 2.3.1", + "half 2.4.0", ] [[package]] @@ -1303,7 +1306,7 @@ dependencies = [ "proc-macro2", "quote", "strsim", - "syn 2.0.48", + "syn 2.0.58", ] [[package]] @@ -1325,7 +1328,7 @@ checksum = "836a9bbc7ad63342d6d6e7b815ccab164bc77a2d95d84bc3117a8c0d5c98e2d5" dependencies = [ "darling_core 0.20.3", "quote", - "syn 2.0.48", + "syn 2.0.58", ] [[package]] @@ -1339,9 +1342,9 @@ dependencies = [ [[package]] name = "deranged" -version = "0.3.9" +version = "0.3.11" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0f32d04922c60427da6f9fef14d042d9edddef64cb9d4ce0d64d0685fbeb1fd3" +checksum = "b42b6fa04a440b495c8b04d0e71b707c585f83cb9cb28cf8cd0d976c315e31b4" dependencies = [ "powerfmt", "serde", @@ -1355,7 +1358,7 @@ checksum = "67e77553c4162a157adbf834ebae5b415acbecbeafc7a74b0e886657506a7611" dependencies = [ "proc-macro2", "quote", - "syn 2.0.48", + "syn 2.0.58", ] [[package]] @@ -1459,7 +1462,7 @@ dependencies = [ "convert_case 0.6.0", "proc-macro2", "quote", - "syn 2.0.48", + "syn 2.0.58", ] [[package]] @@ -1573,9 +1576,9 @@ dependencies = [ [[package]] name = "either" -version = "1.9.0" +version = "1.10.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a26ae43d7bcc3b814de94796a5e736d4029efb0ee900c12e2d54c993ad1a1e07" +checksum = "11157ac094ffbdde99aa67b23417ebdd801842852b500e395a45a9c0aac03e4a" dependencies = [ "serde", ] @@ -1677,7 +1680,7 @@ dependencies = [ "heck", "proc-macro2", "quote", - "syn 2.0.48", + "syn 2.0.58", ] [[package]] @@ -1697,7 +1700,7 @@ checksum = "03cdc46ec28bd728e67540c528013c6a10eb69a02eb31078a1bda695438cbfb8" dependencies = [ "proc-macro2", "quote", - "syn 2.0.48", + "syn 2.0.58", ] [[package]] @@ -1780,7 +1783,7 @@ dependencies = [ "darling 0.20.3", "proc-macro2", "quote", - "syn 2.0.48", + "syn 2.0.58", "uuid", ] @@ -1912,7 +1915,7 @@ checksum = "87750cf4b7a4c0625b1529e4c543c2182106e4dedc60a2a6455e00d212c489ac" dependencies = [ "proc-macro2", "quote", - "syn 2.0.48", + "syn 2.0.58", ] [[package]] @@ -1973,7 +1976,7 @@ version = "0.6.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "27d12c0aed7f1e24276a241aadc4cb8ea9f83000f34bc062b7cc2d51e3b0fabd" dependencies = [ - "bitflags 2.4.1", + "bitflags 2.5.0", "debugid", "fxhash", "serde", @@ -1982,9 +1985,9 @@ dependencies = [ [[package]] name = "gemm" -version = "0.17.0" +version = "0.17.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e97d506c68f4fb12325b52a638e7d54cc87e3593a4ded0de60218b6dfd65f645" +checksum = "6ab24cc62135b40090e31a76a9b2766a501979f3070fa27f689c27ec04377d32" dependencies = [ "dyn-stack", "gemm-c32", @@ -2002,9 +2005,9 @@ dependencies = [ [[package]] name = "gemm-c32" -version = "0.17.0" +version = "0.17.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0dd16f26e8f34661edc906d8c9522b59ec1655c865a98a58950d0246eeaca9da" +checksum = "b9c030d0b983d1e34a546b86e08f600c11696fde16199f971cd46c12e67512c0" dependencies = [ "dyn-stack", "gemm-common", @@ -2017,9 +2020,9 @@ dependencies = [ [[package]] name = "gemm-c64" -version = "0.17.0" +version = "0.17.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a8e34381bc060b47fbd25522a281799ef763cd27f43bbd1783d935774659242a" +checksum = "fbb5f2e79fefb9693d18e1066a557b4546cd334b226beadc68b11a8f9431852a" dependencies = [ "dyn-stack", "gemm-common", @@ -2032,13 +2035,13 @@ dependencies = [ [[package]] name = "gemm-common" -version = "0.17.0" +version = "0.17.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "22518a76339b09276f77c3166c44262e55f633712fe8a44fd0573505887feeab" +checksum = "a2e7ea062c987abcd8db95db917b4ffb4ecdfd0668471d8dc54734fdff2354e8" dependencies = [ "bytemuck", "dyn-stack", - "half 2.3.1", + "half 2.4.0", "num-complex", "num-traits", "once_cell", @@ -2052,14 +2055,14 @@ dependencies = [ [[package]] name = "gemm-f16" -version = "0.17.0" +version = "0.17.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "70409bbf3ef83b38cbe4a58cd4b797c1c27902505bdd926a588ea61b6c550a84" +checksum = "7ca4c06b9b11952071d317604acb332e924e817bd891bec8dfb494168c7cedd4" dependencies = [ "dyn-stack", "gemm-common", "gemm-f32", - "half 2.3.1", + "half 2.4.0", "num-complex", "num-traits", "paste", @@ -2070,9 +2073,9 @@ dependencies = [ [[package]] name = "gemm-f32" -version = "0.17.0" +version = "0.17.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5ea3068edca27f100964157211782eba19e961aa4d0d2bdac3e1775a51aa7680" +checksum = "e9a69f51aaefbd9cf12d18faf273d3e982d9d711f60775645ed5c8047b4ae113" dependencies = [ "dyn-stack", "gemm-common", @@ -2085,9 +2088,9 @@ dependencies = [ [[package]] name = "gemm-f64" -version = "0.17.0" +version = "0.17.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5fd41e8f5a60dce8d8acd852a3f4b22f8e18be957e1937731be692c037652510" +checksum = "aa397a48544fadf0b81ec8741e5c0fba0043008113f71f2034def1935645d2b0" dependencies = [ "dyn-stack", "gemm-common", @@ -2116,9 +2119,9 @@ checksum = "36d244a08113319b5ebcabad2b8b7925732d15eec46d7e7ac3c11734f3b7a6ad" [[package]] name = "getrandom" -version = "0.2.10" +version = "0.2.12" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "be4136b2a15dd319360be1c07d9933517ccf0be8f16bf62a3bee4f0d618df427" +checksum = "190092ea657667030ac6a35e305e62fc4dd69fd98ac98631e5d3a2b1575a12b5" dependencies = [ "cfg-if", "js-sys", @@ -2151,7 +2154,7 @@ version = "0.18.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1b3ba52851e73b46a4c3df1d89343741112003f0f6f13beb0dfac9e457c3fdcd" dependencies = [ - "bitflags 2.4.1", + "bitflags 2.5.0", "libc", "libgit2-sys", "log", @@ -2203,9 +2206,9 @@ checksum = "eabb4a44450da02c90444cf74558da904edde8fb4e9035a9a6a4e15445af0bd7" [[package]] name = "half" -version = "2.3.1" +version = "2.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bc52e53916c08643f1b56ec082790d1e86a32e58dc5268f897f313fbae7b4872" +checksum = "b5eceaaeec696539ddaf7b333340f1af35a5aa87ae3e4f3ead0532f72affab2e" dependencies = [ "bytemuck", "cfg-if", @@ -2259,7 +2262,7 @@ version = "0.20.0-alpha.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9648a50991c86df7d00c56c268c27754fcf4c80be2ba57fc4a00dc928c6fe934" dependencies = [ - "bitflags 2.4.1", + "bitflags 2.5.0", "bytemuck", "byteorder", "heed-traits", @@ -2293,9 +2296,9 @@ dependencies = [ [[package]] name = "hermit-abi" -version = "0.3.2" +version = "0.3.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "443144c8cdadd93ebf52ddb4056d257f5b52c04d3c804e657d19eb73fc33668b" +checksum = "d231dfb89cfffdbc30e7fc41579ed6066ad03abda9e567ccafae602b97ec5024" [[package]] name = "hex" @@ -2474,9 +2477,9 @@ dependencies = [ [[package]] name = "indexmap" -version = "2.2.5" +version = "2.2.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7b0b929d511467233429c45a44ac1dcaa21ba0f5ba11e4879e6ed28ddb4f9df4" +checksum = "168fb715dda47215e360912c096649d23d58bf392ac62f73919e831745e40f26" dependencies = [ "equivalent", "hashbrown", @@ -2588,10 +2591,19 @@ dependencies = [ ] [[package]] -name = "itoa" -version = "1.0.9" +name = "itertools" +version = "0.12.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "af150ab688ff2122fcef229be89cb50dd66af9e01a4ff320cc137eecc9bacc38" +checksum = "ba291022dbbd398a455acf126c1e341954079855bc60dfdda641363bd6922569" +dependencies = [ + "either", +] + +[[package]] +name = "itoa" +version = "1.0.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "49f1f14873335454500d59611f1cf4a4b0f786f9ac11f4312a78e4cf2566695b" [[package]] name = "jieba-rs" @@ -2610,18 +2622,18 @@ dependencies = [ [[package]] name = "jobserver" -version = "0.1.26" +version = "0.1.28" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "936cfd212a0155903bcbc060e316fb6cc7cbf2e1907329391ebadc1fe0ce77c2" +checksum = "ab46a6e9526ddef3ae7f787c06f0f2600639ba80ea3eade3d8e670a2230f51d6" dependencies = [ "libc", ] [[package]] name = "js-sys" -version = "0.3.64" +version = "0.3.69" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c5f195fe497f702db0f318b07fdd68edb16955aed830df8363d837542f8f935a" +checksum = "29c15563dc2726973df627357ce0c9ddddbea194836909d655df6a75d2cf296d" dependencies = [ "wasm-bindgen", ] @@ -2716,9 +2728,9 @@ dependencies = [ [[package]] name = "libm" -version = "0.2.7" +version = "0.2.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f7012b1bbb0719e1097c47611d3898568c546d597c2e74d66f6087edd5233ff4" +checksum = "4ec2a862134d2a7d32d7983ddcdd1c4923530833c9f2ea1a44fc5fa473989058" [[package]] name = "libmimalloc-sys" @@ -3018,7 +3030,7 @@ checksum = "fc2fb41a9bb4257a3803154bdf7e2df7d45197d1941c9b1a90ad815231630721" dependencies = [ "proc-macro2", "quote", - "syn 2.0.48", + "syn 2.0.58", ] [[package]] @@ -3084,9 +3096,9 @@ dependencies = [ [[package]] name = "log" -version = "0.4.20" +version = "0.4.21" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b5e6163cb8c49088c2c36f57875e58ccd8c87c7427f7fbd50ea6710b2f3f2e8f" +checksum = "90ed8c1e510134f979dbc4f070f87d4313098b704861a105fe34231c70a3901c" [[package]] name = "lz4_flex" @@ -3119,7 +3131,7 @@ dependencies = [ "once_cell", "proc-macro2", "quote", - "syn 2.0.48", + "syn 2.0.58", ] [[package]] @@ -3316,9 +3328,9 @@ dependencies = [ [[package]] name = "memmap2" -version = "0.9.3" +version = "0.9.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "45fd3a57831bf88bc63f8cebc0cf956116276e97fef3966103e96416209f7c92" +checksum = "fe751422e4a8caa417e13c3ea66452215d7d63e19e604f4980461212f3ae1322" dependencies = [ "libc", "stable_deref_trait", @@ -3422,9 +3434,9 @@ checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a" [[package]] name = "miniz_oxide" -version = "0.7.1" +version = "0.7.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e7810e0be55b428ada41041c41f32c9f1a42817901b4ccf45fa3d4b6561e74c7" +checksum = "9d811f3e15f28568be3407c8e7fdb6514c1cda3cb30683f15b6a1a1dc4ea14a7" dependencies = [ "adler", ] @@ -3459,7 +3471,7 @@ checksum = "371717c0a5543d6a800cac822eac735aa7d2d2fbb41002e9856a4089532dbdce" dependencies = [ "proc-macro2", "quote", - "syn 2.0.48", + "syn 2.0.58", ] [[package]] @@ -3521,9 +3533,9 @@ dependencies = [ [[package]] name = "num-complex" -version = "0.4.4" +version = "0.4.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1ba157ca0885411de85d6ca030ba7e2a83a28636056c7c699b07c8b6f7383214" +checksum = "23c6602fda94a57c990fe0df199a035d83576b496aa29f4e634a8ac6004e68a6" dependencies = [ "bytemuck", "num-traits", @@ -3547,9 +3559,9 @@ dependencies = [ [[package]] name = "num-traits" -version = "0.2.16" +version = "0.2.18" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f30b0abd723be7e2ffca1272140fac1a2f084c77ec3e123c192b66af1ee9e6c2" +checksum = "da0df0e5185db44f69b44f26786fe401b6c293d1907744beaa7fa62b2e5a517a" dependencies = [ "autocfg", "libm", @@ -3803,7 +3815,7 @@ dependencies = [ "pest_meta", "proc-macro2", "quote", - "syn 2.0.48", + "syn 2.0.58", ] [[package]] @@ -3857,7 +3869,7 @@ dependencies = [ "phf_shared", "proc-macro2", "quote", - "syn 2.0.48", + "syn 2.0.58", ] [[package]] @@ -3886,7 +3898,7 @@ checksum = "266c042b60c9c76b8d53061e52b2e0d1116abc57cefc8c5cd671619a56ac3690" dependencies = [ "proc-macro2", "quote", - "syn 2.0.48", + "syn 2.0.58", ] [[package]] @@ -3909,9 +3921,9 @@ checksum = "16f2611cd06a1ac239a0cea4521de9eb068a6ca110324ee00631aa68daa74fc0" [[package]] name = "pkg-config" -version = "0.3.27" +version = "0.3.30" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "26072860ba924cbfa98ea39c8c19b4dd6a4a25423dbdf219c1eca91aa0cf6964" +checksum = "d231b230927b5e4ad203db57bbcbee2802f6bce620b1e4a9024a07d94e2907ec" [[package]] name = "platform-dirs" @@ -3994,9 +4006,9 @@ dependencies = [ [[package]] name = "proc-macro2" -version = "1.0.76" +version = "1.0.79" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "95fc56cda0b5c3325f5fbbd7ff9fda9e02bb00bb3dac51252d2f1bfa1cb8cc8c" +checksum = "e835ff2298f5721608eb1a980ecaee1aef2c132bf95ecc026a11b7bf3c01c02e" dependencies = [ "unicode-ident", ] @@ -4056,9 +4068,9 @@ dependencies = [ [[package]] name = "pulp" -version = "0.18.4" +version = "0.18.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7057c1435edb390ebfc51743abad043377f1f698ce8e649a9b52a4b378be5e4d" +checksum = "03457ac216146f43f921500bac4e892d5cd32b0479b929cbfc90f95cd6c599c2" dependencies = [ "bytemuck", "libm", @@ -4136,9 +4148,9 @@ dependencies = [ [[package]] name = "rayon" -version = "1.8.1" +version = "1.10.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fa7237101a77a10773db45d62004a272517633fbcc3df19d96455ede1122e051" +checksum = "b418a60154510ca1a002a752ca9714984e21e4241e804d32555251faf8b78ffa" dependencies = [ "either", "rayon-core", @@ -4209,7 +4221,7 @@ dependencies = [ "aho-corasick", "memchr", "regex-automata", - "regex-syntax 0.8.2", + "regex-syntax", ] [[package]] @@ -4220,15 +4232,9 @@ checksum = "5f804c7828047e88b2d32e2d7fe5a105da8ee3264f01902f796c8e067dc2483f" dependencies = [ "aho-corasick", "memchr", - "regex-syntax 0.8.2", + "regex-syntax", ] -[[package]] -name = "regex-syntax" -version = "0.7.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e5ea92a5b6195c6ef2a0295ea818b312502c6fc94dde986c5553242e18fd4ce2" - [[package]] name = "regex-syntax" version = "0.8.2" @@ -4283,12 +4289,6 @@ version = "0.1.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8c31b5c4033f8fdde8700e4657be2c497e7288f01515be52168c631e2e4d4086" -[[package]] -name = "riff" -version = "1.0.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b9b1a3d5f46d53f4a3478e2be4a5a5ce5108ea58b100dcd139830eae7f79a3a1" - [[package]] name = "ring" version = "0.17.8" @@ -4369,7 +4369,7 @@ version = "0.38.31" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ea3e1a662af26cd7a3ba09c0297a31af215563ecf42817c98df621387f4e949" dependencies = [ - "bitflags 2.4.1", + "bitflags 2.5.0", "errno", "libc", "linux-raw-sys 0.4.12", @@ -4446,15 +4446,15 @@ checksum = "7ffc183a10b4478d04cbbbfc96d0873219d962dd5accaff2ffbd4ceb7df837f4" [[package]] name = "ryu" -version = "1.0.15" +version = "1.0.17" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1ad4cc8da4ef723ed60bced201181d83791ad433213d8c24efffda1eec85d741" +checksum = "e86697c916019a8588c99b5fac3cead74ec0b4b819707a682fd4d23fa0ce1ba1" [[package]] name = "safetensors" -version = "0.4.1" +version = "0.4.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1659ef1c27917eb58c2d53664b5506d0b68c9cb9b460d3e0901011cf71269a8e" +checksum = "8d980e6bfb34436fb0a81e42bc41af43f11805bbbca443e7f68e9faaabe669ed" dependencies = [ "serde", "serde_json", @@ -4540,14 +4540,14 @@ checksum = "7eb0b34b42edc17f6b7cac84a52a1c5f0e1bb2227e997ca9011ea3dd34e8610b" dependencies = [ "proc-macro2", "quote", - "syn 2.0.48", + "syn 2.0.58", ] [[package]] name = "serde_json" -version = "1.0.114" +version = "1.0.115" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c5f09b1bd632ef549eaa9f60a1f8de742bdbc698e6cee2095fc84dde5f549ae0" +checksum = "12dc5c46daa8e9fdf4f5e71b6cf9a53f2487da0e86e55808e2d35539666497dd" dependencies = [ "indexmap", "itoa", @@ -4598,9 +4598,9 @@ dependencies = [ [[package]] name = "sha1" -version = "0.10.5" +version = "0.10.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f04293dc80c3993519f2d7f6f511707ee7094fe0c6d3406feb330cdb3540eba3" +checksum = "e3bf829a2d51ab4a5ddf1352d8470c140cadc8301b2ae1789db023f01cedd6ba" dependencies = [ "cfg-if", "cpufeatures", @@ -4814,9 +4814,9 @@ dependencies = [ [[package]] name = "syn" -version = "2.0.48" +version = "2.0.58" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0f3531638e407dfc0814761abb7c00a5b54992b849452a0646b7f65c9f770f3f" +checksum = "44cfb93f38070beee36b3fef7d4f5a16f27751d94b187b666a5cc5e9b0d30687" dependencies = [ "proc-macro2", "quote", @@ -4834,14 +4834,13 @@ dependencies = [ [[package]] name = "synstructure" -version = "0.13.0" +version = "0.13.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "285ba80e733fac80aa4270fbcdf83772a79b80aa35c97075320abfee4a915b06" +checksum = "c8af7666ab7b6390ab78131fb5b0fce11d6b7a6951602017c35fa82800708971" dependencies = [ "proc-macro2", "quote", - "syn 2.0.48", - "unicode-xid", + "syn 2.0.58", ] [[package]] @@ -4850,7 +4849,7 @@ version = "0.5.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ec7dddc5f0fee506baf8b9fdb989e242f17e4b11c61dfbb0635b705217199eea" dependencies = [ - "bitflags 2.4.1", + "bitflags 2.5.0", "byteorder", "enum-as-inner", "libc", @@ -4952,7 +4951,7 @@ checksum = "c61f3ba182994efc43764a46c018c347bc492c79f024e705f46567b418f6d4f7" dependencies = [ "proc-macro2", "quote", - "syn 2.0.48", + "syn 2.0.58", ] [[package]] @@ -5040,14 +5039,14 @@ checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20" [[package]] name = "tokenizers" -version = "0.14.1" -source = "git+https://github.com/huggingface/tokenizers.git?tag=v0.14.1#6357206cdcce4d78ffb1e0372feb456caea09375" +version = "0.15.2" +source = "git+https://github.com/huggingface/tokenizers.git?tag=v0.15.2#701a73b869602b5639589d197e805349cdba3223" dependencies = [ "aho-corasick", "derive_builder 0.12.0", "esaxx-rs", "getrandom", - "itertools 0.11.0", + "itertools 0.12.1", "lazy_static", "log", "macro_rules_attribute", @@ -5058,7 +5057,7 @@ dependencies = [ "rayon", "rayon-cond", "regex", - "regex-syntax 0.7.4", + "regex-syntax", "serde", "serde_json", "spm_precompiled", @@ -5095,7 +5094,7 @@ checksum = "5b8a1e28f2deaa14e508979454cb3a223b10b938b45af148bc0986de36f1923b" dependencies = [ "proc-macro2", "quote", - "syn 2.0.48", + "syn 2.0.58", ] [[package]] @@ -5206,7 +5205,7 @@ checksum = "34704c8d6ebcbc939824180af020566b01a7c01f80641264eba0999f6c2b6be7" dependencies = [ "proc-macro2", "quote", - "syn 2.0.48", + "syn 2.0.58", ] [[package]] @@ -5291,9 +5290,9 @@ checksum = "3528ecfd12c466c6f163363caf2d02a71161dd5e1cc6ae7b34207ea2d42d81ed" [[package]] name = "typenum" -version = "1.16.0" +version = "1.17.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "497961ef93d974e23eb6f433eb5fe1b7930b659f06d12dec6fc44a8f554c0bba" +checksum = "42ff0bf0c66b8238c6f3b578df37d0b7848e55df8577b3f74f92a69acceeb825" [[package]] name = "ucd-trie" @@ -5327,9 +5326,9 @@ checksum = "92888ba5573ff080736b3648696b70cafad7d250551175acbaa4e0385b3e1460" [[package]] name = "unicode-ident" -version = "1.0.11" +version = "1.0.12" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "301abaae475aa91687eb82514b328ab47a211a533026cb25fc3e519b86adfc3c" +checksum = "3354b9ac3fae1ff6755cb6db53683adb661634f67557942dea4facebec0fee4b" [[package]] name = "unicode-normalization" @@ -5361,12 +5360,6 @@ version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e51733f11c9c4f72aa0c160008246859e340b00807569a0da0e7a1079b27ba85" -[[package]] -name = "unicode-xid" -version = "0.2.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f962df74c8c05a667b5ee8bcf162993134c104e96440b663c8daa176dc772d8c" - [[package]] name = "unicode_categories" version = "0.1.1" @@ -5500,9 +5493,9 @@ checksum = "49874b5167b65d7193b8aba1567f5c7d93d001cafc34600cee003eda787e483f" [[package]] name = "walkdir" -version = "2.4.0" +version = "2.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d71d857dc86794ca4c280d616f7da00d2dbfd8cd788846559a6813e6aa4b54ee" +checksum = "29790946404f91d9c5d06f9874efddea1dc06c5efe94541a7d6863108e3a5e4b" dependencies = [ "same-file", "winapi-util", @@ -5536,9 +5529,9 @@ checksum = "9c8d87e72b64a3b4db28d11ce29237c246188f4f51057d65a7eab63b7987e423" [[package]] name = "wasm-bindgen" -version = "0.2.87" +version = "0.2.92" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7706a72ab36d8cb1f80ffbf0e071533974a60d0a308d01a5d0375bf60499a342" +checksum = "4be2531df63900aeb2bca0daaaddec08491ee64ceecbee5076636a3b026795a8" dependencies = [ "cfg-if", "wasm-bindgen-macro", @@ -5546,16 +5539,16 @@ dependencies = [ [[package]] name = "wasm-bindgen-backend" -version = "0.2.87" +version = "0.2.92" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5ef2b6d3c510e9625e5fe6f509ab07d66a760f0885d858736483c32ed7809abd" +checksum = "614d787b966d3989fa7bb98a654e369c762374fd3213d212cfc0251257e747da" dependencies = [ "bumpalo", "log", "once_cell", "proc-macro2", "quote", - "syn 2.0.48", + "syn 2.0.58", "wasm-bindgen-shared", ] @@ -5573,9 +5566,9 @@ dependencies = [ [[package]] name = "wasm-bindgen-macro" -version = "0.2.87" +version = "0.2.92" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dee495e55982a3bd48105a7b947fd2a9b4a8ae3010041b9e0faab3f9cd028f1d" +checksum = "a1f8823de937b71b9460c0c34e25f3da88250760bec0ebac694b49997550d726" dependencies = [ "quote", "wasm-bindgen-macro-support", @@ -5583,22 +5576,22 @@ dependencies = [ [[package]] name = "wasm-bindgen-macro-support" -version = "0.2.87" +version = "0.2.92" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "54681b18a46765f095758388f2d0cf16eb8d4169b639ab575a8f5693af210c7b" +checksum = "e94f17b526d0a461a191c78ea52bbce64071ed5c04c9ffe424dcb38f74171bb7" dependencies = [ "proc-macro2", "quote", - "syn 2.0.48", + "syn 2.0.58", "wasm-bindgen-backend", "wasm-bindgen-shared", ] [[package]] name = "wasm-bindgen-shared" -version = "0.2.87" +version = "0.2.92" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ca6ad05a4870b2bf5fe995117d3728437bd27d7cd5f06f13c17443ef369775a1" +checksum = "af190c94f2773fdb3729c55b007a722abb5384da03bc0986df4c289bf5567e96" [[package]] name = "wasm-streams" @@ -5613,15 +5606,6 @@ dependencies = [ "web-sys", ] -[[package]] -name = "wav" -version = "1.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a65e199c799848b4f997072aa4d673c034f80f40191f97fe2f0a23f410be1609" -dependencies = [ - "riff", -] - [[package]] name = "web-sys" version = "0.3.64" @@ -5675,9 +5659,9 @@ checksum = "ac3b87c63620426dd9b991e5ce0329eff545bccbbb34f3be09ff6fb6ab51b7b6" [[package]] name = "winapi-util" -version = "0.1.5" +version = "0.1.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "70ec6ce85bb158151cae5e5c87f95a8e97d2c0c4b001223f33a334e3ce5de178" +checksum = "f29e6f9198ba0d26b4c9f07dbe6f9ed633e1f3d5b8b414090084349e46a52596" dependencies = [ "winapi", ] @@ -5983,9 +5967,9 @@ dependencies = [ [[package]] name = "yoke" -version = "0.7.2" +version = "0.7.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "61e38c508604d6bbbd292dadb3c02559aa7fff6b654a078a36217cad871636e4" +checksum = "65e71b2e4f287f467794c671e2b8f8a5f3716b3c829079a1c44740148eff07e4" dependencies = [ "serde", "stable_deref_trait", @@ -5995,13 +5979,13 @@ dependencies = [ [[package]] name = "yoke-derive" -version = "0.7.2" +version = "0.7.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d5e19fb6ed40002bab5403ffa37e53e0e56f914a4450c8765f533018db1db35f" +checksum = "9e6936f0cce458098a201c245a11bef556c6a0181129c7034d10d76d1ec3a2b8" dependencies = [ "proc-macro2", "quote", - "syn 2.0.48", + "syn 2.0.58", "synstructure", ] @@ -6022,7 +6006,7 @@ checksum = "9ce1b18ccd8e73a9321186f97e46f9f04b778851177567b1975109d26a08d2a6" dependencies = [ "proc-macro2", "quote", - "syn 2.0.48", + "syn 2.0.58", ] [[package]] @@ -6042,7 +6026,7 @@ checksum = "e6a647510471d372f2e6c2e6b7219e44d8c574d24fdc11c610a61455782f18c3" dependencies = [ "proc-macro2", "quote", - "syn 2.0.48", + "syn 2.0.58", "synstructure", ] @@ -6102,11 +6086,10 @@ dependencies = [ [[package]] name = "zstd-sys" -version = "2.0.8+zstd.1.5.5" +version = "2.0.10+zstd.1.5.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5556e6ee25d32df2586c098bbfa278803692a20d0ab9565e049480d52707ec8c" +checksum = "c253a4914af5bafc8fa8c86ee400827e83cf6ec01195ec1f1ed8441bf00d65aa" dependencies = [ "cc", - "libc", "pkg-config", ] diff --git a/milli/Cargo.toml b/milli/Cargo.toml index 4c3ccc132..4a52c0459 100644 --- a/milli/Cargo.toml +++ b/milli/Cargo.toml @@ -71,10 +71,10 @@ itertools = "0.11.0" puffin = "0.16.0" csv = "1.3.0" -candle-core = { git = "https://github.com/huggingface/candle.git", rev="5270224f407502b82fe90bc2622894ce3871b002", version = "0.3.3" } -candle-transformers = { git = "https://github.com/huggingface/candle.git", rev="5270224f407502b82fe90bc2622894ce3871b002", version = "0.3.3" } -candle-nn = { git = "https://github.com/huggingface/candle.git", rev="5270224f407502b82fe90bc2622894ce3871b002", version = "0.3.3" } -tokenizers = { git = "https://github.com/huggingface/tokenizers.git", tag = "v0.14.1", version = "0.14.1", default_features = false, features = [ +candle-core = { version = "0.4.1" } +candle-transformers = { version = "0.4.1" } +candle-nn = { version = "0.4.1" } +tokenizers = { git = "https://github.com/huggingface/tokenizers.git", tag = "v0.15.2", version = "0.15.2", default_features = false, features = [ "onig", ] } hf-hub = { git = "https://github.com/dureuill/hf-hub.git", branch = "rust_tls", default_features = false, features = [ From 90e812fc0bae7ec1812df325736e0a83196df551 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Cl=C3=A9ment=20Renault?= Date: Thu, 4 Apr 2024 12:01:04 +0200 Subject: [PATCH 72/86] Add some tests --- meilisearch/tests/search/mod.rs | 104 ++++++++++++++++++++++++++++++++ 1 file changed, 104 insertions(+) diff --git a/meilisearch/tests/search/mod.rs b/meilisearch/tests/search/mod.rs index 88470187a..07de60d7c 100644 --- a/meilisearch/tests/search/mod.rs +++ b/meilisearch/tests/search/mod.rs @@ -184,6 +184,110 @@ async fn phrase_search_with_stop_word() { .await; } +#[actix_rt::test] +async fn negative_phrase_search() { + let server = Server::new().await; + let index = server.index("test"); + + let documents = DOCUMENTS.clone(); + index.add_documents(documents, None).await; + index.wait_task(0).await; + + index + .search(json!({"q": "-\"train your dragon\"" }), |response, code| { + assert_eq!(code, 200, "{}", response); + let hits = response["hits"].as_array().unwrap(); + assert_eq!(hits.len(), 4); + assert_eq!(hits[0]["id"], "287947"); + assert_eq!(hits[1]["id"], "299537"); + assert_eq!(hits[2]["id"], "522681"); + assert_eq!(hits[3]["id"], "450465"); + }) + .await; +} + +#[actix_rt::test] +async fn negative_word_search() { + let server = Server::new().await; + let index = server.index("test"); + + let documents = DOCUMENTS.clone(); + index.add_documents(documents, None).await; + index.wait_task(0).await; + + index + .search(json!({"q": "-escape" }), |response, code| { + assert_eq!(code, 200, "{}", response); + let hits = response["hits"].as_array().unwrap(); + assert_eq!(hits.len(), 4); + assert_eq!(hits[0]["id"], "287947"); + assert_eq!(hits[1]["id"], "299537"); + assert_eq!(hits[2]["id"], "166428"); + assert_eq!(hits[3]["id"], "450465"); + }) + .await; + + // Everything that contains derivates of escape but not escape: nothing + index + .search(json!({"q": "-escape escape" }), |response, code| { + assert_eq!(code, 200, "{}", response); + let hits = response["hits"].as_array().unwrap(); + assert_eq!(hits.len(), 0); + }) + .await; +} + +#[actix_rt::test] +async fn non_negative_search() { + let server = Server::new().await; + let index = server.index("test"); + + let documents = DOCUMENTS.clone(); + index.add_documents(documents, None).await; + index.wait_task(0).await; + + index + .search(json!({"q": "- escape" }), |response, code| { + assert_eq!(code, 200, "{}", response); + let hits = response["hits"].as_array().unwrap(); + assert_eq!(hits.len(), 1); + assert_eq!(hits[0]["id"], "522681"); + }) + .await; + + index + .search(json!({"q": "- \"train your dragon\"" }), |response, code| { + assert_eq!(code, 200, "{}", response); + let hits = response["hits"].as_array().unwrap(); + assert_eq!(hits.len(), 1); + assert_eq!(hits[0]["id"], "166428"); + }) + .await; +} + +#[actix_rt::test] +async fn negative_special_cases_search() { + let server = Server::new().await; + let index = server.index("test"); + + let documents = DOCUMENTS.clone(); + index.add_documents(documents, None).await; + index.wait_task(0).await; + + index.update_settings(json!({"synonyms": { "escape": ["glass"] }})).await; + index.wait_task(1).await; + + // There is a synonym for escape -> glass but we don't want "escape", only the derivates: glass + index + .search(json!({"q": "-escape escape" }), |response, code| { + assert_eq!(code, 200, "{}", response); + let hits = response["hits"].as_array().unwrap(); + assert_eq!(hits.len(), 1); + assert_eq!(hits[0]["id"], "450465"); + }) + .await; +} + #[cfg(feature = "default")] #[actix_rt::test] async fn test_kanji_language_detection() { From 928e6e4c059718300231fda19ad1b6bd041ab477 Mon Sep 17 00:00:00 2001 From: Louis Dureuil Date: Wed, 27 Mar 2024 15:36:49 +0100 Subject: [PATCH 73/86] Breaking change: remove vector for score details --- meilisearch/src/search.rs | 6 ++---- milli/src/score_details.rs | 16 ++++++---------- milli/src/search/new/vector_sort.rs | 28 +++++++--------------------- 3 files changed, 15 insertions(+), 35 deletions(-) diff --git a/meilisearch/src/search.rs b/meilisearch/src/search.rs index db58c6102..2e0df18ad 100644 --- a/meilisearch/src/search.rs +++ b/meilisearch/src/search.rs @@ -633,10 +633,8 @@ pub fn perform_search( let mut semantic_score = None; for details in &score { - if let ScoreDetails::Vector(score_details::Vector { - target_vector: _, - value_similarity: Some((_matching_vector, similarity)), - }) = details + if let ScoreDetails::Vector(score_details::Vector { similarity: Some(similarity) }) = + details { semantic_score = Some(*similarity); break; diff --git a/milli/src/score_details.rs b/milli/src/score_details.rs index 08dfcdbb6..0a9b77e2b 100644 --- a/milli/src/score_details.rs +++ b/milli/src/score_details.rs @@ -98,9 +98,9 @@ impl ScoreDetails { ScoreDetails::ExactWords(e) => RankOrValue::Rank(e.rank()), ScoreDetails::Sort(sort) => RankOrValue::Sort(sort), ScoreDetails::GeoSort(geosort) => RankOrValue::GeoSort(geosort), - ScoreDetails::Vector(vector) => RankOrValue::Score( - vector.value_similarity.as_ref().map(|(_, s)| *s as f64).unwrap_or(0.0f64), - ), + ScoreDetails::Vector(vector) => { + RankOrValue::Score(vector.similarity.as_ref().map(|s| *s as f64).unwrap_or(0.0f64)) + } ScoreDetails::Skipped => RankOrValue::Rank(Rank { rank: 0, max_rank: 1 }), } } @@ -249,16 +249,13 @@ impl ScoreDetails { order += 1; } ScoreDetails::Vector(s) => { - let vector = format!("vectorSort({:?})", s.target_vector); - let value = s.value_similarity.as_ref().map(|(v, _)| v); - let similarity = s.value_similarity.as_ref().map(|(_, s)| s); + let similarity = s.similarity.as_ref(); let details = serde_json::json!({ "order": order, - "value": value, "similarity": similarity, }); - details_map.insert(vector, details); + details_map.insert("vectorSort".into(), details); order += 1; } ScoreDetails::Skipped => { @@ -494,8 +491,7 @@ impl PartialOrd for GeoSort { #[derive(Debug, Clone, PartialEq, PartialOrd)] pub struct Vector { - pub target_vector: Vec, - pub value_similarity: Option<(Vec, f32)>, + pub similarity: Option, } impl GeoSort { diff --git a/milli/src/search/new/vector_sort.rs b/milli/src/search/new/vector_sort.rs index b29a72827..476477218 100644 --- a/milli/src/search/new/vector_sort.rs +++ b/milli/src/search/new/vector_sort.rs @@ -12,7 +12,7 @@ pub struct VectorSort { query: Option, target: Vec, vector_candidates: RoaringBitmap, - cached_sorted_docids: std::vec::IntoIter<(DocumentId, f32, Vec)>, + cached_sorted_docids: std::vec::IntoIter<(DocumentId, f32)>, limit: usize, distribution_shift: Option, embedder_index: u8, @@ -70,14 +70,9 @@ impl VectorSort { for reader in readers.iter() { let nns_by_vector = reader.nns_by_vector(ctx.txn, target, self.limit, None, Some(vector_candidates))?; - let vectors: std::result::Result, _> = nns_by_vector - .iter() - .map(|(docid, _)| reader.item_vector(ctx.txn, *docid).transpose().unwrap()) - .collect(); - let vectors = vectors?; - results.extend(nns_by_vector.into_iter().zip(vectors).map(|((x, y), z)| (x, y, z))); + results.extend(nns_by_vector.into_iter()); } - results.sort_unstable_by_key(|(_, distance, _)| OrderedFloat(*distance)); + results.sort_unstable_by_key(|(_, distance)| OrderedFloat(*distance)); self.cached_sorted_docids = results.into_iter(); Ok(()) @@ -118,14 +113,11 @@ impl<'ctx, Q: RankingRuleQueryTrait> RankingRule<'ctx, Q> for VectorSort { return Ok(Some(RankingRuleOutput { query, candidates: universe.clone(), - score: ScoreDetails::Vector(score_details::Vector { - target_vector: self.target.clone(), - value_similarity: None, - }), + score: ScoreDetails::Vector(score_details::Vector { similarity: None }), })); } - for (docid, distance, vector) in self.cached_sorted_docids.by_ref() { + for (docid, distance) in self.cached_sorted_docids.by_ref() { if vector_candidates.contains(docid) { let score = 1.0 - distance; let score = self @@ -135,10 +127,7 @@ impl<'ctx, Q: RankingRuleQueryTrait> RankingRule<'ctx, Q> for VectorSort { return Ok(Some(RankingRuleOutput { query, candidates: RoaringBitmap::from_iter([docid]), - score: ScoreDetails::Vector(score_details::Vector { - target_vector: self.target.clone(), - value_similarity: Some((vector, score)), - }), + score: ScoreDetails::Vector(score_details::Vector { similarity: Some(score) }), })); } } @@ -154,10 +143,7 @@ impl<'ctx, Q: RankingRuleQueryTrait> RankingRule<'ctx, Q> for VectorSort { return Ok(Some(RankingRuleOutput { query, candidates: universe.clone(), - score: ScoreDetails::Vector(score_details::Vector { - target_vector: self.target.clone(), - value_similarity: None, - }), + score: ScoreDetails::Vector(score_details::Vector { similarity: None }), })); } From 190933f6e1986571056b5147a5189c8d53dcb972 Mon Sep 17 00:00:00 2001 From: Louis Dureuil Date: Wed, 27 Mar 2024 15:40:02 +0100 Subject: [PATCH 74/86] Breaking: Remove vector from SearchResult --- meilisearch/src/analytics/segment_analytics.rs | 1 - meilisearch/src/search.rs | 3 --- 2 files changed, 4 deletions(-) diff --git a/meilisearch/src/analytics/segment_analytics.rs b/meilisearch/src/analytics/segment_analytics.rs index 4a0ef5b35..fcf4d9144 100644 --- a/meilisearch/src/analytics/segment_analytics.rs +++ b/meilisearch/src/analytics/segment_analytics.rs @@ -758,7 +758,6 @@ impl SearchAggregator { let SearchResult { hits: _, query: _, - vector: _, processing_time_ms, hits_info: _, facet_distribution: _, diff --git a/meilisearch/src/search.rs b/meilisearch/src/search.rs index 2e0df18ad..a1aa37779 100644 --- a/meilisearch/src/search.rs +++ b/meilisearch/src/search.rs @@ -314,8 +314,6 @@ pub struct SearchHit { pub struct SearchResult { pub hits: Vec, pub query: String, - #[serde(skip_serializing_if = "Option::is_none")] - pub vector: Option>, pub processing_time_ms: u128, #[serde(flatten)] pub hits_info: HitsInfo, @@ -713,7 +711,6 @@ pub fn perform_search( hits: documents, hits_info, query: query.q.unwrap_or_default(), - vector: query.vector, processing_time_ms: before_search.elapsed().as_millis(), facet_distribution, facet_stats, From 00c4ed3bc25455d81695680d164db58e7c310e89 Mon Sep 17 00:00:00 2001 From: Louis Dureuil Date: Thu, 28 Mar 2024 11:49:00 +0100 Subject: [PATCH 75/86] milli: refactor getting embedder and embedder name --- milli/src/vector/mod.rs | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/milli/src/vector/mod.rs b/milli/src/vector/mod.rs index 1cb0a18f7..5aa58da5d 100644 --- a/milli/src/vector/mod.rs +++ b/milli/src/vector/mod.rs @@ -143,7 +143,7 @@ impl EmbeddingConfigs { /// Get the default embedder configuration, if any. pub fn get_default(&self) -> Option<(Arc, Arc)> { - self.get_default_embedder_name().and_then(|default| self.get(&default)) + self.get(self.get_default_embedder_name()) } /// Get the name of the default embedder configuration. @@ -153,14 +153,14 @@ impl EmbeddingConfigs { /// - If there is only one embedder, it is always the default. /// - If there are multiple embedders and one of them is called `default`, then that one is the default embedder. /// - In all other cases, there is no default embedder. - pub fn get_default_embedder_name(&self) -> Option { + pub fn get_default_embedder_name(&self) -> &str { let mut it = self.0.keys(); let first_name = it.next(); let second_name = it.next(); match (first_name, second_name) { - (None, _) => None, - (Some(first), None) => Some(first.to_owned()), - (Some(_), Some(_)) => Some("default".to_owned()), + (None, _) => "default", + (Some(first), None) => first, + (Some(_), Some(_)) => "default", } } } From fabc9cf14af3efe900c4c62d9f936422db830be0 Mon Sep 17 00:00:00 2001 From: Louis Dureuil Date: Thu, 28 Mar 2024 11:49:23 +0100 Subject: [PATCH 76/86] milli: add Embedder::embed_one --- milli/src/vector/error.rs | 7 ++++++- milli/src/vector/mod.rs | 11 +++++++++++ 2 files changed, 17 insertions(+), 1 deletion(-) diff --git a/milli/src/vector/error.rs b/milli/src/vector/error.rs index 1e0bcc7fb..d3369ef3d 100644 --- a/milli/src/vector/error.rs +++ b/milli/src/vector/error.rs @@ -58,7 +58,7 @@ pub enum EmbedErrorKind { RestResponseDeserialization(std::io::Error), #[error("component `{0}` not found in path `{1}` in response: `{2}`")] RestResponseMissingEmbeddings(String, String, String), - #[error("expected a response parseable as a vector or an array of vectors: {0}")] + #[error("unexpected format of the embedding response: {0}")] RestResponseFormat(serde_json::Error), #[error("expected a response containing {0} embeddings, got only {1}")] RestResponseEmbeddingCount(usize, usize), @@ -78,6 +78,8 @@ pub enum EmbedErrorKind { RestNotAnObject(serde_json::Value, Vec), #[error("while embedding tokenized, was expecting embeddings of dimension `{0}`, got embeddings of dimensions `{1}`")] OpenAiUnexpectedDimension(usize, usize), + #[error("no embedding was produced")] + MissingEmbedding, } impl EmbedError { @@ -190,6 +192,9 @@ impl EmbedError { fault: FaultSource::Runtime, } } + pub(crate) fn missing_embedding() -> EmbedError { + Self { kind: EmbedErrorKind::MissingEmbedding, fault: FaultSource::Undecided } + } } #[derive(Debug, thiserror::Error)] diff --git a/milli/src/vector/mod.rs b/milli/src/vector/mod.rs index 5aa58da5d..58f7ba5e1 100644 --- a/milli/src/vector/mod.rs +++ b/milli/src/vector/mod.rs @@ -237,6 +237,17 @@ impl Embedder { } } + pub fn embed_one(&self, text: String) -> std::result::Result { + let mut embeddings = self.embed(vec![text])?; + let embeddings = embeddings.pop().ok_or_else(EmbedError::missing_embedding)?; + Ok(if embeddings.iter().nth(1).is_some() { + tracing::warn!("Ignoring embeddings past the first one in long search query"); + embeddings.iter().next().unwrap().to_vec() + } else { + embeddings.into_inner() + }) + } + /// Embed multiple chunks of texts. /// /// Each chunk is composed of one or multiple texts. From 6ebb6b55a64b266bfad68b8a76ee5ce00435b2d0 Mon Sep 17 00:00:00 2001 From: Louis Dureuil Date: Thu, 28 Mar 2024 11:50:53 +0100 Subject: [PATCH 77/86] Lazily embed, don't fail hybrid search on embedding failure --- .../src/routes/indexes/facet_search.rs | 4 +- meilisearch/src/routes/indexes/search.rs | 119 ++++++--------- meilisearch/src/routes/multi_search.rs | 8 +- meilisearch/src/search.rs | 141 +++++++++++++----- milli/src/index.rs | 8 - milli/src/lib.rs | 2 +- milli/src/search/facet/search.rs | 12 +- milli/src/search/hybrid.rs | 37 +++-- milli/src/search/mod.rs | 93 +++++------- milli/src/search/new/mod.rs | 10 +- milli/src/search/new/vector_sort.rs | 6 +- 11 files changed, 237 insertions(+), 203 deletions(-) diff --git a/meilisearch/src/routes/indexes/facet_search.rs b/meilisearch/src/routes/indexes/facet_search.rs index 272b8156f..56880a472 100644 --- a/meilisearch/src/routes/indexes/facet_search.rs +++ b/meilisearch/src/routes/indexes/facet_search.rs @@ -12,6 +12,7 @@ use tracing::debug; use crate::analytics::{Analytics, FacetSearchAggregator}; use crate::extractors::authentication::policies::*; use crate::extractors::authentication::GuardedData; +use crate::routes::indexes::search::search_kind; use crate::search::{ add_search_rules, perform_facet_search, HybridQuery, MatchingStrategy, SearchQuery, DEFAULT_CROP_LENGTH, DEFAULT_CROP_MARKER, DEFAULT_HIGHLIGHT_POST_TAG, @@ -73,9 +74,10 @@ pub async fn search( let index = index_scheduler.index(&index_uid)?; let features = index_scheduler.features(); + let search_kind = search_kind(&search_query, &index_scheduler, &index)?; let _permit = search_queue.try_get_search_permit().await?; let search_result = tokio::task::spawn_blocking(move || { - perform_facet_search(&index, search_query, facet_query, facet_name, features) + perform_facet_search(&index, search_query, facet_query, facet_name, features, search_kind) }) .await?; diff --git a/meilisearch/src/routes/indexes/search.rs b/meilisearch/src/routes/indexes/search.rs index f16a6c4df..a5fe3c5d6 100644 --- a/meilisearch/src/routes/indexes/search.rs +++ b/meilisearch/src/routes/indexes/search.rs @@ -8,19 +8,19 @@ use meilisearch_types::error::deserr_codes::*; use meilisearch_types::error::ResponseError; use meilisearch_types::index_uid::IndexUid; use meilisearch_types::milli; -use meilisearch_types::milli::vector::DistributionShift; use meilisearch_types::serde_cs::vec::CS; use serde_json::Value; -use tracing::{debug, warn}; +use tracing::debug; use crate::analytics::{Analytics, SearchAggregator}; +use crate::error::MeilisearchHttpError; use crate::extractors::authentication::policies::*; use crate::extractors::authentication::GuardedData; use crate::extractors::sequential_extractor::SeqHandler; use crate::metrics::MEILISEARCH_DEGRADED_SEARCH_REQUESTS; use crate::search::{ - add_search_rules, perform_search, HybridQuery, MatchingStrategy, SearchQuery, SemanticRatio, - DEFAULT_CROP_LENGTH, DEFAULT_CROP_MARKER, DEFAULT_HIGHLIGHT_POST_TAG, + add_search_rules, perform_search, HybridQuery, MatchingStrategy, SearchKind, SearchQuery, + SemanticRatio, DEFAULT_CROP_LENGTH, DEFAULT_CROP_MARKER, DEFAULT_HIGHLIGHT_POST_TAG, DEFAULT_HIGHLIGHT_PRE_TAG, DEFAULT_SEARCH_LIMIT, DEFAULT_SEARCH_OFFSET, DEFAULT_SEMANTIC_RATIO, }; use crate::search_queue::SearchQueue; @@ -204,11 +204,11 @@ pub async fn search_with_url_query( let index = index_scheduler.index(&index_uid)?; let features = index_scheduler.features(); - let distribution = embed(&mut query, index_scheduler.get_ref(), &index)?; + let search_kind = search_kind(&query, index_scheduler.get_ref(), &index)?; let _permit = search_queue.try_get_search_permit().await?; let search_result = - tokio::task::spawn_blocking(move || perform_search(&index, query, features, distribution)) + tokio::task::spawn_blocking(move || perform_search(&index, query, features, search_kind)) .await?; if let Ok(ref search_result) = search_result { aggregate.succeed(search_result); @@ -245,11 +245,11 @@ pub async fn search_with_post( let features = index_scheduler.features(); - let distribution = embed(&mut query, index_scheduler.get_ref(), &index)?; + let search_kind = search_kind(&query, index_scheduler.get_ref(), &index)?; let _permit = search_queue.try_get_search_permit().await?; let search_result = - tokio::task::spawn_blocking(move || perform_search(&index, query, features, distribution)) + tokio::task::spawn_blocking(move || perform_search(&index, query, features, search_kind)) .await?; if let Ok(ref search_result) = search_result { aggregate.succeed(search_result); @@ -265,76 +265,49 @@ pub async fn search_with_post( Ok(HttpResponse::Ok().json(search_result)) } -pub fn embed( - query: &mut SearchQuery, +pub fn search_kind( + query: &SearchQuery, index_scheduler: &IndexScheduler, index: &milli::Index, -) -> Result, ResponseError> { - match (&query.hybrid, &query.vector, &query.q) { - (Some(HybridQuery { semantic_ratio: _, embedder }), None, Some(q)) - if !q.trim().is_empty() => - { - let embedder_configs = index.embedding_configs(&index.read_txn()?)?; - let embedders = index_scheduler.embedders(embedder_configs)?; - - let embedder = if let Some(embedder_name) = embedder { - embedders.get(embedder_name) - } else { - embedders.get_default() - }; - - let embedder = embedder - .ok_or(milli::UserError::InvalidEmbedder("default".to_owned())) - .map_err(milli::Error::from)? - .0; - - let distribution = embedder.distribution(); - - let embeddings = embedder - .embed(vec![q.to_owned()]) - .map_err(milli::vector::Error::from) - .map_err(milli::Error::from)? - .pop() - .expect("No vector returned from embedding"); - - if embeddings.iter().nth(1).is_some() { - warn!("Ignoring embeddings past the first one in long search query"); - query.vector = Some(embeddings.iter().next().unwrap().to_vec()); - } else { - query.vector = Some(embeddings.into_inner()); - } - Ok(distribution) +) -> Result { + // regardless of anything, always do a semantic search when we don't have a vector and the query is whitespace or missing + if query.vector.is_none() { + match &query.q { + Some(q) if q.trim().is_empty() => return Ok(SearchKind::KeywordOnly), + None => return Ok(SearchKind::KeywordOnly), + _ => {} } - (Some(hybrid), vector, _) => { - let embedder_configs = index.embedding_configs(&index.read_txn()?)?; - let embedders = index_scheduler.embedders(embedder_configs)?; + } - let embedder = if let Some(embedder_name) = &hybrid.embedder { - embedders.get(embedder_name) - } else { - embedders.get_default() - }; - - let embedder = embedder - .ok_or(milli::UserError::InvalidEmbedder("default".to_owned())) - .map_err(milli::Error::from)? - .0; - - if let Some(vector) = vector { - if vector.len() != embedder.dimensions() { - return Err(meilisearch_types::milli::Error::UserError( - meilisearch_types::milli::UserError::InvalidVectorDimensions { - expected: embedder.dimensions(), - found: vector.len(), - }, - ) - .into()); - } - } - - Ok(embedder.distribution()) + match &query.hybrid { + Some(HybridQuery { semantic_ratio, embedder }) if **semantic_ratio == 1.0 => { + Ok(SearchKind::semantic( + index_scheduler, + index, + embedder.as_deref(), + query.vector.as_ref().map(Vec::len), + )?) } - _ => Ok(None), + Some(HybridQuery { semantic_ratio, embedder: _ }) if **semantic_ratio == 0.0 => { + Ok(SearchKind::KeywordOnly) + } + Some(HybridQuery { semantic_ratio, embedder }) => Ok(SearchKind::hybrid( + index_scheduler, + index, + embedder.as_deref(), + **semantic_ratio, + query.vector.as_ref().map(Vec::len), + )?), + None => match (query.q.as_deref(), query.vector.as_deref()) { + (_query, None) => Ok(SearchKind::KeywordOnly), + (None, Some(_vector)) => Ok(SearchKind::semantic( + index_scheduler, + index, + None, + query.vector.as_ref().map(Vec::len), + )?), + (Some(_), Some(_)) => Err(MeilisearchHttpError::MissingSearchHybrid.into()), + }, } } diff --git a/meilisearch/src/routes/multi_search.rs b/meilisearch/src/routes/multi_search.rs index b2055fb07..04cd3f637 100644 --- a/meilisearch/src/routes/multi_search.rs +++ b/meilisearch/src/routes/multi_search.rs @@ -13,7 +13,7 @@ use crate::analytics::{Analytics, MultiSearchAggregator}; use crate::extractors::authentication::policies::ActionPolicy; use crate::extractors::authentication::{AuthenticationError, GuardedData}; use crate::extractors::sequential_extractor::SeqHandler; -use crate::routes::indexes::search::embed; +use crate::routes::indexes::search::search_kind; use crate::search::{ add_search_rules, perform_search, SearchQueryWithIndex, SearchResultWithIndex, }; @@ -81,11 +81,11 @@ pub async fn multi_search_with_post( }) .with_index(query_index)?; - let distribution = - embed(&mut query, index_scheduler.get_ref(), &index).with_index(query_index)?; + let search_kind = + search_kind(&query, index_scheduler.get_ref(), &index).with_index(query_index)?; let search_result = tokio::task::spawn_blocking(move || { - perform_search(&index, query, features, distribution) + perform_search(&index, query, features, search_kind) }) .await .with_index(query_index)?; diff --git a/meilisearch/src/search.rs b/meilisearch/src/search.rs index a1aa37779..2a22cb2ce 100644 --- a/meilisearch/src/search.rs +++ b/meilisearch/src/search.rs @@ -1,6 +1,7 @@ use std::cmp::min; use std::collections::{BTreeMap, BTreeSet, HashSet}; use std::str::FromStr; +use std::sync::Arc; use std::time::{Duration, Instant}; use deserr::Deserr; @@ -10,10 +11,11 @@ use indexmap::IndexMap; use meilisearch_auth::IndexSearchRules; use meilisearch_types::deserr::DeserrJsonError; use meilisearch_types::error::deserr_codes::*; +use meilisearch_types::error::ResponseError; use meilisearch_types::heed::RoTxn; use meilisearch_types::index_uid::IndexUid; use meilisearch_types::milli::score_details::{self, ScoreDetails, ScoringStrategy}; -use meilisearch_types::milli::vector::DistributionShift; +use meilisearch_types::milli::vector::Embedder; use meilisearch_types::milli::{FacetValueHit, OrderBy, SearchForFacetValues, TimeBudget}; use meilisearch_types::settings::DEFAULT_PAGINATION_MAX_TOTAL_HITS; use meilisearch_types::{milli, Document}; @@ -90,13 +92,75 @@ pub struct SearchQuery { #[derive(Debug, Clone, Default, PartialEq, Deserr)] #[deserr(error = DeserrJsonError, rename_all = camelCase, deny_unknown_fields)] pub struct HybridQuery { - /// TODO validate that sementic ratio is between 0.0 and 1,0 #[deserr(default, error = DeserrJsonError, default)] pub semantic_ratio: SemanticRatio, #[deserr(default, error = DeserrJsonError, default)] pub embedder: Option, } +pub enum SearchKind { + KeywordOnly, + SemanticOnly { embedder_name: String, embedder: Arc }, + Hybrid { embedder_name: String, embedder: Arc, semantic_ratio: f32 }, +} +impl SearchKind { + pub(crate) fn semantic( + index_scheduler: &index_scheduler::IndexScheduler, + index: &Index, + embedder_name: Option<&str>, + vector_len: Option, + ) -> Result { + let (embedder_name, embedder) = + Self::embedder(index_scheduler, index, embedder_name, vector_len)?; + Ok(Self::SemanticOnly { embedder_name, embedder }) + } + + pub(crate) fn hybrid( + index_scheduler: &index_scheduler::IndexScheduler, + index: &Index, + embedder_name: Option<&str>, + semantic_ratio: f32, + vector_len: Option, + ) -> Result { + let (embedder_name, embedder) = + Self::embedder(index_scheduler, index, embedder_name, vector_len)?; + Ok(Self::Hybrid { embedder_name, embedder, semantic_ratio }) + } + + fn embedder( + index_scheduler: &index_scheduler::IndexScheduler, + index: &Index, + embedder_name: Option<&str>, + vector_len: Option, + ) -> Result<(String, Arc), ResponseError> { + let embedder_configs = index.embedding_configs(&index.read_txn()?)?; + let embedders = index_scheduler.embedders(embedder_configs)?; + + let embedder_name = embedder_name.unwrap_or_else(|| embedders.get_default_embedder_name()); + + let embedder = embedders.get(embedder_name); + + let embedder = embedder + .ok_or(milli::UserError::InvalidEmbedder(embedder_name.to_owned())) + .map_err(milli::Error::from)? + .0; + + if let Some(vector_len) = vector_len { + if vector_len != embedder.dimensions() { + return Err(meilisearch_types::milli::Error::UserError( + meilisearch_types::milli::UserError::InvalidVectorDimensions { + expected: embedder.dimensions(), + found: vector_len, + }, + ) + .into()); + } + } + + Ok((embedder_name.to_owned(), embedder)) + } +} + #[derive(Debug, Clone, Copy, PartialEq, Deserr)] #[deserr(try_from(f32) = TryFrom::try_from -> InvalidSearchSemanticRatio)] pub struct SemanticRatio(f32); @@ -385,7 +449,7 @@ fn prepare_search<'t>( rtxn: &'t RoTxn, query: &'t SearchQuery, features: RoFeatures, - distribution: Option, + search_kind: &SearchKind, time_budget: TimeBudget, ) -> Result<(milli::Search<'t>, bool, usize, usize), MeilisearchHttpError> { let mut search = index.search(rtxn); @@ -399,32 +463,30 @@ fn prepare_search<'t>( features.check_vector("Passing `hybrid` as a query parameter")?; } - if query.hybrid.is_none() && query.q.is_some() && query.vector.is_some() { - return Err(MeilisearchHttpError::MissingSearchHybrid); - } - - search.distribution_shift(distribution); - - if let Some(ref vector) = query.vector { - match &query.hybrid { - // If semantic ratio is 0.0, only the query search will impact the search results, - // skip the vector - Some(hybrid) if *hybrid.semantic_ratio == 0.0 => (), - _otherwise => { - search.vector(vector.clone()); - } - } - } - - if let Some(ref q) = query.q { - match &query.hybrid { - // If semantic ratio is 1.0, only the vector search will impact the search results, - // skip the query - Some(hybrid) if *hybrid.semantic_ratio == 1.0 => (), - _otherwise => { + match search_kind { + SearchKind::KeywordOnly => { + if let Some(q) = &query.q { search.query(q); } } + SearchKind::SemanticOnly { embedder_name, embedder } => { + let vector = match query.vector.clone() { + Some(vector) => vector, + None => embedder + .embed_one(query.q.clone().unwrap()) + .map_err(milli::vector::Error::from) + .map_err(milli::Error::from)?, + }; + + search.semantic(embedder_name.clone(), embedder.clone(), Some(vector)); + } + SearchKind::Hybrid { embedder_name, embedder, semantic_ratio: _ } => { + if let Some(q) = &query.q { + search.query(q); + } + // will be embedded in hybrid search if necessary + search.semantic(embedder_name.clone(), embedder.clone(), query.vector.clone()); + } } if let Some(ref searchable) = query.attributes_to_search_on { @@ -447,10 +509,6 @@ fn prepare_search<'t>( ScoringStrategy::Skip }); - if let Some(HybridQuery { embedder: Some(embedder), .. }) = &query.hybrid { - search.embedder_name(embedder); - } - // compute the offset on the limit depending on the pagination mode. let (offset, limit) = if is_finite_pagination { let limit = query.hits_per_page.unwrap_or_else(DEFAULT_SEARCH_LIMIT); @@ -494,7 +552,7 @@ pub fn perform_search( index: &Index, query: SearchQuery, features: RoFeatures, - distribution: Option, + search_kind: SearchKind, ) -> Result { let before_search = Instant::now(); let rtxn = index.read_txn()?; @@ -504,7 +562,7 @@ pub fn perform_search( }; let (search, is_finite_pagination, max_total_hits, offset) = - prepare_search(index, &rtxn, &query, features, distribution, time_budget)?; + prepare_search(index, &rtxn, &query, features, &search_kind, time_budget)?; let milli::SearchResult { documents_ids, @@ -514,12 +572,9 @@ pub fn perform_search( degraded, used_negative_operator, .. - } = match &query.hybrid { - Some(hybrid) => match *hybrid.semantic_ratio { - ratio if ratio == 0.0 || ratio == 1.0 => search.execute()?, - ratio => search.execute_hybrid(ratio)?, - }, - None => search.execute()?, + } = match &search_kind { + SearchKind::KeywordOnly | SearchKind::SemanticOnly { .. } => search.execute()?, + SearchKind::Hybrid { semantic_ratio, .. } => search.execute_hybrid(*semantic_ratio)?, }; let fields_ids_map = index.fields_ids_map(&rtxn).unwrap(); @@ -726,6 +781,7 @@ pub fn perform_facet_search( facet_query: Option, facet_name: String, features: RoFeatures, + search_kind: SearchKind, ) -> Result { let before_search = Instant::now(); let rtxn = index.read_txn()?; @@ -735,9 +791,12 @@ pub fn perform_facet_search( }; let (search, _, _, _) = - prepare_search(index, &rtxn, &search_query, features, None, time_budget)?; - let mut facet_search = - SearchForFacetValues::new(facet_name, search, search_query.hybrid.is_some()); + prepare_search(index, &rtxn, &search_query, features, &search_kind, time_budget)?; + let mut facet_search = SearchForFacetValues::new( + facet_name, + search, + matches!(search_kind, SearchKind::Hybrid { .. }), + ); if let Some(facet_query) = &facet_query { facet_search.query(facet_query); } diff --git a/milli/src/index.rs b/milli/src/index.rs index 80e524fb1..db31c953a 100644 --- a/milli/src/index.rs +++ b/milli/src/index.rs @@ -1499,14 +1499,6 @@ impl Index { .unwrap_or_default()) } - pub fn default_embedding_name(&self, rtxn: &RoTxn<'_>) -> Result { - let configs = self.embedding_configs(rtxn)?; - Ok(match configs.as_slice() { - [(ref first_name, _)] => first_name.clone(), - _ => "default".to_owned(), - }) - } - pub(crate) fn put_search_cutoff(&self, wtxn: &mut RwTxn<'_>, cutoff: u64) -> heed::Result<()> { self.main.remap_types::().put(wtxn, main_key::SEARCH_CUTOFF, &cutoff) } diff --git a/milli/src/lib.rs b/milli/src/lib.rs index df44ca127..22816787b 100644 --- a/milli/src/lib.rs +++ b/milli/src/lib.rs @@ -61,7 +61,7 @@ pub use self::index::Index; pub use self::search::facet::{FacetValueHit, SearchForFacetValues}; pub use self::search::{ FacetDistribution, Filter, FormatOptions, MatchBounds, MatcherBuilder, MatchingWords, OrderBy, - Search, SearchResult, TermsMatchingStrategy, DEFAULT_VALUES_PER_FACET, + Search, SearchResult, SemanticSearch, TermsMatchingStrategy, DEFAULT_VALUES_PER_FACET, }; pub type Result = std::result::Result; diff --git a/milli/src/search/facet/search.rs b/milli/src/search/facet/search.rs index 0251d6b8d..a6756a7af 100644 --- a/milli/src/search/facet/search.rs +++ b/milli/src/search/facet/search.rs @@ -92,9 +92,15 @@ impl<'a> SearchForFacetValues<'a> { None => return Ok(Vec::new()), }; - let search_candidates = self - .search_query - .execute_for_candidates(self.is_hybrid || self.search_query.vector.is_some())?; + let search_candidates = self.search_query.execute_for_candidates( + self.is_hybrid + || self + .search_query + .semantic + .as_ref() + .and_then(|semantic| semantic.vector.as_ref()) + .is_some(), + )?; let mut results = match index.sort_facet_values_by(rtxn)?.get(&self.facet) { OrderBy::Lexicographic => ValuesCollection::by_lexicographic(self.max_values), diff --git a/milli/src/search/hybrid.rs b/milli/src/search/hybrid.rs index 2a6d9f7a5..e45652206 100644 --- a/milli/src/search/hybrid.rs +++ b/milli/src/search/hybrid.rs @@ -4,6 +4,7 @@ use itertools::Itertools; use roaring::RoaringBitmap; use crate::score_details::{ScoreDetails, ScoreValue, ScoringStrategy}; +use crate::search::SemanticSearch; use crate::{MatchingWords, Result, Search, SearchResult}; struct ScoreWithRatioResult { @@ -126,7 +127,6 @@ impl<'a> Search<'a> { // create separate keyword and semantic searches let mut search = Search { query: self.query.clone(), - vector: self.vector.clone(), filter: self.filter.clone(), offset: 0, limit: self.limit + self.offset, @@ -139,26 +139,41 @@ impl<'a> Search<'a> { exhaustive_number_hits: self.exhaustive_number_hits, rtxn: self.rtxn, index: self.index, - distribution_shift: self.distribution_shift, - embedder_name: self.embedder_name.clone(), + semantic: self.semantic.clone(), time_budget: self.time_budget.clone(), }; - let vector_query = search.vector.take(); + let semantic = search.semantic.take(); let keyword_results = search.execute()?; - // skip semantic search if we don't have a vector query (placeholder search) - let Some(vector_query) = vector_query else { - return Ok(keyword_results); - }; - // completely skip semantic search if the results of the keyword search are good enough if self.results_good_enough(&keyword_results, semantic_ratio) { return Ok(keyword_results); } - search.vector = Some(vector_query); - search.query = None; + // no vector search against placeholder search + let Some(query) = search.query.take() else { return Ok(keyword_results) }; + // no embedder, no semantic search + let Some(SemanticSearch { vector, embedder_name, embedder }) = semantic else { + return Ok(keyword_results); + }; + + let vector_query = match vector { + Some(vector_query) => vector_query, + None => { + // attempt to embed the vector + match embedder.embed_one(query) { + Ok(embedding) => embedding, + Err(error) => { + tracing::error!(error=%error, "Embedding failed"); + return Ok(keyword_results); + } + } + } + }; + + search.semantic = + Some(SemanticSearch { vector: Some(vector_query), embedder_name, embedder }); // TODO: would be better to have two distinct functions at this point let vector_results = search.execute()?; diff --git a/milli/src/search/mod.rs b/milli/src/search/mod.rs index 3c709a647..bab67e6bd 100644 --- a/milli/src/search/mod.rs +++ b/milli/src/search/mod.rs @@ -1,4 +1,5 @@ use std::fmt; +use std::sync::Arc; use levenshtein_automata::{LevenshteinAutomatonBuilder as LevBuilder, DFA}; use once_cell::sync::Lazy; @@ -8,7 +9,7 @@ pub use self::facet::{FacetDistribution, Filter, OrderBy, DEFAULT_VALUES_PER_FAC pub use self::new::matches::{FormatOptions, MatchBounds, MatcherBuilder, MatchingWords}; use self::new::{execute_vector_search, PartialSearchResult}; use crate::score_details::{ScoreDetails, ScoringStrategy}; -use crate::vector::DistributionShift; +use crate::vector::Embedder; use crate::{ execute_search, filtered_universe, AscDesc, DefaultSearchLogger, DocumentId, Index, Result, SearchContext, TimeBudget, @@ -24,9 +25,15 @@ mod fst_utils; pub mod hybrid; pub mod new; +#[derive(Debug, Clone)] +pub struct SemanticSearch { + vector: Option>, + embedder_name: String, + embedder: Arc, +} + pub struct Search<'a> { query: Option, - vector: Option>, // this should be linked to the String in the query filter: Option>, offset: usize, @@ -38,12 +45,9 @@ pub struct Search<'a> { scoring_strategy: ScoringStrategy, words_limit: usize, exhaustive_number_hits: bool, - /// TODO: Add semantic ratio or pass it directly to execute_hybrid() rtxn: &'a heed::RoTxn<'a>, index: &'a Index, - distribution_shift: Option, - embedder_name: Option, - + semantic: Option, time_budget: TimeBudget, } @@ -51,7 +55,6 @@ impl<'a> Search<'a> { pub fn new(rtxn: &'a heed::RoTxn, index: &'a Index) -> Search<'a> { Search { query: None, - vector: None, filter: None, offset: 0, limit: 20, @@ -64,8 +67,7 @@ impl<'a> Search<'a> { words_limit: 10, rtxn, index, - distribution_shift: None, - embedder_name: None, + semantic: None, time_budget: TimeBudget::max(), } } @@ -75,8 +77,13 @@ impl<'a> Search<'a> { self } - pub fn vector(&mut self, vector: Vec) -> &mut Search<'a> { - self.vector = Some(vector); + pub fn semantic( + &mut self, + embedder_name: String, + embedder: Arc, + vector: Option>, + ) -> &mut Search<'a> { + self.semantic = Some(SemanticSearch { embedder_name, embedder, vector }); self } @@ -133,19 +140,6 @@ impl<'a> Search<'a> { self } - pub fn distribution_shift( - &mut self, - distribution_shift: Option, - ) -> &mut Search<'a> { - self.distribution_shift = distribution_shift; - self - } - - pub fn embedder_name(&mut self, embedder_name: impl Into) -> &mut Search<'a> { - self.embedder_name = Some(embedder_name.into()); - self - } - pub fn time_budget(&mut self, time_budget: TimeBudget) -> &mut Search<'a> { self.time_budget = time_budget; self @@ -161,15 +155,6 @@ impl<'a> Search<'a> { } pub fn execute(&self) -> Result { - let embedder_name; - let embedder_name = match &self.embedder_name { - Some(embedder_name) => embedder_name, - None => { - embedder_name = self.index.default_embedding_name(self.rtxn)?; - &embedder_name - } - }; - let mut ctx = SearchContext::new(self.index, self.rtxn); if let Some(searchable_attributes) = self.searchable_attributes { @@ -184,21 +169,23 @@ impl<'a> Search<'a> { document_scores, degraded, used_negative_operator, - } = match self.vector.as_ref() { - Some(vector) => execute_vector_search( - &mut ctx, - vector, - self.scoring_strategy, - universe, - &self.sort_criteria, - self.geo_strategy, - self.offset, - self.limit, - self.distribution_shift, - embedder_name, - self.time_budget.clone(), - )?, - None => execute_search( + } = match self.semantic.as_ref() { + Some(SemanticSearch { vector: Some(vector), embedder_name, embedder }) => { + execute_vector_search( + &mut ctx, + vector, + self.scoring_strategy, + universe, + &self.sort_criteria, + self.geo_strategy, + self.offset, + self.limit, + embedder_name, + embedder, + self.time_budget.clone(), + )? + } + _ => execute_search( &mut ctx, self.query.as_deref(), self.terms_matching_strategy, @@ -237,7 +224,6 @@ impl fmt::Debug for Search<'_> { fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { let Search { query, - vector: _, filter, offset, limit, @@ -250,8 +236,7 @@ impl fmt::Debug for Search<'_> { exhaustive_number_hits, rtxn: _, index: _, - distribution_shift, - embedder_name, + semantic, time_budget, } = self; f.debug_struct("Search") @@ -266,8 +251,10 @@ impl fmt::Debug for Search<'_> { .field("scoring_strategy", scoring_strategy) .field("exhaustive_number_hits", exhaustive_number_hits) .field("words_limit", words_limit) - .field("distribution_shift", distribution_shift) - .field("embedder_name", embedder_name) + .field( + "semantic.embedder_name", + &semantic.as_ref().map(|semantic| &semantic.embedder_name), + ) .field("time_budget", time_budget) .finish() } diff --git a/milli/src/search/new/mod.rs b/milli/src/search/new/mod.rs index 1f0ae7b29..617068ef8 100644 --- a/milli/src/search/new/mod.rs +++ b/milli/src/search/new/mod.rs @@ -52,7 +52,7 @@ use self::vector_sort::VectorSort; use crate::error::FieldIdMapMissingEntry; use crate::score_details::{ScoreDetails, ScoringStrategy}; use crate::search::new::distinct::apply_distinct_rule; -use crate::vector::DistributionShift; +use crate::vector::Embedder; use crate::{ AscDesc, DocumentId, FieldId, Filter, Index, Member, Result, TermsMatchingStrategy, TimeBudget, UserError, @@ -298,8 +298,8 @@ fn get_ranking_rules_for_vector<'ctx>( geo_strategy: geo_sort::Strategy, limit_plus_offset: usize, target: &[f32], - distribution_shift: Option, embedder_name: &str, + embedder: &Embedder, ) -> Result>> { // query graph search @@ -325,8 +325,8 @@ fn get_ranking_rules_for_vector<'ctx>( target.to_vec(), vector_candidates, limit_plus_offset, - distribution_shift, embedder_name, + embedder, )?; ranking_rules.push(Box::new(vector_sort)); vector = true; @@ -548,8 +548,8 @@ pub fn execute_vector_search( geo_strategy: geo_sort::Strategy, from: usize, length: usize, - distribution_shift: Option, embedder_name: &str, + embedder: &Embedder, time_budget: TimeBudget, ) -> Result { check_sort_criteria(ctx, sort_criteria.as_ref())?; @@ -562,8 +562,8 @@ pub fn execute_vector_search( geo_strategy, from + length, vector, - distribution_shift, embedder_name, + embedder, )?; let mut placeholder_search_logger = logger::DefaultSearchLogger; diff --git a/milli/src/search/new/vector_sort.rs b/milli/src/search/new/vector_sort.rs index 476477218..de272ed47 100644 --- a/milli/src/search/new/vector_sort.rs +++ b/milli/src/search/new/vector_sort.rs @@ -5,7 +5,7 @@ use roaring::RoaringBitmap; use super::ranking_rules::{RankingRule, RankingRuleOutput, RankingRuleQueryTrait}; use crate::score_details::{self, ScoreDetails}; -use crate::vector::DistributionShift; +use crate::vector::{DistributionShift, Embedder}; use crate::{DocumentId, Result, SearchContext, SearchLogger}; pub struct VectorSort { @@ -24,8 +24,8 @@ impl VectorSort { target: Vec, vector_candidates: RoaringBitmap, limit: usize, - distribution_shift: Option, embedder_name: &str, + embedder: &Embedder, ) -> Result { let embedder_index = ctx .index @@ -39,7 +39,7 @@ impl VectorSort { vector_candidates, cached_sorted_docids: Default::default(), limit, - distribution_shift, + distribution_shift: embedder.distribution(), embedder_index, }) } From 466d718a05c8c3b69738816594f0ed5a45e63bbe Mon Sep 17 00:00:00 2001 From: Louis Dureuil Date: Thu, 28 Mar 2024 11:51:41 +0100 Subject: [PATCH 78/86] Fix test --- milli/src/update/index_documents/mod.rs | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/milli/src/update/index_documents/mod.rs b/milli/src/update/index_documents/mod.rs index dbacb4002..d534661da 100644 --- a/milli/src/update/index_documents/mod.rs +++ b/milli/src/update/index_documents/mod.rs @@ -2672,7 +2672,16 @@ mod tests { .unwrap(); let rtxn = index.read_txn().unwrap(); - let res = index.search(&rtxn).vector([0.0, 1.0, 2.0].to_vec()).execute().unwrap(); + let mut embedding_configs = index.embedding_configs(&rtxn).unwrap(); + let (embedder_name, embedder) = embedding_configs.pop().unwrap(); + let embedder = + std::sync::Arc::new(crate::vector::Embedder::new(embedder.embedder_options).unwrap()); + assert_eq!("manual", embedder_name); + let res = index + .search(&rtxn) + .semantic(embedder_name, embedder, Some([0.0, 1.0, 2.0].to_vec())) + .execute() + .unwrap(); assert_eq!(res.documents_ids.len(), 3); } From 4564a38ae7d94ce36cc3a20a34b59bea6162d534 Mon Sep 17 00:00:00 2001 From: Louis Dureuil Date: Thu, 28 Mar 2024 12:06:59 +0100 Subject: [PATCH 79/86] Bail earlier when the experimental feature is not enabled --- .../src/routes/indexes/facet_search.rs | 4 ++-- meilisearch/src/routes/indexes/search.rs | 21 ++++++++++++------- meilisearch/src/routes/multi_search.rs | 13 ++++++------ meilisearch/src/search.rs | 17 ++------------- 4 files changed, 24 insertions(+), 31 deletions(-) diff --git a/meilisearch/src/routes/indexes/facet_search.rs b/meilisearch/src/routes/indexes/facet_search.rs index 56880a472..4d6950988 100644 --- a/meilisearch/src/routes/indexes/facet_search.rs +++ b/meilisearch/src/routes/indexes/facet_search.rs @@ -74,10 +74,10 @@ pub async fn search( let index = index_scheduler.index(&index_uid)?; let features = index_scheduler.features(); - let search_kind = search_kind(&search_query, &index_scheduler, &index)?; + let search_kind = search_kind(&search_query, &index_scheduler, &index, features)?; let _permit = search_queue.try_get_search_permit().await?; let search_result = tokio::task::spawn_blocking(move || { - perform_facet_search(&index, search_query, facet_query, facet_name, features, search_kind) + perform_facet_search(&index, search_query, facet_query, facet_name, search_kind) }) .await?; diff --git a/meilisearch/src/routes/indexes/search.rs b/meilisearch/src/routes/indexes/search.rs index a5fe3c5d6..0f7d3b1ee 100644 --- a/meilisearch/src/routes/indexes/search.rs +++ b/meilisearch/src/routes/indexes/search.rs @@ -1,7 +1,7 @@ use actix_web::web::Data; use actix_web::{web, HttpRequest, HttpResponse}; use deserr::actix_web::{AwebJson, AwebQueryParameter}; -use index_scheduler::IndexScheduler; +use index_scheduler::{IndexScheduler, RoFeatures}; use meilisearch_types::deserr::query_params::Param; use meilisearch_types::deserr::{DeserrJsonError, DeserrQueryParamError}; use meilisearch_types::error::deserr_codes::*; @@ -204,12 +204,11 @@ pub async fn search_with_url_query( let index = index_scheduler.index(&index_uid)?; let features = index_scheduler.features(); - let search_kind = search_kind(&query, index_scheduler.get_ref(), &index)?; + let search_kind = search_kind(&query, index_scheduler.get_ref(), &index, features)?; let _permit = search_queue.try_get_search_permit().await?; let search_result = - tokio::task::spawn_blocking(move || perform_search(&index, query, features, search_kind)) - .await?; + tokio::task::spawn_blocking(move || perform_search(&index, query, search_kind)).await?; if let Ok(ref search_result) = search_result { aggregate.succeed(search_result); } @@ -245,12 +244,11 @@ pub async fn search_with_post( let features = index_scheduler.features(); - let search_kind = search_kind(&query, index_scheduler.get_ref(), &index)?; + let search_kind = search_kind(&query, index_scheduler.get_ref(), &index, features)?; let _permit = search_queue.try_get_search_permit().await?; let search_result = - tokio::task::spawn_blocking(move || perform_search(&index, query, features, search_kind)) - .await?; + tokio::task::spawn_blocking(move || perform_search(&index, query, search_kind)).await?; if let Ok(ref search_result) = search_result { aggregate.succeed(search_result); if search_result.degraded { @@ -269,7 +267,16 @@ pub fn search_kind( query: &SearchQuery, index_scheduler: &IndexScheduler, index: &milli::Index, + features: RoFeatures, ) -> Result { + if query.vector.is_some() { + features.check_vector("Passing `vector` as a query parameter")?; + } + + if query.hybrid.is_some() { + features.check_vector("Passing `hybrid` as a query parameter")?; + } + // regardless of anything, always do a semantic search when we don't have a vector and the query is whitespace or missing if query.vector.is_none() { match &query.q { diff --git a/meilisearch/src/routes/multi_search.rs b/meilisearch/src/routes/multi_search.rs index 04cd3f637..7b7cbd265 100644 --- a/meilisearch/src/routes/multi_search.rs +++ b/meilisearch/src/routes/multi_search.rs @@ -81,14 +81,13 @@ pub async fn multi_search_with_post( }) .with_index(query_index)?; - let search_kind = - search_kind(&query, index_scheduler.get_ref(), &index).with_index(query_index)?; + let search_kind = search_kind(&query, index_scheduler.get_ref(), &index, features) + .with_index(query_index)?; - let search_result = tokio::task::spawn_blocking(move || { - perform_search(&index, query, features, search_kind) - }) - .await - .with_index(query_index)?; + let search_result = + tokio::task::spawn_blocking(move || perform_search(&index, query, search_kind)) + .await + .with_index(query_index)?; search_results.push(SearchResultWithIndex { index_uid: index_uid.into_inner(), diff --git a/meilisearch/src/search.rs b/meilisearch/src/search.rs index 2a22cb2ce..7cb860f2e 100644 --- a/meilisearch/src/search.rs +++ b/meilisearch/src/search.rs @@ -6,7 +6,6 @@ use std::time::{Duration, Instant}; use deserr::Deserr; use either::Either; -use index_scheduler::RoFeatures; use indexmap::IndexMap; use meilisearch_auth::IndexSearchRules; use meilisearch_types::deserr::DeserrJsonError; @@ -448,21 +447,12 @@ fn prepare_search<'t>( index: &'t Index, rtxn: &'t RoTxn, query: &'t SearchQuery, - features: RoFeatures, search_kind: &SearchKind, time_budget: TimeBudget, ) -> Result<(milli::Search<'t>, bool, usize, usize), MeilisearchHttpError> { let mut search = index.search(rtxn); search.time_budget(time_budget); - if query.vector.is_some() { - features.check_vector("Passing `vector` as a query parameter")?; - } - - if query.hybrid.is_some() { - features.check_vector("Passing `hybrid` as a query parameter")?; - } - match search_kind { SearchKind::KeywordOnly => { if let Some(q) = &query.q { @@ -551,7 +541,6 @@ fn prepare_search<'t>( pub fn perform_search( index: &Index, query: SearchQuery, - features: RoFeatures, search_kind: SearchKind, ) -> Result { let before_search = Instant::now(); @@ -562,7 +551,7 @@ pub fn perform_search( }; let (search, is_finite_pagination, max_total_hits, offset) = - prepare_search(index, &rtxn, &query, features, &search_kind, time_budget)?; + prepare_search(index, &rtxn, &query, &search_kind, time_budget)?; let milli::SearchResult { documents_ids, @@ -780,7 +769,6 @@ pub fn perform_facet_search( search_query: SearchQuery, facet_query: Option, facet_name: String, - features: RoFeatures, search_kind: SearchKind, ) -> Result { let before_search = Instant::now(); @@ -790,8 +778,7 @@ pub fn perform_facet_search( None => TimeBudget::default(), }; - let (search, _, _, _) = - prepare_search(index, &rtxn, &search_query, features, &search_kind, time_budget)?; + let (search, _, _, _) = prepare_search(index, &rtxn, &search_query, &search_kind, time_budget)?; let mut facet_search = SearchForFacetValues::new( facet_name, search, From 3c6e9851a4bcf85a446f1584b569a056c4139a90 Mon Sep 17 00:00:00 2001 From: Louis Dureuil Date: Thu, 28 Mar 2024 15:26:21 +0100 Subject: [PATCH 80/86] Correct error formatting --- milli/src/error.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/milli/src/error.rs b/milli/src/error.rs index aba80b475..1d61bef63 100644 --- a/milli/src/error.rs +++ b/milli/src/error.rs @@ -196,7 +196,7 @@ only composed of alphanumeric characters (a-z A-Z 0-9), hyphens (-) and undersco InvalidPromptForEmbeddings(String, crate::prompt::error::NewPromptError), #[error("Too many embedders in the configuration. Found {0}, but limited to 256.")] TooManyEmbedders(usize), - #[error("Cannot find embedder with name {0}.")] + #[error("Cannot find embedder with name `{0}`.")] InvalidEmbedder(String), #[error("Too many vectors for document with id {0}: found {1}, but limited to 256.")] TooManyVectors(String, usize), From 1ff2a2d6fb20c25a51f87869b833b17f307f659f Mon Sep 17 00:00:00 2001 From: Louis Dureuil Date: Wed, 3 Apr 2024 09:35:07 +0200 Subject: [PATCH 81/86] Add semanticHitCount --- .../src/analytics/segment_analytics.rs | 1 + meilisearch/src/search.rs | 31 ++++--- meilisearch/tests/search/hybrid.rs | 30 +++++++ milli/src/search/hybrid.rs | 82 +++++++++++++------ 4 files changed, 108 insertions(+), 36 deletions(-) diff --git a/meilisearch/src/analytics/segment_analytics.rs b/meilisearch/src/analytics/segment_analytics.rs index fcf4d9144..c49a04576 100644 --- a/meilisearch/src/analytics/segment_analytics.rs +++ b/meilisearch/src/analytics/segment_analytics.rs @@ -760,6 +760,7 @@ impl SearchAggregator { query: _, processing_time_ms, hits_info: _, + semantic_hit_count: _, facet_distribution: _, facet_stats: _, degraded, diff --git a/meilisearch/src/search.rs b/meilisearch/src/search.rs index 7cb860f2e..85438e816 100644 --- a/meilisearch/src/search.rs +++ b/meilisearch/src/search.rs @@ -385,6 +385,9 @@ pub struct SearchResult { #[serde(skip_serializing_if = "Option::is_none")] pub facet_stats: Option>, + #[serde(skip_serializing_if = "Option::is_none")] + pub semantic_hit_count: Option, + // These fields are only used for analytics purposes #[serde(skip)] pub degraded: bool, @@ -553,16 +556,23 @@ pub fn perform_search( let (search, is_finite_pagination, max_total_hits, offset) = prepare_search(index, &rtxn, &query, &search_kind, time_budget)?; - let milli::SearchResult { - documents_ids, - matching_words, - candidates, - document_scores, - degraded, - used_negative_operator, - .. - } = match &search_kind { - SearchKind::KeywordOnly | SearchKind::SemanticOnly { .. } => search.execute()?, + let ( + milli::SearchResult { + documents_ids, + matching_words, + candidates, + document_scores, + degraded, + used_negative_operator, + }, + semantic_hit_count, + ) = match &search_kind { + SearchKind::KeywordOnly => (search.execute()?, None), + SearchKind::SemanticOnly { .. } => { + let results = search.execute()?; + let semantic_hit_count = results.document_scores.len() as u32; + (results, Some(semantic_hit_count)) + } SearchKind::Hybrid { semantic_ratio, .. } => search.execute_hybrid(*semantic_ratio)?, }; @@ -760,6 +770,7 @@ pub fn perform_search( facet_stats, degraded, used_negative_operator, + semantic_hit_count, }; Ok(result) } diff --git a/meilisearch/tests/search/hybrid.rs b/meilisearch/tests/search/hybrid.rs index 8decb7ded..77c4f30a3 100644 --- a/meilisearch/tests/search/hybrid.rs +++ b/meilisearch/tests/search/hybrid.rs @@ -77,6 +77,16 @@ async fn simple_search() { .await; snapshot!(code, @"200 OK"); snapshot!(response["hits"], @r###"[{"title":"Captain Planet","desc":"He's not part of the Marvel Cinematic Universe","id":"2","_vectors":{"default":[1.0,2.0]}},{"title":"Captain Marvel","desc":"a Shazam ersatz","id":"3","_vectors":{"default":[2.0,3.0]}},{"title":"Shazam!","desc":"a Captain Marvel ersatz","id":"1","_vectors":{"default":[1.0,3.0]}}]"###); + snapshot!(response["semanticHitCount"], @"0"); + + let (response, code) = index + .search_post( + json!({"q": "Captain", "vector": [1.0, 1.0], "hybrid": {"semanticRatio": 0.5}}), + ) + .await; + snapshot!(code, @"200 OK"); + snapshot!(response["hits"], @r###"[{"title":"Captain Planet","desc":"He's not part of the Marvel Cinematic Universe","id":"2","_vectors":{"default":[1.0,2.0]}},{"title":"Captain Marvel","desc":"a Shazam ersatz","id":"3","_vectors":{"default":[2.0,3.0]}},{"title":"Shazam!","desc":"a Captain Marvel ersatz","id":"1","_vectors":{"default":[1.0,3.0]},"_semanticScore":0.9472136}]"###); + snapshot!(response["semanticHitCount"], @"1"); let (response, code) = index .search_post( @@ -85,6 +95,7 @@ async fn simple_search() { .await; snapshot!(code, @"200 OK"); snapshot!(response["hits"], @r###"[{"title":"Captain Marvel","desc":"a Shazam ersatz","id":"3","_vectors":{"default":[2.0,3.0]},"_semanticScore":0.99029034},{"title":"Captain Planet","desc":"He's not part of the Marvel Cinematic Universe","id":"2","_vectors":{"default":[1.0,2.0]},"_semanticScore":0.97434163},{"title":"Shazam!","desc":"a Captain Marvel ersatz","id":"1","_vectors":{"default":[1.0,3.0]},"_semanticScore":0.9472136}]"###); + snapshot!(response["semanticHitCount"], @"3"); } #[actix_rt::test] @@ -136,6 +147,7 @@ async fn highlighter() { .await; snapshot!(code, @"200 OK"); snapshot!(response["hits"], @r###"[{"title":"Captain Marvel","desc":"a Shazam ersatz","id":"3","_vectors":{"default":[2.0,3.0]},"_formatted":{"title":"Captain Marvel","desc":"a Shazam ersatz","id":"3","_vectors":{"default":["2.0","3.0"]}}},{"title":"Shazam!","desc":"a Captain Marvel ersatz","id":"1","_vectors":{"default":[1.0,3.0]},"_formatted":{"title":"Shazam!","desc":"a **BEGIN**Captain**END** **BEGIN**Marvel**END** ersatz","id":"1","_vectors":{"default":["1.0","3.0"]}}},{"title":"Captain Planet","desc":"He's not part of the Marvel Cinematic Universe","id":"2","_vectors":{"default":[1.0,2.0]},"_formatted":{"title":"Captain Planet","desc":"He's not part of the **BEGIN**Marvel**END** Cinematic Universe","id":"2","_vectors":{"default":["1.0","2.0"]}}}]"###); + snapshot!(response["semanticHitCount"], @"0"); let (response, code) = index .search_post(json!({"q": "Captain Marvel", "vector": [1.0, 1.0], @@ -149,6 +161,7 @@ async fn highlighter() { .await; snapshot!(code, @"200 OK"); snapshot!(response["hits"], @r###"[{"title":"Captain Marvel","desc":"a Shazam ersatz","id":"3","_vectors":{"default":[2.0,3.0]},"_formatted":{"title":"Captain Marvel","desc":"a Shazam ersatz","id":"3","_vectors":{"default":["2.0","3.0"]}},"_semanticScore":0.99029034},{"title":"Captain Planet","desc":"He's not part of the Marvel Cinematic Universe","id":"2","_vectors":{"default":[1.0,2.0]},"_formatted":{"title":"Captain Planet","desc":"He's not part of the **BEGIN**Marvel**END** Cinematic Universe","id":"2","_vectors":{"default":["1.0","2.0"]}},"_semanticScore":0.97434163},{"title":"Shazam!","desc":"a Captain Marvel ersatz","id":"1","_vectors":{"default":[1.0,3.0]},"_formatted":{"title":"Shazam!","desc":"a **BEGIN**Captain**END** **BEGIN**Marvel**END** ersatz","id":"1","_vectors":{"default":["1.0","3.0"]}},"_semanticScore":0.9472136}]"###); + snapshot!(response["semanticHitCount"], @"3"); // no highlighting on full semantic let (response, code) = index @@ -163,6 +176,7 @@ async fn highlighter() { .await; snapshot!(code, @"200 OK"); snapshot!(response["hits"], @r###"[{"title":"Captain Marvel","desc":"a Shazam ersatz","id":"3","_vectors":{"default":[2.0,3.0]},"_formatted":{"title":"Captain Marvel","desc":"a Shazam ersatz","id":"3","_vectors":{"default":["2.0","3.0"]}},"_semanticScore":0.99029034},{"title":"Captain Planet","desc":"He's not part of the Marvel Cinematic Universe","id":"2","_vectors":{"default":[1.0,2.0]},"_formatted":{"title":"Captain Planet","desc":"He's not part of the Marvel Cinematic Universe","id":"2","_vectors":{"default":["1.0","2.0"]}},"_semanticScore":0.97434163},{"title":"Shazam!","desc":"a Captain Marvel ersatz","id":"1","_vectors":{"default":[1.0,3.0]},"_formatted":{"title":"Shazam!","desc":"a Captain Marvel ersatz","id":"1","_vectors":{"default":["1.0","3.0"]}}}]"###); + snapshot!(response["semanticHitCount"], @"3"); } #[actix_rt::test] @@ -250,4 +264,20 @@ async fn single_document() { snapshot!(code, @"200 OK"); snapshot!(response["hits"][0], @r###"{"title":"Shazam!","desc":"a Captain Marvel ersatz","id":"1","_vectors":{"default":[1.0,3.0]},"_rankingScore":1.0,"_semanticScore":1.0}"###); + snapshot!(response["semanticHitCount"], @"1"); +} + +#[actix_rt::test] +async fn query_combination() { + let server = Server::new().await; + let index = index_with_documents(&server, &SIMPLE_SEARCH_DOCUMENTS).await; + + // search without query and vector, but with hybrid => still placeholder + let (response, code) = index + .search_post(json!({"hybrid": {"semanticRatio": 1.0}, "showRankingScore": true})) + .await; + + snapshot!(code, @"200 OK"); + snapshot!(response["hits"][0], @r###"{"title":"Shazam!","desc":"a Captain Marvel ersatz","id":"1","_vectors":{"default":[1.0,3.0]},"_rankingScore":1.0,"_semanticScore":1.0}"###); + snapshot!(response["semanticHitCount"], @"1"); } diff --git a/milli/src/search/hybrid.rs b/milli/src/search/hybrid.rs index e45652206..fc13a5e1e 100644 --- a/milli/src/search/hybrid.rs +++ b/milli/src/search/hybrid.rs @@ -84,45 +84,73 @@ impl ScoreWithRatioResult { } } - fn merge(left: Self, right: Self, from: usize, length: usize) -> SearchResult { - let mut documents_ids = - Vec::with_capacity(left.document_scores.len() + right.document_scores.len()); - let mut document_scores = - Vec::with_capacity(left.document_scores.len() + right.document_scores.len()); + fn merge( + vector_results: Self, + keyword_results: Self, + from: usize, + length: usize, + ) -> (SearchResult, u32) { + #[derive(Clone, Copy)] + enum ResultSource { + Semantic, + Keyword, + } + let mut semantic_hit_count = 0; + + let mut documents_ids = Vec::with_capacity( + vector_results.document_scores.len() + keyword_results.document_scores.len(), + ); + let mut document_scores = Vec::with_capacity( + vector_results.document_scores.len() + keyword_results.document_scores.len(), + ); let mut documents_seen = RoaringBitmap::new(); - for (docid, (main_score, _sub_score)) in left + for ((docid, (main_score, _sub_score)), source) in vector_results .document_scores .into_iter() - .merge_by(right.document_scores.into_iter(), |(_, left), (_, right)| { - // the first value is the one with the greatest score - compare_scores(left, right).is_ge() - }) + .zip(std::iter::repeat(ResultSource::Semantic)) + .merge_by( + keyword_results + .document_scores + .into_iter() + .zip(std::iter::repeat(ResultSource::Keyword)), + |((_, left), _), ((_, right), _)| { + // the first value is the one with the greatest score + compare_scores(left, right).is_ge() + }, + ) // remove documents we already saw - .filter(|(docid, _)| documents_seen.insert(*docid)) + .filter(|((docid, _), _)| documents_seen.insert(*docid)) // start skipping **after** the filter .skip(from) // take **after** skipping .take(length) { + if let ResultSource::Semantic = source { + semantic_hit_count += 1; + } documents_ids.push(docid); // TODO: pass both scores to documents_score in some way? document_scores.push(main_score); } - SearchResult { - matching_words: right.matching_words, - candidates: left.candidates | right.candidates, - documents_ids, - document_scores, - degraded: left.degraded | right.degraded, - used_negative_operator: left.used_negative_operator | right.used_negative_operator, - } + ( + SearchResult { + matching_words: keyword_results.matching_words, + candidates: vector_results.candidates | keyword_results.candidates, + documents_ids, + document_scores, + degraded: vector_results.degraded | keyword_results.degraded, + used_negative_operator: vector_results.used_negative_operator + | keyword_results.used_negative_operator, + }, + semantic_hit_count, + ) } } impl<'a> Search<'a> { - pub fn execute_hybrid(&self, semantic_ratio: f32) -> Result { + pub fn execute_hybrid(&self, semantic_ratio: f32) -> Result<(SearchResult, Option)> { // TODO: find classier way to achieve that than to reset vector and query params // create separate keyword and semantic searches let mut search = Search { @@ -148,14 +176,16 @@ impl<'a> Search<'a> { // completely skip semantic search if the results of the keyword search are good enough if self.results_good_enough(&keyword_results, semantic_ratio) { - return Ok(keyword_results); + return Ok((keyword_results, Some(0))); } // no vector search against placeholder search - let Some(query) = search.query.take() else { return Ok(keyword_results) }; + let Some(query) = search.query.take() else { + return Ok((keyword_results, Some(0))); + }; // no embedder, no semantic search let Some(SemanticSearch { vector, embedder_name, embedder }) = semantic else { - return Ok(keyword_results); + return Ok((keyword_results, Some(0))); }; let vector_query = match vector { @@ -166,7 +196,7 @@ impl<'a> Search<'a> { Ok(embedding) => embedding, Err(error) => { tracing::error!(error=%error, "Embedding failed"); - return Ok(keyword_results); + return Ok((keyword_results, Some(0))); } } } @@ -181,10 +211,10 @@ impl<'a> Search<'a> { let keyword_results = ScoreWithRatioResult::new(keyword_results, 1.0 - semantic_ratio); let vector_results = ScoreWithRatioResult::new(vector_results, semantic_ratio); - let merge_results = + let (merge_results, semantic_hit_count) = ScoreWithRatioResult::merge(vector_results, keyword_results, self.offset, self.limit); assert!(merge_results.documents_ids.len() <= self.limit); - Ok(merge_results) + Ok((merge_results, Some(semantic_hit_count))) } fn results_good_enough(&self, keyword_results: &SearchResult, semantic_ratio: f32) -> bool { From 7c27417a5d5d119d79d7dd130c4e9d085222ff70 Mon Sep 17 00:00:00 2001 From: Louis Dureuil Date: Wed, 3 Apr 2024 10:23:01 +0200 Subject: [PATCH 82/86] Add tests --- meilisearch/tests/search/hybrid.rs | 98 +++++++++++++++++++++++++++++- 1 file changed, 96 insertions(+), 2 deletions(-) diff --git a/meilisearch/tests/search/hybrid.rs b/meilisearch/tests/search/hybrid.rs index 77c4f30a3..77cfac3d9 100644 --- a/meilisearch/tests/search/hybrid.rs +++ b/meilisearch/tests/search/hybrid.rs @@ -278,6 +278,100 @@ async fn query_combination() { .await; snapshot!(code, @"200 OK"); - snapshot!(response["hits"][0], @r###"{"title":"Shazam!","desc":"a Captain Marvel ersatz","id":"1","_vectors":{"default":[1.0,3.0]},"_rankingScore":1.0,"_semanticScore":1.0}"###); - snapshot!(response["semanticHitCount"], @"1"); + snapshot!(response["hits"], @r###"[{"title":"Shazam!","desc":"a Captain Marvel ersatz","id":"1","_vectors":{"default":[1.0,3.0]},"_rankingScore":1.0},{"title":"Captain Planet","desc":"He's not part of the Marvel Cinematic Universe","id":"2","_vectors":{"default":[1.0,2.0]},"_rankingScore":1.0},{"title":"Captain Marvel","desc":"a Shazam ersatz","id":"3","_vectors":{"default":[2.0,3.0]},"_rankingScore":1.0}]"###); + snapshot!(response["semanticHitCount"], @"null"); + + // same with a different semantic ratio + let (response, code) = index + .search_post(json!({"hybrid": {"semanticRatio": 0.76}, "showRankingScore": true})) + .await; + + snapshot!(code, @"200 OK"); + snapshot!(response["hits"], @r###"[{"title":"Shazam!","desc":"a Captain Marvel ersatz","id":"1","_vectors":{"default":[1.0,3.0]},"_rankingScore":1.0},{"title":"Captain Planet","desc":"He's not part of the Marvel Cinematic Universe","id":"2","_vectors":{"default":[1.0,2.0]},"_rankingScore":1.0},{"title":"Captain Marvel","desc":"a Shazam ersatz","id":"3","_vectors":{"default":[2.0,3.0]},"_rankingScore":1.0}]"###); + snapshot!(response["semanticHitCount"], @"null"); + + // wrong vector dimensions + let (response, code) = index + .search_post(json!({"vector": [1.0, 0.0, 1.0], "hybrid": {"semanticRatio": 1.0}, "showRankingScore": true})) + .await; + + snapshot!(code, @"400 Bad Request"); + snapshot!(response, @r###" + { + "message": "Invalid vector dimensions: expected: `2`, found: `3`.", + "code": "invalid_vector_dimensions", + "type": "invalid_request", + "link": "https://docs.meilisearch.com/errors#invalid_vector_dimensions" + } + "###); + + // full vector + let (response, code) = index + .search_post(json!({"vector": [1.0, 0.0], "hybrid": {"semanticRatio": 1.0}, "showRankingScore": true})) + .await; + + snapshot!(code, @"200 OK"); + snapshot!(response["hits"], @r###"[{"title":"Captain Marvel","desc":"a Shazam ersatz","id":"3","_vectors":{"default":[2.0,3.0]},"_rankingScore":0.7773500680923462,"_semanticScore":0.77735007},{"title":"Captain Planet","desc":"He's not part of the Marvel Cinematic Universe","id":"2","_vectors":{"default":[1.0,2.0]},"_rankingScore":0.7236068248748779,"_semanticScore":0.7236068},{"title":"Shazam!","desc":"a Captain Marvel ersatz","id":"1","_vectors":{"default":[1.0,3.0]},"_rankingScore":0.6581138968467712,"_semanticScore":0.6581139}]"###); + snapshot!(response["semanticHitCount"], @"3"); + + // full keyword, without a query + let (response, code) = index + .search_post(json!({"vector": [1.0, 0.0], "hybrid": {"semanticRatio": 0.0}, "showRankingScore": true})) + .await; + + snapshot!(code, @"200 OK"); + snapshot!(response["hits"], @r###"[{"title":"Shazam!","desc":"a Captain Marvel ersatz","id":"1","_vectors":{"default":[1.0,3.0]},"_rankingScore":1.0},{"title":"Captain Planet","desc":"He's not part of the Marvel Cinematic Universe","id":"2","_vectors":{"default":[1.0,2.0]},"_rankingScore":1.0},{"title":"Captain Marvel","desc":"a Shazam ersatz","id":"3","_vectors":{"default":[2.0,3.0]},"_rankingScore":1.0}]"###); + snapshot!(response["semanticHitCount"], @"null"); + + // query + vector, full keyword => keyword + let (response, code) = index + .search_post(json!({"q": "Captain", "vector": [1.0, 0.0], "hybrid": {"semanticRatio": 0.0}, "showRankingScore": true})) + .await; + + snapshot!(code, @"200 OK"); + snapshot!(response["hits"], @r###"[{"title":"Captain Planet","desc":"He's not part of the Marvel Cinematic Universe","id":"2","_vectors":{"default":[1.0,2.0]},"_rankingScore":0.996969696969697},{"title":"Captain Marvel","desc":"a Shazam ersatz","id":"3","_vectors":{"default":[2.0,3.0]},"_rankingScore":0.996969696969697},{"title":"Shazam!","desc":"a Captain Marvel ersatz","id":"1","_vectors":{"default":[1.0,3.0]},"_rankingScore":0.8848484848484849}]"###); + snapshot!(response["semanticHitCount"], @"null"); + + // query + vector, no hybrid keyword => + let (response, code) = index + .search_post(json!({"q": "Captain", "vector": [1.0, 0.0], "showRankingScore": true})) + .await; + + snapshot!(code, @"400 Bad Request"); + snapshot!(response, @r###" + { + "message": "Invalid request: missing `hybrid` parameter when both `q` and `vector` are present.", + "code": "missing_search_hybrid", + "type": "invalid_request", + "link": "https://docs.meilisearch.com/errors#missing_search_hybrid" + } + "###); + + // full vector, without a vector => error + let (response, code) = index + .search_post( + json!({"q": "Captain", "hybrid": {"semanticRatio": 1.0}, "showRankingScore": true}), + ) + .await; + + snapshot!(code, @"400 Bad Request"); + snapshot!(response, @r###" + { + "message": "Error while generating embeddings: user error: attempt to embed the following text in a configuration where embeddings must be user provided: \"Captain\"", + "code": "vector_embedding_error", + "type": "invalid_request", + "link": "https://docs.meilisearch.com/errors#vector_embedding_error" + } + "###); + + // hybrid without a vector => full keyword + let (response, code) = index + .search_post( + json!({"q": "Planet", "hybrid": {"semanticRatio": 0.99}, "showRankingScore": true}), + ) + .await; + + snapshot!(code, @"200 OK"); + snapshot!(response["hits"], @r###"[{"title":"Captain Planet","desc":"He's not part of the Marvel Cinematic Universe","id":"2","_vectors":{"default":[1.0,2.0]},"_rankingScore":0.9848484848484848}]"###); + snapshot!(response["semanticHitCount"], @"0"); } From 355e5282b24bbf5e66fe1f18b25620ae70546e61 Mon Sep 17 00:00:00 2001 From: Louis Dureuil Date: Wed, 3 Apr 2024 14:29:17 +0200 Subject: [PATCH 83/86] Remove `_semanticScore` --- meilisearch/src/search.rs | 15 +-------------- meilisearch/tests/search/hybrid.rs | 18 ++++++++++-------- meilisearch/tests/search/mod.rs | 13 ++++++++----- 3 files changed, 19 insertions(+), 27 deletions(-) diff --git a/meilisearch/src/search.rs b/meilisearch/src/search.rs index 85438e816..47c2b5f4b 100644 --- a/meilisearch/src/search.rs +++ b/meilisearch/src/search.rs @@ -13,7 +13,7 @@ use meilisearch_types::error::deserr_codes::*; use meilisearch_types::error::ResponseError; use meilisearch_types::heed::RoTxn; use meilisearch_types::index_uid::IndexUid; -use meilisearch_types::milli::score_details::{self, ScoreDetails, ScoringStrategy}; +use meilisearch_types::milli::score_details::{ScoreDetails, ScoringStrategy}; use meilisearch_types::milli::vector::Embedder; use meilisearch_types::milli::{FacetValueHit, OrderBy, SearchForFacetValues, TimeBudget}; use meilisearch_types::settings::DEFAULT_PAGINATION_MAX_TOTAL_HITS; @@ -368,8 +368,6 @@ pub struct SearchHit { pub ranking_score: Option, #[serde(rename = "_rankingScoreDetails", skip_serializing_if = "Option::is_none")] pub ranking_score_details: Option>, - #[serde(rename = "_semanticScore", skip_serializing_if = "Option::is_none")] - pub semantic_score: Option, } #[derive(Serialize, Debug, Clone, PartialEq)] @@ -683,16 +681,6 @@ pub fn perform_search( insert_geo_distance(sort, &mut document); } - let mut semantic_score = None; - for details in &score { - if let ScoreDetails::Vector(score_details::Vector { similarity: Some(similarity) }) = - details - { - semantic_score = Some(*similarity); - break; - } - } - let ranking_score = query.show_ranking_score.then(|| ScoreDetails::global_score(score.iter())); let ranking_score_details = @@ -704,7 +692,6 @@ pub fn perform_search( matches_position, ranking_score_details, ranking_score, - semantic_score, }; documents.push(hit); } diff --git a/meilisearch/tests/search/hybrid.rs b/meilisearch/tests/search/hybrid.rs index 77cfac3d9..637242bc7 100644 --- a/meilisearch/tests/search/hybrid.rs +++ b/meilisearch/tests/search/hybrid.rs @@ -81,20 +81,20 @@ async fn simple_search() { let (response, code) = index .search_post( - json!({"q": "Captain", "vector": [1.0, 1.0], "hybrid": {"semanticRatio": 0.5}}), + json!({"q": "Captain", "vector": [1.0, 1.0], "hybrid": {"semanticRatio": 0.5}, "showRankingScore": true}), ) .await; snapshot!(code, @"200 OK"); - snapshot!(response["hits"], @r###"[{"title":"Captain Planet","desc":"He's not part of the Marvel Cinematic Universe","id":"2","_vectors":{"default":[1.0,2.0]}},{"title":"Captain Marvel","desc":"a Shazam ersatz","id":"3","_vectors":{"default":[2.0,3.0]}},{"title":"Shazam!","desc":"a Captain Marvel ersatz","id":"1","_vectors":{"default":[1.0,3.0]},"_semanticScore":0.9472136}]"###); + snapshot!(response["hits"], @r###"[{"title":"Captain Planet","desc":"He's not part of the Marvel Cinematic Universe","id":"2","_vectors":{"default":[1.0,2.0]},"_rankingScore":0.996969696969697},{"title":"Captain Marvel","desc":"a Shazam ersatz","id":"3","_vectors":{"default":[2.0,3.0]},"_rankingScore":0.996969696969697},{"title":"Shazam!","desc":"a Captain Marvel ersatz","id":"1","_vectors":{"default":[1.0,3.0]},"_rankingScore":0.9472135901451112}]"###); snapshot!(response["semanticHitCount"], @"1"); let (response, code) = index .search_post( - json!({"q": "Captain", "vector": [1.0, 1.0], "hybrid": {"semanticRatio": 0.8}}), + json!({"q": "Captain", "vector": [1.0, 1.0], "hybrid": {"semanticRatio": 0.8}, "showRankingScore": true}), ) .await; snapshot!(code, @"200 OK"); - snapshot!(response["hits"], @r###"[{"title":"Captain Marvel","desc":"a Shazam ersatz","id":"3","_vectors":{"default":[2.0,3.0]},"_semanticScore":0.99029034},{"title":"Captain Planet","desc":"He's not part of the Marvel Cinematic Universe","id":"2","_vectors":{"default":[1.0,2.0]},"_semanticScore":0.97434163},{"title":"Shazam!","desc":"a Captain Marvel ersatz","id":"1","_vectors":{"default":[1.0,3.0]},"_semanticScore":0.9472136}]"###); + snapshot!(response["hits"], @r###"[{"title":"Captain Marvel","desc":"a Shazam ersatz","id":"3","_vectors":{"default":[2.0,3.0]},"_rankingScore":0.990290343761444},{"title":"Captain Planet","desc":"He's not part of the Marvel Cinematic Universe","id":"2","_vectors":{"default":[1.0,2.0]},"_rankingScore":0.974341630935669},{"title":"Shazam!","desc":"a Captain Marvel ersatz","id":"1","_vectors":{"default":[1.0,3.0]},"_rankingScore":0.9472135901451112}]"###); snapshot!(response["semanticHitCount"], @"3"); } @@ -152,6 +152,7 @@ async fn highlighter() { let (response, code) = index .search_post(json!({"q": "Captain Marvel", "vector": [1.0, 1.0], "hybrid": {"semanticRatio": 0.8}, + "showRankingScore": true, "attributesToHighlight": [ "desc" ], @@ -160,13 +161,14 @@ async fn highlighter() { })) .await; snapshot!(code, @"200 OK"); - snapshot!(response["hits"], @r###"[{"title":"Captain Marvel","desc":"a Shazam ersatz","id":"3","_vectors":{"default":[2.0,3.0]},"_formatted":{"title":"Captain Marvel","desc":"a Shazam ersatz","id":"3","_vectors":{"default":["2.0","3.0"]}},"_semanticScore":0.99029034},{"title":"Captain Planet","desc":"He's not part of the Marvel Cinematic Universe","id":"2","_vectors":{"default":[1.0,2.0]},"_formatted":{"title":"Captain Planet","desc":"He's not part of the **BEGIN**Marvel**END** Cinematic Universe","id":"2","_vectors":{"default":["1.0","2.0"]}},"_semanticScore":0.97434163},{"title":"Shazam!","desc":"a Captain Marvel ersatz","id":"1","_vectors":{"default":[1.0,3.0]},"_formatted":{"title":"Shazam!","desc":"a **BEGIN**Captain**END** **BEGIN**Marvel**END** ersatz","id":"1","_vectors":{"default":["1.0","3.0"]}},"_semanticScore":0.9472136}]"###); + snapshot!(response["hits"], @r###"[{"title":"Captain Marvel","desc":"a Shazam ersatz","id":"3","_vectors":{"default":[2.0,3.0]},"_formatted":{"title":"Captain Marvel","desc":"a Shazam ersatz","id":"3","_vectors":{"default":["2.0","3.0"]}},"_rankingScore":0.990290343761444},{"title":"Captain Planet","desc":"He's not part of the Marvel Cinematic Universe","id":"2","_vectors":{"default":[1.0,2.0]},"_formatted":{"title":"Captain Planet","desc":"He's not part of the **BEGIN**Marvel**END** Cinematic Universe","id":"2","_vectors":{"default":["1.0","2.0"]}},"_rankingScore":0.974341630935669},{"title":"Shazam!","desc":"a Captain Marvel ersatz","id":"1","_vectors":{"default":[1.0,3.0]},"_formatted":{"title":"Shazam!","desc":"a **BEGIN**Captain**END** **BEGIN**Marvel**END** ersatz","id":"1","_vectors":{"default":["1.0","3.0"]}},"_rankingScore":0.9472135901451112}]"###); snapshot!(response["semanticHitCount"], @"3"); // no highlighting on full semantic let (response, code) = index .search_post(json!({"q": "Captain Marvel", "vector": [1.0, 1.0], "hybrid": {"semanticRatio": 1.0}, + "showRankingScore": true, "attributesToHighlight": [ "desc" ], @@ -175,7 +177,7 @@ async fn highlighter() { })) .await; snapshot!(code, @"200 OK"); - snapshot!(response["hits"], @r###"[{"title":"Captain Marvel","desc":"a Shazam ersatz","id":"3","_vectors":{"default":[2.0,3.0]},"_formatted":{"title":"Captain Marvel","desc":"a Shazam ersatz","id":"3","_vectors":{"default":["2.0","3.0"]}},"_semanticScore":0.99029034},{"title":"Captain Planet","desc":"He's not part of the Marvel Cinematic Universe","id":"2","_vectors":{"default":[1.0,2.0]},"_formatted":{"title":"Captain Planet","desc":"He's not part of the Marvel Cinematic Universe","id":"2","_vectors":{"default":["1.0","2.0"]}},"_semanticScore":0.97434163},{"title":"Shazam!","desc":"a Captain Marvel ersatz","id":"1","_vectors":{"default":[1.0,3.0]},"_formatted":{"title":"Shazam!","desc":"a Captain Marvel ersatz","id":"1","_vectors":{"default":["1.0","3.0"]}}}]"###); + snapshot!(response["hits"], @r###"[{"title":"Captain Marvel","desc":"a Shazam ersatz","id":"3","_vectors":{"default":[2.0,3.0]},"_formatted":{"title":"Captain Marvel","desc":"a Shazam ersatz","id":"3","_vectors":{"default":["2.0","3.0"]}},"_rankingScore":0.990290343761444},{"title":"Captain Planet","desc":"He's not part of the Marvel Cinematic Universe","id":"2","_vectors":{"default":[1.0,2.0]},"_formatted":{"title":"Captain Planet","desc":"He's not part of the Marvel Cinematic Universe","id":"2","_vectors":{"default":["1.0","2.0"]}},"_rankingScore":0.974341630935669},{"title":"Shazam!","desc":"a Captain Marvel ersatz","id":"1","_vectors":{"default":[1.0,3.0]},"_formatted":{"title":"Shazam!","desc":"a Captain Marvel ersatz","id":"1","_vectors":{"default":["1.0","3.0"]}},"_rankingScore":0.9472135901451112}]"###); snapshot!(response["semanticHitCount"], @"3"); } @@ -263,7 +265,7 @@ async fn single_document() { .await; snapshot!(code, @"200 OK"); - snapshot!(response["hits"][0], @r###"{"title":"Shazam!","desc":"a Captain Marvel ersatz","id":"1","_vectors":{"default":[1.0,3.0]},"_rankingScore":1.0,"_semanticScore":1.0}"###); + snapshot!(response["hits"][0], @r###"{"title":"Shazam!","desc":"a Captain Marvel ersatz","id":"1","_vectors":{"default":[1.0,3.0]},"_rankingScore":1.0}"###); snapshot!(response["semanticHitCount"], @"1"); } @@ -311,7 +313,7 @@ async fn query_combination() { .await; snapshot!(code, @"200 OK"); - snapshot!(response["hits"], @r###"[{"title":"Captain Marvel","desc":"a Shazam ersatz","id":"3","_vectors":{"default":[2.0,3.0]},"_rankingScore":0.7773500680923462,"_semanticScore":0.77735007},{"title":"Captain Planet","desc":"He's not part of the Marvel Cinematic Universe","id":"2","_vectors":{"default":[1.0,2.0]},"_rankingScore":0.7236068248748779,"_semanticScore":0.7236068},{"title":"Shazam!","desc":"a Captain Marvel ersatz","id":"1","_vectors":{"default":[1.0,3.0]},"_rankingScore":0.6581138968467712,"_semanticScore":0.6581139}]"###); + snapshot!(response["hits"], @r###"[{"title":"Captain Marvel","desc":"a Shazam ersatz","id":"3","_vectors":{"default":[2.0,3.0]},"_rankingScore":0.7773500680923462},{"title":"Captain Planet","desc":"He's not part of the Marvel Cinematic Universe","id":"2","_vectors":{"default":[1.0,2.0]},"_rankingScore":0.7236068248748779},{"title":"Shazam!","desc":"a Captain Marvel ersatz","id":"1","_vectors":{"default":[1.0,3.0]},"_rankingScore":0.6581138968467712}]"###); snapshot!(response["semanticHitCount"], @"3"); // full keyword, without a query diff --git a/meilisearch/tests/search/mod.rs b/meilisearch/tests/search/mod.rs index ce45b0d4f..b4350f686 100644 --- a/meilisearch/tests/search/mod.rs +++ b/meilisearch/tests/search/mod.rs @@ -1040,6 +1040,7 @@ async fn experimental_feature_vector_store() { let (response, code) = index .search_post(json!({ "vector": [1.0, 2.0, 3.0], + "showRankingScore": true })) .await; meili_snap::snapshot!(code, @"400 Bad Request"); @@ -1082,6 +1083,7 @@ async fn experimental_feature_vector_store() { let (response, code) = index .search_post(json!({ "vector": [1.0, 2.0, 3.0], + "showRankingScore": true, })) .await; @@ -1099,7 +1101,7 @@ async fn experimental_feature_vector_store() { 3 ] }, - "_semanticScore": 1.0 + "_rankingScore": 1.0 }, { "title": "Captain Marvel", @@ -1111,7 +1113,7 @@ async fn experimental_feature_vector_store() { 54 ] }, - "_semanticScore": 0.9129112 + "_rankingScore": 0.9129111766815186 }, { "title": "Gläss", @@ -1123,7 +1125,7 @@ async fn experimental_feature_vector_store() { 90 ] }, - "_semanticScore": 0.8106413 + "_rankingScore": 0.8106412887573242 }, { "title": "How to Train Your Dragon: The Hidden World", @@ -1135,7 +1137,7 @@ async fn experimental_feature_vector_store() { 32 ] }, - "_semanticScore": 0.74120104 + "_rankingScore": 0.7412010431289673 }, { "title": "Escape Room", @@ -1146,7 +1148,8 @@ async fn experimental_feature_vector_store() { -23, 32 ] - } + }, + "_rankingScore": 0.6972063183784485 } ] "###); From ca499a03025a2104de90a1a554faa578c6e472b9 Mon Sep 17 00:00:00 2001 From: Louis Dureuil Date: Wed, 3 Apr 2024 15:04:20 +0200 Subject: [PATCH 84/86] Fix test after rebase --- meilisearch/tests/search/hybrid.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/meilisearch/tests/search/hybrid.rs b/meilisearch/tests/search/hybrid.rs index 637242bc7..68ae4c0aa 100644 --- a/meilisearch/tests/search/hybrid.rs +++ b/meilisearch/tests/search/hybrid.rs @@ -106,7 +106,7 @@ async fn distribution_shift() { let search = json!({"q": "Captain", "vector": [1.0, 1.0], "showRankingScore": true, "hybrid": {"semanticRatio": 1.0}}); let (response, code) = index.search_post(search.clone()).await; snapshot!(code, @"200 OK"); - snapshot!(response["hits"], @r###"[{"title":"Captain Marvel","desc":"a Shazam ersatz","id":"3","_vectors":{"default":[2.0,3.0]},"_rankingScore":0.990290343761444,"_semanticScore":0.99029034},{"title":"Captain Planet","desc":"He's not part of the Marvel Cinematic Universe","id":"2","_vectors":{"default":[1.0,2.0]},"_rankingScore":0.974341630935669,"_semanticScore":0.97434163},{"title":"Shazam!","desc":"a Captain Marvel ersatz","id":"1","_vectors":{"default":[1.0,3.0]},"_rankingScore":0.9472135901451112,"_semanticScore":0.9472136}]"###); + snapshot!(response["hits"], @r###"[{"title":"Captain Marvel","desc":"a Shazam ersatz","id":"3","_vectors":{"default":[2.0,3.0]},"_rankingScore":0.990290343761444},{"title":"Captain Planet","desc":"He's not part of the Marvel Cinematic Universe","id":"2","_vectors":{"default":[1.0,2.0]},"_rankingScore":0.974341630935669},{"title":"Shazam!","desc":"a Captain Marvel ersatz","id":"1","_vectors":{"default":[1.0,3.0]},"_rankingScore":0.9472135901451112}]"###); let (response, code) = index .update_settings(json!({ @@ -127,7 +127,7 @@ async fn distribution_shift() { let (response, code) = index.search_post(search).await; snapshot!(code, @"200 OK"); - snapshot!(response["hits"], @r###"[{"title":"Captain Marvel","desc":"a Shazam ersatz","id":"3","_vectors":{"default":[2.0,3.0]},"_rankingScore":0.19161224365234375,"_semanticScore":0.19161224},{"title":"Captain Planet","desc":"He's not part of the Marvel Cinematic Universe","id":"2","_vectors":{"default":[1.0,2.0]},"_rankingScore":1.1920928955078125e-7,"_semanticScore":1.1920929e-7},{"title":"Shazam!","desc":"a Captain Marvel ersatz","id":"1","_vectors":{"default":[1.0,3.0]},"_rankingScore":1.1920928955078125e-7,"_semanticScore":1.1920929e-7}]"###); + snapshot!(response["hits"], @r###"[{"title":"Captain Marvel","desc":"a Shazam ersatz","id":"3","_vectors":{"default":[2.0,3.0]},"_rankingScore":0.19161224365234375},{"title":"Captain Planet","desc":"He's not part of the Marvel Cinematic Universe","id":"2","_vectors":{"default":[1.0,2.0]},"_rankingScore":1.1920928955078125e-7},{"title":"Shazam!","desc":"a Captain Marvel ersatz","id":"1","_vectors":{"default":[1.0,3.0]},"_rankingScore":1.1920928955078125e-7}]"###); } #[actix_rt::test] From a9013ed68316b68b2a4a9e069c16c070aaa161e1 Mon Sep 17 00:00:00 2001 From: Louis Dureuil Date: Thu, 4 Apr 2024 17:21:47 +0200 Subject: [PATCH 85/86] Fix comment mistake Co-authored-by: Tamo --- meilisearch/src/routes/indexes/search.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/meilisearch/src/routes/indexes/search.rs b/meilisearch/src/routes/indexes/search.rs index 0f7d3b1ee..5581e6a68 100644 --- a/meilisearch/src/routes/indexes/search.rs +++ b/meilisearch/src/routes/indexes/search.rs @@ -277,7 +277,7 @@ pub fn search_kind( features.check_vector("Passing `hybrid` as a query parameter")?; } - // regardless of anything, always do a semantic search when we don't have a vector and the query is whitespace or missing + // regardless of anything, always do a keyword search when we don't have a vector and the query is whitespace or missing if query.vector.is_none() { match &query.q { Some(q) if q.trim().is_empty() => return Ok(SearchKind::KeywordOnly), From d6b6cd322c26768180758ac9936b9eca9e4db987 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Cl=C3=A9mentine=20U=2E=20-=20curqui?= Date: Fri, 5 Apr 2024 18:40:28 +0200 Subject: [PATCH 86/86] Update sprint_issue.md (#4556) --- .github/ISSUE_TEMPLATE/sprint_issue.md | 1 - 1 file changed, 1 deletion(-) diff --git a/.github/ISSUE_TEMPLATE/sprint_issue.md b/.github/ISSUE_TEMPLATE/sprint_issue.md index 0a3eb6843..261e11798 100644 --- a/.github/ISSUE_TEMPLATE/sprint_issue.md +++ b/.github/ISSUE_TEMPLATE/sprint_issue.md @@ -9,7 +9,6 @@ assignees: '' Related product team resources: [PRD]() (_internal only_) Related product discussion: -Related spec: WIP ## Motivation