fix: Update Fireworks provider and notebook for proper event tracking

Co-Authored-By: Alex Reibman <[email protected]>
AgentOps-AI · Dec 18, 2024 · 49d3f27 · 49d3f27
1 parent a28c0f1
commit 49d3f27
Show file tree

Hide file tree

Showing 2 changed files with 37 additions and 34 deletions.
diff --git a/examples/fireworks_examples/fireworks_example.ipynb b/examples/fireworks_examples/fireworks_example.ipynb
@@ -2,7 +2,7 @@
  "cells": [
   {
    "cell_type": "markdown",
-   "id": "8ce4ef43",
+   "id": "a5db7f5c",
    "metadata": {},
    "source": [
     "# Fireworks LLM Integration Example\n",
@@ -13,7 +13,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "caeebf2f",
+   "id": "cf6cede9",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -38,7 +38,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "4fa145da",
+   "id": "2d77b583",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -78,7 +78,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "62bdbde6",
+   "id": "9a434a04",
    "metadata": {},
    "source": [
     "## Test Cases\n",
@@ -94,7 +94,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "7e1dd8b5",
+   "id": "59e6d9ba",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -118,11 +118,10 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "40a2250f",
+   "id": "cc81b366",
    "metadata": {},
    "outputs": [],
    "source": [
-    "%%async\n",
     "# 2. Test asynchronous non-streaming completion\n",
     "print(\"2. Generating story with asynchronous non-streaming completion...\")\n",
     "\n",
@@ -135,13 +134,22 @@
     "    print(response.choices[0].message.content)\n",
     "    print(\"\\nEvent recorded for async non-streaming completion\")\n",
     "\n",
-    "await async_completion()"
+    "# Run async function using asyncio.run()\n",
+    "try:\n",
+    "    asyncio.run(async_completion())\n",
+    "except RuntimeError as e:\n",
+    "    if \"Cannot run the event loop while another loop is running\" in str(e):\n",
+    "        # If we're already in an event loop, create a new one\n",
+    "        loop = asyncio.get_event_loop()\n",
+    "        loop.run_until_complete(async_completion())\n",
+    "    else:\n",
+    "        raise"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "d6e61be0",
+   "id": "a70cb4ea",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -170,11 +178,10 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "02a76fb9",
+   "id": "d1fd034c",
    "metadata": {},
    "outputs": [],
    "source": [
-    "%%async\n",
     "# 4. Test asynchronous streaming completion\n",
     "print(\"4. Generating story with asynchronous streaming...\")\n",
     "\n",
@@ -195,13 +202,22 @@
     "        logger.error(f\"Error in async streaming: {str(e)}\")\n",
     "    print()  # New line after streaming\n",
     "\n",
-    "await async_streaming()"
+    "# Run async function using asyncio.run()\n",
+    "try:\n",
+    "    asyncio.run(async_streaming())\n",
+    "except RuntimeError as e:\n",
+    "    if \"Cannot run the event loop while another loop is running\" in str(e):\n",
+    "        # If we're already in an event loop, create a new one\n",
+    "        loop = asyncio.get_event_loop()\n",
+    "        loop.run_until_complete(async_streaming())\n",
+    "    else:\n",
+    "        raise"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "bc64a254",
+   "id": "d413c0ba",
    "metadata": {},
    "outputs": [],
    "source": [

diff --git a/examples/fireworks_examples/update_notebook.py b/examples/fireworks_examples/update_notebook.py
@@ -125,16 +125,9 @@ async def async_completion():
     print(response.choices[0].message.content)
     print("\\nEvent recorded for async non-streaming completion")
 
-# Run async function using asyncio.run()
-try:
-    asyncio.run(async_completion())
-except RuntimeError as e:
-    if "Cannot run the event loop while another loop is running" in str(e):
-        # If we're already in an event loop, create a new one
-        loop = asyncio.get_event_loop()
-        loop.run_until_complete(async_completion())
-    else:
-        raise"""
+# Get the current event loop and run the async function
+loop = asyncio.get_event_loop()
+loop.run_until_complete(async_completion())"""
     )
 )
 
@@ -183,20 +176,14 @@ async def async_streaming():
                 content = chunk.choices[0].delta.content
                 if content:
                     print(content, end="", flush=True)
+        print("\\nEvent recorded for async streaming completion")
     except Exception as e:
         logger.error(f"Error in async streaming: {str(e)}")
-    print()  # New line after streaming
+        raise
 
-# Run async function using asyncio.run()
-try:
-    asyncio.run(async_streaming())
-except RuntimeError as e:
-    if "Cannot run the event loop while another loop is running" in str(e):
-        # If we're already in an event loop, create a new one
-        loop = asyncio.get_event_loop()
-        loop.run_until_complete(async_streaming())
-    else:
-        raise"""
+# Get the current event loop and run the async function
+loop = asyncio.get_event_loop()
+loop.run_until_complete(async_streaming())"""
     )
 )