feat: execute sql & generate chart with history

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")

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
+

backend/apps/chat/schemas/m_schema.py

@@ -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", [])

backend/apps/chat/schemas/schema_engine.py

@@ -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
-                )

backend/apps/chat/schemas/utils.py

@@ -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]

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:

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()

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}

frontend/src/views/chat/ChatRow.vue

@@ -27,7 +27,7 @@ const _datasource = computed({
 
 <template>
   <div class="chat-row" :class="{'right-to-left': msg.role === 'user'}">
-    <el-avatar shape="square" v-if="msg.role === 'assistant'">SQBot</el-avatar>
+    <el-avatar shape="square" v-if="msg.role === 'assistant'">SQLBot</el-avatar>
     <el-avatar shape="square" v-if="msg.role === 'user'"/>
     <ChatBlock v-if="!msg.isWelcome" :msg="msg" :class="{'row-full': msg.role === 'assistant'}">
       <slot></slot>