feat: rag

rag/query.py

@@ -0,0 +1,118 @@
+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}")