如果内置的嵌入函数不能满足你的需要,你可以选择自定义嵌入函数。
这是最灵活也最常用的方式 —— 通常公司里一般有自研的 embedding 服务、或者调用第三方提供的服务,直接包一层就能用,抽取到单独的文件中,其他项目还可以复用。
前面“内置嵌入函数一览”介绍了很多内置函数,自定义最简单的方式就是参考内置实现。
一个嵌入函数的最小实现,只需要一个类,实现 __call__,如下:
from chromadb import Documents, EmbeddingFunction, Embeddings
class MyEmbedding(EmbeddingFunction[Documents]):
def __call__(self, input: Documents) -> Embeddings:
# input 是 List[str]
# 返回 List[List[float]]
return None下面以使用阿里云远程 API 为例,实现自定义嵌入函数:
import chromadb
import os
import tempfile
import numpy as np
import requests
from chromadb import Documents, EmbeddingFunction, Embeddings
from dotenv import load_dotenv
# 持久化示例写到临时目录,工作目录保持干净
os.chdir(tempfile.mkdtemp(prefix="chroma_demo_"))
# 从 .env 读取环境变量 OPENAI_API_KEY
load_dotenv()
# 自定义嵌入函数
class MyEmbedding(EmbeddingFunction[Documents]):
def __call__(self, input: Documents) -> Embeddings:
# input 是 List[str]
# 返回 List[List[float]]
resp = requests.post("https://llm-0sb8vr1k7aav4kn3.cn-beijing.maas.aliyuncs.com/compatible-mode/v1/embeddings",
json={
"model": "qwen3.7-text-embedding-flash",
"input": input,
"encoding_format": "float",
}, headers={
"Authorization": f"Bearer {os.getenv("OPENAI_API_KEY", "")}",
"Content-Type": "application/json",
}, timeout=60)
resp.raise_for_status()
data = resp.json()
return [
np.array(item["embedding"], dtype=np.float32) for item in data["data"]
]
if __name__ == "__main__":
ef = MyEmbedding()
client = chromadb.PersistentClient(path="./chroma_data")
col = client.get_or_create_collection(
name="kb_zh",
embedding_function=ef,
metadata={"hnsw:space": "cosine"}
)
col.add(
ids=["a", "b"],
documents=["如何申请退款", "快递丢了怎么办"]
)
res = col.query(
query_texts=["买的东西不想要了,钱能退回来吗"],
n_results=1
)
print(res)注意,在项目根目录创建 .env,添加如下内容:
# OpenAI API Key
OPENAI_API_KEY=sk-8e20ca1xxxxxxxxxxxxxxxxx012e8ffc3运行示例,输出如下:
D:\xxx\chat07_7.7_demo1.py:45: DeprecationWarning: The class MyEmbedding does not implement __init__. This will be required in a future version.
ef = MyEmbedding()
MyEmbedding 这个类没有实现 __init__ 方法。在未来的版本中,这会变成强制要求(会报错)。 ef = MyEmbedding()
{'ids': [['a']], 'embeddings': None, 'documents': [['如何申请退款']], 'uris': None,
'included': ['metadatas', 'documents', 'distances'], 'data': None,
'metadatas': [[None]], 'distances': [[0.411679744720459]]}从运行效果来看,成功实现了嵌入函数,也能正常工作。但如果只做到这一步,有两个问题:
这个配置不会被持久化,换个进程 get_collection 就找不回来
服务端不知道你用的什么模型
为了解决最小实现的问题,下面给出一个完整版本。补上 name()、get_config()、build_from_config(),再加注册装饰器,模板如下:
from typing import Dict, Any
from chromadb import Documents, EmbeddingFunction, Embeddings
from chromadb.utils.embedding_functions import register_embedding_function
# 类名字根据自己业务需要进行修改,这里叫 AliyunEmbeddingFunction,因为内部调用阿里云服务实现
# @register_embedding_function 用来注册进全局表,否则反序列化时找不到类
@register_embedding_function
class AliyunEmbeddingFunction(EmbeddingFunction[Documents]):
def __init__(self, dim: int = 64):
#
def __call__(self, input: Documents) -> Embeddings:
#
# 下面三个是"可持久化"的关键
@staticmethod
def name() -> str:
# 注册名,全局唯一,存到集合配置里
def get_config(self) -> Dict[str, Any]:
# 要序列化的参数(必须能 JSON 化,所以别塞模型对象)
@staticmethod
def build_from_config(config: Dict[str, Any]) -> "AliyunEmbeddingFunction":
# 从配置反构造实例注意,使用 @register_embedding_function 注册之后,可以通过“"aliyum_df" in known_embedding_functions”验证一下,注册是否成功。
完整代码如下:
import chromadb
import os
import tempfile
import numpy as np
import requests
from typing import Dict, Any
from chromadb import Documents, EmbeddingFunction, Embeddings
from chromadb.utils.embedding_functions import register_embedding_function, known_embedding_functions
from dotenv import load_dotenv
# 持久化示例写到临时目录,工作目录保持干净
os.chdir(tempfile.mkdtemp(prefix="chroma_demo_"))
# 从 .env 读取环境变量 OPENAI_API_KEY
load_dotenv()
@register_embedding_function
class AliyunEmbeddingFunction(EmbeddingFunction[Documents]):
"""演示用:一个确定性的哈希嵌入,不依赖任何模型"""
def __init__(self, dim: int = 64):
self.dim = dim
def __call__(self, input: Documents) -> Embeddings:
print(f"__call__ => {input}")
resp = requests.post("https://llm-0sb8vr1k7aav4kn3.cn-beijing.maas.aliyuncs.com/compatible-mode/v1/embeddings",
json={
"model": "qwen3.7-text-embedding-flash",
"input": input,
"encoding_format": "float",
}, headers={
"Authorization": f"Bearer {os.getenv("OPENAI_API_KEY", "")}",
"Content-Type": "application/json",
}, timeout=60)
resp.raise_for_status()
data = resp.json()
return [
np.array(item["embedding"], dtype=np.float32) for item in data["data"]
]
# 下面三个是"可持久化"的关键
@staticmethod
def name() -> str:
return "aliyum_df"
def get_config(self) -> Dict[str, Any]:
return {"dim": self.dim}
@staticmethod
def build_from_config(config: Dict[str, Any]) -> "AliyunEmbeddingFunction":
return AliyunEmbeddingFunction(dim=config.get("dim", 64))
if __name__ == "__main__":
# 查看自定义嵌入函数是否存在
print(f"自定义嵌入函数是否存在?{"aliyum_df" in known_embedding_functions}")
# 使用自定义的嵌入函数
# 为了演示自定义嵌入函数可以恢复,下面将创建两个客户端
ef = AliyunEmbeddingFunction(dim=32)
# 客户端1,用来添加数据到集合
client = chromadb.PersistentClient(path="./chroma_data")
col = client.get_or_create_collection(
name="kb_zh",
embedding_function=ef,
metadata={"hnsw:space": "cosine"}
)
col.add(
ids=["a", "b"],
documents=["如何申请退款", "快递丢了怎么办"]
)
# 客户端2,重新打开一个客户端,用来从结合中查询数据
client2 = chromadb.PersistentClient(path="./chroma_data")
col2 = client2.get_or_create_collection(name="kb_zh")
res = col2.query(
query_texts=["买的东西不想要了,钱能退回来吗"],
n_results=1
)
print(res)运行代码,输出如下:
自定义嵌入函数是否存在?True
__call__ => ['如何申请退款', '快递丢了怎么办']
__call__ => ['买的东西不想要了,钱能退回来吗']
{'ids': [['a']], 'embeddings': None, 'documents': [['如何申请退款']], 'uris': None,
'included': ['metadatas', 'documents', 'distances'], 'data': None,
'metadatas': [[None]], 'distances': [[0.404982328414917]]}从输出可以看出,__call__ 被调用了两次,第一次是添加数据到集合,第二次是从数据查询数据。
通常,将自定义的嵌入函数单独放进一个 python 文件(比如 myapp/embeddings.py),入口文件里 import 一下,其他地方就可以直接使用了。