daimon 0.16.0

A Rust-native AI agent framework
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91

//! Qdrant vector store retriever.

use std::sync::Arc;

use crate::error::{DaimonError, Result};
use crate::model::ErasedEmbeddingModel;
use crate::retriever::types::Document;
use crate::retriever::traits::Retriever;

/// Retriever backed by a Qdrant vector database.
pub struct QdrantRetriever {
    client: qdrant_client::Qdrant,
    collection: String,
    embedding_model: Arc<dyn ErasedEmbeddingModel>,
    content_field: String,
}

impl QdrantRetriever {
    /// Creates a new Qdrant retriever.
    ///
    /// - `url`: Qdrant server URL (e.g. `http://localhost:6334`)
    /// - `collection`: Name of the Qdrant collection
    /// - `embedding_model`: Model used to embed queries
    pub async fn new(
        url: impl Into<String>,
        collection: impl Into<String>,
        embedding_model: Arc<dyn ErasedEmbeddingModel>,
    ) -> Result<Self> {
        let client = qdrant_client::Qdrant::from_url(&url.into())
            .build()
            .map_err(|e| DaimonError::Other(format!("Qdrant connection error: {e}")))?;

        Ok(Self {
            client,
            collection: collection.into(),
            embedding_model,
            content_field: "content".to_string(),
        })
    }

    /// Sets the payload field name that contains document text. Default: `"content"`.
    pub fn with_content_field(mut self, field: impl Into<String>) -> Self {
        self.content_field = field.into();
        self
    }
}

impl Retriever for QdrantRetriever {
    async fn retrieve(&self, query: &str, top_k: usize) -> Result<Vec<Document>> {
        let texts = [query];
        let embeddings = self.embedding_model.embed_erased(&texts).await?;
        let query_vec = embeddings.into_iter().next().unwrap_or_default();

        let results = self
            .client
            .search_points(
                qdrant_client::qdrant::SearchPointsBuilder::new(
                    &self.collection,
                    query_vec,
                    top_k as u64,
                )
                .with_payload(true),
            )
            .await
            .map_err(|e| DaimonError::Other(format!("Qdrant search error: {e}")))?;

        let mut docs = Vec::with_capacity(results.result.len());
        for point in results.result {
            let content = point
                .payload
                .get(&self.content_field)
                .and_then(|v| v.as_str())
                .map(|s| s.to_string())
                .unwrap_or_default();

            let mut doc = Document::new(content).with_score(point.score as f64);

            for (key, val) in &point.payload {
                if key != &self.content_field {
                    if let Some(s) = val.as_str() {
                        doc = doc.with_metadata(key, serde_json::json!(s));
                    }
                }
            }

            docs.push(doc);
        }

        Ok(docs)
    }
}