perf: Chat origin

This commit is contained in:
fit2cloud-chenyw
2025-07-22 22:28:59 +08:00
parent 0f6a772e80
commit aeee30e5a9
5 changed files with 39 additions and 6 deletions
@@ -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 ###
+1 -1
View File
@@ -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,
+4 -3
View File
@@ -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")
+3 -1
View File
@@ -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):
+1 -1
View File
@@ -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}