maritime-logistics/src/retrieval/vector_store.rs

//! Thin wrapper around the Qdrant gRPC client.
//!
//! Constructed explicitly (see `cli.rs`) and handed to a `DocumentRetriever`
//! by reference -- there is no global client instance anywhere in this
//! crate.

use std::sync::Arc;

use anyhow::Result;
use qdrant_client::qdrant::{
    CreateCollectionBuilder, Distance, PointStruct, QueryPointsBuilder, UpsertPointsBuilder,
    VectorParamsBuilder,
};
use qdrant_client::{Payload, Qdrant};
use serde::{Deserialize, Serialize};
use uuid::Uuid;

use crate::embeddings::EmbeddingProvider;
use crate::knowledge::loader::DocumentRecord;
use crate::models::agent_io::RetrievedPassage;
use crate::retrieval::chunking::chunk_text;

const CHUNK_SIZE: usize = 700;
const CHUNK_OVERLAP: usize = 100;

#[derive(Debug, Serialize, Deserialize)]
struct ChunkPayload {
    source: String,
    text: String,
}

pub struct VectorStore {
    client: Qdrant,
    collection_name: String,
    embeddings: Arc<dyn EmbeddingProvider>,
}

impl VectorStore {
    pub fn new(
        url: &str,
        collection_name: impl Into<String>,
        embeddings: Arc<dyn EmbeddingProvider>,
    ) -> Result<Self> {
        let client = Qdrant::from_url(url).build()?;
        Ok(Self {
            client,
            collection_name: collection_name.into(),
            embeddings,
        })
    }

    /// Creates the collection if it doesn't already exist. Idempotent.
    pub async fn ensure_collection(&self) -> Result<()> {
        if self.client.collection_exists(&self.collection_name).await? {
            return Ok(());
        }
        self.client
            .create_collection(CreateCollectionBuilder::new(&self.collection_name).vectors_config(
                VectorParamsBuilder::new(self.embeddings.dimension() as u64, Distance::Cosine),
            ))
            .await?;
        Ok(())
    }

    /// Chunks and embeds every document, then upserts the chunks as points.
    /// Returns the number of chunks indexed.
    pub async fn index_documents(&self, documents: &[DocumentRecord]) -> Result<usize> {
        let mut chunk_texts = Vec::new();
        let mut sources = Vec::new();
        for doc in documents {
            for chunk in chunk_text(&doc.text, CHUNK_SIZE, CHUNK_OVERLAP) {
                sources.push(doc.source.clone());
                chunk_texts.push(chunk);
            }
        }
        if chunk_texts.is_empty() {
            return Ok(0);
        }

        let vectors = self.embeddings.embed(chunk_texts.clone()).await?;
        let points: Vec<PointStruct> = vectors
            .into_iter()
            .zip(chunk_texts)
            .zip(sources)
            .map(|((vector, text), source)| {
                let payload = ChunkPayload { source, text };
                let payload_value = serde_json::to_value(payload).expect("payload serializes");
                PointStruct::new(
                    Uuid::new_v4().to_string(),
                    vector,
                    Payload::try_from(payload_value).expect("payload is a JSON object"),
                )
            })
            .collect();

        let count = points.len();
        self.client
            .upsert_points(UpsertPointsBuilder::new(&self.collection_name, points))
            .await?;
        Ok(count)
    }

    /// Embeds `query` and returns the `limit` nearest chunks.
    pub async fn search(&self, query: &str, limit: u64) -> Result<Vec<RetrievedPassage>> {
        let vector = self.embeddings.embed_one(query).await?;
        let response = self
            .client
            .query(
                QueryPointsBuilder::new(&self.collection_name)
                    .query(vector)
                    .limit(limit)
                    .with_payload(true),
            )
            .await?;

        let mut passages = Vec::with_capacity(response.result.len());
        for point in response.result {
            let payload: ChunkPayload = Payload::from(point.payload).deserialize()?;
            passages.push(RetrievedPassage {
                source: payload.source,
                text: payload.text,
                score: point.score,
            });
        }
        Ok(passages)
    }
}