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_dotenvclass 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_keyself.base_url = base_urlself.model = modeldef 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 embeddingsdef 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]