quickwit
This commit is contained in:
@@ -2,11 +2,13 @@ use std::collections::HashSet;
|
|||||||
|
|
||||||
use anyhow::{anyhow, Result};
|
use anyhow::{anyhow, Result};
|
||||||
use reqwest::Client;
|
use reqwest::Client;
|
||||||
use serde::Deserialize;
|
use serde::{Deserialize, Serialize};
|
||||||
use serde_json::{json, Value};
|
use serde_json::{json, Value};
|
||||||
use tracing::{debug, error};
|
use tracing::{debug, error};
|
||||||
use uuid::Uuid;
|
use uuid::Uuid;
|
||||||
|
|
||||||
|
use crate::models::{Document, DocumentVersion};
|
||||||
|
|
||||||
pub const QUICKWIT_MAX_HITS: usize = 200;
|
pub const QUICKWIT_MAX_HITS: usize = 200;
|
||||||
|
|
||||||
pub fn build_quickwit_query(input: &str) -> Option<String> {
|
pub fn build_quickwit_query(input: &str) -> Option<String> {
|
||||||
@@ -164,3 +166,69 @@ struct QuickwitSearchResponse {
|
|||||||
#[serde(default)]
|
#[serde(default)]
|
||||||
hits: Vec<Value>,
|
hits: Vec<Value>,
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[derive(Serialize)]
|
||||||
|
pub struct QuickwitIngestRecord {
|
||||||
|
pub document_id: Uuid,
|
||||||
|
pub version_id: Uuid,
|
||||||
|
pub tenant_id: Uuid,
|
||||||
|
pub title: String,
|
||||||
|
pub text: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn build_quickwit_ingest_record(
|
||||||
|
document: &Document,
|
||||||
|
version: &DocumentVersion,
|
||||||
|
tenant_id: Uuid,
|
||||||
|
text: &str,
|
||||||
|
) -> QuickwitIngestRecord {
|
||||||
|
QuickwitIngestRecord {
|
||||||
|
document_id: document.id,
|
||||||
|
version_id: version.id,
|
||||||
|
tenant_id,
|
||||||
|
title: document.title.to_lowercase(),
|
||||||
|
text: text.to_lowercase(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn quickwit_ingest(
|
||||||
|
client: &Client,
|
||||||
|
endpoint: &str,
|
||||||
|
index: &str,
|
||||||
|
records: &[QuickwitIngestRecord],
|
||||||
|
) -> Result<()> {
|
||||||
|
if records.is_empty() {
|
||||||
|
return Ok(());
|
||||||
|
}
|
||||||
|
|
||||||
|
let url = format!(
|
||||||
|
"{}/api/v1/{}/ingest?commit=auto",
|
||||||
|
endpoint.trim_end_matches('/'),
|
||||||
|
index
|
||||||
|
);
|
||||||
|
|
||||||
|
let mut body = String::new();
|
||||||
|
for record in records {
|
||||||
|
let line = serde_json::to_string(record)?;
|
||||||
|
body.push_str(&line);
|
||||||
|
body.push('\n');
|
||||||
|
}
|
||||||
|
|
||||||
|
debug!(%url, lines = records.len(), "sending quickwit ingest request");
|
||||||
|
let response = client
|
||||||
|
.post(url)
|
||||||
|
.header("content-type", "application/x-ndjson")
|
||||||
|
.body(body)
|
||||||
|
.send()
|
||||||
|
.await?;
|
||||||
|
|
||||||
|
if !response.status().is_success() {
|
||||||
|
let status = response.status();
|
||||||
|
let body = response.text().await.unwrap_or_default();
|
||||||
|
error!(%status, %body, "quickwit ingest request failed");
|
||||||
|
return Err(anyhow!("quickwit ingest failed with status {status}: {body}"));
|
||||||
|
}
|
||||||
|
|
||||||
|
debug!("quickwit ingest request succeeded");
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|||||||
@@ -5,12 +5,12 @@ use async_trait::async_trait;
|
|||||||
use diesel::prelude::*;
|
use diesel::prelude::*;
|
||||||
use reqwest::Client;
|
use reqwest::Client;
|
||||||
use serde::Deserialize;
|
use serde::Deserialize;
|
||||||
use serde_json::json;
|
|
||||||
use tokio::task;
|
use tokio::task;
|
||||||
use tracing::{error, warn};
|
use tracing::{error, warn};
|
||||||
use uuid::Uuid;
|
use uuid::Uuid;
|
||||||
|
|
||||||
use crate::{
|
use crate::{
|
||||||
|
documents::search::{build_quickwit_ingest_record, quickwit_ingest},
|
||||||
jobs::JOB_INDEX_DOCUMENT_TEXT,
|
jobs::JOB_INDEX_DOCUMENT_TEXT,
|
||||||
models::{Document, DocumentVersion},
|
models::{Document, DocumentVersion},
|
||||||
schema::{document_asset_objects, document_assets, document_versions, documents},
|
schema::{document_asset_objects, document_assets, document_versions, documents},
|
||||||
@@ -129,7 +129,7 @@ impl JobHandler for IndexDocumentTextJob {
|
|||||||
.await
|
.await
|
||||||
{
|
{
|
||||||
Ok(bytes) => bytes,
|
Ok(bytes) => bytes,
|
||||||
Err(err) => return handle_fetch_error(job, err, "failed to download ocr text"),
|
Err(err) => return handle_fetch_error(&job, err, "failed to download ocr text"),
|
||||||
};
|
};
|
||||||
let text = match String::from_utf8(bytes) {
|
let text = match String::from_utf8(bytes) {
|
||||||
Ok(text) => text,
|
Ok(text) => text,
|
||||||
@@ -148,43 +148,17 @@ impl JobHandler for IndexDocumentTextJob {
|
|||||||
};
|
};
|
||||||
}
|
}
|
||||||
|
|
||||||
let client = client;
|
let record = build_quickwit_ingest_record(
|
||||||
let url = format!(
|
&context.document,
|
||||||
"{}/api/v1/{}/ingest?commit=auto",
|
&context.version,
|
||||||
quickwit_endpoint, quickwit_index
|
job.tenant_id,
|
||||||
|
&text,
|
||||||
);
|
);
|
||||||
let payload = json!({
|
|
||||||
"document_id": context.document.id,
|
|
||||||
"version_id": context.version.id,
|
|
||||||
"tenant_id": job.tenant_id,
|
|
||||||
"title": context.document.title.to_lowercase(),
|
|
||||||
"text": text.to_lowercase()
|
|
||||||
});
|
|
||||||
|
|
||||||
let body = serde_json::to_string(&payload).unwrap();
|
match quickwit_ingest(&client, &quickwit_endpoint, &quickwit_index, &[record]).await {
|
||||||
|
Ok(()) => JobExecution::Success,
|
||||||
match client
|
|
||||||
.post(&url)
|
|
||||||
.header("content-type", "application/x-ndjson")
|
|
||||||
.body(format!("{}\n", body))
|
|
||||||
.send()
|
|
||||||
.await
|
|
||||||
{
|
|
||||||
Ok(response) => {
|
|
||||||
if response.status().is_success() {
|
|
||||||
JobExecution::Success
|
|
||||||
} else {
|
|
||||||
let status = response.status();
|
|
||||||
let body = response.text().await.unwrap_or_default();
|
|
||||||
warn!(job_id = %job.id, %status, %body, "quickwit ingest failed");
|
|
||||||
JobExecution::Retry {
|
|
||||||
delay: Duration::from_secs(30),
|
|
||||||
error: format!("quickwit ingest failed with status {status}"),
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
Err(err) => {
|
Err(err) => {
|
||||||
warn!(job_id = %job.id, error = %err, "quickwit request failed");
|
warn!(job_id = %job.id, error = %err, "quickwit ingest failed");
|
||||||
JobExecution::Retry {
|
JobExecution::Retry {
|
||||||
delay: Duration::from_secs(30),
|
delay: Duration::from_secs(30),
|
||||||
error: err.to_string(),
|
error: err.to_string(),
|
||||||
|
|||||||
@@ -186,14 +186,14 @@ pub(crate) async fn fetch_version_object(
|
|||||||
}
|
}
|
||||||
|
|
||||||
pub(crate) fn handle_fetch_error(
|
pub(crate) fn handle_fetch_error(
|
||||||
job_id: crate::models::Job,
|
job: &crate::models::Job,
|
||||||
err: FetchVersionError,
|
err: FetchVersionError,
|
||||||
message: &str,
|
message: &str,
|
||||||
) -> JobExecution {
|
) -> JobExecution {
|
||||||
match err {
|
match err {
|
||||||
FetchVersionError::TooLarge { size, limit } => {
|
FetchVersionError::TooLarge { size, limit } => {
|
||||||
warn!(
|
warn!(
|
||||||
job_id = %job_id.id,
|
job_id = %job.id,
|
||||||
size_bytes = size,
|
size_bytes = size,
|
||||||
limit_bytes = limit,
|
limit_bytes = limit,
|
||||||
"document exceeds worker size limit"
|
"document exceeds worker size limit"
|
||||||
@@ -205,7 +205,12 @@ pub(crate) fn handle_fetch_error(
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
FetchVersionError::Storage(err) => {
|
FetchVersionError::Storage(err) => {
|
||||||
warn!(job_id = %job_id.id, error = %err, "{message}");
|
warn!(
|
||||||
|
job_id = %job.id,
|
||||||
|
error = %err,
|
||||||
|
context = message,
|
||||||
|
"failed to fetch object for worker"
|
||||||
|
);
|
||||||
JobExecution::Retry {
|
JobExecution::Retry {
|
||||||
delay: Duration::from_secs(30),
|
delay: Duration::from_secs(30),
|
||||||
error: err.to_string(),
|
error: err.to_string(),
|
||||||
|
|||||||
@@ -108,7 +108,7 @@ impl JobHandler for GenerateOcrTextJob {
|
|||||||
.await
|
.await
|
||||||
{
|
{
|
||||||
Ok(bytes) => bytes,
|
Ok(bytes) => bytes,
|
||||||
Err(err) => return handle_fetch_error(job, err, "failed to fetch document for ocr"),
|
Err(err) => return handle_fetch_error(&job, err, "failed to fetch document for ocr"),
|
||||||
};
|
};
|
||||||
|
|
||||||
let doc_meta = PdfDocumentMeta {
|
let doc_meta = PdfDocumentMeta {
|
||||||
|
|||||||
@@ -111,7 +111,7 @@ impl JobHandler for GenerateThumbnailsJob {
|
|||||||
.await
|
.await
|
||||||
{
|
{
|
||||||
Ok(bytes) => bytes,
|
Ok(bytes) => bytes,
|
||||||
Err(err) => return handle_fetch_error(job, err, "thumbnail fetch failed; will retry"),
|
Err(err) => return handle_fetch_error(&job, err, "thumbnail fetch failed; will retry"),
|
||||||
};
|
};
|
||||||
|
|
||||||
let generation = match generate_preview_and_thumbnail(&initial.document, &bytes) {
|
let generation = match generate_preview_and_thumbnail(&initial.document, &bytes) {
|
||||||
|
|||||||
Reference in New Issue
Block a user