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

464 阅读6分钟

前言

在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了,估计在下个版本就可以直接使用了。