From c54902e15c26428bd5efe8d70356824e4cd5944b Mon Sep 17 00:00:00 2001 From: junjun Date: Thu, 9 Oct 2025 14:06:19 +0800 Subject: [PATCH] feat: Vector retrieval matches tables --- backend/alembic/versions/047_table_embedding.py | 3 +-- backend/apps/datasource/crud/table.py | 8 ++++---- backend/apps/datasource/embedding/table_embedding.py | 5 +++-- backend/apps/datasource/models/datasource.py | 3 +-- 4 files changed, 9 insertions(+), 10 deletions(-) diff --git a/backend/alembic/versions/047_table_embedding.py b/backend/alembic/versions/047_table_embedding.py index cdb1766a..1b6c766e 100644 --- a/backend/alembic/versions/047_table_embedding.py +++ b/backend/alembic/versions/047_table_embedding.py @@ -9,7 +9,6 @@ from alembic import op import sqlalchemy as sa import sqlmodel.sql.sqltypes from sqlalchemy.dialects import postgresql -import pgvector # revision identifiers, used by Alembic. revision = 'c1b794a961ce' @@ -20,7 +19,7 @@ depends_on = None def upgrade(): # ### commands auto generated by Alembic - please adjust! ### - op.add_column('core_table', sa.Column('embedding', pgvector.sqlalchemy.vector.VECTOR(), nullable=True)) + op.add_column('core_table', sa.Column('embedding', sa.Text(), nullable=True)) # ### end Alembic commands ### diff --git a/backend/apps/datasource/crud/table.py b/backend/apps/datasource/crud/table.py index b4541d87..220d5e3b 100644 --- a/backend/apps/datasource/crud/table.py +++ b/backend/apps/datasource/crud/table.py @@ -38,7 +38,7 @@ def run_fill_empty_table_embedding(session: Session): SQLBotLogUtil.info('get tables') stmt = select(CoreTable.id).where(and_(CoreTable.embedding.is_(None))) results = session.execute(stmt).scalars().all() - SQLBotLogUtil.info(results) + SQLBotLogUtil.info(json.dumps(results)) save_table_embedding(session, results) @@ -53,8 +53,8 @@ def save_table_embedding(session: Session, ids: List[int]): SQLBotLogUtil.info('start table embedding') start_time = time.time() model = EmbeddingModelCache.get_model() - for id in ids: - table = session.query(CoreTable).filter(CoreTable.id == id).first() + for _id in ids: + table = session.query(CoreTable).filter(CoreTable.id == _id).first() fields = session.query(CoreField).filter(CoreField.table_id == table.id).all() schema_table = '' @@ -82,7 +82,7 @@ def save_table_embedding(session: Session, ids: List[int]): # table_schema.append(schema_table) emb = model.embed_query(schema_table) - stmt = update(CoreTable).where(and_(CoreTable.id == id)).values(embedding=emb) + stmt = update(CoreTable).where(and_(CoreTable.id == id)).values(embedding=json.dumps(emb)) session.execute(stmt) session.commit() diff --git a/backend/apps/datasource/embedding/table_embedding.py b/backend/apps/datasource/embedding/table_embedding.py index ca949a5d..0724d7dd 100644 --- a/backend/apps/datasource/embedding/table_embedding.py +++ b/backend/apps/datasource/embedding/table_embedding.py @@ -56,12 +56,13 @@ def calc_table_embedding(tables: list[dict], question: str): # results = model.embed_documents(text) # end_time = time.time() # SQLBotLogUtil.info(str(end_time - start_time)) - results = [item.get('embedding') if item.get('embedding') else ' ' for item in _list] + results = [item.get('embedding') for item in _list] q_embedding = model.embed_query(question) for index in range(len(results)): item = results[index] - _list[index]['cosine_similarity'] = cosine_similarity(q_embedding, item) + if item: + _list[index]['cosine_similarity'] = cosine_similarity(q_embedding, json.loads(item)) _list.sort(key=lambda x: x['cosine_similarity'], reverse=True) _list = _list[:settings.TABLE_EMBEDDING_COUNT] diff --git a/backend/apps/datasource/models/datasource.py b/backend/apps/datasource/models/datasource.py index fd8f3d5b..9584ed3f 100644 --- a/backend/apps/datasource/models/datasource.py +++ b/backend/apps/datasource/models/datasource.py @@ -1,7 +1,6 @@ from datetime import datetime from typing import List, Optional -from pgvector.sqlalchemy import VECTOR from pydantic import BaseModel from sqlalchemy import Column, Text, BigInteger, DateTime, Identity from sqlalchemy.dialects.postgresql import JSONB @@ -32,7 +31,7 @@ class CoreTable(SQLModel, table=True): table_name: str = Field(sa_column=Column(Text)) table_comment: str = Field(sa_column=Column(Text)) custom_comment: str = Field(sa_column=Column(Text)) - embedding: Optional[List[float]] = Field(sa_column=Column(VECTOR(), nullable=True)) + embedding: str = Field(sa_column=Column(Text, nullable=True)) class CoreField(SQLModel, table=True):