Langgraph实战--自定义embeding

概述

在Langgraph中我想使用第三方的embeding接口来实现文本的embeding。但目前langchain只提供了两个类，一个是AzureOpenAIEmbeddings，一个是：OpenAIEmbeddings。通过ChatOpenAI无法使用第三方的接口，例如：硅基流平台的接口。只能自己封装一个类，继承Embeding接口，从而实现整合第三方平台Embending API的能力。

实现思路

通过继承和实现langchain_core.embeddings的Embeddings类，并实现文本嵌入和查询接口。

在实现嵌入类时，需要实现embed_documents和embed_query两个接口。

import requests
import os
from typing import List
from langchain_core.embeddings import Embeddings
from dotenv import load_dotenv

class CustomSiliconFlowEmbeddings(Embeddings):
    def __init__(
        self,
        api_key: str,
        base_url: str = "https://api.siliconflow.cn/v1/embeddings",
        model: str = "BAAI/bge-large-zh-v1.5"
    ):
        self.api_key = api_key
        self.base_url = base_url
        self.model = model

    def embed_documents(self, texts: List[str]) -> List[List[float]]:
        """Embed a list of documents."""
        embeddings = []
        for text in texts:
            embedding = self.embed_query(text)
            embeddings.append(embedding)
        return embeddings

    def embed_query(self, text: str) -> List[float]:
        """Embed a query."""
        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json"
        }
        
        payload = {
            "model": self.model,
            "input": text,
            "encoding_format": "float"
        }
        
        response = requests.post(
            self.base_url,
            json=payload,
            headers=headers
        )
        
        if response.status_code == 200:
            return response.json()["data"][0]["embedding"]
        else:
            raise Exception(f"Error in embedding: {response.text}")

使用CustomSiliconFlowEmbeddings嵌入类

使用时，需要设置api_key的值，和模型名称，以及base_url等参数。

# Load environment variables
load_dotenv()
SL_API_KEY = os.getenv("SL_API_KEY")

# Initialize embedding model
embedding_model = CustomSiliconFlowEmbeddings(
    base_url="https://api.siliconflow.cn/v1/embeddings",
    api_key=SL_API_KEY,
    model="BAAI/bge-large-zh-v1.5"
)

# Test the embedding
if __name__ == "__main__":
    test_text = "您好世界！"
    result = embedding_model.embed_query(test_text)
    print(f"Embedding dimension: {len(result)}")
    print(f"First few values: {result[:10]}")

    # 获取网页中的数据，并进行分割，然后存储到FAISS中
    urls = [
    "https://lilianweng.github.io/posts/2023-06-23-agent/",
    "https://lilianweng.github.io/posts/2023-03-15-prompt-engineering/",
    "https://lilianweng.github.io/posts/2023-10-25-adv-attack-llm/"
    ]

    docs = [WebBaseLoader(url).load() for url in urls]
    docs_list = [item for sublist in docs for item in sublist]
    text_splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(chunk_size=250, chunk_overlap=0)
    doc_splits = text_splitter.split_documents(docs_list)

    vectorstore = FAISS.from_documents(documents=doc_splits, embedding=embedding_model)
    retriever = vectorstore.as_retriever()

    # 测试检索功能，查询与问题最相关的分块文档
    resp = retriever.invoke("什么是prompt engineering？")
    # 返回的是一个个Document对象
    for doc in resp:
        print(doc.id + ": " + doc.page_content)

输出：

Embedding dimension: 1024
First few values: [0.021915348, 0.0048826355, -0.09566349, -0.010307786, -0.0025656442, 0.043084737, -0.045955546, 0.011641469, 0.02809776, -0.012489148]

参考资料

• https://docs.siliconflow.cn/cn/api-reference/embeddings/create-embeddings