outman/dify
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

add score threshold

Browse Source
This commit is contained in:
jyong 2024-04-25 19:25:08 +08:00
parent a401a73eb7
commit 3f34b7e103
1 changed files with 13 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

18
api/core/rag/datasource/vdb/tencent/tencent_vector.py
View File

@ -67,7 +67,11 @@ class TencentVector(BaseVector):
def _create_collection(self, dimension: int) -> None: def _create_collection(self, dimension: int) -> None:
lock_name = 'vector_indexing_lock_{}'.format(self._collection_name) lock_name = 'vector_indexing_lock_{}'.format(self._collection_name)
with redis_client.lock(lock_name, timeout=20): with redis_client.lock(lock_name, timeout=20):
self.delete() collections = self._db.list_collections()
for collection in collections:
if collection.collection_name == self._collection_name:
self.collection = collection
return
index_type = None index_type = None
for k, v in enum.IndexType.__members__.items(): for k, v in enum.IndexType.__members__.items():
if k == self._client_config.index_type: if k == self._client_config.index_type:
@ -153,12 +157,13 @@ class TencentVector(BaseVector):
limit=kwargs.get('top_k', 4), limit=kwargs.get('top_k', 4),
timeout=self._client_config.timeout, timeout=self._client_config.timeout,
) )
return self._get_search_res(res) score_threshold = kwargs.get("score_threshold", .0) if kwargs.get('score_threshold', .0) else 0.0
return self._get_search_res(res, score_threshold)
def search_by_full_text(self, query: str, **kwargs: Any) -> list[Document]: def search_by_full_text(self, query: str, **kwargs: Any) -> list[Document]:
return [] return []
def _get_search_res(self, res): def _get_search_res(self, res, score_threshold):
docs = [] docs = []
if res is None or len(res) == 0: if res is None or len(res) == 0:
return docs return docs
@ -167,8 +172,11 @@ class TencentVector(BaseVector):
meta = result.get(self.field_metadata) meta = result.get(self.field_metadata)
if meta is not None: if meta is not None:
meta = json.loads(meta) meta = json.loads(meta)
doc = Document(page_content=result.get(self.field_text), metadata=meta) score = 1 - result.get("score")
docs.append(doc) if score > score_threshold:
meta['score'] = score
doc = Document(page_content=result.get(self.field_text), metadata=meta)
docs.append(doc)
return docs return docs
def delete(self) -> None: def delete(self) -> None: