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

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "5d184f91",
   "metadata": {},
   "source": [
    "# MLflow\n",
    "\n",
    ">[MLflow](https://mlflow.org/) is a versatile, open-source platform for managing workflows and artifacts across the machine learning and generative AI lifecycle. It has built-in integrations with many popular AI and ML libraries, but can be used with any library, algorithm, or deployment tool.\n",
    "\n",
    "MLflow's [LangChain integration](https://mlflow.org/docs/latest/llms/langchain/autologging.html) provides the following capabilities:\n",
    "\n",
    "- **[Tracing](https://mlflow.org/docs/latest/llms/langchain/autologging.html)**: Visualize data flows through your LangChain components with one line of code (`mlflow.langchain.autolog()`)\n",
    "- **[Experiment Tracking](https://mlflow.org/docs/latest/llms/langchain/index.html#experiment-tracking)**: Log artifacts, code, and metrics from your LangChain runs\n",
    "- **[Model Management](https://mlflow.org/docs/latest/model-registry.html)**: Version and deploy LangChain applications with dependency tracking\n",
    "- **[Evaluation](https://mlflow.org/docs/latest/llms/langchain/index.html#mlflow-evaluate)**: Measure the performance of your LangChain applications\n",
    "\n",
    "**Note**: MLflow tracing is available in MLflow versions 2.14.0 and later.\n",
    "\n",
    "This short guide focuses on MLflow's tracing capability for LangChain and LangGraph applications. You'll see how to enable tracing with one line of code and view the execution flow of your applications. For information about MLflow's other capabilities and to explore additional tutorials, please refer to the [MLflow documentation for LangChain](https://mlflow.org/docs/latest/llms/langchain/index.html). If you're new to MLflow, check out the [Getting Started with MLflow](https://mlflow.org/docs/latest/getting-started/index.html) guide."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "da942c89",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "To get started with MLflow tracing for LangChain, install the MLflow Python package. We will also use the `langchain-openai` package.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "42406548",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install mlflow langchain-openai langgraph -qU"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8e626bb4",
   "metadata": {},
   "source": [
    "Next, set the MLflow tracking URI and OpenAI API key."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "7e87b21d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "# Set MLflow tracking URI if you have MLflow Tracking Server running\n",
    "os.environ[\"MLFLOW_TRACKING_URI\"] = \"\"\n",
    "os.environ[\"OPENAI_API_KEY\"] = \"\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7bc78387",
   "metadata": {},
   "source": [
    "## MLflow Tracing\n",
    "\n",
    "MLflow's tracing capability helps you visualize the execution flow of your LangChain applications. Here's how to enable it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "155d2a6f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import mlflow\n",
    "\n",
    "# Optional: Set an experiment to organize your traces\n",
    "mlflow.set_experiment(\"LangChain MLflow Integration\")\n",
    "\n",
    "# Enable tracing\n",
    "mlflow.langchain.autolog()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "105a77e5",
   "metadata": {},
   "source": [
    "## Example: Tracing a LangChain Application\n",
    "\n",
    "Here's a complete example showing MLflow tracing with LangChain:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "99cab27c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import mlflow\n",
    "from langchain_core.output_parsers import StrOutputParser\n",
    "from langchain_core.prompts import ChatPromptTemplate\n",
    "from langchain_openai import ChatOpenAI\n",
    "\n",
    "# Enable MLflow tracing\n",
    "mlflow.langchain.autolog()\n",
    "\n",
    "# Create a simple chain\n",
    "llm = ChatOpenAI(model_name=\"gpt-4o\")\n",
    "\n",
    "prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\n",
    "            \"system\",\n",
    "            \"You are a helpful assistant that translates {input_language} to {output_language}.\",\n",
    "        ),\n",
    "        (\"human\", \"{input}\"),\n",
    "    ]\n",
    ")\n",
    "\n",
    "chain = prompt | llm | StrOutputParser()\n",
    "\n",
    "# Run the chain\n",
    "result = chain.invoke(\n",
    "    {\n",
    "        \"input_language\": \"English\",\n",
    "        \"output_language\": \"German\",\n",
    "        \"input\": \"I love programming.\",\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6240cf2d",
   "metadata": {},
   "source": [
    "To view the trace, run `mlflow ui` in your terminal and navigate to the Traces tab in the MLflow UI.\n",
    "\n",
    "## Example: Tracing a LangGraph Application\n",
    "\n",
    "MLflow also supports tracing LangGraph applications:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "f703c505",
   "metadata": {},
   "outputs": [],
   "source": [
    "import mlflow\n",
    "from langchain_core.tools import tool\n",
    "from langgraph.prebuilt import create_react_agent\n",
    "\n",
    "# Enable MLflow tracing\n",
    "mlflow.langchain.autolog()\n",
    "\n",
    "\n",
    "# Define a tool\n",
    "@tool\n",
    "def count_words(text: str) -> str:\n",
    "    \"\"\"Counts the number of words in a text.\"\"\"\n",
    "    word_count = len(text.split())\n",
    "    return f\"This text contains {word_count} words.\"\n",
    "\n",
    "\n",
    "# Create a LangGraph agent\n",
    "llm = ChatOpenAI(model=\"gpt-4o\")\n",
    "tools = [count_words]\n",
    "graph = create_react_agent(llm, tools)\n",
    "\n",
    "# Run the agent\n",
    "result = graph.invoke(\n",
    "    {\"messages\": [{\"role\": \"user\", \"content\": \"Write me a 71-word story about a cat.\"}]}\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "19c2f3ff",
   "metadata": {},
   "source": [
    "To view the trace, run `mlflow ui` in your terminal and navigate to the Traces tab in the MLflow UI."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "48accc76",
   "metadata": {},
   "source": [
    "## Resources\n",
    "\n",
    "For more information on using MLflow with LangChain, please visit:\n",
    "\n",
    "- [MLflow LangChain Integration Documentation](https://mlflow.org/docs/latest/llms/langchain/index.html)\n",
    "- [MLflow Tracing Documentation](https://mlflow.org/docs/latest/llms/tracing/index.html)\n",
    "- [Logging LangChain and LangGraph Models](https://mlflow.org/docs/latest/llms/langchain/index.html#logging-models-from-code)\n",
    "- [Evaluating LangChain and LangGraph Models](https://mlflow.org/docs/latest/llms/langchain/index.html#how-can-i-evaluate-a-langgraph-agent)"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "provenance": []
  },
  "kernelspec": {
   "display_name": ".venv",
   "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.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}