klinsmaya/SQLBot
1
0
Fork 0
mirror of https://gh-proxy.org/https://github.com/klinsmaya/SQLBot.git synced 2026-11-08 01:06:59 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: generate sql with history

Browse Source
This commit is contained in:
ulleo
2025-05-30 17:41:45 +08:00
parent 5de9b7d330
commit 054296741c
25 changed files with 526 additions and 366 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/alembic/versions/1f077c30e476_009_modify_chat.py
+64
View File
@@ -0,0 +1,64 @@
"""009_modify_chat
Revision ID: 1f077c30e476
Revises: 35d925df4568
Create Date: 2025-05-30 16:11:08.020715
"""
from alembic import op
import sqlalchemy as sa
import sqlmodel.sql.sqltypes
from sqlalchemy.dialects import postgresql
# revision identifiers, used by Alembic.
revision = '1f077c30e476'
down_revision = '35d925df4568'
branch_labels = None
depends_on = None
def upgrade():
# ### commands auto generated by Alembic - please adjust! ###
op.drop_table('chat_record')
op.create_table('chat_record',
sa.Column('id', sa.Integer(), sa.Identity(always=True), nullable=False),
sa.Column('chat_id', sa.Integer(), nullable=True),
sa.Column('create_time', sa.DateTime(timezone=True), nullable=True),
sa.Column('finish_time', sa.DateTime(timezone=True), nullable=True),
sa.Column('create_by', sa.BigInteger(), nullable=True),
sa.Column('datasource', sa.Integer(), nullable=False),
sa.Column('engine_type', sqlmodel.sql.sqltypes.AutoString(length=64), nullable=False),
sa.Column('question', sa.Text(), nullable=True),
sa.Column('sql_answer', sa.Text(), nullable=True),
sa.Column('sql', sa.Text(), nullable=True),
sa.Column('sql_exec_result', sa.Text(), nullable=True),
sa.Column('data', sa.Text(), nullable=True),
sa.Column('chart_answer', sa.Text(), nullable=True),
sa.Column('chart', sa.Text(), nullable=True),
sa.Column('full_sql_message', sa.Text(), nullable=True),
sa.Column('full_chart_message', sa.Text(), nullable=True),
sa.Column('finish', sa.Boolean(), nullable=True),
sa.Column('error', sa.Text(), nullable=True),
sa.Column('run_time', sa.Float(), nullable=False),
sa.PrimaryKeyConstraint('id')
)
# ### end Alembic commands ###
def downgrade():
# ### commands auto generated by Alembic - please adjust! ###
op.drop_table('chat_record')
op.create_table('chat_record',
sa.Column('id', sa.Integer(), sa.Identity(always=True), nullable=False),
sa.Column('chat_id', sa.Integer(), nullable=True),
sa.Column('create_time', sa.DateTime(timezone=True), nullable=True),
sa.Column('create_by', sa.BigInteger(), nullable=True),
sa.Column('datasource', sa.Integer(), nullable=False),
sa.Column('engine_type', sqlmodel.sql.sqltypes.AutoString(length=64), nullable=False),
sa.Column('question', sa.Text(), nullable=True),
sa.Column('full_question', sa.Text(), nullable=True),
sa.Column('answer', sa.Text(), nullable=True),
sa.Column('run_time', sa.Float(), nullable=False),
sa.PrimaryKeyConstraint('id')
)
# ### end Alembic commands ###
backend/apps/ai_model/__init__.py
View File
backend/apps/ai_model/llm.py
+1
View File
@@ -0,0 +1 @@
# todo
backend/apps/chat/schemas/chat_base_schema.py → backend/apps/ai_model/model_factory.py
+17 -1
View File
@@ -3,6 +3,10 @@ from typing import Optional, Dict, Any, Type
from abc import ABC, abstractmethod
from langchain_core.language_models import BaseLLM as LangchainBaseLLM
from langchain_openai import ChatOpenAI
from apps.system.models.system_model import AiModelDetail
# from langchain_community.llms import Tongyi, VLLM
class LLMConfig(BaseModel):
@@ -84,4 +88,16 @@ class LLMFactory:
@classmethod
def register_llm(cls, model_type: str, llm_class: Type[BaseLLM]):
"""Register new model type"""
cls._llm_types[model_type] = llm_class
cls._llm_types[model_type] = llm_class
# todo
def get_llm_config(aimodel: AiModelDetail) -> LLMConfig:
config = LLMConfig(
model_type="openai",
model_name=aimodel.name,
api_key=aimodel.api_key,
api_base_url=aimodel.endpoint,
additional_params={"temperature": aimodel.temperature}
)
return config
backend/apps/ai_model/openai/__init__.py
View File
backend/apps/ai_model/openai/llm.py
+1
View File
@@ -0,0 +1 @@
# todo
backend/apps/chat/api/chat.py
+51 -97
View File
@@ -1,19 +1,19 @@
import asyncio
import json
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, \
delete_chat
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
delete_chat, list_records
from apps.chat.models.chat_model import CreateChat, ChatRecord, RenameChat, Chat, ChatQuestion
from apps.chat.task.llm import LLMService
from apps.datasource.crud.datasource import get_table_schema
from apps.datasource.models.datasource import CoreDatasource
from apps.system.models.system_model import AiModelDetail
from common.core.deps import SessionDep, CurrentUser
import json
import asyncio
router = APIRouter(tags=["Data Q&A"], prefix="/chat")
@@ -96,14 +96,7 @@ async def stream_sql(session: SessionDep, current_user: CurrentUser, request_que
detail="No available datasource configuration found"
)
record: ChatRecord
try:
record = save_question(session=session, current_user=current_user, question=request_question)
except Exception as e1:
raise HTTPException(
status_code=400,
detail=str(e1)
)
request_question.engine = ds.type_name
# Get available AI model
aimodel = session.exec(select(AiModelDetail).where(
@@ -112,93 +105,54 @@ async def stream_sql(session: SessionDep, current_user: CurrentUser, request_que
)).first()
if not aimodel:
raise HTTPException(
status_code=400,
status_code=500,
detail="No available AI model configuration found"
)
# Use Tongyi Qianwen
tongyi_config = LLMConfig(
model_type="openai",
model_name=aimodel.name,
api_key=aimodel.api_key,
api_base_url=aimodel.endpoint,
additional_params={"temperature": aimodel.temperature}
)
# llm_service = LLMService(tongyi_config)
llm_service = AgentService(tongyi_config, ds)
# Use Custom VLLM model
""" vllm_config = LLMConfig(
model_type="vllm",
model_name="your-model-path",
additional_params={
"max_new_tokens": 200,
"temperature": 0.3
}
)
vllm_service = LLMService(vllm_config) """
""" result = llm_service.generate_sql(question)
return result """
history_records: List[ChatRecord] = list_records(session=session, current_user=current_user,
chart_id=request_question.chat_id)
# 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"
request_question.db_schema = get_table_schema(session=session, ds=ds)
llm_service = LLMService(request_question, history_records, ds, aimodel)
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'
llm_service.init_record(session=session, current_user=current_user)
print(schema_str)
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'
async def event_generator():
all_text = ''
try:
async for chunk in llm_service.async_generate(question, schema_str):
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)
)
# 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)
# )
# return EventSourceResponse(event_generator(), headers={"Content-Type": "text/event-stream"})
return StreamingResponse(event_generator(), media_type="text/event-stream")
return StreamingResponse(run_task(), media_type="text/event-stream")
backend/apps/chat/curd/chat.py
+32 -7
View File
@@ -1,12 +1,9 @@
import datetime
import json
from typing import List
from sqlalchemy import text, and_
from sqlmodel import select
from sqlalchemy import and_
from apps.chat.models.chat_model import Chat, ChatRecord, CreateChat, ChatInfo, RenameChat
from apps.chat.schemas.chat_schema import ChatQuestion
from apps.chat.models.chat_model import Chat, ChatRecord, CreateChat, ChatInfo, RenameChat, ChatQuestion
from apps.datasource.models.datasource import CoreDatasource
from common.core.deps import SessionDep, CurrentUser
@@ -66,6 +63,12 @@ def get_chat_with_records(session: SessionDep, chart_id: int, current_user: Curr
return chat_info
def list_records(session: SessionDep, chart_id: int, current_user: CurrentUser) -> List[ChatRecord]:
record_list = session.query(ChatRecord).filter(
and_(Chat.create_by == current_user.id, ChatRecord.chat_id == chart_id)).order_by(ChatRecord.create_time).all()
return record_list
def create_chat(session: SessionDep, current_user: CurrentUser, create_chat_obj: CreateChat) -> ChatInfo:
if not create_chat_obj.datasource:
raise Exception("Datasource cannot be None")
@@ -125,11 +128,33 @@ def save_question(session: SessionDep, current_user: CurrentUser, question: Chat
return result
def save_full_question(session: SessionDep, id: int, full_question: str) -> ChatRecord:
def save_full_sql_message(session: SessionDep, record_id: int, full_message: str) -> ChatRecord:
return save_full_sql_message_and_answer(session=session, record_id=record_id, full_message=full_message, answer='')
def save_full_sql_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_sql_message = full_message
record.sql_answer = answer
result = ChatRecord(**record.model_dump())
session.add(record)
session.flush()
session.refresh(record)
session.commit()
return result
def save_full_chart_message(session: SessionDep, id: int, full_message: str) -> ChatRecord:
if not id:
raise Exception("Record id cannot be None")
record = session.query(ChatRecord).filter(ChatRecord.id == id).first()
record.full_question = full_question
record.full_chart_message = full_message
result = ChatRecord(**record.model_dump())
backend/apps/chat/models/chat_model.py
+38 -3
View File
@@ -1,9 +1,12 @@
from sqlmodel import SQLModel, Field
from sqlalchemy import Column, Text, BigInteger, DateTime, Integer, Identity
from sqlalchemy import Column, Text, BigInteger, DateTime, Integer, Identity, Boolean
from datetime import datetime
from pydantic import BaseModel
from typing import List, Optional
from apps.template.generate_chart.generator import get_chart_template
from apps.template.generate_sql.generator import get_sql_template
class Chat(SQLModel, table=True):
__tablename__ = "chat"
@@ -21,12 +24,21 @@ class ChatRecord(SQLModel, table=True):
id: Optional[int] = Field(sa_column=Column(Integer, Identity(always=True), primary_key=True))
chat_id: int = Field(sa_column=Column(Integer))
create_time: datetime = Field(sa_column=Column(DateTime(timezone=True), nullable=True))
finish_time: datetime = Field(sa_column=Column(DateTime(timezone=True), nullable=True))
create_by: int = Field(sa_column=Column(BigInteger, nullable=True))
datasource: int = Field(sa_column=Column(Integer, nullable=False))
engine_type: str = Field(max_length=64)
question: str = Field(sa_column=Column(Text, nullable=True))
full_question: str = Field(sa_column=Column(Text, nullable=True))
answer: 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))
full_sql_message: str = Field(sa_column=Column(Text, nullable=True))
full_chart_message: str = Field(sa_column=Column(Text, nullable=True))
finish: bool = Field(sa_column=Column(Boolean, nullable=True, default=False))
error: str = Field(sa_column=Column(Text, nullable=True))
run_time: float = Field(default=0)
@@ -52,3 +64,26 @@ class ChatInfo(BaseModel):
datasource_name: str = ''
datasource_exists: bool = True
records: List[ChatRecord] = []
class AiModelQuestion(BaseModel):
engine: str = ""
db_schema: str = ""
sql: str = ""
def sql_sys_question(self):
return get_sql_template()['system'].format(engine=self.engine, schema=self.db_schema, question=self.question)
def sql_user_question(self):
return get_sql_template()['user'].format(engine=self.engine, schema=self.db_schema, question=self.question)
def chart_sys_question(self):
return get_chart_template()['system'].format(sql=self.sql, question=self.question)
def chart_user_question(self):
return get_chart_template()['user'].format(sql=self.sql, question=self.question)
class ChatQuestion(AiModelQuestion):
question: str
chat_id: int
backend/apps/chat/schemas/chat_schema.py
-7
View File
@@ -1,7 +0,0 @@
from pydantic import BaseModel
class ChatQuestion(BaseModel):
question: str
chat_id: int
backend/apps/chat/schemas/llm.py
-197
View File
@@ -1,197 +0,0 @@
from langchain_community.utilities import SQLDatabase
from langgraph.prebuilt import create_react_agent
from langchain_core.prompts import ChatPromptTemplate
from apps.chat.schemas.chat_base_schema import LLMConfig, LLMFactory
from apps.chat.schemas.schema_engine import SchemaEngine
from apps.datasource.models.datasource import CoreDatasource
from apps.db.db import exec_sql, get_uri
from common.core.config import settings
import warnings
from langchain.tools import Tool
from functools import partial
import logging
from typing import AsyncGenerator
import json
import asyncio
warnings.filterwarnings("ignore")
class LLMService:
def __init__(self, config: LLMConfig):
# Initialize database connection
self.db = SQLDatabase.from_uri(str(settings.SQLALCHEMY_DATABASE_URI))
# Create LLM instance through factory
llm_instance = LLMFactory.create_llm(config)
self.llm = llm_instance.llm
# Define prompt template
self.prompt = ChatPromptTemplate.from_messages([
("system", """You are a professional SQL engineer. Generate PostgreSQL SELECT queries based on the database schema and user questions.
Data modification or deletion is prohibited. Table structure is as follows:
{schema}
"""),
("human", "{question}")
])
def generate_sql(self, question: str) -> str:
chain = self.prompt | self.llm
schema = self.db.get_table_info()
return chain.invoke({"schema": schema, "question": question})
class AgentService:
def __init__(self, config: LLMConfig, ds: CoreDatasource):
# Initialize database connection
self.ds = ds
db_uri = get_uri(ds)
self.db = SQLDatabase.from_uri(db_uri)
# self.db = SQLDatabase.from_uri(str(settings.SQLALCHEMY_DATABASE_URI))
# Create LLM instance through factory
llm_instance = LLMFactory.create_llm(config)
self.llm = llm_instance.llm
# Create a partial function of execute_sql with preset ds parameter
# bound_execute_sql = partial(execute_sql, self.ds)
bound_execute_sql = partial(execute_sql_with_db, self.db)
# Wrap as Tool object
tools = [
Tool(
name="execute_sql",
func=bound_execute_sql,
description="""A tool for executing SQL queries.
Input: SQL query statement (string)
Output: Query results
Example: "SELECT * FROM table_name LIMIT 5"
"""
)
]
self.agent_executor = create_react_agent(self.llm, tools)
system_prompt = """
You are an intelligent agent capable of data analysis. When users input their data analysis requirements,
you need to first convert the requirements into executable SQL, then execute the SQL through tools to return results,
and finally summarize the SQL query results. When all tasks are completed, you need to generate an HTML format data analysis report.
You can analyze requirements step by step to determine the final SQL query to generate.
To improve SQL generation accuracy, please evaluate the accuracy of the SQL after generation,
if there are issues, regenerate the SQL.
When SQL execution fails, you need to correct the SQL based on the error message and try to execute again.
### Tools ###
execute_sql: Can execute SQL by passing in SQL statements and return execution results
"""
user_prompt = """
Below is the database information I need to query:
{schema}
My requirement is: {question}
"""
# Define prompt template
self.prompt = ChatPromptTemplate.from_messages([
("system", system_prompt),
("human", user_prompt)
])
def generate_sql(self, question: str) -> str:
chain = self.prompt | self.agent_executor
schema = self.db.get_table_info()
return chain.invoke({"schema": schema, "question": question})
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}):
async for chunk in chain.astream({"schema": schema, "question": question}):
if not isinstance(chunk, dict):
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"
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"
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"
await asyncio.sleep(0.1)
yield f"data: {json.dumps({'type': 'complete'})}\n\n"
def execute_sql(ds: CoreDatasource, sql: str) -> str:
"""Execute SQL query
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)
def execute_sql_with_db(db: SQLDatabase, sql: str) -> str:
"""Execute SQL query using SQLDatabase
Args:
db: SQLDatabase instance
sql: SQL query statement
Returns:
str: Query results formatted as string
"""
try:
# Execute query
result = db.run(sql)
if not result:
return "Query executed successfully but returned no results."
# Format results
return str(result)
except Exception as e:
error_msg = f"SQL execution failed: {str(e)}"
logging.error(error_msg)
raise RuntimeError(error_msg)
backend/apps/chat/task/__init__.py
View File
backend/apps/chat/task/llm.py
+210
View File
@@ -0,0 +1,210 @@
import asyncio
import json
import logging
import warnings
from typing import Any, List, Union, AsyncGenerator
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.models.chat_model import ChatQuestion, ChatRecord
from apps.datasource.models.datasource import CoreDatasource
from apps.db.db import exec_sql
from apps.system.models.system_model import AiModelDetail
from common.core.deps import SessionDep, CurrentUser
warnings.filterwarnings("ignore")
class LLMService:
ds: CoreDatasource
chat_question: ChatQuestion
record: ChatRecord
config: LLMConfig
llm: BaseLLM
sql_message: List[Union[BaseMessage, dict[str, Any]]] = []
chart_message: List[Union[BaseMessage, dict[str, Any]]] = []
def __init__(self, chat_question: ChatQuestion, history_records: List[ChatRecord],
ds: CoreDatasource, aimodel: AiModelDetail):
self.ds = ds
self.chat_question = chat_question
self.config = get_llm_config(aimodel)
# Create LLM instance through factory
llm_instance = LLMFactory.create_llm(self.config)
self.llm = llm_instance.llm
# self.agent_executor = create_react_agent(self.llm)
last_sql_messages = list(
filter(lambda r: True if r.full_sql_message is not None and r.full_sql_message.strip() != '' else False,
history_records))
last_sql_message_str = "[]" if last_sql_messages is None or len(last_sql_messages) == 0 else last_sql_messages[
-1].full_sql_message
last_chart_messages = list(
filter(
lambda r: True if r.full_chart_message is not None and r.full_chart_message.strip() != '' else False,
history_records))
last_chart_message_str = "[]" if last_chart_messages is None or len(last_chart_messages) == 0 else \
last_chart_messages[-1].full_chart_message
last_sql_messages: List[dict[str, Any]] = json.loads(last_sql_message_str)
if last_sql_messages is None or len(last_sql_messages) == 0:
# add sys prompt
self.sql_message.append(SystemMessage(content=self.chat_question.sql_sys_question()))
else:
for last_sql_message in last_sql_messages:
_msg: BaseMessage
if last_sql_message['type'] == 'human':
_msg = HumanMessage(content=last_sql_message['content'])
elif last_sql_message['type'] == 'ai':
_msg = AIMessage(content=last_sql_message['content'])
else:
_msg = SystemMessage(content=last_sql_message['content'])
self.sql_message.append(_msg)
last_chart_messages: List[dict[str, Any]] = json.loads(last_chart_message_str)
if last_chart_messages is None or len(last_chart_messages) == 0:
# add sys prompt
self.chart_message.append(SystemMessage(content=self.chat_question.chart_sys_question()))
else:
for last_chart_message in last_chart_messages:
_msg: BaseMessage
if last_chart_message['type'] == 'human':
_msg = HumanMessage(content=last_chart_message['content'])
elif last_chart_message['type'] == 'ai':
_msg = AIMessage(content=last_chart_message['content'])
else:
_msg = SystemMessage(content=last_chart_message['content'])
self.chart_message.append(_msg)
def init_record(self, session: SessionDep, current_user: CurrentUser) -> ChatRecord:
self.record = save_question(session=session, current_user=current_user, question=self.chat_question)
return self.record
def generate_sql(self, session: SessionDep):
# append current question
self.sql_message.append(HumanMessage(self.chat_question.sql_user_question()))
self.record = save_full_sql_message(session=session, record_id=self.record.id,
full_message=json.dumps(
[{'type': msg.type, 'content': msg.content} for msg in
self.sql_message], ensure_ascii=False))
full_text = ''
res = self.llm.stream(self.sql_message)
for chunk in res:
if isinstance(chunk, dict):
full_text += chunk['content']
yield chunk['content']
continue
if isinstance(chunk, AIMessageChunk):
full_text += chunk.content
yield chunk.content
continue
print(full_text)
self.sql_message.append(AIMessage(full_text))
self.record = save_full_sql_message_and_answer(session=session, record_id=self.record.id,
answer=full_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):
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"
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"
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"
await asyncio.sleep(0.1)
yield f"data: {json.dumps({'type': 'complete'})}\n\n"
def execute_sql(ds: CoreDatasource, sql: str) -> str:
"""Execute SQL query
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)
def execute_sql_with_db(db: SQLDatabase, sql: str) -> str:
"""Execute SQL query using SQLDatabase
Args:
db: SQLDatabase instance
sql: SQL query statement
Returns:
str: Query results formatted as string
"""
try:
# Execute query
result = db.run(sql)
if not result:
return "Query executed successfully but returned no results."
# Format results
return str(result)
except Exception as e:
error_msg = f"SQL execution failed: {str(e)}"
logging.error(error_msg)
raise RuntimeError(error_msg)
backend/apps/datasource/crud/datasource.py
+29
View File
@@ -220,3 +220,32 @@ def get_table_obj_by_ds(session: SessionDep, ds: CoreDatasource) -> List[TableAn
fields = session.query(CoreField).filter(and_(CoreField.table_id == table.id, CoreField.checked == True)).all()
_list.append(TableAndFields(schema=schema, table=table, fields=fields))
return _list
def get_table_schema(session: SessionDep, ds: CoreDatasource) -> str:
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'
return schema_str
backend/apps/template/__init__.py
View File
backend/apps/template/generate_analysis/__init__.py
View File
backend/apps/template/generate_chart/__init__.py
View File
backend/apps/template/generate_chart/generator.py
+6
View File
@@ -0,0 +1,6 @@
from apps.template.template import get_base_template
def get_chart_template():
template = get_base_template()
return template['template']['chart']
backend/apps/template/generate_sql/__init__.py
View File
backend/apps/template/generate_sql/generator.py
+6
View File
@@ -0,0 +1,6 @@
from apps.template.template import get_base_template
def get_sql_template():
template = get_base_template()
return template['template']['sql']
backend/apps/template/template.py
+15
View File
@@ -0,0 +1,15 @@
import yaml
base_template = None
def load():
with open('./template.yaml', 'r') as f:
global base_template
base_template = yaml.load(f, Loader=yaml.SafeLoader)
def get_base_template():
if not base_template:
load()
return base_template
backend/pyproject.toml
+2 -1
View File
@@ -32,7 +32,8 @@ dependencies = [
"pandas (>=2.2.3,<3.0.0)",
"openpyxl (>=3.1.5,<4.0.0)",
"psycopg2-binary (>=2.9.10,<3.0.0)",
"oracledb (>=3.1.1,<4.0.0)"
"oracledb (>=3.1.1,<4.0.0)",
"pyyaml (>=6.0.2,<7.0.0)"
]
[[tool.uv.index]]
url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple"
backend/template.yaml
+52
View File
@@ -0,0 +1,52 @@
template:
sql:
system: |
任务:
根据表结构和问题生成符合{engine}数据库引擎规范的sql语句。
你必须遵守以下规则:
- 生成的SQL必须符合{engine}的规范。
- 根据表结构生成SQL语句。
- SQL查询中不能使用星号(*)，必须明确指定字段名.
- SQL查询的字段名和别名必须为英文。
- 计算占比，百分比类型字段，保留两位小数，以%结尾。
- 问题与生成SQL无关时，直接回答：抱歉，我无法回答您的问题。
user: |
表结构:
{schema}
问题:
{question}
chart:
system: |
### 说明：
您的任务是通过给定的问题和SQL生成 JSON 以进行数据可视化。
请遵守以下规则：
- 如果需要表格，则生成的 JSON 格式应为：
{{"type":"table", "title": "标题", "columns": [{{"name":"中文字段名1", "value": "SQL 查询列 1(有别名用别名)"}}, {{"name": "中文字段名 2", "value": "SQL 查询列 2(有别名用别名)"}}]}}
必须从 SQL 查询列中提取“columns”。
- 如果需要条形图，则生成的 JSON 格式应为：
{{"type":"bar", "title": "标题", "axis": {{"x": {{"name":"x轴的中文名称", "value": "SQL 查询 x 轴的列(有别名用别名)"}}, "y": {{"name":"y轴的中文名称","value": "SQL 查询 y 轴的列(有别名用别名)"}}}}
必须从 SQL 查询列中提取“x”和“y”。
- 如果需要折线图，则生成的 JSON 格式应为：
{{"type":"line", "title": "标题", "axis": {{"x": {{"name":"x轴的中文名称","value": "x轴的SQL查询列(有别名用别名)"}}, "y": {{"name":"y轴的中文名称","value": "y轴的SQL查询列(有别名用别名)"}}}}
其中“x”和“y”必须从SQL查询列中提取。
- 如果需要饼图，则生成的 JSON 格式应为：
{{"type":"pie", "title": "标题", "column": {{"name":"中文字段名","value":"SQL查询列1(有别名用别名)"}}}}
其中“column”必须从SQL查询列中提取。
- 如果答案未知或者与生成JSON无关，则生成的 JSON 格式应为：
{{"type":"error", "reason": "抱歉，我无法回答您的问题。"}}
### 示例
如果 SQL 为： SELECT products_sales_data.category, AVG(products_sales_data.price) AS average_price FROM products_sales_data GROUP BY products_sales_data.category;
问题是：每个商品分类的平均价格
则生成的 JSON 可以是： {{"type":"table", "title": "每个商品分类的平均价格", "columns": [{{"name":"商品分类","value":"category"}}, {{"name":"平均价格","value":"average_price"}}]}}
### 响应：
根据您的指示，这是我生成的与 问题 和 sql 匹配的 JSON：
```json
user: |
### SQL:
{sql}
### 问题：
{question}
frontend/components.d.ts
Vendored
-51
View File
@@ -1,51 +0,0 @@
/* eslint-disable */
/* prettier-ignore */
// @ts-nocheck
// Generated by unplugin-vue-components
// Read more: https://github.com/vuejs/core/pull/3399
import '@vue/runtime-core'
export {}
declare module '@vue/runtime-core' {
export interface GlobalComponents {
ElAside: typeof import('element-plus-secondary/es')['ElAside']
ElAvatar: typeof import('element-plus-secondary/es')['ElAvatar']
ElButton: typeof import('element-plus-secondary/es')['ElButton']
ElCheckbox: typeof import('element-plus-secondary/es')['ElCheckbox']
ElContainer: typeof import('element-plus-secondary/es')['ElContainer']
ElDialog: typeof import('element-plus-secondary/es')['ElDialog']
ElDivider: typeof import('element-plus-secondary/es')['ElDivider']
ElDropdown: typeof import('element-plus-secondary/es')['ElDropdown']
ElDropdownItem: typeof import('element-plus-secondary/es')['ElDropdownItem']
ElDropdownMenu: typeof import('element-plus-secondary/es')['ElDropdownMenu']
ElFooter: typeof import('element-plus-secondary/es')['ElFooter']
ElForm: typeof import('element-plus-secondary/es')['ElForm']
ElFormItem: typeof import('element-plus-secondary/es')['ElFormItem']
ElHeader: typeof import('element-plus-secondary/es')['ElHeader']
ElIcon: typeof import('element-plus-secondary/es')['ElIcon']
ElInput: typeof import('element-plus-secondary/es')['ElInput']
ElMain: typeof import('element-plus-secondary/es')['ElMain']
ElMenu: typeof import('element-plus-secondary/es')['ElMenu']
ElMenuItem: typeof import('element-plus-secondary/es')['ElMenuItem']
ElOption: typeof import('element-plus-secondary/es')['ElOption']
ElPagination: typeof import('element-plus-secondary/es')['ElPagination']
ElScrollbar: typeof import('element-plus-secondary/es')['ElScrollbar']
ElSelect: typeof import('element-plus-secondary/es')['ElSelect']
ElSlider: typeof import('element-plus-secondary/es')['ElSlider']
ElSwitch: typeof import('element-plus-secondary/es')['ElSwitch']
ElTable: typeof import('element-plus-secondary/es')['ElTable']
ElTableColumn: typeof import('element-plus-secondary/es')['ElTableColumn']
ElTabPane: typeof import('element-plus-secondary/es')['ElTabPane']
ElTabs: typeof import('element-plus-secondary/es')['ElTabs']
ElTooltip: typeof import('element-plus-secondary/es')['ElTooltip']
Icon: typeof import('./src/components/icon-custom/src/Icon.vue')['default']
Layout: typeof import('./src/components/layout/index.vue')['default']
RouterLink: typeof import('vue-router')['RouterLink']
RouterView: typeof import('vue-router')['RouterView']
SinglePage: typeof import('./src/components/layout/SinglePage.vue')['default']
}
export interface ComponentCustomProperties {
vLoading: typeof import('element-plus-secondary/es')['ElLoadingDirective']
}
}
frontend/src/views/system/model/index.vue
+2 -2
View File
@@ -119,13 +119,13 @@
<el-input v-model="state.form.endpoint" placeholder="Please input endpoint" />
</el-form-item>
<el-form-item label="Maximum Context Window" >
<!-- <el-form-item label="Maximum Context Window" >
<el-input v-model="state.form.max_context_window" type="number" placeholder="Please input max_context_window" />
</el-form-item>
<el-form-item label="Temperature" >
<el-slider v-model="state.form.temperature" :format-tooltip="formatTooltip" />
</el-form-item>
</el-form-item>-->
<el-form-item label="Model Status" >
<el-checkbox v-model="state.form.status" label="" size="large" />