用langchain 通过text-embedding-3-small生成embedding

13 阅读2分钟

OpenAIEmbedding介绍

OpenAIEmbedding提供了两个主要方法:

  • embed_query: 用于生成单个查询文本的向量(通常用于搜索时的查询端)。
  • embed_documents: 用于批量生成多个文档文本的向量(通常用于构建知识库时的文档端)。

实现步骤

    1. 申请openai的key
    1. 安装langchain
    1. 构建生成embedding的json格式的文档
    1. 调用模型生成embedding

具体代码

批量生成向量

from langchain_openai import OpenAIEmbeddings
from app.config import EnvConfig
from common.Logger import getLogger
logger = getLogger("embedding_generator.py")

def bulk_generate_user_embedding(documents:list[str]):
    """
    :param documents: 生成用于生成embeddings的结构化文档
    :return: 返回批量生成的向量列表
    """
    api_key = EnvConfig.OPENAI_API_KEY
    openai_api_base = EnvConfig.OPENAI_API_BASE
    if not api_key or not openai_api_base:
        logger.error("未找到 OPENAI_API_KEY 环境变量。")
        return None

    # 2. 初始化模型
    # model: 指定模型名称,如 'text-embedding-3-small'
    # temperature: 控制输出的随机性,0 为最确定,1 为最随机
    embeddings = OpenAIEmbeddings(
        model="text-embedding-3-small",
        api_key=api_key,
        openai_api_base = openai_api_base
    )

    # 4. 批量生成文档向量
    doc_vectors = embeddings.embed_documents(documents)
    print(f"\n生成了 {len(doc_vectors)} 个文档向量")

    if len(doc_vectors)>0:
        print(f"每个向量维度: {len(doc_vectors[0])}")

    return doc_vectors

生成单个向量

def generate_user_embedding(query_text:str):
    """
    :param query_text: 准备的格式化JSON文件
    :return:
    """
    api_key = EnvConfig.OPENAI_API_KEY
    openai_api_base = EnvConfig.OPENAI_API_BASE
    if not api_key or not openai_api_base:
        logger.error("未找到 OPENAI_API_KEY 环境变量。")
        return None


    # 2. 初始化模型
    # model: 指定模型名称,这里用 'text-embedding-3-small'
    # temperature: 控制输出的随机性,0 为最确定,1 为最随机
    embeddings = OpenAIEmbeddings(
        model="text-embedding-3-small",
        api_key=api_key,
        openai_api_base = openai_api_base
    )


    # 3. Embedding查询
    query_vector = embeddings.embed_query(query_text)
    logger.debug(f"查询向量维度: {len(query_vector)}")
    logger.debug(f"查询向量前5个值: {query_vector[:5]}")
    return query_vector

生成的构建化文档方法如下

def build_profile_text(profile) -> str:
    """构建用于向量化的结构化文本"""
    parts = [
        f"角色: {profile.role}",
        f"岗位: {profile.title}",
        f"介绍: {profile.bio}",
        f"行业: {profile.industry_tags}",
        f"城市: {profile.city}",
        f"技能: {profile.skill_tags}"
    ]
    return " | ".join(parts)