heyong.fu
119 lines
3.1 KiB
Python
119 lines
3.1 KiB
Python
import os
|
|
from typing import Optional, List
|
|
import logging
|
|
|
|
import chromadb
|
|
from sentence_transformers import SentenceTransformer
|
|
|
|
from llm import get_doubao_llm
|
|
|
|
logging.basicConfig(
|
|
level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s"
|
|
)
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# 默认集合的名称
|
|
DEFAULT_COLLECTION_NAME = "rag_system"
|
|
# 返回几条数据
|
|
DEFAULT_N_RESULTS = 2
|
|
# 默认向量化模型的名称
|
|
DEFAULT_MODEL_NAME = "all-MiniLM-L6-v2"
|
|
# 定义向量模型的全局变量
|
|
_mode: Optional[SentenceTransformer] = None
|
|
# 定义chromadb客户端
|
|
_client: Optional[chromadb.PersistentClient] = None
|
|
_collection: Optional[chromadb.Collection] = None
|
|
|
|
# 默认数据库存放路径
|
|
DEFAULT_DB_PATH = "./chroma_db"
|
|
|
|
|
|
def _get_model():
|
|
global _mode
|
|
if _mode is None:
|
|
_mode = SentenceTransformer(DEFAULT_MODEL_NAME)
|
|
return _mode
|
|
|
|
|
|
def _get_client():
|
|
global _client
|
|
if _client is None:
|
|
_client = chromadb.PersistentClient(path=DEFAULT_DB_PATH)
|
|
return _client
|
|
|
|
|
|
def get_query_embedding(query: str) -> List[float]:
|
|
model = _get_model()
|
|
embedding = model.encode([query])[0].tolist()
|
|
return embedding
|
|
|
|
|
|
def _get_collection(collection_name: str = DEFAULT_COLLECTION_NAME):
|
|
global _collection
|
|
if _collection is None:
|
|
client = _get_client()
|
|
_collection = client.get_or_create_collection(collection_name)
|
|
return _collection
|
|
|
|
|
|
def retrieve_relate_chunks(
|
|
query_embedding: List[float],
|
|
n_results: int = DEFAULT_N_RESULTS,
|
|
collection_name: str = DEFAULT_COLLECTION_NAME,
|
|
):
|
|
try:
|
|
collection = _get_collection(collection_name)
|
|
# print(n_results)
|
|
# 去指定集合查找相似度检索，找到数据
|
|
results = collection.query(
|
|
query_embeddings=[query_embedding], n_results=n_results
|
|
)
|
|
related_chunks = results.get("documents")
|
|
if not related_chunks or not related_chunks[0]:
|
|
raise ValueError("未找到相关内容")
|
|
return related_chunks[0]
|
|
|
|
except Exception as e:
|
|
logger.error(f"向量检索失败：{str(e)}")
|
|
raise
|
|
|
|
|
|
def query_rag(
|
|
query: str,
|
|
n_results: int = DEFAULT_N_RESULTS,
|
|
collection_name: str = DEFAULT_COLLECTION_NAME,
|
|
):
|
|
"""
|
|
查询函数：
|
|
query：用户查询的问题
|
|
n_results：检索数量
|
|
collection_name: 集合名字
|
|
"""
|
|
# 1. 将查询问题转为向量
|
|
query_embedding = get_query_embedding(query)
|
|
# print(query_embedding)
|
|
# 基于查询向量做检索
|
|
related_chunks = retrieve_relate_chunks(
|
|
query_embedding, n_results, collection_name=collection_name
|
|
)
|
|
# print("related_chunks", related_chunks)
|
|
content = "\n".join(related_chunks)
|
|
prompt = f"""
|
|
已知信息：{content}
|
|
请根据上述内容回答用户问题:{query}
|
|
"""
|
|
print(prompt)
|
|
answer = get_doubao_llm(prompt)
|
|
return answer
|
|
|
|
|
|
query = "西游记是谁写的"
|
|
|
|
try:
|
|
answer = query_rag(query, n_results=1)
|
|
print(f"答案:", answer)
|
|
except ValueError as e:
|
|
print(f"错误{e}")
|
|
except Exception as e:
|
|
print(f"错误{e}")
|