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

perf: Assistant dynamic ds

Browse Source
This commit is contained in:
fit2cloud-chenyw
2025-07-22 10:16:58 +08:00
parent 517de2fcf6
commit 8dbac2a101
6 changed files with 55 additions and 41 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/apps/chat/task/llm.py
+1 -1
View File
@@ -748,7 +748,7 @@ def run_task(llm_service: LLMService, in_chat: bool = True):
yield orjson.dumps({'id': llm_service.ds.id, 'datasource_name': llm_service.ds.name,
'engine_type': llm_service.ds.type_name or llm_service.ds.type, 'type': 'datasource'}).decode() + '\n\n'
llm_service.chat_question.db_schema = llm_service.out_ds_instance.get_db_schema() if llm_service.out_ds_instance else get_table_schema(session=llm_service.session, current_user=llm_service.current_user, ds=llm_service.ds)
llm_service.chat_question.db_schema = llm_service.out_ds_instance.get_db_schema(llm_service.ds.id) if llm_service.out_ds_instance else get_table_schema(session=llm_service.session, current_user=llm_service.current_user, ds=llm_service.ds)
# generate sql
sql_res = llm_service.generate_sql()
backend/apps/db/db.py
+14 -8
View File
@@ -11,26 +11,32 @@ from apps.datasource.utils.utils import aes_decrypt
from apps.db.engine import get_engine_config
from decimal import Decimal
from apps.system.crud.assistant import get_ds_engine
from apps.system.schemas.system_schema import AssistantOutDsSchema
def get_uri(ds: CoreDatasource):
def get_uri(ds: CoreDatasource) -> str:
conf = DatasourceConf(**json.loads(aes_decrypt(ds.configuration))) if ds.type != "excel" else get_engine_config()
return get_uri_from_config(ds.type, conf)
def get_uri_from_config(type: str,conf: DatasourceConf) -> str:
db_url: str
if ds.type == "mysql":
if type == "mysql":
if conf.extraJdbc is not None and conf.extraJdbc != '':
db_url = f"mysql+pymysql://{urllib.parse.quote(conf.username)}:{urllib.parse.quote(conf.password)}@{conf.host}:{conf.port}/{urllib.parse.quote(conf.database)}?{conf.extraJdbc}"
else:
db_url = f"mysql+pymysql://{urllib.parse.quote(conf.username)}:{urllib.parse.quote(conf.password)}@{conf.host}:{conf.port}/{urllib.parse.quote(conf.database)}"
elif ds.type == "sqlServer":
elif type == "sqlServer":
if conf.extraJdbc is not None and conf.extraJdbc != '':
db_url = f"mssql+pymssql://{urllib.parse.quote(conf.username)}:{urllib.parse.quote(conf.password)}@{conf.host}:{conf.port}/{urllib.parse.quote(conf.database)}?{conf.extraJdbc}"
else:
db_url = f"mssql+pymssql://{urllib.parse.quote(conf.username)}:{urllib.parse.quote(conf.password)}@{conf.host}:{conf.port}/{urllib.parse.quote(conf.database)}"
elif ds.type == "pg" or ds.type == "excel":
elif type == "pg" or type == "excel":
if conf.extraJdbc is not None and conf.extraJdbc != '':
db_url = f"postgresql+psycopg2://{urllib.parse.quote(conf.username)}:{urllib.parse.quote(conf.password)}@{conf.host}:{conf.port}/{urllib.parse.quote(conf.database)}?{conf.extraJdbc}"
else:
db_url = f"postgresql+psycopg2://{urllib.parse.quote(conf.username)}:{urllib.parse.quote(conf.password)}@{conf.host}:{conf.port}/{urllib.parse.quote(conf.database)}"
elif ds.type == "oracle":
elif type == "oracle":
if conf.mode == "service_name":
if conf.extraJdbc is not None and conf.extraJdbc != '':
db_url = f"oracle+oracledb://{urllib.parse.quote(conf.username)}:{urllib.parse.quote(conf.password)}@{conf.host}:{conf.port}?service_name={urllib.parse.quote(conf.database)}&{conf.extraJdbc}"
@@ -62,8 +68,8 @@ def get_engine(ds: CoreDatasource) -> Engine:
return engine
def get_session(ds: CoreDatasource):
engine = get_engine(ds)
def get_session(ds: CoreDatasource | AssistantOutDsSchema):
engine = get_engine(ds) if isinstance(ds, CoreDatasource) else get_ds_engine(ds)
session_maker = sessionmaker(bind=engine)
session = session_maker()
return session
@@ -238,7 +244,7 @@ def get_fields(ds: CoreDatasource, table_name: str = None):
session.close()
def exec_sql(ds: CoreDatasource, sql: str):
def exec_sql(ds: CoreDatasource | AssistantOutDsSchema, sql: str):
session = get_session(ds)
result = session.execute(text(sql))
try:
backend/apps/system/crud/assistant.py
+27 -10
View File
@@ -4,8 +4,10 @@ import json
from typing import Optional
from fastapi import FastAPI
import requests
from sqlalchemy import Engine, create_engine
from sqlmodel import Session, select
from apps.datasource.models.datasource import CoreDatasource
from apps.datasource.models.datasource import CoreDatasource, DatasourceConf
from apps.system.models.system_model import AssistantModel
from apps.system.schemas.auth import CacheName, CacheNamespace
from apps.system.schemas.system_schema import AssistantOutDsSchema, UserInfoDTO
@@ -92,7 +94,7 @@ class AssistantOutDs:
self.get_ds_from_api(certificate)
#@cache(namespace=CacheNamespace.EMBEDDED_INFO, cacheName=CacheName.ASSISTANT_DS, keyExpression="current_user.id")
async def get_ds_from_api(self, certificate: Optional[str] = None):
def get_ds_from_api(self, certificate: Optional[str] = None):
config: dict[any] = json.loads(self.assistant.configuration)
endpoint: str = config['endpoint']
certificateList: list[any] = json.loads(certificate)
@@ -106,7 +108,7 @@ class AssistantOutDs:
res = requests.get(url=endpoint, headers=header, cookies=cookies, timeout=10)
if res.status_code == 200:
result_json: dict[any] = json.loads(res.json())
result_json: dict[any] = json.loads(res.text)
if result_json.get('code') == 0:
temp_list = result_json.get('data', [])
self.ds_list = [
@@ -144,20 +146,35 @@ class AssistantOutDs:
def get_ds(self, ds_id: int):
if self.ds_list:
for ds in self.ds_list:
if ds['id'] == ds_id:
if ds.id == ds_id:
return ds
else:
raise Exception("Datasource list is not found.")
raise Exception(f"Datasource with id {ds_id} not found.")
def get_ds_engine(self, ds_id: int):
ds = self.get_ds(ds_id)
ds_type = ds.get('type') if ds else None
if not ds_type:
raise Exception(f"Datasource with id {ds_id} not found or type is not defined.")
return ds_type
class AssistantOutDsFactory:
@staticmethod
def get_instance(assistant: AssistantModel, certificate: Optional[str] = None) -> AssistantOutDs:
return AssistantOutDs(assistant, certificate)
def get_ds_engine(ds: AssistantOutDsSchema) -> Engine:
timeout: int = 30
connect_args = {"connect_timeout": timeout}
conf = DatasourceConf(
host=ds.host,
port=ds.port,
username=ds.user,
password=ds.password,
database=ds.dataBase,
driver='',
extraJdbc=ds.extraParams,
dbSchema=ds.schema or ''
)
from apps.db.db import get_uri_from_config
uri = get_uri_from_config(ds.type, conf)
if ds.type == "pg" and ds.schema:
connect_args.update({"options": f"-c search_path={ds.schema}"})
engine = create_engine(uri, connect_args=connect_args, pool_timeout=timeout, pool_size=20, max_overflow=10)
return engine
backend/apps/system/schemas/system_schema.py
+3 -18
View File
@@ -125,29 +125,14 @@ class AssistantOutDsBase(BaseModel):
comment: Optional[str] = None
description: Optional[str] = None
def __init__(self, id: Optional[int] = None, name: str = '', type: Optional[str] = None,
type_name: Optional[str] = None, comment: Optional[str] = None):
super().__init__(id=id, name=name, type=type, type_name=type_name, comment=comment)
class AssistantOutDsSchema(AssistantOutDsBase):
host: Optional[str] = None
port: Optional[int] = None
dataBase: Optional[str] = None
user: Optional[str] = None
password: Optional[str] = None
schema: Optional[str] = None
extraParams: Optional[str] = None
tables: Optional[list[AssistantTableSchema]] = None
def __init__(self, id: int, name: str, comment: Optional[str] = None, type: Optional[str] = None,
type_name: Optional[str] = None, host: Optional[str] = None, port: Optional[int] = None,
username: Optional[str] = None, password: Optional[str] = None, schema: Optional[str] = None):
self.id = id
self.name = name
self.comment = comment
self.type = type
self.type_name = type_name
self.host = host
self.port = port
self.username = username
self.password = password
self.schema = schema
self.description = comment
backend/common/core/deps.py
+2 -2
View File
@@ -19,8 +19,8 @@ async def get_current_user(request: Request) -> UserInfoDTO:
CurrentUser = Annotated[UserInfoDTO, Depends(get_current_user)]
async def get_current_assistant(request: Request) -> AssistantModel:
return request.state.assistant
async def get_current_assistant(request: Request) -> AssistantModel | None:
return request.state.assistant if hasattr(request.state, "assistant") else None
CurrentAssistant = Annotated[AssistantModel, Depends(get_current_assistant)]
frontend/src/views/chat/answer/ChartAnswer.vue
+8 -2
View File
@@ -1,7 +1,9 @@
<script setup lang="ts">
import BaseAnswer from './BaseAnswer.vue'
import { Chat, chatApi, ChatInfo, type ChatMessage, ChatRecord, questionApi } from '@/api/chat.ts'
import { useAssistantStore } from '@/stores/assistant'
import { computed, nextTick, ref } from 'vue'
const assistantStore = useAssistantStore()
const props = withDefaults(
defineProps<{
chatList?: Array<ChatInfo>
@@ -96,11 +98,15 @@ const sendMessage = async () => {
try {
const controller: AbortController = new AbortController()
const response = await questionApi.add({
const param = {
question: currentRecord.question,
chat_id: _currentChatId.value,
assistant_certificate: assistantStore.getCertificate,
controller,
})
}
console.log(assistantStore.getCertificate)
console.log(param)
const response = await questionApi.add(param)
const reader = response.body.getReader()
const decoder = new TextDecoder()