{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a9ab2a39-7c2d-4119-9dc7-8035fdfba3cb",
   "metadata": {},
   "source": [
    "# Fine-Tuning on LangSmith LLM Runs\n",
    "\n",
    "This notebook demonstrates how to directly load data from LangSmith's LLM runs and fine-tune a model on that data.\n",
    "The process is simple and comprises 3 steps.\n",
    "\n",
    "1. Select the LLM runs to train on.\n",
    "2. Use the LangSmithRunChatLoader to load runs as chat sessions.\n",
    "3. Fine-tune your model.\n",
    "\n",
    "Then you can use the fine-tuned model in your LangChain app.\n",
    "\n",
    "Before diving in, let's install our prerequisites.\n",
    "\n",
    "## Prerequisites\n",
    "\n",
    "Ensure you've installed langchain >= 0.0.311 and have configured your environment with your LangSmith API key."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "ef488003-514a-48b4-93f1-7de4417abf5d",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install -U langchain openai"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "473adce5-c863-49e6-85c3-049e0ec2222e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import uuid\n",
    "uid = uuid.uuid4().hex[:6]\n",
    "project_name = f\"Run Fine-tuning Walkthrough {uid}\"\n",
    "os.environ[\"LANGCHAIN_TRACING_V2\"] = \"true\"\n",
    "os.environ[\"LANGCHAIN_API_KEY\"] = \"YOUR API KEY\"\n",
    "os.environ[\"LANGCHAIN_PROJECT\"] = project_name"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8533ab63-d437-492a-aaec-ccca31167bf2",
   "metadata": {},
   "source": [
    "## 1. Select Runs\n",
    "The first step is selecting which runs to fine-tune on. A common case would be to select LLM runs within\n",
    "traces that have received positive user feedback. You can find examples of this in the[LangSmith Cookbook](https://github.com/langchain-ai/langsmith-cookbook/blob/main/exploratory-data-analysis/exporting-llm-runs-and-feedback/llm_run_etl.ipynb) and in the [docs](https://docs.smith.langchain.com/tracing/use-cases/export-runs/local).\n",
    "\n",
    "For the sake of this tutorial, we will generate some runs for you to use here. Let's try fine-tuning a\n",
    "simple function-calling chain."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "9a36d27f-2f3b-4148-b94a-9436fe8b00e0",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.pydantic_v1 import BaseModel, Field\n",
    "from enum import Enum\n",
    "\n",
    "\n",
    "class Operation(Enum):\n",
    "    add = \"+\"\n",
    "    subtract = \"-\"\n",
    "    multiply = \"*\"\n",
    "    divide = \"/\"\n",
    "\n",
    "class Calculator(BaseModel):\n",
    "    \"\"\"A calculator function\"\"\"\n",
    "    num1: float\n",
    "    num2: float\n",
    "    operation: Operation = Field(..., description=\"+,-,*,/\")\n",
    "\n",
    "    def calculate(self):\n",
    "        if self.operation == Operation.add:\n",
    "            return self.num1 + self.num2\n",
    "        elif self.operation == Operation.subtract:\n",
    "            return self.num1 - self.num2\n",
    "        elif self.operation == Operation.multiply:\n",
    "            return self.num1 * self.num2\n",
    "        elif self.operation == Operation.divide:\n",
    "            if self.num2 != 0:\n",
    "                return self.num1 / self.num2\n",
    "            else:\n",
    "                return \"Cannot divide by zero\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "89bcc676-27e8-40dc-a4d6-92cf28e0db58",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'description': 'A calculator function',\n",
      " 'name': 'Calculator',\n",
      " 'parameters': {'description': 'A calculator function',\n",
      "                'properties': {'num1': {'title': 'Num1', 'type': 'number'},\n",
      "                               'num2': {'title': 'Num2', 'type': 'number'},\n",
      "                               'operation': {'allOf': [{'description': 'An '\n",
      "                                                                       'enumeration.',\n",
      "                                                        'enum': ['+',\n",
      "                                                                 '-',\n",
      "                                                                 '*',\n",
      "                                                                 '/'],\n",
      "                                                        'title': 'Operation'}],\n",
      "                                             'description': '+,-,*,/'}},\n",
      "                'required': ['num1', 'num2', 'operation'],\n",
      "                'title': 'Calculator',\n",
      "                'type': 'object'}}\n"
     ]
    }
   ],
   "source": [
    "from langchain.utils.openai_functions import convert_pydantic_to_openai_function\n",
    "from langchain.pydantic_v1 import BaseModel\n",
    "from pprint import pprint\n",
    "\n",
    "openai_function_def = convert_pydantic_to_openai_function(Calculator)\n",
    "pprint(openai_function_def)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "cd44ff01-22cf-431a-8bf4-29a758d1fcff",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.prompts import ChatPromptTemplate\n",
    "from langchain.chat_models import ChatOpenAI\n",
    "from langchain.output_parsers.openai_functions import PydanticOutputFunctionsParser\n",
    "\n",
    "prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\"system\", \"You are an accounting assistant.\"),\n",
    "        (\"user\", \"{input}\"),\n",
    "    ]\n",
    ")\n",
    "chain = (\n",
    "    prompt\n",
    "    | ChatOpenAI().bind(functions=[openai_function_def])\n",
    "    | PydanticOutputFunctionsParser(pydantic_schema=Calculator)\n",
    "    | (lambda x: x.calculate())\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "62da7d8f-5cfc-45a6-946e-2bcda2b0ba1f",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Retrying langchain.chat_models.openai.ChatOpenAI.completion_with_retry.<locals>._completion_with_retry in 4.0 seconds as it raised ServiceUnavailableError: The server is overloaded or not ready yet..\n"
     ]
    }
   ],
   "source": [
    "math_questions = [\n",
    "    \"What's 45/9?\",\n",
    "    \"What's 81/9?\",\n",
    "    \"What's 72/8?\",\n",
    "    \"What's 56/7?\",\n",
    "    \"What's 36/6?\",\n",
    "    \"What's 64/8?\",\n",
    "    \"What's 12*6?\",\n",
    "    \"What's 8*8?\",\n",
    "    \"What's 10*10?\",\n",
    "    \"What's 11*11?\",\n",
    "    \"What's 13*13?\",\n",
    "    \"What's 45+30?\",\n",
    "    \"What's 72+28?\",\n",
    "    \"What's 56+44?\",\n",
    "    \"What's 63+37?\",\n",
    "    \"What's 70-35?\",\n",
    "    \"What's 60-30?\",\n",
    "    \"What's 50-25?\",\n",
    "    \"What's 40-20?\",\n",
    "    \"What's 30-15?\"\n",
    "]\n",
    "results = chain.batch([{\"input\": q} for q in math_questions], return_exceptions=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cbb1bcae-b922-4d38-b4bd-4b65be400b88",
   "metadata": {},
   "source": [
    "#### Load runs that did not error\n",
    "\n",
    "Now we can select the successful runs to fine-tune on."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "d6037992-050d-4ada-a061-860c124f0bf1",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langsmith.client import Client\n",
    "\n",
    "client = Client()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "0444919a-6f5a-4726-9916-4603b1420d0e",
   "metadata": {},
   "outputs": [],
   "source": [
    "successful_traces = {\n",
    "    run.trace_id\n",
    "    for run in client.list_runs(\n",
    "        project_name=project_name,\n",
    "        execution_order=1,\n",
    "        error=False,\n",
    "    )\n",
    "}\n",
    "    \n",
    "llm_runs = [\n",
    "    run for run in client.list_runs(\n",
    "        project_name=project_name,\n",
    "        run_type=\"llm\",\n",
    "    ) \n",
    "    if run.trace_id in successful_traces\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f365a359-52f7-47ff-8c36-aadc1070b409",
   "metadata": {},
   "source": [
    "## 2. Prepare data\n",
    "Now we can create an instance of LangSmithRunChatLoader and load the chat sessions using its lazy_load() method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "817bc077-c18a-473b-94a4-a7d810d583a8",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.chat_loaders.langsmith import LangSmithRunChatLoader\n",
    "\n",
    "loader = LangSmithRunChatLoader(runs=llm_runs)\n",
    "\n",
    "chat_sessions = loader.lazy_load()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f21a3bbd-1ed4-481b-9640-206b8bf0d751",
   "metadata": {},
   "source": [
    "#### With the chat sessions loaded, convert them into a format suitable for fine-tuning."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "9e5ac127-b094-4584-9159-5a6d3d7315c7",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.adapters.openai import convert_messages_for_finetuning\n",
    "\n",
    "training_data = convert_messages_for_finetuning(chat_sessions)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "188c4978-d85e-4984-a008-a50f6cd6bb84",
   "metadata": {},
   "source": [
    "## 3. Fine-tune the model\n",
    "Now, initiate the fine-tuning process using the OpenAI library."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "11d19e28-be49-4801-8065-1a58d13cd192",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Status=[running]... 346.26s. 31.70s\r"
     ]
    }
   ],
   "source": [
    "import openai\n",
    "import time\n",
    "import json\n",
    "from io import BytesIO\n",
    "\n",
    "my_file = BytesIO()\n",
    "for dialog in training_data:\n",
    "    my_file.write((json.dumps({\"messages\": dialog}) + \"\\n\").encode('utf-8'))\n",
    "\n",
    "my_file.seek(0)\n",
    "training_file = openai.File.create(\n",
    "  file=my_file,\n",
    "  purpose='fine-tune'\n",
    ")\n",
    "\n",
    "job = openai.FineTuningJob.create(\n",
    "    training_file=training_file.id,\n",
    "    model=\"gpt-3.5-turbo\",\n",
    ")\n",
    "\n",
    "# Wait for the fine-tuning to complete (this may take some time)\n",
    "status = openai.FineTuningJob.retrieve(job.id).status\n",
    "start_time = time.time()\n",
    "while status != \"succeeded\":\n",
    "    print(f\"Status=[{status}]... {time.time() - start_time:.2f}s\", end=\"\\r\", flush=True)\n",
    "    time.sleep(5)\n",
    "    status = openai.FineTuningJob.retrieve(job.id).status\n",
    "\n",
    "# Now your model is fine-tuned!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "54c4cead-500d-41dd-8333-2defde634396",
   "metadata": {},
   "source": [
    "## 4. Use in LangChain\n",
    "\n",
    "After fine-tuning, use the resulting model ID with the ChatOpenAI model class in your LangChain app."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "7f45b281-1dfa-43cb-bd28-99fa7e9f45d1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get the fine-tuned model ID\n",
    "job = openai.FineTuningJob.retrieve(job.id)\n",
    "model_id = job.fine_tuned_model\n",
    "\n",
    "# Use the fine-tuned model in LangChain\n",
    "model = ChatOpenAI(\n",
    "    model=model_id,\n",
    "    temperature=1,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "7d3b5845-6385-42d1-9f7d-5ea798dc2cd9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AIMessage(content='{\\n  \"num1\": 56,\\n  \"num2\": 7,\\n  \"operation\": \"/\"\\n}')"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "(prompt | model).invoke({\"input\": \"What's 56/7?\"})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5b8c2c79-ce27-4f37-b1b2-5977db8c4e84",
   "metadata": {},
   "source": [
    "Now you have successfully fine-tuned a model using data from LangSmith LLM runs!"
   ]
  }
 ],
 "metadata": {
  "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.11.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
