{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "initial_id",
   "metadata": {},
   "outputs": [],
   "source": [
    "# !pip install arize-phoenix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dcbd25c8",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install portpicker"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4331618113a05e5c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from datetime import datetime, timezone"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d4520d5a133d7488",
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import load_dataset\n",
    "\n",
    "import phoenix as px\n",
    "\n",
    "sample_size = 5\n",
    "path = \"nvidia/ChatQA-Training-Data\"\n",
    "name = \"synthetic_convqa\"\n",
    "df = load_dataset(path, name, split=\"train\").to_pandas().sample(sample_size, random_state=42)\n",
    "df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e955fb85754f5c87",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_name = \"nvidia/ChatQA-Training-Data\" + \"-\" + datetime.now(timezone.utc).isoformat()\n",
    "px.Client().upload_dataset(\n",
    "    dataset_name=dataset_name,\n",
    "    dataframe=df,\n",
    "    input_keys=(\"messages\", \"document\"),\n",
    "    output_keys=(\"answers\",),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "30a54e1924e8e192",
   "metadata": {},
   "outputs": [],
   "source": [
    "ds = px.Client().get_dataset(name=dataset_name)\n",
    "type(ds)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e1449f5027787d6b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from contextlib import contextmanager\n",
    "from threading import Thread\n",
    "from time import sleep, time\n",
    "from typing import Awaitable, Callable, Generator\n",
    "\n",
    "from portpicker import pick_unused_port\n",
    "from starlette.applications import Starlette\n",
    "from starlette.responses import JSONResponse, Response\n",
    "from starlette.routing import Request, Route\n",
    "from uvicorn import Config, Server\n",
    "\n",
    "\n",
    "async def hello(_: Request) -> Response:\n",
    "    return JSONResponse(\n",
    "        content={\n",
    "            \"id\": \"chatcmpl-123\",\n",
    "            \"object\": \"chat.completion\",\n",
    "            \"created\": 1677652288,\n",
    "            \"model\": \"gpt-3.5-turbo-0125\",\n",
    "            \"system_fingerprint\": \"fp_44709d6fcb\",\n",
    "            \"choices\": [\n",
    "                {\n",
    "                    \"index\": 0,\n",
    "                    \"message\": {\n",
    "                        \"role\": \"assistant\",\n",
    "                        \"content\": \"\\n\\nHello there, how may I assist you today?\",\n",
    "                    },\n",
    "                    \"finish_reason\": \"stop\",\n",
    "                }\n",
    "            ],\n",
    "            \"usage\": {\"prompt_tokens\": 9, \"completion_tokens\": 12, \"total_tokens\": 21},\n",
    "        }\n",
    "    )\n",
    "\n",
    "\n",
    "class Receiver:\n",
    "    def __init__(self, chat_completion: Callable[[Request], Awaitable[Response]]) -> None:\n",
    "        self.app = Starlette(\n",
    "            routes=[\n",
    "                Route(\"/v1/chat/completions\", chat_completion, methods=[\"POST\"]),\n",
    "            ]\n",
    "        )\n",
    "\n",
    "    def install_signal_handlers(self) -> None:\n",
    "        pass\n",
    "\n",
    "    @contextmanager\n",
    "    def run_in_thread(self, port: int) -> Generator[Thread, None, None]:\n",
    "        \"\"\"A coroutine to keep the server running in a thread.\"\"\"\n",
    "        config = Config(app=self.app, port=port, loop=\"asyncio\", log_level=\"critical\")\n",
    "        server = Server(config=config)\n",
    "        thread = Thread(target=server.run)\n",
    "        thread.start()\n",
    "        time_limit = time() + 5  # 5 seconds\n",
    "        try:\n",
    "            while not server.started and thread.is_alive() and time() < time_limit:\n",
    "                sleep(1e-3)\n",
    "            if time() > time_limit:\n",
    "                raise RuntimeError(\"server took too long to start\")\n",
    "            yield\n",
    "        finally:\n",
    "            server.should_exit = True\n",
    "            thread.join(timeout=5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1885dad6ab417351",
   "metadata": {},
   "outputs": [],
   "source": [
    "import openai\n",
    "\n",
    "port = pick_unused_port()\n",
    "client = openai.OpenAI(api_key=\"sk-\", base_url=f\"http://localhost:{port}/v1/\")\n",
    "\n",
    "\n",
    "def task(input):\n",
    "    response = client.chat.completions.create(\n",
    "        model=\"gpt-3.5-turbo\",\n",
    "        messages=input[\"messages\"],\n",
    "        max_tokens=20,\n",
    "    )\n",
    "    return response.choices[0].message.content"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b7792033bae20707",
   "metadata": {},
   "outputs": [],
   "source": [
    "import nest_asyncio\n",
    "\n",
    "from phoenix.experiments import run_experiment\n",
    "\n",
    "nest_asyncio.apply()\n",
    "\n",
    "with Receiver(hello).run_in_thread(port):\n",
    "    run_experiment(\n",
    "        ds,\n",
    "        task,\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26bf8cef",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}