chore: return search sumamry with source id (#1336)

* chore: return search sumamry with source id

* chore: clippy

Cargo.lock

@@ -4185,6 +4185,7 @@ dependencies = [
  "serde",
  "serde_json",
  "tracing",
+ "uuid",
 ]
 
 [[package]]

libs/llm-client/Cargo.toml

@@ -9,4 +9,5 @@ tracing.workspace = true
 app-error = { workspace = true, features = ["appflowy_ai_error"] }
 serde_json.workspace = true
 serde.workspace = true
-schemars = "0.8.22"
+schemars = "0.8.22"
+uuid.workspace = true

libs/llm-client/src/chat.rs

@@ -7,8 +7,9 @@ use async_openai::types::{
 use async_openai::Client;
 use schemars::{schema_for, JsonSchema};
 use serde::{Deserialize, Serialize};
-use serde_json::{json, Value};
-use tracing::trace;
+use serde_json::json;
+use tracing::{error, info, trace};
+use uuid::Uuid;
 
 pub enum AITool {
   OpenAI(OpenAIChat),
@@ -16,11 +17,11 @@ pub enum AITool {
 }
 
 impl AITool {
-  pub async fn summary_documents(
+  pub async fn summarize_documents(
     &self,
     question: &str,
     model_name: &str,
-    documents: &[LLMDocument],
+    documents: Vec<LLMDocument>,
     only_context: bool,
   ) -> Result<SummarySearchResponse, AppError> {
     trace!(
@@ -77,60 +78,72 @@ impl AzureOpenAIChat {
   }
 }
 
-#[derive(Debug, Serialize, Deserialize, JsonSchema)]
-#[serde(deny_unknown_fields)]
+#[derive(Debug, Serialize, Deserialize)]
 pub struct SearchSummary {
   pub content: String,
-  pub metadata: Value,
-  pub score: String,
+  pub sources: Vec<Uuid>,
+  pub score: f32,
 }
 
-#[derive(Debug, Serialize, Deserialize, JsonSchema)]
-#[serde(deny_unknown_fields)]
+#[derive(Debug, Serialize, Deserialize)]
 pub struct SummarySearchResponse {
   pub summaries: Vec<SearchSummary>,
 }
 
+#[derive(Debug, Serialize, Deserialize, JsonSchema)]
+#[serde(deny_unknown_fields)]
+struct SummarySearchSchema {
+  pub answer: String,
+  pub score: String,
+  pub sources: Vec<String>,
+}
+
 const SYSTEM_PROMPT: &str = r#"
-You are a strict, context-based question answering assistant.
-Provide an answer with appropriate metadata in JSON format.
-For each answer, include:
-1. The answer text. It must be concise summaries and highly precise.
-2. Metadata with empty JSON map
-3. A numeric relevance score (0.0 to 1.0):
-   - 1.0: Completely relevant.
-   - 0.0: Not relevant.
-   - Intermediate values indicate partial relevance.
+You are a concise, intelligent question-answering assistant.
+
+Instructions:
+- Use the provided context as the **primary basis** for your answer.
+- You **may incorporate relevant knowledge** only if it supports or enhances the context meaningfully.
+- The answer must be a **clear and concise**.
+
+Output must include:
+- `answer`: a concise summary.
+- `score`: relevance score (0.0–1.0), where:
+  - 1.0 = fully supported by context,
+  - 0.0 = unsupported,
+  - values in between reflect partial support.
+- `sources`: array of source IDs used for the answer.
 "#;
+
 const ONLY_CONTEXT_SYSTEM_PROMPT: &str = r#"
 You are a strict, context-bound question answering assistant. Answer solely based on the text provided below. If the context lacks sufficient information for a confident response, reply with an empty answer.
 
-Your response must include:
-1. The answer text. It must be concise summaries and highly precise.
-2. Metadata extracted from the context. (If the answer is not relevant, return an empty JSON Map.)
-3. A numeric score (0.0 to 1.0) indicating the answer's relevance to the user's question:
-   - 1.0: Completely relevant.
-   - 0.0: Not relevant at all.
-   - Intermediate values indicate partial relevance.
+Output must include:
+- `answer`: a concise answer.
+- `score`: relevance score (0.0–1.0), where:
+  - 1.0 = fully supported by context,
+  - 0.0 = unsupported,
+  - values in between reflect partial support.
+- `sources`: array of source IDs used for the answer.
 
 Do not reference or use any information beyond what is provided in the context.
 "#;
 
-#[derive(Debug, Serialize, Deserialize)]
+#[derive(Debug, Serialize, Deserialize, Clone)]
 pub struct LLMDocument {
   pub content: String,
-  pub metadata: Value,
+  pub object_id: Uuid,
 }
 
 impl LLMDocument {
-  pub fn new(content: String, metadata: Value) -> Self {
-    Self { content, metadata }
+  pub fn new(content: String, object_id: Uuid) -> Self {
+    Self { content, object_id }
   }
 }
 
-fn convert_documents_to_text(documents: &[LLMDocument]) -> String {
+fn convert_documents_to_text(documents: Vec<LLMDocument>) -> String {
   documents
-    .iter()
+    .into_iter()
     .map(|doc| json!(doc).to_string())
     .collect::<Vec<String>>()
     .join("\n")
@@ -140,7 +153,7 @@ pub async fn summarize_documents<C: Config>(
   client: &Client<C>,
   question: &str,
   model_name: &str,
-  documents: &[LLMDocument],
+  documents: Vec<LLMDocument>,
   only_context: bool,
 ) -> Result<SummarySearchResponse, AppError> {
   let documents_text = convert_documents_to_text(documents);
@@ -153,12 +166,14 @@ pub async fn summarize_documents<C: Config>(
     SYSTEM_PROMPT.to_string()
   };
 
-  let schema = schema_for!(SummarySearchResponse);
+  let schema = schema_for!(SummarySearchSchema);
   let schema_value = serde_json::to_value(&schema)?;
   let response_format = ResponseFormat::JsonSchema {
     json_schema: ResponseFormatJsonSchema {
-      description: Some("A response containing a list of answers, each with the answer text, metadata extracted from context, and relevance score".to_string()),
-      name: "SummarySearchResponse".into(),
+      description: Some(
+        "A response containing a final answer, score and relevance sources".to_string(),
+      ),
+      name: "SummarySearchSchema".into(),
       schema: Some(schema_value),
       strict: Some(true),
     },
@@ -179,20 +194,51 @@ pub async fn summarize_documents<C: Config>(
     .response_format(response_format)
     .build()?;
 
-  let mut response = client
+  let response = client
     .chat()
     .create(request)
     .await?
     .choices
     .first()
     .and_then(|choice| choice.message.content.clone())
-    .and_then(|content| serde_json::from_str::<SummarySearchResponse>(&content).ok())
+    .and_then(|content| serde_json::from_str::<SummarySearchSchema>(&content).ok())
     .ok_or_else(|| AppError::Unhandled("No response from OpenAI".to_string()))?;
 
-  // Remove empty summaries
-  response
-    .summaries
-    .retain(|summary| !summary.content.is_empty());
+  if response.answer.is_empty() {
+    return Ok(SummarySearchResponse { summaries: vec![] });
+  }
 
-  Ok(response)
+  let score = match response.score.parse::<f32>() {
+    Ok(score) => score,
+    Err(err) => {
+      error!(
+        "[Search] Failed to parse AI summary score: {}. Error: {}",
+        response.score, err
+      );
+      0.0
+    },
+  };
+
+  // If only_context is true, we need to ensure the score is above a certain threshold.
+  if only_context && score < 0.4 {
+    info!(
+      "[Search] AI summary score is too low: {}. Returning empty result.",
+      score
+    );
+    return Ok(SummarySearchResponse { summaries: vec![] });
+  }
+
+  let summary = SearchSummary {
+    content: response.answer,
+    sources: response
+      .sources
+      .into_iter()
+      .flat_map(|s| Uuid::parse_str(&s).ok())
+      .collect(),
+    score,
+  };
+
+  Ok(SummarySearchResponse {
+    summaries: vec![summary],
+  })
 }

libs/shared-entity/src/dto/search_dto.rs

@@ -1,6 +1,5 @@
 use chrono::{DateTime, Utc};
 use serde::{Deserialize, Serialize};
-use serde_json::Value;
 use serde_repr::{Deserialize_repr, Serialize_repr};
 use uuid::Uuid;
 
@@ -53,8 +52,7 @@ fn default_search_score_limit() -> f64 {
 #[derive(Clone, Debug, Serialize, Deserialize)]
 pub struct Summary {
   pub content: String,
-  pub metadata: Value,
-  pub score: String,
+  pub sources: Vec<Uuid>,
 }
 
 /// Response array element for the collab vector search query.

src/api/search.rs

@@ -1,4 +1,4 @@
-use crate::biz::search::{search_document, summary_search_results};
+use crate::biz::search::{search_document, summarize_search_results};
 use crate::state::AppState;
 use access_control::act::Action;
 use actix_web::web::{Data, Json, Query};
@@ -67,7 +67,7 @@ async fn summary_search_results_handler(
     .await?;
 
   let ai_tool = create_ai_tool(&state.config.azure_ai_config, &state.config.open_ai_config);
-  let result = summary_search_results(ai_tool, request).await?;
+  let result = summarize_search_results(ai_tool, request).await?;
   Ok(AppResponse::Ok().with_data(result).into())
 }
 

src/biz/search/ops.rs

@@ -13,7 +13,6 @@ use indexer::scheduler::IndexerScheduler;
 use indexer::vector::embedder::{CreateEmbeddingRequestArgs, EmbeddingInput, EncodingFormat};
 use infra::env_util::get_env_var;
 use llm_client::chat::{AITool, LLMDocument};
-use serde_json::json;
 use shared_entity::dto::search_dto::{
   SearchContentType, SearchDocumentRequest, SearchDocumentResponseItem, SearchSummaryResult,
   Summary, SummarySearchResultRequest,
@@ -176,7 +175,7 @@ pub async fn search_document(
   Ok(items)
 }
 
-pub async fn summary_search_results(
+pub async fn summarize_search_results(
   ai_tool: Option<AITool>,
   request: SummarySearchResultRequest,
 ) -> Result<SearchSummaryResult, AppError> {
@@ -206,19 +205,10 @@ pub async fn summary_search_results(
 
   let llm_docs: Vec<LLMDocument> = search_results
     .into_iter()
-    .map(|result| {
-      LLMDocument::new(
-        result.content,
-        json!({
-            "id": result.object_id,
-            "source": "appflowy",
-            "name": "document",
-        }),
-      )
-    })
+    .map(|result| LLMDocument::new(result.content, result.object_id))
     .collect();
   match ai_tool
-    .summary_documents(&query, &model_name, &llm_docs, only_context)
+    .summarize_documents(&query, &model_name, llm_docs, only_context)
     .await
   {
     Ok(resp) => {
@@ -228,8 +218,7 @@ pub async fn summary_search_results(
         .into_iter()
         .map(|s| Summary {
           content: s.content,
-          metadata: s.metadata,
-          score: s.score,
+          sources: s.sources,
         })
         .collect();
     },

tests/ai_test/summary_search_test.rs

@@ -2,7 +2,7 @@ use appflowy_cloud::api::search::create_ai_tool;
 use client_api_test::{ai_test_enabled, load_env};
 use indexer::vector::embedder::get_open_ai_config;
 use llm_client::chat::LLMDocument;
-use serde_json::json;
+use uuid::Uuid;
 
 #[tokio::test]
 async fn chat_with_search_result_simple() {
@@ -13,32 +13,20 @@ async fn chat_with_search_result_simple() {
   let ai_chat = create_ai_tool(&azure_config, &open_ai_config).unwrap();
   let model_name = "gpt-4o-mini";
   let docs = vec![
-    "“GPT-4o-mini” is typically used to suggest a streamlined version of the GPT-4 family. The idea is to create a model that maintains the core language capabilities of GPT-4 while reducing computational requirements",
-    "The name “Llama3.1” hints at an incremental evolution within the LLaMA (Large Language Model Meta AI) series—a family of models designed by Meta for research accessibility and performance efficiency",
-  ].into_iter().map(|content| LLMDocument::new(content.to_string(), json!({
-    "id": content.len().to_string(),
-    "source": if content.contains("GPT-4") { "openai_docs" } else { "meta_docs" },
-    "timestamp": if content.contains("GPT-4") { "2024-01-01" } else { "2024-02-01" }
-  }))).collect::<Vec<_>>();
+    ("GPT-4o-mini is typically used to suggest a streamlined version of the GPT-4 family. The idea is to create a model that maintains the core language capabilities of GPT-4 while reducing computational requirements", Uuid::new_v4()),
+    ("The name “Llama3.1” hints at an incremental evolution within the LLaMA (Large Language Model Meta AI) series—a family of models designed by Meta for research accessibility and performance efficiency",Uuid::new_v4()),
+  ].into_iter().map(|(content, object_id)| LLMDocument::new(content.to_string(), object_id)).collect::<Vec<_>>();
 
   let resp = ai_chat
-    .summary_documents("gpt-4o", model_name, &docs, true)
+    .summarize_documents("gpt-4o", model_name, docs.clone(), true)
     .await
     .unwrap();
   dbg!(&resp);
   assert_eq!(resp.summaries.len(), 1);
-  assert!(resp.summaries[0].score.parse::<f32>().unwrap() > 0.6);
-  assert_eq!(
-    resp.summaries[0].metadata.get("source").unwrap(),
-    "openai_docs"
-  );
-  assert_eq!(
-    resp.summaries[0].metadata.get("timestamp").unwrap(),
-    "2024-01-01"
-  );
+  assert_eq!(resp.summaries[0].sources[0], docs[0].object_id);
 
   let resp = ai_chat
-    .summary_documents("deepseek-r1", model_name, &docs, true)
+    .summarize_documents("deepseek-r1", model_name, docs.clone(), true)
     .await
     .unwrap();
   dbg!(&resp);
@@ -46,13 +34,12 @@ async fn chat_with_search_result_simple() {
 
   // When only_context is false, the llm knowledge base is used to answer the question.
   let resp = ai_chat
-    .summary_documents("deepseek-r1 llm model", model_name, &docs, false)
+    .summarize_documents("deepseek-r1 llm model", model_name, docs, false)
     .await
     .unwrap();
   dbg!(&resp);
   assert_eq!(resp.summaries.len(), 1);
-  assert!(resp.summaries[0].score.parse::<f32>().unwrap() > 0.6);
-  assert_eq!(resp.summaries[0].metadata, json!({}));
+  assert!(resp.summaries[0].sources.is_empty());
 }
 
 #[tokio::test]
@@ -65,22 +52,25 @@ async fn summary_search_result() {
   let ai_chat = create_ai_tool(&azure_config, &open_ai_config).unwrap();
   let model_name = "gpt-4o-mini";
   let docs = vec![
-    ("Rust is a multiplayer survival game developed by Facepunch Studios, first released in early access in December 2013 and fully launched in February 2018. It has since become one of the most popular games in the survival genre, known for its harsh environment, intricate crafting system, and player-driven dynamics. The game is available on Windows, macOS, and PlayStation, with a community-driven approach to updates and content additions.", json!({"id": 1, "source": "test"})),
-    ("Rust is a modern, system-level programming language designed with a focus on performance, safety, and concurrency. It was created by Mozilla and first released in 2010, with its 1.0 version launched in 2015. Rust is known for providing the control and performance of languages like C and C++, but with built-in safety features that prevent common programming errors, such as memory leaks, data races, and buffer overflows.", json!({"id": 2, "source": "test2"})),
-    ("Rust as a Natural Process (Oxidation) refers to the chemical reaction that occurs when metals, primarily iron, come into contact with oxygen and moisture (water) over time, leading to the formation of iron oxide, commonly known as rust. This process is a form of oxidation, where a substance reacts with oxygen in the air or water, resulting in the degradation of the metal.", json!({"id": 3})),
-  ].into_iter().map(|(content, metadata)| LLMDocument::new(content.to_string(), metadata)).collect::<Vec<_>>();
+    ("Rust is a multiplayer survival game developed by Facepunch Studios, first released in early access in December 2013 and fully launched in February 2018. It has since become one of the most popular games in the survival genre, known for its harsh environment, intricate crafting system, and player-driven dynamics. The game is available on Windows, macOS, and PlayStation, with a community-driven approach to updates and content additions.", uuid::Uuid::new_v4()),
+    ("Rust is a modern, system-level programming language designed with a focus on performance, safety, and concurrency. It was created by Mozilla and first released in 2010, with its 1.0 version launched in 2015. Rust is known for providing the control and performance of languages like C and C++, but with built-in safety features that prevent common programming errors, such as memory leaks, data races, and buffer overflows.", uuid::Uuid::new_v4()),
+    ("Rust as a Natural Process (Oxidation) refers to the chemical reaction that occurs when metals, primarily iron, come into contact with oxygen and moisture (water) over time, leading to the formation of iron oxide, commonly known as rust. This process is a form of oxidation, where a substance reacts with oxygen in the air or water, resulting in the degradation of the metal.", uuid::Uuid::new_v4()),
+  ].into_iter().map(|(content, object_id)| LLMDocument::new(content.to_string(), object_id)).collect::<Vec<_>>();
 
   let resp = ai_chat
-    .summary_documents("Rust", model_name, &docs, true)
-    .await
-    .unwrap();
-  dbg!(&resp);
-  assert_eq!(resp.summaries.len(), 3);
-
-  let resp = ai_chat
-    .summary_documents("Play Rust over time", model_name, &docs, true)
+    .summarize_documents("What is Rust", model_name, docs.clone(), true)
     .await
     .unwrap();
   dbg!(&resp);
   assert_eq!(resp.summaries.len(), 1);
+  assert_eq!(resp.summaries[0].sources.len(), 3);
+
+  let resp = ai_chat
+    .summarize_documents("multiplayer game", model_name, docs.clone(), true)
+    .await
+    .unwrap();
+  dbg!(&resp);
+  assert_eq!(resp.summaries.len(), 1);
+  assert_eq!(resp.summaries[0].sources.len(), 1);
+  assert_eq!(resp.summaries[0].sources[0], docs[0].object_id);
 }