github/DB-GPT
1
0
Fork 0
mirror of https://github.com/csunny/DB-GPT.git synced 2025-08-17 15:58:25 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix:db schema search not complete

Browse Source 
Close #660
This commit is contained in:
aries_ckt 2023-10-16 10:07:40 +08:00
parent 66020a1adf
commit a596b42c57
1 changed files with 14 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

16
pilot/summary/db_summary_client.py
View File

@ -162,6 +162,12 @@ class DBSummaryClient:
)
def init_db_profile(self, db_summary_client, dbname, embeddings):
"""db profile initialization
Args:
db_summary_client(DBSummaryClient): DB Summary Client
dbname(str): dbname
embeddings(SourceEmbedding): embedding for read string document
"""
from pilot.embedding_engine.string_embedding import StringEmbedding
vector_store_name = dbname + "_profile"
@ -178,9 +184,15 @@ class DBSummaryClient:
docs = []
docs.extend(embedding.read_batch())
for table_summary in db_summary_client.table_info_json():
from langchain.text_splitter import RecursiveCharacterTextSplitter
text_splitter = RecursiveCharacterTextSplitter(
chunk_size=len(table_summary), chunk_overlap=100
)
embedding = StringEmbedding(
table_summary,
profile_store_config,
file_path=table_summary,
vector_store_config=profile_store_config,
text_splitter=text_splitter,
)
docs.extend(embedding.read_batch())
embedding.index_to_store(docs)