Initial commit

api/events/event_handlers/__init__.py

@@ -0,0 +1,9 @@
+from .create_installed_app_when_app_created import handle
+from .delete_installed_app_when_app_deleted import handle
+from .create_provider_when_tenant_created import handle
+from .create_provider_when_tenant_updated import handle
+from .clean_when_document_deleted import handle
+from .clean_when_dataset_deleted import handle
+from .update_app_dataset_join_when_app_model_config_updated import handle
+from .generate_conversation_name_when_first_message_created import handle
+from .generate_conversation_summary_when_few_message_created import handle

api/events/event_handlers/clean_when_dataset_deleted.py

@@ -0,0 +1,8 @@
+from events.dataset_event import dataset_was_deleted
+from tasks.clean_dataset_task import clean_dataset_task
+
+
+@dataset_was_deleted.connect
+def handle(sender, **kwargs):
+    dataset = sender
+    clean_dataset_task.delay(dataset.id, dataset.tenant_id, dataset.indexing_technique, dataset.index_struct)

api/events/event_handlers/clean_when_document_deleted.py

@@ -0,0 +1,9 @@
+from events.document_event import document_was_deleted
+from tasks.clean_document_task import clean_document_task
+
+
+@document_was_deleted.connect
+def handle(sender, **kwargs):
+    document_id = sender
+    dataset_id = kwargs.get('dataset_id')
+    clean_document_task.delay(document_id, dataset_id)

api/events/event_handlers/create_installed_app_when_app_created.py

@@ -0,0 +1,16 @@
+from events.app_event import app_was_created
+from extensions.ext_database import db
+from models.model import InstalledApp
+
+
+@app_was_created.connect
+def handle(sender, **kwargs):
+    """Create an installed app when an app is created."""
+    app = sender
+    installed_app = InstalledApp(
+        tenant_id=app.tenant_id,
+        app_id=app.id,
+        app_owner_tenant_id=app.tenant_id
+    )
+    db.session.add(installed_app)
+    db.session.commit()

api/events/event_handlers/create_provider_when_tenant_created.py

@@ -0,0 +1,9 @@
+from events.tenant_event import tenant_was_updated
+from services.provider_service import ProviderService
+
+
+@tenant_was_updated.connect
+def handle(sender, **kwargs):
+    tenant = sender
+    if tenant.status == 'normal':
+        ProviderService.create_system_provider(tenant)

api/events/event_handlers/create_provider_when_tenant_updated.py

@@ -0,0 +1,9 @@
+from events.tenant_event import tenant_was_created
+from services.provider_service import ProviderService
+
+
+@tenant_was_created.connect
+def handle(sender, **kwargs):
+    tenant = sender
+    if tenant.status == 'normal':
+        ProviderService.create_system_provider(tenant)

api/events/event_handlers/delete_installed_app_when_app_deleted.py

@@ -0,0 +1,12 @@
+from events.app_event import app_was_deleted
+from extensions.ext_database import db
+from models.model import InstalledApp
+
+
+@app_was_deleted.connect
+def handle(sender, **kwargs):
+    app = sender
+    installed_apps = db.session.query(InstalledApp).filter(InstalledApp.app_id == app.id).all()
+    for installed_app in installed_apps:
+        db.session.delete(installed_app)
+    db.session.commit()

api/events/event_handlers/generate_conversation_name_when_first_message_created.py

@@ -0,0 +1,29 @@
+import logging
+
+from core.generator.llm_generator import LLMGenerator
+from events.message_event import message_was_created
+from extensions.ext_database import db
+
+
+@message_was_created.connect
+def handle(sender, **kwargs):
+    message = sender
+    conversation = kwargs.get('conversation')
+    is_first_message = kwargs.get('is_first_message')
+
+    if is_first_message:
+        if conversation.mode == 'chat':
+            app_model = conversation.app
+            if not app_model:
+                return
+
+            # generate conversation name
+            try:
+                name = LLMGenerator.generate_conversation_name(app_model.tenant_id, message.query, message.answer)
+                conversation.name = name
+            except:
+                conversation.name = 'New Chat'
+                logging.exception('generate_conversation_name failed')
+
+            db.session.add(conversation)
+            db.session.commit()

api/events/event_handlers/generate_conversation_summary_when_few_message_created.py

@@ -0,0 +1,14 @@
+from events.message_event import message_was_created
+from tasks.generate_conversation_summary_task import generate_conversation_summary_task
+
+
+@message_was_created.connect
+def handle(sender, **kwargs):
+    message = sender
+    conversation = kwargs.get('conversation')
+    is_first_message = kwargs.get('is_first_message')
+
+    if not is_first_message and conversation.mode == 'chat' and not conversation.summary:
+        history_message_count = conversation.message_count
+        if history_message_count >= 5:
+            generate_conversation_summary_task.delay(conversation.id)

api/events/event_handlers/update_app_dataset_join_when_app_model_config_updated.py

@@ -0,0 +1,66 @@
+from events.app_event import app_model_config_was_updated
+from extensions.ext_database import db
+from models.dataset import AppDatasetJoin
+from models.model import AppModelConfig
+
+
+@app_model_config_was_updated.connect
+def handle(sender, **kwargs):
+    app_model = sender
+    app_model_config = kwargs.get('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_model.id
+    ).all()
+
+    removed_dataset_ids = []
+    if not app_dataset_joins:
+        added_dataset_ids = dataset_ids
+    else:
+        old_dataset_ids = set()
+        for app_dataset_join in app_dataset_joins:
+            old_dataset_ids.add(app_dataset_join.dataset_id)
+
+        added_dataset_ids = dataset_ids - old_dataset_ids
+        removed_dataset_ids = old_dataset_ids - dataset_ids
+
+    if removed_dataset_ids:
+        for dataset_id in removed_dataset_ids:
+            db.session.query(AppDatasetJoin).filter(
+                AppDatasetJoin.app_id == app_model.id,
+                AppDatasetJoin.dataset_id == dataset_id
+            ).delete()
+
+    if added_dataset_ids:
+        for dataset_id in added_dataset_ids:
+            app_dataset_join = AppDatasetJoin(
+                app_id=app_model.id,
+                dataset_id=dataset_id
+            )
+            db.session.add(app_dataset_join)
+
+    db.session.commit()
+
+
+def get_dataset_ids_from_model_config(app_model_config: AppModelConfig) -> set:
+    dataset_ids = set()
+    if not app_model_config:
+        return dataset_ids
+
+    agent_mode = app_model_config.agent_mode_dict
+    if agent_mode.get('enabled') is False:
+        return dataset_ids
+
+    if not agent_mode.get('tools'):
+        return dataset_ids
+
+    tools = agent_mode.get('tools')
+    for tool in tools:
+        tool_type = list(tool.keys())[0]
+        tool_config = list(tool.values())[0]
+        if tool_type == "dataset":
+            dataset_ids.add(tool_config.get("id"))
+
+    return dataset_ids