diff --git a/backend/apps/chat/api/chat.py b/backend/apps/chat/api/chat.py index 02a76c06..4f4a7087 100644 --- a/backend/apps/chat/api/chat.py +++ b/backend/apps/chat/api/chat.py @@ -4,10 +4,11 @@ from sqlmodel import select from apps.chat.curd.chat import list_chats, get_chat_with_records, create_chat, save_question, save_answer, rename_chat, \ delete_chat -from apps.chat.models.chat_model import CreateChat, ChatRecord, RenameChat +from apps.chat.models.chat_model import CreateChat, ChatRecord, RenameChat, Chat from apps.chat.schemas.chat_base_schema import LLMConfig from apps.chat.schemas.chat_schema import ChatQuestion from apps.chat.schemas.llm import AgentService +from apps.datasource.crud.datasource import get_table_obj_by_ds from apps.datasource.models.datasource import CoreDatasource from apps.system.models.system_model import AiModelDetail from common.core.deps import SessionDep, CurrentUser @@ -80,26 +81,18 @@ async def stream_sql(session: SessionDep, current_user: CurrentUser, request_que """ question = request_question.question - # Get available AI model - aimodel = session.exec(select(AiModelDetail).where( - AiModelDetail.status == True, - AiModelDetail.api_key.is_not(None) - )).first() - - # Get available datasource - ds = session.exec(select(CoreDatasource).where( - CoreDatasource.status == 'Success' - )).first() - - if not aimodel: + chat = session.query(Chat).filter(Chat.id == request_question.chat_id).first() + if not chat: raise HTTPException( status_code=400, - detail="No available AI model configuration found" + detail=f"Chat with id {request_question.chart_id} not found" ) + # Get available datasource + ds = session.query(CoreDatasource).filter(CoreDatasource.id == chat.datasource).first() if not ds: raise HTTPException( - status_code=400, + status_code=500, detail="No available datasource configuration found" ) @@ -112,6 +105,17 @@ async def stream_sql(session: SessionDep, current_user: CurrentUser, request_que detail=str(e1) ) + # Get available AI model + aimodel = session.exec(select(AiModelDetail).where( + AiModelDetail.status == True, + AiModelDetail.api_key.is_not(None) + )).first() + if not aimodel: + raise HTTPException( + status_code=400, + detail="No available AI model configuration found" + ) + # Use Tongyi Qianwen tongyi_config = LLMConfig( model_type="openai", @@ -136,10 +140,39 @@ async def stream_sql(session: SessionDep, current_user: CurrentUser, request_que """ result = llm_service.generate_sql(question) return result """ + # get schema + schema_str = "" + table_objs = get_table_obj_by_ds(session=session, ds=ds) + db_name = table_objs[0].schema + schema_str += f"【DB_ID】 {db_name}\n【Schema】\n" + for obj in table_objs: + schema_str += f"# Table: {db_name}.{obj.table.table_name}" + table_comment = '' + if obj.table.custom_comment: + table_comment = obj.table.custom_comment.strip() + if table_comment == '': + schema_str += '\n[\n' + else: + schema_str += f", {table_comment}\n[\n" + + field_list = [] + for field in obj.fields: + field_comment = '' + if field.custom_comment: + field_comment = field.custom_comment.strip() + if field_comment == '': + field_list.append(f"({field.field_name}:{field.field_type})") + else: + field_list.append(f"({field.field_name}:{field.field_type}, {field_comment})") + schema_str += ",\n".join(field_list) + schema_str += '\n]\n' + + print(schema_str) + async def event_generator(): all_text = '' try: - async for chunk in llm_service.async_generate(question): + async for chunk in llm_service.async_generate(question, schema_str): data = json.loads(chunk.replace('data: ', '')) if data['type'] in ['final', 'tool_result']: diff --git a/backend/apps/chat/schemas/llm.py b/backend/apps/chat/schemas/llm.py index f9815704..fb75926e 100644 --- a/backend/apps/chat/schemas/llm.py +++ b/backend/apps/chat/schemas/llm.py @@ -101,16 +101,17 @@ class AgentService: schema = self.db.get_table_info() return chain.invoke({"schema": schema, "question": question}) - async def async_generate(self, question: str) -> AsyncGenerator[str, None]: + async def async_generate(self, question: str, schema: str) -> AsyncGenerator[str, None]: chain = self.prompt | self.agent_executor # schema = self.db.get_table_info() - schema_engine = SchemaEngine(engine=self.db._engine) - mschema = schema_engine.mschema - mschema_str = mschema.to_mschema() + # schema_engine = SchemaEngine(engine=self.db._engine) + # mschema = schema_engine.mschema + # mschema_str = mschema.to_mschema() - async for chunk in chain.astream({"schema": mschema_str, "question": question}): + # async for chunk in chain.astream({"schema": mschema_str, "question": question}): + async for chunk in chain.astream({"schema": schema, "question": question}): if not isinstance(chunk, dict): continue diff --git a/backend/apps/datasource/crud/datasource.py b/backend/apps/datasource/crud/datasource.py index c15eac4c..db2fc95a 100644 --- a/backend/apps/datasource/crud/datasource.py +++ b/backend/apps/datasource/crud/datasource.py @@ -211,10 +211,12 @@ def preview(session: SessionDep, id: int, data: TableObj): return exec_sql(ds, sql) -def get_tableobj_by_ds(session: SessionDep, id: int) -> List[TableAndFields]: - list: List = [] - tables = session.query(CoreTable).filter(CoreTable.ds_id == id).all() +def get_table_obj_by_ds(session: SessionDep, ds: CoreDatasource) -> List[TableAndFields]: + _list: List = [] + tables = session.query(CoreTable).filter(CoreTable.ds_id == ds.id).all() + conf = DatasourceConf(**json.loads(aes_decrypt(ds.configuration))) if ds.type != "excel" else get_engine_config() + schema = conf.dbSchema if conf.dbSchema is not None and conf.dbSchema != "" else conf.database for table in tables: fields = session.query(CoreField).filter(and_(CoreField.table_id == table.id, CoreField.checked == True)).all() - list.append(TableAndFields(table=table, fields=fields)) - return list + _list.append(TableAndFields(schema=schema, table=table, fields=fields)) + return _list diff --git a/backend/apps/datasource/models/datasource.py b/backend/apps/datasource/models/datasource.py index 41667d6c..f41ca5ce 100644 --- a/backend/apps/datasource/models/datasource.py +++ b/backend/apps/datasource/models/datasource.py @@ -111,9 +111,11 @@ class ColumnSchema: class TableAndFields: - def __init__(self, table, fields): + def __init__(self, schema, table, fields): + self.schema = schema self.table = table self.fields = fields + schema: str table: CoreTable fields: List[CoreField]