Re-introduce HyDE, work around batched calls (#735)

* Re-introduce HyDE logic

This reverts commit f18eb02.

* Add trace logs to hyde doc generation

* Log around qdrant batch search

* Patch batch_search_with using multiple calls to search_points

Author: calyptobai

server/bleep/src/semantic.rs

@@ -18,6 +18,7 @@ use qdrant_client::{
     },
 };
 
+use futures::{stream, StreamExt, TryStreamExt};
 use rayon::prelude::*;
 use thiserror::Error;
 use tracing::{debug, info, trace, warn};
@@ -326,6 +327,62 @@ impl Semantic {
         Ok(response.result)
     }
 
+    pub async fn batch_search_with<'a>(
+        &self,
+        parsed_queries: &[&SemanticQuery<'a>],
+        vectors: Vec<Embedding>,
+        limit: u64,
+        offset: u64,
+    ) -> anyhow::Result<Vec<ScoredPoint>> {
+        // FIXME: This method uses `search_points` internally, and not `search_batch_points`. It's
+        // not clear why, but it seems that the `batch` variant of the `qdrant` calls leads to
+        // HTTP2 errors on some deployment configurations. A typical example error:
+        //
+        // ```
+        // hyper::proto::h2::client: client response error: stream error received: stream no longer needed
+        // ```
+        //
+        // Given that qdrant uses `tonic`, this may be a `tonic` issue, possibly similar to:
+        // https://github.com/hyperium/tonic/issues/222
+
+        // Queries should contain the same filters, so we get the first one
+        let parsed_query = parsed_queries.first().unwrap();
+        let filters = &build_conditions(parsed_query);
+
+        let responses = stream::iter(vectors.into_iter())
+            .map(|vector| async move {
+                let points = SearchPoints {
+                    limit,
+                    vector,
+                    collection_name: COLLECTION_NAME.to_string(),
+                    offset: Some(offset),
+                    score_threshold: Some(SCORE_THRESHOLD),
+                    with_payload: Some(WithPayloadSelector {
+                        selector_options: Some(with_payload_selector::SelectorOptions::Enable(
+                            true,
+                        )),
+                    }),
+                    filter: Some(Filter {
+                        must: filters.clone(),
+                        ..Default::default()
+                    }),
+                    with_vectors: Some(WithVectorsSelector {
+                        selector_options: Some(with_vectors_selector::SelectorOptions::Enable(
+                            true,
+                        )),
+                    }),
+                    ..Default::default()
+                };
+
+                self.qdrant.search_points(&points).await
+            })
+            .buffered(10)
+            .try_collect::<Vec<_>>()
+            .await?;
+
+        Ok(responses.into_iter().flat_map(|r| r.result).collect())
+    }
+
     pub async fn search<'a>(
         &self,
         parsed_query: &SemanticQuery<'a>,
@@ -357,6 +414,46 @@ impl Semantic {
         Ok(deduplicate_snippets(results, vector, limit))
     }
 
+    pub async fn batch_search<'a>(
+        &self,
+        parsed_queries: &[&SemanticQuery<'a>],
+        limit: u64,
+        offset: u64,
+        retrieve_more: bool,
+    ) -> anyhow::Result<Vec<Payload>> {
+        if parsed_queries.iter().any(|q| q.target().is_none()) {
+            anyhow::bail!("no search target for query");
+        };
+
+        let vectors = parsed_queries
+            .iter()
+            .map(|q| self.embed(&q.target().unwrap()))
+            .collect::<anyhow::Result<Vec<_>>>()?;
+
+        tracing::trace!(?parsed_queries, "performing qdrant batch search");
+
+        let result = self
+            .batch_search_with(
+                parsed_queries,
+                vectors.clone(),
+                if retrieve_more { limit * 2 } else { limit }, // Retrieve double `limit` and deduplicate
+                offset,
+            )
+            .await;
+
+        tracing::trace!(?result, "qdrant batch search returned");
+
+        let results = result?
+            .into_iter()
+            .map(Payload::from_qdrant)
+            .collect::<Vec<_>>();
+
+        // deduplicate with mmr with respect to the mean of query vectors
+        // TODO: implement a more robust multi-vector deduplication strategy
+        let target_vector = mean_pool(vectors);
+        Ok(deduplicate_snippets(results, target_vector, limit))
+    }
+
     #[allow(clippy::too_many_arguments)]
     #[tracing::instrument(skip(self, repo_ref, relative_path, buffer))]
     pub async fn insert_points_for_buffer(
@@ -593,6 +690,19 @@ fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 {
     dot(a, b) / (norm(a) * norm(b))
 }
 
+// Calculate the element-wise mean of the embeddings
+fn mean_pool(embeddings: Vec<Vec<f32>>) -> Vec<f32> {
+    let len = embeddings.len() as f32;
+    let mut result = vec![0.0; EMBEDDING_DIM];
+    for embedding in embeddings {
+        for (i, v) in embedding.iter().enumerate() {
+            result[i] += v;
+        }
+    }
+    result.iter_mut().for_each(|v| *v /= len);
+    result
+}
+
 // returns a list of indices to preserve from `snippets`
 //
 // query_embedding: the embedding of the query terms

server/bleep/src/webserver/answer.rs

@@ -738,7 +738,14 @@ impl Agent {
                 self.update(Update::Step(SearchStep::Code(query.clone())))
                     .await?;
 
-                let results = self.semantic_search(query.into(), 10, 0, true).await?;
+                let mut results = self.semantic_search(query.into(), 10, 0, true).await?;
+
+                let hyde_docs = self.hyde(query).await?;
+                if !hyde_docs.is_empty() {
+                    let hyde_docs = hyde_docs.iter().map(|d| d.into()).collect();
+                    let hyde_results = self.batch_semantic_search(hyde_docs, 10, 0, true).await?;
+                    results.extend(hyde_results);
+                }
 
                 let chunks = results
                     .into_iter()
@@ -764,6 +771,7 @@ impl Agent {
                 self.track_query(
                     EventData::input_stage("semantic code search")
                         .with_payload("query", query)
+                        .with_payload("hyde_queries", &hyde_docs)
                         .with_payload("chunks", &chunks)
                         .with_payload("raw_prompt", &prompt),
                 );
@@ -849,6 +857,33 @@ impl Agent {
         Ok(Some(action))
     }
 
+    async fn hyde(&self, query: &str) -> Result<Vec<String>> {
+        let prompt = vec![llm_gateway::api::Message::system(
+            &prompts::hypothetical_document_prompt(query),
+        )];
+
+        tracing::trace!(?query, "generating hyde docs");
+
+        let response = self
+            .llm_gateway
+            .clone()
+            .model("gpt-3.5-turbo-0613")
+            .chat(&prompt, None)
+            .await?
+            .try_collect::<String>()
+            .await?;
+
+        tracing::trace!("parsing hyde response");
+
+        let documents = prompts::try_parse_hypothetical_documents(&response);
+
+        for doc in documents.iter() {
+            info!(?doc, "got hyde doc");
+        }
+
+        Ok(documents)
+    }
+
     async fn proc(&mut self, question: &str, path_aliases: &[usize]) -> Result<String> {
         let paths = path_aliases
             .iter()
@@ -1364,6 +1399,36 @@ impl Agent {
             .await
     }
 
+    async fn batch_semantic_search(
+        &self,
+        queries: Vec<Literal<'_>>,
+        limit: u64,
+        offset: u64,
+        retrieve_more: bool,
+    ) -> Result<Vec<semantic::Payload>> {
+        let queries = queries
+            .iter()
+            .map(|q| SemanticQuery {
+                target: Some(q.clone()),
+                repos: [Literal::Plain(
+                    self.conversation.repo_ref.display_name().into(),
+                )]
+                .into(),
+                ..self.conversation.last_exchange().query.clone()
+            })
+            .collect::<Vec<_>>();
+
+        let queries = queries.iter().collect::<Vec<_>>();
+
+        debug!(?queries, %self.thread_id, "executing semantic query");
+        self.app
+            .semantic
+            .as_ref()
+            .unwrap()
+            .batch_search(queries.as_slice(), limit, offset, retrieve_more)
+            .await
+    }
+
     async fn get_file_content(&self, path: &str) -> Result<Option<ContentDocument>> {
         let branch = self.conversation.last_exchange().query.first_branch();
         let repo_ref = &self.conversation.repo_ref;

server/bleep/src/webserver/answer/prompts.rs

@@ -280,3 +280,100 @@ Respect these rules at all times:
 - Always begin your answer with an appropriate title"#
     )
 }
+
+pub fn hypothetical_document_prompt(query: &str) -> String {
+    format!(
+        r#"Write three code snippets that could hypothetically be returned by a code search engine as the answer to the query: {query}
+
+- Write these snippets in a variety of programming languages
+- The snippets should not be too similar to one another
+- Each snippet should be between 5 and 10 lines long
+- Surround the snippets in triple backticks
+
+For example:
+
+What's the Qdrant threshold?
+
+```rust
+SearchPoints {{
+    limit,
+    vector: vectors.get(idx).unwrap().clone(),
+    collection_name: COLLECTION_NAME.to_string(),
+    offset: Some(offset),
+    score_threshold: Some(0.3),
+    with_payload: Some(WithPayloadSelector {{
+        selector_options: Some(with_payload_selector::SelectorOptions::Enable(true)),
+    }}),
+```
+
+```python
+memories = self.client.search(
+    collection_name=self.collection_name,
+    limit=k,
+    score_threshold=threshold,
+    search_params=SearchParams(
+        quantization=QuantizationSearchParams(
+            ignore=False,
+            rescore=True
+        )
+    )
+```
+
+```typescript
+const res = await qdrant.search(collectionName, {{
+    vector: embedding,
+    score_threshold:
+    distanceMetric === "Cosine" ? 1 - threshold : threshold,
+    with_payload: true,
+}});
+```
+"#
+    )
+}
+
+pub fn try_parse_hypothetical_documents(document: &str) -> Vec<String> {
+    let pattern = r"```([\s\S]*?)```";
+    let re = regex::Regex::new(pattern).unwrap();
+
+    re.captures_iter(document)
+        .map(|m| m[1].trim().to_string())
+        .collect()
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_parse_hypothetical_document() {
+        let document = r#"Here is some pointless text
+
+        ```rust
+pub fn final_explanation_prompt(context: &str, query: &str, query_history: &str) -> String {
+    struct Rule<'a> {
+        title: &'a str,
+        description: &'a str,
+        note: &'a str,
+        schema: &'a str,```
+
+Here is some more pointless text
+
+```
+pub fn functions() -> serde_json::Value {
+    serde_json::json!(
+```"#;
+        let expected = vec![
+            r#"rust
+pub fn final_explanation_prompt(context: &str, query: &str, query_history: &str) -> String {
+    struct Rule<'a> {
+        title: &'a str,
+        description: &'a str,
+        note: &'a str,
+        schema: &'a str,"#,
+            r#"pub fn functions() -> serde_json::Value {
+    serde_json::json!("#,
+        ];
+
+        assert_eq!(try_parse_hypothetical_documents(document), expected);
+    }
+}