langchain-doubao.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "44a22ffe-2303-44e0-9e52-ef7a630cf1d0",
   "metadata": {},
   "source": [
    "#通过langchain的自定义LLM封装Doubao模型\n",
    "\n",
    "自定义LLM，参考 https://blog.csdn.net/qq_39749966/article/details/136108973\n",
    "\n",
    "豆包模型调用参考：\n",
    "```\n",
    "'''\n",
    "Usage:\n",
    "Ark v3 sdk\n",
    "pip install 'volcengine-python-sdk[ark]'\n",
    "'''\n",
    "\n",
    "from volcenginesdkarkruntime import Ark\n",
    "\n",
    "# fetch ak&sk from environmental variables \"VOLC_ACCESSKEY\", \"VOLC_SECRETKEY\"\n",
    "# or specify ak&sk by Ark(ak=\"${YOUR_AK}\", sk=\"${YOUR_SK}\").\n",
    "# you can get ak&sk follow this document(https://www.volcengine.com/docs/6291/65568)\n",
    "client = Ark()\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    # Non-streaming:\n",
    "    print(\"----- standard request -----\")\n",
    "    completion = client.chat.completions.create(\n",
    "        model=\"ep-20240531085255-6tmx9\",\n",
    "        messages=[\n",
    "            {\n",
    "                \"role\": \"user\",\n",
    "                \"content\": \"Say this is a test\",\n",
    "            },\n",
    "        ]\n",
    "    )\n",
    "    print(completion.choices[0].message.content)\n",
    "\n",
    "    # Streaming:\n",
    "    print(\"----- streaming request -----\")\n",
    "    stream = client.chat.completions.create(\n",
    "        model=\"ep-20240531085255-6tmx9\",\n",
    "        messages=[\n",
    "            {\n",
    "                \"role\": \"user\",\n",
    "                \"content\": \"How do I output all files in a directory using Python?\",\n",
    "            },\n",
    "        ],\n",
    "        stream=True\n",
    "    )\n",
    "    for chunk in stream:\n",
    "        if not chunk.choices:\n",
    "            continue\n",
    "\n",
    "        print(chunk.choices[0].delta.content, end=\"\")\n",
    "    print()\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "47097375-3cde-4624-9e07-4e3e8f01e3b8",
   "metadata": {},
   "source": [
    "需要安装的库\n",
    "```\n",
    "pip install langchain\n",
    "pip install volcengine-python-sdk\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "fb6ea4f6-06c4-4faa-91f1-cea0340b9325",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'question': '姚明身高是多少',\n",
       " 'text': '姚明的身高是 2.26 米。\\n\\n姚明是中国著名篮球运动员，他在篮球领域取得了卓越的成就。'}"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from volcenginesdkarkruntime import Ark\n",
    "from langchain.llms.base import LLM\n",
    "from langchain import PromptTemplate, LLMChain\n",
    "import logging\n",
    "from typing import Any\n",
    "\n",
    "class DoubaoLLM(LLM):\n",
    "\n",
    "    model: str = \"ep-20240531085255-6tmx9\"\n",
    "    client: Ark = None\n",
    "\n",
    "    def __init__(self, akstr: str, skstr: str):\n",
    "        super().__init__()\n",
    "        self.client = Ark(ak=akstr, sk=skstr)\n",
    "        \n",
    "    @property\n",
    "    def _llm_type(self) -> str:\n",
    "        return \"Doubao\"\n",
    "\n",
    "    def _call(self, prompt: str, **kwargs: Any) -> str:\n",
    "        try:\n",
    "            logging.info(\"Making API call to Doubao endpoint.\")\n",
    "            return self._make_api_call(prompt)\n",
    "        except Exception as e:\n",
    "            logging.error(f\"Error in DoubaoLLM _call: {e}\", exc_info=True)\n",
    "\n",
    "    def _make_api_call(self, prompt: str) -> str:\n",
    "        completion = self.client.chat.completions.create(\n",
    "            model = self.model,\n",
    "            messages = [\n",
    "                {\n",
    "                    \"role\": \"user\",\n",
    "                    \"content\": prompt,\n",
    "                },\n",
    "            ]\n",
    "        )\n",
    "        return completion.choices[0].message.content\n",
    "\n",
    "\n",
    "llmdoubao = DoubaoLLM(\n",
    "    akstr=\"AKLTM2ZkYzczMThlMWM5NGYxYTg5NDE4N2ZlYmRhOTNlZDA\", \n",
    "    skstr=\"TXpZME9EWTBOVGcwTTJVNU5EQmxPVGc0WkRkak16Rm1ORFJoTjJNd1pEQQ==\"\n",
    ")\n",
    "\n",
    "prompt_template = \"你是一个智能助手，回答以下的问题： {question}\"\n",
    "promptdoubao = PromptTemplate(\n",
    "    input_variables=[\"question\"], template=prompt_template\n",
    ")\n",
    "chat = LLMChain(llm=llmdoubao, prompt=promptdoubao)\n",
    "\n",
    "chat(\"姚明身高是多少\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}