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

Merge remote-tracking branch 'origin/main' into feat_rag_graph

Browse Source
This commit is contained in:
aries_ckt 2023-10-16 12:48:08 +08:00
commit 71c31c3e2e
1 changed files with 17 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

20
pilot/summary/db_summary_client.py
View File

@ -22,8 +22,10 @@ chat_factory = ChatFactory()
class DBSummaryClient: class DBSummaryClient:
"""db summary client, provide db_summary_embedding(put db profile and table profile summary into vector store) """DB Summary client, provide db_summary_embedding(put db profile and table profile summary into vector store)
, get_similar_tables method(get user query related tables info) , get_similar_tables method(get user query related tables info)
Args:
system_app (SystemApp): Main System Application class that manages the lifecycle and registration of components..
""" """
def __init__(self, system_app: SystemApp): def __init__(self, system_app: SystemApp):
@ -160,6 +162,12 @@ class DBSummaryClient:
) )
def init_db_profile(self, db_summary_client, dbname, embeddings): 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 from pilot.embedding_engine.string_embedding import StringEmbedding
vector_store_name = dbname + "_profile" vector_store_name = dbname + "_profile"
@ -176,9 +184,15 @@ class DBSummaryClient:
docs = [] docs = []
docs.extend(embedding.read_batch()) docs.extend(embedding.read_batch())
for table_summary in db_summary_client.table_info_json(): 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( embedding = StringEmbedding(
table_summary, file_path=table_summary,
profile_store_config, vector_store_config=profile_store_config,
text_splitter=text_splitter,
) )
docs.extend(embedding.read_batch()) docs.extend(embedding.read_batch())
embedding.index_to_store(docs) embedding.index_to_store(docs)