Merge remote-tracking branch 'origin/main' into feat/llm-ops-tracing

* origin/main:
  chore: set build system to Poetry and remove unnecessary settings with package mode disabled  (#5263)
  feat: add `flask upgrade-db` command for running db upgrade with redis lock (#5333)
  fix: wrong token usage in iteration node for streaming result (#5336)

.github/workflows/db-migration-test.yml

@@ -38,13 +38,14 @@ jobs:
       - name: Install dependencies
         run: poetry install -C api
 
-      - name: Set up Middleware
+      - name: Set up Middlewares
         uses: hoverkraft-tech/[email protected]
         with:
           compose-file: |
             docker/docker-compose.middleware.yaml
           services: |
             db
+            redis
 
       - name: Prepare configs
         run: |
@@ -54,4 +55,4 @@ jobs:
       - name: Run DB Migration
         run: |
           cd api
-          poetry run python -m flask db upgrade
+          poetry run python -m flask upgrade-db

api/commands.py

@@ -1,5 +1,6 @@
 import base64
 import json
+import logging
 import secrets
 from typing import Optional
 
@@ -12,6 +13,7 @@
 from core.rag.datasource.vdb.vector_type import VectorType
 from core.rag.models.document import Document
 from extensions.ext_database import db
+from extensions.ext_redis import redis_client
 from libs.helper import email as email_validate
 from libs.password import hash_password, password_pattern, valid_password
 from libs.rsa import generate_key_pair
@@ -553,6 +555,28 @@ def create_tenant(email: str, language: Optional[str] = None):
                            'Account: {}\nPassword: {}'.format(email, new_password), fg='green'))
 
 
+@click.command('upgrade-db', help='upgrade the database')
+def upgrade_db():
+    click.echo('Preparing database migration...')
+    lock = redis_client.lock(name='db_upgrade_lock', timeout=60)
+    if lock.acquire(blocking=False):
+        try:
+            click.echo(click.style('Start database migration.', fg='green'))
+
+            # run db migration
+            import flask_migrate
+            flask_migrate.upgrade()
+
+            click.echo(click.style('Database migration successful!', fg='green'))
+
+        except Exception as e:
+            logging.exception(f'Database migration failed, error: {e}')
+        finally:
+            lock.release()
+    else:
+        click.echo('Database migration skipped')
+
+
 def register_commands(app):
     app.cli.add_command(reset_password)
     app.cli.add_command(reset_email)
@@ -561,4 +585,4 @@ def register_commands(app):
     app.cli.add_command(convert_to_agent_apps)
     app.cli.add_command(add_qdrant_doc_id_index)
     app.cli.add_command(create_tenant)
-
+    app.cli.add_command(upgrade_db)

api/core/app/entities/task_entities.py

@@ -367,7 +367,7 @@ class Data(BaseModel):
 
 class IterationNodeCompletedStreamResponse(StreamResponse):
     """
-    NodeStartStreamResponse entity
+    NodeCompletedStreamResponse entity
     """
     class Data(BaseModel):
         """
@@ -385,6 +385,7 @@ class Data(BaseModel):
         error: Optional[str] = None
         elapsed_time: float
         total_tokens: int
+        execution_metadata: Optional[dict] = None
         finished_at: int
         steps: int
 
@@ -545,4 +546,4 @@ class Data(BaseModel):
         total_tokens: int = 0
         node_data: BaseNodeData
 
-    current_iterations: dict[str, Data] = None
+    current_iterations: dict[str, Data] = None

api/core/app/task_pipeline/workflow_iteration_cycle_manage.py

@@ -95,6 +95,9 @@ def _handle_iteration_to_stream_response(self, task_id: str, event: QueueIterati
                     error=None,
                     elapsed_time=time.perf_counter() - current_iteration.started_at,
                     total_tokens=current_iteration.total_tokens,
+                    execution_metadata={
+                        'total_tokens': current_iteration.total_tokens,
+                    },
                     finished_at=int(time.time()),
                     steps=current_iteration.current_index
                 )
@@ -276,7 +279,10 @@ def _handle_iteration_exception(self, task_id: str, error: str) -> Generator[Ite
                     error=error,
                     elapsed_time=time.perf_counter() - current_iteration.started_at,
                     total_tokens=current_iteration.total_tokens,
+                    execution_metadata={
+                        'total_tokens': current_iteration.total_tokens,
+                    },
                     finished_at=int(time.time()),
                     steps=current_iteration.current_index
                 )
-            )
+            )

api/docker/entrypoint.sh

@@ -4,7 +4,7 @@ set -e
 
 if [[ "${MIGRATION_ENABLED}" == "true" ]]; then
   echo "Running migrations"
-  flask db upgrade
+  flask upgrade-db
 fi
 
 if [[ "${MODE}" == "worker" ]]; then

api/migrations/README

@@ -1,2 +1 @@
 Single-database configuration for Flask.
-

api/pyproject.toml

@@ -1,6 +1,10 @@
 [project]
 requires-python = ">=3.10"
 
+[build-system]
+requires = ["poetry-core"]
+build-backend = "poetry.core.masonry.api"
+
 [tool.ruff]
 exclude = [
 ]
@@ -78,17 +82,10 @@ MOCK_SWITCH = "true"
 CODE_MAX_STRING_LENGTH = "80000"
 CODE_EXECUTION_ENDPOINT="http://127.0.0.1:8194"
 CODE_EXECUTION_API_KEY="dify-sandbox"
-
 FIRECRAWL_API_KEY = "fc-"
 
-
-
 [tool.poetry]
 name = "dify-api"
-version = "0.6.11"
-description = ""
-authors = ["Dify <[email protected]>"]
-readme = "README.md"
 package-mode = false
 
 [tool.poetry.dependencies]