{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "cf4fb76d-c534-485b-8b51-a0714ee3b82e",
   "metadata": {},
   "source": [
    "# Routing by semantic similarity\n",
    "\n",
    "With LCEL you can easily add [custom routing logic](/docs/expression_language/how_to/routing#using-a-custom-function) to your chain to dynamically determine the chain logic based on user input. All you need to do is define a function that given an input returns a `Runnable`.\n",
    "\n",
    "One especially useful technique is to use embeddings to route a query to the most relevant prompt. Here's a very simple example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "eef9020a-5f7c-4291-98eb-fa73f17d4b92",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.chat_models import ChatOpenAI\n",
    "from langchain.embeddings import OpenAIEmbeddings\n",
    "from langchain.prompts import PromptTemplate\n",
    "from langchain.schema.output_parser import StrOutputParser\n",
    "from langchain.schema.runnable import RunnableLambda, RunnablePassthrough\n",
    "from langchain.utils.math import cosine_similarity\n",
    "\n",
    "physics_template = \"\"\"You are a very smart physics professor. \\\n",
    "You are great at answering questions about physics in a concise and easy to understand manner. \\\n",
    "When you don't know the answer to a question you admit that you don't know.\n",
    "\n",
    "Here is a question:\n",
    "{query}\"\"\"\n",
    "\n",
    "math_template = \"\"\"You are a very good mathematician. You are great at answering math questions. \\\n",
    "You are so good because you are able to break down hard problems into their component parts, \\\n",
    "answer the component parts, and then put them together to answer the broader question.\n",
    "\n",
    "Here is a question:\n",
    "{query}\"\"\"\n",
    "\n",
    "embeddings = OpenAIEmbeddings()\n",
    "prompt_templates = [physics_template, math_template]\n",
    "prompt_embeddings = embeddings.embed_documents(prompt_templates)\n",
    "\n",
    "\n",
    "def prompt_router(input):\n",
    "    query_embedding = embeddings.embed_query(input[\"query\"])\n",
    "    similarity = cosine_similarity([query_embedding], prompt_embeddings)[0]\n",
    "    most_similar = prompt_templates[similarity.argmax()]\n",
    "    print(\"Using MATH\" if most_similar == math_template else \"Using PHYSICS\")\n",
    "    return PromptTemplate.from_template(most_similar)\n",
    "\n",
    "\n",
    "chain = (\n",
    "    {\"query\": RunnablePassthrough()}\n",
    "    | RunnableLambda(prompt_router)\n",
    "    | ChatOpenAI()\n",
    "    | StrOutputParser()\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "4d22b0f3-24f2-4a47-9440-065b57ebcdbd",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using PHYSICS\n",
      "A black hole is a region in space where gravity is extremely strong, so strong that nothing, not even light, can escape its gravitational pull. It is formed when a massive star collapses under its own gravity during a supernova explosion. The collapse causes an incredibly dense mass to be concentrated in a small volume, creating a gravitational field that is so intense that it warps space and time. Black holes have a boundary called the event horizon, which marks the point of no return for anything that gets too close. Beyond the event horizon, the gravitational pull is so strong that even light cannot escape, hence the name \"black hole.\" While we have a good understanding of black holes, there is still much to learn, especially about what happens inside them.\n"
     ]
    }
   ],
   "source": [
    "print(chain.invoke(\"What's a black hole\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "f261910d-1de1-4a01-8c8a-308db02b81de",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using MATH\n",
      "Thank you for your kind words! I will do my best to break down the concept of a path integral for you.\n",
      "\n",
      "In mathematics and physics, a path integral is a mathematical tool used to calculate the probability amplitude or wave function of a particle or system of particles. It was introduced by Richard Feynman and is an integral over all possible paths that a particle can take to go from an initial state to a final state.\n",
      "\n",
      "To understand the concept better, let's consider an example. Suppose we have a particle moving from point A to point B in space. Classically, we would describe this particle's motion using a definite trajectory, but in quantum mechanics, particles can simultaneously take multiple paths from A to B.\n",
      "\n",
      "The path integral formalism considers all possible paths that the particle could take and assigns a probability amplitude to each path. These probability amplitudes are then added up, taking into account the interference effects between different paths.\n",
      "\n",
      "To calculate a path integral, we need to define an action, which is a mathematical function that describes the behavior of the system. The action is usually expressed in terms of the particle's position, velocity, and time.\n",
      "\n",
      "Once we have the action, we can write down the path integral as an integral over all possible paths. Each path is weighted by a factor determined by the action and the principle of least action, which states that a particle takes a path that minimizes the action.\n",
      "\n",
      "Mathematically, the path integral is expressed as:\n",
      "\n",
      "∫ e^(iS/ħ) D[x(t)]\n",
      "\n",
      "Here, S is the action, ħ is the reduced Planck's constant, and D[x(t)] represents the integration over all possible paths x(t) of the particle.\n",
      "\n",
      "By evaluating this integral, we can obtain the probability amplitude for the particle to go from the initial state to the final state. The absolute square of this amplitude gives us the probability of finding the particle in a particular state.\n",
      "\n",
      "Path integrals have proven to be a powerful tool in various areas of physics, including quantum mechanics, quantum field theory, and statistical mechanics. They allow us to study complex systems and calculate probabilities that are difficult to obtain using other methods.\n",
      "\n",
      "I hope this explanation helps you understand the concept of a path integral. If you have any further questions, feel free to ask!\n"
     ]
    }
   ],
   "source": [
    "print(chain.invoke(\"What's a path integral\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f0c1732a-01ca-4d10-977c-29ed7480972b",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.9.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
