use crate::pgvector::{FieldConfig, PgVector, PgVectorBuilder};
use anyhow::{Result, anyhow};
use async_trait::async_trait;
use pgvector::Vector;
use sqlx::{Column, Postgres, QueryBuilder, Row, prelude::FromRow, types::Uuid};
use swiftide_core::{
Retrieve,
document::Document,
indexing::Metadata,
querying::{
Query,
search_strategies::{CustomStrategy, SimilaritySingleEmbedding},
states,
},
};
#[allow(dead_code)]
#[derive(Debug, Clone)]
struct VectorSearchResult {
id: Uuid,
chunk: String,
metadata: Metadata,
}
impl From<VectorSearchResult> for Document {
fn from(val: VectorSearchResult) -> Self {
Document::new(val.chunk, Some(val.metadata))
}
}
impl FromRow<'_, sqlx::postgres::PgRow> for VectorSearchResult {
fn from_row(row: &sqlx::postgres::PgRow) -> Result<Self, sqlx::Error> {
let mut metadata = Metadata::default();
for column in row.columns() {
if column.name().starts_with("meta_") {
row.try_get::<serde_json::Value, _>(column.name())?
.as_object()
.and_then(|object| {
object.keys().collect::<Vec<_>>().first().map(|key| {
metadata.insert(
key.to_owned(),
object.get(key.as_str()).expect("infallible").clone(),
);
})
});
}
}
Ok(VectorSearchResult {
id: row.try_get("id")?,
chunk: row.try_get("chunk")?,
metadata,
})
}
}
#[allow(clippy::redundant_closure_for_method_calls)]
#[async_trait]
impl Retrieve<SimilaritySingleEmbedding<String>> for PgVector {
#[tracing::instrument]
async fn retrieve(
&self,
search_strategy: &SimilaritySingleEmbedding<String>,
query_state: Query<states::Pending>,
) -> Result<Query<states::Retrieved>> {
let embedding = if let Some(embedding) = query_state.embedding.as_ref() {
Vector::from(embedding.clone())
} else {
return Err(anyhow::Error::msg("Missing embedding in query state"));
};
let vector_column_name = self.get_vector_column_name()?;
let pool = self.pool_get_or_initialize().await?;
let default_columns: Vec<_> = PgVectorBuilder::default_fields()
.iter()
.map(|f| f.field_name().to_string())
.chain(
self.fields
.iter()
.filter(|f| matches!(f, FieldConfig::Metadata(_)))
.map(|f| f.field_name().to_string()),
)
.collect();
if !PgVector::is_valid_identifier(&self.table_name)
|| default_columns
.iter()
.any(|column| !PgVector::is_valid_identifier(column))
|| !PgVector::is_valid_identifier(&vector_column_name)
{
return Err(anyhow!("Invalid table or field name"));
}
let mut sql = QueryBuilder::<Postgres>::new("SELECT ");
sql.push(default_columns.join(", "))
.push(" FROM ")
.push(&self.table_name);
if let Some(filter) = search_strategy.filter() {
let (key, value) = filter
.split_once('=')
.ok_or_else(|| anyhow!("Invalid filter format"))?;
let key = key.trim();
let value = value.trim().trim_matches('"');
let metadata_column = format!("meta_{}", PgVector::normalize_field_name(key));
tracing::debug!(
"Filter being applied: key = {:#?}, value = {:#?}",
key,
value
);
sql.push(" WHERE ")
.push(metadata_column)
.push("->>")
.push_bind(key.to_owned())
.push(" = ")
.push_bind(value.to_owned());
}
let top_k = i32::try_from(search_strategy.top_k())
.map_err(|_| anyhow!("Failed to convert top_k to i32"))?;
sql.push(" ORDER BY ")
.push(vector_column_name)
.push(" <=> ")
.push_bind(embedding)
.push(" LIMIT ")
.push_bind(top_k);
tracing::debug!("Running retrieve with SQL: {}", sql.sql().as_str());
let data: Vec<VectorSearchResult> = sql.build_query_as().fetch_all(pool).await?;
let docs = data.into_iter().map(Into::into).collect();
Ok(query_state.retrieved_documents(docs))
}
}
#[async_trait]
impl Retrieve<SimilaritySingleEmbedding> for PgVector {
async fn retrieve(
&self,
search_strategy: &SimilaritySingleEmbedding,
query: Query<states::Pending>,
) -> Result<Query<states::Retrieved>> {
Retrieve::<SimilaritySingleEmbedding<String>>::retrieve(
self,
&search_strategy.into_concrete_filter::<String>(),
query,
)
.await
}
}
#[async_trait]
impl Retrieve<CustomStrategy<sqlx::QueryBuilder<sqlx::Postgres>>> for PgVector {
async fn retrieve(
&self,
search_strategy: &CustomStrategy<sqlx::QueryBuilder<sqlx::Postgres>>,
query: Query<states::Pending>,
) -> Result<Query<states::Retrieved>> {
let pool = self.get_pool().await?;
let mut query_builder = search_strategy.build_query(&query).await?;
let results = query_builder
.build_query_as::<VectorSearchResult>() .fetch_all(pool) .await
.map_err(|e| anyhow!("Failed to execute search query: {e}"))?;
let documents = results.into_iter().map(Into::into).collect();
Ok(query.retrieved_documents(documents))
}
}
#[cfg(test)]
mod tests {
use crate::pgvector::fixtures::TestContext;
use futures_util::TryStreamExt;
use std::collections::HashSet;
use swiftide_core::{Persist, indexing, indexing::EmbeddedField};
use swiftide_core::{
Retrieve,
querying::{Query, search_strategies::SimilaritySingleEmbedding, states},
};
#[test_log::test(tokio::test)]
async fn test_retrieve_multiple_docs_and_filter() {
let test_context = TestContext::setup_with_cfg(
vec!["filter"].into(),
HashSet::from([EmbeddedField::Combined]),
)
.await
.expect("Test setup failed");
let nodes = vec![
indexing::TextNode::new("test_query1").with_metadata(("filter", "true")),
indexing::TextNode::new("test_query2").with_metadata(("filter", "true")),
indexing::TextNode::new("test_query3").with_metadata(("filter", "false")),
]
.into_iter()
.map(|node| {
node.with_vectors([(EmbeddedField::Combined, vec![1.0; 384])]);
node.to_owned()
})
.collect();
test_context
.pgv_storage
.batch_store(nodes)
.await
.try_collect::<Vec<_>>()
.await
.unwrap();
let mut query = Query::<states::Pending>::new("test_query");
query.embedding = Some(vec![1.0; 384]);
let search_strategy = SimilaritySingleEmbedding::<()>::default();
let result = test_context
.pgv_storage
.retrieve(&search_strategy, query.clone())
.await
.unwrap();
assert_eq!(result.documents().len(), 3);
let search_strategy =
SimilaritySingleEmbedding::from_filter("filter = \"true\"".to_string());
let result = test_context
.pgv_storage
.retrieve(&search_strategy, query.clone())
.await
.unwrap();
assert_eq!(result.documents().len(), 2);
let search_strategy =
SimilaritySingleEmbedding::from_filter("filter = \"banana\"".to_string());
let result = test_context
.pgv_storage
.retrieve(&search_strategy, query.clone())
.await
.unwrap();
assert_eq!(result.documents().len(), 0);
}
#[test_log::test(tokio::test)]
async fn test_retrieve_docs_with_metadata() {
let test_context = TestContext::setup_with_cfg(
vec!["other", "text"].into(),
HashSet::from([EmbeddedField::Combined]),
)
.await
.expect("Test setup failed");
let nodes = vec![
indexing::TextNode::new("test_query1")
.with_metadata([
("other", serde_json::Value::from(10)),
("text", serde_json::Value::from("some text")),
])
.with_vectors([(EmbeddedField::Combined, vec![1.0; 384])])
.to_owned(),
];
test_context
.pgv_storage
.batch_store(nodes)
.await
.try_collect::<Vec<_>>()
.await
.unwrap();
let mut query = Query::<states::Pending>::new("test_query");
query.embedding = Some(vec![1.0; 384]);
let search_strategy = SimilaritySingleEmbedding::<()>::default();
let result = test_context
.pgv_storage
.retrieve(&search_strategy, query.clone())
.await
.unwrap();
assert_eq!(result.documents().len(), 1);
let doc = result.documents().first().unwrap();
assert_eq!(
doc.metadata().get("other"),
Some(&serde_json::Value::from(10))
);
assert_eq!(
doc.metadata().get("text"),
Some(&serde_json::Value::from("some text"))
);
}
}