diff --git a/backend/alembic/versions/1f077c30e476_009_modify_chat.py b/backend/alembic/versions/1f077c30e476_009_modify_chat.py
new file mode 100644
index 00000000..8db7b9d9
--- /dev/null
+++ b/backend/alembic/versions/1f077c30e476_009_modify_chat.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 ###
diff --git a/backend/apps/ai_model/__init__.py b/backend/apps/ai_model/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/backend/apps/ai_model/llm.py b/backend/apps/ai_model/llm.py
new file mode 100644
index 00000000..99f10893
--- /dev/null
+++ b/backend/apps/ai_model/llm.py
@@ -0,0 +1 @@
+# todo
\ No newline at end of file
diff --git a/backend/apps/chat/schemas/chat_base_schema.py b/backend/apps/ai_model/model_factory.py
similarity index 86%
rename from backend/apps/chat/schemas/chat_base_schema.py
rename to backend/apps/ai_model/model_factory.py
index 691bd010..b9bcc908 100644
--- a/backend/apps/chat/schemas/chat_base_schema.py
+++ b/backend/apps/ai_model/model_factory.py
@@ -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
\ No newline at end of file
+ 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
\ No newline at end of file
diff --git a/backend/apps/ai_model/openai/__init__.py b/backend/apps/ai_model/openai/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/backend/apps/ai_model/openai/llm.py b/backend/apps/ai_model/openai/llm.py
new file mode 100644
index 00000000..044a4824
--- /dev/null
+++ b/backend/apps/ai_model/openai/llm.py
@@ -0,0 +1 @@
+# todo
diff --git a/backend/apps/chat/api/chat.py b/backend/apps/chat/api/chat.py
index 4f4a7087..0b26ea14 100644
--- a/backend/apps/chat/api/chat.py
+++ b/backend/apps/chat/api/chat.py
@@ -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")
diff --git a/backend/apps/chat/curd/chat.py b/backend/apps/chat/curd/chat.py
index 5f6c2658..5ba0437e 100644
--- a/backend/apps/chat/curd/chat.py
+++ b/backend/apps/chat/curd/chat.py
@@ -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())
diff --git a/backend/apps/chat/models/chat_model.py b/backend/apps/chat/models/chat_model.py
index 1c276caa..b65af0ad 100644
--- a/backend/apps/chat/models/chat_model.py
+++ b/backend/apps/chat/models/chat_model.py
@@ -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
diff --git a/backend/apps/chat/schemas/chat_schema.py b/backend/apps/chat/schemas/chat_schema.py
deleted file mode 100644
index 57b4dee7..00000000
--- a/backend/apps/chat/schemas/chat_schema.py
+++ /dev/null
@@ -1,7 +0,0 @@
-
-from pydantic import BaseModel
-
-
-class ChatQuestion(BaseModel):
- question: str
- chat_id: int
\ No newline at end of file
diff --git a/backend/apps/chat/schemas/llm.py b/backend/apps/chat/schemas/llm.py
deleted file mode 100644
index fb75926e..00000000
--- a/backend/apps/chat/schemas/llm.py
+++ /dev/null
@@ -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)
\ No newline at end of file
diff --git a/backend/apps/chat/task/__init__.py b/backend/apps/chat/task/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/backend/apps/chat/task/llm.py b/backend/apps/chat/task/llm.py
new file mode 100644
index 00000000..da62eab9
--- /dev/null
+++ b/backend/apps/chat/task/llm.py
@@ -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)
diff --git a/backend/apps/datasource/crud/datasource.py b/backend/apps/datasource/crud/datasource.py
index db2fc95a..1e1b597f 100644
--- a/backend/apps/datasource/crud/datasource.py
+++ b/backend/apps/datasource/crud/datasource.py
@@ -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
diff --git a/backend/apps/template/__init__.py b/backend/apps/template/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/backend/apps/template/generate_analysis/__init__.py b/backend/apps/template/generate_analysis/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/backend/apps/template/generate_chart/__init__.py b/backend/apps/template/generate_chart/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/backend/apps/template/generate_chart/generator.py b/backend/apps/template/generate_chart/generator.py
new file mode 100644
index 00000000..0962cba7
--- /dev/null
+++ b/backend/apps/template/generate_chart/generator.py
@@ -0,0 +1,6 @@
+from apps.template.template import get_base_template
+
+
+def get_chart_template():
+ template = get_base_template()
+ return template['template']['chart']
diff --git a/backend/apps/template/generate_sql/__init__.py b/backend/apps/template/generate_sql/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/backend/apps/template/generate_sql/generator.py b/backend/apps/template/generate_sql/generator.py
new file mode 100644
index 00000000..5d08cc9e
--- /dev/null
+++ b/backend/apps/template/generate_sql/generator.py
@@ -0,0 +1,6 @@
+from apps.template.template import get_base_template
+
+
+def get_sql_template():
+ template = get_base_template()
+ return template['template']['sql']
diff --git a/backend/apps/template/template.py b/backend/apps/template/template.py
new file mode 100644
index 00000000..7dd5b1cc
--- /dev/null
+++ b/backend/apps/template/template.py
@@ -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
diff --git a/backend/pyproject.toml b/backend/pyproject.toml
index f1f911ac..9d4acbc1 100644
--- a/backend/pyproject.toml
+++ b/backend/pyproject.toml
@@ -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"
diff --git a/backend/template.yaml b/backend/template.yaml
new file mode 100644
index 00000000..60dee447
--- /dev/null
+++ b/backend/template.yaml
@@ -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}
\ No newline at end of file
diff --git a/frontend/components.d.ts b/frontend/components.d.ts
deleted file mode 100644
index ccfedf7f..00000000
--- a/frontend/components.d.ts
+++ /dev/null
@@ -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']
- }
-}
diff --git a/frontend/src/views/system/model/index.vue b/frontend/src/views/system/model/index.vue
index 72789b8e..02982c2a 100644
--- a/frontend/src/views/system/model/index.vue
+++ b/frontend/src/views/system/model/index.vue
@@ -119,13 +119,13 @@
-
+