mirror of
https://gh-proxy.org/https://github.com/klinsmaya/SQLBot.git
synced 2026-11-08 17:26:55 +08:00
perf: Chat origin
This commit is contained in:
@@ -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 ###
|
||||
@@ -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,
|
||||
|
||||
@@ -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")
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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}
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user