retrieval-pipeline: v1

html-router/src/routes/chat/message_response_stream.rs

@@ -16,7 +16,7 @@ use json_stream_parser::JsonStreamParser;
 use minijinja::Value;
 use retrieval_pipeline::{
     answer_retrieval::{create_chat_request, create_user_message_with_history, LLMResponseFormat},
-    retrieve_entities, retrieved_entities_to_json, RetrievalConfig, StrategyOutput,
+    retrieved_entities_to_json,
 };
 use serde::{Deserialize, Serialize};
 use serde_json::from_str;
@@ -123,23 +123,22 @@ pub async fn get_response_stream(
         None => None,
     };
 
-    let mut retrieval_config = RetrievalConfig::default();
-    retrieval_config.strategy = state.retrieval_strategy();
-    let entities = match retrieve_entities(
+    let strategy = state.retrieval_strategy();
+    let config = retrieval_pipeline::RetrievalConfig::for_chat(strategy);
+
+    let entities = match retrieval_pipeline::retrieve_entities(
         &state.db,
         &state.openai_client,
         &user_message.content,
         &user.id,
-        retrieval_config,
+        config,
         rerank_lease,
     )
     .await
     {
-        Ok(StrategyOutput::Entities(entities)) => entities,
-        Ok(StrategyOutput::Chunks(_)) => {
-            return Sse::new(create_error_stream(
-                "Chunk-only retrieval results are not supported in this route",
-            ))
+        Ok(retrieval_pipeline::StrategyOutput::Entities(entities)) => entities,
+        Ok(retrieval_pipeline::StrategyOutput::Chunks(_chunks)) => {
+            return Sse::new(create_error_stream("Chat retrieval currently only supports Entity-based strategies (Initial). Revised strategy returns Chunks which are not yet supported by this handler."));
         }
         Err(_e) => {
             return Sse::new(create_error_stream("Failed to retrieve knowledge entities"));

html-router/src/routes/knowledge/handlers.rs

@@ -24,7 +24,7 @@ use common::{
     },
     utils::embedding::generate_embedding,
 };
-use retrieval_pipeline::{retrieve_entities, RetrievalConfig, RetrievedEntity, StrategyOutput};
+use retrieval_pipeline;
 use tracing::debug;
 use uuid::Uuid;
 
@@ -284,20 +284,18 @@ pub async fn suggest_knowledge_relationships(
             None => None,
         };
 
-        let mut retrieval_config = RetrievalConfig::default();
-        retrieval_config.strategy = state.retrieval_strategy();
-
-        if let Ok(StrategyOutput::Entities(results)) = retrieve_entities(
+        let config = retrieval_pipeline::RetrievalConfig::for_relationship_suggestion();
+        if let Ok(retrieval_pipeline::StrategyOutput::Entities(results)) = retrieval_pipeline::retrieve_entities(
             &state.db,
             &state.openai_client,
             &query,
             &user.id,
-            retrieval_config,
+            config,
             rerank_lease,
         )
         .await
         {
-            for RetrievedEntity { entity, score, .. } in results {
+            for retrieval_pipeline::RetrievedEntity { entity, score, .. } in results {
                 if suggestion_scores.len() >= MAX_RELATIONSHIP_SUGGESTIONS {
                     break;
                 }