fix(page): add doc query when convert thread to page (#4179)

* fix(page): add doc query when convert thread to page

* chore(page): always return doc and code query event even no doc or code matched

Author: zwpaper

ee/tabby-db/src/pages.rs

@@ -31,7 +31,7 @@ pub struct PageSectionDAO {
 }
 
 impl DbConn {
-    pub async fn create_page(&self, author_id: i64, code_source_id: Option<String>) -> Result<i64> {
+    pub async fn create_page(&self, author_id: i64, code_source_id: Option<&str>) -> Result<i64> {
         let res = query!(
             "INSERT INTO pages(author_id, code_source_id) VALUES (?, ?)",
             author_id,

ee/tabby-webserver/src/service/page.rs

@@ -94,18 +94,25 @@ impl PageService for PageServiceImpl {
             .list_thread_messages(&input.thread_id, None, None, None, None)
             .await?;
 
+        // use the first message source id for page
+        let code_source_id = thread_messages
+            .iter()
+            .find_map(|m| m.code_source_id.as_deref())
+            .map(ToOwned::to_owned);
+
         let debug_option = input
             .debug_option
             .as_ref()
             .map(|d| PageRunDebugOptionInput {
                 return_chat_completion_request: d.return_chat_completion_request,
                 return_query_request: d.return_query_request,
             });
+
         self.page_run(
             policy,
             author_id,
             None,
-            None,
+            code_source_id.as_deref(),
             None,
             Some(&thread_messages),
             debug_option.as_ref(),
@@ -119,11 +126,17 @@ impl PageService for PageServiceImpl {
         author_id: &ID,
         input: &CreatePageRunInput,
     ) -> Result<PageRunStream> {
+        let code_source_id = if let Some(code_query) = input.code_query.as_ref() {
+            get_source_id(self.context.clone(), policy, code_query).await
+        } else {
+            None
+        };
+
         self.page_run(
             policy,
             author_id,
             Some(&input.title_prompt),
-            input.code_query.as_ref(),
+            code_source_id.as_deref(),
             input.doc_query.as_ref(),
             None,
             input.debug_option.as_ref(),
@@ -215,7 +228,7 @@ impl PageService for PageServiceImpl {
                 Some(new_section_prompt),
                 &existing_page_sections,
                 &page_title,
-                code_source_id,
+                code_source_id.as_deref(),
                 doc_query.as_ref(),
                 None,
                 db,
@@ -356,26 +369,14 @@ impl PageServiceImpl {
         policy: &AccessPolicy,
         author_id: &ID,
         title_prompt: Option<&str>,
-        code_query: Option<&CodeQueryInput>,
+        code_source_id: Option<&str>,
         doc_query: Option<&DocQueryInput>,
         thread_messages: Option<&[Message]>,
         debug_option: Option<&PageRunDebugOptionInput>,
     ) -> Result<PageRunStream> {
-        let code_source_id = if let Some(code_query) = code_query {
-            get_source_id(self.context.clone(), policy, code_query).await
-        } else {
-            // use the first message source id for page
-            thread_messages.and_then(|messages| {
-                messages
-                    .iter()
-                    .find_map(|m| m.code_source_id.as_deref())
-                    .map(ToOwned::to_owned)
-            })
-        };
-
         let page_id = self
             .db
-            .create_page(author_id.as_rowid()?, code_source_id.clone())
+            .create_page(author_id.as_rowid()?, code_source_id)
             .await?
             .as_id();
 
@@ -397,6 +398,17 @@ impl PageServiceImpl {
             )
             .await?;
 
+        let code_source_id = code_source_id.map(ToOwned::to_owned);
+        let doc_query = if let Some(doc_query) = doc_query {
+            Some(doc_query.to_owned())
+        } else {
+            code_source_id.as_ref().map(|code_source_id| DocQueryInput {
+                content: title_prompt.unwrap_or_else(|| &page_title).to_owned(),
+                search_public: false,
+                source_ids: Some(vec![code_source_id.clone()]),
+            })
+        };
+
         let db = self.db.clone();
         let policy = policy.clone();
         let chat = self.chat.clone();
@@ -405,7 +417,6 @@ impl PageServiceImpl {
         let retrieval = self.retrieval.clone();
         let config = self.config.clone();
         let auth = self.auth.clone();
-        let doc_query = doc_query.map(ToOwned::to_owned);
         let page_rowid = page_id.as_rowid()?;
         let debug_option = debug_option.map(|x| PageSectionRunDebugOptionInput {
             return_chat_completion_request: x.return_chat_completion_request,
@@ -496,7 +507,7 @@ impl PageServiceImpl {
                     None,
                     &existed_sections,
                     &page_title,
-                    code_source_id.clone(),
+                    code_source_id.as_deref(),
                     doc_query.as_ref(),
                     thread_messages.as_deref(),
                     db.clone(),
@@ -719,7 +730,7 @@ async fn generate_section_with_attachments(
     section_prompt: Option<String>,
     existing_sections: &[PageSection],
     page_title: &str,
-    code_source_id: Option<String>,
+    code_source_id: Option<&str>,
     doc_query: Option<&DocQueryInput>,
     thread_messages: Option<&[Message]>,
     db: DbConn,
@@ -735,6 +746,7 @@ async fn generate_section_with_attachments(
     let page_title = page_title.to_string();
     let thread_messages = thread_messages.map(ToOwned::to_owned);
     let context_info_helper = context.read(Some(&policy)).await?.helper();
+    let code_source_id = code_source_id.map(ToOwned::to_owned);
     let doc_query = doc_query.map(ToOwned::to_owned);
     let debug_chat_request = debug
         .map(|d| d.return_chat_completion_request)
@@ -781,31 +793,29 @@ async fn generate_section_with_attachments(
                 &config.code_search_params,
                 None,
             ).await?;
+            let debug_data = if debug_query_request {
+                Some(AttachmentCodeQueryDebugData {
+                    source_id: source_id.clone(),
+                    query: query_content,
+                })
+            } else {
+                None
+            };
 
             attachment.code = hits.iter().map(|x| x.doc.clone().into()).collect::<Vec<_>>();
 
+            let hits = hits.into_iter().map(|x| x.into()).collect::<Vec<_>>();
             if !hits.is_empty() {
-                let hits = hits.into_iter().map(|x| x.into()).collect::<Vec<_>>();
                 db.update_page_section_code_attachments(section_id, &attachment.code.iter().map(|c| c.into()).collect::<Vec<_>>()).await?;
-
-                let debug_data = if debug_query_request {
-                    Some(AttachmentCodeQueryDebugData {
-                        source_id: source_id.clone(),
-                        query: query_content,
-                    })
-                } else {
-                    None
-                };
-
-                yield Ok(SectionRunItem::PageSectionAttachmentCode(
-                    PageSectionAttachmentCode {
-                        id: section_id.as_id(),
-                        codes: hits,
-
-                        debug_data,
-                    }
-                ));
             }
+            yield Ok(SectionRunItem::PageSectionAttachmentCode(
+                PageSectionAttachmentCode {
+                    id: section_id.as_id(),
+                    codes: hits,
+
+                    debug_data,
+                }
+            ));
         }
 
         if let Some(doc_query) = doc_query {
@@ -820,39 +830,39 @@ async fn generate_section_with_attachments(
                     true
                 }
             }).collect::<Vec<_>>();
-
-            if !hits.is_empty() {
-                let hits = futures::future::join_all(hits.into_iter().map(|x| {
-                    let score = x.score;
-                    let doc = x.doc.clone();
-                    let auth = auth.clone();
-                    async move {
-                        AttachmentDocHit {
-                            score: score as f64,
-                            doc: attachment_doc_from_search(auth, doc).await,
-                        }
+            let debug_data = if debug_query_request {
+                Some(AttachmentDocQueryDebugData {
+                    source_ids: doc_query.source_ids.clone().unwrap_or_default(),
+                    query: doc_query.content.clone(),
+                })
+            } else {
+                None
+            };
+
+            let hits = futures::future::join_all(hits.into_iter().map(|x| {
+                let score = x.score;
+                let doc = x.doc.clone();
+                let auth = auth.clone();
+                async move {
+                    AttachmentDocHit {
+                        score: score as f64,
+                        doc: attachment_doc_from_search(auth, doc).await,
                     }
-                })).await;
+                }
+            })).await;
+            if !hits.is_empty() {
                 attachment.doc = hits.clone().into_iter().map(|x| x.doc).collect::<Vec<_>>();
                 db.update_page_section_doc_attachments(section_id, &attachment.doc.iter().map(|c| c.into()).collect::<Vec<_>>()).await?;
-
-                let debug_data = if debug_query_request {
-                    Some(AttachmentDocQueryDebugData {
-                        source_ids: doc_query.source_ids.clone().unwrap_or_default(),
-                        query: doc_query.content.clone(),
-                    })
-                } else {
-                    None
-                };
-                yield Ok(SectionRunItem::PageSectionAttachmentDoc(
-                    PageSectionAttachmentDoc {
-                        id: section_id.as_id(),
-                        doc: hits,
-
-                        debug_data,
-                    }
-                ));
             }
+
+            yield Ok(SectionRunItem::PageSectionAttachmentDoc(
+                PageSectionAttachmentDoc {
+                    id: section_id.as_id(),
+                    doc: hits,
+
+                    debug_data,
+                }
+            ));
         }
 
         // Generate section content