{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1cR_pjqz1AsF"
      },
      "source": [
        "# Building AI Applications by Customizing DSPy Modules\n",
        "\n",
        "In this guide, we will walk you through how to build a GenAI application by customizing `dspy.Module`.\n",
        "\n",
        "A [DSPy module](https://dspy.ai/learn/programming/modules/) is the building block for DSPy programs.\n",
        "\n",
        "- Each built-in module abstracts a prompting technique (like chain of thought or ReAct). Crucially, they are generalized to handle any signature.\n",
        "\n",
        "- A DSPy module has learnable parameters (i.e., the little pieces comprising the prompt and the LM weights) and can be invoked (called) to process inputs and return outputs.\n",
        "\n",
        "- Multiple modules can be composed into bigger modules (programs). DSPy modules are inspired directly by NN modules in PyTorch, but applied to LM programs.\n",
        "\n",
        "Although you can build a DSPy program without implementing a custom module, we highly recommend putting your logic with a custom module so that you can use other DSPy features, like DSPy optimizer or MLflow DSPy tracing."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KBYjBQtv3Cn5"
      },
      "source": [
        "Before getting started, make sure you have DSPy installed:\n",
        "\n",
        "```\n",
        "!pip install dspy\n",
        "```"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "<details>\n",
        "<summary>Recommended: Set up MLflow Tracing to understand what's happening under the hood.</summary>\n",
        "\n",
        "### MLflow DSPy Integration\n",
        "\n",
        "<a href=\"https://mlflow.org/\">MLflow</a> is an LLMOps tool that natively integrates with DSPy and offer explainability and experiment tracking. In this tutorial, you can use MLflow to visualize prompts and optimization progress as traces to understand the DSPy's behavior better. You can set up MLflow easily by following the four steps below.\n",
        "\n",
        "![MLflow Trace](./mlflow-tracing-custom-module.png)\n",
        "\n",
        "1. Install MLflow\n",
        "\n",
        "```bash\n",
        "%pip install mlflow>=3.0.0\n",
        "```\n",
        "\n",
        "2. Start MLflow UI in a separate terminal\n",
        "```bash\n",
        "mlflow ui --port 5000 --backend-store-uri sqlite:///mlruns.db\n",
        "```\n",
        "\n",
        "3. Connect the notebook to MLflow\n",
        "```python\n",
        "import mlflow\n",
        "\n",
        "mlflow.set_tracking_uri(\"http://localhost:5000\")\n",
        "mlflow.set_experiment(\"DSPy\")\n",
        "```\n",
        "\n",
        "4. Enabling tracing.\n",
        "```python\n",
        "mlflow.dspy.autolog()\n",
        "```\n",
        "\n",
        "\n",
        "To learn more about the integration, visit [MLflow DSPy Documentation](https://mlflow.org/docs/latest/llms/dspy/index.html) as well.\n",
        "</details>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "reQSTM8a8qMf"
      },
      "source": [
        "## Customize DSPy Module\n",
        "\n",
        "You can implement custom prompting logic and integrate external tools or services by customizing a DSPy module. To achieve this, subclass from `dspy.Module` and implement the following two key methods:\n",
        "\n",
        "- `__init__`: This is the constructor, where you define the attributes and sub-modules of your program.\n",
        "- `forward`: This method contains the core logic of your DSPy program.\n",
        "\n",
        "Within the `forward()` method, you are not limited to calling only other DSPy modules; you can also integrate any standard Python functions, such as those for interacting with Langchain/Agno agents, MCP tools, database handlers, and more.\n",
        "\n",
        "The basic structure for a custom DSPy module looks like this:\n",
        "\n",
        "```python\n",
        "class MyProgram(dspy.Module):\n",
        "    \n",
        "    def __init__(self, ...):\n",
        "        # Define attributes and sub-modules here\n",
        "        {constructor_code}\n",
        "\n",
        "    def forward(self, input_name1, input_name2, ...):\n",
        "        # Implement your program's logic here\n",
        "        {custom_logic_code}\n",
        "```"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DziTWwT8_TrY"
      },
      "source": [
        "Let's illustrate this with a practical code example. We will build a simple Retrieval-Augmented Generation (RAG) application with multiple stages:\n",
        "\n",
        "1.  **Query Generation:** Generate a suitable query based on the user's question to retrieve relevant context.\n",
        "2.  **Context Retrieval:** Fetch context using the generated query.\n",
        "3.  **Answer Generation:** Produce a final answer based on the retrieved context and the original question.\n",
        "\n",
        "The code implementation for this multi-stage program is shown below."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "id": "lAoV5_v7YlvN"
      },
      "outputs": [],
      "source": [
        "import dspy\n",
        "\n",
        "class QueryGenerator(dspy.Signature):\n",
        "    \"\"\"Generate a query based on question to fetch relevant context\"\"\"\n",
        "    question: str = dspy.InputField()\n",
        "    query: str = dspy.OutputField()\n",
        "\n",
        "def search_wikipedia(query: str) -> list[str]:\n",
        "    \"\"\"Query ColBERT endpoint, which is a knowledge source based on wikipedia data\"\"\"\n",
        "    results = dspy.ColBERTv2(url='http://20.102.90.50:2017/wiki17_abstracts')(query, k=1)\n",
        "    return [x[\"text\"] for x in results]\n",
        "\n",
        "class RAG(dspy.Module):\n",
        "    def __init__(self):\n",
        "        self.query_generator = dspy.Predict(QueryGenerator)\n",
        "        self.answer_generator = dspy.ChainOfThought(\"question,context->answer\")\n",
        "\n",
        "    def forward(self, question, **kwargs):\n",
        "        query = self.query_generator(question=question).query\n",
        "        context = search_wikipedia(query)[0]\n",
        "        return self.answer_generator(question=question, context=context).answer"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Let's take a look at the `forward` method. We first send the question to `self.query_generator`, which is a `dspy.Predict`, to get the query for context retrieving. Then we use the query to call ColBERT and keep the first context retrieved. Finally, we send the question and context into `self.answer_generator`, which is a `dspy.ChainOfThought` to generate the final answer."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FBq_4e8NamwY"
      },
      "source": [
        "Next, we'll create an instance of our `RAG` module to run the program.\n",
        "\n",
        "**Important:** When invoking a custom DSPy module, you should use the module instance directly (which calls the `__call__` method internally), rather than calling the `forward()` method explicitly. The `__call__` method handles necessary internal processing before executing the `forward` logic."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ZR7xcFSTa596",
        "outputId": "f3427754-8a16-48fe-c540-8c9f31d9a30d"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "The question of whether LeBron James is the basketball GOAT is subjective and depends on personal opinions. Many consider him one of the greatest due to his achievements and impact on the game, but others may argue for different players like Michael Jordan.\n"
          ]
        }
      ],
      "source": [
        "import os\n",
        "\n",
        "os.environ[\"OPENAI_API_KEY\"] = \"{your_openai_api_key}\"\n",
        "\n",
        "dspy.configure(lm=dspy.LM(\"openai/gpt-4o-mini\"))\n",
        "rag = RAG()\n",
        "print(rag(question=\"Is Lebron James the basketball GOAT?\"))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "That's it! In summary, to build your GenAI applications, we just put the custom logic into the `forward()` method, then create a module instance and call the instance itself."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aYAYc-Hg39ri"
      },
      "source": [
        "## Why Customizing Module?\n",
        "\n",
        "DSPy is a lightweight authoring and optimization framework, and our focus is to resolve the mess of prompt engineering by transforming prompting (string in, string out) LLM into programming LLM (structured inputs in, structured outputs out) for robust AI system.\n",
        "\n",
        "While we provide pre-built modules which have custom prompting logic like `dspy.ChainOfThought` for reasoning, `dspy.ReAct` for tool calling agent to facilitate building your AI applications, we don't aim at standardizing how you build agents.\n",
        "\n",
        "In DSPy, your application logic simply goes to the `forward` method of your custom Module, which doesn't have any constraint as long as you are writing python code. With this layout, DSPy is easy to migrate to from other frameworks or vanilla SDK usage, and easy to migrate off because essentially it's just python code.\n"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
