AI开发实战（002）：自定义实现 LangChang ChatModel 来集成 Claude 2 API 调用

先看一下效果，当然 OpenAI 还要留着，增加了我常用的 Claude 2 的 LLM 选择按钮。

接下来一步一步实现：

OpenAI 对接 LangChain

为了统一使用体验，我们后端使用 LangChain 来对接，先不用 OpenAI 的原生接口，稍微复杂一点， 但是也没复杂多少。

安装 LangChain

首先安装 LangChain：

pip install langchain

LangChain + OpenAI + FastAPI

增加接口，使用 LangChain 的 ChatOpenAI来实现问答：

@app.get("/chat/openai")

def ask(prompt: str = Query(..., max_length=40)):

    from langchain.schema import AIMessage, HumanMessage, SystemMessage

    from langchain.chat_models import ChatOpenAI

    chat = ChatOpenAI()

    messages = [

        SystemMessage(content="简单回答，使用中文"),

        HumanMessage(content=prompt)

    ]

    return TextMessage(chat(messages))

PostMan 测试

用 PostMan 测试下，效果很好：

自定义实现 Claude 2 LangChain 对接

Claude API

我们使用开源的 claude-api 来模拟调用 Claude 2 的 Web 接口。

先安装：

pip install claude-api

https://github.com/KoushikNavuluri/Claude-API

获取 Cookie，

把 Cookie 存储到 环境变量里面，后面用。

测试 Claude-API

我们编写 claude.py 来测试下接口：

from claude_api import Client

import os



cookie = os.environ.get('CLAUDE_COOKIE')

claude_api = Client(cookie)



conversation_id = claude_api.create_new_chat()['uuid']



def chat(prompt):

    response = claude_api.send_message(prompt, conversation_id)

    print(response)

    return response



if __name__ == '__main__':

    chat('帮我根据颜色  # 389e0d 生成一组颜色，只返回 json')

接口连接成功，得到如下输出：

实现 LangChain ChatModel 接口

接下来我们参考 ChatOpenAI 实现 ChatClaude。

先安装 pydantic， 代码里面用的到：

pip install pydantic

首先引入需要的包：

from pydantic import Field, root_validator



from langchain.schema.language_model import BaseLanguageModel

from langchain.utils import (

    get_from_dict_or_env,

)



from typing import Any, AsyncIterator, Dict, Iterator, List, Optional



from langchain.callbacks.manager import (

    CallbackManagerForLLMRun,

)

from langchain.chat_models.base import BaseChatModel

from langchain.schema import (

    ChatGeneration,

    ChatResult,

)

from langchain.schema.messages import (

    AIMessage,

    BaseMessage,

    ChatMessage,

    HumanMessage,

    SystemMessage,

)

然后实现基础类 _ClaudeCommon：

class _ClaudeCommon(BaseLanguageModel):

    HUMAN_PROMPT: Optional[str] = ''

    AI_PROMPT: Optional[str] = '简单回答'



    @property

    def lc_secrets(self) -> Dict[str, str]:

        return {"claude_cookie": "CLAUDE_COOKIE"}



    @property

    def lc_serializable(self) -> bool:

        return True



    @root_validator()

    def validate_environment(cls, values: Dict) -> Dict:

        values["claude_cookie"] = get_from_dict_or_env(values, "claude_cookie", "CLAUDE_COOKIE")

        values["conversation_id"] = get_from_dict_or_env(values, "conversation_id", "CONVERSATION_ID")



        try:

            from claude_api import Client

            claude_api = Client(values["claude_cookie"])

            values["client"] = claude_api

        except ImportError:

            raise ImportError(

                "Could not import ChatClaude python package. "

            )

        return values

实现 ChatClaude 类：

class ChatClaude(BaseChatModel, _ClaudeCommon):

    @property

    def _llm_type(self) -> str:

        """Return type of chat model."""

        return "claude-chat"



    def _convert_messages_to_text(self, messages: List[BaseMessage]) -> str:

        return "".join(

            self._convert_one_message_to_text(message) for message in messages

        )



    def _convert_messages_to_prompt(self, messages: List[BaseMessage]) -> str:

        messages = messages.copy()  # don't mutate the original list

        if not isinstance(messages[-1], AIMessage):

            messages.append(AIMessage(content=""))

        text = self._convert_messages_to_text(messages)

        return (

            text.rstrip()

        )  # trim off the trailing ' ' that might come from the "Assistant: "



    def _generate(

            self,

            messages: List[BaseMessage],

            stop: Optional[List[str]] = None,

            run_manager: Optional[CallbackManagerForLLMRun] = None,

            **kwargs: Any,

    ) -> ChatResult:

        prompt = self._convert_messages_to_prompt(messages)

        completion = self.client.send_message(prompt, self.conversation_id)

        print(completion)

        message = AIMessage(content=completion)

        return ChatResult(generations=[ChatGeneration(message=message)])

Claude2 + LangChain + FastAPI

增加 Claude 的接口如下：

@app.get("/chat/claude")

def ask(prompt: str = Query(..., max_length=40)):

    from langchain.schema import AIMessage, HumanMessage, SystemMessage

    from chat_models import ChatClaude

    chat = ChatClaude()

    messages = [

        SystemMessage(content="简单回答，使用中文"),

        HumanMessage(content=prompt)

    ]

    return TextMessage(chat(messages))

使用 PostMan 测试

使用下面的地址测试：

http://127.0.0.1:8000/chat/claude?prompt=中国的首都是哪里

结果完美。

前端

后端好了，那我们来改造后前端。

当然昨天的 OpenAI 成果是还要留着，我们增加一个按钮来切换 OpenAI 和 Claude 2 的使用。

添加 LLM 选择按钮

先增加两个按钮，用来切换 Claude 和 OpenAI，

<div class="div">

    <button class="button" onclick="claude('OpenAI')">OpenAI</button>

    <button class="button" onclick="claude('Claude')">Claude</button>

</div>

如下：

然后修改下样式，稍微好看点儿：

.div {

    position: fixed;

    top: 10px;

    right: 40px;

    z-index: 9999;

}



.button {

    background-image: linear-gradient(135deg, #42e695, #3bb2b8);

    border: 0;

    color: white;

    width: 100px;

    height: 40px;

    display: block;

    margin-bottom: 10px;

    border-radius: 20px;

    font-weight: 200;

}



.button:last-child {

    background-image: linear-gradient(135deg, #f97794, #623aa2);

}

定义切换函数逻辑

由于 ChatUI 的文档和 API 不全，我们简单点，直接用 localStorage + reload 来实现：

function claude(llm) {

localStorage.setItem('llm', llm)

window.location.reload()

}

修改 setup.js，定义相关变量给后面用：

const llm = localStorage.getItem('llm') || 'OpenAI'

const url = llm === 'OpenAI' ? 'http://localhost:8000/chat/openai/' : 'http://localhost:8000/chat/claude/'

const welcome1 = '小翼 已链接 ' + llm

const welcome2 = 我是小翼的 ${llm} 分身，请问有什么可以帮您？

然后刷新页面看看效果，点击 Claude 成功切换，：

本文章转载微信公众号@数翼