ruchat 0.1.2

ollama/chroma command-line AI chat tool
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
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111

use crate::chroma::create_client;
use crate::error::RuChatError;
use anyhow::Result;
use chromadb::collection::{ChromaCollection, GetOptions, GetResult, QueryOptions, QueryResult};
use clap::Parser;
use serde_json::json;

/// Chroma database similarity search command line arguments.
///
/// This struct defines the arguments required to perform a similarity
/// search in a Chroma database, including query parameters and database
/// connection information.
#[derive(Parser, Debug, Clone, PartialEq)]
pub struct SimilarityArgs {
    /// Query string to search for similar embeddings.
    #[clap(short, long)]
    pub(crate) query: String,

    /// Number of embeddings to return.
    #[clap(short, long, default_value = "1")]
    pub(crate) count: usize,

    /// Number of similar embeddings to return.
    #[clap(short, long, default_value = "5")]
    pub(crate) similarity_count: usize,

    /// Chroma database collection name.
    #[clap(short, long, default_value = "default")]
    pub(crate) collection: String,

    /// Chroma database metadata, comma separated key:value pairs.
    #[clap(short, long)]
    pub(crate) metadata: Option<String>,

    /// Chroma database server address and port.
    #[clap(short = 'C', long, default_value = "http://localhost:8000")]
    pub(crate) chroma_server: String,

    /// Chroma database name.
    #[clap(short = 'd', long, default_value = "default")]
    pub(crate) chroma_database: String,

    /// Chroma token for authentication.
    #[clap(short = 't', long)]
    pub(crate) chroma_token: Option<String>,
}

/// Subcommand to find similar embeddings in a Chroma database.
///
/// This function connects to a Chroma database using the provided
/// arguments, performs a similarity search, and returns the results.
///
/// # Parameters
///
/// - `args`: The command-line arguments for the similarity search.
///
/// # Returns
///
/// A `Result` indicating success or failure.
pub(crate) async fn similarity_search(args: &SimilarityArgs) -> Result<(), RuChatError> {
    // Instantiate a ChromaClient to connect to the Chroma database
    let client = create_client(
        args.chroma_token.as_deref(),
        &args.chroma_server,
        &args.chroma_database,
    )
    .await?;

    // Instantiate a ChromaCollection to perform operations on a collection
    let collection: ChromaCollection = client
        .get_or_create_collection(&args.collection, None)
        .await?;

    let metadata = args.metadata.as_deref().map(|md| md.into());

    // Create a filter object to filter by document content.
    let where_document = json!({
        "$contains": args.query.as_str()
    });

    // Get embeddings from a collection with filters and limit set to 1.
    // An empty IDs vec will return all embeddings.
    let get_query = GetOptions {
        ids: vec![],
        where_metadata: metadata,
        limit: Some(args.count),
        offset: None,
        where_document: Some(where_document),
        include: Some(vec!["documents".into(), "embeddings".into()]),
    };
    let get_result: GetResult = collection.get(get_query).await?;

    // FIXME: This is a placeholder for the actual embeddings
    // Instantiate QueryOptions to perform a similarity search on the collection
    // Alternatively, an embedding_function can also be provided with query_texts to perform the search
    let query = QueryOptions {
        query_texts: None,
        query_embeddings: get_result
            .embeddings
            .map(|embeddings| embeddings.into_iter().flatten().collect()),
        where_metadata: None,
        where_document: None,
        n_results: Some(args.similarity_count),
        include: None,
    };

    let query_result: QueryResult = collection.query(query, None).await?;
    println!("Query result: {:?}", query_result);

    Ok(())
}