feat(context): add support for Pull Request in GitHub Context

crates/tabby-common/src/api/structured_doc.rs

@@ -20,6 +20,7 @@
 pub enum DocSearchDocument {
     Web(DocSearchWebDocument),
     Issue(DocSearchIssueDocument),
+    Pull(DocSearchPullRequest),
 }
 
 #[derive(Error, Debug)]
@@ -65,6 +66,15 @@
     pub closed: bool,
 }
 
+#[derive(Clone)]
+pub struct DocSearchPullRequest {
+    pub title: String,
+    pub link: String,
+    pub body: String,
+    pub patch: String,
+    pub state: String,
+}
+
 pub trait FromTantivyDocument {
     fn from_tantivy_document(doc: &TantivyDocument, chunk: &TantivyDocument) -> Option<Self>
     where
@@ -82,6 +92,9 @@
             }
             "issue" => DocSearchIssueDocument::from_tantivy_document(doc, chunk)
                 .map(DocSearchDocument::Issue),
+            "pull" => {
+                DocSearchPullRequest::from_tantivy_document(doc, chunk).map(DocSearchDocument::Pull)
+            }
             _ => None,
         }
     }
@@ -146,6 +159,44 @@
     }
 }
 
+impl FromTantivyDocument for DocSearchPullRequest {
+    fn from_tantivy_document(doc: &TantivyDocument, _: &TantivyDocument) -> Option<Self> {
+        let schema = IndexSchema::instance();
+        let title = get_json_text_field(
+            doc,
+            schema.field_attributes,
+            structured_doc::fields::pull::TITLE,
+        );
+        let link = get_json_text_field(
+            doc,
+            schema.field_attributes,
+            structured_doc::fields::pull::LINK,
+        );
+        let body = get_json_text_field(
+            doc,
+            schema.field_attributes,
+            structured_doc::fields::pull::BODY,
+        );
+        let diff = get_json_text_field(
+            doc,
+            schema.field_attributes,
+            structured_doc::fields::pull::PATCH,
+        );
+        let state = get_json_text_field(
+            doc,
+            schema.field_attributes,
+            structured_doc::fields::pull::STATE,
+        );
+        Some(Self {
+            title: title.into(),
+            link: link.into(),
+            body: body.into(),
+            patch: diff.into(),
+            state: state.into(),
+        })
+    }
+}
+
 fn get_json_field<'a>(
     doc: &'a TantivyDocument,
     field: schema::Field,

crates/tabby-common/src/index/structured_doc.rs

@@ -13,4 +13,12 @@ pub mod fields {
         pub const BODY: &str = "body";
         pub const CLOSED: &str = "closed";
     }
+
+    pub mod pull {
+        pub const TITLE: &str = "title";
+        pub const LINK: &str = "link";
+        pub const BODY: &str = "body";
+        pub const PATCH: &str = "patch";
+        pub const STATE: &str = "state";
+    }
 }

crates/tabby-index/src/lib.rs

@@ -17,7 +17,7 @@ pub mod public {
         code::CodeIndexer,
         structured_doc::public::{
             StructuredDoc, StructuredDocFields, StructuredDocIndexer, StructuredDocIssueFields,
-            StructuredDocWebFields,
+            StructuredDocPullRequestFields, StructuredDocPullRequestState, StructuredDocWebFields,
         },
     };
 

crates/tabby-index/src/structured_doc/public.rs

@@ -7,7 +7,12 @@ use tabby_common::index::corpus;
 use tabby_inference::Embedding;
 
 pub use super::types::{
-    issue::IssueDocument as StructuredDocIssueFields, web::WebDocument as StructuredDocWebFields,
+    issue::IssueDocument as StructuredDocIssueFields,
+    pull::{
+        PullRequest as StructuredDocPullRequestFields,
+        PullRequestState as StructuredDocPullRequestState,
+    },
+    web::WebDocument as StructuredDocWebFields,
     StructuredDoc, StructuredDocFields,
 };
 use super::{create_structured_doc_builder, types::BuildStructuredDoc};

crates/tabby-index/src/structured_doc/types.rs

@@ -1,4 +1,5 @@
 pub mod issue;
+pub mod pull;
 pub mod web;
 
 use std::sync::Arc;
@@ -21,13 +22,15 @@
         match &self.fields {
             StructuredDocFields::Web(web) => &web.link,
             StructuredDocFields::Issue(issue) => &issue.link,
+            StructuredDocFields::Pull(pull) => &pull.link,
         }
     }
 
     pub fn kind(&self) -> &'static str {
         match &self.fields {
             StructuredDocFields::Web(_) => "web",
             StructuredDocFields::Issue(_) => "issue",
+            StructuredDocFields::Pull(_) => "pull",
         }
     }
 }
@@ -55,6 +58,7 @@
 pub enum StructuredDocFields {
     Web(web::WebDocument),
     Issue(issue::IssueDocument),
+    Pull(pull::PullRequest),
 }
 
 #[async_trait]
@@ -63,13 +67,15 @@
         match &self.fields {
             StructuredDocFields::Web(doc) => doc.should_skip(),
             StructuredDocFields::Issue(doc) => doc.should_skip(),
+            StructuredDocFields::Pull(doc) => doc.should_skip(),
         }
     }
 
     async fn build_attributes(&self) -> serde_json::Value {
         match &self.fields {
             StructuredDocFields::Web(doc) => doc.build_attributes().await,
             StructuredDocFields::Issue(doc) => doc.build_attributes().await,
+            StructuredDocFields::Pull(doc) => doc.build_attributes().await,
         }
     }
 
@@ -80,6 +86,7 @@
         match &self.fields {
             StructuredDocFields::Web(doc) => doc.build_chunk_attributes(embedding).await,
             StructuredDocFields::Issue(doc) => doc.build_chunk_attributes(embedding).await,
+            StructuredDocFields::Pull(doc) => doc.build_chunk_attributes(embedding).await,
         }
     }
 }

crates/tabby-index/src/structured_doc/types/pull.rs

@@ -0,0 +1,69 @@
+use std::{fmt, sync::Arc};
+
+use async_stream::stream;
+use async_trait::async_trait;
+use futures::stream::BoxStream;
+use serde_json::json;
+use tabby_common::index::structured_doc::fields;
+use tabby_inference::Embedding;
+use tokio::task::JoinHandle;
+
+use super::{build_tokens, BuildStructuredDoc};
+
+pub struct PullRequest {
+    pub link: String,
+    pub title: String,
+    pub body: String,
+    pub patch: String,
+    pub state: String,
+}
+
+pub enum PullRequestState {
+    Open,
+    Closed,
+    Merged,
+}
+
+// Implement Display trait
+impl fmt::Display for PullRequestState {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        match self {
+            PullRequestState::Open => write!(f, "Open"),
+            PullRequestState::Closed => write!(f, "Closed"),
+            PullRequestState::Merged => write!(f, "Merged"),
+        }
+    }
+}
+
+#[async_trait]
+impl BuildStructuredDoc for PullRequest {
+    fn should_skip(&self) -> bool {
+        false
+    }
+
+    async fn build_attributes(&self) -> serde_json::Value {
+        json!({
+            fields::pull::LINK: self.link,
+            fields::pull::TITLE: self.title,
+            fields::pull::BODY: self.body,
+            fields::pull::PATCH: self.patch,
+            fields::pull::STATE: self.state,
+        })
+    }
+
+    async fn build_chunk_attributes(
+        &self,
+        embedding: Arc<dyn Embedding>,
+    ) -> BoxStream<JoinHandle<(Vec<String>, serde_json::Value)>> {
+        let text = format!("{}\n\n{}\n\n{}", self.title, self.body, self.patch);
+        let s = stream! {
+            yield tokio::spawn(async move {
+                let tokens = build_tokens(embedding, &text).await;
+                let chunk_attributes = json!({});
+                (tokens, chunk_attributes)
+            })
+        };
+
+        Box::pin(s)
+    }
+}

ee/tabby-db/src/lib.rs

@@ -15,8 +15,8 @@ pub use server_setting::ServerSettingDAO;
 use sqlx::{query, query_scalar, sqlite::SqliteQueryResult, Pool, Sqlite, SqlitePool};
 pub use threads::{
     ThreadDAO, ThreadMessageAttachmentClientCode, ThreadMessageAttachmentCode,
-    ThreadMessageAttachmentDoc, ThreadMessageAttachmentIssueDoc, ThreadMessageAttachmentWebDoc,
-    ThreadMessageDAO,
+    ThreadMessageAttachmentDoc, ThreadMessageAttachmentIssueDoc,
+    ThreadMessageAttachmentPullRequest, ThreadMessageAttachmentWebDoc, ThreadMessageDAO,
 };
 use tokio::sync::Mutex;
 use user_completions::UserCompletionDailyStatsDAO;

ee/tabby-db/src/threads.rs

@@ -36,6 +36,7 @@
 pub enum ThreadMessageAttachmentDoc {
     Web(ThreadMessageAttachmentWebDoc),
     Issue(ThreadMessageAttachmentIssueDoc),
+    Pull(ThreadMessageAttachmentPullRequest),
 }
 
 #[derive(Serialize, Deserialize)]
@@ -53,6 +54,15 @@
     pub closed: bool,
 }
 
+#[derive(Serialize, Deserialize)]
+pub struct ThreadMessageAttachmentPullRequest {
+    pub title: String,
+    pub link: String,
+    pub body: String,
+    pub diff: String,
+    pub state: String,
+}
+
 #[derive(Serialize, Deserialize)]
 pub struct ThreadMessageAttachmentCode {
     pub git_url: String,

ee/tabby-schema/graphql/schema.graphql

@@ -516,6 +516,14 @@ type MessageAttachmentIssueDoc {
   closed: Boolean!
 }
 
+type MessageAttachmentPullRequest {
+  title: String!
+  link: String!
+  body: String!
+  patch: String!
+  state: String!
+}
+
 type MessageAttachmentWebDoc {
   title: String!
   link: String!
@@ -900,7 +908,7 @@ type WebContextSource implements ContextSourceId & ContextSource {
   sourceName: String!
 }
 
-union MessageAttachmentDoc = MessageAttachmentWebDoc | MessageAttachmentIssueDoc
+union MessageAttachmentDoc = MessageAttachmentWebDoc | MessageAttachmentIssueDoc | MessageAttachmentPullRequest
 
 """
   Schema of thread run stream.

ee/tabby-schema/src/dao.rs

@@ -4,8 +4,9 @@
 use tabby_db::{
     EmailSettingDAO, IntegrationDAO, InvitationDAO, JobRunDAO, OAuthCredentialDAO,
     ServerSettingDAO, ThreadDAO, ThreadMessageAttachmentClientCode, ThreadMessageAttachmentCode,
-    ThreadMessageAttachmentDoc, ThreadMessageAttachmentIssueDoc, ThreadMessageAttachmentWebDoc,
-    ThreadMessageDAO, UserEventDAO,
+    ThreadMessageAttachmentDoc, ThreadMessageAttachmentIssueDoc,
+    ThreadMessageAttachmentPullRequest, ThreadMessageAttachmentWebDoc, ThreadMessageDAO,
+    UserEventDAO,
 };
 
 use crate::{
@@ -246,6 +247,15 @@
                     closed: val.closed,
                 })
             }
+            ThreadMessageAttachmentDoc::Pull(val) => {
+                thread::MessageAttachmentDoc::Pull(thread::MessageAttachmentPullRequest {
+                    title: val.title,
+                    link: val.link,
+                    body: val.body,
+                    patch: val.diff,
+                    state: val.state,
+                })
+            }
         }
     }
 }
@@ -268,6 +278,15 @@
                     closed: val.closed,
                 })
             }
+            thread::MessageAttachmentDoc::Pull(val) => {
+                ThreadMessageAttachmentDoc::Pull(ThreadMessageAttachmentPullRequest {
+                    title: val.title.clone(),
+                    link: val.link.clone(),
+                    body: val.body.clone(),
+                    diff: val.patch.clone(),
+                    state: val.state.clone(),
+                })
+            }
         }
     }
 }

ee/tabby-schema/src/schema/thread/types.rs

@@ -125,6 +125,7 @@
 pub enum MessageAttachmentDoc {
     Web(MessageAttachmentWebDoc),
     Issue(MessageAttachmentIssueDoc),
+    Pull(MessageAttachmentPullRequest),
 }
 
 #[derive(GraphQLObject, Clone)]
@@ -142,6 +143,15 @@
     pub closed: bool,
 }
 
+#[derive(GraphQLObject, Clone)]
+pub struct MessageAttachmentPullRequest {
+    pub title: String,
+    pub link: String,
+    pub body: String,
+    pub patch: String,
+    pub state: String,
+}
+
 impl From<DocSearchDocument> for MessageAttachmentDoc {
     fn from(doc: DocSearchDocument) -> Self {
         match doc {
@@ -158,6 +168,15 @@
                     closed: issue.closed,
                 })
             }
+            DocSearchDocument::Pull(pull) => {
+                MessageAttachmentDoc::Pull(MessageAttachmentPullRequest {
+                    title: pull.title,
+                    link: pull.link,
+                    body: pull.body,
+                    patch: pull.patch,
+                    state: pull.state,
+                })
+            }
         }
     }
 }