[AIGC][Dify开发]添加并实现新的模型

前言

在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, HTTPServerclass 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():#获取百度的tokendef 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.textif __name__ == '__main__':host = "0.0.0.0"port = 8899server = 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) #代码1try:results = wenxin_rerank.rerank(model, query, docs, top_n) #代码1rerank_documents = []for result in results["results"]:     #代码3index = 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()  #代码1url = f"{self.api_bases[model]}?access_token={access_token}"  #代码2try:response = httpx.post(url,json={"model": model, "query": query, "documents": docs, "top_n": top_n},headers={"Content-Type": "application/json"},)                                                   #代码3response.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，使用该模型即可。