01-getting-started/end-to-end-example.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "8iPJnDH0HhXB"
   },
   "source": [
    "<img src=\"https://imagedelivery.net/Dr98IMl5gQ9tPkFM5JRcng/3e5f6fbd-9bc6-4aa1-368e-e8bb1d6ca100/Ultra\" alt=\"Image description\" width=\"160\" />\n",
    "\n",
    "<br/>\n",
    "\n",
    "# Introduction to Contextual AI Platform using the Python Client\n",
    "\n",
    "Contextual AI lets you create and use generative AI agents. This notebook introduces an end-to-end example workflow for creating a Retrieval-Augmented Generation (RAG) agent for a financial use case. The agent will answer questions based on the documents provided, but avoid any forward looking statements, e.g., Tell me about sales in 2028. This notebook uses the python client.\n",
    "\n",
    "This notebook covers the following steps:\n",
    "- Creating a Datastore\n",
    "- Ingesting Documents\n",
    "- Creating an RAG Agent\n",
    "- Querying an RAG Agent\n",
    "- Evaluating an RAG Agent\n",
    "- Improving the RAG Agent (Updating prompt and tuning)\n",
    "\n",
    "With the exception of the tuning model, the rest of the notebook can be run in under 15 minutes. The full documentation is available at [docs.contextual.ai](https://docs.contextual.ai/)\n",
    "\n",
    "To run this notebook interactively, you can open it in Google Colab. However, make sure to copy over the additional files.\n",
    "\n",
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/ContextualAI/examples/blob/main/01-getting-started/end-to-end-example.ipynb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install contextual-client"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import requests\n",
    "import json\n",
    "from pathlib import Path\n",
    "from typing import List, Optional, Dict\n",
    "from IPython.display import display, JSON\n",
    "import pandas as pd\n",
    "from contextual import ContextualAI\n",
    "import ast"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "id": "4wJG66VTIQvO"
   },
   "outputs": [],
   "source": [
    "#Setup API key\n",
    "#os.environ[\"CONTEXTUAL_API_KEY\"] = API_KEY  # You can store the API key is stored as the environment variable \n",
    "\n",
    "client = ContextualAI(\n",
    "    api_key=os.environ.get(\"CONTEXTUAL_API_KEY\"),  # This is the default and can be omitted\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "lcqh_-j1MzCn"
   },
   "source": [
    "## Step 1: Create your Datastore\n",
    "\n",
    "\n",
    "You will need to first create a datastore for your agent using the  /datastores endpoint. A datastore is secure storage for data. Each agent will have it's own datastore for storing data securely."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "wlulbIvjdbZA",
    "outputId": "7029ff7d-193a-4cf1-98e8-0e29451ec446"
   },
   "outputs": [],
   "source": [
    "result = client.datastores.create(name=\"Demo_fin_rag\")\n",
    "datastore_id = result.id\n",
    "print(f\"Datastore ID: {datastore_id}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "_IkAep8Vf29_"
   },
   "source": [
    "## Step 2: Ingest Documents into your Datastore\n",
    "\n",
    "You can now ingest documents into your Agent's datastore using the /datastores endpoint. Documents must be a PDF or HTML file.\n",
    "\n",
    "\n",
    "I am using a example PDF. You can also use your own documents here. If you have very long documents (hundreds of pages), processing can take longer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if not os.path.exists('data'):\n",
    "    os.makedirs('data')\n",
    "\n",
    "if not os.path.exists('data/Apple.pdf'):\n",
    "    print(f\"Fetching data/Apple.pdf\")\n",
    "    response = requests.get(\"https://raw.githubusercontent.com/ContextualAI/examples/refs/heads/main/01-getting-started/data/Apple.pdf\")\n",
    "    with open('data/Apple.pdf', 'wb') as f:\n",
    "        f.write(response.content)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's ingest the file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('data/Apple.pdf', 'rb') as f:\n",
    "    ingestion_result = client.datastores.documents.ingest(datastore_id, file=f)\n",
    "    document_id = ingestion_result.id\n",
    "    print(f\"Successfully uploaded to datastore {datastore_id}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "9mZdyD5cK8lc"
   },
   "source": [
    "Once ingested, you can view the list of documents, see their metadata, and also delete documents."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "j78jlN1_lIbM"
   },
   "outputs": [],
   "source": [
    "metadata = client.datastores.documents.metadata(datastore_id = datastore_id, document_id = document_id)\n",
    "print(\"Document metadata:\", metadata)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "IcTyR5QeHw7z"
   },
   "source": [
    "## Step 3: Create your Agent\n",
    "\n",
    "Next let's create the Agent and modify it to our needs.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "sr4jmjF_LtTp"
   },
   "source": [
    "Some additional parameters include setting a system prompt or using a previously tuned model.\n",
    "\n",
    "`system_prompt` is used for the instructions that your RAG system references when generating responses. Note that we do not guarantee that the system will follow these instructions exactly.\n",
    "\n",
    "`llm_model_id` is the optional model ID of a tuned model to use for generation. Contextual AI will use a default model if none is specified."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "id": "r5jsWXpBPLtx"
   },
   "outputs": [],
   "source": [
    "system_prompt = '''\n",
    "You are an AI assistant specialized in financial analysis and reporting. Your responses should be precise, accurate, and sourced exclusively from official financial documentation provided to you. Please follow these guidelines:\n",
    "\n",
    "Data Analysis & Response Quality:\n",
    "* Only use information explicitly stated in provided documentation (e.g., earnings releases, financial statements, investor presentations)\n",
    "* Present comparative analyses using structured formats with tables and bullet points where appropriate\n",
    "* Include specific period-over-period comparisons (quarter-over-quarter, year-over-year) when relevant\n",
    "* Maintain consistency in numerical presentations (e.g., consistent units, decimal places)\n",
    "* Flag any one-time items or special charges that impact comparability\n",
    "\n",
    "Technical Accuracy:\n",
    "* Use industry-standard financial terminology\n",
    "* Define specialized acronyms on first use\n",
    "* Never interchange distinct financial terms (e.g., revenue, profit, income, cash flow)\n",
    "* Always include units with numerical values\n",
    "* Pay attention to fiscal vs. calendar year distinctions\n",
    "* Present monetary values with appropriate scale (millions/billions)\n",
    "\n",
    "Response Format:\n",
    "* Begin with a high-level summary of key findings when analyzing data\n",
    "* Structure detailed analyses in clear, hierarchical formats\n",
    "* Use markdown for lists, tables, and emphasized text\n",
    "* Maintain a professional, analytical tone\n",
    "* Present quantitative data in consistent formats (e.g., basis points for ratios)\n",
    "\n",
    "Critical Guidelines:\n",
    "* Avoid opinions, speculation, or assumptions\n",
    "* If information is unavailable or irrelevant, clearly state this without additional commentary\n",
    "* Answer questions directly, then stop\n",
    "* Do not reference source document names or file types in responses\n",
    "* Focus only on information that directly answers the query\n",
    "\n",
    "For any analysis, provide comprehensive insights using all relevant available information while maintaining strict adherence to these guidelines and focusing on delivering clear, actionable information.\n",
    "'''\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's create our agent. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "ymCcfHR1Lra6",
    "outputId": "310369a2-90a6-4b40-e3f1-fcc90f117d2f"
   },
   "outputs": [],
   "source": [
    "app_response = client.agents.create(\n",
    "    name=\"Demo-Finance Forward Looking\",\n",
    "    description=\"Research Agent using only Historical Information\",\n",
    "    system_prompt=system_prompt,\n",
    "    datastore_ids=[datastore_id]\n",
    ")\n",
    "agent_id= app_response.id\n",
    "print(f\"Agent ID created: {agent_id}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "WFq4oMe9gMuz"
   },
   "source": [
    "## Step 4: Query your Agent\n",
    "\n",
    "Let's query our agent to see if its working. The required information is the agent_id and messages.  \n",
    "\n",
    "Optional information includes parameters for streaming, conversation_id, and model_id if using a different fine tuned model."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "QOW94UdMySbp"
   },
   "source": [
    "**Note:** It may take a few minutes for the document to be ingested and processed. The Assistant will give a detailed answer once the documents are ingested."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "wE4d616-rseT",
    "outputId": "dc58d319-0ef6-4d2a-ea30-5d2d8e109db3"
   },
   "outputs": [],
   "source": [
    "query_result = client.agents.query.create(\n",
    "    agent_id=agent_id,\n",
    "    messages=[{\n",
    "        \"content\": \"what was the sales for Apple in 2022\",\n",
    "        \"role\": \"user\"\n",
    "    }]\n",
    ")\n",
    "print(query_result.message.content)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "There is lots more information you can access from the query result. You can the retrieved documents for example.   "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "2ozRLyVGS9xb",
    "outputId": "3b0547c9-940c-425b-fd19-5fae2fff72c3"
   },
   "outputs": [],
   "source": [
    "import base64\n",
    "import io\n",
    "from PIL import Image\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "def display_base64_image(base64_string):\n",
    "    # Decode base64 string\n",
    "    img_data = base64.b64decode(base64_string)\n",
    "    \n",
    "    # Create PIL Image object\n",
    "    img = Image.open(io.BytesIO(img_data))\n",
    "    \n",
    "    # Display using matplotlib\n",
    "    plt.figure(figsize=(10, 10))\n",
    "    plt.imshow(img)\n",
    "    plt.axis('off')\n",
    "    plt.show()\n",
    "    \n",
    "    return img\n",
    "\n",
    "# Retrieve the first referenced document `retrieval_contents[0]\n",
    "ret_result = client.agents.query.retrieval_info(message_id = query_result.message_id, agent_id=agent_id,content_ids=[query_result.retrieval_contents[0].content_id])\n",
    "print(\"\\nRetrieval Info:\", ret_result)\n",
    "base64_string = ret_result.content_metadatas[0].page_img\n",
    "img = display_base64_image(base64_string)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "5annKXLeNCGT"
   },
   "source": [
    "## Step 5: Evaluate your Agent\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "yLVlvOTfNctA"
   },
   "source": [
    "Evaluation endpoints allow you to evaluate your Agent using a set of prompts (questions) and reference (gold) answers. We support two metrics: equivalence and groundedness.\n",
    "\n",
    "Equivalance evaluates if the Agent response is equivalent to the ground truth (model-driven binary classification).  \n",
    "Groundedness decomposes the Agent response into claims and then evaluates if the claims are grounded by the retrieved documents.\n",
    "\n",
    "For those using unit tests, we also offer our `lmunit` endpoint, get more details [here](https://contextual.ai/blog/lmunit/) \n",
    "\n",
    "Let's start with an evaluation dataset:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "if not os.path.exists('data/eval_short.csv'):\n",
    "    print(f\"Fetching data/eval_short.csv\")\n",
    "    response = requests.get(\"https://raw.githubusercontent.com/ContextualAI/examples/refs/heads/main/01-getting-started/data/eval_short.csv\")\n",
    "    with open('data/eval_short.csv', 'wb') as f:\n",
    "        f.write(response.content)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 293
    },
    "id": "qGGzUpKaFTOC",
    "outputId": "2a15a0b3-61bb-408c-cca5-7ced897537d7"
   },
   "outputs": [],
   "source": [
    "eval = pd.read_csv('data/eval_short.csv')\n",
    "eval.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Start an evaluation job that measures both equivalence and groundedness"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('data/eval_short.csv', 'rb') as f:\n",
    "    eval_result = client.agents.evaluate.create(\n",
    "        agent_id=agent_id,\n",
    "        metrics=[\"equivalence\", \"groundedness\"],\n",
    "        evalset_file=f\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Evaluation jobs can take time, especially longer ones. Here is how you can check on their status. This dataset usually takes a few minutes to run."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eval_status = client.agents.evaluate.jobs.metadata(agent_id=agent_id, job_id=eval_result.id)\n",
    "\n",
    "from tqdm import tqdm\n",
    "\n",
    "progress = tqdm(total=eval_status.job_metadata.num_predictions)\n",
    "progress.update(eval_status.job_metadata.num_processed_predictions)\n",
    "progress.set_description(\"Evaluation Progress\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once the evaluation job is completed you can look at the final results across your evaluation dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "metadata": {},
   "outputs": [],
   "source": [
    "def process_binary_evaluation(binary_response):\n",
    "    \"\"\"\n",
    "    Process BinaryAPIResponse into a pandas DataFrame.\n",
    "    \n",
    "    Args:\n",
    "        binary_response: BinaryAPIResponse from evaluate.retrieve\n",
    "        \n",
    "    Returns:\n",
    "        pd.DataFrame: Processed evaluation data\n",
    "    \"\"\"\n",
    "    # Read the binary content\n",
    "    content = binary_response.read()\n",
    "    \n",
    "    # Now decode the content\n",
    "    lines = content.decode('utf-8').strip().split('\\n')\n",
    "    \n",
    "    # Parse each line and flatten the results\n",
    "    data = []\n",
    "    for line in lines:\n",
    "        try:\n",
    "            entry = json.loads(line)\n",
    "            \n",
    "            # Parse the results string if it exists\n",
    "            if 'results' in entry:\n",
    "                results = ast.literal_eval(entry['results'])\n",
    "                del entry['results']\n",
    "                if isinstance(results, dict):\n",
    "                    for key, value in results.items():\n",
    "                        if isinstance(value, dict):\n",
    "                            for subkey, subvalue in value.items():\n",
    "                                entry[f'{key}_{subkey}'] = subvalue\n",
    "                        else:\n",
    "                            entry[key] = value\n",
    "            \n",
    "            data.append(entry)\n",
    "            \n",
    "        except Exception as e:\n",
    "            print(f\"Error processing line: {e}\")\n",
    "            continue\n",
    "    \n",
    "    return pd.DataFrame(data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eval_results = client.agents.datasets.evaluate.retrieve(dataset_name=eval_status.dataset_name, agent_id = agent_id)\n",
    "df = process_binary_evaluation(eval_results)\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here I am going to save the results of the evaluation to a csv file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "metadata": {},
   "outputs": [],
   "source": [
    "df.to_csv('eval_results_python.csv', index=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "692gSc7ZNyve"
   },
   "source": [
    "## Step 6: Improving your Agent\n",
    "\n",
    "Contexual AI give you multile methods for improving your overall agent performance. Two methods available via the API are modifying the system prompt or tuning the models. It's recommended you start with modifying the system prompt before using tuning. Please reach out to the account team for more best practices around tuning.\n",
    "\n",
    "Let's go through both options:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 6.1 Revising the system prompt\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After initial testing, you may want to revise the system prompt. Here I have an updated prompt with additional information in the critical guidelines section.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "metadata": {},
   "outputs": [],
   "source": [
    "system_prompt2 = '''\n",
    "You are an AI assistant specialized in financial analysis and reporting. Your responses should be precise, accurate, and sourced exclusively from official financial documentation provided to you. Please follow these guidelines:\n",
    "\n",
    "Data Analysis & Response Quality:\n",
    "* Only use information explicitly stated in provided documentation (e.g., earnings releases, financial statements, investor presentations)\n",
    "* Present comparative analyses using structured formats with tables and bullet points where appropriate\n",
    "* Include specific period-over-period comparisons (quarter-over-quarter, year-over-year) when relevant\n",
    "* Maintain consistency in numerical presentations (e.g., consistent units, decimal places)\n",
    "* Flag any one-time items or special charges that impact comparability\n",
    "\n",
    "Technical Accuracy:\n",
    "* Use industry-standard financial terminology\n",
    "* Define specialized acronyms on first use\n",
    "* Never interchange distinct financial terms (e.g., revenue, profit, income, cash flow)\n",
    "* Always include units with numerical values\n",
    "* Pay attention to fiscal vs. calendar year distinctions\n",
    "* Present monetary values with appropriate scale (millions/billions)\n",
    "\n",
    "Response Format:\n",
    "* Begin with a high-level summary of key findings when analyzing data\n",
    "* Structure detailed analyses in clear, hierarchical formats\n",
    "* Use markdown for lists, tables, and emphasized text\n",
    "* Maintain a professional, analytical tone\n",
    "* Present quantitative data in consistent formats (e.g., basis points for ratios)\n",
    "\n",
    "Critical Guidelines:\n",
    "* Do not make forward-looking projections unless directly quoted from source materials\n",
    "* Do not answer any questions for 2024 or later\n",
    "* Avoid opinions, speculation, or assumptions\n",
    "* If information is unavailable or irrelevant, clearly state this without additional commentary\n",
    "* Answer questions directly, then stop\n",
    "* Do not reference source document names or file types in responses\n",
    "* Focus only on information that directly answers the query\n",
    "\n",
    "For any analysis, provide comprehensive insights using all relevant available information while maintaining strict adherence to these guidelines and focusing on delivering clear, actionable information.\n",
    "'''\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's now update the agent. And verify that changes by checking the agent metadata."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "client.agents.update(agent_id=agent_id, system_prompt=system_prompt2)\n",
    "\n",
    "agent_config = client.agents.metadata(agent_id=agent_id)\n",
    "print (agent_config.system_prompt)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that you have updated the agent, go try running another evaluation job. You will see the performance has improved."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "dmRrBTy0OWzc"
   },
   "source": [
    "### 6.2 Tuning the Contextual System\n",
    "\n",
    "To run a tune job, you need to specificy a training file and an optional test file. (If no test file is provided, the tuning job will hold out a portion of the training file as the test set.)\n",
    "\n",
    "A tuning job requires fine tuning models and the expectation should be it will take a couple of hours to run.\n",
    "\n",
    "After the tune job completes, the metadata associated with the tune job will include evaluation results and a model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 6.2.1 Tuning dataset:  \n",
    "\n",
    "The file should be in JSON array format, where each element of the array is a JSON object represents a single training example. The four required fields are guideline, prompt, response, and knowledge.\n",
    "\n",
    "- knowledge field should be an array of strings, each string representing a piece of knowledge that the model should use to generate the response.\n",
    "\n",
    "- reference: The gold-standard answer to the prompt.\n",
    "\n",
    "- guideline field should be guidelines for the expected response.\n",
    "\n",
    "- prompt field should be a question or statement that the model should respond to."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "Rf1DeSP-VXR4",
    "outputId": "21421ffe-50d0-44be-b0b0-6f45b420260d"
   },
   "outputs": [],
   "source": [
    "if not os.path.exists('data/fin_train.jsonl'):\n",
    "    print(f\"Fetching data/fin_train.jsonl\")\n",
    "    response = requests.get(\"https://raw.githubusercontent.com/ContextualAI/examples/refs/heads/main/01-getting-started/data/fin_train.jsonl\")\n",
    "    with open('data/fin_train.jsonl', 'wb') as f:\n",
    "        f.write(response.content)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!head data/fin_train.jsonl"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "R1tcNqTv0gaQ"
   },
   "source": [
    "#### 6.2.2 Starting a tuning model job"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 82,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('data/fin_train.jsonl', 'rb') as f:\n",
    "    tune_job = client.agents.tune.create(\n",
    "    agent_id=agent_id,\n",
    "    training_file=f\n",
    ")\n",
    "    \n",
    "tune_job_id = tune_job.id\n",
    "print(f\"Tune job created: {tune_job_id}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "tHsd1Vzx0LDN",
    "outputId": "902694ff-0b03-4122-c434-7235ee8d66c7"
   },
   "outputs": [],
   "source": [
    "print (agent_id)\n",
    "print (tune_job_id)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "NGao-yw7yLVY"
   },
   "source": [
    "#### 6.2.3 Checking the status.\n",
    "\n",
    " You can check the status of the job using the API. For detailed information, refer to the API documentation. When the tuning job is complete, the status will turn to completed. The response payload will also contain evaluation_results, such as scores for equivalence."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "yBy8pONdyOnu",
    "outputId": "254f5737-e292-46a6-838a-c5f19dc624b5"
   },
   "outputs": [],
   "source": [
    "tune_metadata = client.agents.tune.jobs.metadata(\n",
    "    agent_id=agent_id,\n",
    "    job_id=tune_job_id\n",
    ")\n",
    "print(\"Tuning job metadata:\", tune_metadata)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "sg4YXZdYyR9X"
   },
   "source": [
    "When the tuning job is complete, the metadata would look like the following:\n",
    "```\n",
    "TuneJobMetadata(job_status='completed',\n",
    "                evaluation_results=None,\n",
    "                model_id='registry/tuned-model-101',\n",
    "                id='e44661f0-bagb-4919-b0df-bada36a31',\n",
    "                evaluation_metadata={'status': 'completed',\n",
    "                                     'metrics': {'equivalence_score': {'score': 0.873}},\n",
    "                                     'job_metadata': {'num_predictions': 200,\n",
    "                                                      'num_failed_predictions': 0,\n",
    "                                                      'num_successful_predictions': 200,\n",
    "                                                      'num_processed_predictions': 0},\n",
    "                                     'dataset_name': 'eval-results-101',\n",
    "                                     'model_name': 'registry/tuned-model-101',\n",
    "                                     'tune_job_id': 'e44661f0-bagb-4919-b0df-bada36a31'})\n",
    " ```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "M_77uMH5ybBh"
   },
   "source": [
    "#### 6.2.4 Updating the agent\n",
    "Once the tuned job is complete, you can deploy the tuned model via editing the agent through API. Note that currently a single fine-tuned model deployment is allowed per tenant. Please see the API doc for more information."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "client.agents.update(agent_id=agent_id, llm_model_id=tune_metadata.model_id)\n",
    "\n",
    "print(\"Agent updated with tuned model\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "yKAmujbKgUBj"
   },
   "source": [
    "## Next Steps\n",
    "\n",
    "In this Notebook, we've created a RAG agent in the finance domain, evaluating the agent, and tuned it for better performance. You can learn more at [docs.contextual.ai](https://docs.contextual.ai/). Finally, reach out to your account team if you have further questions or issues."
   ]
  }
 ],
 "metadata": {
  "colab": {
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "py3.10",
   "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.10.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}