{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# **Chat Completions**\n",
    "\n",
    "This notebook demonstrates how to implement **chat completions** using the **Inflection AI API**. \n",
    "\n",
    "With the new OpenAI API compatibility support, one can use OpenAI Client or direct http requests.\n",
    "\n",
    "By default, the Inference API service generates the entire completion before sending it back as a single response, which can result in significant delays for longer completions.\n",
    "To minimize waiting time, you can opt for 'streaming' completions, allowing you to receive and process partial completions as they're generated. This enables you to begin working with the initial part of the completion without waiting for the entire response.\n",
    "\n",
    "Enable streaming by setting *`stream=True`* when using the chat completions or completions endpoints. This returns an object that delivers data-only server-sent events in chunks, accessible through the delta field instead of the standard message field."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from dotenv import load_dotenv\n",
    "\n",
    "load_dotenv()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Retrieve environment variables\n",
    "base_url = os.getenv(\"BASE_URL\")\n",
    "inflection_api_key = os.getenv(\"INFLECTION_API_KEY\")\n",
    "\n",
    "# Set the model you want to use\n",
    "#model = \"inflection_3_productivity\"\n",
    "model = \"inflection_3_pi\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class color:\n",
    "    BOLD = '\\033[1m'\n",
    "    END = '\\033[0m'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### **Test Message** ###\n",
    "We'll use the simple test message below, but modify as needed to experimengt with the responses"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Messages\n",
    "test_messages = [\n",
    "    {\n",
    "        \"role\": \"user\", \n",
    "        \"content\": \"What's the weather in Palo Alto?\"\n",
    "    }\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## **Using OpenAI Client** ##"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from openai import OpenAI\n",
    "\n",
    "url = base_url + \"/external/api/inference/openai/v1/\"\n",
    "client = OpenAI(base_url=url, api_key=inflection_api_key)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### **Simple Message Completion Without Streaming** ###\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"Starting test: chat completion without streaming using {model}\")\n",
    "print(\"+*\"*20)\n",
    "\n",
    "response = client.chat.completions.create(\n",
    "    model=model,\n",
    "    messages=test_messages,\n",
    ")\n",
    "\n",
    "\n",
    "response_content = response.choices[0].message.content\n",
    "\n",
    "print(f\"{color.BOLD} Response: {color.END} {response_content}\")\n",
    "print(\"+*\"*20)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### **Simple Message Completion With Streaming** ###\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"Starting test: chat completion through OpenAI Client with streaming using {model}\")\n",
    "print(\"+*\"*20)\n",
    "\n",
    "response = client.chat.completions.create(\n",
    "    model=model,\n",
    "    messages=test_messages,\n",
    "    stream=True, # Enable streaming\n",
    ")\n",
    "\n",
    "for chunk in response:\n",
    "    if chunk.choices[0].delta.content:\n",
    "        print(f\"{color.BOLD} Chunk: {color.END} {chunk.choices[0].delta.content}\", end=\"\\n\")\n",
    "\n",
    "\n",
    "print(\"+*\"*20)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## **Using HTTP Requests** ##"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import requests\n",
    "\n",
    "url = base_url + \"/external/api/inference/openai/v1/chat/completions\"\n",
    "\n",
    "headers = {\n",
    "    \"Authorization\": f\"Bearer {inflection_api_key}\",\n",
    "    \"Content-Type\": \"application/json\",\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### **Simple Message Completion Without Streaming** ###"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "json_payload = {\n",
    "    \"model\": model, \n",
    "    \"messages\": test_messages,\n",
    "    \"temperature\": 0.8,\n",
    "    \"top_p\": 1,\n",
    "    \"web_search\": True,\n",
    "}\n",
    "\n",
    "# Convert the payload to a JSON string\n",
    "payload = json.dumps(json_payload)\n",
    "\n",
    "# Make the post request\n",
    "print(f\"Starting test: chat completion through an HTTP POST without streaming using {model}\")\n",
    "print(\"+*\"*20)\n",
    "\n",
    "response = requests.post(url, headers=headers, data=payload)\n",
    "\n",
    "# Check if the request was successful\n",
    "if response.status_code == 200:\n",
    "    # Parse the JSON response\n",
    "    json_response = response.json()\n",
    "    # Print the response\n",
    "    print(f\"{color.BOLD} Response: {color.END} {json_response['choices'][0]['message']['content']}\")\n",
    "\n",
    "print(\"+*\"*20)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### **Simple Message Completion With Streaming** ###"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "json_payload = {\n",
    "    \"model\": model, \n",
    "    \"messages\": test_messages,\n",
    "    \"temperature\": 0.8,\n",
    "    \"top_p\": 1,\n",
    "    \"web_search\": True,\n",
    "    \"stream\": True # Enable streaming\n",
    "}\n",
    "\n",
    "# Convert the payload to a JSON string\n",
    "payload = json.dumps(json_payload)\n",
    "\n",
    "# Make the post request\n",
    "print(f\"Starting test: chat completion through an HTTP POST with streaming using {model}\")\n",
    "print(\"+*\"*20)\n",
    "\n",
    "response = requests.post(url, headers=headers, data=payload)\n",
    "\n",
    "# Check if the request was successful\n",
    "if response.status_code == 200:\n",
    "    onse = requests.post(url, headers=headers, data=payload, stream=True)\n",
    "\n",
    "# Check if the request was successful\n",
    "if response.status_code == 200:\n",
    "    # Process the streaming response\n",
    "    for line in response.iter_lines():\n",
    "        if line:\n",
    "            # Skip the \"data: \" prefix and parse the JSON\n",
    "            line = line.decode('utf-8')\n",
    "            if line.startswith('data: '):\n",
    "                if line == 'data: [DONE]':\n",
    "                    break\n",
    "                \n",
    "                json_data = json.loads(line[6:])  # Skip the 'data: ' prefix\n",
    "                \n",
    "                # Extract the content from the delta if it exists\n",
    "                if 'choices' in json_data and json_data['choices'] and 'delta' in json_data['choices'][0]:\n",
    "                    delta = json_data['choices'][0]['delta']\n",
    "                    if 'content' in delta and delta['content']:\n",
    "                        print(f\"{color.BOLD} Chunk: {color.END} {delta['content']}\", end=\"\\n\")\n",
    "\n",
    "print(\"+*\"*20)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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.12.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}