如何优雅地对Python第三方库做二次开发

今天，有同学给我提了一个问题：如何在Simplemind中接入Azure的GPT接口。如下图所示。

在使用Python时经常会出现这样的情况，某一个第三方库，满足我们99%的需求，但碰巧有一个小需求不满足。遇到这种情况，有些同学会忍痛割爱，换一个库；还有一些同学，会继续使用这个第三方库，但是缺的那个功能，他就完全自己单独写；剩下的同学，可能是把这个第三方库下载下来，放到自己项目的根目录中，然后当做项目的一部分来修改并导入使用。今天我们就来讲一下这个问题。

前两个方法不需要多说什么。第三个方法从功能上来说没什么问题，但会给自己的项目引入大量其他代码，导致项目在做安全性检查、静态类型检查、Code Review时变得很麻烦。而且这个第三方库必须放到项目的根目录，否则在导入时，它的导入语句就跟正常pip安装的导入语句不一样，以后如果官方库支持了这个缺失的功能，你得改很多个导入语句，才能再换回来，无形中引入了很多的不确定性和隐患。

我们今天想实现的功能是，调用这个二次开发的第三方库时，我自己的代码不需要做任何修改，甚至包括环境变量也不需要修改，直接像是调用任何pip安装的第三方库一样使用。

实际上，在pip设计的时候，就已经预料到了这种情况。所以pip install有一个-e参数，可以用来指定某个特定文件夹里面的代码为一个可编辑的第三方库。对这个文件夹里面的所有修改会立刻生效，同时对于使用这个第三方库的代码来说，它不需要做任何修改，就像是在用正常的第三方库一样。它原本是用来方便在开发者自己写第三方库时，测试功能调用的，现在我们对现有的第三方库做二次开发，正好也可以使用它。

就以上面这个问题为例，来说明如何对Simplemind进行二次开发。Simplemind目前支持的大模型如下图所示：

其中的openai.py代码如下，可以看到它初始化OpenAI连接对象时，只使用了api_key参数。因此Simplemind目前只支持OpenAI官方的GPT模型，无法使用Azure提供的GPT模型。

要使用Azure的GPT连接对象，我们需要使用如下的代码：

from openai.lib.azure import AzureOpenAI
client = AzureOpenAI(api_key=..., azure_endpoint=..., api_version=...)

因为Azure的GPT和OpenAI的GPT除了初始化的参数不同，其他调用上的代码完全相同，因此我们可以继承openai.py 中的这个OpenAI类，然后自己只需要复写def client这个属性(注意，这里使用了@cached_property，所以它不是方法，而是属性)，就可以让Simplemind支持Azure的GPT了。

来看看具体的实现方法。从Github上面克隆Simplemind的代码到本地，然后把它安装成可编辑的第三方库：

git clone git@github.com:kennethreitz/simplemind.git
cd simplemind
pip install -e .

这三行代码就够了，这个时候，你在PyCharm中输入import simplemind，会发现可以正常导入。如果你有OpenAI官方的API，那么你可以直接使用Simplemind文档中的代码，立刻测试，会发现它和pip安装的没有任何区别。

现在，我们打开刚刚克隆下来的simplemind/simplemind/providers文件夹，创建一个azure_openai.py文件。里面的代码如下：

from .openai import OpenAI
import os
from functools import cached_property

class AzureOpenAI(OpenAI):
    NAME = 'azure_openai'
    def __init__(self, api_key: str | None = None):
        super().__init__(api_key=api_key)
        self.api_key = os.getenv('OPENAI_API_KEY')
        self.azure_endpoint='你的AzureGPT的url'
        self.api_version = '2024-07-01-preview'


    @cached_property
    def client(self):
        """The raw OpenAI client."""
        if not self.api_key:
            raise ValueError("OpenAI API key is required")
        try:
            from openai.lib.azure import AzureOpenAI
        except ImportError as exc:
            raise ImportError(
                "Please install the `openai` package: `pip install openai`"
            ) from exc
        return AzureOpenAI(api_key=self.api_key, azure_endpoint=self.azure_endpoint, api_version=self.api_version)

如下图所示。

然后编辑这个文件夹里面的__init__.py文件，在其中添加上刚创建的这个新类，如下图所示。

改好了，以上就是全部的修改。现在开始编写调用代码，跟官方文档中的示例完全一样：

import simplemind as sm
from dotenv import load_dotenv


load_dotenv()

resp = sm.generate_text(prompt='太阳为什么是圆的？', llm_model='gpt-4o-mini', llm_provider='azure_openai')
print(resp)

运行效果如下图所示，成功接上了Azure的GPT。

再来测试一下文档里面的记忆功能和工具调用，也全部正常运行：

国产大模型基本都支持直接使用openai库调用，因此理论上使用这个方法，稍作修改，可以接入任意国产大模型。如果你改成使用LiteLLM，甚至可以实现支持任意大模型。

题图 from RealPython