{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.llms import OpenAI, Anthropic\n",
    "from langchain.chat_models import ChatOpenAI\n",
    "from langchain.callbacks.base import CallbackManager\n",
    "from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler\n",
    "from langchain.schema import HumanMessage\n",
    "from langchain.callbacks.base import BaseCallbackHandler\n",
    "import sys\n",
    "from typing import Any, Dict, List, Union\n",
    "from langchain.schema import AgentAction, AgentFinish, LLMResult"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "class PrintCallbackHandler(BaseCallbackHandler):\n",
    "    \"\"\"Callback handler for streaming. Only works with LLMs that support streaming.\"\"\"\n",
    "\n",
    "    def on_llm_start(\n",
    "        self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any\n",
    "    ) -> None:\n",
    "        \"\"\"Run when LLM starts running.\"\"\"\n",
    "        print(\"on_llm_start\")\n",
    "\n",
    "    def on_llm_new_token(self, token: str, **kwargs: Any) -> None:\n",
    "        \"\"\"Run on new LLM token. Only available when streaming is enabled.\"\"\"\n",
    "        sys.stdout.write(token)\n",
    "        sys.stdout.flush()\n",
    "\n",
    "    def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:\n",
    "        \"\"\"Run when LLM ends running.\"\"\"\n",
    "        print(\"\\non_llm_end\")\n",
    "\n",
    "    def on_llm_error(\n",
    "        self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any\n",
    "    ) -> None:\n",
    "        \"\"\"Run when LLM errors.\"\"\"\n",
    "\n",
    "    def on_chain_start(\n",
    "        self, serialized: Dict[str, Any], inputs: Dict[str, Any], **kwargs: Any\n",
    "    ) -> None:\n",
    "        \"\"\"Run when chain starts running.\"\"\"\n",
    "        print(\"on_chain_start\")\n",
    "\n",
    "    def on_chain_end(self, outputs: Dict[str, Any], **kwargs: Any) -> None:\n",
    "        \"\"\"Run when chain ends running.\"\"\"\n",
    "        print(\"on_chain_end\")\n",
    "\n",
    "    def on_chain_error(\n",
    "        self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any\n",
    "    ) -> None:\n",
    "        \"\"\"Run when chain errors.\"\"\"\n",
    "\n",
    "    def on_tool_start(\n",
    "        self, serialized: Dict[str, Any], input_str: str, **kwargs: Any\n",
    "    ) -> None:\n",
    "        \"\"\"Run when tool starts running.\"\"\"\n",
    "        print(\"on_tool_start\")\n",
    "\n",
    "    def on_agent_action(self, action: AgentAction, **kwargs: Any) -> Any:\n",
    "        \"\"\"Run on agent action.\"\"\"\n",
    "        pass\n",
    "\n",
    "    def on_tool_end(self, output: str, **kwargs: Any) -> None:\n",
    "        \"\"\"Run when tool ends running.\"\"\"\n",
    "\n",
    "    def on_tool_error(\n",
    "        self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any\n",
    "    ) -> None:\n",
    "        \"\"\"Run when tool errors.\"\"\"\n",
    "\n",
    "    def on_text(self, text: str, **kwargs: Any) -> None:\n",
    "        \"\"\"Run on arbitrary text.\"\"\"\n",
    "        print(\"on_text\")\n",
    "\n",
    "    def on_agent_finish(self, finish: AgentFinish, **kwargs: Any) -> None:\n",
    "        \"\"\"Run on agent end.\"\"\"\n",
    "        print(\"on_agent_finish\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "on_llm_start\n",
      "\n",
      "\n",
      "每当我想起你 \n",
      "心里总是满满的温柔 \n",
      "每当我想起你 \n",
      "总是有一种特别的感觉 \n",
      "\n",
      "每当我想起你 \n",
      "总是有一种深深的思念 \n",
      "每当我想起你 \n",
      "总是有一种温暖的感觉 \n",
      "\n",
      "每当我想起你 \n",
      "总是有一种温柔的爱 \n",
      "每当我想起你 \n",
      "总是有一种深深的情感 \n",
      "\n",
      "每当我想起你 \n",
      "总是有一种温暖的温柔 \n",
      "每当我想起你 \n",
      "总是有一种特别的感动\n",
      "on_llm_end\n"
     ]
    }
   ],
   "source": [
    "# or use StreamingStdOutCallbackHandler\n",
    "llm = OpenAI(max_tokens=1024, streaming=True, callback_manager=CallbackManager([PrintCallbackHandler()]), verbose=True, temperature=0)\n",
    "resp = llm(\"仿写一份周杰伦的歌词\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "on_llm_start\n",
      "\n",
      "\n",
      "Q: What did the fish say when it hit the wall?\n",
      "A: Dam!on_llm_end\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "LLMResult(generations=[[Generation(text='\\n\\nQ: What did the fish say when it hit the wall?\\nA: Dam!', generation_info={'finish_reason': 'stop', 'logprobs': None})]], llm_output={'token_usage': {}, 'model_name': 'text-davinci-003'})"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "llm.generate([\"Tell me a joke.\"])"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.11.3"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}