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"));