diff --git a/backend/apps/chat/api/chat.py b/backend/apps/chat/api/chat.py index 0b26ea14..c1bbae9c 100644 --- a/backend/apps/chat/api/chat.py +++ b/backend/apps/chat/api/chat.py @@ -1,12 +1,12 @@ -import asyncio import json +import traceback from typing import List from fastapi import APIRouter, HTTPException from fastapi.responses import StreamingResponse from sqlmodel import select -from apps.chat.curd.chat import list_chats, get_chat_with_records, create_chat, save_question, save_answer, rename_chat, \ +from apps.chat.curd.chat import list_chats, get_chat_with_records, create_chat, rename_chat, \ delete_chat, list_records from apps.chat.models.chat_model import CreateChat, ChatRecord, RenameChat, Chat, ChatQuestion from apps.chat.task.llm import LLMService @@ -79,7 +79,6 @@ async def stream_sql(session: SessionDep, current_user: CurrentUser, request_que Returns: Streaming response with analysis results """ - question = request_question.question chat = session.query(Chat).filter(Chat.id == request_question.chat_id).first() if not chat: @@ -113,46 +112,51 @@ async def stream_sql(session: SessionDep, current_user: CurrentUser, request_que chart_id=request_question.chat_id) # get schema request_question.db_schema = get_table_schema(session=session, ds=ds) - llm_service = LLMService(request_question, history_records, ds, aimodel) + llm_service = LLMService(request_question, history_records, CoreDatasource(**ds.model_dump()), aimodel) llm_service.init_record(session=session, current_user=current_user) def run_task(): - sql_res = llm_service.generate_sql(session=session) - for chunk in sql_res: - yield json.dumps({'content': chunk, 'type': 'sql'}) + '\n\n' - yield json.dumps({'type': 'info', 'msg': 'sql generated'}) + '\n\n' + try: + # generate sql + sql_res = llm_service.generate_sql(session=session) + full_sql_text = '' + for chunk in sql_res: + full_sql_text += chunk + yield json.dumps({'content': chunk, 'type': 'sql-result'}, ensure_ascii=False) + '\n\n' + yield json.dumps({'type': 'info', 'msg': 'sql generated'}) + '\n\n' - # async def event_generator(): - # all_text = '' - # try: - # async for chunk in llm_service.async_generate(question, request_question.db_schema): - # data = json.loads(chunk.replace('data: ', '')) - # - # if data['type'] in ['final', 'tool_result']: - # content = data['content'] - # print('-- ' + content) - # all_text += content - # for char in content: - # yield f"data: {json.dumps({'type': 'char', 'content': char})}\n\n" - # await asyncio.sleep(0.05) - # - # if 'html' in data: - # yield f"data: {json.dumps({'type': 'html', 'content': data['html']})}\n\n" - # else: - # yield chunk - # - # except Exception as e: - # all_text = 'Exception:' + str(e) - # yield f"data: {json.dumps({'type': 'error', 'content': str(e)})}\n\n" - # - # try: - # save_answer(session=session, id=record.id, answer=all_text) - # except Exception as e: - # raise HTTPException( - # status_code=500, - # detail=str(e) - # ) + # filter sql + print(full_sql_text) + sql = llm_service.check_save_sql(session=session, res=full_sql_text) + print(sql) + yield json.dumps({'content': sql, 'type': 'sql'}) + '\n\n' + + # execute sql + result = llm_service.execute_sql(sql=sql) + llm_service.save_sql_data(session=session, data_obj=result) + yield json.dumps({'content': result, 'type': 'sql-data'}, ensure_ascii=False) + '\n\n' + + # generate chart + chart_res = llm_service.generate_chart(session=session) + full_chart_text = '' + for chunk in chart_res: + full_chart_text += chunk + yield json.dumps({'content': chunk, 'type': 'chart-result'}, ensure_ascii=False) + '\n\n' + yield json.dumps({'type': 'info', 'msg': 'chart generated'}) + '\n\n' + + # filter chart + print(full_chart_text) + chart = llm_service.check_save_chart(session=session, res=full_chart_text) + print(chart) + yield json.dumps({'content': chart, 'type': 'chart'}, ensure_ascii=False) + '\n\n' + + llm_service.finish(session=session) + yield json.dumps({'type': 'finish'}) + + except Exception as e: + traceback.print_exc() + llm_service.save_error(session=session, message=str(e)) + yield json.dumps({'content': str(e), 'type': 'error'}, ensure_ascii=False) + '\n\n' - # return EventSourceResponse(event_generator(), headers={"Content-Type": "text/event-stream"}) return StreamingResponse(run_task(), media_type="text/event-stream") diff --git a/backend/apps/chat/curd/chat.py b/backend/apps/chat/curd/chat.py index 5ba0437e..774e4967 100644 --- a/backend/apps/chat/curd/chat.py +++ b/backend/apps/chat/curd/chat.py @@ -150,11 +150,35 @@ def save_full_sql_message_and_answer(session: SessionDep, record_id: int, answer return result -def save_full_chart_message(session: SessionDep, id: int, full_message: str) -> ChatRecord: - if not id: +def save_sql(session: SessionDep, record_id: int, sql: str) -> ChatRecord: + if not record_id: raise Exception("Record id cannot be None") - record = session.query(ChatRecord).filter(ChatRecord.id == id).first() + record = session.query(ChatRecord).filter(ChatRecord.id == record_id).first() + record.sql = sql + + result = ChatRecord(**record.model_dump()) + + session.add(record) + session.flush() + session.refresh(record) + + session.commit() + + return result + + +def save_full_chart_message(session: SessionDep, record_id: int, full_message: str) -> ChatRecord: + return save_full_chart_message_and_answer(session=session, record_id=record_id, full_message=full_message, + answer='') + + +def save_full_chart_message_and_answer(session: SessionDep, record_id: int, answer: str, + full_message: str) -> ChatRecord: + if not record_id: + raise Exception("Record id cannot be None") + record = session.query(ChatRecord).filter(ChatRecord.id == record_id).first() record.full_chart_message = full_message + record.chart_answer = answer result = ChatRecord(**record.model_dump()) @@ -167,12 +191,11 @@ def save_full_chart_message(session: SessionDep, id: int, full_message: str) -> return result -def save_answer(session: SessionDep, id: int, answer: str) -> ChatRecord: - if not id: +def save_chart(session: SessionDep, record_id: int, chart: str) -> ChatRecord: + if not record_id: raise Exception("Record id cannot be None") - - record = session.query(ChatRecord).filter(ChatRecord.id == id).first() - record.answer = answer + record = session.query(ChatRecord).filter(ChatRecord.id == record_id).first() + record.chart = chart result = ChatRecord(**record.model_dump()) @@ -183,3 +206,55 @@ def save_answer(session: SessionDep, id: int, answer: str) -> ChatRecord: session.commit() return result + + +def save_error_message(session: SessionDep, record_id: int, message: str) -> ChatRecord: + if not record_id: + raise Exception("Record id cannot be None") + record = session.query(ChatRecord).filter(ChatRecord.id == record_id).first() + record.error = message + record.finish = True + + result = ChatRecord(**record.model_dump()) + + session.add(record) + session.flush() + session.refresh(record) + + session.commit() + + return result + + +def save_sql_exec_data(session: SessionDep, record_id: int, data: str) -> ChatRecord: + if not record_id: + raise Exception("Record id cannot be None") + record = session.query(ChatRecord).filter(ChatRecord.id == record_id).first() + record.data = data + + result = ChatRecord(**record.model_dump()) + + session.add(record) + session.flush() + session.refresh(record) + + session.commit() + + return result + +def finish_record(session: SessionDep, record_id: int) -> ChatRecord: + if not record_id: + raise Exception("Record id cannot be None") + record = session.query(ChatRecord).filter(ChatRecord.id == record_id).first() + record.finish = True + + result = ChatRecord(**record.model_dump()) + + session.add(record) + session.flush() + session.refresh(record) + + session.commit() + + return result + diff --git a/backend/apps/chat/schemas/__init__.py b/backend/apps/chat/schemas/__init__.py deleted file mode 100644 index e69de29b..00000000 diff --git a/backend/apps/chat/schemas/m_schema.py b/backend/apps/chat/schemas/m_schema.py deleted file mode 100644 index 058e2f94..00000000 --- a/backend/apps/chat/schemas/m_schema.py +++ /dev/null @@ -1,195 +0,0 @@ -from apps.chat.schemas.utils import examples_to_str, read_json, write_json -from typing import Any, Dict, Iterable, List, Optional, Tuple, Union - - -class MSchema: - def __init__(self, db_id: str = 'Anonymous', schema: Optional[str] = None): - self.db_id = db_id - self.schema = schema - self.tables = {} - self.foreign_keys = [] - - def add_table(self, name, fields={}, comment=None): - self.tables[name] = {"fields": fields.copy(), 'examples': [], 'comment': comment} - - def add_field(self, table_name: str, field_name: str, field_type: str = "", - primary_key: bool = False, nullable: bool = True, default: Any = None, - autoincrement: bool = False, comment: str = "", examples: list = [], **kwargs): - self.tables[table_name]["fields"][field_name] = { - "type": field_type, - "primary_key": primary_key, - "nullable": nullable, - "default": default if default is None else f'{default}', - "autoincrement": autoincrement, - "comment": comment, - "examples": examples.copy(), - **kwargs} - - def add_foreign_key(self, table_name, field_name, ref_schema, ref_table_name, ref_field_name): - self.foreign_keys.append([table_name, field_name, ref_schema, ref_table_name, ref_field_name]) - - def get_field_type(self, field_type, simple_mode=True)->str: - if not simple_mode: - return field_type - else: - return field_type.split("(")[0] - - def has_table(self, table_name: str) -> bool: - if table_name in self.tables.keys(): - return True - else: - return False - - def has_column(self, table_name: str, field_name: str) -> bool: - if self.has_table(table_name): - if field_name in self.tables[table_name]["fields"].keys(): - return True - else: - return False - else: - return False - - def get_field_info(self, table_name: str, field_name: str) -> Dict: - try: - return self.tables[table_name]['fields'][field_name] - except: - return {} - - def single_table_mschema(self, table_name: str, selected_columns: List = None, - example_num=3, show_type_detail=False) -> str: - table_info = self.tables.get(table_name, {}) - output = [] - table_comment = table_info.get('comment', '') - if table_comment is not None and table_comment != 'None' and len(table_comment) > 0: - if self.schema is not None and len(self.schema) > 0: - output.append(f"# Table: {self.schema}.{table_name}, {table_comment}") - else: - output.append(f"# Table: {table_name}, {table_comment}") - else: - if self.schema is not None and len(self.schema) > 0: - output.append(f"# Table: {self.schema}.{table_name}") - else: - output.append(f"# Table: {table_name}") - - field_lines = [] - for field_name, field_info in table_info['fields'].items(): - if selected_columns is not None and field_name.lower() not in selected_columns: - continue - - raw_type = self.get_field_type(field_info['type'], not show_type_detail) - field_line = f"({field_name}:{raw_type.upper()}" - if field_info['comment'] != '': - field_line += f", {field_info['comment'].strip()}" - else: - pass - - is_primary_key = field_info.get('primary_key', False) - if is_primary_key: - field_line += f", Primary Key" - - if len(field_info.get('examples', [])) > 0 and example_num > 0: - examples = field_info['examples'] - examples = [s for s in examples if s is not None] - examples = examples_to_str(examples) - if len(examples) > example_num: - examples = examples[:example_num] - - if raw_type in ['DATE', 'TIME', 'DATETIME', 'TIMESTAMP']: - examples = [examples[0]] - elif len(examples) > 0 and max([len(s) for s in examples]) > 20: - if max([len(s) for s in examples]) > 50: - examples = [] - else: - examples = [examples[0]] - else: - pass - if len(examples) > 0: - example_str = ', '.join([str(example) for example in examples]) - field_line += f", Examples: [{example_str}]" - else: - pass - else: - field_line += "" - field_line += ")" - - field_lines.append(field_line) - output.append('[') - output.append(',\n'.join(field_lines)) - output.append(']') - - return '\n'.join(output) - - def to_mschema(self, selected_tables: List = None, selected_columns: List = None, - example_num=3, show_type_detail=False) -> str: - """ - Convert to a MSchema string. - - Args: - selected_tables: None means select all tables - selected_columns: None means select all columns, format: ['table_name.column_name'] - example_num: Number of examples to show for each field - show_type_detail: Whether to show detailed type information - - Returns: - Formatted schema string - """ - output = [] - - output.append(f"【DB_ID】 {self.db_id}") - output.append(f"【Schema】") - - # Process selected tables and columns - if selected_tables is not None: - selected_tables = [s.lower() for s in selected_tables] - if selected_columns is not None: - selected_columns = [s.lower() for s in selected_columns] - selected_tables = [s.split('.')[0].lower() for s in selected_columns] - - # Process each table - for table_name, table_info in self.tables.items(): - if selected_tables is None or table_name.lower() in selected_tables: - cur_table_type = table_info.get('type', 'table') - column_names = list(table_info['fields'].keys()) - if selected_columns is not None: - cur_selected_columns = [c.lower() for c in column_names if f"{table_name}.{c}".lower() in selected_columns] - else: - cur_selected_columns = selected_columns - output.append(self.single_table_mschema(table_name, cur_selected_columns, example_num, show_type_detail)) - - # Add foreign key information, skip when table_type is view - if self.foreign_keys: - output.append("【Foreign keys】") - for fk in self.foreign_keys: - ref_schema = fk[2] - table1, column1, _, table2, column2 = fk - if selected_tables is None or \ - (table1.lower() in selected_tables and table2.lower() in selected_tables): - output.append(f"{fk[0]}.{fk[1]}={fk[3]}.{fk[4]}") - - return '\n'.join(output) - - def dump(self): - """ - Dump schema information to dictionary - - Returns: - Dictionary containing schema information - """ - schema_dict = { - "db_id": self.db_id, - "schema": self.schema, - "tables": self.tables, - "foreign_keys": self.foreign_keys - } - return schema_dict - - def save(self, file_path: str): - schema_dict = self.dump() - write_json(file_path, schema_dict) - - def load(self, file_path: str): - data = read_json(file_path) - self.db_id = data.get("db_id", "Anonymous") - self.schema = data.get("schema", None) - self.tables = data.get("tables", {}) - self.foreign_keys = data.get("foreign_keys", []) diff --git a/backend/apps/chat/schemas/schema_engine.py b/backend/apps/chat/schemas/schema_engine.py deleted file mode 100644 index 63910503..00000000 --- a/backend/apps/chat/schemas/schema_engine.py +++ /dev/null @@ -1,126 +0,0 @@ -import json, os -from typing import Any, Dict, Iterable, List, Optional, Tuple, Union -from sqlalchemy import create_engine, MetaData, Table, Column, String, Integer, select, text -from sqlalchemy.engine import Engine -from llama_index.core import SQLDatabase -from apps.chat.schemas.utils import read_json, write_json, save_raw_text, examples_to_str -from apps.chat.schemas.m_schema import MSchema - - -class SchemaEngine(SQLDatabase): - def __init__(self, engine: Engine, schema: Optional[str] = None, metadata: Optional[MetaData] = None, - ignore_tables: Optional[List[str]] = None, include_tables: Optional[List[str]] = None, - sample_rows_in_table_info: int = 3, indexes_in_table_info: bool = False, - custom_table_info: Optional[dict] = None, view_support: bool = False, max_string_length: int = 300, - mschema: Optional[MSchema] = None, db_name: Optional[str] = ''): - super().__init__(engine, schema, metadata, ignore_tables, include_tables, sample_rows_in_table_info, - indexes_in_table_info, custom_table_info, view_support, max_string_length) - - self._db_name = db_name - # Dictionary to store table names and their corresponding schema - self._tables_schemas: Dict[str, str] = {} - - # If a schema is specified, filter by that schema and store that value for every table. - """ if schema: - self._usable_tables = [ - table_name for table_name in self._usable_tables - if self._inspector.has_table(table_name, schema) - ] - for table_name in self._usable_tables: - self._tables_schemas[table_name] = schema - else: - all_tables = [] - # Iterate through all available schemas - for s in self.get_schema_names(): - tables = self._inspector.get_table_names(schema=s) - all_tables.extend(tables) - for table in tables: - self._tables_schemas[table] = s - self._usable_tables = all_tables """ - d_schema = self.default_schema_name() - for table_name in self._usable_tables: - self._tables_schemas[table_name] = d_schema - - self._dialect = engine.dialect.name - if mschema is not None: - self._mschema = mschema - else: - self._mschema = MSchema(db_id=db_name) - self.init_mschema() - - @property - def mschema(self) -> MSchema: - """Return M-Schema""" - return self._mschema - - def get_pk_constraint(self, table_name: str) -> Dict: - return self._inspector.get_pk_constraint(table_name, self._tables_schemas[table_name] )['constrained_columns'] - - def get_table_comment(self, table_name: str): - try: - return self._inspector.get_table_comment(table_name, self._tables_schemas[table_name])['text'] - except: # sqlite does not support comments - return '' - - def default_schema_name(self) -> Optional[str]: - return self._inspector.default_schema_name - - def get_schema_names(self) -> List[str]: - return self._inspector.get_schema_names() - - def get_foreign_keys(self, table_name: str): - return self._inspector.get_foreign_keys(table_name, self._tables_schemas[table_name]) - - def get_unique_constraints(self, table_name: str): - return self._inspector.get_unique_constraints(table_name, self._tables_schemas[table_name]) - - def fectch_distinct_values(self, table_name: str, column_name: str, max_num: int = 5): - table = Table(table_name, self.metadata_obj, autoload_with=self._engine, schema=self._tables_schemas[table_name]) - # Construct SELECT DISTINCT query - query = select(table.c[column_name]).distinct().limit(max_num) - values = [] - with self._engine.connect() as connection: - result = connection.execute(query) - distinct_values = result.fetchall() - for value in distinct_values: - if value[0] is not None and value[0] != '': - values.append(value[0]) - return values - - def init_mschema(self): - for table_name in self._usable_tables: - table_comment = self.get_table_comment(table_name) - table_comment = '' if table_comment is None else table_comment.strip() - table_with_schema = self._tables_schemas[table_name] + '.' + table_name - self._mschema.add_table(table_with_schema, fields={}, comment=table_comment) - pks = self.get_pk_constraint(table_name) - - fks = self.get_foreign_keys(table_name) - for fk in fks: - referred_schema = fk['referred_schema'] - for c, r in zip(fk['constrained_columns'], fk['referred_columns']): - self._mschema.add_foreign_key(table_with_schema, c, referred_schema, fk['referred_table'], r) - - fields = self._inspector.get_columns(table_name, schema=self._tables_schemas[table_name]) - for field in fields: - field_type = f"{field['type']!s}" - field_name = field['name'] - primary_key = field_name in pks - field_comment = field.get("comment", None) - field_comment = "" if field_comment is None else field_comment.strip() - autoincrement = field.get('autoincrement', False) - default = field.get('default', None) - if default is not None: - default = f'{default}' - - try: - examples = self.fectch_distinct_values(table_name, field_name, 5) - except: - examples = [] - examples = examples_to_str(examples) - - self._mschema.add_field( - table_with_schema, field_name, field_type=field_type, primary_key=primary_key, - nullable=field['nullable'], default=default, autoincrement=autoincrement, - comment=field_comment, examples=examples - ) \ No newline at end of file diff --git a/backend/apps/chat/schemas/utils.py b/backend/apps/chat/schemas/utils.py deleted file mode 100644 index efa6ecf3..00000000 --- a/backend/apps/chat/schemas/utils.py +++ /dev/null @@ -1,80 +0,0 @@ -import datetime -import decimal -import re -import json - - -def write_json(path, data): - with open(path, 'w', encoding='utf-8') as f: - json.dump(data, f, ensure_ascii=False, indent=2) - - -def read_json(path): - with open(path, 'r', encoding='utf-8') as f: - data = json.load(f) - return data - - -def read_text(filename)->str: - data = [] - with open(filename, 'r', encoding='utf-8') as file: - for line in file.readlines(): - line = line.strip() - data.append(line) - return data - - -def save_raw_text(filename, content): - with open(filename, 'w', encoding='utf-8') as file: - file.write(content) - - -def read_map_file(path): - data = {} - with open(path, 'r', encoding='utf-8') as f: - for line in f.readlines(): - line = line.strip().split('\t') - data[line[0]] = line[1].split('、') - data[line[0]].append(line[0]) - return data - - -def save_json(target_file,js,indent=4): - with open(target_file, 'w', encoding='utf-8') as f: - json.dump(js, f, ensure_ascii=False, indent=indent) - -def is_email(string): - pattern = r'^[\w\.-]+@[\w\.-]+\.\w+$' - match = re.match(pattern, string) - if match: - return True - else: - return False - - -def examples_to_str(examples: list) -> list[str]: - """ - from examples to a list of str - """ - values = examples - for i in range(len(values)): - if isinstance(values[i], datetime.date): - values = [values[i]] - break - elif isinstance(values[i], datetime.datetime): - values = [values[i]] - break - elif isinstance(values[i], decimal.Decimal): - values[i] = str(float(values[i])) - elif is_email(str(values[i])): - values = [] - break - elif 'http://' in str(values[i]) or 'https://' in str(values[i]): - values = [] - break - elif values[i] is not None and not isinstance(values[i], str): - pass - elif values[i] is not None and '.com' in values[i]: - pass - - return [str(v) for v in values if v is not None and len(str(v)) > 0] \ No newline at end of file diff --git a/backend/apps/chat/task/llm.py b/backend/apps/chat/task/llm.py index da62eab9..b40478a5 100644 --- a/backend/apps/chat/task/llm.py +++ b/backend/apps/chat/task/llm.py @@ -1,15 +1,18 @@ import asyncio import json import logging +import re import warnings -from typing import Any, List, Union, AsyncGenerator +from typing import Any, List, Union, AsyncGenerator, Dict from langchain_community.utilities import SQLDatabase from langchain_core.language_models import BaseLLM from langchain_core.messages import BaseMessage, SystemMessage, HumanMessage, AIMessage, AIMessageChunk from apps.ai_model.model_factory import LLMConfig, LLMFactory, get_llm_config -from apps.chat.curd.chat import save_question, save_full_sql_message, save_full_sql_message_and_answer +from apps.chat.curd.chat import save_question, save_full_sql_message, save_full_sql_message_and_answer, save_sql, \ + save_error_message, save_sql_exec_data, save_full_chart_message, save_full_chart_message_and_answer, save_chart, \ + finish_record from apps.chat.models.chat_model import ChatQuestion, ChatRecord from apps.datasource.models.datasource import CoreDatasource from apps.db.db import exec_sql @@ -92,96 +95,151 @@ class LLMService: full_message=json.dumps( [{'type': msg.type, 'content': msg.content} for msg in self.sql_message], ensure_ascii=False)) - full_text = '' + full_sql_text = '' res = self.llm.stream(self.sql_message) for chunk in res: if isinstance(chunk, dict): - full_text += chunk['content'] + full_sql_text += chunk['content'] yield chunk['content'] continue if isinstance(chunk, AIMessageChunk): - full_text += chunk.content + full_sql_text += chunk.content yield chunk.content continue - print(full_text) - self.sql_message.append(AIMessage(full_text)) + self.sql_message.append(AIMessage(full_sql_text)) self.record = save_full_sql_message_and_answer(session=session, record_id=self.record.id, - answer=full_text, + answer=full_sql_text, full_message=json.dumps( [{'type': msg.type, 'content': msg.content} for msg in self.sql_message], ensure_ascii=False)) - 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() - - # async for chunk in chain.astream({"schema": mschema_str, "question": question}): - for chunk in chain.stream({"schema": schema, "question": question}): - if not isinstance(chunk, dict): + def generate_chart(self, session: SessionDep): + # append current question + self.chart_message.append(HumanMessage(self.chat_question.chart_user_question())) + self.record = save_full_chart_message(session=session, record_id=self.record.id, + full_message=json.dumps( + [{'type': msg.type, 'content': msg.content} for msg in + self.chart_message], ensure_ascii=False)) + full_chart_text = '' + res = self.llm.stream(self.chart_message) + for chunk in res: + if isinstance(chunk, dict): + full_chart_text += chunk['content'] + yield chunk['content'] + continue + if isinstance(chunk, AIMessageChunk): + full_chart_text += chunk.content + yield chunk.content continue - if "agent" in chunk: - messages = chunk["agent"].get("messages", []) - for msg in messages: - if tool_calls := msg.additional_kwargs.get("tool_calls"): - for tool_call in tool_calls: - response = { - "type": "tool_call", - "tool": tool_call["function"]["name"], - "args": tool_call["function"]["arguments"] - } - yield f"data: {json.dumps(response, ensure_ascii=False)}\n\n" + self.chart_message.append(AIMessage(full_chart_text)) + self.record = save_full_chart_message_and_answer(session=session, record_id=self.record.id, + answer=full_chart_text, + full_message=json.dumps( + [{'type': msg.type, 'content': msg.content} for msg in + self.chart_message], ensure_ascii=False)) - if content := msg.content: - html_start = content.find("```html") - html_end = content.find("```", html_start + 6) - if html_start != -1 and html_end != -1: - html_content = content[html_start + 7:html_end].strip() - response = { - "type": "final", - "content": content.split("```html")[0].strip(), - "html": html_content - } - else: - response = { - "type": "final", - "content": content - } - yield f"data: {json.dumps(response, ensure_ascii=False)}\n\n" + def check_save_sql(self, session: SessionDep, res: str) -> str: - if "tools" in chunk: - messages = chunk["tools"].get("messages", []) - for msg in messages: - response = { - "type": "tool_result", - "tool": msg.name, - "content": msg.content - } - yield f"data: {json.dumps(response, ensure_ascii=False)}\n\n" + json_str = extract_nested_json(res) + data = json.loads(json_str) - await asyncio.sleep(0.1) + sql = '' + message = '' + error = False - yield f"data: {json.dumps({'type': 'complete'})}\n\n" + if data['success']: + sql = data['sql'] + else: + message = data['message'] + error = True + + if error: + raise Exception(message) + if sql.strip() == '': + raise Exception("SQL query is empty") + + save_sql(session=session, sql=sql, record_id=self.record.id) + + self.chat_question.sql = sql + + return sql + + def check_save_chart(self, session: SessionDep, res: str) -> Dict[str, Any]: + + json_str = extract_nested_json(res) + data = json.loads(json_str) + + chart: Dict[str, Any] = {} + message = '' + error = False + + if data['type'] and data['type'] != 'error': + # todo type check + chart = data + elif data['type'] == 'error': + message = data['reason'] + error = True + else: + message = 'Chart is empty' + error = True + + if error: + raise Exception(message) + + save_chart(session=session, chart=json.dumps(chart, ensure_ascii=False), record_id=self.record.id) + + return chart + + def save_error(self, session: SessionDep, message: str): + return save_error_message(session=session, record_id=self.record.id, message=message) + + def save_sql_data(self, session: SessionDep, data_obj: Dict[str, Any]): + return save_sql_exec_data(session=session, record_id=self.record.id, data=json.dumps(data_obj, ensure_ascii=False)) + + def finish(self,session: SessionDep): + return finish_record(session=session, record_id=self.record.id) + + def execute_sql(self, sql: str): + """Execute SQL query + + Args: + ds: Data source instance + sql: SQL query statement + + Returns: + Query results + """ + print(f"Executing SQL on ds_id {self.ds.id}: {sql}") + return exec_sql(self.ds, sql) -def execute_sql(ds: CoreDatasource, sql: str) -> str: - """Execute SQL query +def extract_nested_json(text): + stack = [] + start_index = -1 + results = [] - Args: - ds: Data source instance - sql: SQL query statement - - Returns: - Query results - """ - print(f"Executing SQL on ds_id {ds.id}: {sql}") - return exec_sql(ds, sql) + for i, char in enumerate(text): + if char in '{[': + if not stack: # 记录起始位置 + start_index = i + stack.append(char) + elif char in '}]': + if stack and ((char == '}' and stack[-1] == '{') or (char == ']' and stack[-1] == '[')): + stack.pop() + if not stack: # 栈空时截取完整JSON + json_str = text[start_index:i + 1] + try: + json.loads(json_str) # 验证有效性 + results.append(json_str) + except: + pass + else: + stack = [] # 括号不匹配则重置 + if results[0]: + return results[0] + return None def execute_sql_with_db(db: SQLDatabase, sql: str) -> str: diff --git a/backend/apps/db/db.py b/backend/apps/db/db.py index 9378c27e..8bde9578 100644 --- a/backend/apps/db/db.py +++ b/backend/apps/db/db.py @@ -96,22 +96,19 @@ def get_tables(ds: CoreDatasource): """ elif ds.type == "pg" or ds.type == "excel": sql = """ - SELECT - c.relname AS TABLE_NAME, - COALESCE(d.description, obj_description(c.oid)) AS TABLE_COMMENT - FROM - pg_class c - LEFT JOIN - pg_namespace n ON n.oid = c.relnamespace - LEFT JOIN - pg_description d ON d.objoid = c.oid AND d.objsubid = 0 - WHERE - n.nspname = current_schema() - AND c.relkind IN ('r', 'v') - AND c.relname NOT LIKE 'pg_%' - AND c.relname NOT LIKE 'sql_%' - ORDER BY c.relname - """ + SELECT c.relname AS TABLE_NAME, + COALESCE(d.description, obj_description(c.oid)) AS TABLE_COMMENT + FROM pg_class c + LEFT JOIN + pg_namespace n ON n.oid = c.relnamespace + LEFT JOIN + pg_description d ON d.objoid = c.oid AND d.objsubid = 0 + WHERE n.nspname = current_schema() + AND c.relkind IN ('r', 'v') + AND c.relname NOT LIKE 'pg_%' + AND c.relname NOT LIKE 'sql_%' + ORDER BY c.relname \ + """ elif ds.type == "oracle": sql = f""" SELECT @@ -183,21 +180,18 @@ def get_fields(ds: CoreDatasource, table_name: str = None): sql = sql1 + sql2 elif ds.type == "pg" or ds.type == "excel": sql1 = """ - SELECT - a.attname AS COLUMN_NAME, - pg_catalog.format_type(a.atttypid, a.atttypmod) AS DATA_TYPE, - col_description(c.oid, a.attnum) AS COLUMN_COMMENT - FROM - pg_catalog.pg_attribute a - JOIN + SELECT a.attname AS COLUMN_NAME, + pg_catalog.format_type(a.atttypid, a.atttypmod) AS DATA_TYPE, + col_description(c.oid, a.attnum) AS COLUMN_COMMENT + FROM pg_catalog.pg_attribute a + JOIN pg_catalog.pg_class c ON a.attrelid = c.oid - JOIN - pg_catalog.pg_namespace n ON n.oid = c.relnamespace - WHERE - n.nspname = current_schema() - AND a.attnum > 0 - AND NOT a.attisdropped - """ + JOIN + pg_catalog.pg_namespace n ON n.oid = c.relnamespace + WHERE n.nspname = current_schema() + AND a.attnum > 0 + AND NOT a.attisdropped \ + """ sql2 = f" AND c.relname = '{table_name}'" if table_name is not None and table_name != "" else "" sql = sql1 + sql2 elif ds.type == "oracle": @@ -248,6 +242,8 @@ def exec_sql(ds: CoreDatasource, sql: str): for tuple_item in res ] return {"fields": columns, "data": result_list, "sql": bytes.decode(base64.b64encode(bytes(sql, 'utf-8')))} + except Exception as ex: + raise ex finally: if result is not None: result.close() diff --git a/backend/template.yaml b/backend/template.yaml index 60dee447..2a38f9ce 100644 --- a/backend/template.yaml +++ b/backend/template.yaml @@ -9,7 +9,12 @@ template: - SQL查询中不能使用星号(*)，必须明确指定字段名. - SQL查询的字段名和别名必须为英文。 - 计算占比，百分比类型字段，保留两位小数，以%结尾。 - - 问题与生成SQL无关时，直接回答：抱歉，我无法回答您的问题。 + - 生成的SQL使用JSON格式返回： + {{"success":true,"sql":"生成的SQL语句"}} + - 问题与生成SQL无关时，直接回答： + {{"success":false,"message":"抱歉，我无法回答您的问题。"}} + - 如果根据提供的表结构不能生成符合问题与条件的SQL，回答： + {{"success":false,"message":"无法生成SQL的原因"}} user: | 表结构: {schema} diff --git a/frontend/src/views/chat/ChatRow.vue b/frontend/src/views/chat/ChatRow.vue index 3ffe1449..2230a3ca 100644 --- a/frontend/src/views/chat/ChatRow.vue +++ b/frontend/src/views/chat/ChatRow.vue @@ -27,7 +27,7 @@ const _datasource = computed({