microagents-cli 0.1.0

Batteries-included terminal-based agent built on top of the microagents framework
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
92
93
94
95
96

use qdrant_edge::{
    AnyVariants, Condition, FieldCondition, Filter, Fusion, JsonPath, Match, NamedQuery, Payload,
    Prefetch, QueryEnum, QueryRequest, ScoringQuery, SparseVector, WithPayloadInterface,
    WithVector, external::ordered_float::OrderedFloat,
};
use std::str::FromStr;

use crate::init_env::{EmbeddingPayload, SPARSE_VECTORS_NAME, VECTORS_NAME, load_qdrant_edge};

#[derive(Debug, Clone)]
pub struct ResultWithScore {
    pub content: String,
    pub document_path: String,
    pub score: f32,
}

/// Convert Qdrant Payload back to DocMeta
fn payload_to_struct(payload: &Payload) -> Result<EmbeddingPayload, Box<dyn std::error::Error>> {
    let ep: EmbeddingPayload = serde_json::from_value(serde_json::to_value(payload)?)?;
    Ok(ep)
}

pub fn search(
    query_embedding: Vec<f32>,
    sparse_query_embedding: SparseVector,
    document_paths: Option<Vec<String>>,
    limit: Option<usize>,
    score_threshold: Option<f32>,
) -> Result<Vec<ResultWithScore>, Box<dyn std::error::Error>> {
    let edge_shard = load_qdrant_edge()?;
    let mut all_results: Vec<ResultWithScore> = vec![];
    let threshold: Option<OrderedFloat<f32>> = score_threshold.map(OrderedFloat);
    let top_k = limit.unwrap_or(10);
    let stmt_filter = match document_paths {
        Some(d) => Some(Filter::new_must(Condition::Field(
            FieldCondition::new_match(
                JsonPath::from_str("document_path").map_err(|_| "invalid json path")?,
                Match::from(AnyVariants::Strings(d.iter().cloned().collect())),
            ),
        ))),
        None => None,
    };
    let shard_query = QueryRequest {
        prefetches: vec![
            Prefetch {
                prefetches: vec![],
                query: Some(ScoringQuery::Vector(QueryEnum::Nearest(NamedQuery {
                    query: sparse_query_embedding.into(),
                    using: Some(SPARSE_VECTORS_NAME.to_string()),
                }))),
                limit: top_k,
                params: None,
                filter: None,
                score_threshold: threshold,
            },
            Prefetch {
                prefetches: vec![],
                query: Some(ScoringQuery::Vector(QueryEnum::Nearest(NamedQuery {
                    query: query_embedding.into(),
                    using: Some(VECTORS_NAME.to_string()),
                }))),
                limit: top_k,
                params: None,
                filter: None,
                score_threshold: None,
            },
        ],
        query: Some(ScoringQuery::Fusion(Fusion::Rrf {
            k: 2,
            weights: Some(vec![OrderedFloat(0.75), OrderedFloat(0.25)]),
        })),
        filter: stmt_filter,
        score_threshold: None,
        limit: top_k,
        offset: 0,
        params: None,
        with_vector: WithVector::Bool(false),
        with_payload: WithPayloadInterface::Bool(true),
    };
    let results = edge_shard.query(shard_query)?;
    for r in results {
        let payload = match r.payload {
            Some(p) => p,
            None => return Err("Found a None payload when searching".into()),
        };
        let embd_payload = payload_to_struct(&payload)?;
        let scored_result = ResultWithScore {
            content: embd_payload.content,
            document_path: embd_payload.document_path,
            score: r.score,
        };
        all_results.push(scored_result);
    }

    Ok(all_results)
}