langchain/docs/docs/integrations/providers/mlflow_tracking.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "5d184f91",
   "metadata": {},
   "source": [
    "# MLflow\n",
    "\n",
    ">[MLflow](https://www.mlflow.org/docs/latest/what-is-mlflow) is a versatile, expandable, open-source platform for managing workflows and artifacts across the machine learning lifecycle. It has built-in integrations with many popular ML libraries, but can be used with any library, algorithm, or deployment tool. It is designed to be extensible, so you can write plugins to support new workflows, libraries, and tools.\n",
    "\n",
    "This notebook goes over how to track your LangChain experiments into your `MLflow Server`"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ea73efae-7182-4a89-a492-c865b1fcf981",
   "metadata": {},
   "source": [
    "## External examples"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "97361a84-4e8f-45ba-b291-814cf73cd8f2",
   "metadata": {},
   "source": [
    "`MLflow` provides [several examples](https://github.com/mlflow/mlflow/tree/master/examples/langchain) for the `LangChain` integration:\n",
    "- [simple_chain](https://github.com/mlflow/mlflow/blob/master/examples/langchain/simple_chain.py)\n",
    "- [simple_agent](https://github.com/mlflow/mlflow/blob/master/examples/langchain/simple_agent.py)\n",
    "- [retriever_chain](https://github.com/mlflow/mlflow/blob/master/examples/langchain/retriever_chain.py)\n",
    "- [retrieval_qa_chain](https://github.com/mlflow/mlflow/blob/master/examples/langchain/retrieval_qa_chain.py)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e0cbd74b-1542-45a4-a72b-b2eedeffd2e0",
   "metadata": {},
   "source": [
    "## Example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ca7bd72f",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install --upgrade --quiet  azureml-mlflow\n",
    "%pip install --upgrade --quiet  pandas\n",
    "%pip install --upgrade --quiet  textstat\n",
    "%pip install --upgrade --quiet  spacy\n",
    "%pip install --upgrade --quiet  langchain-openai\n",
    "%pip install --upgrade --quiet  google-search-results\n",
    "!python -m spacy download en_core_web_sm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bf8e1f5c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ[\"MLFLOW_TRACKING_URI\"] = \"\"\n",
    "os.environ[\"OPENAI_API_KEY\"] = \"\"\n",
    "os.environ[\"SERPAPI_API_KEY\"] = \"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fd49fd45",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.callbacks import MlflowCallbackHandler\n",
    "from langchain_openai import OpenAI"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "578cac8c",
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"Main function.\n",
    "\n",
    "This function is used to try the callback handler.\n",
    "Scenarios:\n",
    "1. OpenAI LLM\n",
    "2. Chain with multiple SubChains on multiple generations\n",
    "3. Agent with Tools\n",
    "\"\"\"\n",
    "mlflow_callback = MlflowCallbackHandler()\n",
    "llm = OpenAI(\n",
    "    model_name=\"gpt-3.5-turbo\", temperature=0, callbacks=[mlflow_callback], verbose=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9b20acae",
   "metadata": {},
   "outputs": [],
   "source": [
    "# SCENARIO 1 - LLM\n",
    "llm_result = llm.generate([\"Tell me a joke\"])\n",
    "\n",
    "mlflow_callback.flush_tracker(llm)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8b872046",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.chains import LLMChain\n",
    "from langchain.prompts import PromptTemplate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1b2627ef",
   "metadata": {},
   "outputs": [],
   "source": [
    "# SCENARIO 2 - Chain\n",
    "template = \"\"\"You are a playwright. Given the title of play, it is your job to write a synopsis for that title.\n",
    "Title: {title}\n",
    "Playwright: This is a synopsis for the above play:\"\"\"\n",
    "prompt_template = PromptTemplate(input_variables=[\"title\"], template=template)\n",
    "synopsis_chain = LLMChain(llm=llm, prompt=prompt_template, callbacks=[mlflow_callback])\n",
    "\n",
    "test_prompts = [\n",
    "    {\n",
    "        \"title\": \"documentary about good video games that push the boundary of game design\"\n",
    "    },\n",
    "]\n",
    "synopsis_chain.apply(test_prompts)\n",
    "mlflow_callback.flush_tracker(synopsis_chain)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e002823a",
   "metadata": {
    "id": "_jN73xcPVEpI"
   },
   "outputs": [],
   "source": [
    "from langchain.agents import AgentType, initialize_agent, load_tools"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "655bd47e",
   "metadata": {
    "id": "Gpq4rk6VT9cu"
   },
   "outputs": [],
   "source": [
    "# SCENARIO 3 - Agent with Tools\n",
    "tools = load_tools([\"serpapi\", \"llm-math\"], llm=llm, callbacks=[mlflow_callback])\n",
    "agent = initialize_agent(\n",
    "    tools,\n",
    "    llm,\n",
    "    agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,\n",
    "    callbacks=[mlflow_callback],\n",
    "    verbose=True,\n",
    ")\n",
    "agent.run(\n",
    "    \"Who is Leo DiCaprio's girlfriend? What is her current age raised to the 0.43 power?\"\n",
    ")\n",
    "mlflow_callback.flush_tracker(agent, finish=True)"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "provenance": []
  },
  "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.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}