Merge remote-tracking branch 'origin/main'

# Conflicts:
#	frontend/components.d.ts
This commit is contained in:
wangjiahao
2025-05-30 17:42:33 +08:00
33 changed files with 705 additions and 433 deletions
+12
View File
@@ -0,0 +1,12 @@
# Python
backend/__pycache__
backend/app.egg-info
backend/*.pyc
backend/.mypy_cache
backend/.coverage
backend/htmlcov
backend/.venv
backend/.gitignore
backend/README.md
backend/scripts
backend/.DS_Store
+1 -1
View File
@@ -48,7 +48,7 @@ jobs:
- name: Build and push Docker image
uses: docker/build-push-action@v4
with:
context: ./backend
context: .
push: true
tags: |
${{ secrets.DOCKERHUB_REGISTRY }}/dataease-v2/sqlbot:latest
+3 -1
View File
@@ -184,4 +184,6 @@ cython_debug/
.pypirc
.DS_Store
test.py
test.py
+63
View File
@@ -0,0 +1,63 @@
# Build stage
FROM registry.cn-qingdao.aliyuncs.com/dataease/python:3.11-slim-bookworm AS builder
# Set build environment variables
ENV PYTHONUNBUFFERED=1
ENV SQLBOT_HOME=/opt/sqlbot
ENV APP_HOME=${SQLBOT_HOME}/app
ENV UI_HOME=${SQLBOT_HOME}/frontend
ENV PYTHONPATH=${SQLBOT_HOME}/app
ENV PATH="${APP_HOME}/.venv/bin:$PATH"
ENV UV_COMPILE_BYTECODE=1
ENV UV_LINK_MODE=copy
ENV DEBIAN_FRONTEND=noninteractive
# Create necessary directories
RUN mkdir -p ${APP_HOME} ${UI_HOME}
WORKDIR ${APP_HOME}
# Install uv tool
COPY --from=ghcr.io/astral-sh/uv:0.7.8 /uv /uvx /bin/
# COPY ../frontend/dist /opt/sqlbot/frontend/dist
COPY frontend/dist ${UI_HOME}/dist
# Install dependencies
RUN test -f "./uv.lock" && \
--mount=type=cache,target=/root/.cache/uv \
--mount=type=bind,source=backend/uv.lock,target=uv.lock \
--mount=type=bind,source=backend/pyproject.toml,target=pyproject.toml \
uv sync --frozen --no-install-project || echo "uv.lock file not found, skipping intermediate-layers"
COPY ./backend ${APP_HOME}
# Final sync to ensure all dependencies are installed
RUN --mount=type=cache,target=/root/.cache/uv \
uv sync
# Runtime stage
FROM registry.cn-qingdao.aliyuncs.com/dataease/python:3.11-slim-bookworm
# Set runtime environment variables
ENV PYTHONUNBUFFERED=1
ENV SQLBOT_HOME=/opt/sqlbot
ENV PYTHONPATH=${SQLBOT_HOME}/app
ENV PATH="${SQLBOT_HOME}/app/.venv/bin:$PATH"
RUN apt-get update && apt-get install -y --no-install-recommends \
curl \
&& apt-get clean \
&& rm -rf /var/lib/apt/lists/*
# Copy necessary files from builder
COPY --from=builder ${SQLBOT_HOME} ${SQLBOT_HOME}
WORKDIR ${SQLBOT_HOME}/app
# Add health check
HEALTHCHECK --interval=30s --timeout=10s --start-period=5s --retries=3 \
CMD curl -f http://localhost:8000 || exit 1
# Run with uvicorn
CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "8000", "--workers", "4", "--proxy-headers"]
-8
View File
@@ -1,8 +0,0 @@
# Python
__pycache__
app.egg-info
*.pyc
.mypy_cache
.coverage
htmlcov
.venv
-52
View File
@@ -1,52 +0,0 @@
FROM python:3.11-slim-bookworm
ENV PYTHONUNBUFFERED=1
ENV SQLBOT_HOME=/opt/sqlbot
ENV APP_HOME=$SQLBOT_HOME/app
ENV UI_HOME=$SQLBOT_HOME/frontend
RUN mkdir -p $APP_HOME
RUN mkdir -p $UI_HOME
WORKDIR $APP_HOME
ENV PYTHONPATH=$APP_HOME
# COPY ../frontend/dist $UI_HOME/dist
COPY ./dist $UI_HOME/dist
# Install uv
# Ref: https://docs.astral.sh/uv/guides/integration/docker/#installing-uv
COPY --from=ghcr.io/astral-sh/uv:0.5.11 /uv /uvx /bin/
# Place executables in the environment at the front of the path
# Ref: https://docs.astral.sh/uv/guides/integration/docker/#using-the-environment
ENV PATH="$APP_HOME/.venv/bin:$PATH"
# Compile bytecode
# Ref: https://docs.astral.sh/uv/guides/integration/docker/#compiling-bytecode
ENV UV_COMPILE_BYTECODE=1
# uv Cache
# Ref: https://docs.astral.sh/uv/guides/integration/docker/#caching
ENV UV_LINK_MODE=copy
# Install dependencies
# Ref: https://docs.astral.sh/uv/guides/integration/docker/#intermediate-layers
RUN test -f "./uv.lock" && \
--mount=type=cache,target=/root/.cache/uv \
--mount=type=bind,source=uv.lock,target=uv.lock \
--mount=type=bind,source=pyproject.toml,target=pyproject.toml \
uv sync --frozen --no-install-project || echo "uv.lock file not found, skipping intermediate-layers"
COPY . $APP_HOME
# Sync the project
# Ref: https://docs.astral.sh/uv/guides/integration/docker/#intermediate-layers
RUN --mount=type=cache,target=/root/.cache/uv \
uv sync
CMD ["fastapi", "run", "--workers", "4", "main.py"]
@@ -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 ###
View File
+1
View File
@@ -0,0 +1 @@
# todo
@@ -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
+1
View File
@@ -0,0 +1 @@
# todo
+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")
+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())
+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
-7
View File
@@ -1,7 +0,0 @@
from pydantic import BaseModel
class ChatQuestion(BaseModel):
question: str
chat_id: int
-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)
View File
+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)
@@ -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
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']
@@ -0,0 +1,6 @@
from apps.template.template import get_base_template
def get_sql_template():
template = get_base_template()
return template['template']['sql']
+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
+12
View File
@@ -1,4 +1,5 @@
import logging
from fastapi.concurrency import asynccontextmanager
from fastapi.responses import FileResponse
from fastapi.staticfiles import StaticFiles
import os
@@ -11,6 +12,16 @@ from apps.api import api_router
from apps.system.middleware.auth import TokenMiddleware
from common.core.config import settings
from common.core.response_middleware import ResponseMiddleware, exception_handler
from alembic.config import Config
from alembic import command
def run_migrations():
alembic_cfg = Config("alembic.ini")
command.upgrade(alembic_cfg, "head")
@asynccontextmanager
async def lifespan(app: FastAPI):
run_migrations()
yield
def custom_generate_unique_id(route: APIRoute) -> str:
tag = route.tags[0] if route.tags and len(route.tags) > 0 else ""
@@ -24,6 +35,7 @@ app = FastAPI(
title=settings.PROJECT_NAME,
openapi_url=f"{settings.API_V1_STR}/openapi.json",
generate_unique_id_function=custom_generate_unique_id,
lifespan=lifespan
)
# Set all CORS enabled origins
+3 -2
View File
@@ -32,10 +32,11 @@ 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.aliyun.com/pypi/simple"
url = "https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple"
default = true
[tool.uv]
+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}
+2 -2
View File
@@ -6,13 +6,13 @@ export interface SelectOption {
export const modelTypeOptions: SelectOption[] = [
{ label: 'OpenAI', value: 0 },
{ label: 'Anthropic', value: 1 },
/* { label: 'Anthropic', value: 1 },
{ label: 'Baidu', value: 2 },
{ label: 'iFLYTEK', value: 3 },
{ label: 'Zhipu AI', value: 4 },
{ label: 'MiniMax', value: 5 },
{ label: 'Tencent', value: 6 },
{ label: 'Other', value: 7 },
{ label: 'Other', value: 7 }, */
]
export const getModelTypeName = (value: any) => {
+81 -53
View File
@@ -38,7 +38,7 @@
<el-form-item label="Description">
<el-input v-model="form.description" :rows="2" type="textarea" />
</el-form-item>
<el-form-item label="Type">
<el-form-item label="Type" prop="type">
<el-select
v-model="form.type"
placeholder="Select Type"
@@ -72,32 +72,32 @@
</el-form-item>
</div>
<div v-else>
<el-form-item label="Host/Ip">
<el-input v-model="config.host" />
<el-form-item label="Host/Ip" prop="host">
<el-input v-model="form.host" />
</el-form-item>
<el-form-item label="Port">
<el-input v-model="config.port" />
<el-form-item label="Port" prop="port">
<el-input v-model="form.port" />
</el-form-item>
<el-form-item label="Username">
<el-input v-model="config.username" />
<el-input v-model="form.username" />
</el-form-item>
<el-form-item label="Password">
<el-input v-model="config.password" type="password" show-password />
<el-input v-model="form.password" type="password" show-password />
</el-form-item>
<el-form-item label="Database">
<el-input v-model="config.database" />
<el-form-item label="Database" prop="database">
<el-input v-model="form.database" />
</el-form-item>
<el-form-item label="Connect Mode" v-if="form.type === 'oracle'">
<el-radio-group v-model="config.mode">
<el-form-item label="Connect Mode" v-if="form.type === 'oracle'" prop="mode">
<el-radio-group v-model="form.mode">
<el-radio value="service_name">Service Name</el-radio>
<el-radio value="sid">SID</el-radio>
</el-radio-group>
</el-form-item>
<el-form-item label="Extra JDBC String">
<el-input v-model="config.extraJdbc" />
<el-input v-model="form.extraJdbc" />
</el-form-item>
<el-form-item label="Schema" v-if="haveSchema.includes(form.type)">
<el-input v-model="config.dbSchema" />
<el-form-item label="Schema" v-if="haveSchema.includes(form.type)" prop="dbSchema">
<el-input v-model="form.dbSchema" />
<el-button link type="primary" :icon="Plus" v-if="false"
>Get Schema</el-button
>
@@ -189,6 +189,24 @@ const rules = reactive<FormRules>({
{ required: true, message: "Please input name", trigger: "blur" },
{ min: 1, max: 50, message: "Length should be 1 to 50", trigger: "blur" },
],
type: [
{ required: true, message: "Please choose database type" ,trigger: "change"}
],
host: [
{ required: true, message: "Please input host" ,trigger: "blur"}
],
port: [
{ required: true, message: "Please input port" ,trigger: "blur"}
],
database: [
{ required: true, message: "Please input database" ,trigger: "blur"}
],
mode: [
{ required: true, message: "Please choose mode" ,trigger: "change"}
],
dbSchema: [
{ required: true, message: "Please input schema" ,trigger: "blur"}
]
});
const dialogVisible = ref<boolean>(false);
@@ -197,8 +215,6 @@ const form = ref<any>({
description: "",
type: "mysql",
configuration: "",
});
const config = ref<any>({
driver: "",
host: "",
port: 0,
@@ -234,16 +250,16 @@ const open = (item: any, editTable: boolean = false) => {
form.value.configuration = item.configuration;
if (item.configuration) {
const configuration = JSON.parse(decrypted(item.configuration));
config.value.host = configuration.host;
config.value.port = configuration.port;
config.value.username = configuration.username;
config.value.password = configuration.password;
config.value.database = configuration.database;
config.value.extraJdbc = configuration.extraJdbc;
config.value.dbSchema = configuration.dbSchema;
config.value.filename = configuration.filename;
config.value.sheets = configuration.sheets;
config.value.mode = configuration.mode;
form.value.host = configuration.host;
form.value.port = configuration.port;
form.value.username = configuration.username;
form.value.password = configuration.password;
form.value.database = configuration.database;
form.value.extraJdbc = configuration.extraJdbc;
form.value.dbSchema = configuration.dbSchema;
form.value.filename = configuration.filename;
form.value.sheets = configuration.sheets;
form.value.mode = configuration.mode;
}
if (editTable) {
@@ -258,11 +274,12 @@ const open = (item: any, editTable: boolean = false) => {
return ele.table_name;
});
if (item.type === "excel") {
tableList.value = config.value.sheets;
tableList.value = form.value.sheets;
} else {
tableListLoading.value = true;
const requestObj = buildConf()
datasourceApi
.getTablesByConf(form.value)
.getTablesByConf(requestObj)
.then((table) => {
tableList.value = table;
checkList.value = checkList.value.filter((ele: string) => {
@@ -290,8 +307,6 @@ const open = (item: any, editTable: boolean = false) => {
description: "",
type: "mysql",
configuration: "",
};
config.value = {
driver: "",
host: "",
port: 0,
@@ -303,7 +318,7 @@ const open = (item: any, editTable: boolean = false) => {
filename: "",
sheets: [],
mode: "service_name",
};
}
}
dialogVisible.value = true;
};
@@ -320,11 +335,11 @@ const save = async (formEl: FormInstance | undefined) => {
return { table_name: ele.tableName, table_comment: ele.tableComment };
});
buildConf();
const requestObj = buildConf();
if (form.value.id) {
if (!isEditTable.value) {
// only update datasource config info
datasourceApi.update(form.value).then((res) => {
datasourceApi.update(requestObj).then((res) => {
console.log(res);
close();
emit("refresh");
@@ -336,8 +351,8 @@ const save = async (formEl: FormInstance | undefined) => {
});
}
} else {
form.value.tables = list;
datasourceApi.add(form.value).then((res: any) => {
requestObj.tables = list;
datasourceApi.add(requestObj).then((res: any) => {
console.log(res);
close();
emit("refresh");
@@ -350,23 +365,36 @@ const save = async (formEl: FormInstance | undefined) => {
const buildConf = () => {
form.value.configuration = encrypted(
JSON.stringify({
host: config.value.host,
port: config.value.port,
username: config.value.username,
password: config.value.password,
database: config.value.database,
extraJdbc: config.value.extraJdbc,
dbSchema: config.value.dbSchema,
filename: config.value.filename,
sheets: config.value.sheets,
mode: config.value.mode,
host: form.value.host,
port: form.value.port,
username: form.value.username,
password: form.value.password,
database: form.value.database,
extraJdbc: form.value.extraJdbc,
dbSchema: form.value.dbSchema,
filename: form.value.filename,
sheets: form.value.sheets,
mode: form.value.mode,
})
);
)
const obj = JSON.parse(JSON.stringify(form.value))
delete obj.driver
delete obj.host
delete obj.port
delete obj.username
delete obj.password
delete obj.database
delete obj.extraJdbc
delete obj.dbSchema
delete obj.filename
delete obj.sheets
delete obj.mode
return obj
};
const check = () => {
buildConf();
datasourceApi.check(form.value).then((res: any) => {
const requestObj = buildConf();
datasourceApi.check(requestObj).then((res: any) => {
if (res) {
ElMessage({
message: "Connect success",
@@ -394,12 +422,12 @@ const next = async (formEl: FormInstance | undefined) => {
}
} else {
// check status if success do next
buildConf();
datasourceApi.check(form.value).then((res: boolean) => {
const requestObj = buildConf();
datasourceApi.check(requestObj).then((res: boolean) => {
if (res) {
active.value++;
// request tables
datasourceApi.getTablesByConf(form.value).then((res) => {
datasourceApi.getTablesByConf(requestObj).then((res) => {
tableList.value = res;
});
} else {
@@ -428,8 +456,8 @@ const beforeUpload = (rawFile: any) => {
};
const onSuccess = (response: any) => {
config.value.filename = response.data.filename;
config.value.sheets = response.data.sheets;
form.value.filename = response.data.filename;
form.value.sheets = response.data.sheets;
tableList.value = response.data.sheets;
excelUploadSuccess.value = true;
};
+6 -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" />
@@ -191,6 +191,7 @@ const handleSearch = (e: any) => {
const editHandler = (id: any) => {
modelApi.query(id).then((res: any) => {
state.form = res
state.form.temperature = state.form.temperature * 100
dialogFormVisible.value = true
})
}
@@ -259,6 +260,7 @@ const addTerm = () => {
...state.form,
}
delete param.id
param.temperature = param.temperature / 100
modelApi.add(param).then(() => {
dialogFormVisible.value = false
search()
@@ -269,6 +271,8 @@ const addTerm = () => {
})
}
const editTerm = () => {
const param = state.form
param.temperature = param.temperature / 100
modelApi.edit(state.form).then((res: any) => {
console.log('edit term', res)
dialogFormVisible.value = false