Add intermediate outputs to QueryPipeline (run-llama#12683)

Author: xpbowler

docs/docs/examples/pipeline/query_pipeline.ipynb

@@ -77,6 +77,18 @@
     "llama_index.core.set_global_handler(\"arize_phoenix\")"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "abc3f6b1",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "\n",
+    "os.environ[\"OPENAI_API_KEY\"] = \"sk-...\""
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -180,10 +192,10 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "\u001b[1;3;38;2;155;135;227m> Running module 43554ab5-002f-4705-b5f1-f3fb4fd9bd44 with input: \n",
+      "\u001b[1;3;38;2;155;135;227m> Running module 8dc57d24-9691-4d8d-87d7-151865a7cd1b with input: \n",
       "movie_name: The Departed\n",
       "\n",
-      "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running module 4b086131-f82b-4cb9-b929-427a60aca9c8 with input: \n",
+      "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running module 7ed9e26c-a704-4b0b-9cfd-991266e754c0 with input: \n",
       "messages: Please generate related movies to The Departed\n",
       "\n",
       "\u001b[0m"
@@ -204,30 +216,93 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "assistant: 1. Infernal Affairs (2002) - This is the original Hong Kong film that inspired The Departed. It follows a similar storyline of undercover cops infiltrating a criminal organization.\n",
-      "\n",
-      "2. Internal Affairs (1990) - This American crime thriller film shares a similar premise to The Departed, focusing on the corrupt practices within the police force.\n",
-      "\n",
-      "3. The Town (2010) - Directed by Ben Affleck, this crime drama revolves around a group of bank robbers in Boston and the FBI agent determined to bring them down.\n",
-      "\n",
-      "4. American Gangster (2007) - Based on a true story, this crime film explores the rise and fall of a drug lord in Harlem, New York, and the detective who is determined to bring him to justice.\n",
-      "\n",
-      "5. Donnie Brasco (1997) - Starring Johnny Depp and Al Pacino, this crime drama tells the story of an undercover FBI agent who infiltrates the mob and forms a close bond with a mafia hitman.\n",
-      "\n",
-      "6. The Departed (2006) - Although already mentioned, it is worth noting that The Departed itself is a highly acclaimed crime thriller directed by Martin Scorsese, exploring the intertwining lives of undercover cops and mobsters in Boston.\n",
-      "\n",
-      "7. Training Day (2001) - Denzel Washington won an Academy Award for his role in this crime thriller, where he plays a corrupt narcotics detective who takes a rookie cop under his wing for a day of training.\n",
-      "\n",
-      "8. Heat (1995) - Directed by Michael Mann, this crime film features an intense cat-and-mouse game between a skilled detective and a professional thief, both equally obsessed with their work.\n",
+      "assistant: 1. Infernal Affairs (2002) - The original Hong Kong film that inspired The Departed\n",
+      "2. The Town (2010) - A crime thriller directed by and starring Ben Affleck\n",
+      "3. Mystic River (2003) - A crime drama directed by Clint Eastwood\n",
+      "4. Goodfellas (1990) - A classic mobster film directed by Martin Scorsese\n",
+      "5. The Irishman (2019) - Another crime drama directed by Martin Scorsese, starring Robert De Niro and Al Pacino\n",
+      "6. The Departed (2006) - The Departed is a 2006 American crime film directed by Martin Scorsese and written by William Monahan. It is a remake of the 2002 Hong Kong film Infernal Affairs. The film stars Leonardo DiCaprio, Matt Damon, Jack Nicholson, and Mark Wahlberg, with Martin Sheen, Ray Winstone, Vera Farmiga, and Alec Baldwin in supporting roles.\n"
+     ]
+    }
+   ],
+   "source": [
+    "print(str(output))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "9e84db05",
+   "metadata": {},
+   "source": [
+    "### View Intermediate Inputs/Outputs\n",
+    "\n",
+    "For debugging and other purposes, we can also view the inputs and outputs at each step."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6083e322",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\u001b[1;3;38;2;155;135;227m> Running module 8dc57d24-9691-4d8d-87d7-151865a7cd1b with input: \n",
+      "movie_name: The Departed\n",
       "\n",
-      "9. The Godfather (1972) - Francis Ford Coppola's iconic crime saga follows the Corleone family's rise to power in the world of organized crime, with themes of loyalty, betrayal, and family ties.\n",
+      "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running module 7ed9e26c-a704-4b0b-9cfd-991266e754c0 with input: \n",
+      "messages: Please generate related movies to The Departed\n",
       "\n",
-      "10. Casino (1995) - Directed by Martin Scorsese, this film delves into the dark underbelly of Las Vegas casinos, exploring the corrupt and violent world of mob-controlled gambling.\n"
+      "\u001b[0m"
      ]
     }
    ],
    "source": [
-    "print(str(output))"
+    "output, intermediates = p.run_with_intermediates(movie_name=\"The Departed\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "1c0d5874",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "ComponentIntermediates(inputs={'movie_name': 'The Departed'}, outputs={'prompt': 'Please generate related movies to The Departed'})"
+      ]
+     },
+     "execution_count": null,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "intermediates[\"8dc57d24-9691-4d8d-87d7-151865a7cd1b\"]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "02cb9574",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "ComponentIntermediates(inputs={'messages': 'Please generate related movies to The Departed'}, outputs={'output': ChatResponse(message=ChatMessage(role=<MessageRole.ASSISTANT: 'assistant'>, content='1. Infernal Affairs (2002) - The original Hong Kong film that inspired The Departed\\n2. The Town (2010) - A crime thriller directed by Ben Affleck\\n3. Mystic River (2003) - A crime drama directed by Clint Eastwood\\n4. Goodfellas (1990) - A classic crime film directed by Martin Scorsese\\n5. The Irishman (2019) - Another crime film directed by Martin Scorsese, starring Robert De Niro and Al Pacino\\n6. The Godfather (1972) - A classic crime film directed by Francis Ford Coppola\\n7. Heat (1995) - A crime thriller directed by Michael Mann, starring Al Pacino and Robert De Niro\\n8. The Departed (2006) - A crime thriller directed by Martin Scorsese, starring Leonardo DiCaprio and Matt Damon.', additional_kwargs={}), raw={'id': 'chatcmpl-9EKf2nZ4latFJvHy0gzOUZbaB8xwY', 'choices': [Choice(finish_reason='stop', index=0, logprobs=None, message=ChatCompletionMessage(content='1. Infernal Affairs (2002) - The original Hong Kong film that inspired The Departed\\n2. The Town (2010) - A crime thriller directed by Ben Affleck\\n3. Mystic River (2003) - A crime drama directed by Clint Eastwood\\n4. Goodfellas (1990) - A classic crime film directed by Martin Scorsese\\n5. The Irishman (2019) - Another crime film directed by Martin Scorsese, starring Robert De Niro and Al Pacino\\n6. The Godfather (1972) - A classic crime film directed by Francis Ford Coppola\\n7. Heat (1995) - A crime thriller directed by Michael Mann, starring Al Pacino and Robert De Niro\\n8. The Departed (2006) - A crime thriller directed by Martin Scorsese, starring Leonardo DiCaprio and Matt Damon.', role='assistant', function_call=None, tool_calls=None))], 'created': 1713203040, 'model': 'gpt-3.5-turbo-0125', 'object': 'chat.completion', 'system_fingerprint': 'fp_c2295e73ad', 'usage': CompletionUsage(completion_tokens=184, prompt_tokens=15, total_tokens=199)}, delta=None, logprobs=None, additional_kwargs={})})"
+      ]
+     },
+     "execution_count": null,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "intermediates[\"7ed9e26c-a704-4b0b-9cfd-991266e754c0\"]"
    ]
   },
   {
@@ -1154,9 +1229,9 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "llama_index_v2",
+   "display_name": "venv",
    "language": "python",
-   "name": "llama_index_v2"
+   "name": "python3"
   },
   "language_info": {
    "codemirror_mode": {

docs/docs/module_guides/querying/pipeline/usage_pattern.md

@@ -113,6 +113,20 @@ p = QueryPipeline(chain=[prompt_tmpl, llm, summarizer_c])
 p.run(topic="YC")
 ```
 
+### Intermediate outputs
+
+If you wish to obtain the intermediate outputs of modules in QueryPipeline, you can use `run_with_intermediates` or `run_multi_with_intermediates` for single-input and multi-input, respectively.
+
+The output will be a tuple of the normal output and a dictionary containing module key -> `ComponentIntermediates`. ComponentIntermediates has 2 fields: `inputs` dict and `outputs` dict.
+
+```python
+output, intermediates = p.run_with_intermediates(topic="YC")
+print(output)
+print(intermediates)
+
+# output is (Response, {"module_key": ComponentIntermediates("inputs": {}, "outputs": {})})
+```
+
 ## Defining a Custom Query Component
 
 You can easily define a custom component: Either passing a function to a `FnComponent` or subclassing a `CustomQueryComponent`.

llama-index-core/llama_index/core/base/query_pipeline/query.py

@@ -333,6 +333,28 @@ def __init__(
         )
 
 
+class ComponentIntermediates:
+    """Component intermediate inputs and outputs."""
+
+    def __init__(
+        self,
+        inputs: Dict[str, Any],
+        outputs: Dict[str, Any],
+    ) -> None:
+        """Initialize."""
+        self.inputs = inputs
+        self.outputs = outputs
+
+    def __repr__(self) -> str:
+        return (
+            f"ComponentIntermediates(inputs={self.inputs!s}, "
+            f"outputs={self.outputs!s})"
+        )
+
+    def __str__(self) -> str:
+        return self.__repr__()
+
+
 # accept both QueryComponent and ChainableMixin as inputs to query pipeline
 # ChainableMixin modules will be converted to components via `as_query_component`
 QUERY_COMPONENT_TYPE = Union[QueryComponent, ChainableMixin]