from datetime import datetime from typing import List, Optional from fastapi import Body from pydantic import BaseModel from sqlalchemy import Column, Integer, Text, BigInteger, DateTime, Identity, Boolean from sqlmodel import SQLModel, Field from apps.template.filter.generator import get_permissions_template from apps.template.generate_analysis.generator import get_analysis_template from apps.template.generate_chart.generator import get_chart_template from apps.template.generate_guess_question.generator import get_guess_question_template from apps.template.generate_predict.generator import get_predict_template from apps.template.generate_sql.generator import get_sql_template from apps.template.select_datasource.generator import get_datasource_template class Chat(SQLModel, table=True): __tablename__ = "chat" id: Optional[int] = Field(sa_column=Column(BigInteger, Identity(always=True), primary_key=True)) oid: Optional[int] = Field(sa_column=Column(BigInteger, nullable=True, default=1)) create_time: datetime = Field(sa_column=Column(DateTime(timezone=False), nullable=True)) create_by: int = Field(sa_column=Column(BigInteger, nullable=True)) brief: str = Field(max_length=64, nullable=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): __tablename__ = "chat_record" id: Optional[int] = Field(sa_column=Column(BigInteger, Identity(always=True), primary_key=True)) chat_id: int = Field(sa_column=Column(BigInteger, nullable=False)) ai_modal_id: Optional[int] = Field(sa_column=Column(BigInteger)) first_chat: bool = Field(sa_column=Column(Boolean, nullable=True, default=False)) create_time: datetime = Field(sa_column=Column(DateTime(timezone=False), nullable=True)) finish_time: datetime = Field(sa_column=Column(DateTime(timezone=False), nullable=True)) create_by: int = Field(sa_column=Column(BigInteger, nullable=True)) datasource: int = Field(sa_column=Column(BigInteger, nullable=True)) engine_type: str = Field(max_length=64, nullable=True) question: str = Field(sa_column=Column(Text, nullable=True)) sql_answer: str = Field(sa_column=Column(Text, nullable=True)) sql: str = Field(sa_column=Column(Text, nullable=True)) sql_exec_result: str = Field(sa_column=Column(Text, nullable=True)) data: str = Field(sa_column=Column(Text, nullable=True)) chart_answer: str = Field(sa_column=Column(Text, nullable=True)) chart: str = Field(sa_column=Column(Text, nullable=True)) analysis: str = Field(sa_column=Column(Text, nullable=True)) predict: str = Field(sa_column=Column(Text, nullable=True)) predict_data: str = Field(sa_column=Column(Text, nullable=True)) recommended_question_answer: str = Field(sa_column=Column(Text, nullable=True)) recommended_question: str = Field(sa_column=Column(Text, nullable=True)) datasource_select_answer: str = Field(sa_column=Column(Text, nullable=True)) full_sql_message: str = Field(sa_column=Column(Text, nullable=True)) token_sql: str = Field(max_length=256, nullable=True) full_chart_message: str = Field(sa_column=Column(Text, nullable=True)) token_chart: str = Field(max_length=256, nullable=True) full_analysis_message: str = Field(sa_column=Column(Text, nullable=True)) token_analysis: str = Field(max_length=256, nullable=True) full_predict_message: str = Field(sa_column=Column(Text, nullable=True)) token_predict: str = Field(max_length=256, nullable=True) full_recommended_question_message: str = Field(sa_column=Column(Text, nullable=True)) token_recommended_question: str = Field(max_length=256, nullable=True) full_select_datasource_message: str = Field(sa_column=Column(Text, nullable=True)) token_select_datasource_question: str = Field(max_length=256, nullable=True) finish: bool = Field(sa_column=Column(Boolean, nullable=True, default=False)) error: str = Field(sa_column=Column(Text, nullable=True)) analysis_record_id: int = Field(sa_column=Column(BigInteger, nullable=True)) predict_record_id: int = Field(sa_column=Column(BigInteger, nullable=True)) class CreateChat(BaseModel): id: int = None question: str = None datasource: int = None origin: Optional[int] = 0 class RenameChat(BaseModel): id: int = None brief: str = '' class ChatInfo(BaseModel): id: Optional[int] = None create_time: datetime = None create_by: int = None brief: str = '' chat_type: str = "chat" datasource: Optional[int] = None engine_type: str = '' ds_type: str = '' datasource_name: str = '' datasource_exists: bool = True records: List[ChatRecord | dict] = [] class AiModelQuestion(BaseModel): ai_modal_id: int = None engine: str = "" db_schema: str = "" sql: str = "" rule: str = "" fields: str = "" data: str = "" lang: str = "简体中文" filter: str = [] def sql_sys_question(self): return get_sql_template()['system'].format(engine=self.engine, schema=self.db_schema, question=self.question, lang=self.lang) def sql_user_question(self): return get_sql_template()['user'].format(engine=self.engine, schema=self.db_schema, question=self.question, rule=self.rule) def chart_sys_question(self): return get_chart_template()['system'].format(sql=self.sql, question=self.question, lang=self.lang) def chart_user_question(self): return get_chart_template()['user'].format(sql=self.sql, question=self.question, rule=self.rule) def analysis_sys_question(self): return get_analysis_template()['system'].format(lang=self.lang) def analysis_user_question(self): return get_analysis_template()['user'].format(fields=self.fields, data=self.data) def predict_sys_question(self): return get_predict_template()['system'].format(lang=self.lang) def predict_user_question(self): return get_predict_template()['user'].format(fields=self.fields, data=self.data) def datasource_sys_question(self): return get_datasource_template()['system'].format(lang=self.lang) def datasource_user_question(self, datasource_list: str = "[]"): return get_datasource_template()['user'].format(question=self.question, data=datasource_list) def guess_sys_question(self): return get_guess_question_template()['system'].format(lang=self.lang) def guess_user_question(self, old_questions: str = "[]"): return get_guess_question_template()['user'].format(question=self.question, schema=self.db_schema, old_questions=old_questions) def filter_sys_question(self): return get_permissions_template()['system'].format(lang=self.lang, engine=self.engine) def filter_user_question(self): return get_permissions_template()['user'].format(sql=self.sql, filter=self.filter) class ChatQuestion(AiModelQuestion): question: str = Body(description='用户提问') chat_id: int = Body(description='会话ID') class ChatMcp(ChatQuestion): token: str = Body(description='token') class ChatStart(BaseModel): username: str = Body(description='用户名') password: str = Body(description='密码')