Signed-off-by: -LAN- <laipz8200@outlook.com> Co-authored-by: -LAN- <laipz8200@outlook.com> Co-authored-by: Claude <noreply@anthropic.com>tags/1.7.1
| click.echo(click.style("Passwords do not match.", fg="red")) | click.echo(click.style("Passwords do not match.", fg="red")) | ||||
| return | return | ||||
| account = db.session.query(Account).filter(Account.email == email).one_or_none() | |||||
| account = db.session.query(Account).where(Account.email == email).one_or_none() | |||||
| if not account: | if not account: | ||||
| click.echo(click.style("Account not found for email: {}".format(email), fg="red")) | click.echo(click.style("Account not found for email: {}".format(email), fg="red")) | ||||
| click.echo(click.style("New emails do not match.", fg="red")) | click.echo(click.style("New emails do not match.", fg="red")) | ||||
| return | return | ||||
| account = db.session.query(Account).filter(Account.email == email).one_or_none() | |||||
| account = db.session.query(Account).where(Account.email == email).one_or_none() | |||||
| if not account: | if not account: | ||||
| click.echo(click.style("Account not found for email: {}".format(email), fg="red")) | click.echo(click.style("Account not found for email: {}".format(email), fg="red")) | ||||
| tenant.encrypt_public_key = generate_key_pair(tenant.id) | tenant.encrypt_public_key = generate_key_pair(tenant.id) | ||||
| db.session.query(Provider).filter(Provider.provider_type == "custom", Provider.tenant_id == tenant.id).delete() | |||||
| db.session.query(ProviderModel).filter(ProviderModel.tenant_id == tenant.id).delete() | |||||
| db.session.query(Provider).where(Provider.provider_type == "custom", Provider.tenant_id == tenant.id).delete() | |||||
| db.session.query(ProviderModel).where(ProviderModel.tenant_id == tenant.id).delete() | |||||
| db.session.commit() | db.session.commit() | ||||
| click.echo( | click.echo( | ||||
| per_page = 50 | per_page = 50 | ||||
| apps = ( | apps = ( | ||||
| db.session.query(App) | db.session.query(App) | ||||
| .filter(App.status == "normal") | |||||
| .where(App.status == "normal") | |||||
| .order_by(App.created_at.desc()) | .order_by(App.created_at.desc()) | ||||
| .limit(per_page) | .limit(per_page) | ||||
| .offset((page - 1) * per_page) | .offset((page - 1) * per_page) | ||||
| try: | try: | ||||
| click.echo("Creating app annotation index: {}".format(app.id)) | click.echo("Creating app annotation index: {}".format(app.id)) | ||||
| app_annotation_setting = ( | app_annotation_setting = ( | ||||
| db.session.query(AppAnnotationSetting).filter(AppAnnotationSetting.app_id == app.id).first() | |||||
| db.session.query(AppAnnotationSetting).where(AppAnnotationSetting.app_id == app.id).first() | |||||
| ) | ) | ||||
| if not app_annotation_setting: | if not app_annotation_setting: | ||||
| # get dataset_collection_binding info | # get dataset_collection_binding info | ||||
| dataset_collection_binding = ( | dataset_collection_binding = ( | ||||
| db.session.query(DatasetCollectionBinding) | db.session.query(DatasetCollectionBinding) | ||||
| .filter(DatasetCollectionBinding.id == app_annotation_setting.collection_binding_id) | |||||
| .where(DatasetCollectionBinding.id == app_annotation_setting.collection_binding_id) | |||||
| .first() | .first() | ||||
| ) | ) | ||||
| if not dataset_collection_binding: | if not dataset_collection_binding: | ||||
| click.echo("App annotation collection binding not found: {}".format(app.id)) | click.echo("App annotation collection binding not found: {}".format(app.id)) | ||||
| continue | continue | ||||
| annotations = db.session.query(MessageAnnotation).filter(MessageAnnotation.app_id == app.id).all() | |||||
| annotations = db.session.query(MessageAnnotation).where(MessageAnnotation.app_id == app.id).all() | |||||
| dataset = Dataset( | dataset = Dataset( | ||||
| id=app.id, | id=app.id, | ||||
| tenant_id=app.tenant_id, | tenant_id=app.tenant_id, | ||||
| while True: | while True: | ||||
| try: | try: | ||||
| stmt = ( | stmt = ( | ||||
| select(Dataset).filter(Dataset.indexing_technique == "high_quality").order_by(Dataset.created_at.desc()) | |||||
| select(Dataset).where(Dataset.indexing_technique == "high_quality").order_by(Dataset.created_at.desc()) | |||||
| ) | ) | ||||
| datasets = db.paginate(select=stmt, page=page, per_page=50, max_per_page=50, error_out=False) | datasets = db.paginate(select=stmt, page=page, per_page=50, max_per_page=50, error_out=False) | ||||
| if dataset.collection_binding_id: | if dataset.collection_binding_id: | ||||
| dataset_collection_binding = ( | dataset_collection_binding = ( | ||||
| db.session.query(DatasetCollectionBinding) | db.session.query(DatasetCollectionBinding) | ||||
| .filter(DatasetCollectionBinding.id == dataset.collection_binding_id) | |||||
| .where(DatasetCollectionBinding.id == dataset.collection_binding_id) | |||||
| .one_or_none() | .one_or_none() | ||||
| ) | ) | ||||
| if dataset_collection_binding: | if dataset_collection_binding: | ||||
| dataset_documents = ( | dataset_documents = ( | ||||
| db.session.query(DatasetDocument) | db.session.query(DatasetDocument) | ||||
| .filter( | |||||
| .where( | |||||
| DatasetDocument.dataset_id == dataset.id, | DatasetDocument.dataset_id == dataset.id, | ||||
| DatasetDocument.indexing_status == "completed", | DatasetDocument.indexing_status == "completed", | ||||
| DatasetDocument.enabled == True, | DatasetDocument.enabled == True, | ||||
| for dataset_document in dataset_documents: | for dataset_document in dataset_documents: | ||||
| segments = ( | segments = ( | ||||
| db.session.query(DocumentSegment) | db.session.query(DocumentSegment) | ||||
| .filter( | |||||
| .where( | |||||
| DocumentSegment.document_id == dataset_document.id, | DocumentSegment.document_id == dataset_document.id, | ||||
| DocumentSegment.status == "completed", | DocumentSegment.status == "completed", | ||||
| DocumentSegment.enabled == True, | DocumentSegment.enabled == True, | ||||
| app_id = str(i.id) | app_id = str(i.id) | ||||
| if app_id not in proceeded_app_ids: | if app_id not in proceeded_app_ids: | ||||
| proceeded_app_ids.append(app_id) | proceeded_app_ids.append(app_id) | ||||
| app = db.session.query(App).filter(App.id == app_id).first() | |||||
| app = db.session.query(App).where(App.id == app_id).first() | |||||
| if app is not None: | if app is not None: | ||||
| apps.append(app) | apps.append(app) | ||||
| db.session.commit() | db.session.commit() | ||||
| # update conversation mode to agent | # update conversation mode to agent | ||||
| db.session.query(Conversation).filter(Conversation.app_id == app.id).update( | |||||
| db.session.query(Conversation).where(Conversation.app_id == app.id).update( | |||||
| {Conversation.mode: AppMode.AGENT_CHAT.value} | {Conversation.mode: AppMode.AGENT_CHAT.value} | ||||
| ) | ) | ||||
| try: | try: | ||||
| stmt = ( | stmt = ( | ||||
| select(DatasetDocument) | select(DatasetDocument) | ||||
| .filter(DatasetDocument.doc_metadata.is_not(None)) | |||||
| .where(DatasetDocument.doc_metadata.is_not(None)) | |||||
| .order_by(DatasetDocument.created_at.desc()) | .order_by(DatasetDocument.created_at.desc()) | ||||
| ) | ) | ||||
| documents = db.paginate(select=stmt, page=page, per_page=50, max_per_page=50, error_out=False) | documents = db.paginate(select=stmt, page=page, per_page=50, max_per_page=50, error_out=False) | ||||
| else: | else: | ||||
| dataset_metadata = ( | dataset_metadata = ( | ||||
| db.session.query(DatasetMetadata) | db.session.query(DatasetMetadata) | ||||
| .filter(DatasetMetadata.dataset_id == document.dataset_id, DatasetMetadata.name == key) | |||||
| .where(DatasetMetadata.dataset_id == document.dataset_id, DatasetMetadata.name == key) | |||||
| .first() | .first() | ||||
| ) | ) | ||||
| if not dataset_metadata: | if not dataset_metadata: | ||||
| else: | else: | ||||
| dataset_metadata_binding = ( | dataset_metadata_binding = ( | ||||
| db.session.query(DatasetMetadataBinding) # type: ignore | db.session.query(DatasetMetadataBinding) # type: ignore | ||||
| .filter( | |||||
| .where( | |||||
| DatasetMetadataBinding.dataset_id == document.dataset_id, | DatasetMetadataBinding.dataset_id == document.dataset_id, | ||||
| DatasetMetadataBinding.document_id == document.id, | DatasetMetadataBinding.document_id == document.id, | ||||
| DatasetMetadataBinding.metadata_id == dataset_metadata.id, | DatasetMetadataBinding.metadata_id == dataset_metadata.id, | ||||
| continue | continue | ||||
| try: | try: | ||||
| app = db.session.query(App).filter(App.id == app_id).first() | |||||
| app = db.session.query(App).where(App.id == app_id).first() | |||||
| if not app: | if not app: | ||||
| print(f"App {app_id} not found") | print(f"App {app_id} not found") | ||||
| continue | continue |
| parser.add_argument("position", type=int, required=True, nullable=False, location="json") | parser.add_argument("position", type=int, required=True, nullable=False, location="json") | ||||
| args = parser.parse_args() | args = parser.parse_args() | ||||
| app = db.session.execute(select(App).filter(App.id == args["app_id"])).scalar_one_or_none() | |||||
| app = db.session.execute(select(App).where(App.id == args["app_id"])).scalar_one_or_none() | |||||
| if not app: | if not app: | ||||
| raise NotFound(f"App '{args['app_id']}' is not found") | raise NotFound(f"App '{args['app_id']}' is not found") | ||||
| with Session(db.engine) as session: | with Session(db.engine) as session: | ||||
| recommended_app = session.execute( | recommended_app = session.execute( | ||||
| select(RecommendedApp).filter(RecommendedApp.app_id == args["app_id"]) | |||||
| select(RecommendedApp).where(RecommendedApp.app_id == args["app_id"]) | |||||
| ).scalar_one_or_none() | ).scalar_one_or_none() | ||||
| if not recommended_app: | if not recommended_app: | ||||
| def delete(self, app_id): | def delete(self, app_id): | ||||
| with Session(db.engine) as session: | with Session(db.engine) as session: | ||||
| recommended_app = session.execute( | recommended_app = session.execute( | ||||
| select(RecommendedApp).filter(RecommendedApp.app_id == str(app_id)) | |||||
| select(RecommendedApp).where(RecommendedApp.app_id == str(app_id)) | |||||
| ).scalar_one_or_none() | ).scalar_one_or_none() | ||||
| if not recommended_app: | if not recommended_app: | ||||
| return {"result": "success"}, 204 | return {"result": "success"}, 204 | ||||
| with Session(db.engine) as session: | with Session(db.engine) as session: | ||||
| app = session.execute(select(App).filter(App.id == recommended_app.app_id)).scalar_one_or_none() | |||||
| app = session.execute(select(App).where(App.id == recommended_app.app_id)).scalar_one_or_none() | |||||
| if app: | if app: | ||||
| app.is_public = False | app.is_public = False | ||||
| with Session(db.engine) as session: | with Session(db.engine) as session: | ||||
| installed_apps = session.execute( | installed_apps = session.execute( | ||||
| select(InstalledApp).filter( | |||||
| select(InstalledApp).where( | |||||
| InstalledApp.app_id == recommended_app.app_id, | InstalledApp.app_id == recommended_app.app_id, | ||||
| InstalledApp.tenant_id != InstalledApp.app_owner_tenant_id, | InstalledApp.tenant_id != InstalledApp.app_owner_tenant_id, | ||||
| ) | ) |
| _get_resource(resource_id, current_user.current_tenant_id, self.resource_model) | _get_resource(resource_id, current_user.current_tenant_id, self.resource_model) | ||||
| keys = ( | keys = ( | ||||
| db.session.query(ApiToken) | db.session.query(ApiToken) | ||||
| .filter(ApiToken.type == self.resource_type, getattr(ApiToken, self.resource_id_field) == resource_id) | |||||
| .where(ApiToken.type == self.resource_type, getattr(ApiToken, self.resource_id_field) == resource_id) | |||||
| .all() | .all() | ||||
| ) | ) | ||||
| return {"items": keys} | return {"items": keys} | ||||
| current_key_count = ( | current_key_count = ( | ||||
| db.session.query(ApiToken) | db.session.query(ApiToken) | ||||
| .filter(ApiToken.type == self.resource_type, getattr(ApiToken, self.resource_id_field) == resource_id) | |||||
| .where(ApiToken.type == self.resource_type, getattr(ApiToken, self.resource_id_field) == resource_id) | |||||
| .count() | .count() | ||||
| ) | ) | ||||
| key = ( | key = ( | ||||
| db.session.query(ApiToken) | db.session.query(ApiToken) | ||||
| .filter( | |||||
| .where( | |||||
| getattr(ApiToken, self.resource_id_field) == resource_id, | getattr(ApiToken, self.resource_id_field) == resource_id, | ||||
| ApiToken.type == self.resource_type, | ApiToken.type == self.resource_type, | ||||
| ApiToken.id == api_key_id, | ApiToken.id == api_key_id, | ||||
| if key is None: | if key is None: | ||||
| flask_restful.abort(404, message="API key not found") | flask_restful.abort(404, message="API key not found") | ||||
| db.session.query(ApiToken).filter(ApiToken.id == api_key_id).delete() | |||||
| db.session.query(ApiToken).where(ApiToken.id == api_key_id).delete() | |||||
| db.session.commit() | db.session.commit() | ||||
| return {"result": "success"}, 204 | return {"result": "success"}, 204 |
| query = db.select(Conversation).where(Conversation.app_id == app_model.id, Conversation.mode == "completion") | query = db.select(Conversation).where(Conversation.app_id == app_model.id, Conversation.mode == "completion") | ||||
| if args["keyword"]: | if args["keyword"]: | ||||
| query = query.join(Message, Message.conversation_id == Conversation.id).filter( | |||||
| query = query.join(Message, Message.conversation_id == Conversation.id).where( | |||||
| or_( | or_( | ||||
| Message.query.ilike("%{}%".format(args["keyword"])), | Message.query.ilike("%{}%".format(args["keyword"])), | ||||
| Message.answer.ilike("%{}%".format(args["keyword"])), | Message.answer.ilike("%{}%".format(args["keyword"])), | ||||
| conversation = ( | conversation = ( | ||||
| db.session.query(Conversation) | db.session.query(Conversation) | ||||
| .filter(Conversation.id == conversation_id, Conversation.app_id == app_model.id) | |||||
| .where(Conversation.id == conversation_id, Conversation.app_id == app_model.id) | |||||
| .first() | .first() | ||||
| ) | ) | ||||
| Message.conversation_id == Conversation.id, | Message.conversation_id == Conversation.id, | ||||
| ) | ) | ||||
| .join(subquery, subquery.c.conversation_id == Conversation.id) | .join(subquery, subquery.c.conversation_id == Conversation.id) | ||||
| .filter( | |||||
| .where( | |||||
| or_( | or_( | ||||
| Message.query.ilike(keyword_filter), | Message.query.ilike(keyword_filter), | ||||
| Message.answer.ilike(keyword_filter), | Message.answer.ilike(keyword_filter), | ||||
| conversation = ( | conversation = ( | ||||
| db.session.query(Conversation) | db.session.query(Conversation) | ||||
| .filter(Conversation.id == conversation_id, Conversation.app_id == app_model.id) | |||||
| .where(Conversation.id == conversation_id, Conversation.app_id == app_model.id) | |||||
| .first() | .first() | ||||
| ) | ) | ||||
| def _get_conversation(app_model, conversation_id): | def _get_conversation(app_model, conversation_id): | ||||
| conversation = ( | conversation = ( | ||||
| db.session.query(Conversation) | db.session.query(Conversation) | ||||
| .filter(Conversation.id == conversation_id, Conversation.app_id == app_model.id) | |||||
| .where(Conversation.id == conversation_id, Conversation.app_id == app_model.id) | |||||
| .first() | .first() | ||||
| ) | ) | ||||
| @get_app_model | @get_app_model | ||||
| @marshal_with(app_server_fields) | @marshal_with(app_server_fields) | ||||
| def get(self, app_model): | def get(self, app_model): | ||||
| server = db.session.query(AppMCPServer).filter(AppMCPServer.app_id == app_model.id).first() | |||||
| server = db.session.query(AppMCPServer).where(AppMCPServer.app_id == app_model.id).first() | |||||
| return server | return server | ||||
| @setup_required | @setup_required | ||||
| parser.add_argument("parameters", type=dict, required=True, location="json") | parser.add_argument("parameters", type=dict, required=True, location="json") | ||||
| parser.add_argument("status", type=str, required=False, location="json") | parser.add_argument("status", type=str, required=False, location="json") | ||||
| args = parser.parse_args() | args = parser.parse_args() | ||||
| server = db.session.query(AppMCPServer).filter(AppMCPServer.id == args["id"]).first() | |||||
| server = db.session.query(AppMCPServer).where(AppMCPServer.id == args["id"]).first() | |||||
| if not server: | if not server: | ||||
| raise NotFound() | raise NotFound() | ||||
| raise NotFound() | raise NotFound() | ||||
| server = ( | server = ( | ||||
| db.session.query(AppMCPServer) | db.session.query(AppMCPServer) | ||||
| .filter(AppMCPServer.id == server_id) | |||||
| .filter(AppMCPServer.tenant_id == current_user.current_tenant_id) | |||||
| .where(AppMCPServer.id == server_id) | |||||
| .where(AppMCPServer.tenant_id == current_user.current_tenant_id) | |||||
| .first() | .first() | ||||
| ) | ) | ||||
| if not server: | if not server: |
| conversation = ( | conversation = ( | ||||
| db.session.query(Conversation) | db.session.query(Conversation) | ||||
| .filter(Conversation.id == args["conversation_id"], Conversation.app_id == app_model.id) | |||||
| .where(Conversation.id == args["conversation_id"], Conversation.app_id == app_model.id) | |||||
| .first() | .first() | ||||
| ) | ) | ||||
| if args["first_id"]: | if args["first_id"]: | ||||
| first_message = ( | first_message = ( | ||||
| db.session.query(Message) | db.session.query(Message) | ||||
| .filter(Message.conversation_id == conversation.id, Message.id == args["first_id"]) | |||||
| .where(Message.conversation_id == conversation.id, Message.id == args["first_id"]) | |||||
| .first() | .first() | ||||
| ) | ) | ||||
| history_messages = ( | history_messages = ( | ||||
| db.session.query(Message) | db.session.query(Message) | ||||
| .filter( | |||||
| .where( | |||||
| Message.conversation_id == conversation.id, | Message.conversation_id == conversation.id, | ||||
| Message.created_at < first_message.created_at, | Message.created_at < first_message.created_at, | ||||
| Message.id != first_message.id, | Message.id != first_message.id, | ||||
| else: | else: | ||||
| history_messages = ( | history_messages = ( | ||||
| db.session.query(Message) | db.session.query(Message) | ||||
| .filter(Message.conversation_id == conversation.id) | |||||
| .where(Message.conversation_id == conversation.id) | |||||
| .order_by(Message.created_at.desc()) | .order_by(Message.created_at.desc()) | ||||
| .limit(args["limit"]) | .limit(args["limit"]) | ||||
| .all() | .all() | ||||
| current_page_first_message = history_messages[-1] | current_page_first_message = history_messages[-1] | ||||
| rest_count = ( | rest_count = ( | ||||
| db.session.query(Message) | db.session.query(Message) | ||||
| .filter( | |||||
| .where( | |||||
| Message.conversation_id == conversation.id, | Message.conversation_id == conversation.id, | ||||
| Message.created_at < current_page_first_message.created_at, | Message.created_at < current_page_first_message.created_at, | ||||
| Message.id != current_page_first_message.id, | Message.id != current_page_first_message.id, | ||||
| @account_initialization_required | @account_initialization_required | ||||
| @get_app_model | @get_app_model | ||||
| def get(self, app_model): | def get(self, app_model): | ||||
| count = db.session.query(MessageAnnotation).filter(MessageAnnotation.app_id == app_model.id).count() | |||||
| count = db.session.query(MessageAnnotation).where(MessageAnnotation.app_id == app_model.id).count() | |||||
| return {"count": count} | return {"count": count} | ||||
| def get(self, app_model, message_id): | def get(self, app_model, message_id): | ||||
| message_id = str(message_id) | message_id = str(message_id) | ||||
| message = db.session.query(Message).filter(Message.id == message_id, Message.app_id == app_model.id).first() | |||||
| message = db.session.query(Message).where(Message.id == message_id, Message.app_id == app_model.id).first() | |||||
| if not message: | if not message: | ||||
| raise NotFound("Message Not Exists.") | raise NotFound("Message Not Exists.") |
| if app_model.mode == AppMode.AGENT_CHAT.value or app_model.is_agent: | if app_model.mode == AppMode.AGENT_CHAT.value or app_model.is_agent: | ||||
| # get original app model config | # get original app model config | ||||
| original_app_model_config = ( | original_app_model_config = ( | ||||
| db.session.query(AppModelConfig).filter(AppModelConfig.id == app_model.app_model_config_id).first() | |||||
| db.session.query(AppModelConfig).where(AppModelConfig.id == app_model.app_model_config_id).first() | |||||
| ) | ) | ||||
| if original_app_model_config is None: | if original_app_model_config is None: | ||||
| raise ValueError("Original app model config not found") | raise ValueError("Original app model config not found") |
| if not current_user.is_editor: | if not current_user.is_editor: | ||||
| raise Forbidden() | raise Forbidden() | ||||
| site = db.session.query(Site).filter(Site.app_id == app_model.id).first() | |||||
| site = db.session.query(Site).where(Site.app_id == app_model.id).first() | |||||
| if not site: | if not site: | ||||
| raise NotFound | raise NotFound | ||||
| if not current_user.is_admin_or_owner: | if not current_user.is_admin_or_owner: | ||||
| raise Forbidden() | raise Forbidden() | ||||
| site = db.session.query(Site).filter(Site.app_id == app_model.id).first() | |||||
| site = db.session.query(Site).where(Site.app_id == app_model.id).first() | |||||
| if not site: | if not site: | ||||
| raise NotFound | raise NotFound |
| def _load_app_model(app_id: str) -> Optional[App]: | def _load_app_model(app_id: str) -> Optional[App]: | ||||
| app_model = ( | app_model = ( | ||||
| db.session.query(App) | db.session.query(App) | ||||
| .filter(App.id == app_id, App.tenant_id == current_user.current_tenant_id, App.status == "normal") | |||||
| .where(App.id == app_id, App.tenant_id == current_user.current_tenant_id, App.status == "normal") | |||||
| .first() | .first() | ||||
| ) | ) | ||||
| return app_model | return app_model |
| # get workspace data source integrates | # get workspace data source integrates | ||||
| data_source_integrates = ( | data_source_integrates = ( | ||||
| db.session.query(DataSourceOauthBinding) | db.session.query(DataSourceOauthBinding) | ||||
| .filter( | |||||
| .where( | |||||
| DataSourceOauthBinding.tenant_id == current_user.current_tenant_id, | DataSourceOauthBinding.tenant_id == current_user.current_tenant_id, | ||||
| DataSourceOauthBinding.disabled == False, | DataSourceOauthBinding.disabled == False, | ||||
| ) | ) | ||||
| page_id = str(page_id) | page_id = str(page_id) | ||||
| with Session(db.engine) as session: | with Session(db.engine) as session: | ||||
| data_source_binding = session.execute( | data_source_binding = session.execute( | ||||
| select(DataSourceOauthBinding).filter( | |||||
| select(DataSourceOauthBinding).where( | |||||
| db.and_( | db.and_( | ||||
| DataSourceOauthBinding.tenant_id == current_user.current_tenant_id, | DataSourceOauthBinding.tenant_id == current_user.current_tenant_id, | ||||
| DataSourceOauthBinding.provider == "notion", | DataSourceOauthBinding.provider == "notion", |
| file_ids = args["info_list"]["file_info_list"]["file_ids"] | file_ids = args["info_list"]["file_info_list"]["file_ids"] | ||||
| file_details = ( | file_details = ( | ||||
| db.session.query(UploadFile) | db.session.query(UploadFile) | ||||
| .filter(UploadFile.tenant_id == current_user.current_tenant_id, UploadFile.id.in_(file_ids)) | |||||
| .where(UploadFile.tenant_id == current_user.current_tenant_id, UploadFile.id.in_(file_ids)) | |||||
| .all() | .all() | ||||
| ) | ) | ||||
| dataset_id = str(dataset_id) | dataset_id = str(dataset_id) | ||||
| documents = ( | documents = ( | ||||
| db.session.query(Document) | db.session.query(Document) | ||||
| .filter(Document.dataset_id == dataset_id, Document.tenant_id == current_user.current_tenant_id) | |||||
| .where(Document.dataset_id == dataset_id, Document.tenant_id == current_user.current_tenant_id) | |||||
| .all() | .all() | ||||
| ) | ) | ||||
| documents_status = [] | documents_status = [] | ||||
| for document in documents: | for document in documents: | ||||
| completed_segments = ( | completed_segments = ( | ||||
| db.session.query(DocumentSegment) | db.session.query(DocumentSegment) | ||||
| .filter( | |||||
| .where( | |||||
| DocumentSegment.completed_at.isnot(None), | DocumentSegment.completed_at.isnot(None), | ||||
| DocumentSegment.document_id == str(document.id), | DocumentSegment.document_id == str(document.id), | ||||
| DocumentSegment.status != "re_segment", | DocumentSegment.status != "re_segment", | ||||
| ) | ) | ||||
| total_segments = ( | total_segments = ( | ||||
| db.session.query(DocumentSegment) | db.session.query(DocumentSegment) | ||||
| .filter(DocumentSegment.document_id == str(document.id), DocumentSegment.status != "re_segment") | |||||
| .where(DocumentSegment.document_id == str(document.id), DocumentSegment.status != "re_segment") | |||||
| .count() | .count() | ||||
| ) | ) | ||||
| # Create a dictionary with document attributes and additional fields | # Create a dictionary with document attributes and additional fields | ||||
| def get(self): | def get(self): | ||||
| keys = ( | keys = ( | ||||
| db.session.query(ApiToken) | db.session.query(ApiToken) | ||||
| .filter(ApiToken.type == self.resource_type, ApiToken.tenant_id == current_user.current_tenant_id) | |||||
| .where(ApiToken.type == self.resource_type, ApiToken.tenant_id == current_user.current_tenant_id) | |||||
| .all() | .all() | ||||
| ) | ) | ||||
| return {"items": keys} | return {"items": keys} | ||||
| current_key_count = ( | current_key_count = ( | ||||
| db.session.query(ApiToken) | db.session.query(ApiToken) | ||||
| .filter(ApiToken.type == self.resource_type, ApiToken.tenant_id == current_user.current_tenant_id) | |||||
| .where(ApiToken.type == self.resource_type, ApiToken.tenant_id == current_user.current_tenant_id) | |||||
| .count() | .count() | ||||
| ) | ) | ||||
| key = ( | key = ( | ||||
| db.session.query(ApiToken) | db.session.query(ApiToken) | ||||
| .filter( | |||||
| .where( | |||||
| ApiToken.tenant_id == current_user.current_tenant_id, | ApiToken.tenant_id == current_user.current_tenant_id, | ||||
| ApiToken.type == self.resource_type, | ApiToken.type == self.resource_type, | ||||
| ApiToken.id == api_key_id, | ApiToken.id == api_key_id, | ||||
| if key is None: | if key is None: | ||||
| flask_restful.abort(404, message="API key not found") | flask_restful.abort(404, message="API key not found") | ||||
| db.session.query(ApiToken).filter(ApiToken.id == api_key_id).delete() | |||||
| db.session.query(ApiToken).where(ApiToken.id == api_key_id).delete() | |||||
| db.session.commit() | db.session.commit() | ||||
| return {"result": "success"}, 204 | return {"result": "success"}, 204 |
| # get the latest process rule | # get the latest process rule | ||||
| dataset_process_rule = ( | dataset_process_rule = ( | ||||
| db.session.query(DatasetProcessRule) | db.session.query(DatasetProcessRule) | ||||
| .filter(DatasetProcessRule.dataset_id == document.dataset_id) | |||||
| .where(DatasetProcessRule.dataset_id == document.dataset_id) | |||||
| .order_by(DatasetProcessRule.created_at.desc()) | .order_by(DatasetProcessRule.created_at.desc()) | ||||
| .limit(1) | .limit(1) | ||||
| .one_or_none() | .one_or_none() | ||||
| if search: | if search: | ||||
| search = f"%{search}%" | search = f"%{search}%" | ||||
| query = query.filter(Document.name.like(search)) | |||||
| query = query.where(Document.name.like(search)) | |||||
| if sort.startswith("-"): | if sort.startswith("-"): | ||||
| sort_logic = desc | sort_logic = desc | ||||
| for document in documents: | for document in documents: | ||||
| completed_segments = ( | completed_segments = ( | ||||
| db.session.query(DocumentSegment) | db.session.query(DocumentSegment) | ||||
| .filter( | |||||
| .where( | |||||
| DocumentSegment.completed_at.isnot(None), | DocumentSegment.completed_at.isnot(None), | ||||
| DocumentSegment.document_id == str(document.id), | DocumentSegment.document_id == str(document.id), | ||||
| DocumentSegment.status != "re_segment", | DocumentSegment.status != "re_segment", | ||||
| ) | ) | ||||
| total_segments = ( | total_segments = ( | ||||
| db.session.query(DocumentSegment) | db.session.query(DocumentSegment) | ||||
| .filter(DocumentSegment.document_id == str(document.id), DocumentSegment.status != "re_segment") | |||||
| .where(DocumentSegment.document_id == str(document.id), DocumentSegment.status != "re_segment") | |||||
| .count() | .count() | ||||
| ) | ) | ||||
| document.completed_segments = completed_segments | document.completed_segments = completed_segments | ||||
| file = ( | file = ( | ||||
| db.session.query(UploadFile) | db.session.query(UploadFile) | ||||
| .filter(UploadFile.tenant_id == document.tenant_id, UploadFile.id == file_id) | |||||
| .where(UploadFile.tenant_id == document.tenant_id, UploadFile.id == file_id) | |||||
| .first() | .first() | ||||
| ) | ) | ||||
| file_id = data_source_info["upload_file_id"] | file_id = data_source_info["upload_file_id"] | ||||
| file_detail = ( | file_detail = ( | ||||
| db.session.query(UploadFile) | db.session.query(UploadFile) | ||||
| .filter(UploadFile.tenant_id == current_user.current_tenant_id, UploadFile.id == file_id) | |||||
| .where(UploadFile.tenant_id == current_user.current_tenant_id, UploadFile.id == file_id) | |||||
| .first() | .first() | ||||
| ) | ) | ||||
| for document in documents: | for document in documents: | ||||
| completed_segments = ( | completed_segments = ( | ||||
| db.session.query(DocumentSegment) | db.session.query(DocumentSegment) | ||||
| .filter( | |||||
| .where( | |||||
| DocumentSegment.completed_at.isnot(None), | DocumentSegment.completed_at.isnot(None), | ||||
| DocumentSegment.document_id == str(document.id), | DocumentSegment.document_id == str(document.id), | ||||
| DocumentSegment.status != "re_segment", | DocumentSegment.status != "re_segment", | ||||
| ) | ) | ||||
| total_segments = ( | total_segments = ( | ||||
| db.session.query(DocumentSegment) | db.session.query(DocumentSegment) | ||||
| .filter(DocumentSegment.document_id == str(document.id), DocumentSegment.status != "re_segment") | |||||
| .where(DocumentSegment.document_id == str(document.id), DocumentSegment.status != "re_segment") | |||||
| .count() | .count() | ||||
| ) | ) | ||||
| # Create a dictionary with document attributes and additional fields | # Create a dictionary with document attributes and additional fields | ||||
| completed_segments = ( | completed_segments = ( | ||||
| db.session.query(DocumentSegment) | db.session.query(DocumentSegment) | ||||
| .filter( | |||||
| .where( | |||||
| DocumentSegment.completed_at.isnot(None), | DocumentSegment.completed_at.isnot(None), | ||||
| DocumentSegment.document_id == str(document_id), | DocumentSegment.document_id == str(document_id), | ||||
| DocumentSegment.status != "re_segment", | DocumentSegment.status != "re_segment", | ||||
| ) | ) | ||||
| total_segments = ( | total_segments = ( | ||||
| db.session.query(DocumentSegment) | db.session.query(DocumentSegment) | ||||
| .filter(DocumentSegment.document_id == str(document_id), DocumentSegment.status != "re_segment") | |||||
| .where(DocumentSegment.document_id == str(document_id), DocumentSegment.status != "re_segment") | |||||
| .count() | .count() | ||||
| ) | ) | ||||
| query = ( | query = ( | ||||
| select(DocumentSegment) | select(DocumentSegment) | ||||
| .filter( | |||||
| .where( | |||||
| DocumentSegment.document_id == str(document_id), | DocumentSegment.document_id == str(document_id), | ||||
| DocumentSegment.tenant_id == current_user.current_tenant_id, | DocumentSegment.tenant_id == current_user.current_tenant_id, | ||||
| ) | ) | ||||
| ) | ) | ||||
| if status_list: | if status_list: | ||||
| query = query.filter(DocumentSegment.status.in_(status_list)) | |||||
| query = query.where(DocumentSegment.status.in_(status_list)) | |||||
| if hit_count_gte is not None: | if hit_count_gte is not None: | ||||
| query = query.filter(DocumentSegment.hit_count >= hit_count_gte) | |||||
| query = query.where(DocumentSegment.hit_count >= hit_count_gte) | |||||
| if keyword: | if keyword: | ||||
| query = query.where(DocumentSegment.content.ilike(f"%{keyword}%")) | query = query.where(DocumentSegment.content.ilike(f"%{keyword}%")) | ||||
| if args["enabled"].lower() != "all": | if args["enabled"].lower() != "all": | ||||
| if args["enabled"].lower() == "true": | if args["enabled"].lower() == "true": | ||||
| query = query.filter(DocumentSegment.enabled == True) | |||||
| query = query.where(DocumentSegment.enabled == True) | |||||
| elif args["enabled"].lower() == "false": | elif args["enabled"].lower() == "false": | ||||
| query = query.filter(DocumentSegment.enabled == False) | |||||
| query = query.where(DocumentSegment.enabled == False) | |||||
| segments = db.paginate(select=query, page=page, per_page=limit, max_per_page=100, error_out=False) | segments = db.paginate(select=query, page=page, per_page=limit, max_per_page=100, error_out=False) | ||||
| segment_id = str(segment_id) | segment_id = str(segment_id) | ||||
| segment = ( | segment = ( | ||||
| db.session.query(DocumentSegment) | db.session.query(DocumentSegment) | ||||
| .filter(DocumentSegment.id == str(segment_id), DocumentSegment.tenant_id == current_user.current_tenant_id) | |||||
| .where(DocumentSegment.id == str(segment_id), DocumentSegment.tenant_id == current_user.current_tenant_id) | |||||
| .first() | .first() | ||||
| ) | ) | ||||
| if not segment: | if not segment: | ||||
| segment_id = str(segment_id) | segment_id = str(segment_id) | ||||
| segment = ( | segment = ( | ||||
| db.session.query(DocumentSegment) | db.session.query(DocumentSegment) | ||||
| .filter(DocumentSegment.id == str(segment_id), DocumentSegment.tenant_id == current_user.current_tenant_id) | |||||
| .where(DocumentSegment.id == str(segment_id), DocumentSegment.tenant_id == current_user.current_tenant_id) | |||||
| .first() | .first() | ||||
| ) | ) | ||||
| if not segment: | if not segment: | ||||
| segment_id = str(segment_id) | segment_id = str(segment_id) | ||||
| segment = ( | segment = ( | ||||
| db.session.query(DocumentSegment) | db.session.query(DocumentSegment) | ||||
| .filter(DocumentSegment.id == str(segment_id), DocumentSegment.tenant_id == current_user.current_tenant_id) | |||||
| .where(DocumentSegment.id == str(segment_id), DocumentSegment.tenant_id == current_user.current_tenant_id) | |||||
| .first() | .first() | ||||
| ) | ) | ||||
| if not segment: | if not segment: | ||||
| segment_id = str(segment_id) | segment_id = str(segment_id) | ||||
| segment = ( | segment = ( | ||||
| db.session.query(DocumentSegment) | db.session.query(DocumentSegment) | ||||
| .filter(DocumentSegment.id == str(segment_id), DocumentSegment.tenant_id == current_user.current_tenant_id) | |||||
| .where(DocumentSegment.id == str(segment_id), DocumentSegment.tenant_id == current_user.current_tenant_id) | |||||
| .first() | .first() | ||||
| ) | ) | ||||
| if not segment: | if not segment: | ||||
| segment_id = str(segment_id) | segment_id = str(segment_id) | ||||
| segment = ( | segment = ( | ||||
| db.session.query(DocumentSegment) | db.session.query(DocumentSegment) | ||||
| .filter(DocumentSegment.id == str(segment_id), DocumentSegment.tenant_id == current_user.current_tenant_id) | |||||
| .where(DocumentSegment.id == str(segment_id), DocumentSegment.tenant_id == current_user.current_tenant_id) | |||||
| .first() | .first() | ||||
| ) | ) | ||||
| if not segment: | if not segment: | ||||
| segment_id = str(segment_id) | segment_id = str(segment_id) | ||||
| segment = ( | segment = ( | ||||
| db.session.query(DocumentSegment) | db.session.query(DocumentSegment) | ||||
| .filter(DocumentSegment.id == str(segment_id), DocumentSegment.tenant_id == current_user.current_tenant_id) | |||||
| .where(DocumentSegment.id == str(segment_id), DocumentSegment.tenant_id == current_user.current_tenant_id) | |||||
| .first() | .first() | ||||
| ) | ) | ||||
| if not segment: | if not segment: | ||||
| child_chunk_id = str(child_chunk_id) | child_chunk_id = str(child_chunk_id) | ||||
| child_chunk = ( | child_chunk = ( | ||||
| db.session.query(ChildChunk) | db.session.query(ChildChunk) | ||||
| .filter(ChildChunk.id == str(child_chunk_id), ChildChunk.tenant_id == current_user.current_tenant_id) | |||||
| .where(ChildChunk.id == str(child_chunk_id), ChildChunk.tenant_id == current_user.current_tenant_id) | |||||
| .first() | .first() | ||||
| ) | ) | ||||
| if not child_chunk: | if not child_chunk: | ||||
| segment_id = str(segment_id) | segment_id = str(segment_id) | ||||
| segment = ( | segment = ( | ||||
| db.session.query(DocumentSegment) | db.session.query(DocumentSegment) | ||||
| .filter(DocumentSegment.id == str(segment_id), DocumentSegment.tenant_id == current_user.current_tenant_id) | |||||
| .where(DocumentSegment.id == str(segment_id), DocumentSegment.tenant_id == current_user.current_tenant_id) | |||||
| .first() | .first() | ||||
| ) | ) | ||||
| if not segment: | if not segment: | ||||
| child_chunk_id = str(child_chunk_id) | child_chunk_id = str(child_chunk_id) | ||||
| child_chunk = ( | child_chunk = ( | ||||
| db.session.query(ChildChunk) | db.session.query(ChildChunk) | ||||
| .filter(ChildChunk.id == str(child_chunk_id), ChildChunk.tenant_id == current_user.current_tenant_id) | |||||
| .where(ChildChunk.id == str(child_chunk_id), ChildChunk.tenant_id == current_user.current_tenant_id) | |||||
| .first() | .first() | ||||
| ) | ) | ||||
| if not child_chunk: | if not child_chunk: |
| if app_id: | if app_id: | ||||
| installed_apps = ( | installed_apps = ( | ||||
| db.session.query(InstalledApp) | db.session.query(InstalledApp) | ||||
| .filter(and_(InstalledApp.tenant_id == current_tenant_id, InstalledApp.app_id == app_id)) | |||||
| .where(and_(InstalledApp.tenant_id == current_tenant_id, InstalledApp.app_id == app_id)) | |||||
| .all() | .all() | ||||
| ) | ) | ||||
| else: | else: | ||||
| installed_apps = db.session.query(InstalledApp).filter(InstalledApp.tenant_id == current_tenant_id).all() | |||||
| installed_apps = db.session.query(InstalledApp).where(InstalledApp.tenant_id == current_tenant_id).all() | |||||
| current_user.role = TenantService.get_user_role(current_user, current_user.current_tenant) | current_user.role = TenantService.get_user_role(current_user, current_user.current_tenant) | ||||
| installed_app_list: list[dict[str, Any]] = [ | installed_app_list: list[dict[str, Any]] = [ | ||||
| parser.add_argument("app_id", type=str, required=True, help="Invalid app_id") | parser.add_argument("app_id", type=str, required=True, help="Invalid app_id") | ||||
| args = parser.parse_args() | args = parser.parse_args() | ||||
| recommended_app = db.session.query(RecommendedApp).filter(RecommendedApp.app_id == args["app_id"]).first() | |||||
| recommended_app = db.session.query(RecommendedApp).where(RecommendedApp.app_id == args["app_id"]).first() | |||||
| if recommended_app is None: | if recommended_app is None: | ||||
| raise NotFound("App not found") | raise NotFound("App not found") | ||||
| current_tenant_id = current_user.current_tenant_id | current_tenant_id = current_user.current_tenant_id | ||||
| app = db.session.query(App).filter(App.id == args["app_id"]).first() | |||||
| app = db.session.query(App).where(App.id == args["app_id"]).first() | |||||
| if app is None: | if app is None: | ||||
| raise NotFound("App not found") | raise NotFound("App not found") | ||||
| installed_app = ( | installed_app = ( | ||||
| db.session.query(InstalledApp) | db.session.query(InstalledApp) | ||||
| .filter(and_(InstalledApp.app_id == args["app_id"], InstalledApp.tenant_id == current_tenant_id)) | |||||
| .where(and_(InstalledApp.app_id == args["app_id"], InstalledApp.tenant_id == current_tenant_id)) | |||||
| .first() | .first() | ||||
| ) | ) | ||||
| installed_app = ( | installed_app = ( | ||||
| db.session.query(InstalledApp) | db.session.query(InstalledApp) | ||||
| .filter( | |||||
| .where( | |||||
| InstalledApp.id == str(installed_app_id), InstalledApp.tenant_id == current_user.current_tenant_id | InstalledApp.id == str(installed_app_id), InstalledApp.tenant_id == current_user.current_tenant_id | ||||
| ) | ) | ||||
| .first() | .first() |
| with Session(db.engine) as session: | with Session(db.engine) as session: | ||||
| permission = ( | permission = ( | ||||
| session.query(TenantPluginPermission) | session.query(TenantPluginPermission) | ||||
| .filter( | |||||
| .where( | |||||
| TenantPluginPermission.tenant_id == tenant_id, | TenantPluginPermission.tenant_id == tenant_id, | ||||
| ) | ) | ||||
| .first() | .first() |
| # check invitation code | # check invitation code | ||||
| invitation_code = ( | invitation_code = ( | ||||
| db.session.query(InvitationCode) | db.session.query(InvitationCode) | ||||
| .filter( | |||||
| .where( | |||||
| InvitationCode.code == args["invitation_code"], | InvitationCode.code == args["invitation_code"], | ||||
| InvitationCode.status == "unused", | InvitationCode.status == "unused", | ||||
| ) | ) | ||||
| def get(self): | def get(self): | ||||
| account = current_user | account = current_user | ||||
| account_integrates = db.session.query(AccountIntegrate).filter(AccountIntegrate.account_id == account.id).all() | |||||
| account_integrates = db.session.query(AccountIntegrate).where(AccountIntegrate.account_id == account.id).all() | |||||
| base_url = request.url_root.rstrip("/") | base_url = request.url_root.rstrip("/") | ||||
| oauth_base_path = "/console/api/oauth/login" | oauth_base_path = "/console/api/oauth/login" |
| @login_required | @login_required | ||||
| @account_initialization_required | @account_initialization_required | ||||
| def delete(self, member_id): | def delete(self, member_id): | ||||
| member = db.session.query(Account).filter(Account.id == str(member_id)).first() | |||||
| member = db.session.query(Account).where(Account.id == str(member_id)).first() | |||||
| if member is None: | if member is None: | ||||
| abort(404) | abort(404) | ||||
| else: | else: |
| user_id = "DEFAULT-USER" | user_id = "DEFAULT-USER" | ||||
| if user_id == "DEFAULT-USER": | if user_id == "DEFAULT-USER": | ||||
| user_model = session.query(EndUser).filter(EndUser.session_id == "DEFAULT-USER").first() | |||||
| user_model = session.query(EndUser).where(EndUser.session_id == "DEFAULT-USER").first() | |||||
| if not user_model: | if not user_model: | ||||
| user_model = EndUser( | user_model = EndUser( | ||||
| tenant_id=tenant_id, | tenant_id=tenant_id, | ||||
| else: | else: | ||||
| user_model = AccountService.load_user(user_id) | user_model = AccountService.load_user(user_id) | ||||
| if not user_model: | if not user_model: | ||||
| user_model = session.query(EndUser).filter(EndUser.id == user_id).first() | |||||
| user_model = session.query(EndUser).where(EndUser.id == user_id).first() | |||||
| if not user_model: | if not user_model: | ||||
| raise ValueError("user not found") | raise ValueError("user not found") | ||||
| except Exception: | except Exception: | ||||
| try: | try: | ||||
| tenant_model = ( | tenant_model = ( | ||||
| db.session.query(Tenant) | db.session.query(Tenant) | ||||
| .filter( | |||||
| .where( | |||||
| Tenant.id == tenant_id, | Tenant.id == tenant_id, | ||||
| ) | ) | ||||
| .first() | .first() |
| if signature_base64 != token: | if signature_base64 != token: | ||||
| return view(*args, **kwargs) | return view(*args, **kwargs) | ||||
| kwargs["user"] = db.session.query(EndUser).filter(EndUser.id == user_id).first() | |||||
| kwargs["user"] = db.session.query(EndUser).where(EndUser.id == user_id).first() | |||||
| return view(*args, **kwargs) | return view(*args, **kwargs) | ||||
| request_id = args.get("id") | request_id = args.get("id") | ||||
| server = db.session.query(AppMCPServer).filter(AppMCPServer.server_code == server_code).first() | |||||
| server = db.session.query(AppMCPServer).where(AppMCPServer.server_code == server_code).first() | |||||
| if not server: | if not server: | ||||
| return helper.compact_generate_response( | return helper.compact_generate_response( | ||||
| create_mcp_error_response(request_id, types.INVALID_REQUEST, "Server Not Found") | create_mcp_error_response(request_id, types.INVALID_REQUEST, "Server Not Found") | ||||
| create_mcp_error_response(request_id, types.INVALID_REQUEST, "Server is not active") | create_mcp_error_response(request_id, types.INVALID_REQUEST, "Server is not active") | ||||
| ) | ) | ||||
| app = db.session.query(App).filter(App.id == server.app_id).first() | |||||
| app = db.session.query(App).where(App.id == server.app_id).first() | |||||
| if not app: | if not app: | ||||
| return helper.compact_generate_response( | return helper.compact_generate_response( | ||||
| create_mcp_error_response(request_id, types.INVALID_REQUEST, "App Not Found") | create_mcp_error_response(request_id, types.INVALID_REQUEST, "App Not Found") |
| @marshal_with(fields.site_fields) | @marshal_with(fields.site_fields) | ||||
| def get(self, app_model: App): | def get(self, app_model: App): | ||||
| """Retrieve app site info.""" | """Retrieve app site info.""" | ||||
| site = db.session.query(Site).filter(Site.app_id == app_model.id).first() | |||||
| site = db.session.query(Site).where(Site.app_id == app_model.id).first() | |||||
| if not site: | if not site: | ||||
| raise Forbidden() | raise Forbidden() |
| dataset_id = str(dataset_id) | dataset_id = str(dataset_id) | ||||
| tenant_id = str(tenant_id) | tenant_id = str(tenant_id) | ||||
| dataset = db.session.query(Dataset).filter(Dataset.tenant_id == tenant_id, Dataset.id == dataset_id).first() | |||||
| dataset = db.session.query(Dataset).where(Dataset.tenant_id == tenant_id, Dataset.id == dataset_id).first() | |||||
| if not dataset: | if not dataset: | ||||
| raise ValueError("Dataset does not exist.") | raise ValueError("Dataset does not exist.") | ||||
| args = parser.parse_args() | args = parser.parse_args() | ||||
| dataset_id = str(dataset_id) | dataset_id = str(dataset_id) | ||||
| tenant_id = str(tenant_id) | tenant_id = str(tenant_id) | ||||
| dataset = db.session.query(Dataset).filter(Dataset.tenant_id == tenant_id, Dataset.id == dataset_id).first() | |||||
| dataset = db.session.query(Dataset).where(Dataset.tenant_id == tenant_id, Dataset.id == dataset_id).first() | |||||
| if not dataset: | if not dataset: | ||||
| raise ValueError("Dataset does not exist.") | raise ValueError("Dataset does not exist.") | ||||
| # get dataset info | # get dataset info | ||||
| dataset_id = str(dataset_id) | dataset_id = str(dataset_id) | ||||
| tenant_id = str(tenant_id) | tenant_id = str(tenant_id) | ||||
| dataset = db.session.query(Dataset).filter(Dataset.tenant_id == tenant_id, Dataset.id == dataset_id).first() | |||||
| dataset = db.session.query(Dataset).where(Dataset.tenant_id == tenant_id, Dataset.id == dataset_id).first() | |||||
| if not dataset: | if not dataset: | ||||
| raise ValueError("Dataset does not exist.") | raise ValueError("Dataset does not exist.") | ||||
| # get dataset info | # get dataset info | ||||
| dataset_id = str(dataset_id) | dataset_id = str(dataset_id) | ||||
| tenant_id = str(tenant_id) | tenant_id = str(tenant_id) | ||||
| dataset = db.session.query(Dataset).filter(Dataset.tenant_id == tenant_id, Dataset.id == dataset_id).first() | |||||
| dataset = db.session.query(Dataset).where(Dataset.tenant_id == tenant_id, Dataset.id == dataset_id).first() | |||||
| if not dataset: | if not dataset: | ||||
| raise ValueError("Dataset does not exist.") | raise ValueError("Dataset does not exist.") | ||||
| tenant_id = str(tenant_id) | tenant_id = str(tenant_id) | ||||
| # get dataset info | # get dataset info | ||||
| dataset = db.session.query(Dataset).filter(Dataset.tenant_id == tenant_id, Dataset.id == dataset_id).first() | |||||
| dataset = db.session.query(Dataset).where(Dataset.tenant_id == tenant_id, Dataset.id == dataset_id).first() | |||||
| if not dataset: | if not dataset: | ||||
| raise ValueError("Dataset does not exist.") | raise ValueError("Dataset does not exist.") | ||||
| page = request.args.get("page", default=1, type=int) | page = request.args.get("page", default=1, type=int) | ||||
| limit = request.args.get("limit", default=20, type=int) | limit = request.args.get("limit", default=20, type=int) | ||||
| search = request.args.get("keyword", default=None, type=str) | search = request.args.get("keyword", default=None, type=str) | ||||
| dataset = db.session.query(Dataset).filter(Dataset.tenant_id == tenant_id, Dataset.id == dataset_id).first() | |||||
| dataset = db.session.query(Dataset).where(Dataset.tenant_id == tenant_id, Dataset.id == dataset_id).first() | |||||
| if not dataset: | if not dataset: | ||||
| raise NotFound("Dataset not found.") | raise NotFound("Dataset not found.") | ||||
| if search: | if search: | ||||
| search = f"%{search}%" | search = f"%{search}%" | ||||
| query = query.filter(Document.name.like(search)) | |||||
| query = query.where(Document.name.like(search)) | |||||
| query = query.order_by(desc(Document.created_at), desc(Document.position)) | query = query.order_by(desc(Document.created_at), desc(Document.position)) | ||||
| batch = str(batch) | batch = str(batch) | ||||
| tenant_id = str(tenant_id) | tenant_id = str(tenant_id) | ||||
| # get dataset | # get dataset | ||||
| dataset = db.session.query(Dataset).filter(Dataset.tenant_id == tenant_id, Dataset.id == dataset_id).first() | |||||
| dataset = db.session.query(Dataset).where(Dataset.tenant_id == tenant_id, Dataset.id == dataset_id).first() | |||||
| if not dataset: | if not dataset: | ||||
| raise NotFound("Dataset not found.") | raise NotFound("Dataset not found.") | ||||
| # get documents | # get documents | ||||
| for document in documents: | for document in documents: | ||||
| completed_segments = ( | completed_segments = ( | ||||
| db.session.query(DocumentSegment) | db.session.query(DocumentSegment) | ||||
| .filter( | |||||
| .where( | |||||
| DocumentSegment.completed_at.isnot(None), | DocumentSegment.completed_at.isnot(None), | ||||
| DocumentSegment.document_id == str(document.id), | DocumentSegment.document_id == str(document.id), | ||||
| DocumentSegment.status != "re_segment", | DocumentSegment.status != "re_segment", | ||||
| ) | ) | ||||
| total_segments = ( | total_segments = ( | ||||
| db.session.query(DocumentSegment) | db.session.query(DocumentSegment) | ||||
| .filter(DocumentSegment.document_id == str(document.id), DocumentSegment.status != "re_segment") | |||||
| .where(DocumentSegment.document_id == str(document.id), DocumentSegment.status != "re_segment") | |||||
| .count() | .count() | ||||
| ) | ) | ||||
| # Create a dictionary with document attributes and additional fields | # Create a dictionary with document attributes and additional fields |
| # check dataset | # check dataset | ||||
| dataset_id = str(dataset_id) | dataset_id = str(dataset_id) | ||||
| tenant_id = str(tenant_id) | tenant_id = str(tenant_id) | ||||
| dataset = db.session.query(Dataset).filter(Dataset.tenant_id == tenant_id, Dataset.id == dataset_id).first() | |||||
| dataset = db.session.query(Dataset).where(Dataset.tenant_id == tenant_id, Dataset.id == dataset_id).first() | |||||
| if not dataset: | if not dataset: | ||||
| raise NotFound("Dataset not found.") | raise NotFound("Dataset not found.") | ||||
| # check document | # check document | ||||
| tenant_id = str(tenant_id) | tenant_id = str(tenant_id) | ||||
| page = request.args.get("page", default=1, type=int) | page = request.args.get("page", default=1, type=int) | ||||
| limit = request.args.get("limit", default=20, type=int) | limit = request.args.get("limit", default=20, type=int) | ||||
| dataset = db.session.query(Dataset).filter(Dataset.tenant_id == tenant_id, Dataset.id == dataset_id).first() | |||||
| dataset = db.session.query(Dataset).where(Dataset.tenant_id == tenant_id, Dataset.id == dataset_id).first() | |||||
| if not dataset: | if not dataset: | ||||
| raise NotFound("Dataset not found.") | raise NotFound("Dataset not found.") | ||||
| # check document | # check document | ||||
| # check dataset | # check dataset | ||||
| dataset_id = str(dataset_id) | dataset_id = str(dataset_id) | ||||
| tenant_id = str(tenant_id) | tenant_id = str(tenant_id) | ||||
| dataset = db.session.query(Dataset).filter(Dataset.tenant_id == tenant_id, Dataset.id == dataset_id).first() | |||||
| dataset = db.session.query(Dataset).where(Dataset.tenant_id == tenant_id, Dataset.id == dataset_id).first() | |||||
| if not dataset: | if not dataset: | ||||
| raise NotFound("Dataset not found.") | raise NotFound("Dataset not found.") | ||||
| # check user's model setting | # check user's model setting | ||||
| # check dataset | # check dataset | ||||
| dataset_id = str(dataset_id) | dataset_id = str(dataset_id) | ||||
| tenant_id = str(tenant_id) | tenant_id = str(tenant_id) | ||||
| dataset = db.session.query(Dataset).filter(Dataset.tenant_id == tenant_id, Dataset.id == dataset_id).first() | |||||
| dataset = db.session.query(Dataset).where(Dataset.tenant_id == tenant_id, Dataset.id == dataset_id).first() | |||||
| if not dataset: | if not dataset: | ||||
| raise NotFound("Dataset not found.") | raise NotFound("Dataset not found.") | ||||
| # check user's model setting | # check user's model setting | ||||
| # check dataset | # check dataset | ||||
| dataset_id = str(dataset_id) | dataset_id = str(dataset_id) | ||||
| tenant_id = str(tenant_id) | tenant_id = str(tenant_id) | ||||
| dataset = db.session.query(Dataset).filter(Dataset.tenant_id == tenant_id, Dataset.id == dataset_id).first() | |||||
| dataset = db.session.query(Dataset).where(Dataset.tenant_id == tenant_id, Dataset.id == dataset_id).first() | |||||
| if not dataset: | if not dataset: | ||||
| raise NotFound("Dataset not found.") | raise NotFound("Dataset not found.") | ||||
| # check user's model setting | # check user's model setting | ||||
| # check dataset | # check dataset | ||||
| dataset_id = str(dataset_id) | dataset_id = str(dataset_id) | ||||
| tenant_id = str(tenant_id) | tenant_id = str(tenant_id) | ||||
| dataset = db.session.query(Dataset).filter(Dataset.tenant_id == tenant_id, Dataset.id == dataset_id).first() | |||||
| dataset = db.session.query(Dataset).where(Dataset.tenant_id == tenant_id, Dataset.id == dataset_id).first() | |||||
| if not dataset: | if not dataset: | ||||
| raise NotFound("Dataset not found.") | raise NotFound("Dataset not found.") | ||||
| # check dataset | # check dataset | ||||
| dataset_id = str(dataset_id) | dataset_id = str(dataset_id) | ||||
| tenant_id = str(tenant_id) | tenant_id = str(tenant_id) | ||||
| dataset = db.session.query(Dataset).filter(Dataset.tenant_id == tenant_id, Dataset.id == dataset_id).first() | |||||
| dataset = db.session.query(Dataset).where(Dataset.tenant_id == tenant_id, Dataset.id == dataset_id).first() | |||||
| if not dataset: | if not dataset: | ||||
| raise NotFound("Dataset not found.") | raise NotFound("Dataset not found.") | ||||
| # check dataset | # check dataset | ||||
| dataset_id = str(dataset_id) | dataset_id = str(dataset_id) | ||||
| tenant_id = str(tenant_id) | tenant_id = str(tenant_id) | ||||
| dataset = db.session.query(Dataset).filter(Dataset.tenant_id == tenant_id, Dataset.id == dataset_id).first() | |||||
| dataset = db.session.query(Dataset).where(Dataset.tenant_id == tenant_id, Dataset.id == dataset_id).first() | |||||
| if not dataset: | if not dataset: | ||||
| raise NotFound("Dataset not found.") | raise NotFound("Dataset not found.") | ||||
| # check dataset | # check dataset | ||||
| dataset_id = str(dataset_id) | dataset_id = str(dataset_id) | ||||
| tenant_id = str(tenant_id) | tenant_id = str(tenant_id) | ||||
| dataset = db.session.query(Dataset).filter(Dataset.tenant_id == tenant_id, Dataset.id == dataset_id).first() | |||||
| dataset = db.session.query(Dataset).where(Dataset.tenant_id == tenant_id, Dataset.id == dataset_id).first() | |||||
| if not dataset: | if not dataset: | ||||
| raise NotFound("Dataset not found.") | raise NotFound("Dataset not found.") | ||||
| # check dataset | # check dataset | ||||
| dataset_id = str(dataset_id) | dataset_id = str(dataset_id) | ||||
| tenant_id = str(tenant_id) | tenant_id = str(tenant_id) | ||||
| dataset = db.session.query(Dataset).filter(Dataset.tenant_id == tenant_id, Dataset.id == dataset_id).first() | |||||
| dataset = db.session.query(Dataset).where(Dataset.tenant_id == tenant_id, Dataset.id == dataset_id).first() | |||||
| if not dataset: | if not dataset: | ||||
| raise NotFound("Dataset not found.") | raise NotFound("Dataset not found.") | ||||
| # check document | # check document | ||||
| data_source_info = document.data_source_info_dict | data_source_info = document.data_source_info_dict | ||||
| if data_source_info and "upload_file_id" in data_source_info: | if data_source_info and "upload_file_id" in data_source_info: | ||||
| file_id = data_source_info["upload_file_id"] | file_id = data_source_info["upload_file_id"] | ||||
| upload_file = db.session.query(UploadFile).filter(UploadFile.id == file_id).first() | |||||
| upload_file = db.session.query(UploadFile).where(UploadFile.id == file_id).first() | |||||
| if not upload_file: | if not upload_file: | ||||
| raise NotFound("UploadFile not found.") | raise NotFound("UploadFile not found.") | ||||
| else: | else: |
| def decorated_view(*args, **kwargs): | def decorated_view(*args, **kwargs): | ||||
| api_token = validate_and_get_api_token("app") | api_token = validate_and_get_api_token("app") | ||||
| app_model = db.session.query(App).filter(App.id == api_token.app_id).first() | |||||
| app_model = db.session.query(App).where(App.id == api_token.app_id).first() | |||||
| if not app_model: | if not app_model: | ||||
| raise Forbidden("The app no longer exists.") | raise Forbidden("The app no longer exists.") | ||||
| if not app_model.enable_api: | if not app_model.enable_api: | ||||
| raise Forbidden("The app's API service has been disabled.") | raise Forbidden("The app's API service has been disabled.") | ||||
| tenant = db.session.query(Tenant).filter(Tenant.id == app_model.tenant_id).first() | |||||
| tenant = db.session.query(Tenant).where(Tenant.id == app_model.tenant_id).first() | |||||
| if tenant is None: | if tenant is None: | ||||
| raise ValueError("Tenant does not exist.") | raise ValueError("Tenant does not exist.") | ||||
| if tenant.status == TenantStatus.ARCHIVE: | if tenant.status == TenantStatus.ARCHIVE: | ||||
| tenant_account_join = ( | tenant_account_join = ( | ||||
| db.session.query(Tenant, TenantAccountJoin) | db.session.query(Tenant, TenantAccountJoin) | ||||
| .filter(Tenant.id == api_token.tenant_id) | |||||
| .filter(TenantAccountJoin.tenant_id == Tenant.id) | |||||
| .filter(TenantAccountJoin.role.in_(["owner"])) | |||||
| .filter(Tenant.status == TenantStatus.NORMAL) | |||||
| .where(Tenant.id == api_token.tenant_id) | |||||
| .where(TenantAccountJoin.tenant_id == Tenant.id) | |||||
| .where(TenantAccountJoin.role.in_(["owner"])) | |||||
| .where(Tenant.status == TenantStatus.NORMAL) | |||||
| .one_or_none() | .one_or_none() | ||||
| ) # TODO: only owner information is required, so only one is returned. | ) # TODO: only owner information is required, so only one is returned. | ||||
| if tenant_account_join: | if tenant_account_join: | ||||
| tenant, ta = tenant_account_join | tenant, ta = tenant_account_join | ||||
| account = db.session.query(Account).filter(Account.id == ta.account_id).first() | |||||
| account = db.session.query(Account).where(Account.id == ta.account_id).first() | |||||
| # Login admin | # Login admin | ||||
| if account: | if account: | ||||
| account.current_tenant = tenant | account.current_tenant = tenant | ||||
| api_token = validate_and_get_api_token("dataset") | api_token = validate_and_get_api_token("dataset") | ||||
| tenant_account_join = ( | tenant_account_join = ( | ||||
| db.session.query(Tenant, TenantAccountJoin) | db.session.query(Tenant, TenantAccountJoin) | ||||
| .filter(Tenant.id == api_token.tenant_id) | |||||
| .filter(TenantAccountJoin.tenant_id == Tenant.id) | |||||
| .filter(TenantAccountJoin.role.in_(["owner"])) | |||||
| .filter(Tenant.status == TenantStatus.NORMAL) | |||||
| .where(Tenant.id == api_token.tenant_id) | |||||
| .where(TenantAccountJoin.tenant_id == Tenant.id) | |||||
| .where(TenantAccountJoin.role.in_(["owner"])) | |||||
| .where(Tenant.status == TenantStatus.NORMAL) | |||||
| .one_or_none() | .one_or_none() | ||||
| ) # TODO: only owner information is required, so only one is returned. | ) # TODO: only owner information is required, so only one is returned. | ||||
| if tenant_account_join: | if tenant_account_join: | ||||
| tenant, ta = tenant_account_join | tenant, ta = tenant_account_join | ||||
| account = db.session.query(Account).filter(Account.id == ta.account_id).first() | |||||
| account = db.session.query(Account).where(Account.id == ta.account_id).first() | |||||
| # Login admin | # Login admin | ||||
| if account: | if account: | ||||
| account.current_tenant = tenant | account.current_tenant = tenant | ||||
| end_user = ( | end_user = ( | ||||
| db.session.query(EndUser) | db.session.query(EndUser) | ||||
| .filter( | |||||
| .where( | |||||
| EndUser.tenant_id == app_model.tenant_id, | EndUser.tenant_id == app_model.tenant_id, | ||||
| EndUser.app_id == app_model.id, | EndUser.app_id == app_model.id, | ||||
| EndUser.session_id == user_id, | EndUser.session_id == user_id, | ||||
| method_decorators = [validate_dataset_token] | method_decorators = [validate_dataset_token] | ||||
| def get_dataset(self, dataset_id: str, tenant_id: str) -> Dataset: | def get_dataset(self, dataset_id: str, tenant_id: str) -> Dataset: | ||||
| dataset = db.session.query(Dataset).filter(Dataset.id == dataset_id, Dataset.tenant_id == tenant_id).first() | |||||
| dataset = db.session.query(Dataset).where(Dataset.id == dataset_id, Dataset.tenant_id == tenant_id).first() | |||||
| if not dataset: | if not dataset: | ||||
| raise NotFound("Dataset not found.") | raise NotFound("Dataset not found.") |
| from flask import request | from flask import request | ||||
| from flask_restful import Resource | from flask_restful import Resource | ||||
| from sqlalchemy import func, select | |||||
| from werkzeug.exceptions import NotFound, Unauthorized | from werkzeug.exceptions import NotFound, Unauthorized | ||||
| from configs import dify_config | from configs import dify_config | ||||
| raise WebAppAuthRequiredError() | raise WebAppAuthRequiredError() | ||||
| # get site from db and check if it is normal | # get site from db and check if it is normal | ||||
| site = db.session.query(Site).filter(Site.code == app_code, Site.status == "normal").first() | |||||
| site = db.session.scalar(select(Site).where(Site.code == app_code, Site.status == "normal")) | |||||
| if not site: | if not site: | ||||
| raise NotFound() | raise NotFound() | ||||
| # get app from db and check if it is normal and enable_site | # get app from db and check if it is normal and enable_site | ||||
| app_model = db.session.query(App).filter(App.id == site.app_id).first() | |||||
| app_model = db.session.scalar(select(App).where(App.id == site.app_id)) | |||||
| if not app_model or app_model.status != "normal" or not app_model.enable_site: | if not app_model or app_model.status != "normal" or not app_model.enable_site: | ||||
| raise NotFound() | raise NotFound() | ||||
| if user_id: | if user_id: | ||||
| end_user = ( | |||||
| db.session.query(EndUser).filter(EndUser.app_id == app_model.id, EndUser.session_id == user_id).first() | |||||
| end_user = db.session.scalar( | |||||
| select(EndUser).where(EndUser.app_id == app_model.id, EndUser.session_id == user_id) | |||||
| ) | ) | ||||
| if end_user: | if end_user: | ||||
| if not user_auth_type: | if not user_auth_type: | ||||
| raise Unauthorized("Missing auth_type in the token.") | raise Unauthorized("Missing auth_type in the token.") | ||||
| site = db.session.query(Site).filter(Site.code == app_code, Site.status == "normal").first() | |||||
| site = db.session.scalar(select(Site).where(Site.code == app_code, Site.status == "normal")) | |||||
| if not site: | if not site: | ||||
| raise NotFound() | raise NotFound() | ||||
| app_model = db.session.query(App).filter(App.id == site.app_id).first() | |||||
| app_model = db.session.scalar(select(App).where(App.id == site.app_id)) | |||||
| if not app_model or app_model.status != "normal" or not app_model.enable_site: | if not app_model or app_model.status != "normal" or not app_model.enable_site: | ||||
| raise NotFound() | raise NotFound() | ||||
| end_user = None | end_user = None | ||||
| if end_user_id: | if end_user_id: | ||||
| end_user = db.session.query(EndUser).filter(EndUser.id == end_user_id).first() | |||||
| end_user = db.session.scalar(select(EndUser).where(EndUser.id == end_user_id)) | |||||
| if session_id: | if session_id: | ||||
| end_user = ( | |||||
| db.session.query(EndUser) | |||||
| .filter( | |||||
| end_user = db.session.scalar( | |||||
| select(EndUser).where( | |||||
| EndUser.session_id == session_id, | EndUser.session_id == session_id, | ||||
| EndUser.tenant_id == app_model.tenant_id, | EndUser.tenant_id == app_model.tenant_id, | ||||
| EndUser.app_id == app_model.id, | EndUser.app_id == app_model.id, | ||||
| ) | ) | ||||
| .first() | |||||
| ) | ) | ||||
| if not end_user: | if not end_user: | ||||
| if not session_id: | if not session_id: | ||||
| user_id = token_decoded.get("user_id") | user_id = token_decoded.get("user_id") | ||||
| end_user = None | end_user = None | ||||
| if user_id: | if user_id: | ||||
| end_user = ( | |||||
| db.session.query(EndUser).filter(EndUser.app_id == app_model.id, EndUser.session_id == user_id).first() | |||||
| end_user = db.session.scalar( | |||||
| select(EndUser).where(EndUser.app_id == app_model.id, EndUser.session_id == user_id) | |||||
| ) | ) | ||||
| if not end_user: | if not end_user: | ||||
| """ | """ | ||||
| while True: | while True: | ||||
| session_id = str(uuid.uuid4()) | session_id = str(uuid.uuid4()) | ||||
| existing_count = db.session.query(EndUser).filter(EndUser.session_id == session_id).count() | |||||
| existing_count = db.session.scalar( | |||||
| select(func.count()).select_from(EndUser).where(EndUser.session_id == session_id) | |||||
| ) | |||||
| if existing_count == 0: | if existing_count == 0: | ||||
| return session_id | return session_id |
| def get(self, app_model, end_user): | def get(self, app_model, end_user): | ||||
| """Retrieve app site info.""" | """Retrieve app site info.""" | ||||
| # get site | # get site | ||||
| site = db.session.query(Site).filter(Site.app_id == app_model.id).first() | |||||
| site = db.session.query(Site).where(Site.app_id == app_model.id).first() | |||||
| if not site: | if not site: | ||||
| raise Forbidden() | raise Forbidden() |
| from flask import request | from flask import request | ||||
| from flask_restful import Resource | from flask_restful import Resource | ||||
| from sqlalchemy import select | |||||
| from werkzeug.exceptions import BadRequest, NotFound, Unauthorized | from werkzeug.exceptions import BadRequest, NotFound, Unauthorized | ||||
| from controllers.web.error import WebAppAuthAccessDeniedError, WebAppAuthRequiredError | from controllers.web.error import WebAppAuthAccessDeniedError, WebAppAuthRequiredError | ||||
| decoded = PassportService().verify(tk) | decoded = PassportService().verify(tk) | ||||
| app_code = decoded.get("app_code") | app_code = decoded.get("app_code") | ||||
| app_id = decoded.get("app_id") | app_id = decoded.get("app_id") | ||||
| app_model = db.session.query(App).filter(App.id == app_id).first() | |||||
| site = db.session.query(Site).filter(Site.code == app_code).first() | |||||
| app_model = db.session.scalar(select(App).where(App.id == app_id)) | |||||
| site = db.session.scalar(select(Site).where(Site.code == app_code)) | |||||
| if not app_model: | if not app_model: | ||||
| raise NotFound() | raise NotFound() | ||||
| if not app_code or not site: | if not app_code or not site: | ||||
| if app_model.enable_site is False: | if app_model.enable_site is False: | ||||
| raise BadRequest("Site is disabled.") | raise BadRequest("Site is disabled.") | ||||
| end_user_id = decoded.get("end_user_id") | end_user_id = decoded.get("end_user_id") | ||||
| end_user = db.session.query(EndUser).filter(EndUser.id == end_user_id).first() | |||||
| end_user = db.session.scalar(select(EndUser).where(EndUser.id == end_user_id)) | |||||
| if not end_user: | if not end_user: | ||||
| raise NotFound() | raise NotFound() | ||||
| # get how many agent thoughts have been created | # get how many agent thoughts have been created | ||||
| self.agent_thought_count = ( | self.agent_thought_count = ( | ||||
| db.session.query(MessageAgentThought) | db.session.query(MessageAgentThought) | ||||
| .filter( | |||||
| .where( | |||||
| MessageAgentThought.message_id == self.message.id, | MessageAgentThought.message_id == self.message.id, | ||||
| ) | ) | ||||
| .count() | .count() | ||||
| Save agent thought | Save agent thought | ||||
| """ | """ | ||||
| updated_agent_thought = ( | updated_agent_thought = ( | ||||
| db.session.query(MessageAgentThought).filter(MessageAgentThought.id == agent_thought.id).first() | |||||
| db.session.query(MessageAgentThought).where(MessageAgentThought.id == agent_thought.id).first() | |||||
| ) | ) | ||||
| if not updated_agent_thought: | if not updated_agent_thought: | ||||
| raise ValueError("agent thought not found") | raise ValueError("agent thought not found") | ||||
| return result | return result | ||||
| def organize_agent_user_prompt(self, message: Message) -> UserPromptMessage: | def organize_agent_user_prompt(self, message: Message) -> UserPromptMessage: | ||||
| files = db.session.query(MessageFile).filter(MessageFile.message_id == message.id).all() | |||||
| files = db.session.query(MessageFile).where(MessageFile.message_id == message.id).all() | |||||
| if not files: | if not files: | ||||
| return UserPromptMessage(content=message.query) | return UserPromptMessage(content=message.query) | ||||
| if message.app_model_config: | if message.app_model_config: |
| app_config = self.application_generate_entity.app_config | app_config = self.application_generate_entity.app_config | ||||
| app_config = cast(AdvancedChatAppConfig, app_config) | app_config = cast(AdvancedChatAppConfig, app_config) | ||||
| app_record = db.session.query(App).filter(App.id == app_config.app_id).first() | |||||
| app_record = db.session.query(App).where(App.id == app_config.app_id).first() | |||||
| if not app_record: | if not app_record: | ||||
| raise ValueError("App not found") | raise ValueError("App not found") | ||||
| app_config = application_generate_entity.app_config | app_config = application_generate_entity.app_config | ||||
| app_config = cast(AgentChatAppConfig, app_config) | app_config = cast(AgentChatAppConfig, app_config) | ||||
| app_record = db.session.query(App).filter(App.id == app_config.app_id).first() | |||||
| app_record = db.session.query(App).where(App.id == app_config.app_id).first() | |||||
| if not app_record: | if not app_record: | ||||
| raise ValueError("App not found") | raise ValueError("App not found") | ||||
| if {ModelFeature.MULTI_TOOL_CALL, ModelFeature.TOOL_CALL}.intersection(model_schema.features or []): | if {ModelFeature.MULTI_TOOL_CALL, ModelFeature.TOOL_CALL}.intersection(model_schema.features or []): | ||||
| agent_entity.strategy = AgentEntity.Strategy.FUNCTION_CALLING | agent_entity.strategy = AgentEntity.Strategy.FUNCTION_CALLING | ||||
| conversation_result = db.session.query(Conversation).filter(Conversation.id == conversation.id).first() | |||||
| conversation_result = db.session.query(Conversation).where(Conversation.id == conversation.id).first() | |||||
| if conversation_result is None: | if conversation_result is None: | ||||
| raise ValueError("Conversation not found") | raise ValueError("Conversation not found") | ||||
| message_result = db.session.query(Message).filter(Message.id == message.id).first() | |||||
| message_result = db.session.query(Message).where(Message.id == message.id).first() | |||||
| if message_result is None: | if message_result is None: | ||||
| raise ValueError("Message not found") | raise ValueError("Message not found") | ||||
| db.session.close() | db.session.close() |
| app_config = application_generate_entity.app_config | app_config = application_generate_entity.app_config | ||||
| app_config = cast(ChatAppConfig, app_config) | app_config = cast(ChatAppConfig, app_config) | ||||
| app_record = db.session.query(App).filter(App.id == app_config.app_id).first() | |||||
| app_record = db.session.query(App).where(App.id == app_config.app_id).first() | |||||
| if not app_record: | if not app_record: | ||||
| raise ValueError("App not found") | raise ValueError("App not found") | ||||
| """ | """ | ||||
| message = ( | message = ( | ||||
| db.session.query(Message) | db.session.query(Message) | ||||
| .filter( | |||||
| .where( | |||||
| Message.id == message_id, | Message.id == message_id, | ||||
| Message.app_id == app_model.id, | Message.app_id == app_model.id, | ||||
| Message.from_source == ("api" if isinstance(user, EndUser) else "console"), | Message.from_source == ("api" if isinstance(user, EndUser) else "console"), |
| app_config = application_generate_entity.app_config | app_config = application_generate_entity.app_config | ||||
| app_config = cast(CompletionAppConfig, app_config) | app_config = cast(CompletionAppConfig, app_config) | ||||
| app_record = db.session.query(App).filter(App.id == app_config.app_id).first() | |||||
| app_record = db.session.query(App).where(App.id == app_config.app_id).first() | |||||
| if not app_record: | if not app_record: | ||||
| raise ValueError("App not found") | raise ValueError("App not found") | ||||
| if conversation: | if conversation: | ||||
| app_model_config = ( | app_model_config = ( | ||||
| db.session.query(AppModelConfig) | db.session.query(AppModelConfig) | ||||
| .filter(AppModelConfig.id == conversation.app_model_config_id, AppModelConfig.app_id == app_model.id) | |||||
| .where(AppModelConfig.id == conversation.app_model_config_id, AppModelConfig.app_id == app_model.id) | |||||
| .first() | .first() | ||||
| ) | ) | ||||
| :param conversation_id: conversation id | :param conversation_id: conversation id | ||||
| :return: conversation | :return: conversation | ||||
| """ | """ | ||||
| conversation = db.session.query(Conversation).filter(Conversation.id == conversation_id).first() | |||||
| conversation = db.session.query(Conversation).where(Conversation.id == conversation_id).first() | |||||
| if not conversation: | if not conversation: | ||||
| raise ConversationNotExistsError("Conversation not exists") | raise ConversationNotExistsError("Conversation not exists") | ||||
| :param message_id: message id | :param message_id: message id | ||||
| :return: message | :return: message | ||||
| """ | """ | ||||
| message = db.session.query(Message).filter(Message.id == message_id).first() | |||||
| message = db.session.query(Message).where(Message.id == message_id).first() | |||||
| if message is None: | if message is None: | ||||
| raise MessageNotExistsError("Message not exists") | raise MessageNotExistsError("Message not exists") |
| :return: | :return: | ||||
| """ | """ | ||||
| annotation_setting = ( | annotation_setting = ( | ||||
| db.session.query(AppAnnotationSetting).filter(AppAnnotationSetting.app_id == app_record.id).first() | |||||
| db.session.query(AppAnnotationSetting).where(AppAnnotationSetting.app_id == app_record.id).first() | |||||
| ) | ) | ||||
| if not annotation_setting: | if not annotation_setting: |
| :return: | :return: | ||||
| """ | """ | ||||
| agent_thought: Optional[MessageAgentThought] = ( | agent_thought: Optional[MessageAgentThought] = ( | ||||
| db.session.query(MessageAgentThought).filter(MessageAgentThought.id == event.agent_thought_id).first() | |||||
| db.session.query(MessageAgentThought).where(MessageAgentThought.id == event.agent_thought_id).first() | |||||
| ) | ) | ||||
| if agent_thought: | if agent_thought: |
| def _generate_conversation_name_worker(self, flask_app: Flask, conversation_id: str, query: str): | def _generate_conversation_name_worker(self, flask_app: Flask, conversation_id: str, query: str): | ||||
| with flask_app.app_context(): | with flask_app.app_context(): | ||||
| # get conversation and message | # get conversation and message | ||||
| conversation = db.session.query(Conversation).filter(Conversation.id == conversation_id).first() | |||||
| conversation = db.session.query(Conversation).where(Conversation.id == conversation_id).first() | |||||
| if not conversation: | if not conversation: | ||||
| return | return | ||||
| :param event: event | :param event: event | ||||
| :return: | :return: | ||||
| """ | """ | ||||
| message_file = db.session.query(MessageFile).filter(MessageFile.id == event.message_file_id).first() | |||||
| message_file = db.session.query(MessageFile).where(MessageFile.id == event.message_file_id).first() | |||||
| if message_file and message_file.url is not None: | if message_file and message_file.url is not None: | ||||
| # get tool file id | # get tool file id |
| for document in documents: | for document in documents: | ||||
| if document.metadata is not None: | if document.metadata is not None: | ||||
| document_id = document.metadata["document_id"] | document_id = document.metadata["document_id"] | ||||
| dataset_document = db.session.query(DatasetDocument).filter(DatasetDocument.id == document_id).first() | |||||
| dataset_document = db.session.query(DatasetDocument).where(DatasetDocument.id == document_id).first() | |||||
| if not dataset_document: | if not dataset_document: | ||||
| _logger.warning( | _logger.warning( | ||||
| "Expected DatasetDocument record to exist, but none was found, document_id=%s", | "Expected DatasetDocument record to exist, but none was found, document_id=%s", | ||||
| if dataset_document.doc_form == IndexType.PARENT_CHILD_INDEX: | if dataset_document.doc_form == IndexType.PARENT_CHILD_INDEX: | ||||
| child_chunk = ( | child_chunk = ( | ||||
| db.session.query(ChildChunk) | db.session.query(ChildChunk) | ||||
| .filter( | |||||
| .where( | |||||
| ChildChunk.index_node_id == document.metadata["doc_id"], | ChildChunk.index_node_id == document.metadata["doc_id"], | ||||
| ChildChunk.dataset_id == dataset_document.dataset_id, | ChildChunk.dataset_id == dataset_document.dataset_id, | ||||
| ChildChunk.document_id == dataset_document.id, | ChildChunk.document_id == dataset_document.id, | ||||
| if child_chunk: | if child_chunk: | ||||
| segment = ( | segment = ( | ||||
| db.session.query(DocumentSegment) | db.session.query(DocumentSegment) | ||||
| .filter(DocumentSegment.id == child_chunk.segment_id) | |||||
| .where(DocumentSegment.id == child_chunk.segment_id) | |||||
| .update( | .update( | ||||
| {DocumentSegment.hit_count: DocumentSegment.hit_count + 1}, synchronize_session=False | {DocumentSegment.hit_count: DocumentSegment.hit_count + 1}, synchronize_session=False | ||||
| ) | ) | ||||
| ) | ) | ||||
| else: | else: | ||||
| query = db.session.query(DocumentSegment).filter( | |||||
| query = db.session.query(DocumentSegment).where( | |||||
| DocumentSegment.index_node_id == document.metadata["doc_id"] | DocumentSegment.index_node_id == document.metadata["doc_id"] | ||||
| ) | ) | ||||
| if "dataset_id" in document.metadata: | if "dataset_id" in document.metadata: | ||||
| query = query.filter(DocumentSegment.dataset_id == document.metadata["dataset_id"]) | |||||
| query = query.where(DocumentSegment.dataset_id == document.metadata["dataset_id"]) | |||||
| # add hit count to document segment | # add hit count to document segment | ||||
| query.update({DocumentSegment.hit_count: DocumentSegment.hit_count + 1}, synchronize_session=False) | query.update({DocumentSegment.hit_count: DocumentSegment.hit_count + 1}, synchronize_session=False) |
| provider_record = ( | provider_record = ( | ||||
| db.session.query(Provider) | db.session.query(Provider) | ||||
| .filter( | |||||
| .where( | |||||
| Provider.tenant_id == self.tenant_id, | Provider.tenant_id == self.tenant_id, | ||||
| Provider.provider_type == ProviderType.CUSTOM.value, | Provider.provider_type == ProviderType.CUSTOM.value, | ||||
| Provider.provider_name.in_(provider_names), | Provider.provider_name.in_(provider_names), | ||||
| provider_model_record = ( | provider_model_record = ( | ||||
| db.session.query(ProviderModel) | db.session.query(ProviderModel) | ||||
| .filter( | |||||
| .where( | |||||
| ProviderModel.tenant_id == self.tenant_id, | ProviderModel.tenant_id == self.tenant_id, | ||||
| ProviderModel.provider_name.in_(provider_names), | ProviderModel.provider_name.in_(provider_names), | ||||
| ProviderModel.model_name == model, | ProviderModel.model_name == model, | ||||
| return ( | return ( | ||||
| db.session.query(ProviderModelSetting) | db.session.query(ProviderModelSetting) | ||||
| .filter( | |||||
| .where( | |||||
| ProviderModelSetting.tenant_id == self.tenant_id, | ProviderModelSetting.tenant_id == self.tenant_id, | ||||
| ProviderModelSetting.provider_name.in_(provider_names), | ProviderModelSetting.provider_name.in_(provider_names), | ||||
| ProviderModelSetting.model_type == model_type.to_origin_model_type(), | ProviderModelSetting.model_type == model_type.to_origin_model_type(), | ||||
| return ( | return ( | ||||
| db.session.query(LoadBalancingModelConfig) | db.session.query(LoadBalancingModelConfig) | ||||
| .filter( | |||||
| .where( | |||||
| LoadBalancingModelConfig.tenant_id == self.tenant_id, | LoadBalancingModelConfig.tenant_id == self.tenant_id, | ||||
| LoadBalancingModelConfig.provider_name.in_(provider_names), | LoadBalancingModelConfig.provider_name.in_(provider_names), | ||||
| LoadBalancingModelConfig.model_type == model_type.to_origin_model_type(), | LoadBalancingModelConfig.model_type == model_type.to_origin_model_type(), | ||||
| load_balancing_config_count = ( | load_balancing_config_count = ( | ||||
| db.session.query(LoadBalancingModelConfig) | db.session.query(LoadBalancingModelConfig) | ||||
| .filter( | |||||
| .where( | |||||
| LoadBalancingModelConfig.tenant_id == self.tenant_id, | LoadBalancingModelConfig.tenant_id == self.tenant_id, | ||||
| LoadBalancingModelConfig.provider_name.in_(provider_names), | LoadBalancingModelConfig.provider_name.in_(provider_names), | ||||
| LoadBalancingModelConfig.model_type == model_type.to_origin_model_type(), | LoadBalancingModelConfig.model_type == model_type.to_origin_model_type(), | ||||
| model_setting = ( | model_setting = ( | ||||
| db.session.query(ProviderModelSetting) | db.session.query(ProviderModelSetting) | ||||
| .filter( | |||||
| .where( | |||||
| ProviderModelSetting.tenant_id == self.tenant_id, | ProviderModelSetting.tenant_id == self.tenant_id, | ||||
| ProviderModelSetting.provider_name.in_(provider_names), | ProviderModelSetting.provider_name.in_(provider_names), | ||||
| ProviderModelSetting.model_type == model_type.to_origin_model_type(), | ProviderModelSetting.model_type == model_type.to_origin_model_type(), | ||||
| preferred_model_provider = ( | preferred_model_provider = ( | ||||
| db.session.query(TenantPreferredModelProvider) | db.session.query(TenantPreferredModelProvider) | ||||
| .filter( | |||||
| .where( | |||||
| TenantPreferredModelProvider.tenant_id == self.tenant_id, | TenantPreferredModelProvider.tenant_id == self.tenant_id, | ||||
| TenantPreferredModelProvider.provider_name.in_(provider_names), | TenantPreferredModelProvider.provider_name.in_(provider_names), | ||||
| ) | ) |
| # get api_based_extension | # get api_based_extension | ||||
| api_based_extension = ( | api_based_extension = ( | ||||
| db.session.query(APIBasedExtension) | db.session.query(APIBasedExtension) | ||||
| .filter(APIBasedExtension.tenant_id == tenant_id, APIBasedExtension.id == api_based_extension_id) | |||||
| .where(APIBasedExtension.tenant_id == tenant_id, APIBasedExtension.id == api_based_extension_id) | |||||
| .first() | .first() | ||||
| ) | ) | ||||
| # get api_based_extension | # get api_based_extension | ||||
| api_based_extension = ( | api_based_extension = ( | ||||
| db.session.query(APIBasedExtension) | db.session.query(APIBasedExtension) | ||||
| .filter(APIBasedExtension.tenant_id == self.tenant_id, APIBasedExtension.id == api_based_extension_id) | |||||
| .where(APIBasedExtension.tenant_id == self.tenant_id, APIBasedExtension.id == api_based_extension_id) | |||||
| .first() | .first() | ||||
| ) | ) | ||||
| from models.account import Tenant | from models.account import Tenant | ||||
| from models.engine import db | from models.engine import db | ||||
| if not (tenant := db.session.query(Tenant).filter(Tenant.id == tenant_id).first()): | |||||
| if not (tenant := db.session.query(Tenant).where(Tenant.id == tenant_id).first()): | |||||
| raise ValueError(f"Tenant with id {tenant_id} not found") | raise ValueError(f"Tenant with id {tenant_id} not found") | ||||
| encrypted_token = rsa.encrypt(token, tenant.encrypt_public_key) | encrypted_token = rsa.encrypt(token, tenant.encrypt_public_key) | ||||
| return base64.b64encode(encrypted_token).decode() | return base64.b64encode(encrypted_token).decode() |
| # get the process rule | # get the process rule | ||||
| processing_rule = ( | processing_rule = ( | ||||
| db.session.query(DatasetProcessRule) | db.session.query(DatasetProcessRule) | ||||
| .filter(DatasetProcessRule.id == dataset_document.dataset_process_rule_id) | |||||
| .where(DatasetProcessRule.id == dataset_document.dataset_process_rule_id) | |||||
| .first() | .first() | ||||
| ) | ) | ||||
| if not processing_rule: | if not processing_rule: | ||||
| db.session.delete(document_segment) | db.session.delete(document_segment) | ||||
| if dataset_document.doc_form == IndexType.PARENT_CHILD_INDEX: | if dataset_document.doc_form == IndexType.PARENT_CHILD_INDEX: | ||||
| # delete child chunks | # delete child chunks | ||||
| db.session.query(ChildChunk).filter(ChildChunk.segment_id == document_segment.id).delete() | |||||
| db.session.query(ChildChunk).where(ChildChunk.segment_id == document_segment.id).delete() | |||||
| db.session.commit() | db.session.commit() | ||||
| # get the process rule | # get the process rule | ||||
| processing_rule = ( | processing_rule = ( | ||||
| db.session.query(DatasetProcessRule) | db.session.query(DatasetProcessRule) | ||||
| .filter(DatasetProcessRule.id == dataset_document.dataset_process_rule_id) | |||||
| .where(DatasetProcessRule.id == dataset_document.dataset_process_rule_id) | |||||
| .first() | .first() | ||||
| ) | ) | ||||
| if not processing_rule: | if not processing_rule: | ||||
| # get the process rule | # get the process rule | ||||
| processing_rule = ( | processing_rule = ( | ||||
| db.session.query(DatasetProcessRule) | db.session.query(DatasetProcessRule) | ||||
| .filter(DatasetProcessRule.id == dataset_document.dataset_process_rule_id) | |||||
| .where(DatasetProcessRule.id == dataset_document.dataset_process_rule_id) | |||||
| .first() | .first() | ||||
| ) | ) | ||||
| # delete image files and related db records | # delete image files and related db records | ||||
| image_upload_file_ids = get_image_upload_file_ids(document.page_content) | image_upload_file_ids = get_image_upload_file_ids(document.page_content) | ||||
| for upload_file_id in image_upload_file_ids: | for upload_file_id in image_upload_file_ids: | ||||
| image_file = db.session.query(UploadFile).filter(UploadFile.id == upload_file_id).first() | |||||
| image_file = db.session.query(UploadFile).where(UploadFile.id == upload_file_id).first() | |||||
| if image_file is None: | if image_file is None: | ||||
| continue | continue | ||||
| try: | try: | ||||
| raise ValueError("no upload file found") | raise ValueError("no upload file found") | ||||
| file_detail = ( | file_detail = ( | ||||
| db.session.query(UploadFile).filter(UploadFile.id == data_source_info["upload_file_id"]).one_or_none() | |||||
| db.session.query(UploadFile).where(UploadFile.id == data_source_info["upload_file_id"]).one_or_none() | |||||
| ) | ) | ||||
| if file_detail: | if file_detail: | ||||
| keyword.create(documents) | keyword.create(documents) | ||||
| if dataset.indexing_technique != "high_quality": | if dataset.indexing_technique != "high_quality": | ||||
| document_ids = [document.metadata["doc_id"] for document in documents] | document_ids = [document.metadata["doc_id"] for document in documents] | ||||
| db.session.query(DocumentSegment).filter( | |||||
| db.session.query(DocumentSegment).where( | |||||
| DocumentSegment.document_id == document_id, | DocumentSegment.document_id == document_id, | ||||
| DocumentSegment.dataset_id == dataset_id, | DocumentSegment.dataset_id == dataset_id, | ||||
| DocumentSegment.index_node_id.in_(document_ids), | DocumentSegment.index_node_id.in_(document_ids), | ||||
| index_processor.load(dataset, chunk_documents, with_keywords=False) | index_processor.load(dataset, chunk_documents, with_keywords=False) | ||||
| document_ids = [document.metadata["doc_id"] for document in chunk_documents] | document_ids = [document.metadata["doc_id"] for document in chunk_documents] | ||||
| db.session.query(DocumentSegment).filter( | |||||
| db.session.query(DocumentSegment).where( | |||||
| DocumentSegment.document_id == dataset_document.id, | DocumentSegment.document_id == dataset_document.id, | ||||
| DocumentSegment.dataset_id == dataset.id, | DocumentSegment.dataset_id == dataset.id, | ||||
| DocumentSegment.index_node_id.in_(document_ids), | DocumentSegment.index_node_id.in_(document_ids), |
| ): | ): | ||||
| self.app = app | self.app = app | ||||
| self.request = request | self.request = request | ||||
| mcp_server = db.session.query(AppMCPServer).filter(AppMCPServer.app_id == self.app.id).first() | |||||
| mcp_server = db.session.query(AppMCPServer).where(AppMCPServer.app_id == self.app.id).first() | |||||
| if not mcp_server: | if not mcp_server: | ||||
| raise ValueError("MCP server not found") | raise ValueError("MCP server not found") | ||||
| self.mcp_server: AppMCPServer = mcp_server | self.mcp_server: AppMCPServer = mcp_server | ||||
| def retrieve_end_user(self): | def retrieve_end_user(self): | ||||
| return ( | return ( | ||||
| db.session.query(EndUser) | db.session.query(EndUser) | ||||
| .filter(EndUser.external_user_id == self.mcp_server.id, EndUser.type == "mcp") | |||||
| .where(EndUser.external_user_id == self.mcp_server.id, EndUser.type == "mcp") | |||||
| .first() | .first() | ||||
| ) | ) | ||||
| prompt_messages: list[PromptMessage] = [] | prompt_messages: list[PromptMessage] = [] | ||||
| for message in messages: | for message in messages: | ||||
| files = db.session.query(MessageFile).filter(MessageFile.message_id == message.id).all() | |||||
| files = db.session.query(MessageFile).where(MessageFile.message_id == message.id).all() | |||||
| if files: | if files: | ||||
| file_extra_config = None | file_extra_config = None | ||||
| if self.conversation.mode in {AppMode.AGENT_CHAT, AppMode.COMPLETION, AppMode.CHAT}: | if self.conversation.mode in {AppMode.AGENT_CHAT, AppMode.COMPLETION, AppMode.CHAT}: |
| def _get_api_based_extension(tenant_id: str, api_based_extension_id: str) -> Optional[APIBasedExtension]: | def _get_api_based_extension(tenant_id: str, api_based_extension_id: str) -> Optional[APIBasedExtension]: | ||||
| extension = ( | extension = ( | ||||
| db.session.query(APIBasedExtension) | db.session.query(APIBasedExtension) | ||||
| .filter(APIBasedExtension.tenant_id == tenant_id, APIBasedExtension.id == api_based_extension_id) | |||||
| .where(APIBasedExtension.tenant_id == tenant_id, APIBasedExtension.id == api_based_extension_id) | |||||
| .first() | .first() | ||||
| ) | ) | ||||
| user_id = message_data.from_account_id | user_id = message_data.from_account_id | ||||
| if message_data.from_end_user_id: | if message_data.from_end_user_id: | ||||
| end_user_data: Optional[EndUser] = ( | end_user_data: Optional[EndUser] = ( | ||||
| db.session.query(EndUser).filter(EndUser.id == message_data.from_end_user_id).first() | |||||
| db.session.query(EndUser).where(EndUser.id == message_data.from_end_user_id).first() | |||||
| ) | ) | ||||
| if end_user_data is not None: | if end_user_data is not None: | ||||
| user_id = end_user_data.session_id | user_id = end_user_data.session_id | ||||
| if not app_id: | if not app_id: | ||||
| raise ValueError("No app_id found in trace_info metadata") | raise ValueError("No app_id found in trace_info metadata") | ||||
| app = session.query(App).filter(App.id == app_id).first() | |||||
| app = session.query(App).where(App.id == app_id).first() | |||||
| if not app: | if not app: | ||||
| raise ValueError(f"App with id {app_id} not found") | raise ValueError(f"App with id {app_id} not found") | ||||
| if not app.created_by: | if not app.created_by: | ||||
| raise ValueError(f"App with id {app_id} has no creator (created_by is None)") | raise ValueError(f"App with id {app_id} has no creator (created_by is None)") | ||||
| service_account = session.query(Account).filter(Account.id == app.created_by).first() | |||||
| service_account = session.query(Account).where(Account.id == app.created_by).first() | |||||
| if not service_account: | if not service_account: | ||||
| raise ValueError(f"Creator account with id {app.created_by} not found for app {app_id}") | raise ValueError(f"Creator account with id {app.created_by} not found for app {app_id}") | ||||
| current_tenant = ( | current_tenant = ( |
| # Add end user data if available | # Add end user data if available | ||||
| if trace_info.message_data.from_end_user_id: | if trace_info.message_data.from_end_user_id: | ||||
| end_user_data: Optional[EndUser] = ( | end_user_data: Optional[EndUser] = ( | ||||
| db.session.query(EndUser).filter(EndUser.id == trace_info.message_data.from_end_user_id).first() | |||||
| db.session.query(EndUser).where(EndUser.id == trace_info.message_data.from_end_user_id).first() | |||||
| ) | ) | ||||
| if end_user_data is not None: | if end_user_data is not None: | ||||
| message_metadata["end_user_id"] = end_user_data.session_id | message_metadata["end_user_id"] = end_user_data.session_id | ||||
| WorkflowNodeExecutionModel.process_data, | WorkflowNodeExecutionModel.process_data, | ||||
| WorkflowNodeExecutionModel.execution_metadata, | WorkflowNodeExecutionModel.execution_metadata, | ||||
| ) | ) | ||||
| .filter(WorkflowNodeExecutionModel.workflow_run_id == workflow_run_id) | |||||
| .where(WorkflowNodeExecutionModel.workflow_run_id == workflow_run_id) | |||||
| .all() | .all() | ||||
| ) | ) | ||||
| return workflow_nodes | return workflow_nodes |
| """ | """ | ||||
| with Session(db.engine, expire_on_commit=False) as session: | with Session(db.engine, expire_on_commit=False) as session: | ||||
| # Get the app to find its creator | # Get the app to find its creator | ||||
| app = session.query(App).filter(App.id == app_id).first() | |||||
| app = session.query(App).where(App.id == app_id).first() | |||||
| if not app: | if not app: | ||||
| raise ValueError(f"App with id {app_id} not found") | raise ValueError(f"App with id {app_id} not found") | ||||
| if not app.created_by: | if not app.created_by: | ||||
| raise ValueError(f"App with id {app_id} has no creator (created_by is None)") | raise ValueError(f"App with id {app_id} has no creator (created_by is None)") | ||||
| service_account = session.query(Account).filter(Account.id == app.created_by).first() | |||||
| service_account = session.query(Account).where(Account.id == app.created_by).first() | |||||
| if not service_account: | if not service_account: | ||||
| raise ValueError(f"Creator account with id {app.created_by} not found for app {app_id}") | raise ValueError(f"Creator account with id {app.created_by} not found for app {app_id}") | ||||
| user_id = message_data.from_account_id | user_id = message_data.from_account_id | ||||
| if message_data.from_end_user_id: | if message_data.from_end_user_id: | ||||
| end_user_data: Optional[EndUser] = ( | end_user_data: Optional[EndUser] = ( | ||||
| db.session.query(EndUser).filter(EndUser.id == message_data.from_end_user_id).first() | |||||
| db.session.query(EndUser).where(EndUser.id == message_data.from_end_user_id).first() | |||||
| ) | ) | ||||
| if end_user_data is not None: | if end_user_data is not None: | ||||
| user_id = end_user_data.session_id | user_id = end_user_data.session_id |
| if message_data.from_end_user_id: | if message_data.from_end_user_id: | ||||
| end_user_data: Optional[EndUser] = ( | end_user_data: Optional[EndUser] = ( | ||||
| db.session.query(EndUser).filter(EndUser.id == message_data.from_end_user_id).first() | |||||
| db.session.query(EndUser).where(EndUser.id == message_data.from_end_user_id).first() | |||||
| ) | ) | ||||
| if end_user_data is not None: | if end_user_data is not None: | ||||
| end_user_id = end_user_data.session_id | end_user_id = end_user_data.session_id |
| if message_data.from_end_user_id: | if message_data.from_end_user_id: | ||||
| end_user_data: Optional[EndUser] = ( | end_user_data: Optional[EndUser] = ( | ||||
| db.session.query(EndUser).filter(EndUser.id == message_data.from_end_user_id).first() | |||||
| db.session.query(EndUser).where(EndUser.id == message_data.from_end_user_id).first() | |||||
| ) | ) | ||||
| if end_user_data is not None: | if end_user_data is not None: | ||||
| end_user_id = end_user_data.session_id | end_user_id = end_user_data.session_id |
| """ | """ | ||||
| trace_config_data: Optional[TraceAppConfig] = ( | trace_config_data: Optional[TraceAppConfig] = ( | ||||
| db.session.query(TraceAppConfig) | db.session.query(TraceAppConfig) | ||||
| .filter(TraceAppConfig.app_id == app_id, TraceAppConfig.tracing_provider == tracing_provider) | |||||
| .where(TraceAppConfig.app_id == app_id, TraceAppConfig.tracing_provider == tracing_provider) | |||||
| .first() | .first() | ||||
| ) | ) | ||||
| return None | return None | ||||
| # decrypt_token | # decrypt_token | ||||
| app = db.session.query(App).filter(App.id == app_id).first() | |||||
| app = db.session.query(App).where(App.id == app_id).first() | |||||
| if not app: | if not app: | ||||
| raise ValueError("App not found") | raise ValueError("App not found") | ||||
| if app_id is None: | if app_id is None: | ||||
| return None | return None | ||||
| app: Optional[App] = db.session.query(App).filter(App.id == app_id).first() | |||||
| app: Optional[App] = db.session.query(App).where(App.id == app_id).first() | |||||
| if app is None: | if app is None: | ||||
| return None | return None | ||||
| @classmethod | @classmethod | ||||
| def get_app_config_through_message_id(cls, message_id: str): | def get_app_config_through_message_id(cls, message_id: str): | ||||
| app_model_config = None | app_model_config = None | ||||
| message_data = db.session.query(Message).filter(Message.id == message_id).first() | |||||
| message_data = db.session.query(Message).where(Message.id == message_id).first() | |||||
| if not message_data: | if not message_data: | ||||
| return None | return None | ||||
| conversation_id = message_data.conversation_id | conversation_id = message_data.conversation_id | ||||
| conversation_data = db.session.query(Conversation).filter(Conversation.id == conversation_id).first() | |||||
| conversation_data = db.session.query(Conversation).where(Conversation.id == conversation_id).first() | |||||
| if not conversation_data: | if not conversation_data: | ||||
| return None | return None | ||||
| if conversation_data.app_model_config_id: | if conversation_data.app_model_config_id: | ||||
| app_model_config = ( | app_model_config = ( | ||||
| db.session.query(AppModelConfig) | db.session.query(AppModelConfig) | ||||
| .filter(AppModelConfig.id == conversation_data.app_model_config_id) | |||||
| .where(AppModelConfig.id == conversation_data.app_model_config_id) | |||||
| .first() | .first() | ||||
| ) | ) | ||||
| elif conversation_data.app_model_config_id is None and conversation_data.override_model_configs: | elif conversation_data.app_model_config_id is None and conversation_data.override_model_configs: | ||||
| if tracing_provider is not None: | if tracing_provider is not None: | ||||
| raise ValueError(f"Invalid tracing provider: {tracing_provider}") | raise ValueError(f"Invalid tracing provider: {tracing_provider}") | ||||
| app_config: Optional[App] = db.session.query(App).filter(App.id == app_id).first() | |||||
| app_config: Optional[App] = db.session.query(App).where(App.id == app_id).first() | |||||
| if not app_config: | if not app_config: | ||||
| raise ValueError("App not found") | raise ValueError("App not found") | ||||
| app_config.tracing = json.dumps( | app_config.tracing = json.dumps( | ||||
| :param app_id: app id | :param app_id: app id | ||||
| :return: | :return: | ||||
| """ | """ | ||||
| app: Optional[App] = db.session.query(App).filter(App.id == app_id).first() | |||||
| app: Optional[App] = db.session.query(App).where(App.id == app_id).first() | |||||
| if not app: | if not app: | ||||
| raise ValueError("App not found") | raise ValueError("App not found") | ||||
| if not app.tracing: | if not app.tracing: |
| from typing import Optional, Union | from typing import Optional, Union | ||||
| from urllib.parse import urlparse | from urllib.parse import urlparse | ||||
| from sqlalchemy import select | |||||
| from extensions.ext_database import db | from extensions.ext_database import db | ||||
| from models.model import Message | from models.model import Message | ||||
| def get_message_data(message_id: str): | def get_message_data(message_id: str): | ||||
| return db.session.query(Message).filter(Message.id == message_id).first() | |||||
| return db.session.scalar(select(Message).where(Message.id == message_id)) | |||||
| @contextmanager | @contextmanager |
| if message_data.from_end_user_id: | if message_data.from_end_user_id: | ||||
| end_user_data: Optional[EndUser] = ( | end_user_data: Optional[EndUser] = ( | ||||
| db.session.query(EndUser).filter(EndUser.id == message_data.from_end_user_id).first() | |||||
| db.session.query(EndUser).where(EndUser.id == message_data.from_end_user_id).first() | |||||
| ) | ) | ||||
| if end_user_data is not None: | if end_user_data is not None: | ||||
| end_user_id = end_user_data.session_id | end_user_id = end_user_data.session_id |
| get the user by user id | get the user by user id | ||||
| """ | """ | ||||
| user = db.session.query(EndUser).filter(EndUser.id == user_id).first() | |||||
| user = db.session.query(EndUser).where(EndUser.id == user_id).first() | |||||
| if not user: | if not user: | ||||
| user = db.session.query(Account).filter(Account.id == user_id).first() | |||||
| user = db.session.query(Account).where(Account.id == user_id).first() | |||||
| if not user: | if not user: | ||||
| raise ValueError("user not found") | raise ValueError("user not found") | ||||
| get app | get app | ||||
| """ | """ | ||||
| try: | try: | ||||
| app = db.session.query(App).filter(App.id == app_id).filter(App.tenant_id == tenant_id).first() | |||||
| app = db.session.query(App).where(App.id == app_id).where(App.tenant_id == tenant_id).first() | |||||
| except Exception: | except Exception: | ||||
| raise ValueError("app not found") | raise ValueError("app not found") | ||||
| # Get the corresponding TenantDefaultModel record | # Get the corresponding TenantDefaultModel record | ||||
| default_model = ( | default_model = ( | ||||
| db.session.query(TenantDefaultModel) | db.session.query(TenantDefaultModel) | ||||
| .filter( | |||||
| .where( | |||||
| TenantDefaultModel.tenant_id == tenant_id, | TenantDefaultModel.tenant_id == tenant_id, | ||||
| TenantDefaultModel.model_type == model_type.to_origin_model_type(), | TenantDefaultModel.model_type == model_type.to_origin_model_type(), | ||||
| ) | ) | ||||
| # Get the list of available models from get_configurations and check if it is LLM | # Get the list of available models from get_configurations and check if it is LLM | ||||
| default_model = ( | default_model = ( | ||||
| db.session.query(TenantDefaultModel) | db.session.query(TenantDefaultModel) | ||||
| .filter( | |||||
| .where( | |||||
| TenantDefaultModel.tenant_id == tenant_id, | TenantDefaultModel.tenant_id == tenant_id, | ||||
| TenantDefaultModel.model_type == model_type.to_origin_model_type(), | TenantDefaultModel.model_type == model_type.to_origin_model_type(), | ||||
| ) | ) | ||||
| db.session.rollback() | db.session.rollback() | ||||
| existed_provider_record = ( | existed_provider_record = ( | ||||
| db.session.query(Provider) | db.session.query(Provider) | ||||
| .filter( | |||||
| .where( | |||||
| Provider.tenant_id == tenant_id, | Provider.tenant_id == tenant_id, | ||||
| Provider.provider_name == ModelProviderID(provider_name).provider_name, | Provider.provider_name == ModelProviderID(provider_name).provider_name, | ||||
| Provider.provider_type == ProviderType.SYSTEM.value, | Provider.provider_type == ProviderType.SYSTEM.value, |
| documents = [] | documents = [] | ||||
| for chunk_index in sorted_chunk_indices: | for chunk_index in sorted_chunk_indices: | ||||
| segment_query = db.session.query(DocumentSegment).filter( | |||||
| segment_query = db.session.query(DocumentSegment).where( | |||||
| DocumentSegment.dataset_id == self.dataset.id, DocumentSegment.index_node_id == chunk_index | DocumentSegment.dataset_id == self.dataset.id, DocumentSegment.index_node_id == chunk_index | ||||
| ) | ) | ||||
| if document_ids_filter: | if document_ids_filter: | ||||
| segment_query = segment_query.filter(DocumentSegment.document_id.in_(document_ids_filter)) | |||||
| segment_query = segment_query.where(DocumentSegment.document_id.in_(document_ids_filter)) | |||||
| segment = segment_query.first() | segment = segment_query.first() | ||||
| if segment: | if segment: | ||||
| def _update_segment_keywords(self, dataset_id: str, node_id: str, keywords: list[str]): | def _update_segment_keywords(self, dataset_id: str, node_id: str, keywords: list[str]): | ||||
| document_segment = ( | document_segment = ( | ||||
| db.session.query(DocumentSegment) | db.session.query(DocumentSegment) | ||||
| .filter(DocumentSegment.dataset_id == dataset_id, DocumentSegment.index_node_id == node_id) | |||||
| .where(DocumentSegment.dataset_id == dataset_id, DocumentSegment.index_node_id == node_id) | |||||
| .first() | .first() | ||||
| ) | ) | ||||
| if document_segment: | if document_segment: |
| external_retrieval_model: Optional[dict] = None, | external_retrieval_model: Optional[dict] = None, | ||||
| metadata_filtering_conditions: Optional[dict] = None, | metadata_filtering_conditions: Optional[dict] = None, | ||||
| ): | ): | ||||
| dataset = db.session.query(Dataset).filter(Dataset.id == dataset_id).first() | |||||
| dataset = db.session.query(Dataset).where(Dataset.id == dataset_id).first() | |||||
| if not dataset: | if not dataset: | ||||
| return [] | return [] | ||||
| metadata_condition = ( | metadata_condition = ( | ||||
| @classmethod | @classmethod | ||||
| def _get_dataset(cls, dataset_id: str) -> Optional[Dataset]: | def _get_dataset(cls, dataset_id: str) -> Optional[Dataset]: | ||||
| with Session(db.engine) as session: | with Session(db.engine) as session: | ||||
| return session.query(Dataset).filter(Dataset.id == dataset_id).first() | |||||
| return session.query(Dataset).where(Dataset.id == dataset_id).first() | |||||
| @classmethod | @classmethod | ||||
| def keyword_search( | def keyword_search( | ||||
| dataset_documents = { | dataset_documents = { | ||||
| doc.id: doc | doc.id: doc | ||||
| for doc in db.session.query(DatasetDocument) | for doc in db.session.query(DatasetDocument) | ||||
| .filter(DatasetDocument.id.in_(document_ids)) | |||||
| .where(DatasetDocument.id.in_(document_ids)) | |||||
| .options(load_only(DatasetDocument.id, DatasetDocument.doc_form, DatasetDocument.dataset_id)) | .options(load_only(DatasetDocument.id, DatasetDocument.doc_form, DatasetDocument.dataset_id)) | ||||
| .all() | .all() | ||||
| } | } | ||||
| child_index_node_id = document.metadata.get("doc_id") | child_index_node_id = document.metadata.get("doc_id") | ||||
| child_chunk = ( | child_chunk = ( | ||||
| db.session.query(ChildChunk).filter(ChildChunk.index_node_id == child_index_node_id).first() | |||||
| db.session.query(ChildChunk).where(ChildChunk.index_node_id == child_index_node_id).first() | |||||
| ) | ) | ||||
| if not child_chunk: | if not child_chunk: | ||||
| segment = ( | segment = ( | ||||
| db.session.query(DocumentSegment) | db.session.query(DocumentSegment) | ||||
| .filter( | |||||
| .where( | |||||
| DocumentSegment.dataset_id == dataset_document.dataset_id, | DocumentSegment.dataset_id == dataset_document.dataset_id, | ||||
| DocumentSegment.enabled == True, | DocumentSegment.enabled == True, | ||||
| DocumentSegment.status == "completed", | DocumentSegment.status == "completed", | ||||
| segment = ( | segment = ( | ||||
| db.session.query(DocumentSegment) | db.session.query(DocumentSegment) | ||||
| .filter( | |||||
| .where( | |||||
| DocumentSegment.dataset_id == dataset_document.dataset_id, | DocumentSegment.dataset_id == dataset_document.dataset_id, | ||||
| DocumentSegment.enabled == True, | DocumentSegment.enabled == True, | ||||
| DocumentSegment.status == "completed", | DocumentSegment.status == "completed", |
| if dataset.collection_binding_id: | if dataset.collection_binding_id: | ||||
| dataset_collection_binding = ( | dataset_collection_binding = ( | ||||
| db.session.query(DatasetCollectionBinding) | db.session.query(DatasetCollectionBinding) | ||||
| .filter(DatasetCollectionBinding.id == dataset.collection_binding_id) | |||||
| .where(DatasetCollectionBinding.id == dataset.collection_binding_id) | |||||
| .one_or_none() | .one_or_none() | ||||
| ) | ) | ||||
| if dataset_collection_binding: | if dataset_collection_binding: |
| class TidbOnQdrantVectorFactory(AbstractVectorFactory): | class TidbOnQdrantVectorFactory(AbstractVectorFactory): | ||||
| def init_vector(self, dataset: Dataset, attributes: list, embeddings: Embeddings) -> TidbOnQdrantVector: | def init_vector(self, dataset: Dataset, attributes: list, embeddings: Embeddings) -> TidbOnQdrantVector: | ||||
| tidb_auth_binding = ( | tidb_auth_binding = ( | ||||
| db.session.query(TidbAuthBinding).filter(TidbAuthBinding.tenant_id == dataset.tenant_id).one_or_none() | |||||
| db.session.query(TidbAuthBinding).where(TidbAuthBinding.tenant_id == dataset.tenant_id).one_or_none() | |||||
| ) | ) | ||||
| if not tidb_auth_binding: | if not tidb_auth_binding: | ||||
| with redis_client.lock("create_tidb_serverless_cluster_lock", timeout=900): | with redis_client.lock("create_tidb_serverless_cluster_lock", timeout=900): | ||||
| tidb_auth_binding = ( | tidb_auth_binding = ( | ||||
| db.session.query(TidbAuthBinding) | db.session.query(TidbAuthBinding) | ||||
| .filter(TidbAuthBinding.tenant_id == dataset.tenant_id) | |||||
| .where(TidbAuthBinding.tenant_id == dataset.tenant_id) | |||||
| .one_or_none() | .one_or_none() | ||||
| ) | ) | ||||
| if tidb_auth_binding: | if tidb_auth_binding: | ||||
| else: | else: | ||||
| idle_tidb_auth_binding = ( | idle_tidb_auth_binding = ( | ||||
| db.session.query(TidbAuthBinding) | db.session.query(TidbAuthBinding) | ||||
| .filter(TidbAuthBinding.active == False, TidbAuthBinding.status == "ACTIVE") | |||||
| .where(TidbAuthBinding.active == False, TidbAuthBinding.status == "ACTIVE") | |||||
| .limit(1) | .limit(1) | ||||
| .one_or_none() | .one_or_none() | ||||
| ) | ) |
| if dify_config.VECTOR_STORE_WHITELIST_ENABLE: | if dify_config.VECTOR_STORE_WHITELIST_ENABLE: | ||||
| whitelist = ( | whitelist = ( | ||||
| db.session.query(Whitelist) | db.session.query(Whitelist) | ||||
| .filter(Whitelist.tenant_id == self._dataset.tenant_id, Whitelist.category == "vector_db") | |||||
| .where(Whitelist.tenant_id == self._dataset.tenant_id, Whitelist.category == "vector_db") | |||||
| .one_or_none() | .one_or_none() | ||||
| ) | ) | ||||
| if whitelist: | if whitelist: |
| @property | @property | ||||
| def docs(self) -> dict[str, Document]: | def docs(self) -> dict[str, Document]: | ||||
| document_segments = ( | document_segments = ( | ||||
| db.session.query(DocumentSegment).filter(DocumentSegment.dataset_id == self._dataset.id).all() | |||||
| db.session.query(DocumentSegment).where(DocumentSegment.dataset_id == self._dataset.id).all() | |||||
| ) | ) | ||||
| output = {} | output = {} | ||||
| def add_documents(self, docs: Sequence[Document], allow_update: bool = True, save_child: bool = False) -> None: | def add_documents(self, docs: Sequence[Document], allow_update: bool = True, save_child: bool = False) -> None: | ||||
| max_position = ( | max_position = ( | ||||
| db.session.query(func.max(DocumentSegment.position)) | db.session.query(func.max(DocumentSegment.position)) | ||||
| .filter(DocumentSegment.document_id == self._document_id) | |||||
| .where(DocumentSegment.document_id == self._document_id) | |||||
| .scalar() | .scalar() | ||||
| ) | ) | ||||
| segment_document.tokens = tokens | segment_document.tokens = tokens | ||||
| if save_child and doc.children: | if save_child and doc.children: | ||||
| # delete the existing child chunks | # delete the existing child chunks | ||||
| db.session.query(ChildChunk).filter( | |||||
| db.session.query(ChildChunk).where( | |||||
| ChildChunk.tenant_id == self._dataset.tenant_id, | ChildChunk.tenant_id == self._dataset.tenant_id, | ||||
| ChildChunk.dataset_id == self._dataset.id, | ChildChunk.dataset_id == self._dataset.id, | ||||
| ChildChunk.document_id == self._document_id, | ChildChunk.document_id == self._document_id, | ||||
| def get_document_segment(self, doc_id: str) -> Optional[DocumentSegment]: | def get_document_segment(self, doc_id: str) -> Optional[DocumentSegment]: | ||||
| document_segment = ( | document_segment = ( | ||||
| db.session.query(DocumentSegment) | db.session.query(DocumentSegment) | ||||
| .filter(DocumentSegment.dataset_id == self._dataset.id, DocumentSegment.index_node_id == doc_id) | |||||
| .where(DocumentSegment.dataset_id == self._dataset.id, DocumentSegment.index_node_id == doc_id) | |||||
| .first() | .first() | ||||
| ) | ) | ||||
| def _get_access_token(cls, tenant_id: str, notion_workspace_id: str) -> str: | def _get_access_token(cls, tenant_id: str, notion_workspace_id: str) -> str: | ||||
| data_source_binding = ( | data_source_binding = ( | ||||
| db.session.query(DataSourceOauthBinding) | db.session.query(DataSourceOauthBinding) | ||||
| .filter( | |||||
| .where( | |||||
| db.and_( | db.and_( | ||||
| DataSourceOauthBinding.tenant_id == tenant_id, | DataSourceOauthBinding.tenant_id == tenant_id, | ||||
| DataSourceOauthBinding.provider == "notion", | DataSourceOauthBinding.provider == "notion", |
| child_node_ids = ( | child_node_ids = ( | ||||
| db.session.query(ChildChunk.index_node_id) | db.session.query(ChildChunk.index_node_id) | ||||
| .join(DocumentSegment, ChildChunk.segment_id == DocumentSegment.id) | .join(DocumentSegment, ChildChunk.segment_id == DocumentSegment.id) | ||||
| .filter( | |||||
| .where( | |||||
| DocumentSegment.dataset_id == dataset.id, | DocumentSegment.dataset_id == dataset.id, | ||||
| DocumentSegment.index_node_id.in_(node_ids), | DocumentSegment.index_node_id.in_(node_ids), | ||||
| ChildChunk.dataset_id == dataset.id, | ChildChunk.dataset_id == dataset.id, | ||||
| child_node_ids = [child_node_id[0] for child_node_id in child_node_ids] | child_node_ids = [child_node_id[0] for child_node_id in child_node_ids] | ||||
| vector.delete_by_ids(child_node_ids) | vector.delete_by_ids(child_node_ids) | ||||
| if delete_child_chunks: | if delete_child_chunks: | ||||
| db.session.query(ChildChunk).filter( | |||||
| db.session.query(ChildChunk).where( | |||||
| ChildChunk.dataset_id == dataset.id, ChildChunk.index_node_id.in_(child_node_ids) | ChildChunk.dataset_id == dataset.id, ChildChunk.index_node_id.in_(child_node_ids) | ||||
| ).delete() | ).delete() | ||||
| db.session.commit() | db.session.commit() | ||||
| vector.delete() | vector.delete() | ||||
| if delete_child_chunks: | if delete_child_chunks: | ||||
| db.session.query(ChildChunk).filter(ChildChunk.dataset_id == dataset.id).delete() | |||||
| db.session.query(ChildChunk).where(ChildChunk.dataset_id == dataset.id).delete() | |||||
| db.session.commit() | db.session.commit() | ||||
| def retrieve( | def retrieve( |
| available_datasets = [] | available_datasets = [] | ||||
| for dataset_id in dataset_ids: | for dataset_id in dataset_ids: | ||||
| # get dataset from dataset id | # get dataset from dataset id | ||||
| dataset = db.session.query(Dataset).filter(Dataset.tenant_id == tenant_id, Dataset.id == dataset_id).first() | |||||
| dataset = db.session.query(Dataset).where(Dataset.tenant_id == tenant_id, Dataset.id == dataset_id).first() | |||||
| # pass if dataset is not available | # pass if dataset is not available | ||||
| if not dataset: | if not dataset: | ||||
| dataset = db.session.query(Dataset).filter_by(id=segment.dataset_id).first() | dataset = db.session.query(Dataset).filter_by(id=segment.dataset_id).first() | ||||
| document = ( | document = ( | ||||
| db.session.query(DatasetDocument) | db.session.query(DatasetDocument) | ||||
| .filter( | |||||
| .where( | |||||
| DatasetDocument.id == segment.document_id, | DatasetDocument.id == segment.document_id, | ||||
| DatasetDocument.enabled == True, | DatasetDocument.enabled == True, | ||||
| DatasetDocument.archived == False, | DatasetDocument.archived == False, | ||||
| if dataset_id: | if dataset_id: | ||||
| # get retrieval model config | # get retrieval model config | ||||
| dataset = db.session.query(Dataset).filter(Dataset.id == dataset_id).first() | |||||
| dataset = db.session.query(Dataset).where(Dataset.id == dataset_id).first() | |||||
| if dataset: | if dataset: | ||||
| results = [] | results = [] | ||||
| if dataset.provider == "external": | if dataset.provider == "external": | ||||
| if document.metadata is not None: | if document.metadata is not None: | ||||
| dataset_document = ( | dataset_document = ( | ||||
| db.session.query(DatasetDocument) | db.session.query(DatasetDocument) | ||||
| .filter(DatasetDocument.id == document.metadata["document_id"]) | |||||
| .where(DatasetDocument.id == document.metadata["document_id"]) | |||||
| .first() | .first() | ||||
| ) | ) | ||||
| if dataset_document: | if dataset_document: | ||||
| if dataset_document.doc_form == IndexType.PARENT_CHILD_INDEX: | if dataset_document.doc_form == IndexType.PARENT_CHILD_INDEX: | ||||
| child_chunk = ( | child_chunk = ( | ||||
| db.session.query(ChildChunk) | db.session.query(ChildChunk) | ||||
| .filter( | |||||
| .where( | |||||
| ChildChunk.index_node_id == document.metadata["doc_id"], | ChildChunk.index_node_id == document.metadata["doc_id"], | ||||
| ChildChunk.dataset_id == dataset_document.dataset_id, | ChildChunk.dataset_id == dataset_document.dataset_id, | ||||
| ChildChunk.document_id == dataset_document.id, | ChildChunk.document_id == dataset_document.id, | ||||
| if child_chunk: | if child_chunk: | ||||
| segment = ( | segment = ( | ||||
| db.session.query(DocumentSegment) | db.session.query(DocumentSegment) | ||||
| .filter(DocumentSegment.id == child_chunk.segment_id) | |||||
| .where(DocumentSegment.id == child_chunk.segment_id) | |||||
| .update( | .update( | ||||
| {DocumentSegment.hit_count: DocumentSegment.hit_count + 1}, | {DocumentSegment.hit_count: DocumentSegment.hit_count + 1}, | ||||
| synchronize_session=False, | synchronize_session=False, | ||||
| ) | ) | ||||
| db.session.commit() | db.session.commit() | ||||
| else: | else: | ||||
| query = db.session.query(DocumentSegment).filter( | |||||
| query = db.session.query(DocumentSegment).where( | |||||
| DocumentSegment.index_node_id == document.metadata["doc_id"] | DocumentSegment.index_node_id == document.metadata["doc_id"] | ||||
| ) | ) | ||||
| # if 'dataset_id' in document.metadata: | # if 'dataset_id' in document.metadata: | ||||
| if "dataset_id" in document.metadata: | if "dataset_id" in document.metadata: | ||||
| query = query.filter(DocumentSegment.dataset_id == document.metadata["dataset_id"]) | |||||
| query = query.where(DocumentSegment.dataset_id == document.metadata["dataset_id"]) | |||||
| # add hit count to document segment | # add hit count to document segment | ||||
| query.update( | query.update( | ||||
| ): | ): | ||||
| with flask_app.app_context(): | with flask_app.app_context(): | ||||
| with Session(db.engine) as session: | with Session(db.engine) as session: | ||||
| dataset = session.query(Dataset).filter(Dataset.id == dataset_id).first() | |||||
| dataset = session.query(Dataset).where(Dataset.id == dataset_id).first() | |||||
| if not dataset: | if not dataset: | ||||
| return [] | return [] | ||||
| available_datasets = [] | available_datasets = [] | ||||
| for dataset_id in dataset_ids: | for dataset_id in dataset_ids: | ||||
| # get dataset from dataset id | # get dataset from dataset id | ||||
| dataset = db.session.query(Dataset).filter(Dataset.tenant_id == tenant_id, Dataset.id == dataset_id).first() | |||||
| dataset = db.session.query(Dataset).where(Dataset.tenant_id == tenant_id, Dataset.id == dataset_id).first() | |||||
| # pass if dataset is not available | # pass if dataset is not available | ||||
| if not dataset: | if not dataset: | ||||
| metadata_filtering_conditions: Optional[MetadataFilteringCondition], | metadata_filtering_conditions: Optional[MetadataFilteringCondition], | ||||
| inputs: dict, | inputs: dict, | ||||
| ) -> tuple[Optional[dict[str, list[str]]], Optional[MetadataCondition]]: | ) -> tuple[Optional[dict[str, list[str]]], Optional[MetadataCondition]]: | ||||
| document_query = db.session.query(DatasetDocument).filter( | |||||
| document_query = db.session.query(DatasetDocument).where( | |||||
| DatasetDocument.dataset_id.in_(dataset_ids), | DatasetDocument.dataset_id.in_(dataset_ids), | ||||
| DatasetDocument.indexing_status == "completed", | DatasetDocument.indexing_status == "completed", | ||||
| DatasetDocument.enabled == True, | DatasetDocument.enabled == True, | ||||
| raise ValueError("Invalid metadata filtering mode") | raise ValueError("Invalid metadata filtering mode") | ||||
| if filters: | if filters: | ||||
| if metadata_filtering_conditions and metadata_filtering_conditions.logical_operator == "and": # type: ignore | if metadata_filtering_conditions and metadata_filtering_conditions.logical_operator == "and": # type: ignore | ||||
| document_query = document_query.filter(and_(*filters)) | |||||
| document_query = document_query.where(and_(*filters)) | |||||
| else: | else: | ||||
| document_query = document_query.filter(or_(*filters)) | |||||
| document_query = document_query.where(or_(*filters)) | |||||
| documents = document_query.all() | documents = document_query.all() | ||||
| # group by dataset_id | # group by dataset_id | ||||
| metadata_filter_document_ids = defaultdict(list) if documents else None # type: ignore | metadata_filter_document_ids = defaultdict(list) if documents else None # type: ignore | ||||
| self, dataset_ids: list, query: str, tenant_id: str, user_id: str, metadata_model_config: ModelConfig | self, dataset_ids: list, query: str, tenant_id: str, user_id: str, metadata_model_config: ModelConfig | ||||
| ) -> Optional[list[dict[str, Any]]]: | ) -> Optional[list[dict[str, Any]]]: | ||||
| # get all metadata field | # get all metadata field | ||||
| metadata_fields = db.session.query(DatasetMetadata).filter(DatasetMetadata.dataset_id.in_(dataset_ids)).all() | |||||
| metadata_fields = db.session.query(DatasetMetadata).where(DatasetMetadata.dataset_id.in_(dataset_ids)).all() | |||||
| all_metadata_fields = [metadata_field.name for metadata_field in metadata_fields] | all_metadata_fields = [metadata_field.name for metadata_field in metadata_fields] | ||||
| # get metadata model config | # get metadata model config | ||||
| if metadata_model_config is None: | if metadata_model_config is None: |
| # get tenant api providers | # get tenant api providers | ||||
| db_providers: list[ApiToolProvider] = ( | db_providers: list[ApiToolProvider] = ( | ||||
| db.session.query(ApiToolProvider) | db.session.query(ApiToolProvider) | ||||
| .filter(ApiToolProvider.tenant_id == tenant_id, ApiToolProvider.name == self.entity.identity.name) | |||||
| .where(ApiToolProvider.tenant_id == tenant_id, ApiToolProvider.name == self.entity.identity.name) | |||||
| .all() | .all() | ||||
| ) | ) | ||||
| with Session(self._engine, expire_on_commit=False) as session: | with Session(self._engine, expire_on_commit=False) as session: | ||||
| tool_file: ToolFile | None = ( | tool_file: ToolFile | None = ( | ||||
| session.query(ToolFile) | session.query(ToolFile) | ||||
| .filter( | |||||
| .where( | |||||
| ToolFile.id == id, | ToolFile.id == id, | ||||
| ) | ) | ||||
| .first() | .first() | ||||
| with Session(self._engine, expire_on_commit=False) as session: | with Session(self._engine, expire_on_commit=False) as session: | ||||
| message_file: MessageFile | None = ( | message_file: MessageFile | None = ( | ||||
| session.query(MessageFile) | session.query(MessageFile) | ||||
| .filter( | |||||
| .where( | |||||
| MessageFile.id == id, | MessageFile.id == id, | ||||
| ) | ) | ||||
| .first() | .first() | ||||
| tool_file: ToolFile | None = ( | tool_file: ToolFile | None = ( | ||||
| session.query(ToolFile) | session.query(ToolFile) | ||||
| .filter( | |||||
| .where( | |||||
| ToolFile.id == tool_file_id, | ToolFile.id == tool_file_id, | ||||
| ) | ) | ||||
| .first() | .first() | ||||
| with Session(self._engine, expire_on_commit=False) as session: | with Session(self._engine, expire_on_commit=False) as session: | ||||
| tool_file: ToolFile | None = ( | tool_file: ToolFile | None = ( | ||||
| session.query(ToolFile) | session.query(ToolFile) | ||||
| .filter( | |||||
| .where( | |||||
| ToolFile.id == tool_file_id, | ToolFile.id == tool_file_id, | ||||
| ) | ) | ||||
| .first() | .first() |
| raise ValueError("Unsupported tool type") | raise ValueError("Unsupported tool type") | ||||
| # delete old labels | # delete old labels | ||||
| db.session.query(ToolLabelBinding).filter(ToolLabelBinding.tool_id == provider_id).delete() | |||||
| db.session.query(ToolLabelBinding).where(ToolLabelBinding.tool_id == provider_id).delete() | |||||
| # insert new labels | # insert new labels | ||||
| for label in labels: | for label in labels: | ||||
| labels = ( | labels = ( | ||||
| db.session.query(ToolLabelBinding.label_name) | db.session.query(ToolLabelBinding.label_name) | ||||
| .filter( | |||||
| .where( | |||||
| ToolLabelBinding.tool_id == provider_id, | ToolLabelBinding.tool_id == provider_id, | ||||
| ToolLabelBinding.tool_type == controller.provider_type.value, | ToolLabelBinding.tool_type == controller.provider_type.value, | ||||
| ) | ) | ||||
| provider_ids.append(controller.provider_id) | provider_ids.append(controller.provider_id) | ||||
| labels: list[ToolLabelBinding] = ( | labels: list[ToolLabelBinding] = ( | ||||
| db.session.query(ToolLabelBinding).filter(ToolLabelBinding.tool_id.in_(provider_ids)).all() | |||||
| db.session.query(ToolLabelBinding).where(ToolLabelBinding.tool_id.in_(provider_ids)).all() | |||||
| ) | ) | ||||
| tool_labels: dict[str, list[str]] = {label.tool_id: [] for label in labels} | tool_labels: dict[str, list[str]] = {label.tool_id: [] for label in labels} |
| try: | try: | ||||
| builtin_provider = ( | builtin_provider = ( | ||||
| db.session.query(BuiltinToolProvider) | db.session.query(BuiltinToolProvider) | ||||
| .filter( | |||||
| .where( | |||||
| BuiltinToolProvider.tenant_id == tenant_id, | BuiltinToolProvider.tenant_id == tenant_id, | ||||
| BuiltinToolProvider.id == credential_id, | BuiltinToolProvider.id == credential_id, | ||||
| ) | ) | ||||
| # use the default provider | # use the default provider | ||||
| builtin_provider = ( | builtin_provider = ( | ||||
| db.session.query(BuiltinToolProvider) | db.session.query(BuiltinToolProvider) | ||||
| .filter( | |||||
| .where( | |||||
| BuiltinToolProvider.tenant_id == tenant_id, | BuiltinToolProvider.tenant_id == tenant_id, | ||||
| (BuiltinToolProvider.provider == str(provider_id_entity)) | (BuiltinToolProvider.provider == str(provider_id_entity)) | ||||
| | (BuiltinToolProvider.provider == provider_id_entity.provider_name), | | (BuiltinToolProvider.provider == provider_id_entity.provider_name), | ||||
| else: | else: | ||||
| builtin_provider = ( | builtin_provider = ( | ||||
| db.session.query(BuiltinToolProvider) | db.session.query(BuiltinToolProvider) | ||||
| .filter(BuiltinToolProvider.tenant_id == tenant_id, (BuiltinToolProvider.provider == provider_id)) | |||||
| .where(BuiltinToolProvider.tenant_id == tenant_id, (BuiltinToolProvider.provider == provider_id)) | |||||
| .order_by(BuiltinToolProvider.is_default.desc(), BuiltinToolProvider.created_at.asc()) | .order_by(BuiltinToolProvider.is_default.desc(), BuiltinToolProvider.created_at.asc()) | ||||
| .first() | .first() | ||||
| ) | ) | ||||
| elif provider_type == ToolProviderType.WORKFLOW: | elif provider_type == ToolProviderType.WORKFLOW: | ||||
| workflow_provider = ( | workflow_provider = ( | ||||
| db.session.query(WorkflowToolProvider) | db.session.query(WorkflowToolProvider) | ||||
| .filter(WorkflowToolProvider.tenant_id == tenant_id, WorkflowToolProvider.id == provider_id) | |||||
| .where(WorkflowToolProvider.tenant_id == tenant_id, WorkflowToolProvider.id == provider_id) | |||||
| .first() | .first() | ||||
| ) | ) | ||||
| ORDER BY tenant_id, provider, is_default DESC, created_at DESC | ORDER BY tenant_id, provider, is_default DESC, created_at DESC | ||||
| """ | """ | ||||
| ids = [row.id for row in db.session.execute(db.text(sql), {"tenant_id": tenant_id}).all()] | ids = [row.id for row in db.session.execute(db.text(sql), {"tenant_id": tenant_id}).all()] | ||||
| return db.session.query(BuiltinToolProvider).filter(BuiltinToolProvider.id.in_(ids)).all() | |||||
| return db.session.query(BuiltinToolProvider).where(BuiltinToolProvider.id.in_(ids)).all() | |||||
| @classmethod | @classmethod | ||||
| def list_providers_from_api( | def list_providers_from_api( | ||||
| # get db api providers | # get db api providers | ||||
| if "api" in filters: | if "api" in filters: | ||||
| db_api_providers: list[ApiToolProvider] = ( | db_api_providers: list[ApiToolProvider] = ( | ||||
| db.session.query(ApiToolProvider).filter(ApiToolProvider.tenant_id == tenant_id).all() | |||||
| db.session.query(ApiToolProvider).where(ApiToolProvider.tenant_id == tenant_id).all() | |||||
| ) | ) | ||||
| api_provider_controllers: list[dict[str, Any]] = [ | api_provider_controllers: list[dict[str, Any]] = [ | ||||
| if "workflow" in filters: | if "workflow" in filters: | ||||
| # get workflow providers | # get workflow providers | ||||
| workflow_providers: list[WorkflowToolProvider] = ( | workflow_providers: list[WorkflowToolProvider] = ( | ||||
| db.session.query(WorkflowToolProvider).filter(WorkflowToolProvider.tenant_id == tenant_id).all() | |||||
| db.session.query(WorkflowToolProvider).where(WorkflowToolProvider.tenant_id == tenant_id).all() | |||||
| ) | ) | ||||
| workflow_provider_controllers: list[WorkflowToolProviderController] = [] | workflow_provider_controllers: list[WorkflowToolProviderController] = [] | ||||
| """ | """ | ||||
| provider: ApiToolProvider | None = ( | provider: ApiToolProvider | None = ( | ||||
| db.session.query(ApiToolProvider) | db.session.query(ApiToolProvider) | ||||
| .filter( | |||||
| .where( | |||||
| ApiToolProvider.id == provider_id, | ApiToolProvider.id == provider_id, | ||||
| ApiToolProvider.tenant_id == tenant_id, | ApiToolProvider.tenant_id == tenant_id, | ||||
| ) | ) | ||||
| """ | """ | ||||
| provider: MCPToolProvider | None = ( | provider: MCPToolProvider | None = ( | ||||
| db.session.query(MCPToolProvider) | db.session.query(MCPToolProvider) | ||||
| .filter( | |||||
| .where( | |||||
| MCPToolProvider.server_identifier == provider_id, | MCPToolProvider.server_identifier == provider_id, | ||||
| MCPToolProvider.tenant_id == tenant_id, | MCPToolProvider.tenant_id == tenant_id, | ||||
| ) | ) | ||||
| provider_name = provider | provider_name = provider | ||||
| provider_obj: ApiToolProvider | None = ( | provider_obj: ApiToolProvider | None = ( | ||||
| db.session.query(ApiToolProvider) | db.session.query(ApiToolProvider) | ||||
| .filter( | |||||
| .where( | |||||
| ApiToolProvider.tenant_id == tenant_id, | ApiToolProvider.tenant_id == tenant_id, | ||||
| ApiToolProvider.name == provider, | ApiToolProvider.name == provider, | ||||
| ) | ) | ||||
| try: | try: | ||||
| workflow_provider: WorkflowToolProvider | None = ( | workflow_provider: WorkflowToolProvider | None = ( | ||||
| db.session.query(WorkflowToolProvider) | db.session.query(WorkflowToolProvider) | ||||
| .filter(WorkflowToolProvider.tenant_id == tenant_id, WorkflowToolProvider.id == provider_id) | |||||
| .where(WorkflowToolProvider.tenant_id == tenant_id, WorkflowToolProvider.id == provider_id) | |||||
| .first() | .first() | ||||
| ) | ) | ||||
| try: | try: | ||||
| api_provider: ApiToolProvider | None = ( | api_provider: ApiToolProvider | None = ( | ||||
| db.session.query(ApiToolProvider) | db.session.query(ApiToolProvider) | ||||
| .filter(ApiToolProvider.tenant_id == tenant_id, ApiToolProvider.id == provider_id) | |||||
| .where(ApiToolProvider.tenant_id == tenant_id, ApiToolProvider.id == provider_id) | |||||
| .first() | .first() | ||||
| ) | ) | ||||
| try: | try: | ||||
| mcp_provider: MCPToolProvider | None = ( | mcp_provider: MCPToolProvider | None = ( | ||||
| db.session.query(MCPToolProvider) | db.session.query(MCPToolProvider) | ||||
| .filter(MCPToolProvider.tenant_id == tenant_id, MCPToolProvider.server_identifier == provider_id) | |||||
| .where(MCPToolProvider.tenant_id == tenant_id, MCPToolProvider.server_identifier == provider_id) | |||||
| .first() | .first() | ||||
| ) | ) | ||||
| index_node_ids = [document.metadata["doc_id"] for document in all_documents if document.metadata] | index_node_ids = [document.metadata["doc_id"] for document in all_documents if document.metadata] | ||||
| segments = ( | segments = ( | ||||
| db.session.query(DocumentSegment) | db.session.query(DocumentSegment) | ||||
| .filter( | |||||
| .where( | |||||
| DocumentSegment.dataset_id.in_(self.dataset_ids), | DocumentSegment.dataset_id.in_(self.dataset_ids), | ||||
| DocumentSegment.completed_at.isnot(None), | DocumentSegment.completed_at.isnot(None), | ||||
| DocumentSegment.status == "completed", | DocumentSegment.status == "completed", | ||||
| dataset = db.session.query(Dataset).filter_by(id=segment.dataset_id).first() | dataset = db.session.query(Dataset).filter_by(id=segment.dataset_id).first() | ||||
| document = ( | document = ( | ||||
| db.session.query(Document) | db.session.query(Document) | ||||
| .filter( | |||||
| .where( | |||||
| Document.id == segment.document_id, | Document.id == segment.document_id, | ||||
| Document.enabled == True, | Document.enabled == True, | ||||
| Document.archived == False, | Document.archived == False, | ||||
| ): | ): | ||||
| with flask_app.app_context(): | with flask_app.app_context(): | ||||
| dataset = ( | dataset = ( | ||||
| db.session.query(Dataset).filter(Dataset.tenant_id == self.tenant_id, Dataset.id == dataset_id).first() | |||||
| db.session.query(Dataset).where(Dataset.tenant_id == self.tenant_id, Dataset.id == dataset_id).first() | |||||
| ) | ) | ||||
| if not dataset: | if not dataset: |
| def _run(self, query: str) -> str: | def _run(self, query: str) -> str: | ||||
| dataset = ( | dataset = ( | ||||
| db.session.query(Dataset).filter(Dataset.tenant_id == self.tenant_id, Dataset.id == self.dataset_id).first() | |||||
| db.session.query(Dataset).where(Dataset.tenant_id == self.tenant_id, Dataset.id == self.dataset_id).first() | |||||
| ) | ) | ||||
| if not dataset: | if not dataset: | ||||
| dataset = db.session.query(Dataset).filter_by(id=segment.dataset_id).first() | dataset = db.session.query(Dataset).filter_by(id=segment.dataset_id).first() | ||||
| document = ( | document = ( | ||||
| db.session.query(DatasetDocument) # type: ignore | db.session.query(DatasetDocument) # type: ignore | ||||
| .filter( | |||||
| .where( | |||||
| DatasetDocument.id == segment.document_id, | DatasetDocument.id == segment.document_id, | ||||
| DatasetDocument.enabled == True, | DatasetDocument.enabled == True, | ||||
| DatasetDocument.archived == False, | DatasetDocument.archived == False, |
| """ | """ | ||||
| workflow: Workflow | None = ( | workflow: Workflow | None = ( | ||||
| db.session.query(Workflow) | db.session.query(Workflow) | ||||
| .filter(Workflow.app_id == db_provider.app_id, Workflow.version == db_provider.version) | |||||
| .where(Workflow.app_id == db_provider.app_id, Workflow.version == db_provider.version) | |||||
| .first() | .first() | ||||
| ) | ) | ||||
| db_providers: WorkflowToolProvider | None = ( | db_providers: WorkflowToolProvider | None = ( | ||||
| db.session.query(WorkflowToolProvider) | db.session.query(WorkflowToolProvider) | ||||
| .filter( | |||||
| .where( | |||||
| WorkflowToolProvider.tenant_id == tenant_id, | WorkflowToolProvider.tenant_id == tenant_id, | ||||
| WorkflowToolProvider.app_id == self.provider_id, | WorkflowToolProvider.app_id == self.provider_id, | ||||
| ) | ) |
| if not version: | if not version: | ||||
| workflow = ( | workflow = ( | ||||
| db.session.query(Workflow) | db.session.query(Workflow) | ||||
| .filter(Workflow.app_id == app_id, Workflow.version != "draft") | |||||
| .where(Workflow.app_id == app_id, Workflow.version != "draft") | |||||
| .order_by(Workflow.created_at.desc()) | .order_by(Workflow.created_at.desc()) | ||||
| .first() | .first() | ||||
| ) | ) | ||||
| else: | else: | ||||
| workflow = db.session.query(Workflow).filter(Workflow.app_id == app_id, Workflow.version == version).first() | |||||
| workflow = db.session.query(Workflow).where(Workflow.app_id == app_id, Workflow.version == version).first() | |||||
| if not workflow: | if not workflow: | ||||
| raise ValueError("workflow not found or not published") | raise ValueError("workflow not found or not published") | ||||
| """ | """ | ||||
| get the app by app id | get the app by app id | ||||
| """ | """ | ||||
| app = db.session.query(App).filter(App.id == app_id).first() | |||||
| app = db.session.query(App).where(App.id == app_id).first() | |||||
| if not app: | if not app: | ||||
| raise ValueError("app not found") | raise ValueError("app not found") | ||||
| # Subquery: Count the number of available documents for each dataset | # Subquery: Count the number of available documents for each dataset | ||||
| subquery = ( | subquery = ( | ||||
| db.session.query(Document.dataset_id, func.count(Document.id).label("available_document_count")) | db.session.query(Document.dataset_id, func.count(Document.id).label("available_document_count")) | ||||
| .filter( | |||||
| .where( | |||||
| Document.indexing_status == "completed", | Document.indexing_status == "completed", | ||||
| Document.enabled == True, | Document.enabled == True, | ||||
| Document.archived == False, | Document.archived == False, | ||||
| results = ( | results = ( | ||||
| db.session.query(Dataset) | db.session.query(Dataset) | ||||
| .outerjoin(subquery, Dataset.id == subquery.c.dataset_id) | .outerjoin(subquery, Dataset.id == subquery.c.dataset_id) | ||||
| .filter(Dataset.tenant_id == self.tenant_id, Dataset.id.in_(dataset_ids)) | |||||
| .filter((subquery.c.available_document_count > 0) | (Dataset.provider == "external")) | |||||
| .where(Dataset.tenant_id == self.tenant_id, Dataset.id.in_(dataset_ids)) | |||||
| .where((subquery.c.available_document_count > 0) | (Dataset.provider == "external")) | |||||
| .all() | .all() | ||||
| ) | ) | ||||
| dataset = db.session.query(Dataset).filter_by(id=segment.dataset_id).first() # type: ignore | dataset = db.session.query(Dataset).filter_by(id=segment.dataset_id).first() # type: ignore | ||||
| document = ( | document = ( | ||||
| db.session.query(Document) | db.session.query(Document) | ||||
| .filter( | |||||
| .where( | |||||
| Document.id == segment.document_id, | Document.id == segment.document_id, | ||||
| Document.enabled == True, | Document.enabled == True, | ||||
| Document.archived == False, | Document.archived == False, | ||||
| def _get_metadata_filter_condition( | def _get_metadata_filter_condition( | ||||
| self, dataset_ids: list, query: str, node_data: KnowledgeRetrievalNodeData | self, dataset_ids: list, query: str, node_data: KnowledgeRetrievalNodeData | ||||
| ) -> tuple[Optional[dict[str, list[str]]], Optional[MetadataCondition]]: | ) -> tuple[Optional[dict[str, list[str]]], Optional[MetadataCondition]]: | ||||
| document_query = db.session.query(Document).filter( | |||||
| document_query = db.session.query(Document).where( | |||||
| Document.dataset_id.in_(dataset_ids), | Document.dataset_id.in_(dataset_ids), | ||||
| Document.indexing_status == "completed", | Document.indexing_status == "completed", | ||||
| Document.enabled == True, | Document.enabled == True, | ||||
| node_data.metadata_filtering_conditions | node_data.metadata_filtering_conditions | ||||
| and node_data.metadata_filtering_conditions.logical_operator == "and" | and node_data.metadata_filtering_conditions.logical_operator == "and" | ||||
| ): # type: ignore | ): # type: ignore | ||||
| document_query = document_query.filter(and_(*filters)) | |||||
| document_query = document_query.where(and_(*filters)) | |||||
| else: | else: | ||||
| document_query = document_query.filter(or_(*filters)) | |||||
| document_query = document_query.where(or_(*filters)) | |||||
| documents = document_query.all() | documents = document_query.all() | ||||
| # group by dataset_id | # group by dataset_id | ||||
| metadata_filter_document_ids = defaultdict(list) if documents else None # type: ignore | metadata_filter_document_ids = defaultdict(list) if documents else None # type: ignore | ||||
| self, dataset_ids: list, query: str, node_data: KnowledgeRetrievalNodeData | self, dataset_ids: list, query: str, node_data: KnowledgeRetrievalNodeData | ||||
| ) -> list[dict[str, Any]]: | ) -> list[dict[str, Any]]: | ||||
| # get all metadata field | # get all metadata field | ||||
| metadata_fields = db.session.query(DatasetMetadata).filter(DatasetMetadata.dataset_id.in_(dataset_ids)).all() | |||||
| metadata_fields = db.session.query(DatasetMetadata).where(DatasetMetadata.dataset_id.in_(dataset_ids)).all() | |||||
| all_metadata_fields = [metadata_field.name for metadata_field in metadata_fields] | all_metadata_fields = [metadata_field.name for metadata_field in metadata_fields] | ||||
| if node_data.metadata_model_config is None: | if node_data.metadata_model_config is None: | ||||
| raise ValueError("metadata_model_config is required") | raise ValueError("metadata_model_config is required") |
| document = ( | document = ( | ||||
| db.session.query(Document) | db.session.query(Document) | ||||
| .filter( | |||||
| .where( | |||||
| Document.id == document_id, | Document.id == document_id, | ||||
| Document.dataset_id == dataset_id, | Document.dataset_id == dataset_id, | ||||
| ) | ) |
| dataset_ids = get_dataset_ids_from_model_config(app_model_config) | dataset_ids = get_dataset_ids_from_model_config(app_model_config) | ||||
| app_dataset_joins = db.session.query(AppDatasetJoin).filter(AppDatasetJoin.app_id == app.id).all() | |||||
| app_dataset_joins = db.session.query(AppDatasetJoin).where(AppDatasetJoin.app_id == app.id).all() | |||||
| removed_dataset_ids: set[str] = set() | removed_dataset_ids: set[str] = set() | ||||
| if not app_dataset_joins: | if not app_dataset_joins: | ||||
| if removed_dataset_ids: | if removed_dataset_ids: | ||||
| for dataset_id in removed_dataset_ids: | for dataset_id in removed_dataset_ids: | ||||
| db.session.query(AppDatasetJoin).filter( | |||||
| db.session.query(AppDatasetJoin).where( | |||||
| AppDatasetJoin.app_id == app.id, AppDatasetJoin.dataset_id == dataset_id | AppDatasetJoin.app_id == app.id, AppDatasetJoin.dataset_id == dataset_id | ||||
| ).delete() | ).delete() | ||||
| published_workflow = cast(Workflow, published_workflow) | published_workflow = cast(Workflow, published_workflow) | ||||
| dataset_ids = get_dataset_ids_from_workflow(published_workflow) | dataset_ids = get_dataset_ids_from_workflow(published_workflow) | ||||
| app_dataset_joins = db.session.query(AppDatasetJoin).filter(AppDatasetJoin.app_id == app.id).all() | |||||
| app_dataset_joins = db.session.query(AppDatasetJoin).where(AppDatasetJoin.app_id == app.id).all() | |||||
| removed_dataset_ids: set[str] = set() | removed_dataset_ids: set[str] = set() | ||||
| if not app_dataset_joins: | if not app_dataset_joins: | ||||
| if removed_dataset_ids: | if removed_dataset_ids: | ||||
| for dataset_id in removed_dataset_ids: | for dataset_id in removed_dataset_ids: | ||||
| db.session.query(AppDatasetJoin).filter( | |||||
| db.session.query(AppDatasetJoin).where( | |||||
| AppDatasetJoin.app_id == app.id, AppDatasetJoin.dataset_id == dataset_id | AppDatasetJoin.app_id == app.id, AppDatasetJoin.dataset_id == dataset_id | ||||
| ).delete() | ).delete() | ||||
| if workspace_id: | if workspace_id: | ||||
| tenant_account_join = ( | tenant_account_join = ( | ||||
| db.session.query(Tenant, TenantAccountJoin) | db.session.query(Tenant, TenantAccountJoin) | ||||
| .filter(Tenant.id == workspace_id) | |||||
| .filter(TenantAccountJoin.tenant_id == Tenant.id) | |||||
| .filter(TenantAccountJoin.role == "owner") | |||||
| .where(Tenant.id == workspace_id) | |||||
| .where(TenantAccountJoin.tenant_id == Tenant.id) | |||||
| .where(TenantAccountJoin.role == "owner") | |||||
| .one_or_none() | .one_or_none() | ||||
| ) | ) | ||||
| if tenant_account_join: | if tenant_account_join: | ||||
| end_user_id = decoded.get("end_user_id") | end_user_id = decoded.get("end_user_id") | ||||
| if not end_user_id: | if not end_user_id: | ||||
| raise Unauthorized("Invalid Authorization token.") | raise Unauthorized("Invalid Authorization token.") | ||||
| end_user = db.session.query(EndUser).filter(EndUser.id == decoded["end_user_id"]).first() | |||||
| end_user = db.session.query(EndUser).where(EndUser.id == decoded["end_user_id"]).first() | |||||
| if not end_user: | if not end_user: | ||||
| raise NotFound("End user not found.") | raise NotFound("End user not found.") | ||||
| return end_user | return end_user | ||||
| server_code = request.view_args.get("server_code") if request.view_args else None | server_code = request.view_args.get("server_code") if request.view_args else None | ||||
| if not server_code: | if not server_code: | ||||
| raise Unauthorized("Invalid Authorization token.") | raise Unauthorized("Invalid Authorization token.") | ||||
| app_mcp_server = db.session.query(AppMCPServer).filter(AppMCPServer.server_code == server_code).first() | |||||
| app_mcp_server = db.session.query(AppMCPServer).where(AppMCPServer.server_code == server_code).first() | |||||
| if not app_mcp_server: | if not app_mcp_server: | ||||
| raise NotFound("App MCP server not found.") | raise NotFound("App MCP server not found.") | ||||
| end_user = ( | end_user = ( | ||||
| db.session.query(EndUser) | db.session.query(EndUser) | ||||
| .filter(EndUser.external_user_id == app_mcp_server.id, EndUser.type == "mcp") | |||||
| .where(EndUser.external_user_id == app_mcp_server.id, EndUser.type == "mcp") | |||||
| .first() | .first() | ||||
| ) | ) | ||||
| if not end_user: | if not end_user: |
| transfer_method: FileTransferMethod, | transfer_method: FileTransferMethod, | ||||
| strict_type_validation: bool = False, | strict_type_validation: bool = False, | ||||
| ) -> File: | ) -> File: | ||||
| tool_file = ( | |||||
| db.session.query(ToolFile) | |||||
| .filter( | |||||
| tool_file = db.session.scalar( | |||||
| select(ToolFile).where( | |||||
| ToolFile.id == mapping.get("tool_file_id"), | ToolFile.id == mapping.get("tool_file_id"), | ||||
| ToolFile.tenant_id == tenant_id, | ToolFile.tenant_id == tenant_id, | ||||
| ) | ) | ||||
| .first() | |||||
| ) | ) | ||||
| if tool_file is None: | if tool_file is None: | ||||
| extension = "." + tool_file.file_key.split(".")[-1] if "." in tool_file.file_key else ".bin" | extension = "." + tool_file.file_key.split(".")[-1] if "." in tool_file.file_key else ".bin" | ||||
| detected_file_type = _standardize_file_type(extension="." + extension, mime_type=tool_file.mimetype) | |||||
| detected_file_type = _standardize_file_type(extension=extension, mime_type=tool_file.mimetype) | |||||
| specified_type = mapping.get("type") | specified_type = mapping.get("type") | ||||
| import requests | import requests | ||||
| from flask_login import current_user | from flask_login import current_user | ||||
| from sqlalchemy import select | |||||
| from extensions.ext_database import db | from extensions.ext_database import db | ||||
| from libs.datetime_utils import naive_utc_now | from libs.datetime_utils import naive_utc_now | ||||
| "total": len(pages), | "total": len(pages), | ||||
| } | } | ||||
| # save data source binding | # save data source binding | ||||
| data_source_binding = ( | |||||
| db.session.query(DataSourceOauthBinding) | |||||
| .filter( | |||||
| db.and_( | |||||
| DataSourceOauthBinding.tenant_id == current_user.current_tenant_id, | |||||
| DataSourceOauthBinding.provider == "notion", | |||||
| DataSourceOauthBinding.access_token == access_token, | |||||
| ) | |||||
| data_source_binding = db.session.scalar( | |||||
| select(DataSourceOauthBinding).where( | |||||
| DataSourceOauthBinding.tenant_id == current_user.current_tenant_id, | |||||
| DataSourceOauthBinding.provider == "notion", | |||||
| DataSourceOauthBinding.access_token == access_token, | |||||
| ) | ) | ||||
| .first() | |||||
| ) | ) | ||||
| if data_source_binding: | if data_source_binding: | ||||
| data_source_binding.source_info = source_info | data_source_binding.source_info = source_info | ||||
| "total": len(pages), | "total": len(pages), | ||||
| } | } | ||||
| # save data source binding | # save data source binding | ||||
| data_source_binding = ( | |||||
| db.session.query(DataSourceOauthBinding) | |||||
| .filter( | |||||
| db.and_( | |||||
| DataSourceOauthBinding.tenant_id == current_user.current_tenant_id, | |||||
| DataSourceOauthBinding.provider == "notion", | |||||
| DataSourceOauthBinding.access_token == access_token, | |||||
| ) | |||||
| data_source_binding = db.session.scalar( | |||||
| select(DataSourceOauthBinding).where( | |||||
| DataSourceOauthBinding.tenant_id == current_user.current_tenant_id, | |||||
| DataSourceOauthBinding.provider == "notion", | |||||
| DataSourceOauthBinding.access_token == access_token, | |||||
| ) | ) | ||||
| .first() | |||||
| ) | ) | ||||
| if data_source_binding: | if data_source_binding: | ||||
| data_source_binding.source_info = source_info | data_source_binding.source_info = source_info | ||||
| def sync_data_source(self, binding_id: str): | def sync_data_source(self, binding_id: str): | ||||
| # save data source binding | # save data source binding | ||||
| data_source_binding = ( | |||||
| db.session.query(DataSourceOauthBinding) | |||||
| .filter( | |||||
| db.and_( | |||||
| DataSourceOauthBinding.tenant_id == current_user.current_tenant_id, | |||||
| DataSourceOauthBinding.provider == "notion", | |||||
| DataSourceOauthBinding.id == binding_id, | |||||
| DataSourceOauthBinding.disabled == False, | |||||
| ) | |||||
| data_source_binding = db.session.scalar( | |||||
| select(DataSourceOauthBinding).where( | |||||
| DataSourceOauthBinding.tenant_id == current_user.current_tenant_id, | |||||
| DataSourceOauthBinding.provider == "notion", | |||||
| DataSourceOauthBinding.id == binding_id, | |||||
| DataSourceOauthBinding.disabled == False, | |||||
| ) | ) | ||||
| .first() | |||||
| ) | ) | ||||
| if data_source_binding: | if data_source_binding: | ||||
| # get all authorized pages | # get all authorized pages | ||||
| pages = self.get_authorized_pages(data_source_binding.access_token) | pages = self.get_authorized_pages(data_source_binding.access_token) |
| from typing import Optional, cast | from typing import Optional, cast | ||||
| from flask_login import UserMixin # type: ignore | from flask_login import UserMixin # type: ignore | ||||
| from sqlalchemy import func | |||||
| from sqlalchemy import func, select | |||||
| from sqlalchemy.orm import Mapped, mapped_column, reconstructor | from sqlalchemy.orm import Mapped, mapped_column, reconstructor | ||||
| from models.base import Base | from models.base import Base | ||||
| @current_tenant.setter | @current_tenant.setter | ||||
| def current_tenant(self, tenant: "Tenant"): | def current_tenant(self, tenant: "Tenant"): | ||||
| ta = db.session.query(TenantAccountJoin).filter_by(tenant_id=tenant.id, account_id=self.id).first() | |||||
| ta = db.session.scalar(select(TenantAccountJoin).filter_by(tenant_id=tenant.id, account_id=self.id).limit(1)) | |||||
| if ta: | if ta: | ||||
| self.role = TenantAccountRole(ta.role) | self.role = TenantAccountRole(ta.role) | ||||
| self._current_tenant = tenant | self._current_tenant = tenant | ||||
| tuple[Tenant, TenantAccountJoin], | tuple[Tenant, TenantAccountJoin], | ||||
| ( | ( | ||||
| db.session.query(Tenant, TenantAccountJoin) | db.session.query(Tenant, TenantAccountJoin) | ||||
| .filter(Tenant.id == tenant_id) | |||||
| .filter(TenantAccountJoin.tenant_id == Tenant.id) | |||||
| .filter(TenantAccountJoin.account_id == self.id) | |||||
| .where(Tenant.id == tenant_id) | |||||
| .where(TenantAccountJoin.tenant_id == Tenant.id) | |||||
| .where(TenantAccountJoin.account_id == self.id) | |||||
| .one_or_none() | .one_or_none() | ||||
| ), | ), | ||||
| ) | ) | ||||
| def get_by_openid(cls, provider: str, open_id: str): | def get_by_openid(cls, provider: str, open_id: str): | ||||
| account_integrate = ( | account_integrate = ( | ||||
| db.session.query(AccountIntegrate) | db.session.query(AccountIntegrate) | ||||
| .filter(AccountIntegrate.provider == provider, AccountIntegrate.open_id == open_id) | |||||
| .where(AccountIntegrate.provider == provider, AccountIntegrate.open_id == open_id) | |||||
| .one_or_none() | .one_or_none() | ||||
| ) | ) | ||||
| if account_integrate: | if account_integrate: | ||||
| return db.session.query(Account).filter(Account.id == account_integrate.account_id).one_or_none() | |||||
| return db.session.query(Account).where(Account.id == account_integrate.account_id).one_or_none() | |||||
| return None | return None | ||||
| # check current_user.current_tenant.current_role in ['admin', 'owner'] | # check current_user.current_tenant.current_role in ['admin', 'owner'] | ||||
| def get_accounts(self) -> list[Account]: | def get_accounts(self) -> list[Account]: | ||||
| return ( | return ( | ||||
| db.session.query(Account) | db.session.query(Account) | ||||
| .filter(Account.id == TenantAccountJoin.account_id, TenantAccountJoin.tenant_id == self.id) | |||||
| .where(Account.id == TenantAccountJoin.account_id, TenantAccountJoin.tenant_id == self.id) | |||||
| .all() | .all() | ||||
| ) | ) | ||||
| from json import JSONDecodeError | from json import JSONDecodeError | ||||
| from typing import Any, Optional, cast | from typing import Any, Optional, cast | ||||
| from sqlalchemy import func | |||||
| from sqlalchemy import func, select | |||||
| from sqlalchemy.dialects.postgresql import JSONB | from sqlalchemy.dialects.postgresql import JSONB | ||||
| from sqlalchemy.orm import Mapped, mapped_column | from sqlalchemy.orm import Mapped, mapped_column | ||||
| @property | @property | ||||
| def dataset_keyword_table(self): | def dataset_keyword_table(self): | ||||
| dataset_keyword_table = ( | dataset_keyword_table = ( | ||||
| db.session.query(DatasetKeywordTable).filter(DatasetKeywordTable.dataset_id == self.id).first() | |||||
| db.session.query(DatasetKeywordTable).where(DatasetKeywordTable.dataset_id == self.id).first() | |||||
| ) | ) | ||||
| if dataset_keyword_table: | if dataset_keyword_table: | ||||
| return dataset_keyword_table | return dataset_keyword_table | ||||
| def latest_process_rule(self): | def latest_process_rule(self): | ||||
| return ( | return ( | ||||
| db.session.query(DatasetProcessRule) | db.session.query(DatasetProcessRule) | ||||
| .filter(DatasetProcessRule.dataset_id == self.id) | |||||
| .where(DatasetProcessRule.dataset_id == self.id) | |||||
| .order_by(DatasetProcessRule.created_at.desc()) | .order_by(DatasetProcessRule.created_at.desc()) | ||||
| .first() | .first() | ||||
| ) | ) | ||||
| def app_count(self): | def app_count(self): | ||||
| return ( | return ( | ||||
| db.session.query(func.count(AppDatasetJoin.id)) | db.session.query(func.count(AppDatasetJoin.id)) | ||||
| .filter(AppDatasetJoin.dataset_id == self.id, App.id == AppDatasetJoin.app_id) | |||||
| .where(AppDatasetJoin.dataset_id == self.id, App.id == AppDatasetJoin.app_id) | |||||
| .scalar() | .scalar() | ||||
| ) | ) | ||||
| @property | @property | ||||
| def document_count(self): | def document_count(self): | ||||
| return db.session.query(func.count(Document.id)).filter(Document.dataset_id == self.id).scalar() | |||||
| return db.session.query(func.count(Document.id)).where(Document.dataset_id == self.id).scalar() | |||||
| @property | @property | ||||
| def available_document_count(self): | def available_document_count(self): | ||||
| return ( | return ( | ||||
| db.session.query(func.count(Document.id)) | db.session.query(func.count(Document.id)) | ||||
| .filter( | |||||
| .where( | |||||
| Document.dataset_id == self.id, | Document.dataset_id == self.id, | ||||
| Document.indexing_status == "completed", | Document.indexing_status == "completed", | ||||
| Document.enabled == True, | Document.enabled == True, | ||||
| def available_segment_count(self): | def available_segment_count(self): | ||||
| return ( | return ( | ||||
| db.session.query(func.count(DocumentSegment.id)) | db.session.query(func.count(DocumentSegment.id)) | ||||
| .filter( | |||||
| .where( | |||||
| DocumentSegment.dataset_id == self.id, | DocumentSegment.dataset_id == self.id, | ||||
| DocumentSegment.status == "completed", | DocumentSegment.status == "completed", | ||||
| DocumentSegment.enabled == True, | DocumentSegment.enabled == True, | ||||
| return ( | return ( | ||||
| db.session.query(Document) | db.session.query(Document) | ||||
| .with_entities(func.coalesce(func.sum(Document.word_count), 0)) | .with_entities(func.coalesce(func.sum(Document.word_count), 0)) | ||||
| .filter(Document.dataset_id == self.id) | |||||
| .where(Document.dataset_id == self.id) | |||||
| .scalar() | .scalar() | ||||
| ) | ) | ||||
| @property | @property | ||||
| def doc_form(self): | def doc_form(self): | ||||
| document = db.session.query(Document).filter(Document.dataset_id == self.id).first() | |||||
| document = db.session.query(Document).where(Document.dataset_id == self.id).first() | |||||
| if document: | if document: | ||||
| return document.doc_form | return document.doc_form | ||||
| return None | return None | ||||
| tags = ( | tags = ( | ||||
| db.session.query(Tag) | db.session.query(Tag) | ||||
| .join(TagBinding, Tag.id == TagBinding.tag_id) | .join(TagBinding, Tag.id == TagBinding.tag_id) | ||||
| .filter( | |||||
| .where( | |||||
| TagBinding.target_id == self.id, | TagBinding.target_id == self.id, | ||||
| TagBinding.tenant_id == self.tenant_id, | TagBinding.tenant_id == self.tenant_id, | ||||
| Tag.tenant_id == self.tenant_id, | Tag.tenant_id == self.tenant_id, | ||||
| if self.provider != "external": | if self.provider != "external": | ||||
| return None | return None | ||||
| external_knowledge_binding = ( | external_knowledge_binding = ( | ||||
| db.session.query(ExternalKnowledgeBindings).filter(ExternalKnowledgeBindings.dataset_id == self.id).first() | |||||
| db.session.query(ExternalKnowledgeBindings).where(ExternalKnowledgeBindings.dataset_id == self.id).first() | |||||
| ) | ) | ||||
| if not external_knowledge_binding: | if not external_knowledge_binding: | ||||
| return None | return None | ||||
| external_knowledge_api = ( | |||||
| db.session.query(ExternalKnowledgeApis) | |||||
| .filter(ExternalKnowledgeApis.id == external_knowledge_binding.external_knowledge_api_id) | |||||
| .first() | |||||
| external_knowledge_api = db.session.scalar( | |||||
| select(ExternalKnowledgeApis).where( | |||||
| ExternalKnowledgeApis.id == external_knowledge_binding.external_knowledge_api_id | |||||
| ) | |||||
| ) | ) | ||||
| if not external_knowledge_api: | if not external_knowledge_api: | ||||
| return None | return None | ||||
| @property | @property | ||||
| def doc_metadata(self): | def doc_metadata(self): | ||||
| dataset_metadatas = db.session.query(DatasetMetadata).filter(DatasetMetadata.dataset_id == self.id).all() | |||||
| dataset_metadatas = db.session.query(DatasetMetadata).where(DatasetMetadata.dataset_id == self.id).all() | |||||
| doc_metadata = [ | doc_metadata = [ | ||||
| { | { | ||||
| data_source_info_dict = json.loads(self.data_source_info) | data_source_info_dict = json.loads(self.data_source_info) | ||||
| file_detail = ( | file_detail = ( | ||||
| db.session.query(UploadFile) | db.session.query(UploadFile) | ||||
| .filter(UploadFile.id == data_source_info_dict["upload_file_id"]) | |||||
| .where(UploadFile.id == data_source_info_dict["upload_file_id"]) | |||||
| .one_or_none() | .one_or_none() | ||||
| ) | ) | ||||
| if file_detail: | if file_detail: | ||||
| @property | @property | ||||
| def dataset(self): | def dataset(self): | ||||
| return db.session.query(Dataset).filter(Dataset.id == self.dataset_id).one_or_none() | |||||
| return db.session.query(Dataset).where(Dataset.id == self.dataset_id).one_or_none() | |||||
| @property | @property | ||||
| def segment_count(self): | def segment_count(self): | ||||
| return db.session.query(DocumentSegment).filter(DocumentSegment.document_id == self.id).count() | |||||
| return db.session.query(DocumentSegment).where(DocumentSegment.document_id == self.id).count() | |||||
| @property | @property | ||||
| def hit_count(self): | def hit_count(self): | ||||
| return ( | return ( | ||||
| db.session.query(DocumentSegment) | db.session.query(DocumentSegment) | ||||
| .with_entities(func.coalesce(func.sum(DocumentSegment.hit_count), 0)) | .with_entities(func.coalesce(func.sum(DocumentSegment.hit_count), 0)) | ||||
| .filter(DocumentSegment.document_id == self.id) | |||||
| .where(DocumentSegment.document_id == self.id) | |||||
| .scalar() | .scalar() | ||||
| ) | ) | ||||
| @property | @property | ||||
| def uploader(self): | def uploader(self): | ||||
| user = db.session.query(Account).filter(Account.id == self.created_by).first() | |||||
| user = db.session.query(Account).where(Account.id == self.created_by).first() | |||||
| return user.name if user else None | return user.name if user else None | ||||
| @property | @property | ||||
| document_metadatas = ( | document_metadatas = ( | ||||
| db.session.query(DatasetMetadata) | db.session.query(DatasetMetadata) | ||||
| .join(DatasetMetadataBinding, DatasetMetadataBinding.metadata_id == DatasetMetadata.id) | .join(DatasetMetadataBinding, DatasetMetadataBinding.metadata_id == DatasetMetadata.id) | ||||
| .filter( | |||||
| .where( | |||||
| DatasetMetadataBinding.dataset_id == self.dataset_id, DatasetMetadataBinding.document_id == self.id | DatasetMetadataBinding.dataset_id == self.dataset_id, DatasetMetadataBinding.document_id == self.id | ||||
| ) | ) | ||||
| .all() | .all() | ||||
| @property | @property | ||||
| def dataset(self): | def dataset(self): | ||||
| return db.session.query(Dataset).filter(Dataset.id == self.dataset_id).first() | |||||
| return db.session.scalar(select(Dataset).where(Dataset.id == self.dataset_id)) | |||||
| @property | @property | ||||
| def document(self): | def document(self): | ||||
| return db.session.query(Document).filter(Document.id == self.document_id).first() | |||||
| return db.session.scalar(select(Document).where(Document.id == self.document_id)) | |||||
| @property | @property | ||||
| def previous_segment(self): | def previous_segment(self): | ||||
| return ( | |||||
| db.session.query(DocumentSegment) | |||||
| .filter(DocumentSegment.document_id == self.document_id, DocumentSegment.position == self.position - 1) | |||||
| .first() | |||||
| return db.session.scalar( | |||||
| select(DocumentSegment).where( | |||||
| DocumentSegment.document_id == self.document_id, DocumentSegment.position == self.position - 1 | |||||
| ) | |||||
| ) | ) | ||||
| @property | @property | ||||
| def next_segment(self): | def next_segment(self): | ||||
| return ( | |||||
| db.session.query(DocumentSegment) | |||||
| .filter(DocumentSegment.document_id == self.document_id, DocumentSegment.position == self.position + 1) | |||||
| .first() | |||||
| return db.session.scalar( | |||||
| select(DocumentSegment).where( | |||||
| DocumentSegment.document_id == self.document_id, DocumentSegment.position == self.position + 1 | |||||
| ) | |||||
| ) | ) | ||||
| @property | @property | ||||
| if rules.parent_mode and rules.parent_mode != ParentMode.FULL_DOC: | if rules.parent_mode and rules.parent_mode != ParentMode.FULL_DOC: | ||||
| child_chunks = ( | child_chunks = ( | ||||
| db.session.query(ChildChunk) | db.session.query(ChildChunk) | ||||
| .filter(ChildChunk.segment_id == self.id) | |||||
| .where(ChildChunk.segment_id == self.id) | |||||
| .order_by(ChildChunk.position.asc()) | .order_by(ChildChunk.position.asc()) | ||||
| .all() | .all() | ||||
| ) | ) | ||||
| if rules.parent_mode: | if rules.parent_mode: | ||||
| child_chunks = ( | child_chunks = ( | ||||
| db.session.query(ChildChunk) | db.session.query(ChildChunk) | ||||
| .filter(ChildChunk.segment_id == self.id) | |||||
| .where(ChildChunk.segment_id == self.id) | |||||
| .order_by(ChildChunk.position.asc()) | .order_by(ChildChunk.position.asc()) | ||||
| .all() | .all() | ||||
| ) | ) | ||||
| @property | @property | ||||
| def dataset(self): | def dataset(self): | ||||
| return db.session.query(Dataset).filter(Dataset.id == self.dataset_id).first() | |||||
| return db.session.query(Dataset).where(Dataset.id == self.dataset_id).first() | |||||
| @property | @property | ||||
| def document(self): | def document(self): | ||||
| return db.session.query(Document).filter(Document.id == self.document_id).first() | |||||
| return db.session.query(Document).where(Document.id == self.document_id).first() | |||||
| @property | @property | ||||
| def segment(self): | def segment(self): | ||||
| return db.session.query(DocumentSegment).filter(DocumentSegment.id == self.segment_id).first() | |||||
| return db.session.query(DocumentSegment).where(DocumentSegment.id == self.segment_id).first() | |||||
| class AppDatasetJoin(Base): | class AppDatasetJoin(Base): | ||||
| def dataset_bindings(self): | def dataset_bindings(self): | ||||
| external_knowledge_bindings = ( | external_knowledge_bindings = ( | ||||
| db.session.query(ExternalKnowledgeBindings) | db.session.query(ExternalKnowledgeBindings) | ||||
| .filter(ExternalKnowledgeBindings.external_knowledge_api_id == self.id) | |||||
| .where(ExternalKnowledgeBindings.external_knowledge_api_id == self.id) | |||||
| .all() | .all() | ||||
| ) | ) | ||||
| dataset_ids = [binding.dataset_id for binding in external_knowledge_bindings] | dataset_ids = [binding.dataset_id for binding in external_knowledge_bindings] | ||||
| datasets = db.session.query(Dataset).filter(Dataset.id.in_(dataset_ids)).all() | |||||
| datasets = db.session.query(Dataset).where(Dataset.id.in_(dataset_ids)).all() | |||||
| dataset_bindings = [] | dataset_bindings = [] | ||||
| for dataset in datasets: | for dataset in datasets: | ||||
| dataset_bindings.append({"id": dataset.id, "name": dataset.name}) | dataset_bindings.append({"id": dataset.id, "name": dataset.name}) |
| @property | @property | ||||
| def site(self): | def site(self): | ||||
| site = db.session.query(Site).filter(Site.app_id == self.id).first() | |||||
| site = db.session.query(Site).where(Site.app_id == self.id).first() | |||||
| return site | return site | ||||
| @property | @property | ||||
| def app_model_config(self): | def app_model_config(self): | ||||
| if self.app_model_config_id: | if self.app_model_config_id: | ||||
| return db.session.query(AppModelConfig).filter(AppModelConfig.id == self.app_model_config_id).first() | |||||
| return db.session.query(AppModelConfig).where(AppModelConfig.id == self.app_model_config_id).first() | |||||
| return None | return None | ||||
| if self.workflow_id: | if self.workflow_id: | ||||
| from .workflow import Workflow | from .workflow import Workflow | ||||
| return db.session.query(Workflow).filter(Workflow.id == self.workflow_id).first() | |||||
| return db.session.query(Workflow).where(Workflow.id == self.workflow_id).first() | |||||
| return None | return None | ||||
| @property | @property | ||||
| def tenant(self): | def tenant(self): | ||||
| tenant = db.session.query(Tenant).filter(Tenant.id == self.tenant_id).first() | |||||
| tenant = db.session.query(Tenant).where(Tenant.id == self.tenant_id).first() | |||||
| return tenant | return tenant | ||||
| @property | @property | ||||
| tags = ( | tags = ( | ||||
| db.session.query(Tag) | db.session.query(Tag) | ||||
| .join(TagBinding, Tag.id == TagBinding.tag_id) | .join(TagBinding, Tag.id == TagBinding.tag_id) | ||||
| .filter( | |||||
| .where( | |||||
| TagBinding.target_id == self.id, | TagBinding.target_id == self.id, | ||||
| TagBinding.tenant_id == self.tenant_id, | TagBinding.tenant_id == self.tenant_id, | ||||
| Tag.tenant_id == self.tenant_id, | Tag.tenant_id == self.tenant_id, | ||||
| @property | @property | ||||
| def author_name(self): | def author_name(self): | ||||
| if self.created_by: | if self.created_by: | ||||
| account = db.session.query(Account).filter(Account.id == self.created_by).first() | |||||
| account = db.session.query(Account).where(Account.id == self.created_by).first() | |||||
| if account: | if account: | ||||
| return account.name | return account.name | ||||
| @property | @property | ||||
| def app(self): | def app(self): | ||||
| app = db.session.query(App).filter(App.id == self.app_id).first() | |||||
| app = db.session.query(App).where(App.id == self.app_id).first() | |||||
| return app | return app | ||||
| @property | @property | ||||
| @property | @property | ||||
| def annotation_reply_dict(self) -> dict: | def annotation_reply_dict(self) -> dict: | ||||
| annotation_setting = ( | annotation_setting = ( | ||||
| db.session.query(AppAnnotationSetting).filter(AppAnnotationSetting.app_id == self.app_id).first() | |||||
| db.session.query(AppAnnotationSetting).where(AppAnnotationSetting.app_id == self.app_id).first() | |||||
| ) | ) | ||||
| if annotation_setting: | if annotation_setting: | ||||
| collection_binding_detail = annotation_setting.collection_binding_detail | collection_binding_detail = annotation_setting.collection_binding_detail | ||||
| @property | @property | ||||
| def app(self): | def app(self): | ||||
| app = db.session.query(App).filter(App.id == self.app_id).first() | |||||
| app = db.session.query(App).where(App.id == self.app_id).first() | |||||
| return app | return app | ||||
| @property | @property | ||||
| def app(self): | def app(self): | ||||
| app = db.session.query(App).filter(App.id == self.app_id).first() | |||||
| app = db.session.query(App).where(App.id == self.app_id).first() | |||||
| return app | return app | ||||
| @property | @property | ||||
| def tenant(self): | def tenant(self): | ||||
| tenant = db.session.query(Tenant).filter(Tenant.id == self.tenant_id).first() | |||||
| tenant = db.session.query(Tenant).where(Tenant.id == self.tenant_id).first() | |||||
| return tenant | return tenant | ||||
| model_config["configs"] = override_model_configs | model_config["configs"] = override_model_configs | ||||
| else: | else: | ||||
| app_model_config = ( | app_model_config = ( | ||||
| db.session.query(AppModelConfig).filter(AppModelConfig.id == self.app_model_config_id).first() | |||||
| db.session.query(AppModelConfig).where(AppModelConfig.id == self.app_model_config_id).first() | |||||
| ) | ) | ||||
| if app_model_config: | if app_model_config: | ||||
| model_config = app_model_config.to_dict() | model_config = app_model_config.to_dict() | ||||
| @property | @property | ||||
| def annotated(self): | def annotated(self): | ||||
| return db.session.query(MessageAnnotation).filter(MessageAnnotation.conversation_id == self.id).count() > 0 | |||||
| return db.session.query(MessageAnnotation).where(MessageAnnotation.conversation_id == self.id).count() > 0 | |||||
| @property | @property | ||||
| def annotation(self): | def annotation(self): | ||||
| return db.session.query(MessageAnnotation).filter(MessageAnnotation.conversation_id == self.id).first() | |||||
| return db.session.query(MessageAnnotation).where(MessageAnnotation.conversation_id == self.id).first() | |||||
| @property | @property | ||||
| def message_count(self): | def message_count(self): | ||||
| return db.session.query(Message).filter(Message.conversation_id == self.id).count() | |||||
| return db.session.query(Message).where(Message.conversation_id == self.id).count() | |||||
| @property | @property | ||||
| def user_feedback_stats(self): | def user_feedback_stats(self): | ||||
| like = ( | like = ( | ||||
| db.session.query(MessageFeedback) | db.session.query(MessageFeedback) | ||||
| .filter( | |||||
| .where( | |||||
| MessageFeedback.conversation_id == self.id, | MessageFeedback.conversation_id == self.id, | ||||
| MessageFeedback.from_source == "user", | MessageFeedback.from_source == "user", | ||||
| MessageFeedback.rating == "like", | MessageFeedback.rating == "like", | ||||
| dislike = ( | dislike = ( | ||||
| db.session.query(MessageFeedback) | db.session.query(MessageFeedback) | ||||
| .filter( | |||||
| .where( | |||||
| MessageFeedback.conversation_id == self.id, | MessageFeedback.conversation_id == self.id, | ||||
| MessageFeedback.from_source == "user", | MessageFeedback.from_source == "user", | ||||
| MessageFeedback.rating == "dislike", | MessageFeedback.rating == "dislike", | ||||
| def admin_feedback_stats(self): | def admin_feedback_stats(self): | ||||
| like = ( | like = ( | ||||
| db.session.query(MessageFeedback) | db.session.query(MessageFeedback) | ||||
| .filter( | |||||
| .where( | |||||
| MessageFeedback.conversation_id == self.id, | MessageFeedback.conversation_id == self.id, | ||||
| MessageFeedback.from_source == "admin", | MessageFeedback.from_source == "admin", | ||||
| MessageFeedback.rating == "like", | MessageFeedback.rating == "like", | ||||
| dislike = ( | dislike = ( | ||||
| db.session.query(MessageFeedback) | db.session.query(MessageFeedback) | ||||
| .filter( | |||||
| .where( | |||||
| MessageFeedback.conversation_id == self.id, | MessageFeedback.conversation_id == self.id, | ||||
| MessageFeedback.from_source == "admin", | MessageFeedback.from_source == "admin", | ||||
| MessageFeedback.rating == "dislike", | MessageFeedback.rating == "dislike", | ||||
| @property | @property | ||||
| def status_count(self): | def status_count(self): | ||||
| messages = db.session.query(Message).filter(Message.conversation_id == self.id).all() | |||||
| messages = db.session.query(Message).where(Message.conversation_id == self.id).all() | |||||
| status_counts = { | status_counts = { | ||||
| WorkflowExecutionStatus.RUNNING: 0, | WorkflowExecutionStatus.RUNNING: 0, | ||||
| WorkflowExecutionStatus.SUCCEEDED: 0, | WorkflowExecutionStatus.SUCCEEDED: 0, | ||||
| def first_message(self): | def first_message(self): | ||||
| return ( | return ( | ||||
| db.session.query(Message) | db.session.query(Message) | ||||
| .filter(Message.conversation_id == self.id) | |||||
| .where(Message.conversation_id == self.id) | |||||
| .order_by(Message.created_at.asc()) | .order_by(Message.created_at.asc()) | ||||
| .first() | .first() | ||||
| ) | ) | ||||
| @property | @property | ||||
| def app(self): | def app(self): | ||||
| return db.session.query(App).filter(App.id == self.app_id).first() | |||||
| return db.session.query(App).where(App.id == self.app_id).first() | |||||
| @property | @property | ||||
| def from_end_user_session_id(self): | def from_end_user_session_id(self): | ||||
| if self.from_end_user_id: | if self.from_end_user_id: | ||||
| end_user = db.session.query(EndUser).filter(EndUser.id == self.from_end_user_id).first() | |||||
| end_user = db.session.query(EndUser).where(EndUser.id == self.from_end_user_id).first() | |||||
| if end_user: | if end_user: | ||||
| return end_user.session_id | return end_user.session_id | ||||
| @property | @property | ||||
| def from_account_name(self): | def from_account_name(self): | ||||
| if self.from_account_id: | if self.from_account_id: | ||||
| account = db.session.query(Account).filter(Account.id == self.from_account_id).first() | |||||
| account = db.session.query(Account).where(Account.id == self.from_account_id).first() | |||||
| if account: | if account: | ||||
| return account.name | return account.name | ||||
| def user_feedback(self): | def user_feedback(self): | ||||
| feedback = ( | feedback = ( | ||||
| db.session.query(MessageFeedback) | db.session.query(MessageFeedback) | ||||
| .filter(MessageFeedback.message_id == self.id, MessageFeedback.from_source == "user") | |||||
| .where(MessageFeedback.message_id == self.id, MessageFeedback.from_source == "user") | |||||
| .first() | .first() | ||||
| ) | ) | ||||
| return feedback | return feedback | ||||
| def admin_feedback(self): | def admin_feedback(self): | ||||
| feedback = ( | feedback = ( | ||||
| db.session.query(MessageFeedback) | db.session.query(MessageFeedback) | ||||
| .filter(MessageFeedback.message_id == self.id, MessageFeedback.from_source == "admin") | |||||
| .where(MessageFeedback.message_id == self.id, MessageFeedback.from_source == "admin") | |||||
| .first() | .first() | ||||
| ) | ) | ||||
| return feedback | return feedback | ||||
| @property | @property | ||||
| def feedbacks(self): | def feedbacks(self): | ||||
| feedbacks = db.session.query(MessageFeedback).filter(MessageFeedback.message_id == self.id).all() | |||||
| feedbacks = db.session.query(MessageFeedback).where(MessageFeedback.message_id == self.id).all() | |||||
| return feedbacks | return feedbacks | ||||
| @property | @property | ||||
| def annotation(self): | def annotation(self): | ||||
| annotation = db.session.query(MessageAnnotation).filter(MessageAnnotation.message_id == self.id).first() | |||||
| annotation = db.session.query(MessageAnnotation).where(MessageAnnotation.message_id == self.id).first() | |||||
| return annotation | return annotation | ||||
| @property | @property | ||||
| def annotation_hit_history(self): | def annotation_hit_history(self): | ||||
| annotation_history = ( | annotation_history = ( | ||||
| db.session.query(AppAnnotationHitHistory).filter(AppAnnotationHitHistory.message_id == self.id).first() | |||||
| db.session.query(AppAnnotationHitHistory).where(AppAnnotationHitHistory.message_id == self.id).first() | |||||
| ) | ) | ||||
| if annotation_history: | if annotation_history: | ||||
| annotation = ( | annotation = ( | ||||
| db.session.query(MessageAnnotation) | db.session.query(MessageAnnotation) | ||||
| .filter(MessageAnnotation.id == annotation_history.annotation_id) | |||||
| .where(MessageAnnotation.id == annotation_history.annotation_id) | |||||
| .first() | .first() | ||||
| ) | ) | ||||
| return annotation | return annotation | ||||
| @property | @property | ||||
| def app_model_config(self): | def app_model_config(self): | ||||
| conversation = db.session.query(Conversation).filter(Conversation.id == self.conversation_id).first() | |||||
| conversation = db.session.query(Conversation).where(Conversation.id == self.conversation_id).first() | |||||
| if conversation: | if conversation: | ||||
| return ( | |||||
| db.session.query(AppModelConfig).filter(AppModelConfig.id == conversation.app_model_config_id).first() | |||||
| ) | |||||
| return db.session.query(AppModelConfig).where(AppModelConfig.id == conversation.app_model_config_id).first() | |||||
| return None | return None | ||||
| def agent_thoughts(self): | def agent_thoughts(self): | ||||
| return ( | return ( | ||||
| db.session.query(MessageAgentThought) | db.session.query(MessageAgentThought) | ||||
| .filter(MessageAgentThought.message_id == self.id) | |||||
| .where(MessageAgentThought.message_id == self.id) | |||||
| .order_by(MessageAgentThought.position.asc()) | .order_by(MessageAgentThought.position.asc()) | ||||
| .all() | .all() | ||||
| ) | ) | ||||
| def message_files(self): | def message_files(self): | ||||
| from factories import file_factory | from factories import file_factory | ||||
| message_files = db.session.query(MessageFile).filter(MessageFile.message_id == self.id).all() | |||||
| current_app = db.session.query(App).filter(App.id == self.app_id).first() | |||||
| message_files = db.session.query(MessageFile).where(MessageFile.message_id == self.id).all() | |||||
| current_app = db.session.query(App).where(App.id == self.app_id).first() | |||||
| if not current_app: | if not current_app: | ||||
| raise ValueError(f"App {self.app_id} not found") | raise ValueError(f"App {self.app_id} not found") | ||||
| if self.workflow_run_id: | if self.workflow_run_id: | ||||
| from .workflow import WorkflowRun | from .workflow import WorkflowRun | ||||
| return db.session.query(WorkflowRun).filter(WorkflowRun.id == self.workflow_run_id).first() | |||||
| return db.session.query(WorkflowRun).where(WorkflowRun.id == self.workflow_run_id).first() | |||||
| return None | return None | ||||
| @property | @property | ||||
| def from_account(self): | def from_account(self): | ||||
| account = db.session.query(Account).filter(Account.id == self.from_account_id).first() | |||||
| account = db.session.query(Account).where(Account.id == self.from_account_id).first() | |||||
| return account | return account | ||||
| def to_dict(self): | def to_dict(self): | ||||
| @property | @property | ||||
| def account(self): | def account(self): | ||||
| account = db.session.query(Account).filter(Account.id == self.account_id).first() | |||||
| account = db.session.query(Account).where(Account.id == self.account_id).first() | |||||
| return account | return account | ||||
| @property | @property | ||||
| def annotation_create_account(self): | def annotation_create_account(self): | ||||
| account = db.session.query(Account).filter(Account.id == self.account_id).first() | |||||
| account = db.session.query(Account).where(Account.id == self.account_id).first() | |||||
| return account | return account | ||||
| account = ( | account = ( | ||||
| db.session.query(Account) | db.session.query(Account) | ||||
| .join(MessageAnnotation, MessageAnnotation.account_id == Account.id) | .join(MessageAnnotation, MessageAnnotation.account_id == Account.id) | ||||
| .filter(MessageAnnotation.id == self.annotation_id) | |||||
| .where(MessageAnnotation.id == self.annotation_id) | |||||
| .first() | .first() | ||||
| ) | ) | ||||
| return account | return account | ||||
| @property | @property | ||||
| def annotation_create_account(self): | def annotation_create_account(self): | ||||
| account = db.session.query(Account).filter(Account.id == self.account_id).first() | |||||
| account = db.session.query(Account).where(Account.id == self.account_id).first() | |||||
| return account | return account | ||||
| collection_binding_detail = ( | collection_binding_detail = ( | ||||
| db.session.query(DatasetCollectionBinding) | db.session.query(DatasetCollectionBinding) | ||||
| .filter(DatasetCollectionBinding.id == self.collection_binding_id) | |||||
| .where(DatasetCollectionBinding.id == self.collection_binding_id) | |||||
| .first() | .first() | ||||
| ) | ) | ||||
| return collection_binding_detail | return collection_binding_detail | ||||
| def generate_server_code(n): | def generate_server_code(n): | ||||
| while True: | while True: | ||||
| result = generate_string(n) | result = generate_string(n) | ||||
| while db.session.query(AppMCPServer).filter(AppMCPServer.server_code == result).count() > 0: | |||||
| while db.session.query(AppMCPServer).where(AppMCPServer.server_code == result).count() > 0: | |||||
| result = generate_string(n) | result = generate_string(n) | ||||
| return result | return result | ||||
| def generate_code(n): | def generate_code(n): | ||||
| while True: | while True: | ||||
| result = generate_string(n) | result = generate_string(n) | ||||
| while db.session.query(Site).filter(Site.code == result).count() > 0: | |||||
| while db.session.query(Site).where(Site.code == result).count() > 0: | |||||
| result = generate_string(n) | result = generate_string(n) | ||||
| return result | return result | ||||
| def generate_api_key(prefix, n): | def generate_api_key(prefix, n): | ||||
| while True: | while True: | ||||
| result = prefix + generate_string(n) | result = prefix + generate_string(n) | ||||
| if db.session.query(ApiToken).filter(ApiToken.token == result).count() > 0: | |||||
| if db.session.query(ApiToken).where(ApiToken.token == result).count() > 0: | |||||
| continue | continue | ||||
| return result | return result | ||||
| def user(self) -> Account | None: | def user(self) -> Account | None: | ||||
| if not self.user_id: | if not self.user_id: | ||||
| return None | return None | ||||
| return db.session.query(Account).filter(Account.id == self.user_id).first() | |||||
| return db.session.query(Account).where(Account.id == self.user_id).first() | |||||
| @property | @property | ||||
| def tenant(self) -> Tenant | None: | def tenant(self) -> Tenant | None: | ||||
| return db.session.query(Tenant).filter(Tenant.id == self.tenant_id).first() | |||||
| return db.session.query(Tenant).where(Tenant.id == self.tenant_id).first() | |||||
| class ToolLabelBinding(Base): | class ToolLabelBinding(Base): | ||||
| @property | @property | ||||
| def user(self) -> Account | None: | def user(self) -> Account | None: | ||||
| return db.session.query(Account).filter(Account.id == self.user_id).first() | |||||
| return db.session.query(Account).where(Account.id == self.user_id).first() | |||||
| @property | @property | ||||
| def tenant(self) -> Tenant | None: | def tenant(self) -> Tenant | None: | ||||
| return db.session.query(Tenant).filter(Tenant.id == self.tenant_id).first() | |||||
| return db.session.query(Tenant).where(Tenant.id == self.tenant_id).first() | |||||
| @property | @property | ||||
| def parameter_configurations(self) -> list[WorkflowToolParameterConfiguration]: | def parameter_configurations(self) -> list[WorkflowToolParameterConfiguration]: | ||||
| @property | @property | ||||
| def app(self) -> App | None: | def app(self) -> App | None: | ||||
| return db.session.query(App).filter(App.id == self.app_id).first() | |||||
| return db.session.query(App).where(App.id == self.app_id).first() | |||||
| class MCPToolProvider(Base): | class MCPToolProvider(Base): | ||||
| ) | ) | ||||
| def load_user(self) -> Account | None: | def load_user(self) -> Account | None: | ||||
| return db.session.query(Account).filter(Account.id == self.user_id).first() | |||||
| return db.session.query(Account).where(Account.id == self.user_id).first() | |||||
| @property | @property | ||||
| def tenant(self) -> Tenant | None: | def tenant(self) -> Tenant | None: | ||||
| return db.session.query(Tenant).filter(Tenant.id == self.tenant_id).first() | |||||
| return db.session.query(Tenant).where(Tenant.id == self.tenant_id).first() | |||||
| @property | @property | ||||
| def credentials(self) -> dict: | def credentials(self) -> dict: |
| @property | @property | ||||
| def message(self): | def message(self): | ||||
| return db.session.query(Message).filter(Message.id == self.message_id).first() | |||||
| return db.session.query(Message).where(Message.id == self.message_id).first() | |||||
| class PinnedConversation(Base): | class PinnedConversation(Base): |
| return ( | return ( | ||||
| db.session.query(WorkflowToolProvider) | db.session.query(WorkflowToolProvider) | ||||
| .filter(WorkflowToolProvider.tenant_id == self.tenant_id, WorkflowToolProvider.app_id == self.app_id) | |||||
| .where(WorkflowToolProvider.tenant_id == self.tenant_id, WorkflowToolProvider.app_id == self.app_id) | |||||
| .count() | .count() | ||||
| > 0 | > 0 | ||||
| ) | ) | ||||
| from models.model import Message | from models.model import Message | ||||
| return ( | return ( | ||||
| db.session.query(Message).filter(Message.app_id == self.app_id, Message.workflow_run_id == self.id).first() | |||||
| db.session.query(Message).where(Message.app_id == self.app_id, Message.workflow_run_id == self.id).first() | |||||
| ) | ) | ||||
| @property | @property | ||||
| def workflow(self): | def workflow(self): | ||||
| return db.session.query(Workflow).filter(Workflow.id == self.workflow_id).first() | |||||
| return db.session.query(Workflow).where(Workflow.id == self.workflow_id).first() | |||||
| def to_dict(self): | def to_dict(self): | ||||
| return { | return { |
| try: | try: | ||||
| embedding_ids = ( | embedding_ids = ( | ||||
| db.session.query(Embedding.id) | db.session.query(Embedding.id) | ||||
| .filter(Embedding.created_at < thirty_days_ago) | |||||
| .where(Embedding.created_at < thirty_days_ago) | |||||
| .order_by(Embedding.created_at.desc()) | .order_by(Embedding.created_at.desc()) | ||||
| .limit(100) | .limit(100) | ||||
| .all() | .all() |
| # Main query with join and filter | # Main query with join and filter | ||||
| messages = ( | messages = ( | ||||
| db.session.query(Message) | db.session.query(Message) | ||||
| .filter(Message.created_at < plan_sandbox_clean_message_day) | |||||
| .where(Message.created_at < plan_sandbox_clean_message_day) | |||||
| .order_by(Message.created_at.desc()) | .order_by(Message.created_at.desc()) | ||||
| .limit(100) | .limit(100) | ||||
| .all() | .all() | ||||
| plan = plan_cache.decode() | plan = plan_cache.decode() | ||||
| if plan == "sandbox": | if plan == "sandbox": | ||||
| # clean related message | # clean related message | ||||
| db.session.query(MessageFeedback).filter(MessageFeedback.message_id == message.id).delete( | |||||
| db.session.query(MessageFeedback).where(MessageFeedback.message_id == message.id).delete( | |||||
| synchronize_session=False | synchronize_session=False | ||||
| ) | ) | ||||
| db.session.query(MessageAnnotation).filter(MessageAnnotation.message_id == message.id).delete( | |||||
| db.session.query(MessageAnnotation).where(MessageAnnotation.message_id == message.id).delete( | |||||
| synchronize_session=False | synchronize_session=False | ||||
| ) | ) | ||||
| db.session.query(MessageChain).filter(MessageChain.message_id == message.id).delete( | |||||
| db.session.query(MessageChain).where(MessageChain.message_id == message.id).delete( | |||||
| synchronize_session=False | synchronize_session=False | ||||
| ) | ) | ||||
| db.session.query(MessageAgentThought).filter(MessageAgentThought.message_id == message.id).delete( | |||||
| db.session.query(MessageAgentThought).where(MessageAgentThought.message_id == message.id).delete( | |||||
| synchronize_session=False | synchronize_session=False | ||||
| ) | ) | ||||
| db.session.query(MessageFile).filter(MessageFile.message_id == message.id).delete( | |||||
| db.session.query(MessageFile).where(MessageFile.message_id == message.id).delete( | |||||
| synchronize_session=False | synchronize_session=False | ||||
| ) | ) | ||||
| db.session.query(SavedMessage).filter(SavedMessage.message_id == message.id).delete( | |||||
| db.session.query(SavedMessage).where(SavedMessage.message_id == message.id).delete( | |||||
| synchronize_session=False | synchronize_session=False | ||||
| ) | ) | ||||
| db.session.query(Message).filter(Message.id == message.id).delete() | |||||
| db.session.query(Message).where(Message.id == message.id).delete() | |||||
| db.session.commit() | db.session.commit() | ||||
| end_at = time.perf_counter() | end_at = time.perf_counter() | ||||
| click.echo(click.style("Cleaned messages from db success latency: {}".format(end_at - start_at), fg="green")) | click.echo(click.style("Cleaned messages from db success latency: {}".format(end_at - start_at), fg="green")) |
| # Subquery for counting new documents | # Subquery for counting new documents | ||||
| document_subquery_new = ( | document_subquery_new = ( | ||||
| db.session.query(Document.dataset_id, func.count(Document.id).label("document_count")) | db.session.query(Document.dataset_id, func.count(Document.id).label("document_count")) | ||||
| .filter( | |||||
| .where( | |||||
| Document.indexing_status == "completed", | Document.indexing_status == "completed", | ||||
| Document.enabled == True, | Document.enabled == True, | ||||
| Document.archived == False, | Document.archived == False, | ||||
| # Subquery for counting old documents | # Subquery for counting old documents | ||||
| document_subquery_old = ( | document_subquery_old = ( | ||||
| db.session.query(Document.dataset_id, func.count(Document.id).label("document_count")) | db.session.query(Document.dataset_id, func.count(Document.id).label("document_count")) | ||||
| .filter( | |||||
| .where( | |||||
| Document.indexing_status == "completed", | Document.indexing_status == "completed", | ||||
| Document.enabled == True, | Document.enabled == True, | ||||
| Document.archived == False, | Document.archived == False, | ||||
| select(Dataset) | select(Dataset) | ||||
| .outerjoin(document_subquery_new, Dataset.id == document_subquery_new.c.dataset_id) | .outerjoin(document_subquery_new, Dataset.id == document_subquery_new.c.dataset_id) | ||||
| .outerjoin(document_subquery_old, Dataset.id == document_subquery_old.c.dataset_id) | .outerjoin(document_subquery_old, Dataset.id == document_subquery_old.c.dataset_id) | ||||
| .filter( | |||||
| .where( | |||||
| Dataset.created_at < plan_sandbox_clean_day, | Dataset.created_at < plan_sandbox_clean_day, | ||||
| func.coalesce(document_subquery_new.c.document_count, 0) == 0, | func.coalesce(document_subquery_new.c.document_count, 0) == 0, | ||||
| func.coalesce(document_subquery_old.c.document_count, 0) > 0, | func.coalesce(document_subquery_old.c.document_count, 0) > 0, | ||||
| for dataset in datasets: | for dataset in datasets: | ||||
| dataset_query = ( | dataset_query = ( | ||||
| db.session.query(DatasetQuery) | db.session.query(DatasetQuery) | ||||
| .filter(DatasetQuery.created_at > plan_sandbox_clean_day, DatasetQuery.dataset_id == dataset.id) | |||||
| .where(DatasetQuery.created_at > plan_sandbox_clean_day, DatasetQuery.dataset_id == dataset.id) | |||||
| .all() | .all() | ||||
| ) | ) | ||||
| if not dataset_query or len(dataset_query) == 0: | if not dataset_query or len(dataset_query) == 0: | ||||
| # add auto disable log | # add auto disable log | ||||
| documents = ( | documents = ( | ||||
| db.session.query(Document) | db.session.query(Document) | ||||
| .filter( | |||||
| .where( | |||||
| Document.dataset_id == dataset.id, | Document.dataset_id == dataset.id, | ||||
| Document.enabled == True, | Document.enabled == True, | ||||
| Document.archived == False, | Document.archived == False, | ||||
| # Subquery for counting new documents | # Subquery for counting new documents | ||||
| document_subquery_new = ( | document_subquery_new = ( | ||||
| db.session.query(Document.dataset_id, func.count(Document.id).label("document_count")) | db.session.query(Document.dataset_id, func.count(Document.id).label("document_count")) | ||||
| .filter( | |||||
| .where( | |||||
| Document.indexing_status == "completed", | Document.indexing_status == "completed", | ||||
| Document.enabled == True, | Document.enabled == True, | ||||
| Document.archived == False, | Document.archived == False, | ||||
| # Subquery for counting old documents | # Subquery for counting old documents | ||||
| document_subquery_old = ( | document_subquery_old = ( | ||||
| db.session.query(Document.dataset_id, func.count(Document.id).label("document_count")) | db.session.query(Document.dataset_id, func.count(Document.id).label("document_count")) | ||||
| .filter( | |||||
| .where( | |||||
| Document.indexing_status == "completed", | Document.indexing_status == "completed", | ||||
| Document.enabled == True, | Document.enabled == True, | ||||
| Document.archived == False, | Document.archived == False, | ||||
| select(Dataset) | select(Dataset) | ||||
| .outerjoin(document_subquery_new, Dataset.id == document_subquery_new.c.dataset_id) | .outerjoin(document_subquery_new, Dataset.id == document_subquery_new.c.dataset_id) | ||||
| .outerjoin(document_subquery_old, Dataset.id == document_subquery_old.c.dataset_id) | .outerjoin(document_subquery_old, Dataset.id == document_subquery_old.c.dataset_id) | ||||
| .filter( | |||||
| .where( | |||||
| Dataset.created_at < plan_pro_clean_day, | Dataset.created_at < plan_pro_clean_day, | ||||
| func.coalesce(document_subquery_new.c.document_count, 0) == 0, | func.coalesce(document_subquery_new.c.document_count, 0) == 0, | ||||
| func.coalesce(document_subquery_old.c.document_count, 0) > 0, | func.coalesce(document_subquery_old.c.document_count, 0) > 0, | ||||
| for dataset in datasets: | for dataset in datasets: | ||||
| dataset_query = ( | dataset_query = ( | ||||
| db.session.query(DatasetQuery) | db.session.query(DatasetQuery) | ||||
| .filter(DatasetQuery.created_at > plan_pro_clean_day, DatasetQuery.dataset_id == dataset.id) | |||||
| .where(DatasetQuery.created_at > plan_pro_clean_day, DatasetQuery.dataset_id == dataset.id) | |||||
| .all() | .all() | ||||
| ) | ) | ||||
| if not dataset_query or len(dataset_query) == 0: | if not dataset_query or len(dataset_query) == 0: |
| try: | try: | ||||
| # check the number of idle tidb serverless | # check the number of idle tidb serverless | ||||
| idle_tidb_serverless_number = ( | idle_tidb_serverless_number = ( | ||||
| db.session.query(TidbAuthBinding).filter(TidbAuthBinding.active == False).count() | |||||
| db.session.query(TidbAuthBinding).where(TidbAuthBinding.active == False).count() | |||||
| ) | ) | ||||
| if idle_tidb_serverless_number >= tidb_serverless_number: | if idle_tidb_serverless_number >= tidb_serverless_number: | ||||
| break | break |
| # send document clean notify mail | # send document clean notify mail | ||||
| try: | try: | ||||
| dataset_auto_disable_logs = ( | dataset_auto_disable_logs = ( | ||||
| db.session.query(DatasetAutoDisableLog).filter(DatasetAutoDisableLog.notified == False).all() | |||||
| db.session.query(DatasetAutoDisableLog).where(DatasetAutoDisableLog.notified == False).all() | |||||
| ) | ) | ||||
| # group by tenant_id | # group by tenant_id | ||||
| dataset_auto_disable_logs_map: dict[str, list[DatasetAutoDisableLog]] = defaultdict(list) | dataset_auto_disable_logs_map: dict[str, list[DatasetAutoDisableLog]] = defaultdict(list) | ||||
| if plan != "sandbox": | if plan != "sandbox": | ||||
| knowledge_details = [] | knowledge_details = [] | ||||
| # check tenant | # check tenant | ||||
| tenant = db.session.query(Tenant).filter(Tenant.id == tenant_id).first() | |||||
| tenant = db.session.query(Tenant).where(Tenant.id == tenant_id).first() | |||||
| if not tenant: | if not tenant: | ||||
| continue | continue | ||||
| # check current owner | # check current owner | ||||
| ) | ) | ||||
| if not current_owner_join: | if not current_owner_join: | ||||
| continue | continue | ||||
| account = db.session.query(Account).filter(Account.id == current_owner_join.account_id).first() | |||||
| account = db.session.query(Account).where(Account.id == current_owner_join.account_id).first() | |||||
| if not account: | if not account: | ||||
| continue | continue | ||||
| ) | ) | ||||
| for dataset_id, document_ids in dataset_auto_dataset_map.items(): | for dataset_id, document_ids in dataset_auto_dataset_map.items(): | ||||
| dataset = db.session.query(Dataset).filter(Dataset.id == dataset_id).first() | |||||
| dataset = db.session.query(Dataset).where(Dataset.id == dataset_id).first() | |||||
| if dataset: | if dataset: | ||||
| document_count = len(document_ids) | document_count = len(document_ids) | ||||
| knowledge_details.append(rf"Knowledge base {dataset.name}: {document_count} documents") | knowledge_details.append(rf"Knowledge base {dataset.name}: {document_count} documents") |
| # check the number of idle tidb serverless | # check the number of idle tidb serverless | ||||
| tidb_serverless_list = ( | tidb_serverless_list = ( | ||||
| db.session.query(TidbAuthBinding) | db.session.query(TidbAuthBinding) | ||||
| .filter(TidbAuthBinding.active == False, TidbAuthBinding.status == "CREATING") | |||||
| .where(TidbAuthBinding.active == False, TidbAuthBinding.status == "CREATING") | |||||
| .all() | .all() | ||||
| ) | ) | ||||
| if len(tidb_serverless_list) == 0: | if len(tidb_serverless_list) == 0: |
| ) | ) | ||||
| ) | ) | ||||
| account = db.session.query(Account).filter(Account.email == email).first() | |||||
| account = db.session.query(Account).where(Account.email == email).first() | |||||
| if not account: | if not account: | ||||
| return None | return None | ||||
| return ( | return ( | ||||
| db.session.query(Tenant) | db.session.query(Tenant) | ||||
| .join(TenantAccountJoin, Tenant.id == TenantAccountJoin.tenant_id) | .join(TenantAccountJoin, Tenant.id == TenantAccountJoin.tenant_id) | ||||
| .filter(TenantAccountJoin.account_id == account.id, Tenant.status == TenantStatus.NORMAL) | |||||
| .where(TenantAccountJoin.account_id == account.id, Tenant.status == TenantStatus.NORMAL) | |||||
| .all() | .all() | ||||
| ) | ) | ||||
| tenant_account_join = ( | tenant_account_join = ( | ||||
| db.session.query(TenantAccountJoin) | db.session.query(TenantAccountJoin) | ||||
| .join(Tenant, TenantAccountJoin.tenant_id == Tenant.id) | .join(Tenant, TenantAccountJoin.tenant_id == Tenant.id) | ||||
| .filter( | |||||
| .where( | |||||
| TenantAccountJoin.account_id == account.id, | TenantAccountJoin.account_id == account.id, | ||||
| TenantAccountJoin.tenant_id == tenant_id, | TenantAccountJoin.tenant_id == tenant_id, | ||||
| Tenant.status == TenantStatus.NORMAL, | Tenant.status == TenantStatus.NORMAL, | ||||
| if not tenant_account_join: | if not tenant_account_join: | ||||
| raise AccountNotLinkTenantError("Tenant not found or account is not a member of the tenant.") | raise AccountNotLinkTenantError("Tenant not found or account is not a member of the tenant.") | ||||
| else: | else: | ||||
| db.session.query(TenantAccountJoin).filter( | |||||
| db.session.query(TenantAccountJoin).where( | |||||
| TenantAccountJoin.account_id == account.id, TenantAccountJoin.tenant_id != tenant_id | TenantAccountJoin.account_id == account.id, TenantAccountJoin.tenant_id != tenant_id | ||||
| ).update({"current": False}) | ).update({"current": False}) | ||||
| tenant_account_join.current = True | tenant_account_join.current = True | ||||
| db.session.query(Account, TenantAccountJoin.role) | db.session.query(Account, TenantAccountJoin.role) | ||||
| .select_from(Account) | .select_from(Account) | ||||
| .join(TenantAccountJoin, Account.id == TenantAccountJoin.account_id) | .join(TenantAccountJoin, Account.id == TenantAccountJoin.account_id) | ||||
| .filter(TenantAccountJoin.tenant_id == tenant.id) | |||||
| .where(TenantAccountJoin.tenant_id == tenant.id) | |||||
| ) | ) | ||||
| # Initialize an empty list to store the updated accounts | # Initialize an empty list to store the updated accounts | ||||
| db.session.query(Account, TenantAccountJoin.role) | db.session.query(Account, TenantAccountJoin.role) | ||||
| .select_from(Account) | .select_from(Account) | ||||
| .join(TenantAccountJoin, Account.id == TenantAccountJoin.account_id) | .join(TenantAccountJoin, Account.id == TenantAccountJoin.account_id) | ||||
| .filter(TenantAccountJoin.tenant_id == tenant.id) | |||||
| .filter(TenantAccountJoin.role == "dataset_operator") | |||||
| .where(TenantAccountJoin.tenant_id == tenant.id) | |||||
| .where(TenantAccountJoin.role == "dataset_operator") | |||||
| ) | ) | ||||
| # Initialize an empty list to store the updated accounts | # Initialize an empty list to store the updated accounts | ||||
| return ( | return ( | ||||
| db.session.query(TenantAccountJoin) | db.session.query(TenantAccountJoin) | ||||
| .filter( | |||||
| TenantAccountJoin.tenant_id == tenant.id, TenantAccountJoin.role.in_([role.value for role in roles]) | |||||
| ) | |||||
| .where(TenantAccountJoin.tenant_id == tenant.id, TenantAccountJoin.role.in_([role.value for role in roles])) | |||||
| .first() | .first() | ||||
| is not None | is not None | ||||
| ) | ) | ||||
| """Get the role of the current account for a given tenant""" | """Get the role of the current account for a given tenant""" | ||||
| join = ( | join = ( | ||||
| db.session.query(TenantAccountJoin) | db.session.query(TenantAccountJoin) | ||||
| .filter(TenantAccountJoin.tenant_id == tenant.id, TenantAccountJoin.account_id == account.id) | |||||
| .where(TenantAccountJoin.tenant_id == tenant.id, TenantAccountJoin.account_id == account.id) | |||||
| .first() | .first() | ||||
| ) | ) | ||||
| return TenantAccountRole(join.role) if join else None | return TenantAccountRole(join.role) if join else None | ||||
| tenant = ( | tenant = ( | ||||
| db.session.query(Tenant) | db.session.query(Tenant) | ||||
| .filter(Tenant.id == invitation_data["workspace_id"], Tenant.status == "normal") | |||||
| .where(Tenant.id == invitation_data["workspace_id"], Tenant.status == "normal") | |||||
| .first() | .first() | ||||
| ) | ) | ||||
| tenant_account = ( | tenant_account = ( | ||||
| db.session.query(Account, TenantAccountJoin.role) | db.session.query(Account, TenantAccountJoin.role) | ||||
| .join(TenantAccountJoin, Account.id == TenantAccountJoin.account_id) | .join(TenantAccountJoin, Account.id == TenantAccountJoin.account_id) | ||||
| .filter(Account.email == invitation_data["email"], TenantAccountJoin.tenant_id == tenant.id) | |||||
| .where(Account.email == invitation_data["email"], TenantAccountJoin.tenant_id == tenant.id) | |||||
| .first() | .first() | ||||
| ) | ) | ||||
| conversation: Conversation | None = ( | conversation: Conversation | None = ( | ||||
| db.session.query(Conversation) | db.session.query(Conversation) | ||||
| .filter( | |||||
| .where( | |||||
| Conversation.id == conversation_id, | Conversation.id == conversation_id, | ||||
| Conversation.app_id == app_model.id, | Conversation.app_id == app_model.id, | ||||
| ) | ) | ||||
| message: Optional[Message] = ( | message: Optional[Message] = ( | ||||
| db.session.query(Message) | db.session.query(Message) | ||||
| .filter( | |||||
| .where( | |||||
| Message.id == message_id, | Message.id == message_id, | ||||
| Message.conversation_id == conversation_id, | Message.conversation_id == conversation_id, | ||||
| ) | ) | ||||
| if conversation.from_end_user_id: | if conversation.from_end_user_id: | ||||
| # only select name field | # only select name field | ||||
| executor = ( | executor = ( | ||||
| db.session.query(EndUser, EndUser.name).filter(EndUser.id == conversation.from_end_user_id).first() | |||||
| db.session.query(EndUser, EndUser.name).where(EndUser.id == conversation.from_end_user_id).first() | |||||
| ) | ) | ||||
| else: | else: | ||||
| executor = ( | |||||
| db.session.query(Account, Account.name).filter(Account.id == conversation.from_account_id).first() | |||||
| ) | |||||
| executor = db.session.query(Account, Account.name).where(Account.id == conversation.from_account_id).first() | |||||
| if executor: | if executor: | ||||
| executor = executor.name | executor = executor.name |
| # get app info | # get app info | ||||
| app = ( | app = ( | ||||
| db.session.query(App) | db.session.query(App) | ||||
| .filter(App.id == app_id, App.tenant_id == current_user.current_tenant_id, App.status == "normal") | |||||
| .where(App.id == app_id, App.tenant_id == current_user.current_tenant_id, App.status == "normal") | |||||
| .first() | .first() | ||||
| ) | ) | ||||
| if args.get("message_id"): | if args.get("message_id"): | ||||
| message_id = str(args["message_id"]) | message_id = str(args["message_id"]) | ||||
| # get message info | # get message info | ||||
| message = db.session.query(Message).filter(Message.id == message_id, Message.app_id == app.id).first() | |||||
| message = db.session.query(Message).where(Message.id == message_id, Message.app_id == app.id).first() | |||||
| if not message: | if not message: | ||||
| raise NotFound("Message Not Exists.") | raise NotFound("Message Not Exists.") | ||||
| db.session.add(annotation) | db.session.add(annotation) | ||||
| db.session.commit() | db.session.commit() | ||||
| # if annotation reply is enabled , add annotation to index | # if annotation reply is enabled , add annotation to index | ||||
| annotation_setting = ( | |||||
| db.session.query(AppAnnotationSetting).filter(AppAnnotationSetting.app_id == app_id).first() | |||||
| ) | |||||
| annotation_setting = db.session.query(AppAnnotationSetting).where(AppAnnotationSetting.app_id == app_id).first() | |||||
| if annotation_setting: | if annotation_setting: | ||||
| add_annotation_to_index_task.delay( | add_annotation_to_index_task.delay( | ||||
| annotation.id, | annotation.id, | ||||
| # get app info | # get app info | ||||
| app = ( | app = ( | ||||
| db.session.query(App) | db.session.query(App) | ||||
| .filter(App.id == app_id, App.tenant_id == current_user.current_tenant_id, App.status == "normal") | |||||
| .where(App.id == app_id, App.tenant_id == current_user.current_tenant_id, App.status == "normal") | |||||
| .first() | .first() | ||||
| ) | ) | ||||
| if keyword: | if keyword: | ||||
| stmt = ( | stmt = ( | ||||
| select(MessageAnnotation) | select(MessageAnnotation) | ||||
| .filter(MessageAnnotation.app_id == app_id) | |||||
| .filter( | |||||
| .where(MessageAnnotation.app_id == app_id) | |||||
| .where( | |||||
| or_( | or_( | ||||
| MessageAnnotation.question.ilike("%{}%".format(keyword)), | MessageAnnotation.question.ilike("%{}%".format(keyword)), | ||||
| MessageAnnotation.content.ilike("%{}%".format(keyword)), | MessageAnnotation.content.ilike("%{}%".format(keyword)), | ||||
| else: | else: | ||||
| stmt = ( | stmt = ( | ||||
| select(MessageAnnotation) | select(MessageAnnotation) | ||||
| .filter(MessageAnnotation.app_id == app_id) | |||||
| .where(MessageAnnotation.app_id == app_id) | |||||
| .order_by(MessageAnnotation.created_at.desc(), MessageAnnotation.id.desc()) | .order_by(MessageAnnotation.created_at.desc(), MessageAnnotation.id.desc()) | ||||
| ) | ) | ||||
| annotations = db.paginate(select=stmt, page=page, per_page=limit, max_per_page=100, error_out=False) | annotations = db.paginate(select=stmt, page=page, per_page=limit, max_per_page=100, error_out=False) | ||||
| # get app info | # get app info | ||||
| app = ( | app = ( | ||||
| db.session.query(App) | db.session.query(App) | ||||
| .filter(App.id == app_id, App.tenant_id == current_user.current_tenant_id, App.status == "normal") | |||||
| .where(App.id == app_id, App.tenant_id == current_user.current_tenant_id, App.status == "normal") | |||||
| .first() | .first() | ||||
| ) | ) | ||||
| raise NotFound("App not found") | raise NotFound("App not found") | ||||
| annotations = ( | annotations = ( | ||||
| db.session.query(MessageAnnotation) | db.session.query(MessageAnnotation) | ||||
| .filter(MessageAnnotation.app_id == app_id) | |||||
| .where(MessageAnnotation.app_id == app_id) | |||||
| .order_by(MessageAnnotation.created_at.desc()) | .order_by(MessageAnnotation.created_at.desc()) | ||||
| .all() | .all() | ||||
| ) | ) | ||||
| # get app info | # get app info | ||||
| app = ( | app = ( | ||||
| db.session.query(App) | db.session.query(App) | ||||
| .filter(App.id == app_id, App.tenant_id == current_user.current_tenant_id, App.status == "normal") | |||||
| .where(App.id == app_id, App.tenant_id == current_user.current_tenant_id, App.status == "normal") | |||||
| .first() | .first() | ||||
| ) | ) | ||||
| db.session.add(annotation) | db.session.add(annotation) | ||||
| db.session.commit() | db.session.commit() | ||||
| # if annotation reply is enabled , add annotation to index | # if annotation reply is enabled , add annotation to index | ||||
| annotation_setting = ( | |||||
| db.session.query(AppAnnotationSetting).filter(AppAnnotationSetting.app_id == app_id).first() | |||||
| ) | |||||
| annotation_setting = db.session.query(AppAnnotationSetting).where(AppAnnotationSetting.app_id == app_id).first() | |||||
| if annotation_setting: | if annotation_setting: | ||||
| add_annotation_to_index_task.delay( | add_annotation_to_index_task.delay( | ||||
| annotation.id, | annotation.id, | ||||
| # get app info | # get app info | ||||
| app = ( | app = ( | ||||
| db.session.query(App) | db.session.query(App) | ||||
| .filter(App.id == app_id, App.tenant_id == current_user.current_tenant_id, App.status == "normal") | |||||
| .where(App.id == app_id, App.tenant_id == current_user.current_tenant_id, App.status == "normal") | |||||
| .first() | .first() | ||||
| ) | ) | ||||
| if not app: | if not app: | ||||
| raise NotFound("App not found") | raise NotFound("App not found") | ||||
| annotation = db.session.query(MessageAnnotation).filter(MessageAnnotation.id == annotation_id).first() | |||||
| annotation = db.session.query(MessageAnnotation).where(MessageAnnotation.id == annotation_id).first() | |||||
| if not annotation: | if not annotation: | ||||
| raise NotFound("Annotation not found") | raise NotFound("Annotation not found") | ||||
| db.session.commit() | db.session.commit() | ||||
| # if annotation reply is enabled , add annotation to index | # if annotation reply is enabled , add annotation to index | ||||
| app_annotation_setting = ( | app_annotation_setting = ( | ||||
| db.session.query(AppAnnotationSetting).filter(AppAnnotationSetting.app_id == app_id).first() | |||||
| db.session.query(AppAnnotationSetting).where(AppAnnotationSetting.app_id == app_id).first() | |||||
| ) | ) | ||||
| if app_annotation_setting: | if app_annotation_setting: | ||||
| # get app info | # get app info | ||||
| app = ( | app = ( | ||||
| db.session.query(App) | db.session.query(App) | ||||
| .filter(App.id == app_id, App.tenant_id == current_user.current_tenant_id, App.status == "normal") | |||||
| .where(App.id == app_id, App.tenant_id == current_user.current_tenant_id, App.status == "normal") | |||||
| .first() | .first() | ||||
| ) | ) | ||||
| if not app: | if not app: | ||||
| raise NotFound("App not found") | raise NotFound("App not found") | ||||
| annotation = db.session.query(MessageAnnotation).filter(MessageAnnotation.id == annotation_id).first() | |||||
| annotation = db.session.query(MessageAnnotation).where(MessageAnnotation.id == annotation_id).first() | |||||
| if not annotation: | if not annotation: | ||||
| raise NotFound("Annotation not found") | raise NotFound("Annotation not found") | ||||
| annotation_hit_histories = ( | annotation_hit_histories = ( | ||||
| db.session.query(AppAnnotationHitHistory) | db.session.query(AppAnnotationHitHistory) | ||||
| .filter(AppAnnotationHitHistory.annotation_id == annotation_id) | |||||
| .where(AppAnnotationHitHistory.annotation_id == annotation_id) | |||||
| .all() | .all() | ||||
| ) | ) | ||||
| if annotation_hit_histories: | if annotation_hit_histories: | ||||
| db.session.commit() | db.session.commit() | ||||
| # if annotation reply is enabled , delete annotation index | # if annotation reply is enabled , delete annotation index | ||||
| app_annotation_setting = ( | app_annotation_setting = ( | ||||
| db.session.query(AppAnnotationSetting).filter(AppAnnotationSetting.app_id == app_id).first() | |||||
| db.session.query(AppAnnotationSetting).where(AppAnnotationSetting.app_id == app_id).first() | |||||
| ) | ) | ||||
| if app_annotation_setting: | if app_annotation_setting: | ||||
| # get app info | # get app info | ||||
| app = ( | app = ( | ||||
| db.session.query(App) | db.session.query(App) | ||||
| .filter(App.id == app_id, App.tenant_id == current_user.current_tenant_id, App.status == "normal") | |||||
| .where(App.id == app_id, App.tenant_id == current_user.current_tenant_id, App.status == "normal") | |||||
| .first() | .first() | ||||
| ) | ) | ||||
| # get app info | # get app info | ||||
| app = ( | app = ( | ||||
| db.session.query(App) | db.session.query(App) | ||||
| .filter(App.id == app_id, App.tenant_id == current_user.current_tenant_id, App.status == "normal") | |||||
| .where(App.id == app_id, App.tenant_id == current_user.current_tenant_id, App.status == "normal") | |||||
| .first() | .first() | ||||
| ) | ) | ||||
| if not app: | if not app: | ||||
| raise NotFound("App not found") | raise NotFound("App not found") | ||||
| annotation = db.session.query(MessageAnnotation).filter(MessageAnnotation.id == annotation_id).first() | |||||
| annotation = db.session.query(MessageAnnotation).where(MessageAnnotation.id == annotation_id).first() | |||||
| if not annotation: | if not annotation: | ||||
| raise NotFound("Annotation not found") | raise NotFound("Annotation not found") | ||||
| stmt = ( | stmt = ( | ||||
| select(AppAnnotationHitHistory) | select(AppAnnotationHitHistory) | ||||
| .filter( | |||||
| .where( | |||||
| AppAnnotationHitHistory.app_id == app_id, | AppAnnotationHitHistory.app_id == app_id, | ||||
| AppAnnotationHitHistory.annotation_id == annotation_id, | AppAnnotationHitHistory.annotation_id == annotation_id, | ||||
| ) | ) | ||||
| @classmethod | @classmethod | ||||
| def get_annotation_by_id(cls, annotation_id: str) -> MessageAnnotation | None: | def get_annotation_by_id(cls, annotation_id: str) -> MessageAnnotation | None: | ||||
| annotation = db.session.query(MessageAnnotation).filter(MessageAnnotation.id == annotation_id).first() | |||||
| annotation = db.session.query(MessageAnnotation).where(MessageAnnotation.id == annotation_id).first() | |||||
| if not annotation: | if not annotation: | ||||
| return None | return None | ||||
| score: float, | score: float, | ||||
| ): | ): | ||||
| # add hit count to annotation | # add hit count to annotation | ||||
| db.session.query(MessageAnnotation).filter(MessageAnnotation.id == annotation_id).update( | |||||
| db.session.query(MessageAnnotation).where(MessageAnnotation.id == annotation_id).update( | |||||
| {MessageAnnotation.hit_count: MessageAnnotation.hit_count + 1}, synchronize_session=False | {MessageAnnotation.hit_count: MessageAnnotation.hit_count + 1}, synchronize_session=False | ||||
| ) | ) | ||||
| # get app info | # get app info | ||||
| app = ( | app = ( | ||||
| db.session.query(App) | db.session.query(App) | ||||
| .filter(App.id == app_id, App.tenant_id == current_user.current_tenant_id, App.status == "normal") | |||||
| .where(App.id == app_id, App.tenant_id == current_user.current_tenant_id, App.status == "normal") | |||||
| .first() | .first() | ||||
| ) | ) | ||||
| if not app: | if not app: | ||||
| raise NotFound("App not found") | raise NotFound("App not found") | ||||
| annotation_setting = ( | |||||
| db.session.query(AppAnnotationSetting).filter(AppAnnotationSetting.app_id == app_id).first() | |||||
| ) | |||||
| annotation_setting = db.session.query(AppAnnotationSetting).where(AppAnnotationSetting.app_id == app_id).first() | |||||
| if annotation_setting: | if annotation_setting: | ||||
| collection_binding_detail = annotation_setting.collection_binding_detail | collection_binding_detail = annotation_setting.collection_binding_detail | ||||
| return { | return { | ||||
| # get app info | # get app info | ||||
| app = ( | app = ( | ||||
| db.session.query(App) | db.session.query(App) | ||||
| .filter(App.id == app_id, App.tenant_id == current_user.current_tenant_id, App.status == "normal") | |||||
| .where(App.id == app_id, App.tenant_id == current_user.current_tenant_id, App.status == "normal") | |||||
| .first() | .first() | ||||
| ) | ) | ||||
| annotation_setting = ( | annotation_setting = ( | ||||
| db.session.query(AppAnnotationSetting) | db.session.query(AppAnnotationSetting) | ||||
| .filter( | |||||
| .where( | |||||
| AppAnnotationSetting.app_id == app_id, | AppAnnotationSetting.app_id == app_id, | ||||
| AppAnnotationSetting.id == annotation_setting_id, | AppAnnotationSetting.id == annotation_setting_id, | ||||
| ) | ) |
| db.session.query(APIBasedExtension) | db.session.query(APIBasedExtension) | ||||
| .filter_by(tenant_id=extension_data.tenant_id) | .filter_by(tenant_id=extension_data.tenant_id) | ||||
| .filter_by(name=extension_data.name) | .filter_by(name=extension_data.name) | ||||
| .filter(APIBasedExtension.id != extension_data.id) | |||||
| .where(APIBasedExtension.id != extension_data.id) | |||||
| .first() | .first() | ||||
| ) | ) | ||||
| elif provider_type == "api": | elif provider_type == "api": | ||||
| try: | try: | ||||
| provider: Optional[ApiToolProvider] = ( | provider: Optional[ApiToolProvider] = ( | ||||
| db.session.query(ApiToolProvider).filter(ApiToolProvider.id == provider_id).first() | |||||
| db.session.query(ApiToolProvider).where(ApiToolProvider.id == provider_id).first() | |||||
| ) | ) | ||||
| if provider is None: | if provider is None: | ||||
| raise ValueError(f"provider not found for tool {tool_name}") | raise ValueError(f"provider not found for tool {tool_name}") | ||||
| :param app_id: app id | :param app_id: app id | ||||
| :return: app code | :return: app code | ||||
| """ | """ | ||||
| site = db.session.query(Site).filter(Site.app_id == app_id).first() | |||||
| site = db.session.query(Site).where(Site.app_id == app_id).first() | |||||
| if not site: | if not site: | ||||
| raise ValueError(f"App with id {app_id} not found") | raise ValueError(f"App with id {app_id} not found") | ||||
| return str(site.code) | return str(site.code) | ||||
| :param app_code: app code | :param app_code: app code | ||||
| :return: app id | :return: app id | ||||
| """ | """ | ||||
| site = db.session.query(Site).filter(Site.code == app_code).first() | |||||
| site = db.session.query(Site).where(Site.code == app_code).first() | |||||
| if not site: | if not site: | ||||
| raise ValueError(f"App with code {app_code} not found") | raise ValueError(f"App with code {app_code} not found") | ||||
| return str(site.app_id) | return str(site.app_id) |