langchain新版本PGVector传入connection,插入、删除操作会回滚

This commit is contained in:
lijia 2024-01-16 16:58:19 +08:00
parent 67ae61d567
commit 2c92cd00ae

View File

@ -23,11 +23,11 @@ class PGKBService(KBService):
self.pg_vector = PGVector(embedding_function=EmbeddingsFunAdapter(self.embed_model),
collection_name=self.kb_name,
distance_strategy=DistanceStrategy.EUCLIDEAN,
connection=self.engine,
connection=PGKBService.engine,
connection_string=kbs_config.get("pg").get("connection_uri"))
def get_doc_by_ids(self, ids: List[str]) -> List[Document]:
with Session(self.engine) as session:
with Session(PGKBService.engine) as session:
stmt = text("SELECT document, cmetadata FROM langchain_pg_embedding WHERE collection_id in :ids")
results = [Document(page_content=row[0], metadata=row[1]) for row in
session.execute(stmt, {'ids': ids}).fetchall()]
@ -47,7 +47,7 @@ class PGKBService(KBService):
return SupportedVSType.PG
def do_drop_kb(self):
with Session(self.engine) as session:
with Session(PGKBService.engine) as session:
session.execute(text(f'''
-- 删除 langchain_pg_embedding 表中关联到 langchain_pg_collection 表中 的记录
DELETE FROM langchain_pg_embedding
@ -72,7 +72,7 @@ class PGKBService(KBService):
return doc_infos
def do_delete_doc(self, kb_file: KnowledgeFile, **kwargs):
with Session(self.engine) as session:
with Session(PGKBService.engine) as session:
filepath = kb_file.filepath.replace('\\', '\\\\')
session.execute(
text(