diff --git a/backend/alembic/versions/033_chat_origin_ddl.py b/backend/alembic/versions/033_chat_origin_ddl.py new file mode 100644 index 00000000..911835a4 --- /dev/null +++ b/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 ### diff --git a/backend/apps/chat/api/chat.py b/backend/apps/chat/api/chat.py index 9977eae7..68e5cb21 100644 --- a/backend/apps/chat/api/chat.py +++ b/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, diff --git a/backend/apps/chat/curd/chat.py b/backend/apps/chat/curd/chat.py index 220ad6d0..dd22de57 100644 --- a/backend/apps/chat/curd/chat.py +++ b/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") diff --git a/backend/apps/chat/models/chat_model.py b/backend/apps/chat/models/chat_model.py index db72e80d..60fb777d 100644 --- a/backend/apps/chat/models/chat_model.py +++ b/backend/apps/chat/models/chat_model.py @@ -2,7 +2,7 @@ from datetime import datetime 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): diff --git a/backend/apps/mcp/mcp.py b/backend/apps/mcp/mcp.py index d5b2dd94..806f1949 100644 --- a/backend/apps/mcp/mcp.py +++ b/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}