ZhouYI/02_RAG/agent/rag/embedding.py

from langchain_community.embeddings import DashScopeEmbeddings
from agent.config import load_config
from langchain_community.vectorstores import Chroma
import os
def get_embedding_model():
    """
    获取 Embedding 模型。

    """
    config = load_config()
    return DashScopeEmbeddings(
        model="text-embedding-v4",
        dashscope_api_key=config.api_key,  # 替换为你的 API Key
    )

def get_vectorstore(embedding_model, documents: list):
    if "./chroma_db" in os.listdir():
        db = Chroma(
            persist_directory="./chroma_db",
            embedding_function=embedding_model,
        )
    else:
        db = Chroma.from_documents(
            documents=documents,
            embedding=embedding_model,
            collection_metadata={"hnsw:space": "cosine"},
            persist_directory="./chroma_db"
        )
        db.persist()
    return db

def query_vectorstore(query: str)-> list:
    """
    查询向量数据库。

    参数:
        query (str): 用户的查询问题。

    返回:
        list: 与查询最相关的文档列表。
    """
    embedding_model = get_embedding_model()
    db = get_vectorstore(embedding_model, [])
    retriever = db.as_retriever(search_kwargs={"k": 3})
    return retriever.get_relevant_documents(query)