前言
在dify中预置了很多模型供应商和众多主流模型,但是现今AI发展迅速,新的模型不停的涌现,如何在dify中使用这些新的模型?
OneApi
如果你有自己的OneApi平台,那么任何LLM和TextEmbedding等模型都可以通过OneApi平台来添加到dify中。
先在OneApi的渠道中添加新模型,然后在令牌中也添加上。
在dify的模型供应商那里选择 OpenAI-API-compatible,然后添加模型即可。模型名称填OneAPI中的模型名称,API Key填OneApi的令牌,API endpoint URL填OneApi的地址(注意带/v1)。这样在dify中就可以使用这个新模型了。
其他类似的平台也可以实现。
接口转发
OneApi因为是基于OpenAI接口的,所以有它的局限性,比如不支持rerank模型。另外可能也没有自己的OneApi平台,那么第二种方法就是自己实现一个简易的服务来进行转发。
这里以文心中的rerank模型为例。
文心提供了一款rerank模型:bce-reranker-base,它是网易有道的。在dify当前(0.9.2版本)版本中文心这个模型提供商下并没有rerank这个类型,所以无法直接添加。
但是看了一遍所有供应商,发现LocalAI是支持rerank模型的,那么就可以自己实现一个服务,实现对bce-reranker-base这个模型的调用,并且将输入和输出进行重新格式化,以符合LocalAI的rerank接口格式即可。
经过比对发现输入格式是一致的,只不过输出有一点不一样,只要简单处理一下就可以了。简单的实例代码如下:
import json
import requests
from http.server import BaseHTTPRequestHandler, HTTPServer
class PostHandler(BaseHTTPRequestHandler):
def do_POST(self):
if self.path == '/rerank':
content_length = int(self.headers['Content-Length'])
post_data = self.rfile.read(content_length)
result = rerank(post_data)
resultJson = json.loads(result)
results = resultJson['results']
for item in results:
doc = item['document']
del item['document']
item['document'] = {'text':doc}
result = json.dumps(resultJson)
print(resultJson)
self.send_response(200)
self.send_header("Content-type", "application/json")
self.end_headers()
self.wfile.write(bytes(result, encoding='utf-8'))
def get_access_token():
#获取百度的token
def rerank(payload):
url = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/reranker/bce_reranker_base?access_token=" + get_access_token()
headers = {
'Content-Type': 'application/json'
}
response = requests.request("POST", url, headers=headers, data=payload)
return response.text
if __name__ == '__main__':
host = "0.0.0.0"
port = 8899
server = HTTPServer((host, port), PostHandler)
print(f"Server is running on http://{host}:{port}")
server.serve_forever()
然后在dify中添加LocalAi模型,注意模型名称是rerank。
注意上面的代码只是简单的示例,实用中要考虑并发性,所以最好通过aiohttp和线程池搭配来实现服务。
dify二次开发
最直接的方法是对dify进行二次开发,然后本地部署。那么该怎么开发?这里还是以文心的rerank模型:bce-reranker-base为例。
模型供应商的源码在项目的api/core/model_runtime/model_providers/目录下,在这里可以看到每个供应商有一个目录,其中百度文心的目录是 wenxin。
添加rerank类型
在这个目录下可以看到两个目录llm和text_embedding。所以我们首先要添加rerank这个类型,创建一个名字为rerank的目录。
在rerank目录下创建一个空的__init__.py文件和一个rerank.py文件,我们需要在rerank.py文件中来实现对rerank模型的调用。
RerankModel
在rerank.py中新建一个class:WenxinRerankModel,它继承RerankModel,并实现它(以及它继承的AIModel)的几个函数。
- _invoke:实现对模型的请求和返回
- validate_credentials:验证模型供应商,可以做一次模型请求,也可以验证密钥(比如文心就验证通过密钥获取token)。当添加模型后会自动执行这一步来验证,验证成功才添加成功。
- _invoke_error_mapping:枚举各种错误
所以最主要的就是_invoke函数,这里实现了请求,不同供应商不一样,wenxin的代码如下:
def _invoke(
self,
model: str,
credentials: dict,
query: str,
docs: list[str],
score_threshold: Optional[float] = None,
top_n: Optional[int] = None,
user: Optional[str] = None,
) -> RerankResult:
if len(docs) == 0:
return RerankResult(model=model, docs=[])
api_key = credentials["api_key"]
secret_key = credentials["secret_key"]
wenxin_rerank: WenxinRerank = WenxinRerank(api_key, secret_key) #代码1
try:
results = wenxin_rerank.rerank(model, query, docs, top_n) #代码1
rerank_documents = []
for result in results["results"]: #代码3
index = result["index"]
if "document" in result:
text = result["document"]
else:
text = docs[index]
rerank_document = RerankDocument(
index=index,
text=text,
score=result["relevance_score"],
)
if score_threshold is None or result["relevance_score"] >= score_threshold:
rerank_documents.append(rerank_document)
return RerankResult(model=model, docs=rerank_documents)
except httpx.HTTPStatusError as e:
raise InternalServerError(str(e))
代码1创建了一个WenxinRerank的对象,这个类包装了请求,我们后面再说。
代码2执行了rerank方法,实际上就是发送请求并获取结果。
代码3这个for循环就是对结果进行重新格式化,转成dify需要的格式。
最后返回RerankResult即可。
WenxinRerank
上面代码中的WenxinRerank类的代码如下:
class WenxinRerank(_CommonWenxin):
def rerank(self, model: str, query: str, docs: list[str], top_n: Optional[int] = None):
access_token = self._get_access_token() #代码1
url = f"{self.api_bases[model]}?access_token={access_token}" #代码2
try:
response = httpx.post(
url,
json={"model": model, "query": query, "documents": docs, "top_n": top_n},
headers={"Content-Type": "application/json"},
) #代码3
response.raise_for_status()
return response.json()
except httpx.HTTPStatusError as e:
raise InternalServerError(str(e))
这个类需要继承_CommonWenxin,这是一个已有的类,里面封装了百度文心的一些基本api和工具。
代码1调用_CommonWenxin的_get_access_token获取token
代码2组合模型的url,这里的api_bases是在_CommonWenxin中,是一个Map,key模型名,value是模型的地址url。
代码3执行请求得到返回。
启动rerank类型
最后要修改wenxin这个目录下的wenxin.yaml,在supported_model_types加上rerank,如下:
supported_model_types:
- llm
- text-embedding
- rerank
这样重新编译运行api代码后,在文心这个模型提供商下就可以看到rerank类型了。
添加模型
上面我们为百度文心添加了rerank类型,实现了这类模型的请求,但是还没有任何模型,下面就是添加模型。
我们添加的模型是bce-reranker-base,它的模型名是 bce-reranker-base_v1,所以在rerank目录下创建一个bce-reranker-base_v1.yaml文件,内容如下:
model: bce-reranker-base_v1
model_type: rerank
model_properties:
context_size: 4096
pricing:
input: '0.0005'
unit: '0.001'
currency: RMB
比较直观,就不一个一个的说了。
然后需要在_CommonWenxin的api_bases中添加模型的地址url,如下:
class _CommonWenxin:
api_bases = {
...
"bce-reranker-base_v1": "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/reranker/bce_reranker_base",
}
注意bce-reranker-base_v1这个模型名要与刚才yaml中的保持一致。
然后我们再重新编译运行api,就可以在百度文心这个供应商下看到这个模型了,启动后就可以使用了。比如在知识检索功能中将召回设置为rerank,使用该模型即可。
总结
这里我们是对dify已有的模型供应商添加新分类和新模型,所以有不少现成的工具来使用。如果要添加一个全新的模型供应商,会稍微更复杂一点,但是其实原理差不多,其实就是实现模型的请求。
另外,百度文心的rerank模型我已经提交pr给dify官方了,目前已经merge了,估计在下个版本就可以直接使用了。