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

// main_rag.rs

use crate::chroma::embed::{embed, EmbedArgs};
use crate::chroma::similarity::{similarity_search, SimilarityArgs};
use crate::ollama::chat::conversation_tree::ConversationTree;
use crate::ollama::Ollama;
use anyhow::Result;
use tokio_stream::StreamExt;

#[tokio::main]
async fn main() -> Result<()> {
    // Initialize retrieval and generation components
    println!("Initializing RAG system...");

    // Step 1: Perform Embedding
    let ollama = Ollama::new(); // Initialize Ollama client
    let embed_args = EmbedArgs {
        model: "nomic-embed-text:latest".to_string(),
        prompt: "example query".to_string(),
        chroma_server: "http://localhost:8000".to_string(),
        chroma_database: "default".to_string(),
        chroma_token: None,
        collection: "default".to_string(),
        collection_metadata: None,
        entries_metadata: None,
    };
    embed(ollama.clone(), &embed_args).await?;

    // Step 2: Perform Similarity Search
    let similarity_args = SimilarityArgs {
        query: "example query".to_string(),
        count: 1,
        similarity_count: 5,
        collection: "default".to_string(),
        metadata: None,
        chroma_server: "http://localhost:8000".to_string(),
        chroma_database: "default".to_string(),
        chroma_token: None,
    };
    similarity_search(&similarity_args).await?;

    // Step 3: Perform Generation
    let mut conversation_tree = ConversationTree::new();
    let question_id = conversation_tree.question(vec![similarity_args.query.clone()])?;
    conversation_tree.add_answer(question_id, vec!["Generated response based on similar entries.".to_string()])?;
    println!("Generated response: {:?}", conversation_tree.get_qa(question_id, conversation_tree.get_current_answer_id(question_id)));

    // Step 4: Integrate and Output
    println!("Integration complete. Outputting response...");

    Ok(())
}