kwaak 0.16.2

Run a team of autonomous agents on your code, right from your terminal
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
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135

use anyhow::Result;
use indoc::formatdoc;
use swiftide::traits::EvaluateQuery;
use swiftide::{
    query::{
        self, answers, query_transformers, search_strategies::SimilaritySingleEmbedding, states,
        Query,
    },
    traits::{EmbeddingModel, Persist, Retrieve, SimplePrompt},
};

use crate::{repository::Repository, templates::Templates, util::strip_markdown_tags};

#[tracing::instrument(skip_all, err)]
pub async fn query<S>(
    repository: &Repository,
    storage: &S,
    query: impl AsRef<str>,
) -> Result<String>
where
    S: Retrieve<SimilaritySingleEmbedding> + Persist + Clone + 'static,
{
    // Ensure the table exists to avoid dumb errors
    let _ = storage.setup().await;

    let answer = build_query_pipeline(repository, storage, None)?
        .query(query.as_ref())
        .await?
        .answer()
        .to_string();
    Ok(strip_markdown_tags(&answer))
}

/// Builds a query pipeline
///
/// # Panics
///
/// Should be infallible
pub fn build_query_pipeline<'b, S>(
    repository: &Repository,
    storage: &S,
    evaluator: Option<Box<dyn EvaluateQuery>>,
) -> Result<query::Pipeline<'b, SimilaritySingleEmbedding, states::Answered>>
where
    S: Retrieve<SimilaritySingleEmbedding> + Clone + 'static,
{
    let backoff = repository.config().backoff;
    let query_provider: Box<dyn SimplePrompt> = repository
        .config()
        .query_provider()
        .get_simple_prompt_model(backoff)?;
    let embedding_provider: Box<dyn EmbeddingModel> = repository
        .config()
        .embedding_provider()
        .get_embedding_model(backoff)?;

    let search_strategy: SimilaritySingleEmbedding<()> = SimilaritySingleEmbedding::default()
        .with_top_k(30)
        .to_owned();

    let prompt_template = Templates::from_file("agentic_answer_prompt.md")?;
    let document_template = Templates::from_file("indexing_document.md")?;

    // NOTE: Changed a lot to tailor it for agentic flows, might be worth upstreaming
    // Simple takes the retrieved documents, formats them with a template, then throws it into a
    // prompt with to answer the original question properly. It's really simple.
    let simple = answers::Simple::builder()
        .client(query_provider.clone())
        .prompt_template(prompt_template.into())
        .document_template(document_template)
        .build()
        .expect("infallible");

    let language = repository.config().language.to_string();
    let project = repository.config().project_name.clone();

    let mut pipeline = query::Pipeline::from_search_strategy(search_strategy);

    if let Some(evaluator) = evaluator {
        pipeline = pipeline.evaluate_with(evaluator);
    }

    Ok(pipeline
        .then_transform_query(move |mut query: Query<states::Pending>| {
            let current = query.current();
            query.transformed_query(formatdoc! {"
                {current}

                The project is written in: {language}
                The project is called: {project}

            "});

            Ok(query)
        })
        .then_transform_query(query_transformers::GenerateSubquestions::from_client(
            query_provider.clone(),
        ))
        .then_transform_query(query_transformers::Embed::from_client(
            embedding_provider.clone(),
        ))
        .then_retrieve(storage.clone())
        // .then_transform_response(response_transformers::Summary::from_client(
        //     query_provider.clone(),
        // ))
        .then_answer(simple))
}

#[cfg(test)]
mod tests {

    #[test]
    fn test_rendering_document() {
        use insta::assert_snapshot;
        use swiftide::{indexing::Metadata, query::Document};
        use tera::Context;

        use crate::templates::Templates;

        let document = Document::new(
            "This is a test document",
            Some(Metadata::from([
                ("path", serde_json::Value::from("my file")),
                ("soups", serde_json::Value::from(["tomato", "snert"])),
                ("empty", serde_json::Value::from("")),
            ])),
        );

        assert_snapshot!(Templates::render(
            "indexing_document.md",
            &Context::from_serialize(document).unwrap(),
        )
        .unwrap());
    }
}