mirror of
https://gh-proxy.org/https://github.com/klinsmaya/SQLBot.git
synced 2026-11-08 01:06:59 +08:00
feat: Vector retrieval matches tables
This commit is contained in:
@@ -406,7 +406,7 @@ class LLMService:
|
||||
full_thinking_text = ''
|
||||
full_text = ''
|
||||
if not ignore_auto_select:
|
||||
if settings.EMBEDDING_ENABLED:
|
||||
if settings.TABLE_EMBEDDING_ENABLED:
|
||||
ds = get_ds_embedding(self.session, self.current_user, _ds_list, self.out_ds_instance,
|
||||
self.chat_question.question, self.current_assistant)
|
||||
yield {'content': '{"id":' + str(ds.get('id')) + '}'}
|
||||
@@ -496,7 +496,7 @@ class LLMService:
|
||||
except Exception as e:
|
||||
_error = e
|
||||
|
||||
if not ignore_auto_select and not settings.EMBEDDING_ENABLED:
|
||||
if not ignore_auto_select and not settings.TABLE_EMBEDDING_ENABLED:
|
||||
self.record = save_select_datasource_answer(session=self.session, record_id=self.record.id,
|
||||
answer=orjson.dumps({'content': full_text}).decode(),
|
||||
datasource=_datasource,
|
||||
|
||||
@@ -13,6 +13,7 @@ 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 deepcopy_ignore_extra
|
||||
from .table import get_tables_by_ds_id
|
||||
@@ -402,7 +403,7 @@ def get_table_schema(session: SessionDep, current_user: CurrentUser, ds: CoreDat
|
||||
all_tables.append(t_obj)
|
||||
|
||||
# do table embedding
|
||||
if embedding and tables:
|
||||
if embedding and tables and settings.TABLE_EMBEDDING_ENABLED:
|
||||
tables = get_table_embedding(session, current_user, tables, question)
|
||||
# splice schema
|
||||
if tables:
|
||||
|
||||
@@ -105,6 +105,7 @@ class Settings(BaseSettings):
|
||||
PG_POOL_RECYCLE: int = 3600
|
||||
PG_POOL_PRE_PING: bool = True
|
||||
|
||||
TABLE_EMBEDDING_ENABLED = False
|
||||
TABLE_EMBEDDING_COUNT: int = 10
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user