02-hands-on-lab/lab3_improve_agent.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Lab 3: Improving your Agent\n",
    "\n",
    "Contexual AI gives you multiple 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": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import requests\n",
    "from contextual import ContextualAI\n",
    "from IPython.display import display, Markdown"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "🔑 Replace \"your_api_key\" with your actual API key 👇🏼"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "API_KEY = \"key-...\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "client = ContextualAI(\n",
    "    api_key=API_KEY,  # This is the default and can be omitted\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Important**: Replace \"agent_id\" with the Agent you created in Lab 1 👇🏼"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "agent_id=\"...\""
   ]
  },
  {
   "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": 14,
   "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": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\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"
     ]
    }
   ],
   "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": "code",
   "execution_count": null,
   "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/02-hands-on-lab/data/eval_short.csv\")\n",
    "    with open('data/eval_short.csv', 'wb') as f:\n",
    "        f.write(response.content)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "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": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "EvaluationJobMetadata(dataset_name='eval_short_2ae7ed23-9456-4277-9c80-ff5dabf0f2c2_results', job_metadata=JobMetadata(num_failed_predictions=0, num_predictions=12, num_successful_predictions=12), metrics={'equivalence_score': {'score': 0.8333333333333334}, 'groundedness_score': {'score': 0.5833333333333334}}, status='completed')"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "eval_status = client.agents.evaluate.jobs.metadata(agent_id=agent_id, job_id=eval_result.id)\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": {
    "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": 37,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "Rf1DeSP-VXR4",
    "outputId": "21421ffe-50d0-44be-b0b0-6f45b420260d"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[\n",
      "{\"guideline\": \"The response should clearly communicate strategic priorities and potential risks associated with forward-looking statements.\", \"prompt\": \"What are the strategic priorities outlined in Apple's 2024 investor call?\", \"knowledge\": [\"Apple plans to expand its services segment, targeting 15% YoY growth.\", \"Investments in AR/VR technology are expected to increase by 30%.\", \"Sustainability goals include achieving carbon neutrality across its supply chain by 2030.\", \"Risks include potential regulatory scrutiny in the EU and economic headwinds impacting consumer spending.\"], \"reference\": \"Apple's 2024 strategic priorities focus on expanding its services segment (targeting 15% YoY growth), increasing investments in AR/VR by 30%, and advancing sustainability goals to achieve carbon neutrality by 2030. However, the company faces risks such as EU regulatory challenges and economic pressures on consumer spending.\"},\n",
      "{\"guideline\": \"The response should focus on actionable insights based on future projections.\", \"prompt\": \"What is the revenue outlook for Tesla's energy storage segment in 2025?\", \"knowledge\": [\"Tesla expects its energy storage segment revenue to grow 50% YoY in 2025.\", \"New gigafactories coming online in Europe and Asia will boost production capacity.\", \"Potential risks include supply chain constraints and geopolitical tensions in key markets.\"], \"reference\": \"Tesla projects 50% YoY revenue growth for its energy storage segment in 2025, driven by expanded production capacity from new gigafactories in Europe and Asia. However, supply chain challenges and geopolitical risks may impact these projections.\"},\n",
      "{\"guideline\": \"The response should address uncertainties inherent in forward-looking projections.\", \"prompt\": \"What are the growth drivers for Microsoft in the next five years?\", \"knowledge\": [\"Microsoft plans to focus on cloud computing and AI services, aiming for a 20% CAGR by 2028.\", \"Significant R&D investments in quantum computing are anticipated to unlock new revenue streams.\", \"The gaming division, bolstered by the Activision Blizzard acquisition, is expected to see double-digit growth.\", \"Challenges include evolving regulatory landscapes and competitive pressures from AWS and Google Cloud.\"], \"reference\": \"Microsoft's growth over the next five years is expected to be driven by cloud computing and AI services (targeting a 20% CAGR by 2028), R&D in quantum computing, and gaming division expansion post-Activision Blizzard acquisition. Regulatory and competitive challenges may pose risks to these goals.\"},\n",
      "{\"guideline\": \"The response should align with compliance requirements for forward-looking statements.\", \"prompt\": \"What is Alphabet's strategy for navigating AI competition in the coming decade?\", \"knowledge\": [\"Alphabet aims to maintain leadership in AI with a focus on responsible AI principles.\", \"Planned investments include $10 billion annually in AI research and infrastructure.\", \"Strategic partnerships with academic institutions will support innovation and talent development.\", \"Potential risks include ethical concerns and increased scrutiny from regulators worldwide.\"], \"reference\": \"Alphabet's long-term strategy to stay competitive in AI includes $10 billion in annual investments for research and infrastructure, fostering innovation through partnerships, and adhering to responsible AI principles. Regulatory and ethical challenges remain key considerations.\"},\n",
      "{\"guideline\": \"The response should be accurate, concise, and provide actionable insights for financial decision-making.\", \"prompt\": \"What are the key financial metrics from Apple's Q3 2023 earnings report?\", \"knowledge\": [\"Revenue: $81.8 billion, down 1% YoY.\", \"Net income: $19.9 billion, flat YoY.\", \"Earnings per share (EPS): $1.26, up 5% YoY.\", \"Strong growth in services revenue offset declines in product sales.\"], \"reference\": \"Apple's Q3 2023 earnings report highlighted $81.8 billion in revenue, a 1% YoY decline, with $19.9 billion in net income (flat YoY). EPS grew 5% YoY to $1.26, driven by services revenue growth despite weaker product sales.\"},\n",
      "{\"guideline\": \"The response should prioritize clarity and use financial terminology correctly.\", \"prompt\": \"What factors drove Tesla's stock price increase in November 2023?\", \"knowledge\": [\"Tesla announced record-breaking deliveries of 450,000 units in Q3 2023.\", \"Margins improved to 21% due to cost reductions in battery production.\", \"Market optimism surrounding Tesla's AI-based driver-assistance technology.\", \"Tesla's Cybertruck production officially started in late October 2023.\"], \"reference\": \"Tesla's stock price increase in November 2023 was driven by record Q3 deliveries of 450,000 units, improved margins of 21%, and positive sentiment around its AI-driven driver-assistance technology. Cybertruck production also began, boosting investor confidence.\"},\n",
      "{\"guideline\": \"The response should avoid speculation and focus only on the information provided.\", \"prompt\": \"Is it a good time to invest in emerging market equities?\", \"knowledge\": [\"Emerging markets (EM) are experiencing mixed economic conditions.\", \"India and Brazil show strong GDP growth and stable inflation rates.\", \"China's recovery post-COVID remains slow, with muted industrial output.\", \"EM equities underperform relative to developed markets in 2023.\"], \"reference\": \"The decision to invest in emerging market equities depends on specific regions. India and Brazil show strong economic performance, but China's slow recovery and overall EM underperformance compared to developed markets in 2023 warrant caution.\"},\n",
      "{\"guideline\": \"The response should provide comparative analysis using numerical data.\", \"prompt\": \"How does Meta's revenue growth compare to Alphabet's in Q2 2023?\", \"knowledge\": [\"Meta reported Q2 2023 revenue growth of 11% YoY, reaching $32 billion.\", \"Alphabet posted revenue growth of 7% YoY, totaling $74.6 billion.\", \"Both companies saw strong performance in digital advertising.\", \"Meta's growth outpaced Alphabet's due to increased ad spend on Instagram and Reels.\"], \"reference\": \"In Q2 2023, Meta's revenue grew 11% YoY to $32 billion, outpacing Alphabet's 7% YoY growth to $74.6 billion. Meta's faster growth was driven by higher ad spending on Instagram and Reels.\"},\n",
      "{\"guideline\": \"The response should maintain neutrality and avoid subjective interpretations.\", \"prompt\": \"What risks are highlighted in Microsoft's Q1 2024 earnings report?\", \"knowledge\": [\"The report cites slower Azure growth, now at 24% YoY (vs. 35% YoY a year earlier).\", \"Macroeconomic concerns impacting enterprise IT budgets.\", \"Regulatory scrutiny over Activision Blizzard acquisition.\", \"Rising competition in AI cloud services from AWS and Google Cloud.\"], \"reference\": \"Microsoft's Q1 2024 earnings report highlights risks including slowing Azure growth (24% YoY vs. 35% YoY last year), macroeconomic pressures on IT budgets, regulatory challenges with the Activision Blizzard deal, and competition in AI cloud services.\"},\n"
     ]
    }
   ],
   "source": [
    "with open('data/fin_train.jsonl', 'r') as f:\n",
    "    for i, line in enumerate(f):\n",
    "        if i < 10:\n",
    "            print(line.strip())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "R1tcNqTv0gaQ"
   },
   "source": [
    "#### 6.2.2 Starting a tuning model job"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Tune job created: 70e28e9f-e471-4876-afc6-d431c57c4c39\n"
     ]
    }
   ],
   "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": 39,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "tHsd1Vzx0LDN",
    "outputId": "902694ff-0b03-4122-c434-7235ee8d66c7"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "faf2cc13-a503-40e1-adc9-432b977d9b4a\n",
      "70e28e9f-e471-4876-afc6-d431c57c4c39\n"
     ]
    }
   ],
   "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": 40,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "yBy8pONdyOnu",
    "outputId": "254f5737-e292-46a6-838a-c5f19dc624b5"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Tuning job metadata: TuneJobMetadata(job_status='pending', evaluation_results=None, model_id=None)\n"
     ]
    }
   ],
   "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": {},
   "source": [
    "#### 6.2.5 Evaluate the agent\n",
    "Once the new model is deployed we can evaluate our agent again."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "EvaluationJobMetadata(dataset_name='eval_short_2ae7ed23-9456-4277-9c80-ff5dabf0f2c2_results', job_metadata=JobMetadata(num_failed_predictions=0, num_predictions=12, num_successful_predictions=12), metrics={'equivalence_score': {'score': 0.8333333333333334}, 'groundedness_score': {'score': 0.5833333333333334}}, status='completed')"
      ]
     },
     "execution_count": 41,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "eval_status = client.agents.evaluate.jobs.metadata(agent_id=agent_id, job_id=eval_result.id)\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": {
    "id": "yKAmujbKgUBj"
   },
   "source": [
    "## Next Steps\n",
    "\n",
    "In this workshop, we've created a RAG agent in the finance domain, evaluated 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. Thanks for coming! 👋"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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": 4
}