344 lines
15 KiB
Python
344 lines
15 KiB
Python
# coding=utf-8
|
||
"""
|
||
@project: maxkb
|
||
@Author：虎
|
||
@file： pg_vector.py
|
||
@date：2023/10/19 15:28
|
||
@desc:
|
||
"""
|
||
import json
|
||
import os
|
||
from abc import ABC, abstractmethod
|
||
from typing import Dict, List
|
||
|
||
import uuid_utils.compat as uuid
|
||
from django.contrib.postgres.search import SearchVector
|
||
from django.db.models import QuerySet, Value
|
||
from langchain_core.embeddings import Embeddings
|
||
|
||
from common.db.search import generate_sql_by_query_dict
|
||
from common.db.sql_execute import select_list
|
||
from common.utils.common import get_file_content
|
||
from common.utils.ts_vecto_util import to_ts_vector, to_query
|
||
from knowledge.models import Embedding, SearchMode, SourceType
|
||
from knowledge.vector.base_vector import BaseVectorStore
|
||
from maxkb.conf import PROJECT_DIR
|
||
|
||
|
||
class PGVector(BaseVectorStore):
|
||
|
||
def delete_by_source_ids(self, source_ids: List[str], source_type: str):
|
||
if len(source_ids) == 0:
|
||
return
|
||
QuerySet(Embedding).filter(source_id__in=source_ids, source_type=source_type).delete()
|
||
|
||
def update_by_source_ids(self, source_ids: List[str], instance: Dict):
|
||
QuerySet(Embedding).filter(source_id__in=source_ids).update(**instance)
|
||
|
||
def vector_is_create(self) -> bool:
|
||
# 项目启动默认是创建好的 不需要再创建
|
||
return True
|
||
|
||
def vector_create(self):
|
||
return True
|
||
|
||
def _save(self, text, source_type: SourceType, knowledge_id: str, document_id: str, paragraph_id: str,
|
||
source_id: str,
|
||
is_active: bool,
|
||
embedding: Embeddings):
|
||
text_embedding = [float(x) for x in embedding.embed_query(text)]
|
||
embedding = Embedding(
|
||
id=uuid.uuid7(),
|
||
knowledge_id=knowledge_id,
|
||
document_id=document_id,
|
||
is_active=is_active,
|
||
paragraph_id=paragraph_id,
|
||
source_id=source_id,
|
||
embedding=text_embedding,
|
||
source_type=source_type,
|
||
search_vector=to_ts_vector(text)
|
||
)
|
||
embedding.save()
|
||
return True
|
||
|
||
def _batch_save(self, text_list: List[Dict], embedding: Embeddings, is_the_task_interrupted):
|
||
from common.utils.logger import maxkb_logger
|
||
texts = [row.get('text') for row in text_list]
|
||
maxkb_logger.info(f"PGVector batch_save: Processing {len(texts)} texts")
|
||
|
||
# Log details of first few items for debugging
|
||
for i, item in enumerate(text_list[:3]):
|
||
maxkb_logger.debug(f"Item {i}: document_id={item.get('document_id')}, "
|
||
f"paragraph_id={item.get('paragraph_id')}, "
|
||
f"is_active={item.get('is_active', True)}, "
|
||
f"text_preview='{item.get('text', '')[:50]}...'")
|
||
|
||
embeddings = embedding.embed_documents(texts)
|
||
embedding_list = [
|
||
Embedding(
|
||
id=uuid.uuid7(),
|
||
document_id=text_list[index].get('document_id'),
|
||
paragraph_id=text_list[index].get('paragraph_id'),
|
||
knowledge_id=text_list[index].get('knowledge_id'),
|
||
is_active=text_list[index].get('is_active', True),
|
||
source_id=text_list[index].get('source_id'),
|
||
source_type=text_list[index].get('source_type'),
|
||
embedding=[float(x) for x in embeddings[index]],
|
||
search_vector=SearchVector(Value(to_ts_vector(text_list[index]['text'])))
|
||
) for index in range(0, len(texts))]
|
||
|
||
maxkb_logger.info(f"PGVector batch_save: Created {len(embedding_list)} embedding objects")
|
||
|
||
if not is_the_task_interrupted():
|
||
if len(embedding_list) > 0:
|
||
QuerySet(Embedding).bulk_create(embedding_list)
|
||
maxkb_logger.info(f"PGVector batch_save: Successfully saved {len(embedding_list)} embeddings to database")
|
||
else:
|
||
maxkb_logger.warning("PGVector batch_save: No embeddings to save")
|
||
else:
|
||
maxkb_logger.warning("PGVector batch_save: Task interrupted, embeddings not saved")
|
||
return True
|
||
|
||
def hit_test(self, query_text, knowledge_id_list: list[str], exclude_document_id_list: list[str], top_number: int,
|
||
similarity: float,
|
||
search_mode: SearchMode,
|
||
embedding: Embeddings):
|
||
if knowledge_id_list is None or len(knowledge_id_list) == 0:
|
||
return []
|
||
exclude_dict = {}
|
||
embedding_query = embedding.embed_query(query_text)
|
||
query_set = QuerySet(Embedding).filter(knowledge_id__in=knowledge_id_list, is_active=True)
|
||
if exclude_document_id_list is not None and len(exclude_document_id_list) > 0:
|
||
exclude_dict.__setitem__('document_id__in', exclude_document_id_list)
|
||
query_set = query_set.exclude(**exclude_dict)
|
||
for search_handle in search_handle_list:
|
||
if search_handle.support(search_mode):
|
||
return search_handle.handle(query_set, query_text, embedding_query, top_number, similarity, search_mode)
|
||
|
||
def query(self, query_text: str, query_embedding: List[float], knowledge_id_list: list[str],
|
||
exclude_document_id_list: list[str],
|
||
exclude_paragraph_list: list[str], is_active: bool, top_n: int, similarity: float,
|
||
search_mode: SearchMode):
|
||
from common.utils.logger import maxkb_logger
|
||
exclude_dict = {}
|
||
if knowledge_id_list is None or len(knowledge_id_list) == 0:
|
||
maxkb_logger.warning("Vector query: knowledge_id_list is empty")
|
||
return []
|
||
|
||
maxkb_logger.info(f"Vector query starting: query_text='{query_text[:50]}...', knowledge_ids={knowledge_id_list}, "
|
||
f"is_active={is_active}, top_n={top_n}, similarity={similarity}, search_mode={search_mode.value}")
|
||
|
||
query_set = QuerySet(Embedding).filter(knowledge_id__in=knowledge_id_list, is_active=is_active)
|
||
initial_count = query_set.count()
|
||
maxkb_logger.info(f"Initial embedding count: {initial_count}")
|
||
|
||
if exclude_document_id_list is not None and len(exclude_document_id_list) > 0:
|
||
query_set = query_set.exclude(document_id__in=exclude_document_id_list)
|
||
maxkb_logger.info(f"After excluding documents: {query_set.count()} embeddings")
|
||
if exclude_paragraph_list is not None and len(exclude_paragraph_list) > 0:
|
||
query_set = query_set.exclude(paragraph_id__in=exclude_paragraph_list)
|
||
maxkb_logger.info(f"After excluding paragraphs: {query_set.count()} embeddings")
|
||
query_set = query_set.exclude(**exclude_dict)
|
||
|
||
final_count = query_set.count()
|
||
maxkb_logger.info(f"Final embedding count before search: {final_count}")
|
||
|
||
for search_handle in search_handle_list:
|
||
if search_handle.support(search_mode):
|
||
maxkb_logger.info(f"Using search handler: {search_handle.__class__.__name__}")
|
||
results = search_handle.handle(query_set, query_text, query_embedding, top_n, similarity, search_mode)
|
||
maxkb_logger.info(f"Search results: {len(results)} items found")
|
||
if len(results) > 0:
|
||
maxkb_logger.info(f"Top result similarity: {results[0].get('similarity', 'N/A')}")
|
||
return results
|
||
|
||
maxkb_logger.warning("No suitable search handler found")
|
||
return []
|
||
|
||
def update_by_source_id(self, source_id: str, instance: Dict):
|
||
QuerySet(Embedding).filter(source_id=source_id).update(**instance)
|
||
|
||
def update_by_paragraph_id(self, paragraph_id: str, instance: Dict):
|
||
QuerySet(Embedding).filter(paragraph_id=paragraph_id).update(**instance)
|
||
|
||
def update_by_paragraph_ids(self, paragraph_id: str, instance: Dict):
|
||
QuerySet(Embedding).filter(paragraph_id__in=paragraph_id).update(**instance)
|
||
|
||
def delete_by_knowledge_id(self, knowledge_id: str):
|
||
QuerySet(Embedding).filter(knowledge_id=knowledge_id).delete()
|
||
|
||
def delete_by_knowledge_id_list(self, knowledge_id_list: List[str]):
|
||
QuerySet(Embedding).filter(knowledge_id__in=knowledge_id_list).delete()
|
||
|
||
def delete_by_document_id(self, document_id: str):
|
||
QuerySet(Embedding).filter(document_id=document_id).delete()
|
||
return True
|
||
|
||
def delete_by_document_id_list(self, document_id_list: List[str]):
|
||
if len(document_id_list) == 0:
|
||
return True
|
||
return QuerySet(Embedding).filter(document_id__in=document_id_list).delete()
|
||
|
||
def delete_by_source_id(self, source_id: str, source_type: str):
|
||
QuerySet(Embedding).filter(source_id=source_id, source_type=source_type).delete()
|
||
return True
|
||
|
||
def delete_by_paragraph_id(self, paragraph_id: str):
|
||
QuerySet(Embedding).filter(paragraph_id=paragraph_id).delete()
|
||
|
||
def delete_by_paragraph_ids(self, paragraph_ids: List[str]):
|
||
QuerySet(Embedding).filter(paragraph_id__in=paragraph_ids).delete()
|
||
|
||
|
||
class ISearch(ABC):
|
||
@abstractmethod
|
||
def support(self, search_mode: SearchMode):
|
||
pass
|
||
|
||
@abstractmethod
|
||
def handle(self, query_set, query_text, query_embedding, top_number: int,
|
||
similarity: float, search_mode: SearchMode):
|
||
pass
|
||
|
||
|
||
class EmbeddingSearch(ISearch):
|
||
def handle(self,
|
||
query_set,
|
||
query_text,
|
||
query_embedding,
|
||
top_number: int,
|
||
similarity: float,
|
||
search_mode: SearchMode):
|
||
from common.utils.logger import maxkb_logger
|
||
maxkb_logger.info(f"EmbeddingSearch: Executing search with similarity threshold={similarity}, top_n={top_number}")
|
||
|
||
# 先查询所有结果不设置相似度阈值，看看实际的相似度是多少
|
||
test_sql = """
|
||
SELECT
|
||
paragraph_id,
|
||
comprehensive_score,
|
||
comprehensive_score as similarity
|
||
FROM
|
||
(
|
||
SELECT DISTINCT ON
|
||
("paragraph_id") ( 1 - distince ),* ,(1 - distince) AS comprehensive_score
|
||
FROM
|
||
( SELECT *, ( embedding.embedding::vector(%s) <=> %s ) AS distince FROM embedding ${embedding_query} ORDER BY distince) TEMP
|
||
ORDER BY
|
||
paragraph_id,
|
||
distince
|
||
) DISTINCT_TEMP
|
||
ORDER BY comprehensive_score DESC
|
||
LIMIT %s
|
||
"""
|
||
|
||
test_exec_sql, test_exec_params = generate_sql_by_query_dict({'embedding_query': query_set},
|
||
select_string=test_sql,
|
||
with_table_name=True)
|
||
|
||
# 查询不带阈值的结果
|
||
test_results = select_list(test_exec_sql, [
|
||
len(query_embedding),
|
||
json.dumps(query_embedding),
|
||
*test_exec_params,
|
||
10 # 获取前10个结果
|
||
])
|
||
|
||
if len(test_results) > 0:
|
||
test_similarities = [r.get('similarity', 0) for r in test_results[:5]]
|
||
maxkb_logger.info(f"Actual similarities (no threshold): {test_similarities}")
|
||
maxkb_logger.info(f"Highest similarity: {test_similarities[0] if test_similarities else 0}, Required threshold: {similarity}")
|
||
if test_similarities[0] < similarity:
|
||
maxkb_logger.warning(f"Best similarity {test_similarities[0]} is below threshold {similarity}")
|
||
# 获取段落内容看看
|
||
if len(test_results) > 0:
|
||
paragraph_id = test_results[0].get('paragraph_id')
|
||
from knowledge.models import Paragraph
|
||
para = QuerySet(Paragraph).filter(id=paragraph_id).first()
|
||
if para:
|
||
maxkb_logger.info(f"Top paragraph content preview (first 200 chars): {para.content[:200]}...")
|
||
maxkb_logger.info(f"Paragraph title: {para.title}, length: {len(para.content)}")
|
||
else:
|
||
maxkb_logger.warning("No embeddings found even without similarity threshold")
|
||
|
||
# 正常查询（带相似度阈值）
|
||
exec_sql, exec_params = generate_sql_by_query_dict({'embedding_query': query_set},
|
||
select_string=get_file_content(
|
||
os.path.join(PROJECT_DIR, "apps", "knowledge", 'sql',
|
||
'embedding_search.sql')),
|
||
with_table_name=True)
|
||
|
||
maxkb_logger.debug(f"EmbeddingSearch SQL params count: {len(exec_params)}")
|
||
|
||
embedding_model = select_list(exec_sql, [
|
||
len(query_embedding),
|
||
json.dumps(query_embedding),
|
||
*exec_params,
|
||
similarity,
|
||
top_number
|
||
])
|
||
|
||
maxkb_logger.info(f"EmbeddingSearch results: {len(embedding_model)} embeddings found (with threshold)")
|
||
if len(embedding_model) > 0:
|
||
similarities = [e.get('similarity', 0) for e in embedding_model[:3]]
|
||
maxkb_logger.info(f"Top 3 similarities above threshold: {similarities}")
|
||
|
||
return embedding_model
|
||
|
||
def support(self, search_mode: SearchMode):
|
||
return search_mode.value == SearchMode.embedding.value
|
||
|
||
|
||
class KeywordsSearch(ISearch):
|
||
def handle(self,
|
||
query_set,
|
||
query_text,
|
||
query_embedding,
|
||
top_number: int,
|
||
similarity: float,
|
||
search_mode: SearchMode):
|
||
exec_sql, exec_params = generate_sql_by_query_dict({'keywords_query': query_set},
|
||
select_string=get_file_content(
|
||
os.path.join(PROJECT_DIR, "apps", "knowledge", 'sql',
|
||
'keywords_search.sql')),
|
||
with_table_name=True)
|
||
embedding_model = select_list(exec_sql, [
|
||
to_query(query_text),
|
||
*exec_params,
|
||
similarity,
|
||
top_number
|
||
])
|
||
return embedding_model
|
||
|
||
def support(self, search_mode: SearchMode):
|
||
return search_mode.value == SearchMode.keywords.value
|
||
|
||
|
||
class BlendSearch(ISearch):
|
||
def handle(self,
|
||
query_set,
|
||
query_text,
|
||
query_embedding,
|
||
top_number: int,
|
||
similarity: float,
|
||
search_mode: SearchMode):
|
||
exec_sql, exec_params = generate_sql_by_query_dict({'embedding_query': query_set},
|
||
select_string=get_file_content(
|
||
os.path.join(PROJECT_DIR, "apps", "knowledge", 'sql',
|
||
'blend_search.sql')),
|
||
with_table_name=True)
|
||
embedding_model = select_list(exec_sql, [
|
||
len(query_embedding),
|
||
json.dumps(query_embedding),
|
||
to_query(query_text),
|
||
*exec_params, similarity,
|
||
top_number
|
||
])
|
||
return embedding_model
|
||
|
||
def support(self, search_mode: SearchMode):
|
||
return search_mode.value == SearchMode.blend.value
|
||
|
||
|
||
search_handle_list = [EmbeddingSearch(), KeywordsSearch(), BlendSearch()]
|