mirror of
https://gh-proxy.org/https://github.com/klinsmaya/SQLBot.git
synced 2026-11-08 09:17:06 +08:00
211 lines
8.8 KiB
Python
211 lines
8.8 KiB
Python
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)
|