{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "NTaMBjChz7m5"
   },
   "source": [
    "# Optimzing using Synthetic Q&A Data from Opik Traces\n",
    "\n",
    "You will need:\n",
    "\n",
    "1. A Comet account, for seeing Opik visualizations (free!) - [comet.com](https://comet.com)\n",
    "2. An OpenAI account, for using an LLM\n",
    "[platform.openai.com/settings/organization/api-keys](https://platform.openai.com/settings/organization/api-keys)\n",
    "\n",
    "This example will use:\n",
    "\n",
    "- [tinyqabenchmarkpp](https://pypi.org/project/tinyqabenchmarkpp/) to generate synthetic test dataset\n",
    "- [opik-optimizer](https://pypi.org/project/opik-optimizer/) to optimize prompts\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "yM1lU0dBBnJs"
   },
   "source": [
    "## Setup\n",
    "\n",
    "This pip-install takes about a minute."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "2Tx6HwuU1rB4"
   },
   "outputs": [],
   "source": [
    "%pip install opik-optimizer tinyqabenchmarkpp --upgrade"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This step configures the Opik library for your session. It will prompt for your Comet API key if not already set in your environment or through Opik's configuration."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "H0DNm-un_0Np"
   },
   "outputs": [],
   "source": [
    "import opik\n",
    "\n",
    "opik.configure()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For this example, we'll use OpenAI models, so we need to set our OpenAI API key:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "vN72mHQy_7Ou"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import getpass\n",
    "\n",
    "if \"OPENAI_API_KEY\" not in os.environ:\n",
    "    os.environ[\"OPENAI_API_KEY\"] = getpass.getpass(\"Enter your OpenAI API key: \")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Fetching Traces from Opik"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next up we will \"fetch\" existing traces of our AI application within Opik (we will use the demo project that ships with every Opik installation). \n",
    "\n",
    "This fetches traces from the demo project and formats them as a context string. Returns a formatted string with explanatory text and cleaned traces."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "OPIK_PROJECT_NAME = \"Demo chatbot 🤖\"\n",
    "\n",
    "# Will prompt for API key if not set\n",
    "opik.configure()\n",
    "\n",
    "# Fetch traces from the demo project\n",
    "#\n",
    "# Commented out the project name to\n",
    "# fetch all traces across all projects\n",
    "client = opik.Opik()\n",
    "traces = client.search_traces(\n",
    "    # project_name=OPIK_PROJECT_NAME,\n",
    "    max_results=40\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"Found {len(traces)} traces\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will define some helper functions to clean and traverse the traces as we don't wan to send noise to the LLM and break the input to the synthetic data generation step."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import re\n",
    "\n",
    "\n",
    "def extract_text_from_dict(d: dict) -> list[str]:\n",
    "    \"\"\"\n",
    "    Recursively extracts text from a dictionary.\n",
    "    \"\"\"\n",
    "    texts = []\n",
    "    for key, value in d.items():\n",
    "        if isinstance(value, str):\n",
    "            texts.append(value)\n",
    "        elif isinstance(value, dict):\n",
    "            texts.extend(extract_text_from_dict(value))\n",
    "        elif isinstance(value, list):\n",
    "            for item in value:\n",
    "                if isinstance(item, str):\n",
    "                    texts.append(item)\n",
    "                elif isinstance(item, dict):\n",
    "                    texts.extend(extract_text_from_dict(item))\n",
    "    return texts\n",
    "\n",
    "\n",
    "def clean_text(text: str) -> str:\n",
    "    \"\"\"\n",
    "    Cleans text by removing special characters and normalizing whitespace.\n",
    "    \"\"\"\n",
    "    if not text:\n",
    "        return \"\"\n",
    "    # Replace special characters with spaces, but keep basic punctuation\n",
    "    text = re.sub(r'[^\\w\\s.,!?;:\\'\"-]', \" \", text)\n",
    "    # Normalize whitespace\n",
    "    text = re.sub(r\"\\s+\", \" \", text)\n",
    "    # Remove any leading/trailing whitespace\n",
    "    text = text.strip()\n",
    "    return text"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We are now ready to extract and clean the text from the traces."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Extract and clean text from traces\n",
    "cleaned_texts = []\n",
    "for i, trace in enumerate(traces):\n",
    "    # Extract from input\n",
    "    if trace.input:\n",
    "        if isinstance(trace.input, dict):\n",
    "            texts = extract_text_from_dict(trace.input)\n",
    "            for text in texts:\n",
    "                cleaned = clean_text(text)\n",
    "                if cleaned:\n",
    "                    cleaned_texts.append(cleaned)\n",
    "        elif isinstance(trace.input, str):\n",
    "            cleaned = clean_text(trace.input)\n",
    "            if cleaned:\n",
    "                cleaned_texts.append(cleaned)\n",
    "\n",
    "    # Extract from output\n",
    "    if trace.output:\n",
    "        if isinstance(trace.output, dict):\n",
    "            texts = extract_text_from_dict(trace.output)\n",
    "            for text in texts:\n",
    "                cleaned = clean_text(text)\n",
    "                if cleaned:\n",
    "                    cleaned_texts.append(cleaned)\n",
    "        elif isinstance(trace.output, str):\n",
    "            cleaned = clean_text(trace.output)\n",
    "            if cleaned:\n",
    "                cleaned_texts.append(cleaned)\n",
    "\n",
    "    # Extract from metadata if it exists\n",
    "    if trace.metadata:\n",
    "        if isinstance(trace.metadata, dict):\n",
    "            texts = extract_text_from_dict(trace.metadata)\n",
    "            for text in texts:\n",
    "                cleaned = clean_text(text)\n",
    "                if cleaned:\n",
    "                    cleaned_texts.append(cleaned)\n",
    "\n",
    "if not cleaned_texts:\n",
    "    print(\"Debug: No text content found in traces. Here's what we got:\")\n",
    "    for i, trace in enumerate(traces[:5]):  # Show first 5 traces for debugging\n",
    "        print(f\"\\nTrace {i}:\")\n",
    "        print(f\"Input: {trace.input}\")\n",
    "        print(f\"Output: {trace.output}\")\n",
    "        print(f\"Metadata: {trace.metadata}\")\n",
    "    raise ValueError(\"No valid text content found in traces\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's quickly inspect the traces we have."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "display(traces[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Remove any duplicates while preserving the order"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Remove duplicates while preserving order\n",
    "seen = set()\n",
    "unique_texts = []\n",
    "for text in cleaned_texts:\n",
    "    if text not in seen:\n",
    "        seen.add(text)\n",
    "        unique_texts.append(text)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we need to create the `context` to pass to the synthetic data generation step"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "context = f\"\"\"\n",
    "This is a collection of AI/LLM conversation traces from\n",
    "a given Comet Opik observability project. The following\n",
    "text contains various interactions and responses that\n",
    "can be used to generate relevant questions and answers.\n",
    "<input>\n",
    "{chr(10).join(unique_texts)}\n",
    "</input>\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"Found and cleaned {len(unique_texts)} unique text segments from traces\")\n",
    "print(f\"Total context length: {len(context)} characters\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generating Synthethic Data\n",
    "\n",
    "We are now ready to generate the synthethic data using `tinyqabenchmarkpp`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Model for tinyqabenchmarkpp\n",
    "TQB_GENERATOR_MODEL = \"openai/gpt-4o-mini\"\n",
    "\n",
    "# Number of questions to generate\n",
    "TQB_NUM_QUESTIONS = 20\n",
    "\n",
    "# Languages to generate questions in\n",
    "TQB_LANGUAGES = \"en\"\n",
    "\n",
    "# Categories to generate questions in\n",
    "TQB_CATEGORIES = (\n",
    "    \"use context provided and elaborate on it to generate a more detailed answers\"\n",
    ")\n",
    "\n",
    "# Difficulty of the questions to generate\n",
    "TQB_DIFFICULTY = \"medium\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Command to generate the synthetic data\n",
    "command = [\n",
    "    \"python\",\n",
    "    \"-m\",\n",
    "    \"tinyqabenchmarkpp.generate\",\n",
    "    \"--num\",\n",
    "    str(TQB_NUM_QUESTIONS),\n",
    "    \"--languages\",\n",
    "    TQB_LANGUAGES,\n",
    "    \"--categories\",\n",
    "    TQB_CATEGORIES,\n",
    "    \"--difficulty\",\n",
    "    TQB_DIFFICULTY,\n",
    "    \"--model\",\n",
    "    TQB_GENERATOR_MODEL,\n",
    "    \"--str-output\",\n",
    "    \"--context\",\n",
    "    context,\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we run the synthetic data generation step, please be patient as the language model is called."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Use a subprocess to run the command\n",
    "import subprocess\n",
    "\n",
    "process = subprocess.run(command, capture_output=True, text=True, check=True)\n",
    "\n",
    "if process.stderr:\n",
    "    # Print the errors\n",
    "    print(\"tinyqabenchmarkpp errors:\")\n",
    "    print(process.stderr)\n",
    "else:\n",
    "    # Print the output\n",
    "    print(\"Synthetic data generated successfully\")\n",
    "    print(process.stdout)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Store New Dataset in Opik\n",
    "We can use the Opik SDK to push this dataset to Opik"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "generated_data = process.stdout"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Helper function to process the JSONL response and push to Opik\n",
    "Once we have defined we will be able to run this"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "\n",
    "def load_synthetic_data_to_opik(data_str):\n",
    "    \"\"\"Load JSONL synthetic data into Opik as a dataset.\"\"\"\n",
    "    items = []\n",
    "    for line in data_str.strip().split(\"\\n\"):\n",
    "        try:\n",
    "            data = json.loads(line)\n",
    "            if not isinstance(data, dict):\n",
    "                continue\n",
    "            item = {\n",
    "                \"question\": data.get(\"text\"),\n",
    "                \"answer\": data.get(\"label\"),\n",
    "                \"generated_context\": data.get(\"context\"),\n",
    "                \"category\": data.get(\"tags\", {}).get(\"category\"),\n",
    "                \"difficulty\": data.get(\"tags\", {}).get(\"difficulty\"),\n",
    "            }\n",
    "            if item[\"question\"] and item[\"answer\"]:\n",
    "                items.append(item)\n",
    "        except Exception:\n",
    "            continue\n",
    "\n",
    "    if not items:\n",
    "        print(\"No valid items found.\")\n",
    "        return None\n",
    "\n",
    "    dataset_name = (\n",
    "        f\"demo-tinyqab-dataset-{TQB_CATEGORIES.replace(',', '_')}-{TQB_NUM_QUESTIONS}\"\n",
    "    )\n",
    "    dataset_name = \"\".join(\n",
    "        c if c.isalnum() or c in [\"-\", \"_\"] else \"_\" for c in dataset_name\n",
    "    )\n",
    "\n",
    "    opik_client = opik.Opik()\n",
    "    dataset = opik_client.get_or_create_dataset(\n",
    "        name=dataset_name,\n",
    "        description=f\"Synthetic QA from tinyqabenchmarkpp for {TQB_CATEGORIES}\",\n",
    "    )\n",
    "    dataset.insert(items)\n",
    "    print(f\"Opik Dataset '{dataset.name}' created with ID: {dataset.id}\")\n",
    "    return dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Push the data to Opik using helper function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "opik_synthetic_dataset = load_synthetic_data_to_opik(generated_data)\n",
    "if not opik_synthetic_dataset:\n",
    "    print(\"Failed to load synthetic data into Opik. Exiting.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Agent Optimization Using Synthetic Data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Lets import the required packages for the Opik Agent Optimizer SDK"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from opik_optimizer import MetaPromptOptimizer\n",
    "from opik.evaluation.metrics import LevenshteinRatio"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We need to setup some intputs to our optimizer such as our starting prompt and some other configuration items."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from opik_optimizer import ChatPrompt\n",
    "\n",
    "# Initial prompt for the optimizer\n",
    "OPTIMIZER_INITIAL_PROMPT = ChatPrompt(\n",
    "    messages=[\n",
    "        {\"role\": \"system\", \"content\": \"You are a helpful assistant.\"},\n",
    "        {\"role\": \"user\", \"content\": \"{question}\"},\n",
    "    ],\n",
    "    project_name=OPIK_PROJECT_NAME,\n",
    ")\n",
    "\n",
    "# Model for Opik Agent Optimizer\n",
    "OPTIMIZER_MODEL = \"openai/gpt-4o-mini\"\n",
    "\n",
    "# Population size for the optimizer\n",
    "# Reduced for quicker demo\n",
    "OPTIMIZER_POPULATION_SIZE = 5\n",
    "\n",
    "# Number of generations for the optimizer\n",
    "# Reduced for quicker demo\n",
    "OPTIMIZER_NUM_GENERATIONS = 2\n",
    "\n",
    "# Number of samples from dataset for optimization eval\n",
    "OPTIMIZER_N_SAMPLES_OPTIMIZATION = 10"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can setup the metric configuration used for the evaluation, as well as he task_config for passing in the dataset headings and initial prompt.\n",
    "\n",
    "We finally pass this to our optimizer to set this up. We are opting to use the `MetaPromptOptimizer` optimizer in the SDK."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Metric Configuration\n",
    "def levenshtein_ratio(dataset_item, llm_output):\n",
    "    return LevenshteinRatio().score(reference=dataset_item[\"answer\"], output=llm_output)\n",
    "\n",
    "\n",
    "# Initialize the optimizer\n",
    "optimizer = MetaPromptOptimizer(\n",
    "    model=OPTIMIZER_MODEL,\n",
    "    population_size=OPTIMIZER_POPULATION_SIZE,\n",
    "    num_generations=OPTIMIZER_NUM_GENERATIONS,\n",
    "    infer_output_style=True,\n",
    "    verbose=1,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can run the optimizer on the dataset and initial starting prompt to find the best prompt based on our synthetic data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Commented out if you wan to pull the dataset from Opik without having\n",
    "# to generate the synthetic data again\n",
    "\n",
    "# import opik\n",
    "# opik_client = opik.Opik()\n",
    "# opik_synthetic_dataset = opik_client.get_or_create_dataset(\"demo-tinyqab-dataset-use_context_provided_and_elaborate_on_it_to_generate_a_more_detailed_answers-20\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run the optimizer\n",
    "result = optimizer.optimize_prompt(\n",
    "    prompt=OPTIMIZER_INITIAL_PROMPT,\n",
    "    dataset=opik_synthetic_dataset,\n",
    "    metric=levenshtein_ratio,\n",
    "    n_samples=OPTIMIZER_N_SAMPLES_OPTIMIZATION,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Optimization process finished\n",
    "We can output our results to show"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "result.display()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "uSDJ1bFx51kd"
   },
   "source": [
    "# Next Steps\n",
    "\n",
    "You can try out other optimizers. More details can be found in the [Opik Agent Optimizer documentation](https://www.comet.com/docs/opik/agent_optimization/overview)."
   ]
  }
 ],
 "metadata": {
  "colab": {
   "provenance": []
  },
  "kernelspec": {
   "display_name": "py312_llm_eval",
   "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.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}