feat: Vector retrieval matches tables

Author: XiaJunjie2020

backend/alembic/versions/047_table_embedding.py

@@ -9,7 +9,6 @@
 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 @@
 
 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 ###
 
 

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()
 

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]

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):