{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a target=\"_blank\" href=\"https://colab.research.google.com/github/okareo-ai/okareo-python-sdk/blob/main/examples/generation_eval.ipynb\">\n",
    "  <img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/>\n",
    "</a>\n",
    "\n",
    "## Welcome to Okareo!\n",
    "\n",
    "Get your API token from [https://app.okareo.com/](https://app.okareo.com/) and set it in the cell below. 👇\n",
    "   (Note: You will also need an OpenAI key.)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "OKAREO_API_KEY = \"<YOUR-OKAREO-API-KEY>\"\n",
    "OPENAI_API_KEY = \"<YOUR-OPENAI-API-KEY>\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install okareo\n",
    "%pip install openai"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Simple summarization prompt using OpenAI'a GPT 3.5 Turbo model\n",
    "from openai import OpenAI\n",
    "\n",
    "client = OpenAI(api_key=OPENAI_API_KEY)\n",
    "\n",
    "def get_turbo_summary(messages, model=\"gpt-3.5-turbo\", \n",
    "  temperature=0, max_tokens=500):\n",
    "  response = client.chat.completions.create(\n",
    "    model=model,\n",
    "    messages=messages,\n",
    "    temperature=temperature, \n",
    "    max_tokens=max_tokens,\n",
    "  )\n",
    "  return response.choices[0].message.content\n",
    "\n",
    "USER_PROMPT_TEMPLATE = \"{scenario_input}\"\n",
    "\n",
    "SUMMARIZATION_CONTEXT_TEMPLATE = \"\"\"\n",
    "You will be provided with text.\n",
    "Summarize the text in 1 simple sentence.\n",
    "\"\"\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Evaluate the scenario and model combination and then get a link to the results on Okareo\n",
    "import random\n",
    "import string\n",
    "import os\n",
    "import tempfile\n",
    "\n",
    "from okareo import Okareo\n",
    "from okareo.model_under_test import OpenAIModel\n",
    "from okareo_api_client.models.test_run_type import TestRunType\n",
    "\n",
    "okareo = Okareo(OKAREO_API_KEY)\n",
    "random_string = ''.join(random.choices(string.ascii_letters, k=5))\n",
    "mut_name = f\"OpenAI Summarization Model - {random_string}\"\n",
    "eval_name = f\"Summarization Run - {random_string}\"\n",
    "\n",
    "# Register the model to use in the test run\n",
    "model_under_test = okareo.register_model(\n",
    "    name=mut_name,\n",
    "    model=OpenAIModel(\n",
    "        model_id=\"gpt-3.5-turbo\",\n",
    "        temperature=0,\n",
    "        system_prompt_template=SUMMARIZATION_CONTEXT_TEMPLATE,\n",
    "        user_prompt_template=USER_PROMPT_TEMPLATE,\n",
    "    ),\n",
    ")\n",
    "\n",
    "# Create a scenario to evaluate the model with\n",
    "# Get jsonl file from Okareo's SDK repo\n",
    "webbizz_articles = os.popen('curl https://raw.githubusercontent.com/okareo-ai/okareo-python-sdk/main/examples/webbizz_10_articles.jsonl').read()\n",
    "temp_dir = tempfile.gettempdir()\n",
    "file_path = os.path.join(temp_dir, \"webbizz_10_articles.jsonl\")\n",
    "with open(file_path, \"w+\") as file:\n",
    "    lines = webbizz_articles.split('\\n')\n",
    "    # Use the first 3 json objects to make a scenario set with 3 scenarios\n",
    "    for i in range(3):\n",
    "        file.write(f\"{lines[i]}\\n\")\n",
    "scenario = okareo.upload_scenario_set(file_path=file_path, scenario_name=\"Webbizz Articles Scenario\")\n",
    "\n",
    "# make sure to clean up tmp file\n",
    "os.remove(file_path)\n",
    "\n",
    "# Run the evaluation\n",
    "evaluation = model_under_test.run_test(\n",
    "    name=eval_name,\n",
    "    scenario=scenario,\n",
    "    api_key=OPENAI_API_KEY,\n",
    "    test_run_type=TestRunType.NL_GENERATION,\n",
    "    calculate_metrics=True,\n",
    ")\n",
    "\n",
    "print(f\"See results in Okareo: {evaluation.app_link}\")"
   ]
  }
 ],
 "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.11.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}