目录
- 前言
- OneApi
- 接口转发
- dify二次开发
- 添加rerank类型
- RerankModel
- WenxinRerank
- 启动rerank类型
- 添加模型
- 总结
前言
在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,使用该模型即可。
总结
这里我们是对dify已有的模型供应商添加新分类和新模型,所以有不少现成的工具来使用。如果要添加一个全新的模型供应商,会稍微更复杂一点,但是其实原理差不多,其实就是实现模型的请求。
另外,百度文心的rerank模型我已经提交pr给dify官方了,目前已经merge了,估计在下个版本就可以直接使用了。