mirror of
https://gh-proxy.org/https://github.com/klinsmaya/SQLBot.git
synced 2026-11-08 01:06:59 +08:00
Merge branch 'main' of https://github.com/dataease/SQLBot
This commit is contained in:
dataeaseShu
@@ -14,6 +14,12 @@ FIRST_SUPERUSER_PASSWORD=123456 # Change this to your pwd
|
||||
|
||||
TOKEN_KEY="X-SQLBOT-TOKEN"
|
||||
DEFAULT_PWD="SQLBot@123456"
|
||||
ASSISTANT_TOKEN_KEY="X-SQLBOT-ASSISTANT-TOKEN"
|
||||
|
||||
LOG_LEVEL="INFO" # DEBUG, INFO, WARNING, ERROR
|
||||
LOG_DIR="/opt/sqlbot/logs"
|
||||
LOG_FORMAT="%(asctime)s - %(name)s - %(levelname)s:%(lineno)d - %(message)s"
|
||||
SQL_DEBUG=False
|
||||
|
||||
# Postgres
|
||||
POSTGRES_SERVER=localhost
|
||||
|
||||
@@ -5,7 +5,6 @@ from os.path import abspath, dirname
|
||||
sys.path.insert(0, dirname(dirname(abspath(__file__))))
|
||||
|
||||
import os
|
||||
from logging.config import fileConfig
|
||||
|
||||
from alembic import context
|
||||
from sqlalchemy import engine_from_config, pool
|
||||
@@ -16,7 +15,6 @@ config = context.config
|
||||
|
||||
# Interpret the config file for Python logging.
|
||||
# This line sets up loggers basically.
|
||||
fileConfig(config.config_file_name)
|
||||
|
||||
# add your model's MetaData object here
|
||||
# for 'autogenerate' support
|
||||
|
||||
@@ -1,6 +1,5 @@
|
||||
import json
|
||||
import sqlparse
|
||||
import logging
|
||||
import traceback
|
||||
import warnings
|
||||
from typing import Any, List, Optional, Union, Dict
|
||||
@@ -35,7 +34,7 @@ from apps.system.crud.assistant import AssistantOutDs, AssistantOutDsFactory, ge
|
||||
from apps.system.schemas.system_schema import AssistantOutDsSchema
|
||||
from common.core.config import settings
|
||||
from common.core.deps import CurrentAssistant, SessionDep, CurrentUser
|
||||
from common.utils.utils import extract_nested_json
|
||||
from common.utils.utils import SQLBotLogUtil, extract_nested_json
|
||||
|
||||
warnings.filterwarnings("ignore")
|
||||
|
||||
@@ -223,7 +222,7 @@ class LLMService:
|
||||
res = self.llm.stream(analysis_msg)
|
||||
token_usage = {}
|
||||
for chunk in res:
|
||||
print(chunk)
|
||||
SQLBotLogUtil.info(chunk)
|
||||
reasoning_content_chunk = ''
|
||||
if 'reasoning_content' in chunk.additional_kwargs:
|
||||
reasoning_content_chunk = chunk.additional_kwargs.get('reasoning_content', '')
|
||||
@@ -273,7 +272,7 @@ class LLMService:
|
||||
res = self.llm.stream(predict_msg)
|
||||
token_usage = {}
|
||||
for chunk in res:
|
||||
print(chunk)
|
||||
SQLBotLogUtil.info(chunk)
|
||||
reasoning_content_chunk = ''
|
||||
if 'reasoning_content' in chunk.additional_kwargs:
|
||||
reasoning_content_chunk = chunk.additional_kwargs.get('reasoning_content', '')
|
||||
@@ -323,7 +322,7 @@ class LLMService:
|
||||
token_usage = {}
|
||||
res = self.llm.stream(guess_msg)
|
||||
for chunk in res:
|
||||
print(chunk)
|
||||
SQLBotLogUtil.info(chunk)
|
||||
reasoning_content_chunk = ''
|
||||
if 'reasoning_content' in chunk.additional_kwargs:
|
||||
reasoning_content_chunk = chunk.additional_kwargs.get('reasoning_content', '')
|
||||
@@ -390,7 +389,7 @@ class LLMService:
|
||||
token_usage = {}
|
||||
res = self.llm.stream(datasource_msg)
|
||||
for chunk in res:
|
||||
print(chunk)
|
||||
SQLBotLogUtil.info(chunk)
|
||||
reasoning_content_chunk = ''
|
||||
if 'reasoning_content' in chunk.additional_kwargs:
|
||||
reasoning_content_chunk = chunk.additional_kwargs.get('reasoning_content', '')
|
||||
@@ -474,7 +473,7 @@ class LLMService:
|
||||
token_usage = {}
|
||||
res = self.llm.stream(self.sql_message)
|
||||
for chunk in res:
|
||||
print(chunk)
|
||||
SQLBotLogUtil.info(chunk)
|
||||
reasoning_content_chunk = ''
|
||||
if 'reasoning_content' in chunk.additional_kwargs:
|
||||
reasoning_content_chunk = chunk.additional_kwargs.get('reasoning_content', '')
|
||||
@@ -543,7 +542,7 @@ class LLMService:
|
||||
res = self.llm.stream(msg)
|
||||
token_usage = {}
|
||||
for chunk in res:
|
||||
print(chunk)
|
||||
SQLBotLogUtil.info(chunk)
|
||||
reasoning_content_chunk = ''
|
||||
if 'reasoning_content' in chunk.additional_kwargs:
|
||||
reasoning_content_chunk = chunk.additional_kwargs.get('reasoning_content', '')
|
||||
@@ -567,7 +566,7 @@ class LLMService:
|
||||
# 'content': msg.content} for msg
|
||||
# in
|
||||
# analysis_msg]).decode())
|
||||
print(full_filter_text)
|
||||
SQLBotLogUtil.info(full_filter_text)
|
||||
return full_filter_text
|
||||
|
||||
def generate_chart(self):
|
||||
@@ -582,7 +581,7 @@ class LLMService:
|
||||
token_usage = {}
|
||||
res = self.llm.stream(self.chart_message)
|
||||
for chunk in res:
|
||||
print(chunk)
|
||||
SQLBotLogUtil.info(chunk)
|
||||
reasoning_content_chunk = ''
|
||||
if 'reasoning_content' in chunk.additional_kwargs:
|
||||
reasoning_content_chunk = chunk.additional_kwargs.get('reasoning_content', '')
|
||||
@@ -691,7 +690,7 @@ class LLMService:
|
||||
Returns:
|
||||
Query results
|
||||
"""
|
||||
print(f"Executing SQL on ds_id {self.ds.id}: {sql}")
|
||||
SQLBotLogUtil.info(f"Executing SQL on ds_id {self.ds.id}: {sql}")
|
||||
return exec_sql(self.ds, sql)
|
||||
|
||||
|
||||
@@ -717,7 +716,7 @@ def execute_sql_with_db(db: SQLDatabase, sql: str) -> str:
|
||||
|
||||
except Exception as e:
|
||||
error_msg = f"SQL execution failed: {str(e)}"
|
||||
logging.error(error_msg)
|
||||
SQLBotLogUtil.exception(error_msg)
|
||||
raise RuntimeError(error_msg)
|
||||
|
||||
|
||||
@@ -741,7 +740,7 @@ def run_task(llm_service: LLMService, in_chat: bool = True):
|
||||
ds_res = llm_service.select_datasource()
|
||||
|
||||
for chunk in ds_res:
|
||||
print(chunk)
|
||||
SQLBotLogUtil.info(chunk)
|
||||
if in_chat:
|
||||
yield orjson.dumps(
|
||||
{'content': chunk.get('content'), 'reasoning_content': chunk.get('reasoning_content'),
|
||||
@@ -765,7 +764,7 @@ def run_task(llm_service: LLMService, in_chat: bool = True):
|
||||
yield orjson.dumps({'type': 'info', 'msg': 'sql generated'}).decode() + '\n\n'
|
||||
|
||||
# filter sql
|
||||
print(full_sql_text)
|
||||
SQLBotLogUtil.info(full_sql_text)
|
||||
|
||||
# todo row permission
|
||||
sql_json_str = extract_nested_json(full_sql_text)
|
||||
@@ -785,11 +784,11 @@ def run_task(llm_service: LLMService, in_chat: bool = True):
|
||||
raise Exception("SQL query is empty")
|
||||
|
||||
sql_result = llm_service.generate_filter(data.get('sql'), data.get('tables')) # maybe no sql and tables
|
||||
print(sql_result)
|
||||
SQLBotLogUtil.info(sql_result)
|
||||
sql = llm_service.check_save_sql(res=sql_result)
|
||||
# sql = llm_service.check_save_sql(res=full_sql_text)
|
||||
|
||||
print(sql)
|
||||
SQLBotLogUtil.info(sql)
|
||||
format_sql = sqlparse.format(sql, reindent=True)
|
||||
if in_chat:
|
||||
yield orjson.dumps({'content': format_sql, 'type': 'sql'}).decode() + '\n\n'
|
||||
@@ -815,9 +814,9 @@ def run_task(llm_service: LLMService, in_chat: bool = True):
|
||||
yield orjson.dumps({'type': 'info', 'msg': 'chart generated'}).decode() + '\n\n'
|
||||
|
||||
# filter chart
|
||||
print(full_chart_text)
|
||||
SQLBotLogUtil.info(full_chart_text)
|
||||
chart = llm_service.check_save_chart(res=full_chart_text)
|
||||
print(chart)
|
||||
SQLBotLogUtil.info(chart)
|
||||
if in_chat:
|
||||
yield orjson.dumps({'content': orjson.dumps(chart).decode(), 'type': 'chart'}).decode() + '\n\n'
|
||||
else:
|
||||
@@ -856,7 +855,7 @@ def run_task(llm_service: LLMService, in_chat: bool = True):
|
||||
if chart['type'] != 'table':
|
||||
yield '### generated chart picture\n\n'
|
||||
image_url = request_picture(llm_service.record.chat_id, llm_service.record.id, chart, result)
|
||||
print(image_url)
|
||||
SQLBotLogUtil.info(image_url)
|
||||
yield f'![{chart["type"]}]({image_url})'
|
||||
except Exception as e:
|
||||
traceback.print_exc()
|
||||
|
||||
@@ -15,7 +15,7 @@ from apps.db.db import get_engine, get_tables, get_fields, exec_sql
|
||||
from apps.db.engine import get_engine_config, get_engine_conn
|
||||
from apps.db.type import db_type_relation
|
||||
from common.core.deps import SessionDep, CurrentUser, Trans
|
||||
from common.utils.utils import deepcopy_ignore_extra
|
||||
from common.utils.utils import SQLBotLogUtil, deepcopy_ignore_extra
|
||||
from .table import get_tables_by_ds_id
|
||||
from ..crud.field import delete_field_by_ds_id, update_field
|
||||
from ..crud.table import delete_table_by_ds_id, update_table
|
||||
@@ -39,12 +39,12 @@ def check_status(session: SessionDep, ds: CoreDatasource, is_raise: bool = False
|
||||
conn = get_engine(ds)
|
||||
try:
|
||||
with conn.connect() as connection:
|
||||
print("success")
|
||||
SQLBotLogUtil.info("success")
|
||||
return True
|
||||
except Exception as e:
|
||||
SQLBotLogUtil.error(f"Datasource {ds.id} connection failed: {e}")
|
||||
if is_raise:
|
||||
raise HTTPException(status_code=500, detail=f'Connect Failed: {e.args}')
|
||||
print("Fail:", e)
|
||||
return False
|
||||
|
||||
|
||||
@@ -86,6 +86,7 @@ def create_ds(session: SessionDep, trans: Trans, user: CurrentUser, create_ds: C
|
||||
|
||||
def chooseTables(session: SessionDep, id: int, tables: List[CoreTable]):
|
||||
ds = session.query(CoreDatasource).filter(CoreDatasource.id == id).first()
|
||||
check_status(session, ds, True)
|
||||
sync_table(session, ds, tables)
|
||||
updateNum(session, ds)
|
||||
|
||||
@@ -131,6 +132,7 @@ def getTables(session: SessionDep, id: int):
|
||||
|
||||
|
||||
def getTablesByDs(session: SessionDep, ds: CoreDatasource):
|
||||
check_status(session, ds, True)
|
||||
tables = get_tables(ds)
|
||||
return tables
|
||||
|
||||
@@ -244,6 +246,7 @@ def preview(session: SessionDep, id: int, data: TableObj):
|
||||
return {"fields": [], "data": [], "sql": ''}
|
||||
|
||||
ds = session.query(CoreDatasource).filter(CoreDatasource.id == id).first()
|
||||
check_status(session, ds, True)
|
||||
conf = DatasourceConf(**json.loads(aes_decrypt(ds.configuration))) if ds.type != "excel" else get_engine_config()
|
||||
sql: str = ""
|
||||
if ds.type == "mysql":
|
||||
|
||||
@@ -1,250 +0,0 @@
|
||||
import logging
|
||||
from typing import Dict, List
|
||||
from langchain.schema import Document
|
||||
from langchain_openai import OpenAIEmbeddings
|
||||
from langchain.vectorstores.pgvector import PGVector
|
||||
from sqlalchemy import Tuple
|
||||
|
||||
from apps.chat.schemas.schema_engine import SchemaEngine
|
||||
from apps.datasource.models.datasource import CoreDatasource
|
||||
from apps.db.db import get_uri
|
||||
from langchain_community.utilities import SQLDatabase
|
||||
from sqlmodel import Session, select
|
||||
from apps.system.models.system_model import AiModelDetail
|
||||
from common.core.db import engine
|
||||
class DsEmbeddingManage:
|
||||
def __init__(self, ds: CoreDatasource):
|
||||
self.ds = ds
|
||||
db_uri = get_uri(ds)
|
||||
self.db = SQLDatabase.from_uri(db_uri)
|
||||
|
||||
# Get AI model configuration from database
|
||||
with Session(engine) as session:
|
||||
aimodel = session.exec(select(AiModelDetail).where(
|
||||
AiModelDetail.status == True,
|
||||
AiModelDetail.api_key.is_not(None)
|
||||
)).first()
|
||||
|
||||
if not aimodel:
|
||||
raise RuntimeError("No available AI model configuration found")
|
||||
|
||||
# Initialize embedding model with AI model config
|
||||
self.embedding_model = OpenAIEmbeddings(
|
||||
model=aimodel.name,
|
||||
openai_api_base=aimodel.endpoint if aimodel.endpoint else None,
|
||||
api_key=aimodel.api_key,
|
||||
)
|
||||
|
||||
def getSchema(self):
|
||||
schema_engine = SchemaEngine(engine=self.db._engine)
|
||||
mschema = schema_engine.mschema
|
||||
mschema_str = mschema.to_mschema()
|
||||
return mschema_str
|
||||
|
||||
def mschema_to_documents(self, mschema_str: str) -> List[Document]:
|
||||
"""Convert M-Schema string to list of Documents
|
||||
|
||||
Args:
|
||||
mschema_str: M-Schema format string
|
||||
|
||||
Returns:
|
||||
List of Documents for vector storage
|
||||
"""
|
||||
docs = []
|
||||
current_table = ""
|
||||
table_content = []
|
||||
table_line = ""
|
||||
current_table_col_docs = []
|
||||
# Split schema into lines
|
||||
lines = mschema_str.split('\n')
|
||||
|
||||
for line in lines:
|
||||
# Skip empty lines and header sections
|
||||
if not line.strip() or line.startswith('【'):
|
||||
continue
|
||||
|
||||
# Process table header
|
||||
if line.startswith('# Table:'):
|
||||
# Store previous table if exists
|
||||
if current_table and table_content:
|
||||
doc_content = f"Table: {current_table}\n" + '\n'.join(table_content)
|
||||
table_doc = Document(
|
||||
page_content=doc_content,
|
||||
metadata={
|
||||
"type": "table",
|
||||
"name": current_table,
|
||||
"origin": table_line
|
||||
}
|
||||
)
|
||||
docs.append(table_doc)
|
||||
docs.extend(current_table_col_docs)
|
||||
|
||||
# Start new table
|
||||
current_table = line.split('# Table: ')[1].split(',')[0].strip()
|
||||
table_content = []
|
||||
current_table_col_docs = []
|
||||
table_line = line
|
||||
# If there's a table comment, add it
|
||||
if ',' in line:
|
||||
table_content.append(f"Description: {line.split(',')[1].strip()}. Contains columns: ")
|
||||
continue
|
||||
|
||||
# Add column information to current table
|
||||
if line.strip() and not line.startswith('[') and not line.startswith(']') and line.startswith('('):
|
||||
# Clean up the line
|
||||
column_line = line.strip('(),')
|
||||
column_info_list = column_line.split(',')
|
||||
column_name = column_info_list[0].split(':')[0]
|
||||
table_content.append(column_name + ",\n")
|
||||
column_content = column_line
|
||||
|
||||
column_doc = Document(
|
||||
page_content=column_content,
|
||||
metadata = {
|
||||
"type": "column",
|
||||
"name": column_name,
|
||||
"table": current_table,
|
||||
"origin": line
|
||||
}
|
||||
)
|
||||
current_table_col_docs.append(column_doc)
|
||||
|
||||
|
||||
# Don't forget to add the last table
|
||||
if current_table and table_content:
|
||||
doc_content = f"Table: {current_table}\n" + '\n'.join(table_content)
|
||||
docs.append(Document(
|
||||
page_content=doc_content,
|
||||
metadata={
|
||||
"type": "table",
|
||||
"name": current_table,
|
||||
"origin": table_line
|
||||
}
|
||||
))
|
||||
docs.extend(current_table_col_docs)
|
||||
|
||||
return docs
|
||||
|
||||
def store_schema_embeddings(self):
|
||||
"""Store schema embeddings in PGVector"""
|
||||
try:
|
||||
# Get schema string
|
||||
mschema_str = self.getSchema()
|
||||
|
||||
# Convert to documents
|
||||
docs = self.mschema_to_documents(mschema_str)
|
||||
# print(docs)
|
||||
# Store in PGVector
|
||||
PGVector.from_documents(
|
||||
embedding=self.embedding_model,
|
||||
documents=docs,
|
||||
collection_name="schema_embeddings",
|
||||
connection_string=str(self.db._engine.url),
|
||||
pre_delete_collection=True
|
||||
)
|
||||
|
||||
return f"Successfully stored {len(docs)} schema documents in vector store"
|
||||
|
||||
except Exception as e:
|
||||
raise RuntimeError(f"Failed to store schema embeddings: {str(e)}")
|
||||
|
||||
def get_relevant_schema(
|
||||
self,
|
||||
query: str,
|
||||
collection_name: str = "schema_embeddings",
|
||||
top_k: int = 5
|
||||
) -> List[Dict]:
|
||||
"""Get relevant schema information based on user query
|
||||
|
||||
Args:
|
||||
query: User query text
|
||||
collection_name: Vector collection name
|
||||
top_k: Number of most relevant items to return
|
||||
|
||||
Returns:
|
||||
List of dictionaries containing relevant schema information
|
||||
"""
|
||||
try:
|
||||
store = PGVector(
|
||||
collection_name=collection_name,
|
||||
connection_string=str(self.db._engine.url),
|
||||
embedding_function=self.embedding_model,
|
||||
)
|
||||
|
||||
# Execute similarity search
|
||||
docs = store.similarity_search_with_score(query, k=top_k)
|
||||
|
||||
relevant_tables = {}
|
||||
|
||||
for doc, score in docs:
|
||||
metadata = doc.metadata
|
||||
item = {
|
||||
"name": metadata["name"],
|
||||
"type": metadata["type"],
|
||||
"origin": metadata["origin"],
|
||||
"score": float(score)
|
||||
}
|
||||
|
||||
if metadata["type"] == "table":
|
||||
table_name = metadata["name"]
|
||||
item["columns"] = []
|
||||
if table_name not in relevant_tables:
|
||||
relevant_tables[table_name] = item
|
||||
else:
|
||||
relevant_tables[table_name]["score"] = max(relevant_tables[table_name]["score"], item["score"])
|
||||
relevant_tables[table_name]["origin"] = item["origin"]
|
||||
|
||||
elif metadata["type"] == "column":
|
||||
table_name = f"{metadata['table']}"
|
||||
if table_name not in relevant_tables:
|
||||
relevant_tables[table_name] = {
|
||||
"name": table_name,
|
||||
"type": "table",
|
||||
"score": float(score) * 0.9,
|
||||
"columns": [item]
|
||||
}
|
||||
else:
|
||||
relevant_tables[table_name]["columns"].append(item)
|
||||
|
||||
# Convert dictionary to sorted list
|
||||
relevant_tables_list = list(relevant_tables.values())
|
||||
# Sort by score in descending order
|
||||
relevant_tables_list.sort(key=lambda x: x["score"], reverse=True)
|
||||
|
||||
return relevant_tables_list
|
||||
|
||||
except Exception as e:
|
||||
logging.error(f"Error during schema search: {str(e)}")
|
||||
raise
|
||||
|
||||
def generate_schema_context(self, relevant_tables_list: List[Dict]) -> str:
|
||||
|
||||
context = "## Relevant Database Schema Information:\n\n【Schema】\n"
|
||||
|
||||
if relevant_tables_list:
|
||||
for table in sorted(relevant_tables_list, key=lambda x: -x["score"]):
|
||||
if table['origin'] is not None:
|
||||
context += f"\n{table['origin']}"
|
||||
else:
|
||||
context += f"# Table: {table['name']}"
|
||||
context += f" (relevance score: {table['score']:.2f})\n"
|
||||
|
||||
columns: List[Dict] = table["columns"]
|
||||
# context += f"[\n{',\n'.join(f'({column["origin"].strip('(),')} , relevance score: {column["score"]:.2f})' for column in columns)}]\n"
|
||||
context += "\n[\n" + ",\n".join(column["origin"] for column in columns) + "\n]\n"
|
||||
return context
|
||||
|
||||
def get_schema_with_query(self,
|
||||
query: str,
|
||||
collection_name: str = "schema_embeddings",
|
||||
top_k: int = 5):
|
||||
list = self.get_relevant_schema(query, collection_name, top_k)
|
||||
context = self.generate_schema_context(list)
|
||||
print(context)
|
||||
return context
|
||||
|
||||
def execute(self):
|
||||
"""Main execution method"""
|
||||
return self.store_schema_embeddings()
|
||||
|
||||
|
||||
@@ -126,7 +126,7 @@ def get_tables(ds: CoreDatasource):
|
||||
LEFT JOIN
|
||||
pg_description d ON d.objoid = c.oid AND d.objsubid = 0
|
||||
WHERE n.nspname = current_schema()
|
||||
AND c.relkind IN ('r', 'v', 'p')
|
||||
AND c.relkind IN ('r', 'v', 'p', 'm')
|
||||
AND c.relname NOT LIKE 'pg_%'
|
||||
AND c.relname NOT LIKE 'sql_%'
|
||||
ORDER BY c.relname \
|
||||
|
||||
@@ -1,13 +1,13 @@
|
||||
|
||||
from typing import Optional
|
||||
from sqlmodel import Session, select, delete as sqlmodel_delete
|
||||
import logging
|
||||
from apps.system.models.system_model import UserWsModel, WorkspaceModel
|
||||
from apps.system.schemas.auth import CacheName, CacheNamespace
|
||||
from apps.system.schemas.system_schema import BaseUserDTO, UserInfoDTO, UserWs
|
||||
from common.core.deps import SessionDep
|
||||
from common.core.sqlbot_cache import cache, clear_cache
|
||||
from common.utils.locale import I18n
|
||||
from common.utils.utils import SQLBotLogUtil
|
||||
from ..models.user import UserModel
|
||||
from common.core.security import verify_md5pwd
|
||||
|
||||
@@ -68,4 +68,4 @@ async def single_delete(session: SessionDep, id: int):
|
||||
|
||||
#@clear_cache(namespace=CacheNamespace.AUTH_INFO, cacheName=CacheName.USER_INFO, keyExpression="id")
|
||||
async def clean_user_cache(id: int):
|
||||
logging.info(f"User cache for [{id}] has been cleaned")
|
||||
SQLBotLogUtil.info(f"User cache for [{id}] has been cleaned")
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
|
||||
from typing import Optional
|
||||
from fastapi import HTTPException, Request
|
||||
from fastapi import Request
|
||||
from fastapi.responses import JSONResponse
|
||||
import jwt
|
||||
from sqlmodel import Session
|
||||
@@ -13,6 +13,7 @@ from apps.system.schemas.system_schema import AssistantHeader, UserInfoDTO
|
||||
from common.core import security
|
||||
from common.core.config import settings
|
||||
from common.core.schemas import TokenPayload
|
||||
from common.utils.utils import SQLBotLogUtil
|
||||
from common.utils.whitelist import whiteUtils
|
||||
from fastapi.security.utils import get_authorization_scheme_param
|
||||
class TokenMiddleware(BaseHTTPMiddleware):
|
||||
@@ -35,7 +36,7 @@ class TokenMiddleware(BaseHTTPMiddleware):
|
||||
request.state.current_user = validator[1]
|
||||
request.state.assistant = validator[2]
|
||||
return await call_next(request)
|
||||
return JSONResponse({"msg": f"Unauthorized:[{validator[1]}]"}, status_code=401)
|
||||
return JSONResponse({"msg": f"Unauthorized:[{validator[1]}]"}, status_code=401, headers={"Access-Control-Allow-Origin": "*"})
|
||||
#validate pass
|
||||
tokenkey = settings.TOKEN_KEY
|
||||
token = request.headers.get(tokenkey)
|
||||
@@ -43,7 +44,7 @@ class TokenMiddleware(BaseHTTPMiddleware):
|
||||
if validate_pass:
|
||||
request.state.current_user = data
|
||||
return await call_next(request)
|
||||
return JSONResponse({"msg": f"Unauthorized:[{data}]"}, status_code=401)
|
||||
return JSONResponse({"msg": f"Unauthorized:[{data}]"}, status_code=401, headers={"Access-Control-Allow-Origin": "*"})
|
||||
|
||||
def is_options(self, request: Request):
|
||||
return request.method == "OPTIONS"
|
||||
@@ -70,6 +71,7 @@ class TokenMiddleware(BaseHTTPMiddleware):
|
||||
) """
|
||||
return True, session_user
|
||||
except Exception as e:
|
||||
SQLBotLogUtil.exception(f"Token validation error: {str(e)}", exc_info=True)
|
||||
return False, e
|
||||
|
||||
|
||||
@@ -95,4 +97,6 @@ class TokenMiddleware(BaseHTTPMiddleware):
|
||||
assistant_info = AssistantHeader.model_validate(assistant_info.model_dump(exclude_unset=True))
|
||||
return True, session_user, assistant_info
|
||||
except Exception as e:
|
||||
SQLBotLogUtil.exception(f"Assistant validation error: {str(e)}", exc_info=True)
|
||||
# Return False and the exception message
|
||||
return False, e
|
||||
@@ -60,6 +60,11 @@ class Settings(BaseSettings):
|
||||
TOKEN_KEY: str = "X-SQLBOT-TOKEN"
|
||||
DEFAULT_PWD: str = "SQLBot@123456"
|
||||
ASSISTANT_TOKEN_KEY: str = "X-SQLBOT-ASSISTANT-TOKEN"
|
||||
|
||||
LOG_LEVEL: str = "INFO" # DEBUG, INFO, WARNING, ERROR
|
||||
LOG_DIR: str = "logs"
|
||||
LOG_FORMAT: str = "%(asctime)s - %(name)s - %(levelname)s:%(lineno)d - %(message)s"
|
||||
SQL_DEBUG: bool = False
|
||||
|
||||
@computed_field # type: ignore[prop-decorator]
|
||||
@property
|
||||
|
||||
@@ -1,10 +1,10 @@
|
||||
import json
|
||||
import logging
|
||||
from starlette.middleware.base import BaseHTTPMiddleware
|
||||
from starlette.responses import JSONResponse
|
||||
from starlette.exceptions import HTTPException
|
||||
from starlette.requests import Request
|
||||
from common.core.config import settings
|
||||
from common.utils.utils import SQLBotLogUtil
|
||||
class ResponseMiddleware(BaseHTTPMiddleware):
|
||||
def __init__(self, app):
|
||||
super().__init__(app)
|
||||
@@ -48,13 +48,17 @@ class ResponseMiddleware(BaseHTTPMiddleware):
|
||||
}
|
||||
)
|
||||
except Exception as e:
|
||||
logging.error(f"Response processing error: {str(e)}", exc_info=True)
|
||||
SQLBotLogUtil.error(f"Response processing error: {str(e)}", exc_info=True)
|
||||
return JSONResponse(
|
||||
status_code=500,
|
||||
content={
|
||||
"code": 500,
|
||||
"data": None,
|
||||
"msg": str(e)
|
||||
},
|
||||
headers={
|
||||
k: v for k, v in response.headers.items()
|
||||
if k.lower() not in ("content-length", "content-type")
|
||||
}
|
||||
)
|
||||
|
||||
@@ -64,26 +68,28 @@ class ResponseMiddleware(BaseHTTPMiddleware):
|
||||
class exception_handler():
|
||||
@staticmethod
|
||||
async def http_exception_handler(request: Request, exc: HTTPException):
|
||||
SQLBotLogUtil.exception(f"HTTP Exception: {exc.detail}", exc_info=True)
|
||||
return JSONResponse(
|
||||
status_code=exc.status_code,
|
||||
content={
|
||||
"code": exc.status_code,
|
||||
"msg": exc.detail,
|
||||
"data": None
|
||||
}
|
||||
},
|
||||
headers={"Access-Control-Allow-Origin": "*"}
|
||||
)
|
||||
|
||||
|
||||
@staticmethod
|
||||
async def global_exception_handler(request: Request, exc: Exception):
|
||||
logging.error(f"Error info: {str(exc)}", exc_info=True)
|
||||
|
||||
SQLBotLogUtil.exception(f"Unhandled Exception: {str(exc)}", exc_info=True)
|
||||
return JSONResponse(
|
||||
status_code=500,
|
||||
content={
|
||||
"code": 500,
|
||||
"msg": str(exc),
|
||||
"data": None
|
||||
}
|
||||
},
|
||||
headers={"Access-Control-Allow-Origin": "*"}
|
||||
)
|
||||
|
||||
|
||||
@@ -1,12 +1,13 @@
|
||||
import logging
|
||||
from fastapi_cache import FastAPICache
|
||||
from fastapi_cache.decorator import cache as original_cache
|
||||
from functools import partial, wraps
|
||||
from typing import Optional, Set, Any, Dict, Tuple, Callable
|
||||
from typing import Optional, Set, Any, Dict, Tuple
|
||||
from inspect import Parameter, signature
|
||||
from contextlib import asynccontextmanager
|
||||
import asyncio
|
||||
|
||||
from common.utils.utils import SQLBotLogUtil
|
||||
|
||||
|
||||
# 锁管理
|
||||
_cache_locks = {}
|
||||
@@ -65,7 +66,7 @@ def custom_key_builder(
|
||||
base_key += f"{value}:"
|
||||
|
||||
except (IndexError, KeyError, AttributeError) as e:
|
||||
logging.warning(f"Failed to evaluate keyExpression '{keyExpression}': {str(e)}")
|
||||
SQLBotLogUtil.warning(f"Failed to evaluate keyExpression '{keyExpression}': {str(e)}")
|
||||
|
||||
return base_key
|
||||
|
||||
@@ -127,19 +128,16 @@ def cache(
|
||||
)
|
||||
|
||||
async with _get_cache_lock(cache_key):
|
||||
logging.info(f"Using cache key: {cache_key}")
|
||||
print(f"Using cache key: {cache_key}")
|
||||
SQLBotLogUtil.info(f"Using cache key: {cache_key}")
|
||||
backend = FastAPICache.get_backend()
|
||||
cached_value = await backend.get(cache_key)
|
||||
if cached_value is not None:
|
||||
logging.info(f"Cache hit for key: {cache_key}, the value is: {cached_value}")
|
||||
print(f"Cache hit for key: {cache_key}, the value is: {cached_value}")
|
||||
SQLBotLogUtil.info(f"Cache hit for key: {cache_key}, the value is: {cached_value}")
|
||||
return cached_value
|
||||
|
||||
result = await func(*args, **kwargs)
|
||||
await backend.set(cache_key, result, expire)
|
||||
logging.info(f"Cache miss for key: {cache_key}, result cached.")
|
||||
print(f"Cache miss for key: {cache_key}, result cached.")
|
||||
SQLBotLogUtil.info(f"Cache miss for key: {cache_key}, result cached.")
|
||||
return result
|
||||
|
||||
return wrapper
|
||||
@@ -167,8 +165,7 @@ def clear_cache(
|
||||
if await FastAPICache.get_backend().get(cache_key):
|
||||
await FastAPICache.clear(key=cache_key)
|
||||
result = await func(*args, **kwargs)
|
||||
logging.info(f"Clearing cache for key: {cache_key}")
|
||||
print(f"Clearing cache for key: {cache_key}")
|
||||
SQLBotLogUtil.info(f"Clearing cache for key: {cache_key}")
|
||||
return result
|
||||
return await func(*args, **kwargs)
|
||||
|
||||
|
||||
@@ -1,6 +1,10 @@
|
||||
import hashlib
|
||||
import inspect
|
||||
import logging
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from logging.handlers import RotatingFileHandler
|
||||
from pathlib import Path
|
||||
from common.core.config import settings
|
||||
from typing import Optional
|
||||
|
||||
import jwt
|
||||
@@ -10,9 +14,6 @@ from jwt.exceptions import InvalidTokenError
|
||||
from common.core import security
|
||||
from common.core.config import settings
|
||||
|
||||
logging.basicConfig(level=logging.DEBUG)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def generate_password_reset_token(email: str) -> str:
|
||||
delta = timedelta(hours=settings.EMAIL_RESET_TOKEN_EXPIRE_HOURS)
|
||||
@@ -76,3 +77,142 @@ def extract_nested_json(text):
|
||||
def string_to_numeric_hash(text: str, bits: Optional[int] = 64) -> int:
|
||||
hash_bytes = hashlib.sha256(text.encode()).digest()
|
||||
return int.from_bytes(hash_bytes[:bits//8], byteorder='big')
|
||||
|
||||
|
||||
def setup_logging():
|
||||
# 确保日志目录存在
|
||||
log_dir = Path(settings.LOG_DIR)
|
||||
log_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# 日志格式
|
||||
formatter = logging.Formatter(
|
||||
f'{settings.LOG_FORMAT}'
|
||||
)
|
||||
|
||||
# 控制台日志
|
||||
console_handler = logging.StreamHandler()
|
||||
console_handler.setLevel(settings.LOG_LEVEL)
|
||||
console_handler.setFormatter(formatter)
|
||||
|
||||
# 文件日志处理器
|
||||
file_handlers = {
|
||||
'debug': logging.DEBUG,
|
||||
'info': logging.INFO,
|
||||
'warn': logging.WARNING,
|
||||
'error': logging.ERROR
|
||||
}
|
||||
|
||||
# 主日志记录器
|
||||
root_logger = logging.getLogger()
|
||||
root_logger.setLevel(logging.DEBUG) # 设置最低级别
|
||||
|
||||
# 添加控制台处理器
|
||||
root_logger.addHandler(console_handler)
|
||||
|
||||
# 为每个级别创建文件处理器
|
||||
for level_name, level in file_handlers.items():
|
||||
file_path = log_dir / f"{level_name}.log"
|
||||
handler = RotatingFileHandler(
|
||||
file_path,
|
||||
maxBytes=10 * 1024 * 1024, # 10MB
|
||||
backupCount=5,
|
||||
encoding='utf-8'
|
||||
)
|
||||
handler.setLevel(level)
|
||||
handler.setFormatter(formatter)
|
||||
|
||||
# 添加过滤器只处理特定级别日志
|
||||
if level_name == 'debug':
|
||||
handler.addFilter(lambda record: record.levelno == logging.DEBUG)
|
||||
elif level_name == 'info':
|
||||
handler.addFilter(lambda record: record.levelno == logging.INFO)
|
||||
elif level_name == 'warn':
|
||||
handler.addFilter(lambda record: record.levelno == logging.WARNING)
|
||||
elif level_name == 'error':
|
||||
handler.addFilter(lambda record: record.levelno >= logging.ERROR)
|
||||
|
||||
root_logger.addHandler(handler)
|
||||
|
||||
# SQL日志特殊处理
|
||||
if settings.LOG_LEVEL == "DEBUG" and settings.SQL_DEBUG:
|
||||
sql_logger = logging.getLogger('sqlalchemy.engine')
|
||||
sql_logger.setLevel(logging.DEBUG)
|
||||
|
||||
sql_handler = RotatingFileHandler(
|
||||
log_dir / "sql.log",
|
||||
maxBytes=10 * 1024 * 1024,
|
||||
backupCount=2,
|
||||
encoding='utf-8'
|
||||
)
|
||||
sql_handler.setFormatter(formatter)
|
||||
sql_logger.addHandler(sql_handler)
|
||||
|
||||
setup_logging()
|
||||
|
||||
|
||||
class CallerLogger(logging.Logger):
|
||||
def __init__(self, logger: logging.Logger):
|
||||
self.logger = logger
|
||||
super().__init__(logger.name, logger.level)
|
||||
|
||||
def _log(self, level, msg, args, exc_info=None, extra=None, stacklevel=3):
|
||||
if self.logger.isEnabledFor(level):
|
||||
self.logger._log(level, msg, args, exc_info=exc_info, extra=extra, stacklevel=stacklevel)
|
||||
|
||||
class SQLBotLogUtil:
|
||||
|
||||
@staticmethod
|
||||
def _get_logger() -> logging.Logger:
|
||||
frame = inspect.currentframe()
|
||||
try:
|
||||
caller_frame = frame.f_back.f_back
|
||||
module_name = caller_frame.f_globals.get('__name__', '__main__')
|
||||
return CallerLogger(logging.getLogger(module_name))
|
||||
finally:
|
||||
del frame
|
||||
|
||||
|
||||
@staticmethod
|
||||
def debug(msg: str, *args, **kwargs):
|
||||
logger = SQLBotLogUtil._get_logger()
|
||||
if logger.isEnabledFor(logging.DEBUG):
|
||||
logger._log(logging.DEBUG, msg, args, **kwargs)
|
||||
|
||||
@staticmethod
|
||||
def info(msg: str, *args, **kwargs):
|
||||
logger = SQLBotLogUtil._get_logger()
|
||||
if logger.isEnabledFor(logging.INFO):
|
||||
logger._log(logging.INFO, msg, args, **kwargs)
|
||||
|
||||
@staticmethod
|
||||
def warning(msg: str, *args, **kwargs):
|
||||
logger = SQLBotLogUtil._get_logger()
|
||||
if logger.isEnabledFor(logging.WARNING):
|
||||
logger._log(logging.WARNING, msg, args, **kwargs)
|
||||
|
||||
@staticmethod
|
||||
def error(msg: str, *args, exc_info: Optional[bool] = None, **kwargs):
|
||||
logger = SQLBotLogUtil._get_logger()
|
||||
if logger.isEnabledFor(logging.ERROR):
|
||||
logger._log(
|
||||
logging.ERROR,
|
||||
msg,
|
||||
args,
|
||||
exc_info=exc_info if exc_info is not None else True,
|
||||
**kwargs
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def exception(msg: str, *args, **kwargs):
|
||||
logger = SQLBotLogUtil._get_logger()
|
||||
if logger.isEnabledFor(logging.ERROR):
|
||||
logger._log(logging.ERROR, msg, args, exc_info=True, **kwargs)
|
||||
|
||||
@staticmethod
|
||||
def critical(msg: str, *args, **kwargs):
|
||||
logger = SQLBotLogUtil._get_logger()
|
||||
if logger.isEnabledFor(logging.CRITICAL):
|
||||
logger._log(logging.CRITICAL, msg, args, **kwargs)
|
||||
|
||||
|
||||
|
||||
|
||||
@@ -2,6 +2,7 @@
|
||||
import re
|
||||
from typing import List, Pattern
|
||||
from common.core.config import settings
|
||||
from common.utils.utils import SQLBotLogUtil
|
||||
wlist = [
|
||||
"/",
|
||||
"/docs",
|
||||
@@ -50,7 +51,7 @@ class WhitelistChecker:
|
||||
try:
|
||||
self._compiled_patterns.append(re.compile(regex_pattern))
|
||||
except re.error:
|
||||
print(f"Invalid regular expression pattern: {pattern}")
|
||||
SQLBotLogUtil.error(f"Invalid regex pattern: {regex_pattern}")
|
||||
|
||||
def is_whitelisted(self, path: str) -> bool:
|
||||
prefix = settings.API_V1_STR
|
||||
|
||||
+4
-3
@@ -1,4 +1,3 @@
|
||||
import logging
|
||||
from fastapi.concurrency import asynccontextmanager
|
||||
import sentry_sdk
|
||||
from fastapi import FastAPI
|
||||
@@ -18,6 +17,8 @@ import sqlbot_xpack
|
||||
from fastapi_cache import FastAPICache
|
||||
from fastapi_cache.backends.inmemory import InMemoryBackend
|
||||
|
||||
from common.utils.utils import SQLBotLogUtil
|
||||
|
||||
def run_migrations():
|
||||
alembic_cfg = Config("alembic.ini")
|
||||
command.upgrade(alembic_cfg, "head")
|
||||
@@ -27,10 +28,10 @@ def run_migrations():
|
||||
async def lifespan(app: FastAPI):
|
||||
run_migrations()
|
||||
FastAPICache.init(InMemoryBackend())
|
||||
print("✅ FastAPICache 初始化完成")
|
||||
init_dynamic_cors(app)
|
||||
SQLBotLogUtil.info("✅ SQLBot 初始化完成")
|
||||
yield
|
||||
|
||||
SQLBotLogUtil.info("SQLBot 应用关闭")
|
||||
|
||||
def custom_generate_unique_id(route: APIRoute) -> str:
|
||||
tag = route.tags[0] if route.tags and len(route.tags) > 0 else ""
|
||||
|
||||
@@ -158,7 +158,7 @@ const afterTreeInit = () => {
|
||||
}
|
||||
|
||||
const copyLoading = ref(false)
|
||||
const emit = defineEmits(['nodeClick'])
|
||||
const emit = defineEmits(['nodeClick', 'deleteCurResource'])
|
||||
|
||||
function createNewObject() {
|
||||
addOperation({ opt: 'newLeaf' })
|
||||
@@ -230,7 +230,8 @@ const operation = (opt: string, data: SQTreeNode) => {
|
||||
dashboardApi.delete_resource({ id: data.id }).then(() => {
|
||||
ElMessage.success(t('dashboard.delete_success'))
|
||||
getTree()
|
||||
dashboardStore.setCurComponent(null)
|
||||
dashboardStore.canvasDataInit()
|
||||
emit('deleteCurResource')
|
||||
})
|
||||
})
|
||||
} else if (opt === 'rename') {
|
||||
|
||||
@@ -48,6 +48,13 @@ const mounted = computed(() => {
|
||||
function createNew() {
|
||||
resourceTreeRef.value?.createNewObject()
|
||||
}
|
||||
|
||||
const stateInit = () => {
|
||||
state.canvasDataPreview = []
|
||||
state.canvasStylePreview = {}
|
||||
state.canvasViewInfoPreview = {}
|
||||
state.dashboardInfo = {}
|
||||
}
|
||||
const loadCanvasData = (params: any) => {
|
||||
dataInitState.value = false
|
||||
load_resource_prepare(
|
||||
@@ -100,6 +107,7 @@ defineExpose({
|
||||
:cur-canvas-type="'dashboard'"
|
||||
:show-position="showPosition"
|
||||
@node-click="resourceNodeClick"
|
||||
@delete-cur-resource="stateInit"
|
||||
/>
|
||||
</el-aside>
|
||||
<el-container
|
||||
|
||||
@@ -95,16 +95,12 @@ const handleEditDatasource = (res: any) => {
|
||||
})
|
||||
}
|
||||
|
||||
const handleQuestion = (ele: any) => {
|
||||
datasourceApi.check(ele).then((res: any) => {
|
||||
if (res) {
|
||||
router.push({
|
||||
path: '/chat/index',
|
||||
query: {
|
||||
start_chat: ele.id,
|
||||
},
|
||||
})
|
||||
}
|
||||
const handleQuestion = (id: string) => {
|
||||
router.push({
|
||||
path: '/chat/index',
|
||||
query: {
|
||||
start_chat: id,
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
@@ -294,7 +290,7 @@ const back = () => {
|
||||
:type-name="ele.type_name"
|
||||
:num="ele.num"
|
||||
:description="ele.description"
|
||||
@question="handleQuestion(ele)"
|
||||
@question="handleQuestion"
|
||||
@edit="handleEditDatasource(ele)"
|
||||
@del="deleteHandler(ele)"
|
||||
@data-table-detail="dataTableDetail(ele)"
|
||||
|
||||
Reference in New Issue
Block a user