Cleaner implementation of streaming.

Author: johann-petrak

llms_wrapper/llms.py

@@ -697,9 +697,26 @@ def query(
                 messages=messages,
                 **completion_kwargs)
             if stream:
-                # TODO: for now we take a shortcut here and simply return the original response
-                #   as "response".
-                ret["response"] = response
+                def chunk_generator(model_generator, retobj):
+                    try:
+                        for chunk in model_generator:
+                            choice0 = chunk["choices"][0]
+                            if choice0.finish_reason == "stop":
+                                break
+                            content = choice0["delta"].get("content", "")
+                            yield dict(error="", answer=content, ok=True)
+                    except Exception as e:
+                        yield dict(error=str(e), answer="", ok=False)
+                    finally:
+                        # TODO: add cost and elapsed time information into retobj
+                        # litellm does not support cost on streaming responses
+                        # response.__hidden_params["response_cost"] is 0.0
+                        ret["cost"] = None
+                        ret["elapsed_time"] = time.time() - start
+                        pass
+                if return_response:
+                    ret["response"] = response
+                ret["chunks"] = chunk_generator(response, ret)
                 ret["ok"] = True
                 ret["error"] = ""
                 return ret
@@ -714,6 +731,9 @@ def query(
                     del completion_kwargs["api_key"]
                 ret["kwargs"] = completion_kwargs
             if return_cost:
+                # TODO: replace with response._hidden_params["response_cost"] ? 
+                #     but what if cost not supported for the model?
+                
                 try:
                     ret["cost"] = completion_cost(
                         completion_response=response,

llms_wrapper/version.py

@@ -1,3 +1,3 @@
 import importlib.metadata
-__version__ = "0.3.0"
+__version__ = "0.3.1"
 

notebooks/test-streaming.ipynb

@@ -110,7 +110,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 14,
+   "execution_count": 3,
    "id": "a89258b6",
    "metadata": {},
    "outputs": [
@@ -121,7 +121,7 @@
        "  'role': 'user'}]"
       ]
      },
-     "execution_count": 14,
+     "execution_count": 3,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -133,45 +133,55 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 15,
+   "execution_count": 4,
    "id": "491b0ddb",
    "metadata": {},
+   "outputs": [],
+   "source": [
+    "ret = llms.query(\"openai/gpt-4o\", msgs, temperature=0.5, max_tokens=1000, stream=True)\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "1c795bd0",
+   "metadata": {},
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "A monoid is an algebraic structure with a single associative binary operation and an identity element. Specifically, a set M is a monoid if it is equipped with a binary operation (let's call it *) that satisfies the following properties:\n",
+      "A monoid is an algebraic structure with a single associative binary operation and an identity element. Formally, a set M equipped with an operation * (often called multiplication) is a monoid if it satisfies two properties:\n",
       "\n",
       "1. **Associativity**: For all elements a, b, and c in M, the equation (a * b) * c = a * (b * c) holds.\n",
-      "2. **Identity Element**: There exists an element e in M such that for every element a in M, the equation e * a = a * e = a holds.\n",
+      "2. **Identity element**: There exists an element e in M such that for every element a in M, the equations e * a = a * e = a hold.\n",
+      "\n",
+      "A simple example of a monoid is the set of natural numbers (including zero) with the operation of addition. In this case:\n",
       "\n",
-      "A simple example of a monoid is the set of natural numbers (including zero) with the operation of addition. \n",
+      "- The binary operation is addition (+), which is associative. For any natural numbers a, b, and c, we have (a + b) + c = a + (b + c).\n",
+      "- The identity element is 0, because for any natural number a, a + 0 = 0 + a = a.\n",
       "\n",
-      "- The set is {0, 1, 2, 3, ...}.\n",
-      "- The binary operation is addition (+).\n",
-      "- The identity element is 0 because adding 0 to any natural number does not change the number (0 + a = a + 0 = a).\n",
-      "- Addition is associative because for any natural numbers a, b, and c, the equation (a + b) + c = a + (b + c) is always true."
+      "Thus, the natural numbers with addition form a monoid.\n"
      ]
     }
    ],
    "source": [
-    "ret = llms.query(\"openai/gpt-4o\", msgs, temperature=0.5, max_tokens=1000, stream=True)\n",
     "if ret[\"ok\"]:\n",
-    "    for chunk in ret[\"response\"]:\n",
-    "        choice0 = chunk.choices[0]\n",
-    "        if choice0.finish_reason == \"stop\":\n",
-    "            break       \n",
-    "        content = choice0.delta.content    \n",
-    "        print(content, end=\"\", flush=True)\n",
+    "    for chunk in ret[\"chunks\"]:\n",
+    "        if chunk[\"error\"]:\n",
+    "            print()\n",
+    "            print(\"Error:\", chunk[\"error\"])\n",
+    "            break\n",
+    "        print(chunk[\"answer\"], end=\"\")\n",
+    "    print()\n",
     "else:\n",
     "    print(\"Error:\", ret[\"error\"])"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "0eeea159",
+   "id": "f7fa724b",
    "metadata": {},
    "outputs": [],
    "source": []