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

Merge branch 'main' of https://github.com/dataease/SQLBot

Browse Source
This commit is contained in:
dataeaseShu
2025-09-03 18:01:58 +08:00
parent c254f462dd 3139b62d76
commit 3a64c8718b
4 changed files with 59 additions and 39 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Dockerfile
+6
View File
@@ -54,6 +54,12 @@ ENV SQLBOT_HOME=/opt/sqlbot
ENV PYTHONPATH=${SQLBOT_HOME}/app
ENV PATH="${SQLBOT_HOME}/app/.venv/bin:$PATH"
ENV SQLBOT_DB_HOST=localhost
ENV SQLBOT_DB_PORT=5432
ENV SQLBOT_DB_DB=sqlbot
ENV SQLBOT_DB_USER=root
ENV SQLBOT_DB_PASSWORD=Password123@pg
# Copy necessary files from builder
COPY start.sh /opt/sqlbot/app/start.sh
COPY g2-ssr/*.ttf /usr/share/fonts/truetype/liberation/
backend/apps/chat/api/chat.py
+35 -38
View File
@@ -3,6 +3,7 @@ import io
import traceback
import numpy as np
import orjson
import pandas as pd
from fastapi import APIRouter, HTTPException
from fastapi.responses import StreamingResponse
@@ -147,10 +148,11 @@ async def stream_sql(session: SessionDep, current_user: CurrentUser, request_que
llm_service.run_task_async()
except Exception as e:
traceback.print_exc()
raise HTTPException(
status_code=500,
detail=str(e)
)
def _err(_e: Exception):
yield 'data:' + orjson.dumps({'content': str(_e), 'type': 'error'}).decode() + '\n\n'
return StreamingResponse(_err(e), media_type="text/event-stream")
return StreamingResponse(llm_service.await_result(), media_type="text/event-stream")
@@ -158,45 +160,40 @@ async def stream_sql(session: SessionDep, current_user: CurrentUser, request_que
@router.post("/record/{chat_record_id}/{action_type}")
async def analysis_or_predict(session: SessionDep, current_user: CurrentUser, chat_record_id: int, action_type: str,
current_assistant: CurrentAssistant):
if action_type != 'analysis' and action_type != 'predict':
raise HTTPException(
status_code=404,
detail="Not Found"
)
record: ChatRecord | None = None
stmt = select(ChatRecord.id, ChatRecord.question, ChatRecord.chat_id, ChatRecord.datasource, ChatRecord.engine_type,
ChatRecord.ai_modal_id, ChatRecord.create_by, ChatRecord.chart, ChatRecord.data).where(
and_(ChatRecord.id == chat_record_id))
result = session.execute(stmt)
for r in result:
record = ChatRecord(id=r.id, question=r.question, chat_id=r.chat_id, datasource=r.datasource,
engine_type=r.engine_type, ai_modal_id=r.ai_modal_id, create_by=r.create_by, chart=r.chart,
data=r.data)
if not record:
raise HTTPException(
status_code=400,
detail=f"Chat record with id {chat_record_id} not found"
)
if not record.chart:
raise HTTPException(
status_code=500,
detail=f"Chat record with id {chat_record_id} has not generated chart, do not support to analyze it"
)
request_question = ChatQuestion(chat_id=record.chat_id, question=record.question)
try:
if action_type != 'analysis' and action_type != 'predict':
raise Exception(f"Type {action_type} Not Found")
record: ChatRecord | None = None
stmt = select(ChatRecord.id, ChatRecord.question, ChatRecord.chat_id, ChatRecord.datasource,
ChatRecord.engine_type,
ChatRecord.ai_modal_id, ChatRecord.create_by, ChatRecord.chart, ChatRecord.data).where(
and_(ChatRecord.id == chat_record_id))
result = session.execute(stmt)
for r in result:
record = ChatRecord(id=r.id, question=r.question, chat_id=r.chat_id, datasource=r.datasource,
engine_type=r.engine_type, ai_modal_id=r.ai_modal_id, create_by=r.create_by,
chart=r.chart,
data=r.data)
if not record:
raise Exception(f"Chat record with id {chat_record_id} not found")
if not record.chart:
raise Exception(
f"Chat record with id {chat_record_id} has not generated chart, do not support to analyze it")
request_question = ChatQuestion(chat_id=record.chat_id, question=record.question)
llm_service = await LLMService.create(current_user, request_question, current_assistant)
llm_service.run_analysis_or_predict_task_async(action_type, record)
except Exception as e:
traceback.print_exc()
raise HTTPException(
status_code=500,
detail=str(e)
)
def _err(_e: Exception):
yield 'data:' + orjson.dumps({'content': str(_e), 'type': 'error'}).decode() + '\n\n'
return StreamingResponse(_err(e), media_type="text/event-stream")
return StreamingResponse(llm_service.await_result(), media_type="text/event-stream")
backend/apps/system/api/assistant.py
+7 -1
View File
@@ -64,9 +64,15 @@ async def validator(session: SessionDep, id: int, virtual: Optional[int] = Query
if not db_model:
return AssistantValidator()
db_model = AssistantModel.model_validate(db_model)
assistant_oid = 1
if(db_model.type == 0):
configuration = db_model.configuration
config_obj = json.loads(configuration) if configuration else {}
assistant_oid = config_obj.get('oid', 1)
access_token_expires = timedelta(minutes=settings.ACCESS_TOKEN_EXPIRE_MINUTES)
assistantDict = {
"id": virtual, "account": 'sqlbot-inner-assistant', "oid": 1, "assistant_id": id
"id": virtual, "account": 'sqlbot-inner-assistant', "oid": assistant_oid, "assistant_id": id
}
access_token = create_access_token(
assistantDict, expires_delta=access_token_expires
backend/apps/system/middleware/auth.py
+11
View File
@@ -1,5 +1,6 @@
import base64
import json
from typing import Optional
from fastapi import Request
from fastapi.responses import JSONResponse
@@ -112,6 +113,16 @@ class TokenMiddleware(BaseHTTPMiddleware):
assistant_info = await get_assistant_info(session=session, assistant_id=payload['assistant_id'])
assistant_info = AssistantModel.model_validate(assistant_info)
assistant_info = AssistantHeader.model_validate(assistant_info.model_dump(exclude_unset=True))
if assistant_info and assistant_info.type == 0:
if payload['oid']:
session_user.oid = payload['oid']
else:
assistant_oid = 1
configuration = assistant_info.configuration
config_obj = json.loads(configuration) if configuration else {}
assistant_oid = config_obj.get('oid', 1)
session_user.oid = assistant_oid
return True, session_user, assistant_info
except Exception as e:
SQLBotLogUtil.exception(f"Assistant validation error: {str(e)}")