perf: Chat origin

Author: fit2cloud-chenyw

backend/alembic/versions/033_chat_origin_ddl.py

@@ -0,0 +1,30 @@
+"""033_chat_origin_ddl
+
+Revision ID: 3cb5d6a54f2e
+Revises: 6549e47f9adc
+Create Date: 2025-07-22 22:00:48.599729
+
+"""
+from alembic import op
+import sqlalchemy as sa
+import sqlmodel.sql.sqltypes
+from sqlalchemy.dialects import postgresql
+
+# revision identifiers, used by Alembic.
+revision = '3cb5d6a54f2e'
+down_revision = '6549e47f9adc'
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.add_column('chat', sa.Column('origin', sa.Integer(), nullable=True, default=0))  # 0: default, 1: mcp, 2: assistant
+    op.execute('UPDATE chat SET origin = 0 WHERE origin IS NULL')
+    # ### end Alembic commands ###
+
+
+def downgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_column('chat', 'origin')
+    # ### end Alembic commands ###

backend/apps/chat/api/chat.py

@@ -85,7 +85,7 @@ async def start_chat(session: SessionDep, current_user: CurrentUser, create_chat
 @router.post("/assistant/start")
 async def start_chat(session: SessionDep, current_user: CurrentUser):
     try:
-        return create_chat(session, current_user, CreateChat(), False)
+        return create_chat(session, current_user, CreateChat(origin=2), False)
     except Exception as e:
         raise HTTPException(
             status_code=500,

backend/apps/chat/curd/chat.py

@@ -21,7 +21,7 @@ def list_chats(session: SessionDep, current_user: CurrentUser) -> List[Chat]:
 
 
 def rename_chat(session: SessionDep, rename_object: RenameChat) -> str:
-    chat = session.query(Chat).filter(Chat.id == rename_object.id).first()
+    chat = session.get(Chat, rename_object.id)
     if not chat:
         raise Exception(f"Chat with id {rename_object.id} not found")
 
@@ -163,11 +163,12 @@ def create_chat(session: SessionDep, current_user: CurrentUser, create_chat_obj:
     chat = Chat(create_time=datetime.datetime.now(),
                 create_by=current_user.id,
                 oid=current_user.oid if current_user.oid is not None else 1,
-                brief=create_chat_obj.question.strip()[:20])
+                brief=create_chat_obj.question.strip()[:20],
+                origin=create_chat_obj.origin if create_chat_obj.origin is not None else 0)
     ds: CoreDatasource | None = None
     if create_chat_obj.datasource:
         chat.datasource = create_chat_obj.datasource
-        ds = session.query(CoreDatasource).filter(CoreDatasource.id == create_chat_obj.datasource).first()
+        ds = session.get(CoreDatasource, create_chat_obj.datasource)
 
         if not ds:
             raise Exception(f"Datasource with id {create_chat_obj.datasource} not found")

backend/apps/chat/models/chat_model.py

@@ -2,7 +2,7 @@
 from typing import List, Optional
 
 from pydantic import BaseModel
-from sqlalchemy import Column, Text, BigInteger, DateTime, Identity, Boolean
+from sqlalchemy import Column, Integer, Text, BigInteger, DateTime, Identity, Boolean
 from sqlmodel import SQLModel, Field
 
 from apps.template.generate_analysis.generator import get_analysis_template
@@ -24,6 +24,7 @@ class Chat(SQLModel, table=True):
     chat_type: str = Field(max_length=20, default="chat")  # chat, datasource
     datasource: int = Field(sa_column=Column(BigInteger, nullable=True))
     engine_type: str = Field(max_length=64)
+    origin: Optional[int] = Field(sa_column=Column(Integer, nullable=False, default=0))  # 0: default, 1: mcp, 2: assistant
 
 
 class ChatRecord(SQLModel, table=True):
@@ -72,6 +73,7 @@ class CreateChat(BaseModel):
     id: int = None
     question: str = None
     datasource: int = None
+    origin: Optional[int] = 0
 
 
 class RenameChat(BaseModel):

backend/apps/mcp/mcp.py

@@ -71,7 +71,7 @@ async def mcp_start(session: SessionDep, chat: ChatStart):
     t = Token(access_token=create_access_token(
         user_dict, expires_delta=access_token_expires
     ))
-    c = create_chat(session, user, CreateChat(), False)
+    c = create_chat(session, user, CreateChat(origin=1), False)
     return {"access_token": t.access_token, "chat_id": c.id}