From 054296741cd52545f78b9afe95792423501e16b2 Mon Sep 17 00:00:00 2001 From: ulleo Date: Fri, 30 May 2025 17:41:45 +0800 Subject: [PATCH] feat: generate sql with history --- .../versions/1f077c30e476_009_modify_chat.py | 64 ++++++ backend/apps/ai_model/__init__.py | 0 backend/apps/ai_model/llm.py | 1 + .../model_factory.py} | 18 +- backend/apps/ai_model/openai/__init__.py | 0 backend/apps/ai_model/openai/llm.py | 1 + backend/apps/chat/api/chat.py | 148 +++++------- backend/apps/chat/curd/chat.py | 39 +++- backend/apps/chat/models/chat_model.py | 41 +++- backend/apps/chat/schemas/chat_schema.py | 7 - backend/apps/chat/schemas/llm.py | 197 ---------------- backend/apps/chat/task/__init__.py | 0 backend/apps/chat/task/llm.py | 210 ++++++++++++++++++ backend/apps/datasource/crud/datasource.py | 29 +++ backend/apps/template/__init__.py | 0 .../template/generate_analysis/__init__.py | 0 .../apps/template/generate_chart/__init__.py | 0 .../apps/template/generate_chart/generator.py | 6 + .../apps/template/generate_sql/__init__.py | 0 .../apps/template/generate_sql/generator.py | 6 + backend/apps/template/template.py | 15 ++ backend/pyproject.toml | 3 +- backend/template.yaml | 52 +++++ frontend/components.d.ts | 51 ----- frontend/src/views/system/model/index.vue | 4 +- 25 files changed, 526 insertions(+), 366 deletions(-) create mode 100644 backend/alembic/versions/1f077c30e476_009_modify_chat.py create mode 100644 backend/apps/ai_model/__init__.py create mode 100644 backend/apps/ai_model/llm.py rename backend/apps/{chat/schemas/chat_base_schema.py => ai_model/model_factory.py} (86%) create mode 100644 backend/apps/ai_model/openai/__init__.py create mode 100644 backend/apps/ai_model/openai/llm.py delete mode 100644 backend/apps/chat/schemas/chat_schema.py delete mode 100644 backend/apps/chat/schemas/llm.py create mode 100644 backend/apps/chat/task/__init__.py create mode 100644 backend/apps/chat/task/llm.py create mode 100644 backend/apps/template/__init__.py create mode 100644 backend/apps/template/generate_analysis/__init__.py create mode 100644 backend/apps/template/generate_chart/__init__.py create mode 100644 backend/apps/template/generate_chart/generator.py create mode 100644 backend/apps/template/generate_sql/__init__.py create mode 100644 backend/apps/template/generate_sql/generator.py create mode 100644 backend/apps/template/template.py create mode 100644 backend/template.yaml delete mode 100644 frontend/components.d.ts 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 @@ - +