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

feat: Vector retrieval matches tables

Browse Source
This commit is contained in:
junjun
2025-09-23 15:17:48 +08:00
parent 5acb04a561
commit eb82d4fb56
4 changed files with 58 additions and 49 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/apps/datasource/crud/datasource.py
+1 -32
View File
@@ -1,22 +1,18 @@
import datetime
import json
import traceback
from typing import List, Optional
from fastapi import HTTPException
from sqlalchemy import and_, text
from sqlmodel import select
from apps.ai_model.embedding import EmbeddingModelCache
from apps.datasource.crud.permission import get_column_permission_fields, get_row_permission_filters, is_normal_user
from apps.datasource.embedding.ds_embedding import cosine_similarity
from apps.datasource.embedding.table_embedding import get_table_embedding
from apps.datasource.utils.utils import aes_decrypt
from apps.db.constant import DB
from apps.db.db import get_tables, get_fields, exec_sql, check_connection
from apps.db.engine import get_engine_config, get_engine_conn
from common.core.config import settings
from common.core.deps import SessionDep, CurrentUser, Trans
from common.utils.utils import SQLBotLogUtil
from common.utils.utils import deepcopy_ignore_extra
from .table import get_tables_by_ds_id
from ..crud.field import delete_field_by_ds_id, update_field
@@ -389,30 +385,3 @@ def get_table_schema(session: SessionDep, current_user: CurrentUser, ds: CoreDat
for s in tables:
schema_str += s
return schema_str
def get_table_embedding(session: SessionDep, current_user: CurrentUser, tables: list[str], question: str):
_list = []
for table_schema in tables:
_list.append({"table_schema": table_schema, "cosine_similarity": 0.0})
if _list:
try:
text = [s.get('table_schema') for s in _list]
model = EmbeddingModelCache.get_model()
results = model.embed_documents(text)
q_embedding = model.embed_query(question)
for index in range(len(results)):
item = results[index]
_list[index]['cosine_similarity'] = cosine_similarity(q_embedding, item)
_list.sort(key=lambda x: x['cosine_similarity'], reverse=True)
_list = _list[:settings.TABLE_EMBEDDING_COUNT]
# print(len(_list))
SQLBotLogUtil.info(json.dumps(_list))
return [t.get("table_schema") for t in _list]
except Exception:
traceback.print_exc()
return _list
backend/apps/datasource/embedding/ds_embedding.py
+2 -17
View File
@@ -1,12 +1,12 @@
# Author: Junjun
# Date: 2025/9/18
import json
import math
import traceback
from typing import Optional
from apps.ai_model.embedding import EmbeddingModelCache
from apps.datasource.crud.datasource import get_table_schema
from apps.datasource.embedding.utils import cosine_similarity
from apps.datasource.models.datasource import CoreDatasource
from apps.system.crud.assistant import AssistantOutDs
from common.core.deps import CurrentAssistant
@@ -14,21 +14,6 @@ from common.core.deps import SessionDep, CurrentUser
from common.utils.utils import SQLBotLogUtil
def cosine_similarity(vec_a, vec_b):
if len(vec_a) != len(vec_b):
raise ValueError("The vector dimension must be the same")
dot_product = sum(a * b for a, b in zip(vec_a, vec_b))
norm_a = math.sqrt(sum(a * a for a in vec_a))
norm_b = math.sqrt(sum(b * b for b in vec_b))
if norm_a == 0 or norm_b == 0:
return 0.0
return dot_product / (norm_a * norm_b)
def get_ds_embedding(session: SessionDep, current_user: CurrentUser, _ds_list, out_ds: AssistantOutDs,
question: str,
current_assistant: Optional[CurrentAssistant] = None):
@@ -45,7 +30,7 @@ def get_ds_embedding(session: SessionDep, current_user: CurrentUser, _ds_list, o
for _ds in _ds_list:
if _ds.get('id'):
ds = session.get(CoreDatasource, _ds.get('id'))
table_schema = get_table_schema(session, current_user, ds, embedding=False)
table_schema = get_table_schema(session, current_user, ds, question, embedding=False)
ds_info = f"{ds.name}, {ds.description}\n"
ds_schema = ds_info + table_schema
_list.append({"id": ds.id, "ds_schema": ds_schema, "cosine_similarity": 0.0, "ds": ds})
backend/apps/datasource/embedding/table_embedding.py
+37
View File
@@ -0,0 +1,37 @@
# Author: Junjun
# Date: 2025/9/23
import json
import traceback
from apps.ai_model.embedding import EmbeddingModelCache
from apps.datasource.embedding.utils import cosine_similarity
from common.core.config import settings
from common.core.deps import SessionDep, CurrentUser
from common.utils.utils import SQLBotLogUtil
def get_table_embedding(session: SessionDep, current_user: CurrentUser, tables: list[str], question: str):
_list = []
for table_schema in tables:
_list.append({"table_schema": table_schema, "cosine_similarity": 0.0})
if _list:
try:
text = [s.get('table_schema') for s in _list]
model = EmbeddingModelCache.get_model()
results = model.embed_documents(text)
q_embedding = model.embed_query(question)
for index in range(len(results)):
item = results[index]
_list[index]['cosine_similarity'] = cosine_similarity(q_embedding, item)
_list.sort(key=lambda x: x['cosine_similarity'], reverse=True)
_list = _list[:settings.TABLE_EMBEDDING_COUNT]
# print(len(_list))
SQLBotLogUtil.info(json.dumps(_list))
return [t.get("table_schema") for t in _list]
except Exception:
traceback.print_exc()
return _list
backend/apps/datasource/embedding/utils.py
+18
View File
@@ -0,0 +1,18 @@
# Author: Junjun
# Date: 2025/9/23
import math
def cosine_similarity(vec_a, vec_b):
if len(vec_a) != len(vec_b):
raise ValueError("The vector dimension must be the same")
dot_product = sum(a * b for a, b in zip(vec_a, vec_b))
norm_a = math.sqrt(sum(a * a for a in vec_a))
norm_b = math.sqrt(sum(b * b for b in vec_b))
if norm_a == 0 or norm_b == 0:
return 0.0
return dot_product / (norm_a * norm_b)