add async and streaming support to OpenAIChat (#1378)

title says it all
langchain-ai · Mar 2, 2023 · fe30be6 · fe30be6
1 parent cfed049
commit fe30be6
Show file tree

Hide file tree

Showing 5 changed files with 319 additions and 138 deletions.
diff --git a/docs/modules/llms/async_llm.ipynb b/docs/modules/llms/async_llm.ipynb
@@ -1,7 +1,6 @@
 {
  "cells": [
   {
-   "attachments": {},
    "cell_type": "markdown",
    "id": "f6574496-b360-4ffa-9523-7fd34a590164",
    "metadata": {},
@@ -10,14 +9,14 @@
     "\n",
     "LangChain provides async support for LLMs by leveraging the [asyncio](https://docs.python.org/3/library/asyncio.html) library.\n",
     "\n",
-    "Async support is particularly useful for calling multiple LLMs concurrently, as these calls are network-bound. Currently, only `OpenAI` and `PromptLayerOpenAI` is supported, but async support for other LLMs is on the roadmap.\n",
+    "Async support is particularly useful for calling multiple LLMs concurrently, as these calls are network-bound. Currently, only `OpenAI` `OpenAIChat`, and `PromptLayerOpenAI` are supported, but async support for other LLMs is on the roadmap.\n",
     "\n",
     "You can use the `agenerate` method to call an OpenAI LLM asynchronously."
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": 1,
    "id": "5e49e96c-0f88-466d-b3d3-ea0966bdf19e",
    "metadata": {
     "tags": []
@@ -29,75 +28,77 @@
      "text": [
       "\n",
       "\n",
-      "I'm doing well. How about you?\n",
+      "As an AI language model, I don't have feelings like humans, but I'm functioning properly. How may I assist you?\n",
       "\n",
       "\n",
-      "I'm doing well, thank you. How about you?\n",
+      "I'm an AI language model, so I don't have emotions, but I'm functioning properly. How may I assist you today?\n",
       "\n",
       "\n",
-      "I'm doing well, thank you. How about you?\n",
+      "As an AI language model, I do not have emotions like humans, but I'm functioning normally. How can I assist you today?\n",
       "\n",
       "\n",
-      "I'm doing well, thank you. How about you?\n",
+      "I am an AI language model, so I do not have feelings, but I am here to assist you. How may I help you today?\n",
       "\n",
-      "I am doing quite well. How about you?\n",
       "\n",
+      "As an AI language model, I do not have feelings or emotions but I'm always ready to assist you. How may I assist you today?\n",
       "\n",
-      "I'm doing well, thank you. How about you?\n",
       "\n",
+      "As an AI language model, I don't have feelings, but I'm functioning normally. How may I assist you today?\n",
       "\n",
-      "I'm doing great, thank you! How about you?\n",
       "\n",
+      "As an AI language model, I don't have feelings, but I'm functioning properly. Thank you. How may I assist you today?\n",
       "\n",
-      "I'm doing well, thanks for asking. How about you?\n",
       "\n",
+      "As an AI language model, I don't have emotions, so I don't have a specific feeling or emotion. How can I assist you today?\n",
       "\n",
-      "I'm doing well, thank you. How about you?\n",
       "\n",
+      "As an AI language model, I do not have feelings or emotions. However, I am functioning as intended and ready to assist you with any queries you may have. How can I be of assistance today?\n",
       "\n",
-      "I'm doing well, thank you. How about you?\n",
-      "\u001b[1mConcurrent executed in 1.93 seconds.\u001b[0m\n",
       "\n",
+      "As an AI language model, I do not have feelings, but I am functioning well. Thank you for asking. How can I assist you today?\n",
+      "\u001b[1mConcurrent executed in 0.92 seconds.\u001b[0m\n",
       "\n",
-      "I'm doing well, thank you. How about you?\n",
       "\n",
+      "As an AI language model, I don't have feelings, but I'm functioning well. How can I assist you today?\n",
       "\n",
-      "I'm doing well, thank you. How about you?\n",
       "\n",
+      "As an AI language model, I don't have feelings, but I'm functioning well. Thank you for asking. How may I assist you today?\n",
       "\n",
-      "I'm doing well, thank you. How about you?\n",
       "\n",
+      "I'm an AI language model, so I don't have feelings, but I'm functioning well. How can I assist you today?\n",
       "\n",
-      "I'm doing well, thank you. How about you?\n",
       "\n",
+      "As an AI language model, I don't have feelings, but I'm functioning well. Thank you for asking. How may I assist you today?\n",
       "\n",
-      "I'm doing well, thank you. How about you?\n",
       "\n",
+      "As an AI language model, I don't have feelings, but I am functioning well. How can I assist you today?\n",
       "\n",
-      "I'm doing well, thank you. How about you?\n",
       "\n",
-      "I'm doing well, thank you. How about you?\n",
+      "As an AI language model, I don't have feelings but I'm functioning well. How can I assist you today?\n",
       "\n",
       "\n",
-      "I'm doing well, thank you. How about you?\n",
+      "As an AI language model, I do not have personal emotions. However, I am functioning well and ready to assist you with any queries or tasks you have. How may I assist you today?\n",
       "\n",
       "\n",
-      "I'm doing well, thank you. How about you?\n",
+      "As an AI language model, I do not have feelings or emotions, but I'm functioning well. How can I assist you today?\n",
       "\n",
       "\n",
-      "I'm doing great, thank you. How about you?\n",
-      "\u001b[1mSerial executed in 10.54 seconds.\u001b[0m\n"
+      "I am an AI language model and do not have feelings. But I am functioning properly and ready to assist you with any task. How may I help you today?\n",
+      "\n",
+      "\n",
+      "As an AI language model, I do not have emotions, but I am functioning well. How can I assist you today?\n",
+      "\u001b[1mSerial executed in 5.00 seconds.\u001b[0m\n"
      ]
     }
    ],
    "source": [
     "import time\n",
     "import asyncio\n",
     "\n",
-    "from langchain.llms import OpenAI\n",
+    "from langchain.llms import OpenAIChat\n",
     "\n",
     "def generate_serially():\n",
-    "    llm = OpenAI(temperature=0.9)\n",
+    "    llm = OpenAIChat(temperature=0.9)\n",
     "    for _ in range(10):\n",
     "        resp = llm.generate([\"Hello, how are you?\"])\n",
     "        print(resp.generations[0][0].text)\n",
@@ -109,7 +110,7 @@
     "\n",
     "\n",
     "async def generate_concurrently():\n",
-    "    llm = OpenAI(temperature=0.9)\n",
+    "    llm = OpenAIChat(temperature=0.9)\n",
     "    tasks = [async_generate(llm) for _ in range(10)]\n",
     "    await asyncio.gather(*tasks)\n",
     "\n",
@@ -125,6 +126,14 @@
     "elapsed = time.perf_counter() - s\n",
     "print('\\033[1m' + f\"Serial executed in {elapsed:0.2f} seconds.\" + '\\033[0m')"
    ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "e1d3a966-3a27-44e8-9441-ed72f01b86f4",
+   "metadata": {},
+   "outputs": [],
+   "source": []
   }
  ],
  "metadata": {

diff --git a/docs/modules/llms/streaming_llm.ipynb b/docs/modules/llms/streaming_llm.ipynb
@@ -7,12 +7,12 @@
    "source": [
     "# Streaming with LLMs\n",
     "\n",
-    "LangChain provides streaming support for LLMs. Currently, we only support streaming for the `OpenAI` LLM implementation, but streaming support for other LLM implementations is on the roadmap. To utilize streaming, use a [`CallbackHandler`](https://github.com/hwchase17/langchain/blob/master/langchain/callbacks/base.py) that implements `on_llm_new_token`. In this example, we are using [`StreamingStdOutCallbackHandler`]()."
+    "LangChain provides streaming support for LLMs. Currently, we only support streaming for the `OpenAI` and `OpenAIChat` LLM implementation, but streaming support for other LLM implementations is on the roadmap. To utilize streaming, use a [`CallbackHandler`](https://github.com/hwchase17/langchain/blob/master/langchain/callbacks/base.py) that implements `on_llm_new_token`. In this example, we are using [`StreamingStdOutCallbackHandler`]()."
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": 2,
    "id": "4ac0ff54-540a-4f2b-8d9a-b590fec7fe07",
    "metadata": {
     "tags": []
@@ -27,43 +27,43 @@
       "Verse 1\n",
       "I'm sippin' on sparkling water,\n",
       "It's so refreshing and light,\n",
-      "It's the perfect way to quench my thirst,\n",
+      "It's the perfect way to quench my thirst\n",
       "On a hot summer night.\n",
       "\n",
       "Chorus\n",
       "Sparkling water, sparkling water,\n",
       "It's the best way to stay hydrated,\n",
-      "It's so refreshing and light,\n",
-      "It's the perfect way to stay alive.\n",
+      "It's so crisp and so clean,\n",
+      "It's the perfect way to stay refreshed.\n",
       "\n",
       "Verse 2\n",
       "I'm sippin' on sparkling water,\n",
       "It's so bubbly and bright,\n",
-      "It's the perfect way to cool me down,\n",
+      "It's the perfect way to cool me down\n",
       "On a hot summer night.\n",
       "\n",
       "Chorus\n",
       "Sparkling water, sparkling water,\n",
       "It's the best way to stay hydrated,\n",
-      "It's so refreshing and light,\n",
-      "It's the perfect way to stay alive.\n",
+      "It's so crisp and so clean,\n",
+      "It's the perfect way to stay refreshed.\n",
       "\n",
       "Verse 3\n",
       "I'm sippin' on sparkling water,\n",
-      "It's so crisp and clean,\n",
-      "It's the perfect way to keep me going,\n",
-      "On a hot summer day.\n",
+      "It's so light and so clear,\n",
+      "It's the perfect way to keep me cool\n",
+      "On a hot summer night.\n",
       "\n",
       "Chorus\n",
       "Sparkling water, sparkling water,\n",
       "It's the best way to stay hydrated,\n",
-      "It's so refreshing and light,\n",
-      "It's the perfect way to stay alive."
+      "It's so crisp and so clean,\n",
+      "It's the perfect way to stay refreshed."
      ]
     }
    ],
    "source": [
-    "from langchain.llms import OpenAI\n",
+    "from langchain.llms import OpenAI, OpenAIChat\n",
     "from langchain.callbacks.base import CallbackManager\n",
     "from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler\n",
     "\n",
@@ -84,7 +84,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": 3,
    "id": "a35373f1-9ee6-4753-a343-5aee749b8527",
    "metadata": {
     "tags": []
@@ -103,17 +103,96 @@
     {
      "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': {}})"
+       "LLMResult(generations=[[Generation(text='\\n\\nQ: What did the fish say when it hit the wall?\\nA: Dam!', generation_info={'finish_reason': None, 'logprobs': None})]], llm_output={'token_usage': {}})"
       ]
      },
-     "execution_count": 8,
+     "execution_count": 3,
      "metadata": {},
      "output_type": "execute_result"
     }
    ],
    "source": [
     "llm.generate([\"Tell me a joke.\"])"
    ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "a93a4d61-0476-49db-8321-7de92bd74059",
+   "metadata": {},
+   "source": [
+    "Here's an example with `OpenAIChat`:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "22665f16-e05b-473c-a4bd-ad75744ea024",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\n",
+      "Verse 1:\n",
+      "Bubbles rising to the top\n",
+      "A refreshing drink that never stops\n",
+      "Clear and crisp, it's pure delight\n",
+      "A taste that's sure to excite\n",
+      "\n",
+      "Chorus:\n",
+      "Sparkling water, oh so fine\n",
+      "A drink that's always on my mind\n",
+      "With every sip, I feel alive\n",
+      "Sparkling water, you're my vibe\n",
+      "\n",
+      "Verse 2:\n",
+      "No sugar, no calories, just pure bliss\n",
+      "A drink that's hard to resist\n",
+      "It's the perfect way to quench my thirst\n",
+      "A drink that always comes first\n",
+      "\n",
+      "Chorus:\n",
+      "Sparkling water, oh so fine\n",
+      "A drink that's always on my mind\n",
+      "With every sip, I feel alive\n",
+      "Sparkling water, you're my vibe\n",
+      "\n",
+      "Bridge:\n",
+      "From the mountains to the sea\n",
+      "Sparkling water, you're the key\n",
+      "To a healthy life, a happy soul\n",
+      "A drink that makes me feel whole\n",
+      "\n",
+      "Chorus:\n",
+      "Sparkling water, oh so fine\n",
+      "A drink that's always on my mind\n",
+      "With every sip, I feel alive\n",
+      "Sparkling water, you're my vibe\n",
+      "\n",
+      "Outro:\n",
+      "Sparkling water, you're the one\n",
+      "A drink that's always so much fun\n",
+      "I'll never let you go, my friend\n",
+      "Sparkling water, until the end."
+     ]
+    }
+   ],
+   "source": [
+    "llm = OpenAIChat(streaming=True, callback_manager=CallbackManager([StreamingStdOutCallbackHandler()]), verbose=True, temperature=0)\n",
+    "resp = llm(\"Write me a song about sparkling water.\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "eadae4ba-9f21-4ec8-845d-dd43b0edc2dc",
+   "metadata": {},
+   "outputs": [],
+   "source": []
   }
  ],
  "metadata": {

diff --git a/langchain/llms/base.py b/langchain/llms/base.py
@@ -319,6 +319,10 @@ class LLM(BaseLLM):
     def _call(self, prompt: str, stop: Optional[List[str]] = None) -> str:
         """Run the LLM on the given prompt and input."""
 
+    async def _acall(self, prompt: str, stop: Optional[List[str]] = None) -> str:
+        """Run the LLM on the given prompt and input."""
+        raise NotImplementedError("Async generation not implemented for this LLM.")
+
     def _generate(
         self, prompts: List[str], stop: Optional[List[str]] = None
     ) -> LLMResult:
@@ -334,4 +338,8 @@ async def _agenerate(
         self, prompts: List[str], stop: Optional[List[str]] = None
     ) -> LLMResult:
         """Run the LLM on the given prompt and input."""
-        raise NotImplementedError("Async generation not implemented for this LLM.")
+        generations = []
+        for prompt in prompts:
+            text = await self._acall(prompt, stop=stop)
+            generations.append([Generation(text=text)])
+        return LLMResult(generations=generations)