kwaak 0.1.2

Run a team of autonomous agents on your code, right from your terminal
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

use anyhow::Result;
use swiftide::{
    query::{
        self, answers, query_transformers, search_strategies::SimilaritySingleEmbedding, states,
    },
    traits::{EmbeddingModel, SimplePrompt},
};

use crate::{repository::Repository, storage, util::strip_markdown_tags};

pub async fn query(repository: &Repository, query: &str) -> Result<String> {
    let answer = build_query_pipeline(repository)?
        .query(query)
        .await?
        .answer()
        .to_string();
    Ok(strip_markdown_tags(&answer))
}

pub fn build_query_pipeline<'b>(
    repository: &Repository,
) -> Result<query::Pipeline<'b, SimilaritySingleEmbedding, states::Answered>> {
    let query_provider: Box<dyn SimplePrompt> = repository.config().query_provider().try_into()?;
    let embedding_provider: Box<dyn EmbeddingModel> =
        repository.config().embedding_provider().try_into()?;

    let lancedb = storage::build_lancedb(repository)?;
    let search_strategy: SimilaritySingleEmbedding<()> = SimilaritySingleEmbedding::default()
        .with_top_k(20)
        .to_owned();

    Ok(query::Pipeline::from_search_strategy(search_strategy)
        .then_transform_query(query_transformers::GenerateSubquestions::from_client(
            query_provider.clone(),
        ))
        .then_transform_query(query_transformers::Embed::from_client(
            embedding_provider.clone(),
        ))
        .then_retrieve(lancedb.build()?)
        // .then_transform_response(response_transformers::Summary::from_client(
        //     query_provider.clone(),
        // ))
        .then_answer(answers::Simple::from_client(query_provider.clone())))
}