add streaming

llms_streaming.ipynb

@@ -0,0 +1,186 @@
+{
+ "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
+}

openai_simple.py

@@ -1,12 +1,15 @@
 #!/usr/bin/python3
 
 from langchain.llms import OpenAI
+from langchain.callbacks.base import CallbackManager
+from langchain.callbacks.openai_info import OpenAICallbackHandler
+from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
 import langchain
 import os
 
 print(langchain.__version__)
 
-llm = OpenAI(temperature=0.9)
+llm = OpenAI(temperature=0.9, callback_manager=CallbackManager([OpenAICallbackHandler()]))
 
 text = "Hi, Could you help me choose a Chinese name? For example"
 print(llm(text))