diff --git a/ee/tabby-schema/graphql/schema.graphql b/ee/tabby-schema/graphql/schema.graphql index 4cad4709e920..67c1b82377d1 100644 --- a/ee/tabby-schema/graphql/schema.graphql +++ b/ee/tabby-schema/graphql/schema.graphql @@ -489,12 +489,6 @@ type MessageAttachment { doc: [MessageAttachmentDoc!]! } -type MessageAttachmentAuthor { - id: String! - email: String - name: String -} - type MessageAttachmentClientCode { filepath: String startLine: Int @@ -518,7 +512,7 @@ type MessageAttachmentCodeScores { type MessageAttachmentIssueDoc { title: String! link: String! - author: MessageAttachmentAuthor + author: User body: String! closed: Boolean! } @@ -526,7 +520,7 @@ type MessageAttachmentIssueDoc { type MessageAttachmentPullDoc { title: String! link: String! - author: MessageAttachmentAuthor + author: User body: String! patch: String! merged: Boolean! diff --git a/ee/tabby-schema/src/dao.rs b/ee/tabby-schema/src/dao.rs index 73fa84bebbdc..2b2da7e41805 100644 --- a/ee/tabby-schema/src/dao.rs +++ b/ee/tabby-schema/src/dao.rs @@ -243,11 +243,7 @@ impl From for thread::MessageAttachmentDoc { thread::MessageAttachmentDoc::Issue(thread::MessageAttachmentIssueDoc { title: val.title, link: val.link, - author: val.author_user_id.map(|x| UserValue { - id: x, - email: None, - name: None, - }), + author: None, // will be filled in service layer body: val.body, closed: val.closed, }) @@ -256,11 +252,7 @@ impl From for thread::MessageAttachmentDoc { thread::MessageAttachmentDoc::Pull(thread::MessageAttachmentPullDoc { title: val.title, link: val.link, - author: val.author_user_id.map(|x| UserValue { - id: x, - email: None, - name: None, - }), + author: None, // will be filled in service layer body: val.body, patch: val.diff, merged: val.merged, @@ -284,7 +276,9 @@ impl From<&thread::MessageAttachmentDoc> for ThreadMessageAttachmentDoc { ThreadMessageAttachmentDoc::Issue(ThreadMessageAttachmentIssueDoc { title: val.title.clone(), link: val.link.clone(), - author_user_id: val.author.as_ref().map(|x| x.id.clone()), + author_user_id: val.author.as_ref().map(|x| match x { + UserValue::UserSecured(user) => user.id.to_string(), + }), body: val.body.clone(), closed: val.closed, }) @@ -293,7 +287,9 @@ impl From<&thread::MessageAttachmentDoc> for ThreadMessageAttachmentDoc { ThreadMessageAttachmentDoc::Pull(ThreadMessageAttachmentPullDoc { title: val.title.clone(), link: val.link.clone(), - author_user_id: val.author.as_ref().map(|x| x.id.clone()), + author_user_id: val.author.as_ref().map(|x| match x { + UserValue::UserSecured(user) => user.id.to_string(), + }), body: val.body.clone(), diff: val.patch.clone(), merged: val.merged, diff --git a/ee/tabby-schema/src/schema/thread/types.rs b/ee/tabby-schema/src/schema/thread/types.rs index 4587b00fdca6..960316a43c32 100644 --- a/ee/tabby-schema/src/schema/thread/types.rs +++ b/ee/tabby-schema/src/schema/thread/types.rs @@ -158,8 +158,8 @@ pub struct MessageAttachmentPullDoc { pub merged: bool, } -impl From for MessageAttachmentDoc { - fn from(doc: DocSearchDocument) -> Self { +impl MessageAttachmentDoc { + pub fn from_doc_search_document(doc: DocSearchDocument, author: Option) -> Self { match doc { DocSearchDocument::Web(web) => MessageAttachmentDoc::Web(MessageAttachmentWebDoc { title: web.title, @@ -170,7 +170,7 @@ impl From for MessageAttachmentDoc { MessageAttachmentDoc::Issue(MessageAttachmentIssueDoc { title: issue.title, link: issue.link, - author: None, + author: author, body: issue.body, closed: issue.closed, }) @@ -178,7 +178,7 @@ impl From for MessageAttachmentDoc { DocSearchDocument::Pull(pull) => MessageAttachmentDoc::Pull(MessageAttachmentPullDoc { title: pull.title, link: pull.link, - author: None, + author: author, body: pull.body, patch: pull.diff, merged: pull.merged, @@ -194,15 +194,6 @@ pub struct MessageDocSearchHit { pub score: f64, } -impl From for MessageDocSearchHit { - fn from(hit: DocSearchHit) -> Self { - Self { - doc: hit.doc.into(), - score: hit.score as f64, - } - } -} - #[derive(GraphQLObject)] #[graphql(context = Context)] pub struct Thread { diff --git a/ee/tabby-webserver/src/service/answer.rs b/ee/tabby-webserver/src/service/answer.rs index 55ec60c225f8..625c3f77deda 100644 --- a/ee/tabby-webserver/src/service/answer.rs +++ b/ee/tabby-webserver/src/service/answer.rs @@ -22,18 +22,19 @@ use tabby_common::{ CodeSearch, CodeSearchError, CodeSearchHit, CodeSearchParams, CodeSearchQuery, CodeSearchScores, }, - structured_doc::{DocSearch, DocSearchError, DocSearchHit}, + structured_doc::{DocSearch, DocSearchDocument, DocSearchError, DocSearchHit}, }, config::AnswerConfig, }; use tabby_inference::ChatCompletionStream; use tabby_schema::{ + auth::AuthenticationService, context::{ContextInfoHelper, ContextService}, policy::AccessPolicy, repository::{Repository, RepositoryService}, thread::{ self, CodeQueryInput, CodeSearchParamsOverrideInput, DocQueryInput, MessageAttachment, - MessageAttachmentDoc, ThreadAssistantMessageAttachmentsCode, + MessageAttachmentDoc, MessageDocSearchHit, ThreadAssistantMessageAttachmentsCode, ThreadAssistantMessageAttachmentsDoc, ThreadAssistantMessageContentDelta, ThreadRelevantQuestions, ThreadRunItem, ThreadRunOptionsInput, }, @@ -44,6 +45,7 @@ use crate::bail; pub struct AnswerService { config: AnswerConfig, + auth: Arc, chat: Arc, code: Arc, doc: Arc, @@ -55,6 +57,7 @@ pub struct AnswerService { impl AnswerService { fn new( config: &AnswerConfig, + auth: Arc, chat: Arc, code: Arc, doc: Arc, @@ -64,6 +67,7 @@ impl AnswerService { ) -> Self { Self { config: config.clone(), + auth, chat, code, doc, @@ -122,14 +126,24 @@ impl AnswerService { if let Some(doc_query) = options.doc_query.as_ref() { let hits = self.collect_relevant_docs(&context_info_helper, doc_query) .await; - attachment.doc = hits.iter() - .map(|x| x.doc.clone().into()) - .collect::>(); + attachment.doc = futures::future::join_all(hits.iter().map(|x| async { + Self::new_message_attachment_doc(self.auth.clone(), x.doc.clone()).await + })).await; debug!("doc content: {:?}: {:?}", doc_query.content, attachment.doc.len()); if !attachment.doc.is_empty() { - let hits = hits.into_iter().map(|x| x.into()).collect::>(); + let hits = futures::future::join_all(hits.into_iter().map(|x| { + let score = x.score; + let doc = x.doc.clone(); + let auth = self.auth.clone(); + async move { + MessageDocSearchHit { + score: score as f64, + doc: Self::new_message_attachment_doc(auth, doc).await, + } + } + })).await; yield Ok(ThreadRunItem::ThreadAssistantMessageAttachmentsDoc( ThreadAssistantMessageAttachmentsDoc { hits } )); @@ -201,6 +215,23 @@ impl AnswerService { Ok(Box::pin(s)) } + async fn new_message_attachment_doc( + auth: Arc, + doc: DocSearchDocument, + ) -> MessageAttachmentDoc { + let email = match &doc { + DocSearchDocument::Issue(issue) => issue.author_email.as_deref(), + DocSearchDocument::Pull(pull) => pull.author_email.as_deref(), + _ => None, + }; + let user = if let Some(email) = email { + auth.get_user_by_email(&email).await.ok().map(|x| x.into()) + } else { + None + }; + MessageAttachmentDoc::from_doc_search_document(doc, user) + } + async fn collect_relevant_code( &self, helper: &ContextInfoHelper, @@ -377,6 +408,7 @@ fn trim_bullet(s: &str) -> String { pub fn create( config: &AnswerConfig, + auth: Arc, chat: Arc, code: Arc, doc: Arc, @@ -384,7 +416,7 @@ pub fn create( serper: Option>, repository: Arc, ) -> AnswerService { - AnswerService::new(config, chat, code, doc, context, serper, repository) + AnswerService::new(config, auth, chat, code, doc, context, serper, repository) } fn convert_messages_to_chat_completion_request( diff --git a/ee/tabby-webserver/src/service/mod.rs b/ee/tabby-webserver/src/service/mod.rs index 387c637ad37f..e832bcb9eec6 100644 --- a/ee/tabby-webserver/src/service/mod.rs +++ b/ee/tabby-webserver/src/service/mod.rs @@ -1,17 +1,17 @@ mod access_policy; mod analytic; pub mod answer; -mod auth; +pub mod auth; pub mod background_job; pub mod context; -mod email; +pub mod email; pub mod event_logger; pub mod integration; pub mod job; -mod license; +pub mod license; mod preset_web_documents_data; pub mod repository; -mod setting; +pub mod setting; mod thread; mod user_event; mod user_group; @@ -58,10 +58,9 @@ use tabby_schema::{ AsID, AsRowid, CoreError, Result, ServiceLocator, }; -use self::{ - analytic::new_analytic_service, email::new_email_service, license::new_license_service, -}; +use self::analytic::new_analytic_service; use crate::rate_limit::UserRateLimiter; + struct ServerContext { db_conn: DbConn, mail: Arc, @@ -91,6 +90,7 @@ struct ServerContext { impl ServerContext { pub async fn new( logger: Arc, + auth: Arc, chat: Option>, completion: Option>, code: Arc, @@ -100,21 +100,13 @@ impl ServerContext { answer: Option>, context: Arc, web_documents: Arc, + mail: Arc, + license: Arc, + setting: Arc, db_conn: DbConn, embedding: Arc, ) -> Self { - let mail = Arc::new( - new_email_service(db_conn.clone()) - .await - .expect("failed to initialize mail service"), - ); - let license = Arc::new( - new_license_service(db_conn.clone()) - .await - .expect("failed to initialize license service"), - ); let user_event = Arc::new(user_event::create(db_conn.clone())); - let setting = Arc::new(setting::create(db_conn.clone())); let thread = Arc::new(thread::create(db_conn.clone(), answer.clone())); let user_group = Arc::new(user_group::create(db_conn.clone())); let access_policy = Arc::new(access_policy::create(db_conn.clone(), context.clone())); @@ -132,16 +124,11 @@ impl ServerContext { .await; Self { - mail: mail.clone(), + mail, embedding, chat, completion, - auth: Arc::new(auth::create( - db_conn.clone(), - mail, - license.clone(), - setting.clone(), - )), + auth, web_documents, thread, context, @@ -354,6 +341,7 @@ impl ServiceLocator for ArcServerContext { pub async fn create_service_locator( logger: Arc, + auth: Arc, chat: Option>, completion: Option>, code: Arc, @@ -363,12 +351,16 @@ pub async fn create_service_locator( answer: Option>, context: Arc, web_documents: Arc, + mail: Arc, + license: Arc, + setting: Arc, db: DbConn, embedding: Arc, ) -> Arc { Arc::new(ArcServerContext::new( ServerContext::new( logger, + auth, chat, completion, code, @@ -378,6 +370,9 @@ pub async fn create_service_locator( answer, context, web_documents, + mail, + license, + setting, db, embedding, ) diff --git a/ee/tabby-webserver/src/webserver.rs b/ee/tabby-webserver/src/webserver.rs index a7e360955908..38afc27289b9 100644 --- a/ee/tabby-webserver/src/webserver.rs +++ b/ee/tabby-webserver/src/webserver.rs @@ -23,6 +23,11 @@ use crate::{ }, }; +use crate::service::{ + auth::create as new_auth_service, email::new_email_service, license::new_license_service, + setting::create as new_setting_service, +}; + pub struct Webserver { db: DbConn, logger: Arc, @@ -88,9 +93,28 @@ impl Webserver { serper.is_some(), )); + let mail = Arc::new( + new_email_service(db.clone()) + .await + .expect("failed to initialize mail service"), + ); + let license = Arc::new( + new_license_service(db.clone()) + .await + .expect("failed to initialize license service"), + ); + let setting = Arc::new(new_setting_service(db.clone())); + let auth = Arc::new(new_auth_service( + db.clone(), + mail.clone(), + license.clone(), + setting.clone(), + )); + let answer = chat.as_ref().map(|chat| { Arc::new(crate::service::answer::create( &config.answer, + auth.clone(), chat.clone(), code.clone(), docsearch.clone(), @@ -102,6 +126,7 @@ impl Webserver { let ctx = create_service_locator( self.logger(), + auth, chat.clone(), completion.clone(), code.clone(), @@ -111,6 +136,9 @@ impl Webserver { answer.clone(), context.clone(), web_documents.clone(), + mail, + license, + setting, self.db.clone(), self.embedding.clone(), )