140 lines
4.1 KiB
Rust
140 lines
4.1 KiB
Rust
use crate::collab_indexer::Indexer;
|
|
use crate::vector::embedder::Embedder;
|
|
use crate::vector::open_ai::split_text_by_max_content_len;
|
|
use anyhow::anyhow;
|
|
use app_error::AppError;
|
|
use appflowy_ai_client::dto::{
|
|
EmbeddingEncodingFormat, EmbeddingInput, EmbeddingModel, EmbeddingOutput, EmbeddingRequest,
|
|
};
|
|
use async_trait::async_trait;
|
|
use collab::preclude::Collab;
|
|
use collab_document::document::DocumentBody;
|
|
use collab_document::error::DocumentError;
|
|
use collab_entity::CollabType;
|
|
use database_entity::dto::{AFCollabEmbeddedChunk, AFCollabEmbeddings, EmbeddingContentType};
|
|
use serde_json::json;
|
|
use tracing::trace;
|
|
use uuid::Uuid;
|
|
|
|
pub struct DocumentIndexer;
|
|
|
|
#[async_trait]
|
|
impl Indexer for DocumentIndexer {
|
|
fn create_embedded_chunks_from_collab(
|
|
&self,
|
|
collab: &Collab,
|
|
embedding_model: EmbeddingModel,
|
|
) -> Result<Vec<AFCollabEmbeddedChunk>, AppError> {
|
|
let object_id = collab.object_id().to_string();
|
|
let document = DocumentBody::from_collab(collab).ok_or_else(|| {
|
|
anyhow!(
|
|
"Failed to get document body from collab `{}`: schema is missing required fields",
|
|
object_id
|
|
)
|
|
})?;
|
|
|
|
let result = document.to_plain_text(collab.transact(), false, true);
|
|
match result {
|
|
Ok(content) => self.create_embedded_chunks_from_text(object_id, content, embedding_model),
|
|
Err(err) => {
|
|
if matches!(err, DocumentError::NoRequiredData) {
|
|
Ok(vec![])
|
|
} else {
|
|
Err(AppError::Internal(err.into()))
|
|
}
|
|
},
|
|
}
|
|
}
|
|
|
|
fn create_embedded_chunks_from_text(
|
|
&self,
|
|
object_id: String,
|
|
text: String,
|
|
model: EmbeddingModel,
|
|
) -> Result<Vec<AFCollabEmbeddedChunk>, AppError> {
|
|
split_text_into_chunks(object_id, text, CollabType::Document, &model)
|
|
}
|
|
|
|
fn embed(
|
|
&self,
|
|
embedder: &Embedder,
|
|
mut content: Vec<AFCollabEmbeddedChunk>,
|
|
) -> Result<Option<AFCollabEmbeddings>, AppError> {
|
|
if content.is_empty() {
|
|
return Ok(None);
|
|
}
|
|
|
|
let contents: Vec<_> = content
|
|
.iter()
|
|
.map(|fragment| fragment.content.clone())
|
|
.collect();
|
|
let resp = embedder.embed(EmbeddingRequest {
|
|
input: EmbeddingInput::StringArray(contents),
|
|
model: embedder.model().name().to_string(),
|
|
encoding_format: EmbeddingEncodingFormat::Float,
|
|
dimensions: EmbeddingModel::TextEmbedding3Small.default_dimensions(),
|
|
})?;
|
|
|
|
trace!(
|
|
"[Embedding] request {} embeddings, received {} embeddings",
|
|
content.len(),
|
|
resp.data.len()
|
|
);
|
|
|
|
for embedding in resp.data {
|
|
let param = &mut content[embedding.index as usize];
|
|
let embedding: Vec<f32> = match embedding.embedding {
|
|
EmbeddingOutput::Float(embedding) => embedding.into_iter().map(|f| f as f32).collect(),
|
|
EmbeddingOutput::Base64(_) => {
|
|
return Err(AppError::OpenError(
|
|
"Unexpected base64 encoding".to_string(),
|
|
))
|
|
},
|
|
};
|
|
param.embedding = Some(embedding);
|
|
}
|
|
|
|
Ok(Some(AFCollabEmbeddings {
|
|
tokens_consumed: resp.usage.total_tokens as u32,
|
|
params: content,
|
|
}))
|
|
}
|
|
}
|
|
|
|
fn split_text_into_chunks(
|
|
object_id: String,
|
|
content: String,
|
|
collab_type: CollabType,
|
|
embedding_model: &EmbeddingModel,
|
|
) -> Result<Vec<AFCollabEmbeddedChunk>, AppError> {
|
|
debug_assert!(matches!(
|
|
embedding_model,
|
|
EmbeddingModel::TextEmbedding3Small
|
|
));
|
|
|
|
if content.is_empty() {
|
|
return Ok(vec![]);
|
|
}
|
|
// We assume that every token is ~4 bytes. We're going to split document content into fragments
|
|
// of ~2000 tokens each.
|
|
let split_contents = split_text_by_max_content_len(content, 8000)?;
|
|
let metadata =
|
|
json!({"id": object_id, "source": "appflowy", "name": "document", "collab_type": collab_type });
|
|
Ok(
|
|
split_contents
|
|
.into_iter()
|
|
.enumerate()
|
|
.map(|(index, content)| AFCollabEmbeddedChunk {
|
|
fragment_id: Uuid::new_v4().to_string(),
|
|
object_id: object_id.clone(),
|
|
content_type: EmbeddingContentType::PlainText,
|
|
content,
|
|
embedding: None,
|
|
metadata: metadata.clone(),
|
|
fragment_index: index as i32,
|
|
embedded_type: 0,
|
|
})
|
|
.collect(),
|
|
)
|
|
}
|