From 06ed8a3534b4e899a19366ac7adfec77cde4a19a Mon Sep 17 00:00:00 2001
From: Jerry Liu <jerryjliu98@gmail.com>
Date: Mon, 8 Jan 2024 08:31:34 -0800
Subject: [PATCH] add query pipeline (#9908)
---
docs/examples/pipeline/query_pipeline.ipynb | 955 ++++++++++++++++++
.../querying/pipeline/module_usage.md | 71 ++
.../querying/pipeline/modules.md | 8 +
docs/module_guides/querying/pipeline/root.md | 58 ++
.../querying/pipeline/usage_pattern.md | 178 ++++
docs/module_guides/querying/querying.md | 16 +-
llama_index/core/base_query_engine.py | 53 +-
llama_index/core/base_retriever.py | 57 +-
llama_index/core/query_pipeline/__init__.py | 0
.../core/query_pipeline/query_component.py | 281 ++++++
.../indices/query/query_transform/base.py | 64 +-
llama_index/llms/ai21.py | 8 +-
llama_index/llms/anthropic.py | 14 +-
llama_index/llms/base.py | 19 +-
llama_index/llms/bedrock.py | 22 +-
llama_index/llms/clarifai.py | 16 +-
llama_index/llms/cohere.py | 14 +-
llama_index/llms/custom.py | 24 +-
llama_index/llms/gemini.py | 8 +-
llama_index/llms/gradient.py | 9 +-
llama_index/llms/huggingface.py | 28 +-
llama_index/llms/konko.py | 14 +-
llama_index/llms/langchain.py | 24 +-
llama_index/llms/litellm.py | 14 +-
llama_index/llms/llama_api.py | 8 +-
llama_index/llms/llama_cpp.py | 14 +-
llama_index/llms/llm.py | 148 ++-
llama_index/llms/mistral.py | 14 +-
llama_index/llms/mock.py | 8 +-
llama_index/llms/monsterapi.py | 11 +-
llama_index/llms/nvidia_triton.py | 14 +-
llama_index/llms/ollama.py | 8 +-
llama_index/llms/openai.py | 14 +-
llama_index/llms/openllm.py | 28 +-
llama_index/llms/palm.py | 8 +-
llama_index/llms/perplexity.py | 14 +-
llama_index/llms/portkey.py | 8 +-
llama_index/llms/predibase.py | 8 +-
llama_index/llms/replicate.py | 12 +-
llama_index/llms/rungpt.py | 14 +-
llama_index/llms/vertex.py | 14 +-
llama_index/llms/vllm.py | 28 +-
llama_index/llms/watsonx.py | 14 +-
llama_index/llms/xinference.py | 8 +-
llama_index/output_parsers/base.py | 59 +-
llama_index/output_parsers/guardrails.py | 5 +-
llama_index/output_parsers/langchain.py | 5 +-
llama_index/output_parsers/pydantic.py | 18 +-
llama_index/postprocessor/types.py | 68 +-
llama_index/prompts/base.py | 84 +-
llama_index/query_pipeline/__init__.py | 17 +
llama_index/query_pipeline/query.py | 474 +++++++++
llama_index/response_synthesizers/base.py | 76 +-
poetry.lock | 363 ++++++-
pyproject.toml | 1 +
tests/llms/test_custom.py | 8 +-
tests/query_pipeline/__init__.py | 0
tests/query_pipeline/query.py | 266 +++++
58 files changed, 3584 insertions(+), 210 deletions(-)
create mode 100644 docs/examples/pipeline/query_pipeline.ipynb
create mode 100644 docs/module_guides/querying/pipeline/module_usage.md
create mode 100644 docs/module_guides/querying/pipeline/modules.md
create mode 100644 docs/module_guides/querying/pipeline/root.md
create mode 100644 docs/module_guides/querying/pipeline/usage_pattern.md
create mode 100644 llama_index/core/query_pipeline/__init__.py
create mode 100644 llama_index/core/query_pipeline/query_component.py
create mode 100644 llama_index/query_pipeline/__init__.py
create mode 100644 llama_index/query_pipeline/query.py
create mode 100644 tests/query_pipeline/__init__.py
create mode 100644 tests/query_pipeline/query.py
diff --git a/docs/examples/pipeline/query_pipeline.ipynb b/docs/examples/pipeline/query_pipeline.ipynb
new file mode 100644
index 0000000000..2191d34ab9
--- /dev/null
+++ b/docs/examples/pipeline/query_pipeline.ipynb
@@ -0,0 +1,955 @@
+{
+ "cells": [
+ {
+ "cell_type": "markdown",
+ "id": "8ebf73ae-a632-4d09-b941-6739e35b760d",
+ "metadata": {},
+ "source": [
+ "# An Introduction to LlamaIndex Query Pipelines\n",
+ "\n",
+ "## Overview\n",
+ "LlamaIndex provides a declarative query API that allows you to chain together different modules in order to orchestrate simple-to-advanced workflows over your data.\n",
+ "\n",
+ "This is centered around our `QueryPipeline` abstraction. Load in a variety of modules (from LLMs to prompts to retrievers to other pipelines), connect them all together into a sequential chain or DAG, and run it end2end.\n",
+ "\n",
+ "**NOTE**: You can orchestrate all these workflows without the declarative pipeline abstraction (by using the modules imperatively and writing your own functions). So what are the advantages of `QueryPipeline`? \n",
+ "\n",
+ "- Express common workflows with fewer lines of code/boilerplate\n",
+ "- Greater readability\n",
+ "- Greater parity / better integration points with common low-code / no-code solutions (e.g. LangFlow)\n",
+ "- [In the future] A declarative interface allows easy serializability of pipeline components, providing portability of pipelines/easier deployment to different systems.\n",
+ "\n",
+ "## Cookbook\n",
+ "\n",
+ "In this cookbook we give you an introduction to our `QueryPipeline` interface and show you some basic workflows you can tackle.\n",
+ "\n",
+ "- Chain together prompt and LLM\n",
+ "- Chain together query rewriting (prompt + LLM) with retrieval\n",
+ "- Chain together a full RAG query pipeline (query rewriting, retrieval, reranking, response synthesis)\n",
+ "- Setting up a custom query component"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "1144b2d4-adbc-44da-8c12-bdb5fe4b18bb",
+ "metadata": {},
+ "source": [
+ "## Setup\n",
+ "\n",
+ "Here we setup some data + indexes (from PG's essay) that we'll be using in the rest of the cookbook."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "adc82744-965a-4d79-b357-faf3de7ba2f8",
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "🌍 To view the Phoenix app in your browser, visit http://127.0.0.1:6006/\n",
+ "📺 To view the Phoenix app in a notebook, run `px.active_session().view()`\n",
+ "📖 For more information on how to use Phoenix, check out https://docs.arize.com/phoenix\n"
+ ]
+ }
+ ],
+ "source": [
+ "# setup Arize Phoenix for logging/observability\n",
+ "import phoenix as px\n",
+ "\n",
+ "px.launch_app()\n",
+ "import llama_index\n",
+ "\n",
+ "llama_index.set_global_handler(\"arize_phoenix\")"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "40fcb621-0894-457e-b602-dbf5fb9134ec",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "from llama_index.query_pipeline.query import QueryPipeline\n",
+ "from llama_index.llms import OpenAI\n",
+ "from llama_index.prompts import PromptTemplate\n",
+ "from llama_index import (\n",
+ " VectorStoreIndex,\n",
+ " ServiceContext,\n",
+ " SimpleDirectoryReader,\n",
+ " load_index_from_storage,\n",
+ ")"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "2009af96-59e3-4d14-8272-382203c8b8a7",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "reader = SimpleDirectoryReader(\"../data/paul_graham\")"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "c55d390b-38ca-4176-8cdb-8c2a0af1add8",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "docs = reader.load_data()"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "780ddc81-0783-4fa5-ade0-60700c918011",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "import os\n",
+ "from llama_index.storage import StorageContext\n",
+ "\n",
+ "if not os.path.exists(\"storage\"):\n",
+ " index = VectorStoreIndex.from_documents(docs)\n",
+ " # save index to disk\n",
+ " index.set_index_id(\"vector_index\")\n",
+ " index.storage_context.persist(\"./storage\")\n",
+ "else:\n",
+ " # rebuild storage context\n",
+ " storage_context = StorageContext.from_defaults(persist_dir=\"storage\")\n",
+ " # load index\n",
+ " index = load_index_from_storage(storage_context, index_id=\"vector_index\")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "d5c59b5c-9b18-4dfa-97ef-e39b8069b73c",
+ "metadata": {},
+ "source": [
+ "## 1. Chain Together Prompt and LLM \n",
+ "\n",
+ "In this section we show a super simple workflow of chaining together a prompt with LLM.\n",
+ "\n",
+ "We simply define `chain` on initialization. This is a special case of a query pipeline where the components are purely sequential, and we automatically convert outputs into the right format for the next inputs."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "cb233a0f-2993-4780-a241-6a2299047598",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# try chaining basic prompts\n",
+ "prompt_str = \"Please generate related movies to {movie_name}\"\n",
+ "prompt_tmpl = PromptTemplate(prompt_str)\n",
+ "llm = OpenAI(model=\"gpt-3.5-turbo\")\n",
+ "\n",
+ "p = QueryPipeline(chain=[prompt_tmpl, llm], verbose=True)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "1b26c0c6-b886-42a6-b524-b19b18d1c01c",
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "\u001b[1;3;38;2;155;135;227m> Running module 14b08322-7758-4089-b602-9559a367daff with input: \n",
+ "movie_name: The Departed\n",
+ "\n",
+ "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running module 0b5b116b-f971-48dd-aaa0-a2234ce3132c with input: \n",
+ "messages: Please generate related movies to The Departed\n",
+ "\n",
+ "\u001b[0m"
+ ]
+ }
+ ],
+ "source": [
+ "output = p.run(movie_name=\"The Departed\")"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "12d05cba-8e06-4c05-9bc7-38535814c066",
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "assistant: 1. Infernal Affairs (2002) - The Departed is actually a remake of this Hong Kong crime thriller, which follows a similar storyline of undercover cops infiltrating a criminal organization.\n",
+ "\n",
+ "2. 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",
+ "3. American Gangster (2007) - Based on a true story, this crime film explores the rise and fall of a Harlem drug lord and the detective who becomes obsessed with bringing him to justice.\n",
+ "\n",
+ "4. Heat (1995) - Directed by Michael Mann, this crime thriller features an intense cat-and-mouse game between a skilled detective and a professional thief, both equally obsessed with their respective careers.\n",
+ "\n",
+ "5. Training Day (2001) - Denzel Washington won an Academy Award for his role as a corrupt narcotics detective who takes a rookie cop (Ethan Hawke) on a dangerous journey through the mean streets of Los Angeles.\n",
+ "\n",
+ "6. The Departed (2006) - Although it's the movie you mentioned, it's worth noting that The Departed itself is a highly acclaimed crime film directed by Martin Scorsese, exploring the intertwining lives of an undercover cop and a mole in the police force.\n",
+ "\n",
+ "7. Donnie Brasco (1997) - Based on a true story, this crime drama follows an FBI agent who infiltrates the mob and forms a close bond with a low-level gangster, blurring the lines between loyalty and betrayal.\n",
+ "\n",
+ "8. The Untouchables (1987) - Set during the Prohibition era, this crime film depicts the efforts of a group of law enforcement officers, led by Eliot Ness (Kevin Costner), to bring down the notorious gangster Al Capone (Robert De Niro).\n",
+ "\n",
+ "9. The Godfather (1972) - Francis Ford Coppola's iconic crime saga follows the Corleone family, an organized crime dynasty, as they navigate power struggles, loyalty, and betrayal in post-World War II America.\n",
+ "\n",
+ "10. Eastern Promises (2007) - Directed by David Cronenberg, this crime thriller delves into the Russian mafia in London, as a midwife becomes entangled in a dangerous web of secrets and violence.\n"
+ ]
+ }
+ ],
+ "source": [
+ "print(str(output))"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "4f44bac2-fff8-4578-8179-2cf68d075429",
+ "metadata": {},
+ "source": [
+ "### Try Output Parsing\n",
+ "\n",
+ "Let's parse the outputs into a structured Pydantic object."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "69fc9aa6-d74e-4d02-8101-83ee03d68d52",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "from typing import List\n",
+ "from pydantic import BaseModel, Field\n",
+ "from llama_index.output_parsers import PydanticOutputParser\n",
+ "\n",
+ "\n",
+ "class Movie(BaseModel):\n",
+ " \"\"\"Object representing a single movie.\"\"\"\n",
+ "\n",
+ " name: str = Field(..., description=\"Name of the movie.\")\n",
+ " year: int = Field(..., description=\"Year of the movie.\")\n",
+ "\n",
+ "\n",
+ "class Movies(BaseModel):\n",
+ " \"\"\"Object representing a list of movies.\"\"\"\n",
+ "\n",
+ " movies: List[Movie] = Field(..., description=\"List of movies.\")\n",
+ "\n",
+ "\n",
+ "llm = OpenAI(model=\"gpt-3.5-turbo\")\n",
+ "output_parser = PydanticOutputParser(Movies)\n",
+ "prompt_str = \"\"\"\\\n",
+ "Please generate related movies to {movie_name}. Output with the following JSON format: \n",
+ "\"\"\"\n",
+ "prompt_str = output_parser.format(prompt_str)\n",
+ "# prompt_str = prompt_str + output_parser.format_string"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "546cfd0a-78f8-4494-8088-c9e4c66bdeef",
+ "metadata": {},
+ "outputs": [
+ {
+ "data": {
+ "text/plain": [
+ "'Please generate related movies to test. Output with the following JSON format: \\n\\n\\n\\nHere\\'s a JSON schema to follow:\\n{\"title\": \"Movies\", \"description\": \"Object representing a list of movies.\", \"type\": \"object\", \"properties\": {\"movies\": {\"title\": \"Movies\", \"description\": \"List of movies.\", \"type\": \"array\", \"items\": {\"$ref\": \"#/definitions/Movie\"}}}, \"required\": [\"movies\"], \"definitions\": {\"Movie\": {\"title\": \"Movie\", \"description\": \"Object representing a single movie.\", \"type\": \"object\", \"properties\": {\"name\": {\"title\": \"Name\", \"description\": \"Name of the movie.\", \"type\": \"string\"}, \"year\": {\"title\": \"Year\", \"description\": \"Year of the movie.\", \"type\": \"integer\"}}, \"required\": [\"name\", \"year\"]}}}\\n\\nOutput a valid JSON object but do not repeat the schema.\\n'"
+ ]
+ },
+ "execution_count": null,
+ "metadata": {},
+ "output_type": "execute_result"
+ }
+ ],
+ "source": [
+ "prompt_str.format(movie_name=\"test\")"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "243b55ca-a7bd-43e4-a837-20e29b3bebed",
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "\u001b[1;3;38;2;155;135;227m> Running module cb73b4a9-33d9-4869-8f8c-ce174963bfd8 with input: \n",
+ "movie_name: Toy Story\n",
+ "\n",
+ "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running module 7e3af4a9-0098-4f4e-a03a-1412254ba22f with input: \n",
+ "messages: Please generate related movies to Toy Story. Output with the following JSON format: \n",
+ "\n",
+ "\n",
+ "\n",
+ "Here's a JSON schema to follow:\n",
+ "{\"title\": \"Movies\", \"description\": \"Object representing a list of movies.\", \"typ...\n",
+ "\n",
+ "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running module afe13a71-d06f-4d41-9eb2-11ed2c439eab with input: \n",
+ "input: assistant: {\n",
+ " \"movies\": [\n",
+ " {\n",
+ " \"name\": \"Finding Nemo\",\n",
+ " \"year\": 2003\n",
+ " },\n",
+ " {\n",
+ " \"name\": \"Cars\",\n",
+ " \"year\": 2006\n",
+ " },\n",
+ " {\n",
+ " \"name\": \"Up\",\n",
+ " \"year\": 2009\n",
+ " },\n",
+ " {...\n",
+ "\n",
+ "\u001b[0m"
+ ]
+ }
+ ],
+ "source": [
+ "# add JSON spec to prompt template\n",
+ "prompt_tmpl = PromptTemplate(prompt_str)\n",
+ "\n",
+ "p = QueryPipeline(chain=[prompt_tmpl, llm, output_parser], verbose=True)\n",
+ "output = p.run(movie_name=\"Toy Story\")"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "4b15bef7-399b-4dcc-94d6-6a4ea0066a41",
+ "metadata": {},
+ "outputs": [
+ {
+ "data": {
+ "text/plain": [
+ "Movies(movies=[Movie(name='Finding Nemo', year=2003), Movie(name='Cars', year=2006), Movie(name='Up', year=2009), Movie(name='Inside Out', year=2015), Movie(name='Coco', year=2017)])"
+ ]
+ },
+ "execution_count": null,
+ "metadata": {},
+ "output_type": "execute_result"
+ }
+ ],
+ "source": [
+ "output"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "4572cdd4-2c94-4871-a496-623e9779e0db",
+ "metadata": {},
+ "source": [
+ "## Chain Together Query Rewriting Workflow (prompts + LLM) with Retrieval\n",
+ "\n",
+ "Here we try a slightly more complex workflow where we send the input through two prompts before initiating retrieval.\n",
+ "\n",
+ "1. Generate question about given topic.\n",
+ "2. Hallucinate answer given question, for better retrieval.\n",
+ "\n",
+ "Since each prompt only takes in one input, note that the `QueryPipeline` will automatically chain LLM outputs into the prompt and then into the LLM. \n",
+ "\n",
+ "You'll see how to define links more explicitly in the next section."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "3c6908b1-5819-4f34-a06c-2f8b9fc81c3e",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "from llama_index.postprocessor import CohereRerank\n",
+ "\n",
+ "# generate question regarding topic\n",
+ "prompt_str1 = \"Please generate a concise question about Paul Graham's life regarding the following topic {topic}\"\n",
+ "prompt_tmpl1 = PromptTemplate(prompt_str1)\n",
+ "# use HyDE to hallucinate answer.\n",
+ "prompt_str2 = (\n",
+ " \"Please write a passage to answer the question\\n\"\n",
+ " \"Try to include as many key details as possible.\\n\"\n",
+ " \"\\n\"\n",
+ " \"\\n\"\n",
+ " \"{query_str}\\n\"\n",
+ " \"\\n\"\n",
+ " \"\\n\"\n",
+ " 'Passage:\"\"\"\\n'\n",
+ ")\n",
+ "prompt_tmpl2 = PromptTemplate(prompt_str2)\n",
+ "\n",
+ "llm = OpenAI(model=\"gpt-3.5-turbo\")\n",
+ "retriever = index.as_retriever(similarity_top_k=5)\n",
+ "p = QueryPipeline(\n",
+ " chain=[prompt_tmpl1, llm, prompt_tmpl2, llm, retriever], verbose=True\n",
+ ")"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "03eb2534-a69a-46fd-b539-84ae93e2e5bb",
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "\u001b[1;3;38;2;155;135;227m> Running module 2790a58f-7633-4f4e-a5d5-71a252e79966 with input: \n",
+ "topic: college\n",
+ "\n",
+ "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running module d0762ea6-ce27-43be-b665-35ff8395599e with input: \n",
+ "messages: Please generate a concise question about Paul Graham's life regarding the following topic college\n",
+ "\n",
+ "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running module 4dc1a550-3097-4bdb-a913-b0c25abb91b6 with input: \n",
+ "query_str: assistant: How did Paul Graham's college experience shape his career and entrepreneurial mindset?\n",
+ "\n",
+ "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running module 54b08480-f316-41be-a42f-fb0563e78f5d with input: \n",
+ "messages: Please write a passage to answer the question\n",
+ "Try to include as many key details as possible.\n",
+ "\n",
+ "\n",
+ "assistant: How did Paul Graham's college experience shape his career and entrepreneurial mindset?\n",
+ "\n",
+ "\n",
+ "Pass...\n",
+ "\n",
+ "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running module 211a5c3a-a545-427b-baff-a59f72fd6142 with input: \n",
+ "input: assistant: Paul Graham's college experience played a pivotal role in shaping his career and entrepreneurial mindset. As a student at Cornell University, Graham immersed himself in the world of compute...\n",
+ "\n",
+ "\u001b[0m"
+ ]
+ },
+ {
+ "data": {
+ "text/plain": [
+ "5"
+ ]
+ },
+ "execution_count": null,
+ "metadata": {},
+ "output_type": "execute_result"
+ }
+ ],
+ "source": [
+ "nodes = p.run(topic=\"college\")\n",
+ "len(nodes)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "2c073634-55d4-4066-977f-fbc189089b95",
+ "metadata": {},
+ "source": [
+ "## Create a Full RAG Pipeline as a DAG\n",
+ "\n",
+ "Here we chain together a full RAG pipeline consisting of query rewriting, retrieval, reranking, and response synthesis.\n",
+ "\n",
+ "Here we can't use `chain` syntax because certain modules depend on multiple inputs (for instance, response synthesis expects both the retrieved nodes and the original question). Instead we'll construct a DAG explicitly, through `add_modules` and then `add_link`."
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "66625e37-a0f2-45f4-afaa-e93948444e97",
+ "metadata": {},
+ "source": [
+ "### 1. RAG Pipeline with Query Rewriting\n",
+ "\n",
+ "We use an LLM to rewrite the query first before passing it to our downstream modules - retrieval/reranking/synthesis."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "531677b6-0e6a-4002-9fdd-2096f1a83685",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "from llama_index.postprocessor import CohereRerank\n",
+ "from llama_index.response_synthesizers import TreeSummarize\n",
+ "from llama_index import ServiceContext\n",
+ "\n",
+ "# define modules\n",
+ "prompt_str = \"Please generate a question about Paul Graham's life regarding the following topic {topic}\"\n",
+ "prompt_tmpl = PromptTemplate(prompt_str)\n",
+ "llm = OpenAI(model=\"gpt-3.5-turbo\")\n",
+ "retriever = index.as_retriever(similarity_top_k=3)\n",
+ "reranker = CohereRerank()\n",
+ "summarizer = TreeSummarize(\n",
+ " service_context=ServiceContext.from_defaults(llm=llm)\n",
+ ")"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "4ea2bd1e-8c32-420b-b557-eb7dbe81718f",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# define query pipeline\n",
+ "p = QueryPipeline(verbose=True)\n",
+ "p.add_modules(\n",
+ " {\n",
+ " \"llm\": llm,\n",
+ " \"prompt_tmpl\": prompt_tmpl,\n",
+ " \"retriever\": retriever,\n",
+ " \"summarizer\": summarizer,\n",
+ " \"reranker\": reranker,\n",
+ " }\n",
+ ")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "7e4e2af5-f7c2-4ab8-bb7c-76c369b63650",
+ "metadata": {},
+ "source": [
+ "Next we draw links between modules with `add_link`. `add_link` takes in the source/destination module ids, and optionally the `source_key` and `dest_key`. Specify the `source_key` or `dest_key` if there are multiple outputs/inputs respectively.\n",
+ "\n",
+ "You can view the set of input/output keys for each module through `module.as_query_component().input_keys` and `module.as_query_component().output_keys`. \n",
+ "\n",
+ "Here we explicitly specify `dest_key` for the `reranker` and `summarizer` modules because they take in two inputs (query_str and nodes). "
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "0f0cccaf-8887-48f7-96b6-1d5458987022",
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "required_keys={'nodes', 'query_str'} optional_keys=set()\n"
+ ]
+ }
+ ],
+ "source": [
+ "p.add_link(\"prompt_tmpl\", \"llm\")\n",
+ "p.add_link(\"llm\", \"retriever\")\n",
+ "p.add_link(\"retriever\", \"reranker\", dest_key=\"nodes\")\n",
+ "p.add_link(\"llm\", \"reranker\", dest_key=\"query_str\")\n",
+ "p.add_link(\"reranker\", \"summarizer\", dest_key=\"nodes\")\n",
+ "p.add_link(\"llm\", \"summarizer\", dest_key=\"query_str\")\n",
+ "\n",
+ "# look at summarizer input keys\n",
+ "print(summarizer.as_query_component().input_keys)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "1cf69738-b409-408a-ac01-62bd4c2c2db4",
+ "metadata": {},
+ "source": [
+ "We use `networkx` to store the graph representation. This gives us an easy way to view the DAG! "
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "4ef45868-ce07-4190-b12e-a70b287a491f",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "## create graph\n",
+ "from pyvis.network import Network\n",
+ "\n",
+ "net = Network(notebook=True, cdn_resources=\"in_line\", directed=True)\n",
+ "net.from_nx(p.dag)\n",
+ "net.show(\"rag_dag.html\")\n",
+ "\n",
+ "## another option using `pygraphviz`\n",
+ "# from networkx.drawing.nx_agraph import to_agraph\n",
+ "# from IPython.display import Image\n",
+ "# agraph = to_agraph(p.dag)\n",
+ "# agraph.layout(prog=\"dot\")\n",
+ "# agraph.draw('rag_dag.png')\n",
+ "# display(Image('rag_dag.png'))"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "574ea13f-88e2-4393-8bd3-97fe436ffc0c",
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "\u001b[1;3;38;2;155;135;227m> Running module prompt_tmpl with input: \n",
+ "topic: YC\n",
+ "\n",
+ "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running module llm with input: \n",
+ "messages: Please generate a question about Paul Graham's life regarding the following topic YC\n",
+ "\n",
+ "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running module retriever with input: \n",
+ "input: assistant: What role did Paul Graham play in the founding and development of Y Combinator (YC)?\n",
+ "\n",
+ "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running module reranker with input: \n",
+ "query_str: assistant: What role did Paul Graham play in the founding and development of Y Combinator (YC)?\n",
+ "nodes: [NodeWithScore(node=TextNode(id_='bf639f66-34a2-46d2-a30d-9aff96d130b9', embedding=None, metadata={'file_path': '../data/paul_graham/paul_graham_essay.txt', 'file_name': 'paul_graham_essay.txt', 'file...\n",
+ "\n",
+ "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running module summarizer with input: \n",
+ "query_str: assistant: What role did Paul Graham play in the founding and development of Y Combinator (YC)?\n",
+ "nodes: [NodeWithScore(node=TextNode(id_='7768a74a-5a54-414f-9217-47f976da7891', embedding=None, metadata={'file_path': '../data/paul_graham/paul_graham_essay.txt', 'file_name': 'paul_graham_essay.txt', 'file...\n",
+ "\n",
+ "\u001b[0m"
+ ]
+ }
+ ],
+ "source": [
+ "response = p.run(topic=\"YC\")"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "998375e4-a8b4-4229-b034-6280d61e69f2",
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "Paul Graham played a significant role in the founding and development of Y Combinator (YC). He was one of the co-founders of YC and was actively involved in its early stages. He helped establish the Summer Founders Program (SFP) and was responsible for selecting and funding the initial batch of startups. Graham also played a key role in shaping the funding model for YC, based on previous deals and agreements. As YC grew, Graham's attention and focus shifted more towards YC, and he gradually reduced his involvement in other projects. However, he continued to work on YC's internal software and wrote essays about startups. Eventually, Graham decided to hand over the reins of YC to someone else and recruited Sam Altman as the new president.\n"
+ ]
+ }
+ ],
+ "source": [
+ "print(str(response))"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "2d976a65-1061-4b4a-8fdd-5dcf4c1bece9",
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "\u001b[1;3;38;2;155;135;227m> Running module prompt_tmpl with input: \n",
+ "topic: YC\n",
+ "\n",
+ "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running module llm with input: \n",
+ "messages: Please generate a question about Paul Graham's life regarding the following topic YC\n",
+ "\n",
+ "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running module retriever with input: \n",
+ "input: assistant: What role did Paul Graham play in the founding and development of Y Combinator (YC)?\n",
+ "\n",
+ "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running module reranker with input: \n",
+ "query_str: assistant: What role did Paul Graham play in the founding and development of Y Combinator (YC)?\n",
+ "nodes: [NodeWithScore(node=TextNode(id_='bf639f66-34a2-46d2-a30d-9aff96d130b9', embedding=None, metadata={'file_path': '../data/paul_graham/paul_graham_essay.txt', 'file_name': 'paul_graham_essay.txt', 'file...\n",
+ "\n",
+ "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running module summarizer with input: \n",
+ "query_str: assistant: What role did Paul Graham play in the founding and development of Y Combinator (YC)?\n",
+ "nodes: [NodeWithScore(node=TextNode(id_='7768a74a-5a54-414f-9217-47f976da7891', embedding=None, metadata={'file_path': '../data/paul_graham/paul_graham_essay.txt', 'file_name': 'paul_graham_essay.txt', 'file...\n",
+ "\n",
+ "\u001b[0mPaul Graham played a significant role in the founding and development of Y Combinator (YC). He was one of the co-founders of YC and was actively involved in its early stages. He helped establish the Summer Founders Program (SFP), which was the precursor to YC, and played a key role in selecting and funding the first batch of startups. As YC grew, Graham became more involved and dedicated a significant amount of his time and attention to the organization. He worked on various projects within YC, including writing essays and developing YC's internal software. However, over time, Graham realized that YC was taking up more of his attention and decided to hand over the reins to someone else. He played a crucial role in recruiting Sam Altman as the new president of YC and transitioning the leadership to him.\n"
+ ]
+ }
+ ],
+ "source": [
+ "# you can do async too\n",
+ "response = await p.arun(topic=\"YC\")\n",
+ "print(str(response))"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "ce083e06-7943-4bc9-b134-ce4b7aa11f6a",
+ "metadata": {},
+ "source": [
+ "### 2. RAG Pipeline without Query Rewriting\n",
+ "\n",
+ "Here we setup a RAG pipeline without the query rewriting step. \n",
+ "\n",
+ "Here we need a way to link the input query to both the retriever, reranker, and summarizer. We can do this by defining a special `InputComponent`, allowing us to link the inputs to multiple downstream modules."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "ab324158-1491-46a7-95fd-c80ecd08957f",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "from llama_index.postprocessor import CohereRerank\n",
+ "from llama_index.response_synthesizers import TreeSummarize\n",
+ "from llama_index import ServiceContext\n",
+ "from llama_index.query_pipeline import InputComponent\n",
+ "\n",
+ "retriever = index.as_retriever(similarity_top_k=5)\n",
+ "summarizer = TreeSummarize(\n",
+ " service_context=ServiceContext.from_defaults(\n",
+ " llm=OpenAI(model=\"gpt-3.5-turbo\")\n",
+ " )\n",
+ ")\n",
+ "reranker = CohereRerank()"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "230cc32f-f155-4b72-8407-7618aa07622e",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "p = QueryPipeline(verbose=True)\n",
+ "p.add_modules(\n",
+ " {\n",
+ " \"input\": InputComponent(),\n",
+ " \"retriever\": retriever,\n",
+ " \"summarizer\": summarizer,\n",
+ " }\n",
+ ")\n",
+ "p.add_link(\"input\", \"retriever\")\n",
+ "p.add_link(\"input\", \"summarizer\", dest_key=\"query_str\")\n",
+ "p.add_link(\"retriever\", \"summarizer\", dest_key=\"nodes\")"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "2a64384a-6574-464a-a2a4-fa5f25091025",
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "\u001b[1;3;38;2;155;135;227m> Running module input with input: \n",
+ "input: what did the author do in YC\n",
+ "\n",
+ "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running module retriever with input: \n",
+ "input: what did the author do in YC\n",
+ "\n",
+ "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running module summarizer with input: \n",
+ "query_str: what did the author do in YC\n",
+ "nodes: [NodeWithScore(node=TextNode(id_='abf6284e-6ad9-43f0-9f18-a8ca595b2ba2', embedding=None, metadata={'file_path': '../data/paul_graham/paul_graham_essay.txt', 'file_name': 'paul_graham_essay.txt', 'file...\n",
+ "\n",
+ "\u001b[0m"
+ ]
+ }
+ ],
+ "source": [
+ "output = p.run(input=\"what did the author do in YC\")"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "82f261e6-d4f9-444e-904c-c253f97123d2",
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "The author had a diverse range of responsibilities at YC, including writing essays, working on YC's internal software, funding and supporting startups, dealing with disputes between cofounders, identifying dishonesty, and advocating for the startups.\n"
+ ]
+ }
+ ],
+ "source": [
+ "print(str(output))"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "76c11f28-983d-4b3e-94a1-541c2804b989",
+ "metadata": {},
+ "source": [
+ "## Defining a Custom Component in a Query Pipeline\n",
+ "\n",
+ "You can easily define a custom component. Simply subclass a `QueryComponent`, implement validation/run functions + some helpers, and plug it in.\n",
+ "\n",
+ "Let's wrap the related movie generation prompt+LLM chain from the first example into a custom component."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "6c470d37-a427-44a0-9b66-c83cbcac2545",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "from llama_index.query_pipeline import (\n",
+ " CustomQueryComponent,\n",
+ " InputKeys,\n",
+ " OutputKeys,\n",
+ ")\n",
+ "from typing import Dict, Any\n",
+ "from llama_index.llms.llm import BaseLLM\n",
+ "from pydantic import Field\n",
+ "\n",
+ "\n",
+ "class RelatedMovieComponent(CustomQueryComponent):\n",
+ " \"\"\"Related movie component.\"\"\"\n",
+ "\n",
+ " llm: BaseLLM = Field(..., description=\"OpenAI LLM\")\n",
+ "\n",
+ " def _validate_component_inputs(\n",
+ " self, input: Dict[str, Any]\n",
+ " ) -> Dict[str, Any]:\n",
+ " \"\"\"Validate component inputs during run_component.\"\"\"\n",
+ " # NOTE: this is OPTIONAL but we show you here how to do validation as an example\n",
+ " return input\n",
+ "\n",
+ " @property\n",
+ " def _input_keys(self) -> set:\n",
+ " \"\"\"Input keys dict.\"\"\"\n",
+ " # NOTE: These are required inputs. If you have optional inputs please override\n",
+ " # `optional_input_keys_dict`\n",
+ " return {\"movie\"}\n",
+ "\n",
+ " @property\n",
+ " def _output_keys(self) -> set:\n",
+ " return {\"output\"}\n",
+ "\n",
+ " def _run_component(self, **kwargs) -> Dict[str, Any]:\n",
+ " \"\"\"Run the component.\"\"\"\n",
+ " # use QueryPipeline itself here for convenience\n",
+ " prompt_str = \"Please generate related movies to {movie_name}\"\n",
+ " prompt_tmpl = PromptTemplate(prompt_str)\n",
+ " p = QueryPipeline(chain=[prompt_tmpl, llm])\n",
+ " return {\"output\": p.run(movie_name=kwargs[\"movie\"])}"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "69dc1a62-1e84-45c5-9f12-cac7778bd46f",
+ "metadata": {},
+ "source": [
+ "Let's try the custom component out! We'll also add a step to convert the output to Shakespeare."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "d8ff382d-baf8-46e2-b4c0-477a07a41219",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "llm = OpenAI(model=\"gpt-3.5-turbo\")\n",
+ "component = RelatedMovieComponent(llm=llm)\n",
+ "\n",
+ "# let's add some subsequent prompts for fun\n",
+ "prompt_str = \"\"\"\\\n",
+ "Here's some text:\n",
+ "\n",
+ "{text}\n",
+ "\n",
+ "Can you rewrite this in the voice of Shakespeare?\n",
+ "\"\"\"\n",
+ "prompt_tmpl = PromptTemplate(prompt_str)\n",
+ "\n",
+ "p = QueryPipeline(chain=[component, prompt_tmpl, llm], verbose=True)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "b7f6aa04-4443-4511-966d-c71cbd268efa",
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "\u001b[1;3;38;2;155;135;227m> Running module 55d9172e-5bc5-49c8-9bbd-bffe99986f5e with input: \n",
+ "movie: Love Actually\n",
+ "\n",
+ "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running module c3cd65ad-13bb-42f5-b703-6ff2657dac02 with input: \n",
+ "text: assistant: 1. \"Valentine's Day\" (2010) - A star-studded ensemble cast explores interconnected love stories on Valentine's Day in Los Angeles.\n",
+ "\n",
+ "2. \"New Year's Eve\" (2011) - Similar to \"Love Actually,\" ...\n",
+ "\n",
+ "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running module 9146da4e-8e84-4616-bd97-b7f268a5f74d with input: \n",
+ "messages: Here's some text:\n",
+ "\n",
+ "assistant: 1. \"Valentine's Day\" (2010) - A star-studded ensemble cast explores interconnected love stories on Valentine's Day in Los Angeles.\n",
+ "\n",
+ "2. \"New Year's Eve\" (2011) - Similar t...\n",
+ "\n",
+ "\u001b[0m"
+ ]
+ }
+ ],
+ "source": [
+ "output = p.run(movie=\"Love Actually\")"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "726ee679-e6a3-4e57-9566-00ef9f40b837",
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "assistant: 1. \"Valentine's Daye\" (2010) - A troupe of stars doth explore intertwined tales of love on Valentine's Daye in fair Los Angeles.\n",
+ "\n",
+ "2. \"New Year's Eve\" (2011) - Much like \"Love Actually,\" this play doth followeth diverse characters as they navigate love and relationships on New Year's Eve in fair New York City.\n",
+ "\n",
+ "3. \"He's Just Not That Into Thee\" (2009) - This romantic comedy doth feature intersecting storylines that explore the complexities of modern relationships and the quest for true love.\n",
+ "\n",
+ "4. \"Crazy, Stupid, Love\" (2011) - A man of middle age doth see his life unravel when his wife doth asketh for a divorce, leading him to seeketh guidance from a young bachelor who doth help him rediscover love.\n",
+ "\n",
+ "5. \"The Holiday\" (2006) - Two women, one from fair Los Angeles and the other from England, doth exchange homes during the Christmas season and unexpectedly findeth love in their new surroundings.\n",
+ "\n",
+ "6. \"Four Weddings and a Funeral\" (1994) - This British romantic comedy doth followeth a group of friends as they attendeth various weddings and a funeral, exploring love, friendship, and the complexities of relationships.\n",
+ "\n",
+ "7. \"Notting Hill\" (1999) - A famous actress and a humble bookstore owner doth fall in love in the vibrant neighborhood of Notting Hill, London, despite the challenges posed by their different worlds.\n",
+ "\n",
+ "8. \"Bridget Jones's Diary\" (2001) - Bridget Jones, a quirky and relatable woman, doth navigate her love life whilst documenting her experiences in a diary, leading to unexpected romantic entanglements.\n",
+ "\n",
+ "9. \"About Time\" (2013) - A young man doth discover he hath the ability to travel through time and useth it to find love, but soon realizeth that altering the past can have unforeseen consequences.\n",
+ "\n",
+ "10. \"The Best Exotic Marigold Hotel\" (2011) - A group of British retirees doth decide to spendeth their golden years in a seemingly luxurious hotel in India, where they findeth unexpected love and new beginnings.\n"
+ ]
+ }
+ ],
+ "source": [
+ "print(str(output))"
+ ]
+ }
+ ],
+ "metadata": {
+ "kernelspec": {
+ "display_name": "llama_index_logan",
+ "language": "python",
+ "name": "llama_index_logan"
+ },
+ "language_info": {
+ "codemirror_mode": {
+ "name": "ipython",
+ "version": 3
+ },
+ "file_extension": ".py",
+ "mimetype": "text/x-python",
+ "name": "python",
+ "nbconvert_exporter": "python",
+ "pygments_lexer": "ipython3"
+ }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}
diff --git a/docs/module_guides/querying/pipeline/module_usage.md b/docs/module_guides/querying/pipeline/module_usage.md
new file mode 100644
index 0000000000..1c6950a404
--- /dev/null
+++ b/docs/module_guides/querying/pipeline/module_usage.md
@@ -0,0 +1,71 @@
+# Module Usage
+
+Currently the following LlamaIndex modules are supported within a QueryPipeline. Remember, you can define your own!
+
+### LLMs (both completion and chat)
+
+- Base class: `LLM`
+- [Module Guide](/module_guides/models/llms.md)
+- If chat model:
+ - Input: `messages`. Takes in any `List[ChatMessage]` or any stringable input.
+ - Output: `output`. Outputs `ChatResponse` (stringable)
+- If completion model:
+ - Input: `prompt`. Takes in any stringable input.
+ - Output: `output`. Outputs `CompletionResponse` (stringable)
+
+### Prompts
+
+- Base class: `PromptTemplate`
+- [Module Guide](/module_guides/models/prompts.md)
+- Input: Prompt template variables. Each variable can be a stringable input.
+- Output: `output`. Outputs formatted prompt string (stringable)
+
+### Query Engines
+
+- Base class: `BaseQueryEngine`
+- [Module Guide](/module_guides/deploying/query_engine/root.md)
+- Input: `input`. Takes in any stringable input.
+- Output: `output`. Outputs `Response` (stringable)
+
+### Query Transforms
+
+- Base class: `BaseQueryTransform`
+- [Module Guide](/optimizing/advanced_retrieval/query_transformations.md)
+- Input: `query_str`, `metadata` (optional). `query_str` is any stringable input.
+- Output: `query_str`. Outputs string.
+
+### Retrievers
+
+- Base class: `BaseRetriever`
+- [Module Guide](/module_guides/querying/retriever/root.md)
+- Input: `input`. Takes in any stringable input.
+- Output: `output`. Outputs list of nodes `List[BaseNode]`.
+
+### Output Parsers
+
+- Base class: `BaseOutputParser`
+- [Module Guide](/module_guides/querying/structured_outputs/output_parser.md)
+- Input: `input`. Takes in any stringable input.
+- Output: `output`. Outputs whatever type output parser is supposed to parse out.
+
+### Postprocessors/Rerankers
+
+- Base class: `BaseNodePostprocessor`
+- [Module Guide](/module_guides/querying/node_postprocessors/root.md)
+- Input: `nodes`, `query_str` (optional). `nodes` is `List[BaseNode]`, `query_str` is any stringable input.
+- Output: `nodes`. Outputs list of nodes `List[BaseNode]`.
+
+### Response Synthesizers
+
+- Base class: `BaseSynthesizer`
+- [Module Guide]()
+- Input: `nodes`, `query_str`. `nodes` is `List[BaseNode]`, `query_str` is any stringable input.
+- Output: `output`. Outputs `Response` object (stringable).
+
+### Other QueryPipeline objects
+
+You can define a `QueryPipeline` as a module within another query pipeline. This makes it easy for you to string together complex workflows.
+
+### Custom Components
+
+See our [custom components guide](query-pipeline-custom-component) for more details.
diff --git a/docs/module_guides/querying/pipeline/modules.md b/docs/module_guides/querying/pipeline/modules.md
new file mode 100644
index 0000000000..656246b160
--- /dev/null
+++ b/docs/module_guides/querying/pipeline/modules.md
@@ -0,0 +1,8 @@
+# Module Guides
+
+```{toctree}
+---
+maxdepth: 1
+---
+/examples/pipeline/query_pipeline.ipynb
+```
diff --git a/docs/module_guides/querying/pipeline/root.md b/docs/module_guides/querying/pipeline/root.md
new file mode 100644
index 0000000000..93d60b415a
--- /dev/null
+++ b/docs/module_guides/querying/pipeline/root.md
@@ -0,0 +1,58 @@
+# Query Pipeline
+
+## Concept
+
+LlamaIndex provides a declarative query API that allows you to chain together different modules in order to orchestrate simple-to-advanced workflows over your data.
+
+This is centered around our `QueryPipeline` abstraction. Load in a variety of modules (from LLMs to prompts to retrievers to other pipelines), connect them all together into a sequential chain or DAG, and run it end2end.
+
+**NOTE**: You can orchestrate all these workflows without the declarative pipeline abstraction (by using the modules imperatively and writing your own functions). So what are the advantages of `QueryPipeline`?
+
+- Express common workflows with fewer lines of code/boilerplate
+- Greater readability
+- Greater parity / better integration points with common low-code / no-code solutions (e.g. LangFlow)
+- [In the future] A declarative interface allows easy serializability of pipeline components, providing portability of pipelines/easier deployment to different systems.
+
+Our query pipelines also propagate callbacks throughout all sub-modules, and these integrate with our [observability partners](/module_guides/observability/observability.md).
+
+
+
+## Usage Pattern
+
+Here are two simple ways to setup a query pipeline - through a simplified syntax of setting up a sequential chain to setting up a full compute DAG.
+
+```python
+from llama_index.query_pipeline.query import QueryPipeline
+
+# sequential chain
+p = QueryPipeline(chain=[prompt_tmpl, llm], verbose=True)
+
+# DAG
+p = QueryPipeline(verbose=True)
+p.add_modules({"prompt_tmpl": prompt_tmpl, "llm": llm})
+p.add_link("prompt_tmpl", "llm")
+
+# run pipeline
+p.run(prompt_key1="<input1>", ...)
+```
+
+More information can be found in our usage pattern guides below.
+
+```{toctree}
+---
+maxdepth: 2
+---
+usage_pattern.md
+module_usage.md
+```
+
+## Module Guides
+
+Check out our `QueryPipeline` end-to-end guides to learn standard to advanced ways to setup orchestration over your data.
+
+```{toctree}
+---
+maxdepth: 2
+---
+modules.md
+```
diff --git a/docs/module_guides/querying/pipeline/usage_pattern.md b/docs/module_guides/querying/pipeline/usage_pattern.md
new file mode 100644
index 0000000000..b556ba2b52
--- /dev/null
+++ b/docs/module_guides/querying/pipeline/usage_pattern.md
@@ -0,0 +1,178 @@
+# Usage Pattern
+
+The usage pattern guide covers setup + usage of the `QueryPipeline` more in-depth.
+
+## Setting up a Pipeline
+
+Here we walk through a few different ways of setting up a query pipeline.
+
+### Defining a Sequential Chain
+
+Some simple pipelines are purely linear in nature - the output of the previous module directly goes into the input of the next module.
+
+Some examples:
+
+- prompt -> LLM -> output parsing
+- prompt -> LLM -> prompt -> LLM
+- retriever -> response synthesizer
+
+These workflows can easily be expressed in the `QueryPipeline` through a simplified `chain` syntax.
+
+```python
+from llama_index.query_pipeline.query import QueryPipeline
+
+# try chaining basic prompts
+prompt_str = "Please generate related movies to {movie_name}"
+prompt_tmpl = PromptTemplate(prompt_str)
+llm = OpenAI(model="gpt-3.5-turbo")
+
+p = QueryPipeline(chain=[prompt_tmpl, llm], verbose=True)
+```
+
+### Defining a DAG
+
+Many pipelines will require you to setup a DAG (for instance, if you want to implement all the steps in a standard RAG pipeline).
+
+Here we offer a lower-level API to add modules along with their keys, and define links between previous module outputs to next
+module inputs.
+
+```python
+from llama_index.postprocessor import CohereRerank
+from llama_index.response_synthesizers import TreeSummarize
+from llama_index import ServiceContext
+
+# define modules
+prompt_str = "Please generate a question about Paul Graham's life regarding the following topic {topic}"
+prompt_tmpl = PromptTemplate(prompt_str)
+llm = OpenAI(model="gpt-3.5-turbo")
+retriever = index.as_retriever(similarity_top_k=3)
+reranker = CohereRerank()
+summarizer = TreeSummarize(
+ service_context=ServiceContext.from_defaults(llm=llm)
+)
+
+# define query pipeline
+p = QueryPipeline(verbose=True)
+p.add_modules(
+ {
+ "llm": llm,
+ "prompt_tmpl": prompt_tmpl,
+ "retriever": retriever,
+ "summarizer": summarizer,
+ "reranker": reranker,
+ }
+)
+p.add_link("prompt_tmpl", "llm")
+p.add_link("llm", "retriever")
+p.add_link("retriever", "reranker", dest_key="nodes")
+p.add_link("llm", "reranker", dest_key="query_str")
+p.add_link("reranker", "summarizer", dest_key="nodes")
+p.add_link("llm", "summarizer", dest_key="query_str")
+```
+
+## Running the Pipeline
+
+### Single-Input/Single-Output
+
+The input is the kwargs of the first component.
+
+If the output of the last component is a single object (and not a dictionary of objects), then we return that directly.
+
+Taking the pipeline in the previous example, the output will be a `Response` object since the last step is the `TreeSummarize` response synthesis module.
+
+```python
+output = p.run(topic="YC")
+# output type is Response
+type(output)
+```
+
+### Multi-Input/Multi-Output
+
+If your DAG has multiple root nodes / and-or output nodes, you can try `run_multi`. Pass in an input dictionary containing module key -> input dict. Output is dictionary of module key -> output dict.
+
+If we ran the prev example,
+
+```python
+output_dict = p.run_multi({"llm": {"topic": "YC"}})
+print(output_dict)
+
+# output dict is {"summarizer": {"output": response}}
+```
+
+### Defining partials
+
+If you wish to prefill certain inputs for a module, you can do so with `partial`! Then the DAG would just hook into the unfilled inputs.
+
+You may need to convert a module via `as_query_component`.
+
+Here's an example:
+
+```python
+summarizer = TreeSummarize(
+ service_context=ServiceContext.from_defaults(llm=llm)
+)
+summarizer_c = summarizer.as_query_component(partial={"nodes": nodes})
+# can define a chain because llm output goes into query_str, nodes is pre-filled
+p = QueryPipeline(chain=[prompt_tmpl, llm, summarizer_c])
+# run pipeline
+p.run(topic="YC")
+```
+
+(query-pipeline-custom-component)=
+
+## Defining a Custom Query Component
+
+You can easily define a custom component. Simply subclass a `QueryComponent`, implement validation/run functions + some helpers, and plug it in.
+
+```python
+from llama_index.query_pipeline import CustomQueryComponent
+from typing import Dict, Any
+
+
+class MyComponent(CustomQueryComponent):
+ """My component."""
+
+ # Pydantic class, put any attributes here
+ ...
+
+ def _validate_component_inputs(
+ self, input: Dict[str, Any]
+ ) -> Dict[str, Any]:
+ """Validate component inputs during run_component."""
+ # NOTE: this is OPTIONAL but we show you here how to do validation as an example
+ return input
+
+ @property
+ def _input_keys(self) -> set:
+ """Input keys dict."""
+ return {"input_key1", ...}
+
+ @property
+ def _output_keys(self) -> set:
+ # can do multi-outputs too
+ return {"output_key"}
+
+ def _run_component(self, **kwargs) -> Dict[str, Any]:
+ """Run the component."""
+ # run logic
+ ...
+ return {"output_key": result}
+```
+
+For more details check out our [in-depth query transformations guide](/examples/pipeline/query_pipeline.ipynb).
+
+## Ensuring outputs are compatible
+
+By linking modules within a `QueryPipeline`, the output of one module goes into the input of the next module.
+
+Generally you must make sure that for a link to work, the expected output and input types _roughly_ line up.
+
+We say roughly because we do some magic on existing modules to make sure that "stringable" outputs can be passed into
+inputs that can be queried as a "string". Certain output types are treated as Stringable - `CompletionResponse`, `ChatResponse`, `Response`, `QueryBundle`, etc. Retrievers/query engines will automatically convert `string` inputs to `QueryBundle` objects.
+
+This lets you do certain workflows that would otherwise require boilerplate string conversion if you were writing this yourself, for instance,
+
+- LLM -> prompt, LLM -> retriever, LLM -> query engine
+- query engine -> prompt, query engine -> retriever
+
+If you are defining a custom component, you should use `_validate_component_inputs` to ensure that the inputs are the right type, and throw an error if they're not.
diff --git a/docs/module_guides/querying/querying.md b/docs/module_guides/querying/querying.md
index 9f9007b4e4..595cd67e7a 100644
--- a/docs/module_guides/querying/querying.md
+++ b/docs/module_guides/querying/querying.md
@@ -2,7 +2,21 @@
Querying is the most important part of your LLM application. To learn more about getting a final product that you can deploy, check out the [query engine](/module_guides/deploying/query_engine/root.md), [chat engine](/module_guides/deploying/chat_engines/root.md) and [agents](/module_guides/deploying/agents/root.md) sections.
-## Modules in this section
+## Query Pipeline
+
+You can create query pipelines/chains with ease with our declarative `QueryPipeline` interface. Check out our [query pipeline guide](/module_guides/querying/pipeline/root.md) for more details.
+
+```{toctree}
+---
+maxdepth: 1
+hidden: True
+---
+/module_guides/querying/pipeline/root.md
+```
+
+Otherwise check out how to use our query modules as standalone components 👇.
+
+## Query Modules
```{toctree}
---
diff --git a/llama_index/core/base_query_engine.py b/llama_index/core/base_query_engine.py
index 934b37314e..4713a99f4f 100644
--- a/llama_index/core/base_query_engine.py
+++ b/llama_index/core/base_query_engine.py
@@ -4,7 +4,15 @@ import logging
from abc import abstractmethod
from typing import Any, Dict, List, Optional, Sequence
+from llama_index.bridge.pydantic import Field
from llama_index.callbacks.base import CallbackManager
+from llama_index.core.query_pipeline.query_component import (
+ ChainableMixin,
+ InputKeys,
+ OutputKeys,
+ QueryComponent,
+ validate_and_convert_stringable,
+)
from llama_index.core.response.schema import RESPONSE_TYPE
from llama_index.prompts.mixin import PromptDictType, PromptMixin
from llama_index.schema import NodeWithScore, QueryBundle, QueryType
@@ -12,7 +20,7 @@ from llama_index.schema import NodeWithScore, QueryBundle, QueryType
logger = logging.getLogger(__name__)
-class BaseQueryEngine(PromptMixin):
+class BaseQueryEngine(ChainableMixin, PromptMixin):
"""Base query engine."""
def __init__(self, callback_manager: Optional[CallbackManager]) -> None:
@@ -69,3 +77,46 @@ class BaseQueryEngine(PromptMixin):
@abstractmethod
async def _aquery(self, query_bundle: QueryBundle) -> RESPONSE_TYPE:
pass
+
+ def _as_query_component(self, **kwargs: Any) -> QueryComponent:
+ """Return a query component."""
+ return QueryEngineComponent(query_engine=self)
+
+
+class QueryEngineComponent(QueryComponent):
+ """Query engine component."""
+
+ query_engine: BaseQueryEngine = Field(..., description="Query engine")
+
+ class Config:
+ arbitrary_types_allowed = True
+
+ def set_callback_manager(self, callback_manager: CallbackManager) -> None:
+ """Set callback manager."""
+ self.query_engine.callback_manager = callback_manager
+
+ def _validate_component_inputs(self, input: Dict[str, Any]) -> Dict[str, Any]:
+ """Validate component inputs during run_component."""
+ # make sure input is a string
+ input["input"] = validate_and_convert_stringable(input["input"])
+ return input
+
+ def _run_component(self, **kwargs: Any) -> Any:
+ """Run component."""
+ output = self.query_engine.query(kwargs["input"])
+ return {"output": output}
+
+ async def _arun_component(self, **kwargs: Any) -> Any:
+ """Run component."""
+ output = await self.query_engine.aquery(kwargs["input"])
+ return {"output": output}
+
+ @property
+ def input_keys(self) -> InputKeys:
+ """Input keys."""
+ return InputKeys.from_keys({"input"})
+
+ @property
+ def output_keys(self) -> OutputKeys:
+ """Output keys."""
+ return OutputKeys.from_keys({"output"})
diff --git a/llama_index/core/base_retriever.py b/llama_index/core/base_retriever.py
index e36607c4e3..0b735e0ddd 100644
--- a/llama_index/core/base_retriever.py
+++ b/llama_index/core/base_retriever.py
@@ -1,19 +1,27 @@
"""Base retriever."""
from abc import abstractmethod
-from typing import List, Optional
+from typing import Any, Dict, List, Optional
+from llama_index.bridge.pydantic import Field
from llama_index.callbacks.base import CallbackManager
from llama_index.callbacks.schema import CBEventType, EventPayload
+from llama_index.core.query_pipeline.query_component import (
+ ChainableMixin,
+ InputKeys,
+ OutputKeys,
+ QueryComponent,
+ validate_and_convert_stringable,
+)
from llama_index.prompts.mixin import PromptDictType, PromptMixin, PromptMixinType
from llama_index.schema import NodeWithScore, QueryBundle, QueryType
from llama_index.service_context import ServiceContext
-class BaseRetriever(PromptMixin):
+class BaseRetriever(ChainableMixin, PromptMixin):
"""Base retriever."""
def __init__(self, callback_manager: Optional[CallbackManager] = None) -> None:
- callback_manager = callback_manager or CallbackManager()
+ self.callback_manager = callback_manager or CallbackManager()
def _check_callback_manager(self) -> None:
"""Check callback manager."""
@@ -104,3 +112,46 @@ class BaseRetriever(PromptMixin):
elif hasattr(self, "_index") and hasattr(self._index, "service_context"):
return self._index.service_context
return None
+
+ def _as_query_component(self, **kwargs: Any) -> QueryComponent:
+ """Return a query component."""
+ return RetrieverComponent(retriever=self)
+
+
+class RetrieverComponent(QueryComponent):
+ """Retriever component."""
+
+ retriever: BaseRetriever = Field(..., description="Retriever")
+
+ class Config:
+ arbitrary_types_allowed = True
+
+ def set_callback_manager(self, callback_manager: CallbackManager) -> None:
+ """Set callback manager."""
+ self.retriever.callback_manager = callback_manager
+
+ def _validate_component_inputs(self, input: Dict[str, Any]) -> Dict[str, Any]:
+ """Validate component inputs during run_component."""
+ # make sure input is a string
+ input["input"] = validate_and_convert_stringable(input["input"])
+ return input
+
+ def _run_component(self, **kwargs: Any) -> Any:
+ """Run component."""
+ output = self.retriever.retrieve(kwargs["input"])
+ return {"output": output}
+
+ async def _arun_component(self, **kwargs: Any) -> Any:
+ """Run component."""
+ output = await self.retriever.aretrieve(kwargs["input"])
+ return {"output": output}
+
+ @property
+ def input_keys(self) -> InputKeys:
+ """Input keys."""
+ return InputKeys.from_keys({"input"})
+
+ @property
+ def output_keys(self) -> OutputKeys:
+ """Output keys."""
+ return OutputKeys.from_keys({"output"})
diff --git a/llama_index/core/query_pipeline/__init__.py b/llama_index/core/query_pipeline/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/llama_index/core/query_pipeline/query_component.py b/llama_index/core/query_pipeline/query_component.py
new file mode 100644
index 0000000000..72bca58aa1
--- /dev/null
+++ b/llama_index/core/query_pipeline/query_component.py
@@ -0,0 +1,281 @@
+"""Pipeline schema."""
+
+from abc import ABC, abstractmethod
+from typing import Any, Dict, Optional, Set, Union, get_args
+
+from llama_index.bridge.pydantic import BaseModel, Field
+from llama_index.callbacks.base import CallbackManager
+from llama_index.core.llms.types import ChatResponse, CompletionResponse
+from llama_index.core.response.schema import Response
+from llama_index.schema import QueryBundle
+
+## Define common types used throughout these components
+StringableInput = Union[CompletionResponse, ChatResponse, str, QueryBundle, Response]
+
+
+def validate_and_convert_stringable(input: Any) -> str:
+ """Validate and convert stringable input."""
+ if not isinstance(input, get_args(StringableInput)):
+ raise ValueError(f"Input {input} is not stringable.")
+ return str(input)
+
+
+class InputKeys(BaseModel):
+ """Input keys."""
+
+ required_keys: Set[str] = Field(default_factory=set)
+ optional_keys: Set[str] = Field(default_factory=set)
+
+ @classmethod
+ def from_keys(
+ cls, required_keys: Set[str], optional_keys: Optional[Set[str]] = None
+ ) -> "InputKeys":
+ """Create InputKeys from tuple."""
+ return cls(required_keys=required_keys, optional_keys=optional_keys or set())
+
+ def validate(self, input_keys: Set[str]) -> None:
+ """Validate input keys."""
+ # check if required keys are present, and that keys all are in required or optional
+ if not self.required_keys.issubset(input_keys):
+ raise ValueError(
+ f"Required keys {self.required_keys} are not present in input keys {input_keys}"
+ )
+ if not input_keys.issubset(self.required_keys.union(self.optional_keys)):
+ raise ValueError(
+ f"Input keys {input_keys} contain keys not in required or optional keys {self.required_keys.union(self.optional_keys)}"
+ )
+
+ def __len__(self) -> int:
+ """Length of input keys."""
+ return len(self.required_keys) + len(self.optional_keys)
+
+ def all(self) -> Set[str]:
+ """Get all input keys."""
+ return self.required_keys.union(self.optional_keys)
+
+
+class OutputKeys(BaseModel):
+ """Output keys."""
+
+ required_keys: Set[str] = Field(default_factory=set)
+
+ @classmethod
+ def from_keys(
+ cls,
+ required_keys: Set[str],
+ ) -> "InputKeys":
+ """Create InputKeys from tuple."""
+ return cls(required_keys=required_keys)
+
+ def validate(self, input_keys: Set[str]) -> None:
+ """Validate input keys."""
+ # validate that input keys exactly match required keys
+ if input_keys != self.required_keys:
+ raise ValueError(
+ f"Input keys {input_keys} do not match required keys {self.required_keys}"
+ )
+
+
+class ChainableMixin(ABC):
+ """Chainable mixin.
+
+ A module that can produce a `QueryComponent` from a set of inputs through
+ `as_query_component`.
+
+ If plugged in directly into a `QueryPipeline`, the `ChainableMixin` will be
+ converted into a `QueryComponent` with default parameters.
+
+ """
+
+ @abstractmethod
+ def _as_query_component(self, **kwargs: Any) -> "QueryComponent":
+ """Get query component."""
+
+ def as_query_component(
+ self, partial: Optional[Dict[str, Any]] = None, **kwargs: Any
+ ) -> "QueryComponent":
+ """Get query component."""
+ component = self._as_query_component(**kwargs)
+ component.partial(**(partial or {}))
+ return component
+
+
+class QueryComponent(BaseModel):
+ """Query component.
+
+ Represents a component that can be run in a `QueryPipeline`.
+
+ """
+
+ partial_dict: Dict[str, Any] = Field(
+ default_factory=dict, description="Partial arguments to run_component"
+ )
+
+ # TODO: make this a subclass of BaseComponent (e.g. use Pydantic)
+
+ def partial(self, **kwargs: Any) -> None:
+ """Update with partial arguments."""
+ self.partial_dict.update(kwargs)
+
+ @abstractmethod
+ def set_callback_manager(self, callback_manager: Any) -> None:
+ """Set callback manager."""
+ # TODO: refactor so that callback_manager is always passed in during runtime.
+
+ @property
+ def free_req_input_keys(self) -> Set[str]:
+ """Get free input keys."""
+ return self.input_keys.required_keys.difference(self.partial_dict.keys())
+
+ @abstractmethod
+ def _validate_component_inputs(self, input: Dict[str, Any]) -> Dict[str, Any]:
+ """Validate component inputs during run_component."""
+
+ def _validate_component_outputs(self, output: Dict[str, Any]) -> Dict[str, Any]:
+ """Validate component outputs during run_component."""
+ # override if needed
+ return output
+
+ def validate_component_inputs(self, input: Dict[str, Any]) -> Dict[str, Any]:
+ """Validate component inputs."""
+ # make sure set of input keys == self.input_keys
+ self.input_keys.validate(set(input.keys()))
+ return self._validate_component_inputs(input)
+
+ def validate_component_outputs(self, output: Dict[str, Any]) -> Dict[str, Any]:
+ """Validate component outputs."""
+ # make sure set of output keys == self.output_keys
+ self.output_keys.validate(set(output.keys()))
+ return self._validate_component_outputs(output)
+
+ def run_component(self, **kwargs: Any) -> Dict[str, Any]:
+ """Run component."""
+ kwargs.update(self.partial_dict)
+ kwargs = self.validate_component_inputs(kwargs)
+ component_outputs = self._run_component(**kwargs)
+ return self.validate_component_outputs(component_outputs)
+
+ async def arun_component(self, **kwargs: Any) -> Dict[str, Any]:
+ """Run component."""
+ kwargs.update(self.partial_dict)
+ kwargs = self.validate_component_inputs(kwargs)
+ component_outputs = await self._arun_component(**kwargs)
+ return self.validate_component_outputs(component_outputs)
+
+ @abstractmethod
+ def _run_component(self, **kwargs: Any) -> Dict:
+ """Run component."""
+
+ @abstractmethod
+ async def _arun_component(self, **kwargs: Any) -> Any:
+ """Run component (async)."""
+
+ @property
+ @abstractmethod
+ def input_keys(self) -> InputKeys:
+ """Input keys."""
+
+ @property
+ @abstractmethod
+ def output_keys(self) -> OutputKeys:
+ """Output keys."""
+
+
+class CustomQueryComponent(QueryComponent):
+ """Custom query component."""
+
+ callback_manager: CallbackManager = Field(
+ default_factory=CallbackManager, description="Callback manager"
+ )
+
+ class Config:
+ arbitrary_types_allowed = True
+
+ def set_callback_manager(self, callback_manager: Any) -> None:
+ """Set callback manager."""
+ self.callback_manager = callback_manager
+
+ def _validate_component_inputs(self, input: Dict[str, Any]) -> Dict[str, Any]:
+ """Validate component inputs during run_component."""
+ # NOTE: user can override this method to validate inputs
+ # but we do this by default for convenience
+ return input
+
+ async def _arun_component(self, **kwargs: Any) -> Any:
+ """Run component (async)."""
+ raise NotImplementedError("This component does not support async run.")
+
+ @property
+ def _input_keys(self) -> Set[str]:
+ """Input keys dict."""
+ raise NotImplementedError("Not implemented yet. Please override this method.")
+
+ @property
+ def _optional_input_keys(self) -> Set[str]:
+ """Optional input keys dict."""
+ return set()
+
+ @property
+ def _output_keys(self) -> Set[str]:
+ """Output keys dict."""
+ raise NotImplementedError("Not implemented yet. Please override this method.")
+
+ @property
+ def input_keys(self) -> InputKeys:
+ """Input keys."""
+ # NOTE: user can override this too, but we have them implement an
+ # abstract method to make sure they do it
+
+ return InputKeys.from_keys(
+ required_keys=self._input_keys, optional_keys=self._optional_input_keys
+ )
+
+ @property
+ def output_keys(self) -> OutputKeys:
+ """Output keys."""
+ # NOTE: user can override this too, but we have them implement an
+ # abstract method to make sure they do it
+ return OutputKeys.from_keys(self._output_keys)
+
+
+class InputComponent(QueryComponent):
+ """Input component."""
+
+ def _validate_component_inputs(self, input: Dict[str, Any]) -> Dict[str, Any]:
+ return input
+
+ def _validate_component_outputs(self, input: Dict[str, Any]) -> Dict[str, Any]:
+ return input
+
+ def validate_component_inputs(self, input: Dict[str, Any]) -> Dict[str, Any]:
+ """Validate component inputs."""
+ # NOTE: we override this to do nothing
+ return input
+
+ def validate_component_outputs(self, output: Dict[str, Any]) -> Dict[str, Any]:
+ """Validate component outputs."""
+ # NOTE: we override this to do nothing
+ return output
+
+ def set_callback_manager(self, callback_manager: Any) -> None:
+ """Set callback manager."""
+
+ def _run_component(self, **kwargs: Any) -> Any:
+ """Run component."""
+ return kwargs
+
+ async def _arun_component(self, **kwargs: Any) -> Any:
+ """Run component (async)."""
+ return self._run_component(**kwargs)
+
+ @property
+ def input_keys(self) -> InputKeys:
+ """Input keys."""
+ # NOTE: this shouldn't be used
+ return InputKeys.from_keys(set(), optional_keys=set())
+ # return InputComponentKeys.from_keys(set(), optional_keys=set())
+
+ @property
+ def output_keys(self) -> OutputKeys:
+ """Output keys."""
+ return OutputKeys.from_keys(set())
diff --git a/llama_index/indices/query/query_transform/base.py b/llama_index/indices/query/query_transform/base.py
index c69d2cc316..a7efa69994 100644
--- a/llama_index/indices/query/query_transform/base.py
+++ b/llama_index/indices/query/query_transform/base.py
@@ -2,8 +2,16 @@
import dataclasses
from abc import abstractmethod
-from typing import Dict, Optional, cast
-
+from typing import Any, Dict, Optional, cast
+
+from llama_index.bridge.pydantic import Field
+from llama_index.core.query_pipeline.query_component import (
+ ChainableMixin,
+ InputKeys,
+ OutputKeys,
+ QueryComponent,
+ validate_and_convert_stringable,
+)
from llama_index.core.response.schema import Response
from llama_index.indices.query.query_transform.prompts import (
DEFAULT_DECOMPOSE_QUERY_TRANSFORM_PROMPT,
@@ -22,7 +30,7 @@ from llama_index.schema import QueryBundle, QueryType
from llama_index.utils import print_text
-class BaseQueryTransform(PromptMixin):
+class BaseQueryTransform(ChainableMixin, PromptMixin):
"""Base class for query transform.
A query transform augments a raw query string with associated transformations
@@ -66,6 +74,10 @@ class BaseQueryTransform(PromptMixin):
"""Run query processor."""
return self.run(query_bundle_or_str, metadata=metadata)
+ def _as_query_component(self, **kwargs: Any) -> QueryComponent:
+ """As query component."""
+ return QueryTransformComponent(query_transform=self)
+
class IdentityQueryTransform(BaseQueryTransform):
"""Identity query transform.
@@ -302,3 +314,49 @@ class StepDecomposeQueryTransform(BaseQueryTransform):
query_str=new_query_str,
custom_embedding_strs=query_bundle.custom_embedding_strs,
)
+
+
+class QueryTransformComponent(QueryComponent):
+ """Query transform component."""
+
+ query_transform: BaseQueryTransform = Field(..., description="Query transform.")
+
+ class Config:
+ arbitrary_types_allowed = True
+
+ def set_callback_manager(self, callback_manager: Any) -> None:
+ """Set callback manager."""
+ # TODO: not implemented yet
+
+ def _validate_component_inputs(self, input: Dict[str, Any]) -> Dict[str, Any]:
+ """Validate component inputs during run_component."""
+ if "query_str" not in input:
+ raise ValueError("Input must have key 'query_str'")
+ input["query_str"] = validate_and_convert_stringable(input["query_str"])
+
+ input["metadata"] = input.get("metadata", {})
+
+ return input
+
+ def _run_component(self, **kwargs: Any) -> Any:
+ """Run component."""
+ output = self._query_transform.run(
+ kwargs["query_str"],
+ metadata=kwargs["metadata"],
+ )
+ return {"query_str": output.query_str}
+
+ async def _arun_component(self, **kwargs: Any) -> Any:
+ """Run component."""
+ # TODO: true async not implemented yet
+ return self._run_component(**kwargs)
+
+ @property
+ def input_keys(self) -> InputKeys:
+ """Input keys."""
+ return InputKeys.from_keys({"query_str"}, optional_keys={"metadata"})
+
+ @property
+ def output_keys(self) -> OutputKeys:
+ """Output keys."""
+ return OutputKeys.from_keys({"query_str"})
diff --git a/llama_index/llms/ai21.py b/llama_index/llms/ai21.py
index 860e360345..4526eb4cc1 100644
--- a/llama_index/llms/ai21.py
+++ b/llama_index/llms/ai21.py
@@ -104,7 +104,9 @@ class AI21(CustomLLM):
}
@llm_completion_callback()
- def complete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
+ def complete(
+ self, prompt: str, formatted: bool = False, **kwargs: Any
+ ) -> CompletionResponse:
all_kwargs = self._get_all_kwargs(**kwargs)
import ai21
@@ -118,7 +120,9 @@ class AI21(CustomLLM):
)
@llm_completion_callback()
- def stream_complete(self, prompt: str, **kwargs: Any) -> CompletionResponseGen:
+ def stream_complete(
+ self, prompt: str, formatted: bool = False, **kwargs: Any
+ ) -> CompletionResponseGen:
raise NotImplementedError(
"AI21 does not currently support streaming completion."
)
diff --git a/llama_index/llms/anthropic.py b/llama_index/llms/anthropic.py
index 5cbf2ca487..1538a238a6 100644
--- a/llama_index/llms/anthropic.py
+++ b/llama_index/llms/anthropic.py
@@ -163,7 +163,9 @@ class Anthropic(LLM):
)
@llm_completion_callback()
- def complete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
+ def complete(
+ self, prompt: str, formatted: bool = False, **kwargs: Any
+ ) -> CompletionResponse:
complete_fn = chat_to_completion_decorator(self.chat)
return complete_fn(prompt, **kwargs)
@@ -193,7 +195,9 @@ class Anthropic(LLM):
return gen()
@llm_completion_callback()
- def stream_complete(self, prompt: str, **kwargs: Any) -> CompletionResponseGen:
+ def stream_complete(
+ self, prompt: str, formatted: bool = False, **kwargs: Any
+ ) -> CompletionResponseGen:
stream_complete_fn = stream_chat_to_completion_decorator(self.stream_chat)
return stream_complete_fn(prompt, **kwargs)
@@ -215,7 +219,9 @@ class Anthropic(LLM):
)
@llm_completion_callback()
- async def acomplete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
+ async def acomplete(
+ self, prompt: str, formatted: bool = False, **kwargs: Any
+ ) -> CompletionResponse:
acomplete_fn = achat_to_completion_decorator(self.achat)
return await acomplete_fn(prompt, **kwargs)
@@ -246,7 +252,7 @@ class Anthropic(LLM):
@llm_completion_callback()
async def astream_complete(
- self, prompt: str, **kwargs: Any
+ self, prompt: str, formatted: bool = False, **kwargs: Any
) -> CompletionResponseAsyncGen:
astream_complete_fn = astream_chat_to_completion_decorator(self.astream_chat)
return await astream_complete_fn(prompt, **kwargs)
diff --git a/llama_index/llms/base.py b/llama_index/llms/base.py
index e3a0b1b317..59acd91d77 100644
--- a/llama_index/llms/base.py
+++ b/llama_index/llms/base.py
@@ -22,6 +22,9 @@ from llama_index.core.llms.types import (
CompletionResponseGen,
LLMMetadata,
)
+from llama_index.core.query_pipeline.query_component import (
+ ChainableMixin,
+)
from llama_index.schema import BaseComponent
@@ -276,7 +279,7 @@ def llm_completion_callback() -> Callable:
return wrap
-class BaseLLM(BaseComponent):
+class BaseLLM(ChainableMixin, BaseComponent):
"""LLM interface."""
callback_manager: CallbackManager = Field(
@@ -302,7 +305,9 @@ class BaseLLM(BaseComponent):
"""Chat endpoint for LLM."""
@abstractmethod
- def complete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
+ def complete(
+ self, prompt: str, formatted: bool = False, **kwargs: Any
+ ) -> CompletionResponse:
"""Completion endpoint for LLM."""
@abstractmethod
@@ -312,7 +317,9 @@ class BaseLLM(BaseComponent):
"""Streaming chat endpoint for LLM."""
@abstractmethod
- def stream_complete(self, prompt: str, **kwargs: Any) -> CompletionResponseGen:
+ def stream_complete(
+ self, prompt: str, formatted: bool = False, **kwargs: Any
+ ) -> CompletionResponseGen:
"""Streaming completion endpoint for LLM."""
# ===== Async Endpoints =====
@@ -323,7 +330,9 @@ class BaseLLM(BaseComponent):
"""Async chat endpoint for LLM."""
@abstractmethod
- async def acomplete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
+ async def acomplete(
+ self, prompt: str, formatted: bool = False, **kwargs: Any
+ ) -> CompletionResponse:
"""Async completion endpoint for LLM."""
@abstractmethod
@@ -334,6 +343,6 @@ class BaseLLM(BaseComponent):
@abstractmethod
async def astream_complete(
- self, prompt: str, **kwargs: Any
+ self, prompt: str, formatted: bool = False, **kwargs: Any
) -> CompletionResponseAsyncGen:
"""Async streaming completion endpoint for LLM."""
diff --git a/llama_index/llms/bedrock.py b/llama_index/llms/bedrock.py
index b76d19b91f..32a0bc4bf1 100644
--- a/llama_index/llms/bedrock.py
+++ b/llama_index/llms/bedrock.py
@@ -179,9 +179,10 @@ class Bedrock(LLM):
}
@llm_completion_callback()
- def complete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
- is_formatted = kwargs.pop("formatted", False)
- if not is_formatted:
+ def complete(
+ self, prompt: str, formatted: bool = False, **kwargs: Any
+ ) -> CompletionResponse:
+ if not formatted:
prompt = self.completion_to_prompt(prompt)
all_kwargs = self._get_all_kwargs(**kwargs)
request_body = self._provider.get_request_body(prompt, all_kwargs)
@@ -199,12 +200,15 @@ class Bedrock(LLM):
)
@llm_completion_callback()
- def stream_complete(self, prompt: str, **kwargs: Any) -> CompletionResponseGen:
+ def stream_complete(
+ self, prompt: str, formatted: bool = False, **kwargs: Any
+ ) -> CompletionResponseGen:
if self.model in BEDROCK_FOUNDATION_LLMS and self.model not in STREAMING_MODELS:
raise ValueError(f"Model {self.model} does not support streaming")
- is_formatted = kwargs.pop("formatted", False)
- if not is_formatted:
+
+ if not formatted:
prompt = self.completion_to_prompt(prompt)
+
all_kwargs = self._get_all_kwargs(**kwargs)
request_body = self._provider.get_request_body(prompt, all_kwargs)
request_body_str = json.dumps(request_body)
@@ -247,7 +251,9 @@ class Bedrock(LLM):
# TODO: do synchronous chat for now
return self.chat(messages, **kwargs)
- async def acomplete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
+ async def acomplete(
+ self, prompt: str, formatted: bool = False, **kwargs: Any
+ ) -> CompletionResponse:
raise NotImplementedError
async def astream_chat(
@@ -256,6 +262,6 @@ class Bedrock(LLM):
raise NotImplementedError
async def astream_complete(
- self, prompt: str, **kwargs: Any
+ self, prompt: str, formatted: bool = False, **kwargs: Any
) -> CompletionResponseAsyncGen:
raise NotImplementedError
diff --git a/llama_index/llms/clarifai.py b/llama_index/llms/clarifai.py
index 88950cc0fc..77334fa153 100644
--- a/llama_index/llms/clarifai.py
+++ b/llama_index/llms/clarifai.py
@@ -149,7 +149,11 @@ class Clarifai(LLM):
return ChatResponse(message=ChatMessage(content=response))
def complete(
- self, prompt: str, inference_params: Optional[Dict] = {}, **kwargs: Any
+ self,
+ prompt: str,
+ formatted: bool = False,
+ inference_params: Optional[Dict] = {},
+ **kwargs: Any,
) -> CompletionResponse:
"""Completion endpoint for LLM."""
try:
@@ -173,7 +177,9 @@ class Clarifai(LLM):
"Clarifai does not currently support streaming completion."
)
- def stream_complete(self, prompt: str, **kwargs: Any) -> CompletionResponseGen:
+ def stream_complete(
+ self, prompt: str, formatted: bool = False, **kwargs: Any
+ ) -> CompletionResponseGen:
raise NotImplementedError(
"Clarifai does not currently support streaming completion."
)
@@ -185,7 +191,9 @@ class Clarifai(LLM):
raise NotImplementedError("Currently not supported.")
@llm_completion_callback()
- async def acomplete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
+ async def acomplete(
+ self, prompt: str, formatted: bool = False, **kwargs: Any
+ ) -> CompletionResponse:
return self.complete(prompt, **kwargs)
@llm_chat_callback()
@@ -196,6 +204,6 @@ class Clarifai(LLM):
@llm_completion_callback()
async def astream_complete(
- self, prompt: str, **kwargs: Any
+ self, prompt: str, formatted: bool = False, **kwargs: Any
) -> CompletionResponseAsyncGen:
raise NotImplementedError("Clarifai does not currently support this function.")
diff --git a/llama_index/llms/cohere.py b/llama_index/llms/cohere.py
index d83d6a39e2..0d9fd38448 100644
--- a/llama_index/llms/cohere.py
+++ b/llama_index/llms/cohere.py
@@ -145,7 +145,9 @@ class Cohere(LLM):
)
@llm_completion_callback()
- def complete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
+ def complete(
+ self, prompt: str, formatted: bool = False, **kwargs: Any
+ ) -> CompletionResponse:
all_kwargs = self._get_all_kwargs(**kwargs)
if "stream" in all_kwargs:
warnings.warn(
@@ -203,7 +205,9 @@ class Cohere(LLM):
return gen()
@llm_completion_callback()
- def stream_complete(self, prompt: str, **kwargs: Any) -> CompletionResponseGen:
+ def stream_complete(
+ self, prompt: str, formatted: bool = False, **kwargs: Any
+ ) -> CompletionResponseGen:
all_kwargs = self._get_all_kwargs(**kwargs)
all_kwargs["stream"] = True
@@ -256,7 +260,9 @@ class Cohere(LLM):
)
@llm_completion_callback()
- async def acomplete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
+ async def acomplete(
+ self, prompt: str, formatted: bool = False, **kwargs: Any
+ ) -> CompletionResponse:
all_kwargs = self._get_all_kwargs(**kwargs)
if "stream" in all_kwargs:
warnings.warn(
@@ -315,7 +321,7 @@ class Cohere(LLM):
@llm_completion_callback()
async def astream_complete(
- self, prompt: str, **kwargs: Any
+ self, prompt: str, formatted: bool = False, **kwargs: Any
) -> CompletionResponseAsyncGen:
all_kwargs = self._get_all_kwargs(**kwargs)
all_kwargs["stream"] = True
diff --git a/llama_index/llms/custom.py b/llama_index/llms/custom.py
index 516a5e08c5..76b7bea39b 100644
--- a/llama_index/llms/custom.py
+++ b/llama_index/llms/custom.py
@@ -13,8 +13,8 @@ from llama_index.llms.base import (
llm_completion_callback,
)
from llama_index.llms.generic_utils import (
- completion_to_chat_decorator,
- stream_completion_to_chat_decorator,
+ completion_response_to_chat_response,
+ stream_completion_response_to_chat_response,
)
from llama_index.llms.llm import LLM
@@ -28,15 +28,17 @@ class CustomLLM(LLM):
@llm_chat_callback()
def chat(self, messages: Sequence[ChatMessage], **kwargs: Any) -> ChatResponse:
- chat_fn = completion_to_chat_decorator(self.complete)
- return chat_fn(messages, **kwargs)
+ prompt = self.messages_to_prompt(messages)
+ completion_response = self.complete(prompt, formatted=True, **kwargs)
+ return completion_response_to_chat_response(completion_response)
@llm_chat_callback()
def stream_chat(
self, messages: Sequence[ChatMessage], **kwargs: Any
) -> ChatResponseGen:
- stream_chat_fn = stream_completion_to_chat_decorator(self.stream_complete)
- return stream_chat_fn(messages, **kwargs)
+ prompt = self.messages_to_prompt(messages)
+ completion_response_gen = self.stream_complete(prompt, formatted=True, **kwargs)
+ return stream_completion_response_to_chat_response(completion_response_gen)
@llm_chat_callback()
async def achat(
@@ -60,15 +62,17 @@ class CustomLLM(LLM):
return gen()
@llm_completion_callback()
- async def acomplete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
- return self.complete(prompt, **kwargs)
+ async def acomplete(
+ self, prompt: str, formatted: bool = False, **kwargs: Any
+ ) -> CompletionResponse:
+ return self.complete(prompt, formatted=formatted, **kwargs)
@llm_completion_callback()
async def astream_complete(
- self, prompt: str, **kwargs: Any
+ self, prompt: str, formatted: bool = False, **kwargs: Any
) -> CompletionResponseAsyncGen:
async def gen() -> CompletionResponseAsyncGen:
- for message in self.stream_complete(prompt, **kwargs):
+ for message in self.stream_complete(prompt, formatted=formatted, **kwargs):
yield message
# NOTE: convert generator to async generator
diff --git a/llama_index/llms/gemini.py b/llama_index/llms/gemini.py
index 57eaa8d1db..e000612524 100644
--- a/llama_index/llms/gemini.py
+++ b/llama_index/llms/gemini.py
@@ -138,11 +138,15 @@ class Gemini(CustomLLM):
)
@llm_completion_callback()
- def complete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
+ def complete(
+ self, prompt: str, formatted: bool = False, **kwargs: Any
+ ) -> CompletionResponse:
result = self._model.generate_content(prompt, **kwargs)
return completion_from_gemini_response(result)
- def stream_complete(self, prompt: str, **kwargs: Any) -> CompletionResponseGen:
+ def stream_complete(
+ self, prompt: str, formatted: bool = False, **kwargs: Any
+ ) -> CompletionResponseGen:
it = self._model.generate_content(prompt, stream=True, **kwargs)
yield from map(completion_from_gemini_response, it)
diff --git a/llama_index/llms/gradient.py b/llama_index/llms/gradient.py
index 9928590587..16265dca99 100644
--- a/llama_index/llms/gradient.py
+++ b/llama_index/llms/gradient.py
@@ -89,7 +89,9 @@ class _BaseGradientLLM(CustomLLM):
@llm_completion_callback()
@override
- def complete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
+ def complete(
+ self, prompt: str, formatted: bool = False, **kwargs: Any
+ ) -> CompletionResponse:
return CompletionResponse(
text=self._model.complete(
query=prompt,
@@ -100,7 +102,9 @@ class _BaseGradientLLM(CustomLLM):
@llm_completion_callback()
@override
- async def acomplete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
+ async def acomplete(
+ self, prompt: str, formatted: bool = False, **kwargs: Any
+ ) -> CompletionResponse:
grdt_reponse = await self._model.acomplete(
query=prompt,
max_generated_token_count=self.max_tokens,
@@ -113,6 +117,7 @@ class _BaseGradientLLM(CustomLLM):
def stream_complete(
self,
prompt: str,
+ formatted: bool = False,
**kwargs: Any,
) -> CompletionResponseGen:
raise NotImplementedError
diff --git a/llama_index/llms/huggingface.py b/llama_index/llms/huggingface.py
index bb338a1f25..0be189e1ff 100644
--- a/llama_index/llms/huggingface.py
+++ b/llama_index/llms/huggingface.py
@@ -273,11 +273,12 @@ class HuggingFaceLLM(CustomLLM):
return generic_messages_to_prompt(messages)
@llm_completion_callback()
- def complete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
+ def complete(
+ self, prompt: str, formatted: bool = False, **kwargs: Any
+ ) -> CompletionResponse:
"""Completion endpoint."""
full_prompt = prompt
- is_formatted = kwargs.pop("formatted", False)
- if not is_formatted:
+ if not formatted:
if self.query_wrapper_prompt:
full_prompt = self.query_wrapper_prompt.format(query_str=prompt)
if self.system_prompt:
@@ -303,13 +304,14 @@ class HuggingFaceLLM(CustomLLM):
return CompletionResponse(text=completion, raw={"model_output": tokens})
@llm_completion_callback()
- def stream_complete(self, prompt: str, **kwargs: Any) -> CompletionResponseGen:
+ def stream_complete(
+ self, prompt: str, formatted: bool = False, **kwargs: Any
+ ) -> CompletionResponseGen:
"""Streaming completion endpoint."""
from transformers import TextIteratorStreamer
full_prompt = prompt
- is_formatted = kwargs.pop("formatted", False)
- if not is_formatted:
+ if not formatted:
if self.query_wrapper_prompt:
full_prompt = self.query_wrapper_prompt.format(query_str=prompt)
if self.system_prompt:
@@ -576,7 +578,9 @@ class HuggingFaceInferenceAPI(CustomLLM):
)
)
- def complete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
+ def complete(
+ self, prompt: str, formatted: bool = False, **kwargs: Any
+ ) -> CompletionResponse:
return CompletionResponse(
text=self._sync_client.text_generation(
prompt, **{**{"max_new_tokens": self.num_output}, **kwargs}
@@ -588,7 +592,9 @@ class HuggingFaceInferenceAPI(CustomLLM):
) -> ChatResponseGen:
raise NotImplementedError
- def stream_complete(self, prompt: str, **kwargs: Any) -> CompletionResponseGen:
+ def stream_complete(
+ self, prompt: str, formatted: bool = False, **kwargs: Any
+ ) -> CompletionResponseGen:
raise NotImplementedError
async def achat(
@@ -596,7 +602,9 @@ class HuggingFaceInferenceAPI(CustomLLM):
) -> ChatResponse:
raise NotImplementedError
- async def acomplete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
+ async def acomplete(
+ self, prompt: str, formatted: bool = False, **kwargs: Any
+ ) -> CompletionResponse:
response = await self._async_client.text_generation(
prompt, **{**{"max_new_tokens": self.num_output}, **kwargs}
)
@@ -608,6 +616,6 @@ class HuggingFaceInferenceAPI(CustomLLM):
raise NotImplementedError
async def astream_complete(
- self, prompt: str, **kwargs: Any
+ self, prompt: str, formatted: bool = False, **kwargs: Any
) -> CompletionResponseAsyncGen:
raise NotImplementedError
diff --git a/llama_index/llms/konko.py b/llama_index/llms/konko.py
index 3ab7cd2bdb..487e1841a7 100644
--- a/llama_index/llms/konko.py
+++ b/llama_index/llms/konko.py
@@ -277,7 +277,9 @@ class Konko(LLM):
return gen()
@llm_completion_callback()
- def complete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
+ def complete(
+ self, prompt: str, formatted: bool = False, **kwargs: Any
+ ) -> CompletionResponse:
if self._is_chat_model:
complete_fn = chat_to_completion_decorator(self._chat)
else:
@@ -285,7 +287,9 @@ class Konko(LLM):
return complete_fn(prompt, **kwargs)
@llm_completion_callback()
- def stream_complete(self, prompt: str, **kwargs: Any) -> CompletionResponseGen:
+ def stream_complete(
+ self, prompt: str, formatted: bool = False, **kwargs: Any
+ ) -> CompletionResponseGen:
if self._is_chat_model:
stream_complete_fn = stream_chat_to_completion_decorator(self._stream_chat)
else:
@@ -413,7 +417,9 @@ class Konko(LLM):
return await astream_chat_fn(messages, **kwargs)
@llm_completion_callback()
- async def acomplete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
+ async def acomplete(
+ self, prompt: str, formatted: bool = False, **kwargs: Any
+ ) -> CompletionResponse:
if self._is_chat_model:
acomplete_fn = achat_to_completion_decorator(self._achat)
else:
@@ -422,7 +428,7 @@ class Konko(LLM):
@llm_completion_callback()
async def astream_complete(
- self, prompt: str, **kwargs: Any
+ self, prompt: str, formatted: bool = False, **kwargs: Any
) -> CompletionResponseAsyncGen:
if self._is_chat_model:
astream_complete_fn = astream_chat_to_completion_decorator(
diff --git a/llama_index/llms/langchain.py b/llama_index/llms/langchain.py
index 873ee9ab29..904005c7dc 100644
--- a/llama_index/llms/langchain.py
+++ b/llama_index/llms/langchain.py
@@ -73,7 +73,7 @@ class LangChainLLM(LLM):
if not self.metadata.is_chat_model:
prompt = self.messages_to_prompt(messages)
- completion_response = self.complete(prompt, is_formatted=True, **kwargs)
+ completion_response = self.complete(prompt, formatted=True, **kwargs)
return completion_response_to_chat_response(completion_response)
lc_messages = to_lc_messages(messages)
@@ -83,9 +83,9 @@ class LangChainLLM(LLM):
@llm_completion_callback()
def complete(
- self, prompt: str, is_formatted: bool = False, **kwargs: Any
+ self, prompt: str, formatted: bool = False, **kwargs: Any
) -> CompletionResponse:
- if not is_formatted:
+ if not formatted:
prompt = self.completion_to_prompt(prompt)
output_str = self._llm.predict(prompt, **kwargs)
@@ -97,9 +97,7 @@ class LangChainLLM(LLM):
) -> ChatResponseGen:
if not self.metadata.is_chat_model:
prompt = self.messages_to_prompt(messages)
- stream_completion = self.stream_complete(
- prompt, is_formatted=True, **kwargs
- )
+ stream_completion = self.stream_complete(prompt, formatted=True, **kwargs)
return stream_completion_response_to_chat_response(stream_completion)
from llama_index.langchain_helpers.streaming import (
@@ -134,9 +132,9 @@ class LangChainLLM(LLM):
@llm_completion_callback()
def stream_complete(
- self, prompt: str, is_formatted: bool = False, **kwargs: Any
+ self, prompt: str, formatted: bool = False, **kwargs: Any
) -> CompletionResponseGen:
- if not is_formatted:
+ if not formatted:
prompt = self.completion_to_prompt(prompt)
from llama_index.langchain_helpers.streaming import (
@@ -174,9 +172,11 @@ class LangChainLLM(LLM):
return self.chat(messages, **kwargs)
@llm_completion_callback()
- async def acomplete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
+ async def acomplete(
+ self, prompt: str, formatted: bool = False, **kwargs: Any
+ ) -> CompletionResponse:
# TODO: Implement async complete
- return self.complete(prompt, **kwargs)
+ return self.complete(prompt, formatted=formatted, **kwargs)
@llm_chat_callback()
async def astream_chat(
@@ -192,12 +192,12 @@ class LangChainLLM(LLM):
@llm_completion_callback()
async def astream_complete(
- self, prompt: str, **kwargs: Any
+ self, prompt: str, formatted: bool = False, **kwargs: Any
) -> CompletionResponseAsyncGen:
# TODO: Implement async stream_complete
async def gen() -> CompletionResponseAsyncGen:
- for response in self.stream_complete(prompt, **kwargs):
+ for response in self.stream_complete(prompt, formatted=formatted, **kwargs):
yield response
return gen()
diff --git a/llama_index/llms/litellm.py b/llama_index/llms/litellm.py
index 4ddfea7386..a3fafb7019 100644
--- a/llama_index/llms/litellm.py
+++ b/llama_index/llms/litellm.py
@@ -158,7 +158,9 @@ class LiteLLM(LLM):
return stream_chat_fn(messages, **kwargs)
@llm_completion_callback()
- def complete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
+ def complete(
+ self, prompt: str, formatted: bool = False, **kwargs: Any
+ ) -> CompletionResponse:
# litellm assumes all llms are chat llms
if self._is_chat_model:
complete_fn = chat_to_completion_decorator(self._chat)
@@ -168,7 +170,9 @@ class LiteLLM(LLM):
return complete_fn(prompt, **kwargs)
@llm_completion_callback()
- def stream_complete(self, prompt: str, **kwargs: Any) -> CompletionResponseGen:
+ def stream_complete(
+ self, prompt: str, formatted: bool = False, **kwargs: Any
+ ) -> CompletionResponseGen:
if self._is_chat_model:
stream_complete_fn = stream_chat_to_completion_decorator(self._stream_chat)
else:
@@ -352,7 +356,9 @@ class LiteLLM(LLM):
return await astream_chat_fn(messages, **kwargs)
@llm_completion_callback()
- async def acomplete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
+ async def acomplete(
+ self, prompt: str, formatted: bool = False, **kwargs: Any
+ ) -> CompletionResponse:
if self._is_chat_model:
acomplete_fn = achat_to_completion_decorator(self._achat)
else:
@@ -361,7 +367,7 @@ class LiteLLM(LLM):
@llm_completion_callback()
async def astream_complete(
- self, prompt: str, **kwargs: Any
+ self, prompt: str, formatted: bool = False, **kwargs: Any
) -> CompletionResponseAsyncGen:
if self._is_chat_model:
astream_complete_fn = astream_chat_to_completion_decorator(
diff --git a/llama_index/llms/llama_api.py b/llama_index/llms/llama_api.py
index 1364a4a797..1ade74bbea 100644
--- a/llama_index/llms/llama_api.py
+++ b/llama_index/llms/llama_api.py
@@ -109,12 +109,16 @@ class LlamaAPI(CustomLLM):
return ChatResponse(message=message, raw=response)
@llm_completion_callback()
- def complete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
+ def complete(
+ self, prompt: str, formatted: bool = False, **kwargs: Any
+ ) -> CompletionResponse:
complete_fn = chat_to_completion_decorator(self.chat)
return complete_fn(prompt, **kwargs)
@llm_completion_callback()
- def stream_complete(self, prompt: str, **kwargs: Any) -> CompletionResponseGen:
+ def stream_complete(
+ self, prompt: str, formatted: bool = False, **kwargs: Any
+ ) -> CompletionResponseGen:
raise NotImplementedError("stream_complete is not supported for LlamaAPI")
@llm_chat_callback()
diff --git a/llama_index/llms/llama_cpp.py b/llama_index/llms/llama_cpp.py
index 124554c92a..f090789073 100644
--- a/llama_index/llms/llama_cpp.py
+++ b/llama_index/llms/llama_cpp.py
@@ -221,11 +221,12 @@ class LlamaCPP(CustomLLM):
return stream_completion_response_to_chat_response(completion_response)
@llm_completion_callback()
- def complete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
+ def complete(
+ self, prompt: str, formatted: bool = False, **kwargs: Any
+ ) -> CompletionResponse:
self.generate_kwargs.update({"stream": False})
- is_formatted = kwargs.pop("formatted", False)
- if not is_formatted:
+ if not formatted:
prompt = self.completion_to_prompt(prompt)
response = self._model(prompt=prompt, **self.generate_kwargs)
@@ -233,11 +234,12 @@ class LlamaCPP(CustomLLM):
return CompletionResponse(text=response["choices"][0]["text"], raw=response)
@llm_completion_callback()
- def stream_complete(self, prompt: str, **kwargs: Any) -> CompletionResponseGen:
+ def stream_complete(
+ self, prompt: str, formatted: bool = False, **kwargs: Any
+ ) -> CompletionResponseGen:
self.generate_kwargs.update({"stream": True})
- is_formatted = kwargs.pop("formatted", False)
- if not is_formatted:
+ if not formatted:
prompt = self.completion_to_prompt(prompt)
response_iter = self._model(prompt=prompt, **self.generate_kwargs)
diff --git a/llama_index/llms/llm.py b/llama_index/llms/llm.py
index 850d3340c3..d200d8af59 100644
--- a/llama_index/llms/llm.py
+++ b/llama_index/llms/llm.py
@@ -1,5 +1,14 @@
from collections import ChainMap
-from typing import Any, List, Optional, Protocol, Sequence, runtime_checkable
+from typing import (
+ Any,
+ Dict,
+ List,
+ Optional,
+ Protocol,
+ Sequence,
+ get_args,
+ runtime_checkable,
+)
from llama_index.bridge.pydantic import BaseModel, Field, validator
from llama_index.callbacks import CBEventType, EventPayload
@@ -11,10 +20,20 @@ from llama_index.core.llms.types import (
CompletionResponseGen,
MessageRole,
)
+from llama_index.core.query_pipeline.query_component import (
+ InputKeys,
+ OutputKeys,
+ QueryComponent,
+ StringableInput,
+ validate_and_convert_stringable,
+)
from llama_index.llms.base import BaseLLM
from llama_index.llms.generic_utils import (
messages_to_prompt as generic_messages_to_prompt,
)
+from llama_index.llms.generic_utils import (
+ prompt_to_messages,
+)
from llama_index.prompts import BasePromptTemplate, PromptTemplate
from llama_index.types import (
BaseOutputParser,
@@ -221,7 +240,7 @@ class LLM(BaseLLM):
output = chat_response.message.content or ""
else:
formatted_prompt = self._get_prompt(prompt, **prompt_args)
- response = self.complete(formatted_prompt)
+ response = self.complete(formatted_prompt, formatted=True)
output = response.text
return self._parse_output(output)
@@ -240,7 +259,7 @@ class LLM(BaseLLM):
stream_tokens = stream_chat_response_to_tokens(chat_response)
else:
formatted_prompt = self._get_prompt(prompt, **prompt_args)
- stream_response = self.stream_complete(formatted_prompt)
+ stream_response = self.stream_complete(formatted_prompt, formatted=True)
stream_tokens = stream_completion_response_to_tokens(stream_response)
if prompt.output_parser is not None or self.output_parser is not None:
@@ -262,7 +281,7 @@ class LLM(BaseLLM):
output = chat_response.message.content or ""
else:
formatted_prompt = self._get_prompt(prompt, **prompt_args)
- response = await self.acomplete(formatted_prompt)
+ response = await self.acomplete(formatted_prompt, formatted=True)
output = response.text
return self._parse_output(output)
@@ -281,7 +300,9 @@ class LLM(BaseLLM):
stream_tokens = await astream_chat_response_to_tokens(chat_response)
else:
formatted_prompt = self._get_prompt(prompt, **prompt_args)
- stream_response = await self.astream_complete(formatted_prompt)
+ stream_response = await self.astream_complete(
+ formatted_prompt, formatted=True
+ )
stream_tokens = await astream_completion_response_to_tokens(stream_response)
if prompt.output_parser is not None or self.output_parser is not None:
@@ -314,3 +335,120 @@ class LLM(BaseLLM):
*messages,
]
return messages
+
+ def _as_query_component(self, **kwargs: Any) -> QueryComponent:
+ """Return query component."""
+ if self.metadata.is_chat_model:
+ return LLMChatComponent(llm=self)
+ else:
+ return LLMCompleteComponent(llm=self)
+
+
+class LLMCompleteComponent(QueryComponent):
+ """LLM completion component."""
+
+ llm: LLM = Field(..., description="LLM")
+
+ class Config:
+ arbitrary_types_allowed = True
+
+ def set_callback_manager(self, callback_manager: Any) -> None:
+ """Set callback manager."""
+ self.llm.callback_manager = callback_manager
+
+ def _validate_component_inputs(self, input: Dict[str, Any]) -> Dict[str, Any]:
+ """Validate component inputs during run_component."""
+ if "prompt" not in input:
+ raise ValueError("Prompt must be in input dict.")
+ # do special check to see if prompt is a list of chat messages
+ if isinstance(input["prompt"], get_args(List[ChatMessage])):
+ input["prompt"] = self.llm.messages_to_prompt(input["prompt"])
+ input["prompt"] = validate_and_convert_stringable(input["prompt"])
+ else:
+ input["prompt"] = validate_and_convert_stringable(input["prompt"])
+ input["prompt"] = self.llm.completion_to_prompt(input["prompt"])
+
+ return input
+
+ def _run_component(self, **kwargs: Any) -> Any:
+ """Run component."""
+ # TODO: support only complete for now
+ # non-trivial to figure how to support chat/complete/etc.
+ prompt = kwargs["prompt"]
+ # ignore all other kwargs for now
+ response = self.llm.complete(prompt, formatted=True)
+ return {"output": response}
+
+ async def _arun_component(self, **kwargs: Any) -> Any:
+ """Run component."""
+ # TODO: support only complete for now
+ # non-trivial to figure how to support chat/complete/etc.
+ prompt = kwargs["prompt"]
+ # ignore all other kwargs for now
+ response = await self.llm.acomplete(prompt, formatted=True)
+ return {"output": response}
+
+ @property
+ def input_keys(self) -> InputKeys:
+ """Input keys."""
+ # TODO: support only complete for now
+ return InputKeys.from_keys({"prompt"})
+
+ @property
+ def output_keys(self) -> OutputKeys:
+ """Output keys."""
+ return OutputKeys.from_keys({"output"})
+
+
+class LLMChatComponent(QueryComponent):
+ """LLM chat component."""
+
+ llm: LLM = Field(..., description="LLM")
+
+ class Config:
+ arbitrary_types_allowed = True
+
+ def set_callback_manager(self, callback_manager: Any) -> None:
+ """Set callback manager."""
+ self.llm.callback_manager = callback_manager
+
+ def _validate_component_inputs(self, input: Dict[str, Any]) -> Dict[str, Any]:
+ """Validate component inputs during run_component."""
+ if "messages" not in input:
+ raise ValueError("Messages must be in input dict.")
+
+ # if `messages` is a string, convert to a list of chat message
+ if isinstance(input["messages"], get_args(StringableInput)):
+ input["messages"] = prompt_to_messages(str(input["messages"]))
+
+ for message in input["messages"]:
+ if not isinstance(message, ChatMessage):
+ raise ValueError("Messages must be a list of ChatMessage")
+ return input
+
+ def _run_component(self, **kwargs: Any) -> Any:
+ """Run component."""
+ # TODO: support only complete for now
+ # non-trivial to figure how to support chat/complete/etc.
+ messages = kwargs["messages"]
+ response = self.llm.chat(messages)
+ return {"output": response}
+
+ async def _arun_component(self, **kwargs: Any) -> Any:
+ """Run component."""
+ # TODO: support only complete for now
+ # non-trivial to figure how to support chat/complete/etc.
+ messages = kwargs["messages"]
+ response = await self.llm.achat(messages)
+ return {"output": response}
+
+ @property
+ def input_keys(self) -> InputKeys:
+ """Input keys."""
+ # TODO: support only complete for now
+ return InputKeys.from_keys({"messages"})
+
+ @property
+ def output_keys(self) -> OutputKeys:
+ """Output keys."""
+ return OutputKeys.from_keys({"output"})
diff --git a/llama_index/llms/mistral.py b/llama_index/llms/mistral.py
index e461525376..1c2d480217 100644
--- a/llama_index/llms/mistral.py
+++ b/llama_index/llms/mistral.py
@@ -195,7 +195,9 @@ class MistralAI(LLM):
)
@llm_completion_callback()
- def complete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
+ def complete(
+ self, prompt: str, formatted: bool = False, **kwargs: Any
+ ) -> CompletionResponse:
complete_fn = chat_to_completion_decorator(self.chat)
return complete_fn(prompt, **kwargs)
@@ -231,7 +233,9 @@ class MistralAI(LLM):
return gen()
@llm_completion_callback()
- def stream_complete(self, prompt: str, **kwargs: Any) -> CompletionResponseGen:
+ def stream_complete(
+ self, prompt: str, formatted: bool = False, **kwargs: Any
+ ) -> CompletionResponseGen:
stream_complete_fn = stream_chat_to_completion_decorator(self.stream_chat)
return stream_complete_fn(prompt, **kwargs)
@@ -256,7 +260,9 @@ class MistralAI(LLM):
)
@llm_completion_callback()
- async def acomplete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
+ async def acomplete(
+ self, prompt: str, formatted: bool = False, **kwargs: Any
+ ) -> CompletionResponse:
acomplete_fn = achat_to_completion_decorator(self.achat)
return await acomplete_fn(prompt, **kwargs)
@@ -292,7 +298,7 @@ class MistralAI(LLM):
@llm_completion_callback()
async def astream_complete(
- self, prompt: str, **kwargs: Any
+ self, prompt: str, formatted: bool = False, **kwargs: Any
) -> CompletionResponseAsyncGen:
astream_complete_fn = astream_chat_to_completion_decorator(self.astream_chat)
return await astream_complete_fn(prompt, **kwargs)
diff --git a/llama_index/llms/mock.py b/llama_index/llms/mock.py
index 0cce089a8a..2b19d9ad33 100644
--- a/llama_index/llms/mock.py
+++ b/llama_index/llms/mock.py
@@ -45,7 +45,9 @@ class MockLLM(CustomLLM):
return " ".join(["text" for _ in range(length)])
@llm_completion_callback()
- def complete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
+ def complete(
+ self, prompt: str, formatted: bool = False, **kwargs: Any
+ ) -> CompletionResponse:
response_text = (
self._generate_text(self.max_tokens) if self.max_tokens else prompt
)
@@ -55,7 +57,9 @@ class MockLLM(CustomLLM):
)
@llm_completion_callback()
- def stream_complete(self, prompt: str, **kwargs: Any) -> CompletionResponseGen:
+ def stream_complete(
+ self, prompt: str, formatted: bool = False, **kwargs: Any
+ ) -> CompletionResponseGen:
def gen_prompt() -> CompletionResponseGen:
for ch in prompt:
yield CompletionResponse(
diff --git a/llama_index/llms/monsterapi.py b/llama_index/llms/monsterapi.py
index aaa1090e57..8864f19627 100644
--- a/llama_index/llms/monsterapi.py
+++ b/llama_index/llms/monsterapi.py
@@ -120,9 +120,10 @@ class MonsterLLM(CustomLLM):
return self.complete(prompt, formatted=True, **kwargs)
@llm_completion_callback()
- def complete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
- is_formatted = kwargs.pop("formatted", False)
- if not is_formatted:
+ def complete(
+ self, prompt: str, formatted: bool = False, **kwargs: Any
+ ) -> CompletionResponse:
+ if not formatted:
prompt = self.completion_to_prompt(prompt)
# Validate input args against input Pydantic model
@@ -138,5 +139,7 @@ class MonsterLLM(CustomLLM):
return CompletionResponse(text=result["text"])
@llm_completion_callback()
- def stream_complete(self, prompt: str, **kwargs: Any) -> CompletionResponseGen:
+ def stream_complete(
+ self, prompt: str, formatted: bool = False, **kwargs: Any
+ ) -> CompletionResponseGen:
raise NotImplementedError
diff --git a/llama_index/llms/nvidia_triton.py b/llama_index/llms/nvidia_triton.py
index 468892d520..70f2171b9a 100644
--- a/llama_index/llms/nvidia_triton.py
+++ b/llama_index/llms/nvidia_triton.py
@@ -190,7 +190,9 @@ class NvidiaTriton(LLM):
) -> ChatResponseGen:
raise NotImplementedError
- def complete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
+ def complete(
+ self, prompt: str, formatted: bool = False, **kwargs: Any
+ ) -> CompletionResponse:
from tritonclient.utils import InferenceServerException
client = self._get_client()
@@ -220,7 +222,9 @@ class NvidiaTriton(LLM):
text=response,
)
- def stream_complete(self, prompt: str, **kwargs: Any) -> CompletionResponseGen:
+ def stream_complete(
+ self, prompt: str, formatted: bool = False, **kwargs: Any
+ ) -> CompletionResponseGen:
raise NotImplementedError
async def achat(
@@ -228,7 +232,9 @@ class NvidiaTriton(LLM):
) -> ChatResponse:
raise NotImplementedError
- async def acomplete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
+ async def acomplete(
+ self, prompt: str, formatted: bool = False, **kwargs: Any
+ ) -> CompletionResponse:
raise NotImplementedError
async def astream_chat(
@@ -237,6 +243,6 @@ class NvidiaTriton(LLM):
raise NotImplementedError
async def astream_complete(
- self, prompt: str, **kwargs: Any
+ self, prompt: str, formatted: bool = False, **kwargs: Any
) -> CompletionResponseAsyncGen:
raise NotImplementedError
diff --git a/llama_index/llms/ollama.py b/llama_index/llms/ollama.py
index 1801c49b46..409dd3aec9 100644
--- a/llama_index/llms/ollama.py
+++ b/llama_index/llms/ollama.py
@@ -165,7 +165,9 @@ class Ollama(CustomLLM):
)
@llm_completion_callback()
- def complete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
+ def complete(
+ self, prompt: str, formatted: bool = False, **kwargs: Any
+ ) -> CompletionResponse:
payload = {
self.prompt_key: prompt,
"model": self.model,
@@ -189,7 +191,9 @@ class Ollama(CustomLLM):
)
@llm_completion_callback()
- def stream_complete(self, prompt: str, **kwargs: Any) -> CompletionResponseGen:
+ def stream_complete(
+ self, prompt: str, formatted: bool = False, **kwargs: Any
+ ) -> CompletionResponseGen:
payload = {
self.prompt_key: prompt,
"model": self.model,
diff --git a/llama_index/llms/openai.py b/llama_index/llms/openai.py
index 738e11b8d1..b5c544033c 100644
--- a/llama_index/llms/openai.py
+++ b/llama_index/llms/openai.py
@@ -247,7 +247,9 @@ class OpenAI(LLM):
return stream_chat_fn(messages, **kwargs)
@llm_completion_callback()
- def complete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
+ def complete(
+ self, prompt: str, formatted: bool = False, **kwargs: Any
+ ) -> CompletionResponse:
if self._use_chat_completions(kwargs):
complete_fn = chat_to_completion_decorator(self._chat)
else:
@@ -255,7 +257,9 @@ class OpenAI(LLM):
return complete_fn(prompt, **kwargs)
@llm_completion_callback()
- def stream_complete(self, prompt: str, **kwargs: Any) -> CompletionResponseGen:
+ def stream_complete(
+ self, prompt: str, formatted: bool = False, **kwargs: Any
+ ) -> CompletionResponseGen:
if self._use_chat_completions(kwargs):
stream_complete_fn = stream_chat_to_completion_decorator(self._stream_chat)
else:
@@ -505,7 +509,9 @@ class OpenAI(LLM):
return await astream_chat_fn(messages, **kwargs)
@llm_completion_callback()
- async def acomplete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
+ async def acomplete(
+ self, prompt: str, formatted: bool = False, **kwargs: Any
+ ) -> CompletionResponse:
if self._use_chat_completions(kwargs):
acomplete_fn = achat_to_completion_decorator(self._achat)
else:
@@ -514,7 +520,7 @@ class OpenAI(LLM):
@llm_completion_callback()
async def astream_complete(
- self, prompt: str, **kwargs: Any
+ self, prompt: str, formatted: bool = False, **kwargs: Any
) -> CompletionResponseAsyncGen:
if self._use_chat_completions(kwargs):
astream_complete_fn = astream_chat_to_completion_decorator(
diff --git a/llama_index/llms/openllm.py b/llama_index/llms/openllm.py
index dea6d54d05..7f7d311892 100644
--- a/llama_index/llms/openllm.py
+++ b/llama_index/llms/openllm.py
@@ -167,7 +167,9 @@ class OpenLLM(LLM):
return generic_messages_to_prompt(messages)
@llm_completion_callback()
- def complete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
+ def complete(
+ self, prompt: str, formatted: bool = False, **kwargs: Any
+ ) -> CompletionResponse:
return asyncio.run(self.acomplete(prompt, **kwargs))
@llm_chat_callback()
@@ -183,7 +185,9 @@ class OpenLLM(LLM):
return loop
@llm_completion_callback()
- def stream_complete(self, prompt: str, **kwargs: Any) -> CompletionResponseGen:
+ def stream_complete(
+ self, prompt: str, formatted: bool = False, **kwargs: Any
+ ) -> CompletionResponseGen:
generator = self.astream_complete(prompt, **kwargs)
# Yield items from the queue synchronously
while True:
@@ -214,7 +218,9 @@ class OpenLLM(LLM):
return completion_response_to_chat_response(response)
@llm_completion_callback()
- async def acomplete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
+ async def acomplete(
+ self, prompt: str, formatted: bool = False, **kwargs: Any
+ ) -> CompletionResponse:
response = await self._llm.generate(prompt, **kwargs)
return CompletionResponse(
text=response.outputs[0].text,
@@ -245,7 +251,7 @@ class OpenLLM(LLM):
@llm_completion_callback()
async def astream_complete(
- self, prompt: str, **kwargs: Any
+ self, prompt: str, formatted: bool = False, **kwargs: Any
) -> CompletionResponseAsyncGen:
config = self._llm.config.model_construct_env(**kwargs)
if config["n"] > 1:
@@ -372,7 +378,9 @@ class OpenLLMAPI(LLM):
)
@llm_completion_callback()
- def complete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
+ def complete(
+ self, prompt: str, formatted: bool = False, **kwargs: Any
+ ) -> CompletionResponse:
response = self._sync_client.generate(prompt, **kwargs)
return CompletionResponse(
text=response.outputs[0].text,
@@ -391,7 +399,9 @@ class OpenLLMAPI(LLM):
)
@llm_completion_callback()
- def stream_complete(self, prompt: str, **kwargs: Any) -> CompletionResponseGen:
+ def stream_complete(
+ self, prompt: str, formatted: bool = False, **kwargs: Any
+ ) -> CompletionResponseGen:
for response_chunk in self._sync_client.generate_stream(prompt, **kwargs):
yield CompletionResponse(
text=response_chunk.text,
@@ -416,7 +426,9 @@ class OpenLLMAPI(LLM):
yield completion_response_to_chat_response(response_chunk)
@llm_completion_callback()
- async def acomplete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
+ async def acomplete(
+ self, prompt: str, formatted: bool = False, **kwargs: Any
+ ) -> CompletionResponse:
response = await self._async_client.generate(prompt, **kwargs)
return CompletionResponse(
text=response.outputs[0].text,
@@ -436,7 +448,7 @@ class OpenLLMAPI(LLM):
@llm_completion_callback()
async def astream_complete(
- self, prompt: str, **kwargs: Any
+ self, prompt: str, formatted: bool = False, **kwargs: Any
) -> CompletionResponseAsyncGen:
async for response_chunk in self._async_client.generate_stream(
prompt, **kwargs
diff --git a/llama_index/llms/palm.py b/llama_index/llms/palm.py
index 1e0200001e..5a236e3dcd 100644
--- a/llama_index/llms/palm.py
+++ b/llama_index/llms/palm.py
@@ -101,7 +101,9 @@ class PaLM(CustomLLM):
)
@llm_completion_callback()
- def complete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
+ def complete(
+ self, prompt: str, formatted: bool = False, **kwargs: Any
+ ) -> CompletionResponse:
"""Predict the answer to a query.
Args:
@@ -121,7 +123,9 @@ class PaLM(CustomLLM):
return CompletionResponse(text=completion.result, raw=completion.candidates[0])
@llm_completion_callback()
- def stream_complete(self, prompt: str, **kwargs: Any) -> CompletionResponseGen:
+ def stream_complete(
+ self, prompt: str, formatted: bool = False, **kwargs: Any
+ ) -> CompletionResponseGen:
"""Stream the answer to a query.
NOTE: this is a beta feature. Will try to build or use
diff --git a/llama_index/llms/perplexity.py b/llama_index/llms/perplexity.py
index 005e010ba3..d11c11ff5a 100644
--- a/llama_index/llms/perplexity.py
+++ b/llama_index/llms/perplexity.py
@@ -159,7 +159,9 @@ class Perplexity(LLM):
return CompletionResponse(text=data["choices"][0]["text"], raw=data)
@llm_completion_callback()
- def complete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
+ def complete(
+ self, prompt: str, formatted: bool = False, **kwargs: Any
+ ) -> CompletionResponse:
if self._is_chat_model():
raise ValueError("The complete method is not supported for chat models.")
return self._complete(prompt, **kwargs)
@@ -199,7 +201,9 @@ class Perplexity(LLM):
return CompletionResponse(text=data["choices"][0]["text"], raw=data)
@llm_completion_callback()
- async def acomplete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
+ async def acomplete(
+ self, prompt: str, formatted: bool = False, **kwargs: Any
+ ) -> CompletionResponse:
if self._is_chat_model():
raise ValueError("The complete method is not supported for chat models.")
return await self._acomplete(prompt, **kwargs)
@@ -258,7 +262,9 @@ class Perplexity(LLM):
return gen()
@llm_completion_callback()
- def stream_complete(self, prompt: str, **kwargs: Any) -> CompletionResponseGen:
+ def stream_complete(
+ self, prompt: str, formatted: bool = False, **kwargs: Any
+ ) -> CompletionResponseGen:
if self._is_chat_model():
raise ValueError("The complete method is not supported for chat models.")
stream_complete_fn = self._stream_complete
@@ -296,7 +302,7 @@ class Perplexity(LLM):
@llm_completion_callback()
async def astream_complete(
- self, prompt: str, **kwargs: Any
+ self, prompt: str, formatted: bool = False, **kwargs: Any
) -> CompletionResponseAsyncGen:
if self._is_chat_model():
raise ValueError("The complete method is not supported for chat models.")
diff --git a/llama_index/llms/portkey.py b/llama_index/llms/portkey.py
index 1c1f1ba234..e79ea527be 100644
--- a/llama_index/llms/portkey.py
+++ b/llama_index/llms/portkey.py
@@ -152,7 +152,9 @@ class Portkey(CustomLLM):
return self
@llm_completion_callback()
- def complete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
+ def complete(
+ self, prompt: str, formatted: bool = False, **kwargs: Any
+ ) -> CompletionResponse:
"""Completion endpoint for LLM."""
if self._is_chat_model:
complete_fn = chat_to_completion_decorator(self._chat)
@@ -169,7 +171,9 @@ class Portkey(CustomLLM):
return chat_fn(messages, **kwargs)
@llm_completion_callback()
- def stream_complete(self, prompt: str, **kwargs: Any) -> CompletionResponseGen:
+ def stream_complete(
+ self, prompt: str, formatted: bool = False, **kwargs: Any
+ ) -> CompletionResponseGen:
"""Completion endpoint for LLM."""
if self._is_chat_model:
complete_fn = stream_chat_to_completion_decorator(self._stream_chat)
diff --git a/llama_index/llms/predibase.py b/llama_index/llms/predibase.py
index cca2997a31..7da49f64c0 100644
--- a/llama_index/llms/predibase.py
+++ b/llama_index/llms/predibase.py
@@ -108,7 +108,9 @@ class PredibaseLLM(CustomLLM):
)
@llm_completion_callback()
- def complete(self, prompt: str, **kwargs: Any) -> "CompletionResponse":
+ def complete(
+ self, prompt: str, formatted: bool = False, **kwargs: Any
+ ) -> "CompletionResponse":
llm = self._client.LLM(f"pb://deployments/{self.model_name}")
results = llm.prompt(
prompt, max_new_tokens=self.max_new_tokens, temperature=self.temperature
@@ -116,5 +118,7 @@ class PredibaseLLM(CustomLLM):
return CompletionResponse(text=results.response)
@llm_completion_callback()
- def stream_complete(self, prompt: str, **kwargs: Any) -> "CompletionResponseGen":
+ def stream_complete(
+ self, prompt: str, formatted: bool = False, **kwargs: Any
+ ) -> "CompletionResponseGen":
raise NotImplementedError
diff --git a/llama_index/llms/replicate.py b/llama_index/llms/replicate.py
index bfbd95eaba..0d04da807f 100644
--- a/llama_index/llms/replicate.py
+++ b/llama_index/llms/replicate.py
@@ -96,15 +96,19 @@ class Replicate(CustomLLM):
return stream_completion_response_to_chat_response(completion_response)
@llm_completion_callback()
- def complete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
- response_gen = self.stream_complete(prompt, **kwargs)
+ def complete(
+ self, prompt: str, formatted: bool = False, **kwargs: Any
+ ) -> CompletionResponse:
+ response_gen = self.stream_complete(prompt, formatted=formatted, **kwargs)
response_list = list(response_gen)
final_response = response_list[-1]
final_response.delta = None
return final_response
@llm_completion_callback()
- def stream_complete(self, prompt: str, **kwargs: Any) -> CompletionResponseGen:
+ def stream_complete(
+ self, prompt: str, formatted: bool = False, **kwargs: Any
+ ) -> CompletionResponseGen:
try:
import replicate
except ImportError:
@@ -113,7 +117,7 @@ class Replicate(CustomLLM):
"Please install replicate with `pip install replicate`"
)
- if not kwargs.get("formatted", False):
+ if not formatted:
prompt = self.completion_to_prompt(prompt)
input_dict = self._get_input_dict(prompt, **kwargs)
response_iter = replicate.run(self.model, input=input_dict)
diff --git a/llama_index/llms/rungpt.py b/llama_index/llms/rungpt.py
index e0296ac1f8..9819bf305f 100644
--- a/llama_index/llms/rungpt.py
+++ b/llama_index/llms/rungpt.py
@@ -102,7 +102,9 @@ class RunGptLLM(LLM):
)
@llm_completion_callback()
- def complete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
+ def complete(
+ self, prompt: str, formatted: bool = False, **kwargs: Any
+ ) -> CompletionResponse:
try:
import requests
except ImportError:
@@ -123,7 +125,9 @@ class RunGptLLM(LLM):
)
@llm_completion_callback()
- def stream_complete(self, prompt: str, **kwargs: Any) -> CompletionResponseGen:
+ def stream_complete(
+ self, prompt: str, formatted: bool = False, **kwargs: Any
+ ) -> CompletionResponseGen:
try:
import requests
except ImportError:
@@ -240,12 +244,14 @@ class RunGptLLM(LLM):
return gen()
@llm_completion_callback()
- async def acomplete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
+ async def acomplete(
+ self, prompt: str, formatted: bool = False, **kwargs: Any
+ ) -> CompletionResponse:
return self.complete(prompt, **kwargs)
@llm_completion_callback()
async def astream_complete(
- self, prompt: str, **kwargs: Any
+ self, prompt: str, formatted: bool = False, **kwargs: Any
) -> CompletionResponseAsyncGen:
async def gen() -> CompletionResponseAsyncGen:
for message in self.stream_complete(prompt, **kwargs):
diff --git a/llama_index/llms/vertex.py b/llama_index/llms/vertex.py
index 9abae84e7b..55fe349b12 100644
--- a/llama_index/llms/vertex.py
+++ b/llama_index/llms/vertex.py
@@ -187,7 +187,9 @@ class Vertex(LLM):
)
@llm_completion_callback()
- def complete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
+ def complete(
+ self, prompt: str, formatted: bool = False, **kwargs: Any
+ ) -> CompletionResponse:
kwargs = kwargs if kwargs else {}
params = {**self._model_kwargs, **kwargs}
if self.iscode and "candidate_count" in params:
@@ -248,7 +250,9 @@ class Vertex(LLM):
return gen()
@llm_completion_callback()
- def stream_complete(self, prompt: str, **kwargs: Any) -> CompletionResponseGen:
+ def stream_complete(
+ self, prompt: str, formatted: bool = False, **kwargs: Any
+ ) -> CompletionResponseGen:
kwargs = kwargs if kwargs else {}
params = {**self._model_kwargs, **kwargs}
if "candidate_count" in params:
@@ -311,7 +315,9 @@ class Vertex(LLM):
)
@llm_completion_callback()
- async def acomplete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
+ async def acomplete(
+ self, prompt: str, formatted: bool = False, **kwargs: Any
+ ) -> CompletionResponse:
kwargs = kwargs if kwargs else {}
params = {**self._model_kwargs, **kwargs}
if self.iscode and "candidate_count" in params:
@@ -333,6 +339,6 @@ class Vertex(LLM):
@llm_completion_callback()
async def astream_complete(
- self, prompt: str, **kwargs: Any
+ self, prompt: str, formatted: bool = False, **kwargs: Any
) -> CompletionResponseAsyncGen:
raise (ValueError("Not Implemented"))
diff --git a/llama_index/llms/vllm.py b/llama_index/llms/vllm.py
index 25b38c1970..f7dce9a5f4 100644
--- a/llama_index/llms/vllm.py
+++ b/llama_index/llms/vllm.py
@@ -227,7 +227,9 @@ class Vllm(LLM):
return completion_response_to_chat_response(completion_response)
@llm_completion_callback()
- def complete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
+ def complete(
+ self, prompt: str, formatted: bool = False, **kwargs: Any
+ ) -> CompletionResponse:
kwargs = kwargs if kwargs else {}
params = {**self._model_kwargs, **kwargs}
@@ -245,7 +247,9 @@ class Vllm(LLM):
raise (ValueError("Not Implemented"))
@llm_completion_callback()
- def stream_complete(self, prompt: str, **kwargs: Any) -> CompletionResponseGen:
+ def stream_complete(
+ self, prompt: str, formatted: bool = False, **kwargs: Any
+ ) -> CompletionResponseGen:
raise (ValueError("Not Implemented"))
@llm_chat_callback()
@@ -256,7 +260,9 @@ class Vllm(LLM):
return self.chat(messages, **kwargs)
@llm_completion_callback()
- async def acomplete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
+ async def acomplete(
+ self, prompt: str, formatted: bool = False, **kwargs: Any
+ ) -> CompletionResponse:
raise (ValueError("Not Implemented"))
@llm_chat_callback()
@@ -267,7 +273,7 @@ class Vllm(LLM):
@llm_completion_callback()
async def astream_complete(
- self, prompt: str, **kwargs: Any
+ self, prompt: str, formatted: bool = False, **kwargs: Any
) -> CompletionResponseAsyncGen:
raise (ValueError("Not Implemented"))
@@ -334,7 +340,9 @@ class VllmServer(Vllm):
return "VllmServer"
@llm_completion_callback()
- def complete(self, prompt: str, **kwargs: Any) -> List[CompletionResponse]:
+ def complete(
+ self, prompt: str, formatted: bool = False, **kwargs: Any
+ ) -> List[CompletionResponse]:
kwargs = kwargs if kwargs else {}
params = {**self._model_kwargs, **kwargs}
@@ -349,7 +357,9 @@ class VllmServer(Vllm):
return CompletionResponse(text=output[0])
@llm_completion_callback()
- def stream_complete(self, prompt: str, **kwargs: Any) -> CompletionResponseGen:
+ def stream_complete(
+ self, prompt: str, formatted: bool = False, **kwargs: Any
+ ) -> CompletionResponseGen:
kwargs = kwargs if kwargs else {}
params = {**self._model_kwargs, **kwargs}
@@ -372,13 +382,15 @@ class VllmServer(Vllm):
return gen()
@llm_completion_callback()
- async def acomplete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
+ async def acomplete(
+ self, prompt: str, formatted: bool = False, **kwargs: Any
+ ) -> CompletionResponse:
kwargs = kwargs if kwargs else {}
return self.complete(prompt, **kwargs)
@llm_completion_callback()
async def astream_complete(
- self, prompt: str, **kwargs: Any
+ self, prompt: str, formatted: bool = False, **kwargs: Any
) -> CompletionResponseAsyncGen:
kwargs = kwargs if kwargs else {}
params = {**self._model_kwargs, **kwargs}
diff --git a/llama_index/llms/watsonx.py b/llama_index/llms/watsonx.py
index 765cf0f5a8..22aa158f8b 100644
--- a/llama_index/llms/watsonx.py
+++ b/llama_index/llms/watsonx.py
@@ -149,7 +149,9 @@ class WatsonX(LLM):
return {**self._model_kwargs, **kwargs}
@llm_completion_callback()
- def complete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
+ def complete(
+ self, prompt: str, formatted: bool = False, **kwargs: Any
+ ) -> CompletionResponse:
all_kwargs = self._get_all_kwargs(**kwargs)
response = self._model.generate_text(prompt=prompt, params=all_kwargs)
@@ -157,7 +159,9 @@ class WatsonX(LLM):
return CompletionResponse(text=response)
@llm_completion_callback()
- def stream_complete(self, prompt: str, **kwargs: Any) -> CompletionResponseGen:
+ def stream_complete(
+ self, prompt: str, formatted: bool = False, **kwargs: Any
+ ) -> CompletionResponseGen:
all_kwargs = self._get_all_kwargs(**kwargs)
stream_response = self._model.generate_text_stream(
@@ -191,7 +195,9 @@ class WatsonX(LLM):
# Async Functions
# IBM Watson Machine Learning Package currently does not have Support for Async calls
- async def acomplete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
+ async def acomplete(
+ self, prompt: str, formatted: bool = False, **kwargs: Any
+ ) -> CompletionResponse:
raise NotImplementedError
async def astream_chat(
@@ -205,6 +211,6 @@ class WatsonX(LLM):
raise NotImplementedError
async def astream_complete(
- self, prompt: str, **kwargs: Any
+ self, prompt: str, formatted: bool = False, **kwargs: Any
) -> CompletionResponseAsyncGen:
raise NotImplementedError
diff --git a/llama_index/llms/xinference.py b/llama_index/llms/xinference.py
index f4b970bcff..b0fca9d90f 100644
--- a/llama_index/llms/xinference.py
+++ b/llama_index/llms/xinference.py
@@ -216,7 +216,9 @@ class Xinference(CustomLLM):
return gen()
@llm_completion_callback()
- def complete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
+ def complete(
+ self, prompt: str, formatted: bool = False, **kwargs: Any
+ ) -> CompletionResponse:
assert self._generator is not None
response_text = self._generator.chat(
prompt=prompt,
@@ -233,7 +235,9 @@ class Xinference(CustomLLM):
)
@llm_completion_callback()
- def stream_complete(self, prompt: str, **kwargs: Any) -> CompletionResponseGen:
+ def stream_complete(
+ self, prompt: str, formatted: bool = False, **kwargs: Any
+ ) -> CompletionResponseGen:
assert self._generator is not None
response_iter = self._generator.chat(
prompt=prompt,
diff --git a/llama_index/output_parsers/base.py b/llama_index/output_parsers/base.py
index fb95c5a6d9..f4ae1a864a 100644
--- a/llama_index/output_parsers/base.py
+++ b/llama_index/output_parsers/base.py
@@ -1,7 +1,17 @@
"""Base output parser class."""
from dataclasses import dataclass
-from typing import Any, Optional
+from typing import Any, Dict, Optional
+
+from llama_index.bridge.pydantic import Field
+from llama_index.core.query_pipeline.query_component import (
+ ChainableMixin,
+ InputKeys,
+ OutputKeys,
+ QueryComponent,
+ validate_and_convert_stringable,
+)
+from llama_index.types import BaseOutputParser
@dataclass
@@ -14,3 +24,50 @@ class StructuredOutput:
class OutputParserException(Exception):
pass
+
+
+class ChainableOutputParser(BaseOutputParser, ChainableMixin):
+ """Chainable output parser."""
+
+ # TODO: consolidate with base at some point if possible.
+
+ def _as_query_component(self, **kwargs: Any) -> QueryComponent:
+ """Get query component."""
+ return OutputParserComponent(output_parser=self)
+
+
+class OutputParserComponent(QueryComponent):
+ """Output parser component."""
+
+ output_parser: BaseOutputParser = Field(..., description="Output parser.")
+
+ class Config:
+ arbitrary_types_allowed = True
+
+ def _run_component(self, **kwargs: Any) -> Dict[str, Any]:
+ """Run component."""
+ output = self.output_parser.parse(kwargs["input"])
+ return {"output": output}
+
+ async def _arun_component(self, **kwargs: Any) -> Dict[str, Any]:
+ """Run component."""
+ # NOTE: no native async for output parser
+ return self._run_component(**kwargs)
+
+ def _validate_component_inputs(self, input: Any) -> Any:
+ """Validate component inputs during run_component."""
+ input["input"] = validate_and_convert_stringable(input["input"])
+ return input
+
+ def set_callback_manager(self, callback_manager: Any) -> None:
+ """Set callback manager."""
+
+ @property
+ def input_keys(self) -> Any:
+ """Input keys."""
+ return InputKeys.from_keys({"input"})
+
+ @property
+ def output_keys(self) -> Any:
+ """Output keys."""
+ return OutputKeys.from_keys({"output"})
diff --git a/llama_index/output_parsers/guardrails.py b/llama_index/output_parsers/guardrails.py
index 406ab904ee..cbf70c07f8 100644
--- a/llama_index/output_parsers/guardrails.py
+++ b/llama_index/output_parsers/guardrails.py
@@ -5,6 +5,8 @@ See https://github.com/ShreyaR/guardrails.
"""
from deprecated import deprecated
+from llama_index.output_parsers.base import ChainableOutputParser
+
try:
from guardrails import Guard
except ImportError:
@@ -16,7 +18,6 @@ from typing import TYPE_CHECKING, Any, Callable, Optional
if TYPE_CHECKING:
from llama_index.bridge.langchain import BaseLLM
-from llama_index.types import BaseOutputParser
def get_callable(llm: Optional["BaseLLM"]) -> Optional[Callable]:
@@ -27,7 +28,7 @@ def get_callable(llm: Optional["BaseLLM"]) -> Optional[Callable]:
return llm.__call__
-class GuardrailsOutputParser(BaseOutputParser):
+class GuardrailsOutputParser(ChainableOutputParser):
"""Guardrails output parser."""
def __init__(
diff --git a/llama_index/output_parsers/langchain.py b/llama_index/output_parsers/langchain.py
index a8a86c67d4..5c876ed223 100644
--- a/llama_index/output_parsers/langchain.py
+++ b/llama_index/output_parsers/langchain.py
@@ -3,12 +3,13 @@
from string import Formatter
from typing import TYPE_CHECKING, Any, Optional
+from llama_index.output_parsers.base import ChainableOutputParser
+
if TYPE_CHECKING:
from llama_index.bridge.langchain import BaseOutputParser as LCOutputParser
-from llama_index.types import BaseOutputParser
-class LangchainOutputParser(BaseOutputParser):
+class LangchainOutputParser(ChainableOutputParser):
"""Langchain output parser."""
def __init__(
diff --git a/llama_index/output_parsers/pydantic.py b/llama_index/output_parsers/pydantic.py
index d88247289f..0c4320ae71 100644
--- a/llama_index/output_parsers/pydantic.py
+++ b/llama_index/output_parsers/pydantic.py
@@ -3,8 +3,9 @@
import json
from typing import Any, List, Optional, Type
+from llama_index.output_parsers.base import ChainableOutputParser
from llama_index.output_parsers.utils import extract_json_str
-from llama_index.types import BaseOutputParser, Model
+from llama_index.types import Model
PYDANTIC_FORMAT_TMPL = """
Here's a JSON schema to follow:
@@ -14,7 +15,7 @@ Output a valid JSON object but do not repeat the schema.
"""
-class PydanticOutputParser(BaseOutputParser):
+class PydanticOutputParser(ChainableOutputParser):
"""Pydantic Output Parser.
Args:
@@ -39,12 +40,21 @@ class PydanticOutputParser(BaseOutputParser):
@property
def format_string(self) -> str:
+ """Format string."""
+ return self.get_format_string(escape_json=True)
+
+ def get_format_string(self, escape_json: bool = True) -> str:
+ """Format string."""
schema_dict = self._output_cls.schema()
for key in self._excluded_schema_keys_from_format:
del schema_dict[key]
schema_str = json.dumps(schema_dict)
- return self._pydantic_format_tmpl.format(schema=schema_str)
+ output_str = self._pydantic_format_tmpl.format(schema=schema_str)
+ if escape_json:
+ return output_str.replace("{", "{{").replace("}", "}}")
+ else:
+ return output_str
def parse(self, text: str) -> Any:
"""Parse, validate, and correct errors programmatically."""
@@ -53,4 +63,4 @@ class PydanticOutputParser(BaseOutputParser):
def format(self, query: str) -> str:
"""Format a query with structured output formatting instructions."""
- return query + "\n\n" + self.format_string
+ return query + "\n\n" + self.get_format_string(escape_json=True)
diff --git a/llama_index/postprocessor/types.py b/llama_index/postprocessor/types.py
index abdcef3ee8..2d0486fce4 100644
--- a/llama_index/postprocessor/types.py
+++ b/llama_index/postprocessor/types.py
@@ -1,13 +1,20 @@
from abc import ABC, abstractmethod
-from typing import List, Optional
+from typing import Any, Dict, List, Optional
from llama_index.bridge.pydantic import Field
from llama_index.callbacks import CallbackManager
+from llama_index.core.query_pipeline.query_component import (
+ ChainableMixin,
+ InputKeys,
+ OutputKeys,
+ QueryComponent,
+ validate_and_convert_stringable,
+)
from llama_index.prompts.mixin import PromptDictType, PromptMixinType
from llama_index.schema import BaseComponent, NodeWithScore, QueryBundle
-class BaseNodePostprocessor(BaseComponent, ABC):
+class BaseNodePostprocessor(ChainableMixin, BaseComponent, ABC):
callback_manager: CallbackManager = Field(
default_factory=CallbackManager, exclude=True
)
@@ -54,3 +61,60 @@ class BaseNodePostprocessor(BaseComponent, ABC):
query_bundle: Optional[QueryBundle] = None,
) -> List[NodeWithScore]:
"""Postprocess nodes."""
+
+ def _as_query_component(self, **kwargs: Any) -> QueryComponent:
+ """As query component."""
+ return PostprocessorComponent(postprocessor=self)
+
+
+class PostprocessorComponent(QueryComponent):
+ """Postprocessor component."""
+
+ postprocessor: BaseNodePostprocessor = Field(..., description="Postprocessor")
+
+ class Config:
+ arbitrary_types_allowed = True
+
+ def set_callback_manager(self, callback_manager: CallbackManager) -> None:
+ """Set callback manager."""
+ self.postprocessor.callback_manager = callback_manager
+
+ def _validate_component_inputs(self, input: Dict[str, Any]) -> Dict[str, Any]:
+ """Validate component inputs during run_component."""
+ # make sure `nodes` is a list of nodes
+ if "nodes" not in input:
+ raise ValueError("Input must have key 'nodes'")
+ nodes = input["nodes"]
+ if not isinstance(nodes, list):
+ raise ValueError("Input nodes must be a list")
+ for node in nodes:
+ if not isinstance(node, NodeWithScore):
+ raise ValueError("Input nodes must be a list of NodeWithScore")
+
+ # if query_str exists, make sure `query_str` is stringable
+ if "query_str" in input:
+ input["query_str"] = validate_and_convert_stringable(input["query_str"])
+
+ return input
+
+ def _run_component(self, **kwargs: Any) -> Any:
+ """Run component."""
+ output = self.postprocessor.postprocess_nodes(
+ kwargs["nodes"], query_str=kwargs.get("query_str", None)
+ )
+ return {"nodes": output}
+
+ async def _arun_component(self, **kwargs: Any) -> Any:
+ """Run component (async)."""
+ # NOTE: no native async for postprocessor
+ return self._run_component(**kwargs)
+
+ @property
+ def input_keys(self) -> InputKeys:
+ """Input keys."""
+ return InputKeys.from_keys({"nodes"}, optional_keys={"query_str"})
+
+ @property
+ def output_keys(self) -> OutputKeys:
+ """Output keys."""
+ return OutputKeys.from_keys({"nodes"})
diff --git a/llama_index/prompts/base.py b/llama_index/prompts/base.py
index 4b3a717f55..fb0946f9f3 100644
--- a/llama_index/prompts/base.py
+++ b/llama_index/prompts/base.py
@@ -3,7 +3,17 @@
from abc import ABC, abstractmethod
from copy import deepcopy
-from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Sequence, Tuple
+from typing import (
+ TYPE_CHECKING,
+ Any,
+ Callable,
+ Dict,
+ List,
+ Optional,
+ Sequence,
+ Tuple,
+ Union,
+)
from llama_index.bridge.pydantic import Field
@@ -14,6 +24,13 @@ if TYPE_CHECKING:
)
from llama_index.bridge.pydantic import BaseModel
from llama_index.core.llms.types import ChatMessage
+from llama_index.core.query_pipeline.query_component import (
+ ChainableMixin,
+ InputKeys,
+ OutputKeys,
+ QueryComponent,
+ validate_and_convert_stringable,
+)
from llama_index.llms.base import BaseLLM
from llama_index.llms.generic_utils import (
messages_to_prompt as default_messages_to_prompt,
@@ -26,7 +43,7 @@ from llama_index.prompts.utils import get_template_vars
from llama_index.types import BaseOutputParser
-class BasePromptTemplate(BaseModel, ABC):
+class BasePromptTemplate(ChainableMixin, BaseModel, ABC):
metadata: Dict[str, Any]
template_vars: List[str]
kwargs: Dict[str, str]
@@ -107,6 +124,12 @@ class BasePromptTemplate(BaseModel, ABC):
def get_template(self, llm: Optional[BaseLLM] = None) -> str:
...
+ def _as_query_component(
+ self, llm: Optional[BaseLLM] = None, **kwargs: Any
+ ) -> QueryComponent:
+ """As query component."""
+ return PromptComponent(prompt=self, format_messages=False, llm=llm)
+
class PromptTemplate(BasePromptTemplate):
template: str
@@ -273,6 +296,12 @@ class ChatPromptTemplate(BasePromptTemplate):
def get_template(self, llm: Optional[BaseLLM] = None) -> str:
return default_messages_to_prompt(self.message_templates)
+ def _as_query_component(
+ self, llm: Optional[BaseLLM] = None, **kwargs: Any
+ ) -> QueryComponent:
+ """As query component."""
+ return PromptComponent(prompt=self, format_messages=True, llm=llm)
+
class SelectorPromptTemplate(BasePromptTemplate):
default_template: BasePromptTemplate
@@ -488,3 +517,54 @@ class LangchainPromptTemplate(BasePromptTemplate):
# NOTE: only for backwards compatibility
Prompt = PromptTemplate
+
+
+class PromptComponent(QueryComponent):
+ """Prompt component."""
+
+ prompt: BasePromptTemplate = Field(..., description="Prompt")
+ llm: Optional[BaseLLM] = Field(
+ default=None, description="LLM to use for formatting prompt."
+ )
+ format_messages: bool = Field(
+ default=False,
+ description="Whether to format the prompt into a list of chat messages.",
+ )
+
+ class Config:
+ arbitrary_types_allowed = True
+
+ def set_callback_manager(self, callback_manager: Any) -> None:
+ """Set callback manager."""
+
+ def _validate_component_inputs(self, input: Dict[str, Any]) -> Dict[str, Any]:
+ """Validate component inputs during run_component."""
+ keys = list(input.keys())
+ for k in keys:
+ input[k] = validate_and_convert_stringable(input[k])
+ return input
+
+ def _run_component(self, **kwargs: Any) -> Any:
+ """Run component."""
+ if self.format_messages:
+ output: Union[str, List[ChatMessage]] = self.prompt.format_messages(
+ llm=self.llm, **kwargs
+ )
+ else:
+ output = self.prompt.format(llm=self.llm, **kwargs)
+ return {"prompt": output}
+
+ async def _arun_component(self, **kwargs: Any) -> Any:
+ """Run component."""
+ # NOTE: no native async for prompt
+ return self._run_component(**kwargs)
+
+ @property
+ def input_keys(self) -> InputKeys:
+ """Input keys."""
+ return InputKeys.from_keys(set(self.prompt.template_vars))
+
+ @property
+ def output_keys(self) -> OutputKeys:
+ """Output keys."""
+ return OutputKeys.from_keys({"prompt"})
diff --git a/llama_index/query_pipeline/__init__.py b/llama_index/query_pipeline/__init__.py
new file mode 100644
index 0000000000..94d4540406
--- /dev/null
+++ b/llama_index/query_pipeline/__init__.py
@@ -0,0 +1,17 @@
+"""Init file."""
+
+from llama_index.core.query_pipeline.query_component import (
+ CustomQueryComponent,
+ InputComponent,
+ QueryComponent,
+)
+from llama_index.query_pipeline.query import InputKeys, OutputKeys, QueryPipeline
+
+__all__ = [
+ "QueryPipeline",
+ "InputKeys",
+ "OutputKeys",
+ "QueryComponent",
+ "CustomQueryComponent",
+ "InputComponent",
+]
diff --git a/llama_index/query_pipeline/query.py b/llama_index/query_pipeline/query.py
new file mode 100644
index 0000000000..9890d13556
--- /dev/null
+++ b/llama_index/query_pipeline/query.py
@@ -0,0 +1,474 @@
+"""Query Pipeline."""
+
+import json
+import uuid
+from typing import Any, Dict, List, Optional, Sequence, Tuple, Union, cast
+
+import networkx
+
+from llama_index.bridge.pydantic import Field
+from llama_index.callbacks import CallbackManager
+from llama_index.callbacks.schema import CBEventType, EventPayload
+from llama_index.core.query_pipeline.query_component import (
+ ChainableMixin,
+ InputKeys,
+ OutputKeys,
+ QueryComponent,
+)
+from llama_index.utils import print_text
+
+# 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]
+
+
+def add_output_to_module_inputs(
+ src_key: str,
+ dest_key: str,
+ output_dict: Dict[str, Any],
+ module: QueryComponent,
+ module_inputs: Dict[str, Any],
+) -> None:
+ """Add input to module deps inputs."""
+ # get relevant output from link
+ if src_key is None:
+ # ensure that output_dict only has one key
+ if len(output_dict) != 1:
+ raise ValueError("Output dict must have exactly one key.")
+ output = next(iter(output_dict.values()))
+ else:
+ output = output_dict[src_key]
+
+ # now attach output to relevant input key for module
+ if dest_key is None:
+ free_keys = module.free_req_input_keys
+ # ensure that there is only one remaining key given partials
+ if len(free_keys) != 1:
+ raise ValueError(
+ "Module input keys must have exactly one key if "
+ "dest_key is not specified. Remaining keys: "
+ f"in module: {free_keys}"
+ )
+ module_inputs[next(iter(free_keys))] = output
+ else:
+ module_inputs[dest_key] = output
+
+
+def print_debug_input(
+ module_key: str,
+ input: Dict[str, Any],
+ val_str_len: int = 200,
+) -> None:
+ """Print debug input."""
+ output = f"> Running module {module_key} with input: \n"
+ for key, value in input.items():
+ # stringify and truncate output
+ val_str = (
+ str(value)[:val_str_len] + "..."
+ if len(str(value)) > val_str_len
+ else str(value)
+ )
+ output += f"{key}: {val_str}\n"
+
+ print_text(output + "\n", color="llama_lavender")
+
+
+class QueryPipeline(QueryComponent):
+ """A query pipeline that can allow arbitrary chaining of different modules.
+
+ A pipeline itself is a query component, and can be used as a module in another pipeline.
+
+ """
+
+ callback_manager: CallbackManager = Field(
+ default_factory=lambda: CallbackManager([]), exclude=True
+ )
+
+ module_dict: Dict[str, QueryComponent] = Field(
+ default_factory=dict, description="The modules in the pipeline."
+ )
+ dag: networkx.MultiDiGraph = Field(
+ default_factory=networkx.MultiDiGraph, description="The DAG of the pipeline."
+ )
+ verbose: bool = Field(
+ default=False, description="Whether to print intermediate steps."
+ )
+
+ class Config:
+ arbitrary_types_allowed = True
+
+ def __init__(
+ self,
+ callback_manager: Optional[CallbackManager] = None,
+ chain: Optional[Sequence[QUERY_COMPONENT_TYPE]] = None,
+ **kwargs: Any,
+ ):
+ super().__init__(
+ callback_manager=callback_manager or CallbackManager([]),
+ **kwargs,
+ )
+
+ if chain is not None:
+ # generate implicit link between each item, add
+ self.add_chain(chain)
+
+ def add_chain(self, chain: Sequence[QUERY_COMPONENT_TYPE]) -> None:
+ """Add a chain of modules to the pipeline.
+
+ This is a special form of pipeline that is purely sequential/linear.
+ This allows a more concise way of specifying a pipeline.
+
+ """
+ # first add all modules
+ module_keys = []
+ for module in chain:
+ module_key = str(uuid.uuid4())
+ self.add(module_key, module)
+ module_keys.append(module_key)
+
+ # then add all links
+ for i in range(len(chain) - 1):
+ self.add_link(src=module_keys[i], dest=module_keys[i + 1])
+
+ def add_modules(self, module_dict: Dict[str, QUERY_COMPONENT_TYPE]) -> None:
+ """Add modules to the pipeline."""
+ for module_key, module in module_dict.items():
+ self.add(module_key, module)
+
+ def add(self, module_key: str, module: QUERY_COMPONENT_TYPE) -> None:
+ """Add a module to the pipeline."""
+ # if already exists, raise error
+ if module_key in self.module_dict:
+ raise ValueError(f"Module {module_key} already exists in pipeline.")
+
+ if isinstance(module, ChainableMixin):
+ module = module.as_query_component()
+ else:
+ pass
+
+ self.module_dict[module_key] = cast(QueryComponent, module)
+ self.dag.add_node(module_key)
+
+ def add_link(
+ self,
+ src: str,
+ dest: str,
+ src_key: Optional[str] = None,
+ dest_key: Optional[str] = None,
+ ) -> None:
+ """Add a link between two modules."""
+ if src not in self.module_dict:
+ raise ValueError(f"Module {src} does not exist in pipeline.")
+ self.dag.add_edge(src, dest, src_key=src_key, dest_key=dest_key)
+
+ def _get_root_keys(self) -> List[str]:
+ """Get root keys."""
+ return [v for v, d in self.dag.in_degree() if d == 0]
+
+ def _get_leaf_keys(self) -> List[str]:
+ """Get leaf keys."""
+ # get all modules without downstream dependencies
+ return [v for v, d in self.dag.out_degree() if d == 0]
+
+ def set_callback_manager(self, callback_manager: CallbackManager) -> None:
+ """Set callback manager."""
+ # go through every module in module dict and set callback manager
+ self.callback_manager = callback_manager
+ for module in self.module_dict.values():
+ module.set_callback_manager(callback_manager)
+
+ def run(
+ self,
+ *args: Any,
+ return_values_direct: bool = True,
+ callback_manager: Optional[CallbackManager] = None,
+ **kwargs: Any,
+ ) -> Any:
+ """Run the pipeline."""
+ # first set callback manager
+ callback_manager = callback_manager or self.callback_manager
+ self.set_callback_manager(callback_manager)
+ with self.callback_manager.as_trace("query"):
+ with self.callback_manager.event(
+ CBEventType.QUERY, payload={EventPayload.QUERY_STR: json.dumps(kwargs)}
+ ) as query_event:
+ return self._run(
+ *args, return_values_direct=return_values_direct, **kwargs
+ )
+
+ def run_multi(
+ self,
+ module_input_dict: Dict[str, Any],
+ callback_manager: Optional[CallbackManager] = None,
+ ) -> Dict[str, Any]:
+ """Run the pipeline for multiple roots."""
+ callback_manager = callback_manager or self.callback_manager
+ self.set_callback_manager(callback_manager)
+ with self.callback_manager.as_trace("query"):
+ with self.callback_manager.event(
+ CBEventType.QUERY,
+ payload={EventPayload.QUERY_STR: json.dumps(module_input_dict)},
+ ) as query_event:
+ return self._run_multi(module_input_dict)
+
+ async def arun(
+ self,
+ *args: Any,
+ return_values_direct: bool = True,
+ callback_manager: Optional[CallbackManager] = None,
+ **kwargs: Any,
+ ) -> Any:
+ """Run the pipeline."""
+ # first set callback manager
+ callback_manager = callback_manager or self.callback_manager
+ self.set_callback_manager(callback_manager)
+ with self.callback_manager.as_trace("query"):
+ with self.callback_manager.event(
+ CBEventType.QUERY, payload={EventPayload.QUERY_STR: json.dumps(kwargs)}
+ ) as query_event:
+ return await self._arun(
+ *args, return_values_direct=return_values_direct, **kwargs
+ )
+
+ async def arun_multi(
+ self,
+ module_input_dict: Dict[str, Any],
+ callback_manager: Optional[CallbackManager] = None,
+ ) -> Dict[str, Any]:
+ """Run the pipeline for multiple roots."""
+ callback_manager = callback_manager or self.callback_manager
+ self.set_callback_manager(callback_manager)
+ with self.callback_manager.as_trace("query"):
+ with self.callback_manager.event(
+ CBEventType.QUERY,
+ payload={EventPayload.QUERY_STR: json.dumps(module_input_dict)},
+ ) as query_event:
+ return await self._arun_multi(module_input_dict)
+
+ def _get_root_key_and_kwargs(
+ self, *args: Any, **kwargs: Any
+ ) -> Tuple[str, Dict[str, Any]]:
+ """Get root key and kwargs.
+
+ This is for `_run`.
+
+ """
+ ## run pipeline
+ ## assume there is only one root - for multiple roots, need to specify `run_multi`
+ root_keys = self._get_root_keys()
+ if len(root_keys) != 1:
+ raise ValueError("Only one root is supported.")
+ root_key = root_keys[0]
+
+ root_module = self.module_dict[root_key]
+ if len(args) > 0:
+ # if args is specified, validate. only one arg is allowed, and there can only be one free
+ # input key in the module
+ if len(args) > 1:
+ raise ValueError("Only one arg is allowed.")
+ if len(kwargs) > 0:
+ raise ValueError("No kwargs allowed if args is specified.")
+ if len(root_module.free_req_input_keys) != 1:
+ raise ValueError("Only one free input key is allowed.")
+ # set kwargs
+ kwargs[next(iter(root_module.free_req_input_keys))] = args[0]
+ return root_key, kwargs
+
+ def _get_single_result_output(
+ self,
+ result_outputs: Dict[str, Any],
+ return_values_direct: bool,
+ ) -> Any:
+ """Get result output from a single module.
+
+ If output dict is a single key, return the value directly
+ if return_values_direct is True.
+
+ """
+ if len(result_outputs) != 1:
+ raise ValueError("Only one output is supported.")
+
+ result_output = next(iter(result_outputs.values()))
+ # return_values_direct: if True, return the value directly
+ # without the key
+ # if it's a dict with one key, return the value
+ if (
+ isinstance(result_output, dict)
+ and len(result_output) == 1
+ and return_values_direct
+ ):
+ return next(iter(result_output.values()))
+ else:
+ return result_output
+
+ def _run(self, *args: Any, return_values_direct: bool = True, **kwargs: Any) -> Any:
+ """Run the pipeline.
+
+ Assume that there is a single root module and a single output module.
+
+ For multi-input and multi-outputs, please see `run_multi`.
+
+ """
+ root_key, kwargs = self._get_root_key_and_kwargs(*args, **kwargs)
+ # call run_multi with one root key
+ result_outputs = self._run_multi({root_key: kwargs})
+ return self._get_single_result_output(result_outputs, return_values_direct)
+
+ async def _arun(
+ self, *args: Any, return_values_direct: bool = True, **kwargs: Any
+ ) -> Any:
+ """Run the pipeline.
+
+ Assume that there is a single root module and a single output module.
+
+ For multi-input and multi-outputs, please see `run_multi`.
+
+ """
+ root_key, kwargs = self._get_root_key_and_kwargs(*args, **kwargs)
+ # call run_multi with one root key
+ result_outputs = await self._arun_multi({root_key: kwargs})
+ return self._get_single_result_output(result_outputs, return_values_direct)
+
+ def _validate_inputs(self, module_input_dict: Dict[str, Any]) -> None:
+ root_keys = self._get_root_keys()
+ # if root keys don't match up with kwargs keys, raise error
+ if set(root_keys) != set(module_input_dict.keys()):
+ raise ValueError(
+ "Expected root keys do not match up with input keys.\n"
+ f"Expected root keys: {root_keys}\n"
+ f"Input keys: {module_input_dict.keys()}\n"
+ )
+
+ def _process_component_output(
+ self,
+ output_dict: Dict[str, Any],
+ module_key: str,
+ all_module_inputs: Dict[str, Dict[str, Any]],
+ result_outputs: Dict[str, Any],
+ ) -> None:
+ """Process component output."""
+ # if there's no more edges, add result to output
+ if module_key in self._get_leaf_keys():
+ result_outputs[module_key] = output_dict
+ else:
+ for _, dest, attr in self.dag.edges(module_key, data=True):
+ edge_module = self.module_dict[dest]
+
+ # add input to module_deps_inputs
+ add_output_to_module_inputs(
+ attr.get("src_key"),
+ attr.get("dest_key"),
+ output_dict,
+ edge_module,
+ all_module_inputs[dest],
+ )
+
+ def _run_multi(self, module_input_dict: Dict[str, Any]) -> Dict[str, Any]:
+ """Run the pipeline for multiple roots.
+
+ kwargs is in the form of module_dict -> input_dict
+ input_dict is in the form of input_key -> input
+
+ """
+ self._validate_inputs(module_input_dict)
+ queue = list(networkx.topological_sort(self.dag))
+
+ # module_deps_inputs is a dict to collect inputs for a module
+ # mapping of module_key -> dict of input_key -> input
+ # initialize with blank dict for every module key
+ # the input dict of each module key will be populated as the upstream modules are run
+ all_module_inputs: Dict[str, Dict[str, Any]] = {
+ module_key: {} for module_key in self.module_dict
+ }
+ result_outputs: Dict[str, Any] = {}
+
+ # add root inputs to all_module_inputs
+ for module_key, module_input in module_input_dict.items():
+ all_module_inputs[module_key] = module_input
+
+ while len(queue) > 0:
+ module_key = queue.pop(0)
+ module = self.module_dict[module_key]
+ module_input = all_module_inputs[module_key]
+
+ if self.verbose:
+ print_debug_input(module_key, module_input)
+ output_dict = module.run_component(**module_input)
+
+ # get new nodes and is_leaf
+ self._process_component_output(
+ output_dict, module_key, all_module_inputs, result_outputs
+ )
+
+ return result_outputs
+
+ async def _arun_multi(self, module_input_dict: Dict[str, Any]) -> Dict[str, Any]:
+ """Run the pipeline for multiple roots.
+
+ kwargs is in the form of module_dict -> input_dict
+ input_dict is in the form of input_key -> input
+
+ """
+ self._validate_inputs(module_input_dict)
+ queue = list(networkx.topological_sort(self.dag))
+
+ # module_deps_inputs is a dict to collect inputs for a module
+ # mapping of module_key -> dict of input_key -> input
+ # initialize with blank dict for every module key
+ # the input dict of each module key will be populated as the upstream modules are run
+ all_module_inputs: Dict[str, Dict[str, Any]] = {
+ module_key: {} for module_key in self.module_dict
+ }
+ result_outputs: Dict[str, Any] = {}
+
+ # add root inputs to all_module_inputs
+ for module_key, module_input in module_input_dict.items():
+ all_module_inputs[module_key] = module_input
+
+ while len(queue) > 0:
+ module_key = queue.pop(0)
+ module = self.module_dict[module_key]
+ module_input = all_module_inputs[module_key]
+
+ if self.verbose:
+ print_debug_input(module_key, module_input)
+ output_dict = await module.arun_component(**module_input)
+
+ # get new nodes and is_leaf
+ self._process_component_output(
+ output_dict, module_key, all_module_inputs, result_outputs
+ )
+
+ return result_outputs
+
+ def _validate_component_inputs(self, input: Dict[str, Any]) -> Dict[str, Any]:
+ """Validate component inputs during run_component."""
+ return input
+
+ def _run_component(self, **kwargs: Any) -> Dict[str, Any]:
+ """Run component."""
+ return self.run(return_values_direct=False, **kwargs)
+
+ async def _arun_component(self, **kwargs: Any) -> Dict[str, Any]:
+ """Run component."""
+ return await self.arun(return_values_direct=False, **kwargs)
+
+ @property
+ def input_keys(self) -> InputKeys:
+ """Input keys."""
+ # get input key of first module
+ root_keys = self._get_root_keys()
+ if len(root_keys) != 1:
+ raise ValueError("Only one root is supported.")
+ root_module = self.module_dict[root_keys[0]]
+ return root_module.input_keys
+
+ @property
+ def output_keys(self) -> OutputKeys:
+ """Output keys."""
+ # get output key of last module
+ leaf_keys = self._get_leaf_keys()
+ if len(leaf_keys) != 1:
+ raise ValueError("Only one leaf is supported.")
+ leaf_module = self.module_dict[leaf_keys[0]]
+ return leaf_module.output_keys
diff --git a/llama_index/response_synthesizers/base.py b/llama_index/response_synthesizers/base.py
index 5790b33134..884cc403b4 100644
--- a/llama_index/response_synthesizers/base.py
+++ b/llama_index/response_synthesizers/base.py
@@ -11,8 +11,16 @@ import logging
from abc import abstractmethod
from typing import Any, Dict, Generator, List, Optional, Sequence, Union
-from llama_index.bridge.pydantic import BaseModel
+from llama_index.bridge.pydantic import BaseModel, Field
+from llama_index.callbacks.base import CallbackManager
from llama_index.callbacks.schema import CBEventType, EventPayload
+from llama_index.core.query_pipeline.query_component import (
+ ChainableMixin,
+ InputKeys,
+ OutputKeys,
+ QueryComponent,
+ validate_and_convert_stringable,
+)
from llama_index.core.response.schema import (
RESPONSE_TYPE,
PydanticResponse,
@@ -29,7 +37,7 @@ logger = logging.getLogger(__name__)
QueryTextType = Union[str, QueryBundle]
-class BaseSynthesizer(PromptMixin):
+class BaseSynthesizer(ChainableMixin, PromptMixin):
"""Response builder class."""
def __init__(
@@ -53,6 +61,14 @@ class BaseSynthesizer(PromptMixin):
def service_context(self) -> ServiceContext:
return self._service_context
+ @property
+ def callback_manager(self) -> CallbackManager:
+ return self._callback_manager
+
+ @callback_manager.setter
+ def callback_manager(self, callback_manager: CallbackManager) -> None:
+ self._callback_manager = callback_manager
+
@abstractmethod
def get_response(
self,
@@ -192,3 +208,59 @@ class BaseSynthesizer(PromptMixin):
event.on_end(payload={EventPayload.RESPONSE: response})
return response
+
+ def _as_query_component(self, **kwargs: Any) -> QueryComponent:
+ """As query component."""
+ return SynthesizerComponent(synthesizer=self)
+
+
+class SynthesizerComponent(QueryComponent):
+ """Synthesizer component."""
+
+ synthesizer: BaseSynthesizer = Field(..., description="Synthesizer")
+
+ class Config:
+ arbitrary_types_allowed = True
+
+ def set_callback_manager(self, callback_manager: CallbackManager) -> None:
+ """Set callback manager."""
+ self.synthesizer.callback_manager = callback_manager
+
+ def _validate_component_inputs(self, input: Dict[str, Any]) -> Dict[str, Any]:
+ """Validate component inputs during run_component."""
+ # make sure both query_str and nodes are there
+ if "query_str" not in input:
+ raise ValueError("Input must have key 'query_str'")
+ input["query_str"] = validate_and_convert_stringable(input["query_str"])
+
+ if "nodes" not in input:
+ raise ValueError("Input must have key 'nodes'")
+ nodes = input["nodes"]
+ if not isinstance(nodes, list):
+ raise ValueError("Input nodes must be a list")
+ for node in nodes:
+ if not isinstance(node, NodeWithScore):
+ raise ValueError("Input nodes must be a list of NodeWithScore")
+ return input
+
+ def _run_component(self, **kwargs: Any) -> Dict[str, Any]:
+ """Run component."""
+ output = self.synthesizer.synthesize(kwargs["query_str"], kwargs["nodes"])
+ return {"output": output}
+
+ async def _arun_component(self, **kwargs: Any) -> Dict[str, Any]:
+ """Run component."""
+ output = await self.synthesizer.asynthesize(
+ kwargs["query_str"], kwargs["nodes"]
+ )
+ return {"output": output}
+
+ @property
+ def input_keys(self) -> InputKeys:
+ """Input keys."""
+ return InputKeys.from_keys({"query_str", "nodes"})
+
+ @property
+ def output_keys(self) -> OutputKeys:
+ """Output keys."""
+ return OutputKeys.from_keys({"output"})
diff --git a/poetry.lock b/poetry.lock
index 294b773a33..6a5330fff8 100644
--- a/poetry.lock
+++ b/poetry.lock
@@ -1,9 +1,10 @@
-# This file is automatically @generated by Poetry 1.7.1 and should not be changed by hand.
+# This file is automatically @generated by Poetry 1.4.1 and should not be changed by hand.
[[package]]
name = "accelerate"
version = "0.25.0"
description = "Accelerate"
+category = "main"
optional = true
python-versions = ">=3.8.0"
files = [
@@ -34,6 +35,7 @@ testing = ["bitsandbytes", "datasets", "deepspeed", "evaluate", "parameterized",
name = "aenum"
version = "3.1.15"
description = "Advanced Enumerations (compatible with Python's stdlib Enum), NamedTuples, and NamedConstants"
+category = "main"
optional = true
python-versions = "*"
files = [
@@ -46,6 +48,7 @@ files = [
name = "aiohttp"
version = "3.9.1"
description = "Async http client/server framework (asyncio)"
+category = "main"
optional = false
python-versions = ">=3.8"
files = [
@@ -142,6 +145,7 @@ speedups = ["Brotli", "aiodns", "brotlicffi"]
name = "aiosignal"
version = "1.3.1"
description = "aiosignal: a list of registered asynchronous callbacks"
+category = "main"
optional = false
python-versions = ">=3.7"
files = [
@@ -156,6 +160,7 @@ frozenlist = ">=1.1.0"
name = "alabaster"
version = "0.7.13"
description = "A configurable sidebar-enabled Sphinx theme"
+category = "dev"
optional = false
python-versions = ">=3.6"
files = [
@@ -167,6 +172,7 @@ files = [
name = "anyio"
version = "4.2.0"
description = "High level compatibility layer for multiple asynchronous event loop implementations"
+category = "main"
optional = false
python-versions = ">=3.8"
files = [
@@ -189,6 +195,7 @@ trio = ["trio (>=0.23)"]
name = "appnope"
version = "0.1.3"
description = "Disable App Nap on macOS >= 10.9"
+category = "dev"
optional = false
python-versions = "*"
files = [
@@ -200,6 +207,7 @@ files = [
name = "argon2-cffi"
version = "23.1.0"
description = "Argon2 for Python"
+category = "dev"
optional = false
python-versions = ">=3.7"
files = [
@@ -220,6 +228,7 @@ typing = ["mypy"]
name = "argon2-cffi-bindings"
version = "21.2.0"
description = "Low-level CFFI bindings for Argon2"
+category = "dev"
optional = false
python-versions = ">=3.6"
files = [
@@ -257,6 +266,7 @@ tests = ["pytest"]
name = "arrow"
version = "1.3.0"
description = "Better dates & times for Python"
+category = "dev"
optional = false
python-versions = ">=3.8"
files = [
@@ -270,12 +280,13 @@ types-python-dateutil = ">=2.8.10"
[package.extras]
doc = ["doc8", "sphinx (>=7.0.0)", "sphinx-autobuild", "sphinx-autodoc-typehints", "sphinx_rtd_theme (>=1.3.0)"]
-test = ["dateparser (==1.*)", "pre-commit", "pytest", "pytest-cov", "pytest-mock", "pytz (==2021.1)", "simplejson (==3.*)"]
+test = ["dateparser (>=1.0.0,<2.0.0)", "pre-commit", "pytest", "pytest-cov", "pytest-mock", "pytz (==2021.1)", "simplejson (>=3.0.0,<4.0.0)"]
[[package]]
name = "astroid"
version = "2.13.5"
description = "An abstract syntax tree for Python with inference support."
+category = "dev"
optional = false
python-versions = ">=3.7.2"
files = [
@@ -295,6 +306,7 @@ wrapt = [
name = "asttokens"
version = "2.4.1"
description = "Annotate AST trees with source code positions"
+category = "dev"
optional = false
python-versions = "*"
files = [
@@ -313,6 +325,7 @@ test = ["astroid (>=1,<2)", "astroid (>=2,<4)", "pytest"]
name = "async-lru"
version = "2.0.4"
description = "Simple LRU cache for asyncio"
+category = "dev"
optional = false
python-versions = ">=3.8"
files = [
@@ -327,6 +340,7 @@ typing-extensions = {version = ">=4.0.0", markers = "python_version < \"3.11\""}
name = "async-timeout"
version = "4.0.3"
description = "Timeout context manager for asyncio programs"
+category = "main"
optional = false
python-versions = ">=3.7"
files = [
@@ -338,6 +352,7 @@ files = [
name = "asyncpg"
version = "0.28.0"
description = "An asyncio PostgreSQL driver"
+category = "main"
optional = true
python-versions = ">=3.7.0"
files = [
@@ -391,6 +406,7 @@ test = ["flake8 (>=5.0,<6.0)", "uvloop (>=0.15.3)"]
name = "attrs"
version = "23.2.0"
description = "Classes Without Boilerplate"
+category = "main"
optional = false
python-versions = ">=3.7"
files = [
@@ -410,6 +426,7 @@ tests-no-zope = ["attrs[tests-mypy]", "cloudpickle", "hypothesis", "pympler", "p
name = "autodoc-pydantic"
version = "1.9.0"
description = "Seamlessly integrate pydantic models in your Sphinx documentation."
+category = "dev"
optional = false
python-versions = ">=3.7.1,<4.0.0"
files = [
@@ -431,6 +448,7 @@ test = ["coverage (>=7,<8)", "pytest (>=7,<8)"]
name = "babel"
version = "2.14.0"
description = "Internationalization utilities"
+category = "dev"
optional = false
python-versions = ">=3.7"
files = [
@@ -448,6 +466,7 @@ dev = ["freezegun (>=1.0,<2.0)", "pytest (>=6.0)", "pytest-cov"]
name = "backcall"
version = "0.2.0"
description = "Specifications for callback functions passed in to an API"
+category = "dev"
optional = false
python-versions = "*"
files = [
@@ -459,6 +478,7 @@ files = [
name = "beautifulsoup4"
version = "4.12.2"
description = "Screen-scraping library"
+category = "main"
optional = false
python-versions = ">=3.6.0"
files = [
@@ -477,6 +497,7 @@ lxml = ["lxml"]
name = "black"
version = "23.9.1"
description = "The uncompromising code formatter."
+category = "dev"
optional = false
python-versions = ">=3.8"
files = [
@@ -525,6 +546,7 @@ uvloop = ["uvloop (>=0.15.2)"]
name = "bleach"
version = "6.1.0"
description = "An easy safelist-based HTML-sanitizing tool."
+category = "dev"
optional = false
python-versions = ">=3.8"
files = [
@@ -543,6 +565,7 @@ css = ["tinycss2 (>=1.1.0,<1.3)"]
name = "blis"
version = "0.7.11"
description = "The Blis BLAS-like linear algebra library, as a self-contained C-extension."
+category = "main"
optional = true
python-versions = "*"
files = [
@@ -592,6 +615,7 @@ numpy = [
name = "boto3"
version = "1.33.6"
description = "The AWS SDK for Python"
+category = "dev"
optional = false
python-versions = ">= 3.7"
files = [
@@ -611,6 +635,7 @@ crt = ["botocore[crt] (>=1.21.0,<2.0a0)"]
name = "botocore"
version = "1.33.13"
description = "Low-level, data-driven core of boto 3."
+category = "dev"
optional = false
python-versions = ">= 3.7"
files = [
@@ -633,6 +658,7 @@ crt = ["awscrt (==0.19.17)"]
name = "cachetools"
version = "5.3.2"
description = "Extensible memoizing collections and decorators"
+category = "main"
optional = false
python-versions = ">=3.7"
files = [
@@ -644,6 +670,7 @@ files = [
name = "catalogue"
version = "2.0.10"
description = "Super lightweight function registries for your library"
+category = "main"
optional = true
python-versions = ">=3.6"
files = [
@@ -655,6 +682,7 @@ files = [
name = "certifi"
version = "2023.11.17"
description = "Python package for providing Mozilla's CA Bundle."
+category = "main"
optional = false
python-versions = ">=3.6"
files = [
@@ -666,6 +694,7 @@ files = [
name = "cffi"
version = "1.16.0"
description = "Foreign Function Interface for Python calling C code."
+category = "main"
optional = false
python-versions = ">=3.8"
files = [
@@ -730,6 +759,7 @@ pycparser = "*"
name = "cfgv"
version = "3.4.0"
description = "Validate configuration and produce human readable error messages."
+category = "dev"
optional = false
python-versions = ">=3.8"
files = [
@@ -741,6 +771,7 @@ files = [
name = "charset-normalizer"
version = "3.3.2"
description = "The Real First Universal Charset Detector. Open, modern and actively maintained alternative to Chardet."
+category = "main"
optional = false
python-versions = ">=3.7.0"
files = [
@@ -840,6 +871,7 @@ files = [
name = "click"
version = "8.1.7"
description = "Composable command line interface toolkit"
+category = "main"
optional = false
python-versions = ">=3.7"
files = [
@@ -854,6 +886,7 @@ colorama = {version = "*", markers = "platform_system == \"Windows\""}
name = "cloudpathlib"
version = "0.16.0"
description = "pathlib-style classes for cloud storage services."
+category = "main"
optional = true
python-versions = ">=3.7"
files = [
@@ -874,6 +907,7 @@ s3 = ["boto3"]
name = "codespell"
version = "2.2.6"
description = "Codespell"
+category = "dev"
optional = false
python-versions = ">=3.8"
files = [
@@ -894,6 +928,7 @@ types = ["chardet (>=5.1.0)", "mypy", "pytest", "pytest-cov", "pytest-dependency
name = "colorama"
version = "0.4.6"
description = "Cross-platform colored terminal text."
+category = "main"
optional = false
python-versions = "!=3.0.*,!=3.1.*,!=3.2.*,!=3.3.*,!=3.4.*,!=3.5.*,!=3.6.*,>=2.7"
files = [
@@ -905,6 +940,7 @@ files = [
name = "coloredlogs"
version = "15.0.1"
description = "Colored terminal output for Python's logging module"
+category = "main"
optional = true
python-versions = ">=2.7, !=3.0.*, !=3.1.*, !=3.2.*, !=3.3.*, !=3.4.*"
files = [
@@ -922,6 +958,7 @@ cron = ["capturer (>=2.4)"]
name = "comm"
version = "0.2.1"
description = "Jupyter Python Comm implementation, for usage in ipykernel, xeus-python etc."
+category = "dev"
optional = false
python-versions = ">=3.8"
files = [
@@ -939,6 +976,7 @@ test = ["pytest"]
name = "confection"
version = "0.1.4"
description = "The sweetest config system for Python"
+category = "main"
optional = true
python-versions = ">=3.6"
files = [
@@ -954,6 +992,7 @@ srsly = ">=2.4.0,<3.0.0"
name = "cryptography"
version = "41.0.7"
description = "cryptography is a package which provides cryptographic recipes and primitives to Python developers."
+category = "main"
optional = false
python-versions = ">=3.7"
files = [
@@ -999,6 +1038,7 @@ test-randomorder = ["pytest-randomly"]
name = "cymem"
version = "2.0.8"
description = "Manage calls to calloc/free through Cython"
+category = "main"
optional = true
python-versions = "*"
files = [
@@ -1041,6 +1081,7 @@ files = [
name = "dataclasses-json"
version = "0.6.3"
description = "Easily serialize dataclasses to and from JSON."
+category = "main"
optional = false
python-versions = ">=3.7,<4.0"
files = [
@@ -1056,6 +1097,7 @@ typing-inspect = ">=0.4.0,<1"
name = "datasets"
version = "2.14.4"
description = "HuggingFace community-driven open-source library of datasets"
+category = "main"
optional = true
python-versions = ">=3.8.0"
files = [
@@ -1098,6 +1140,7 @@ vision = ["Pillow (>=6.2.1)"]
name = "debugpy"
version = "1.8.0"
description = "An implementation of the Debug Adapter Protocol for Python"
+category = "dev"
optional = false
python-versions = ">=3.8"
files = [
@@ -1125,6 +1168,7 @@ files = [
name = "decorator"
version = "5.1.1"
description = "Decorators for Humans"
+category = "dev"
optional = false
python-versions = ">=3.5"
files = [
@@ -1136,6 +1180,7 @@ files = [
name = "defusedxml"
version = "0.7.1"
description = "XML bomb protection for Python stdlib modules"
+category = "dev"
optional = false
python-versions = ">=2.7, !=3.0.*, !=3.1.*, !=3.2.*, !=3.3.*, !=3.4.*"
files = [
@@ -1147,6 +1192,7 @@ files = [
name = "deprecated"
version = "1.2.14"
description = "Python @deprecated decorator to deprecate old python classes, functions or methods."
+category = "main"
optional = false
python-versions = ">=2.7, !=3.0.*, !=3.1.*, !=3.2.*, !=3.3.*"
files = [
@@ -1164,6 +1210,7 @@ dev = ["PyTest", "PyTest-Cov", "bump2version (<1)", "sphinx (<2)", "tox"]
name = "dill"
version = "0.3.7"
description = "serialize all of Python"
+category = "main"
optional = false
python-versions = ">=3.7"
files = [
@@ -1178,6 +1225,7 @@ graph = ["objgraph (>=1.7.2)"]
name = "diskcache"
version = "5.6.3"
description = "Disk Cache -- Disk and file backed persistent cache."
+category = "main"
optional = true
python-versions = ">=3"
files = [
@@ -1189,6 +1237,7 @@ files = [
name = "distlib"
version = "0.3.8"
description = "Distribution utilities"
+category = "dev"
optional = false
python-versions = "*"
files = [
@@ -1200,6 +1249,7 @@ files = [
name = "distro"
version = "1.9.0"
description = "Distro - an OS platform information API"
+category = "main"
optional = false
python-versions = ">=3.6"
files = [
@@ -1211,6 +1261,7 @@ files = [
name = "dnspython"
version = "2.4.2"
description = "DNS toolkit"
+category = "dev"
optional = false
python-versions = ">=3.8,<4.0"
files = [
@@ -1230,6 +1281,7 @@ wmi = ["wmi (>=1.5.1,<2.0.0)"]
name = "docutils"
version = "0.16"
description = "Docutils -- Python Documentation Utilities"
+category = "dev"
optional = false
python-versions = ">=2.7, !=3.0.*, !=3.1.*, !=3.2.*, !=3.3.*, !=3.4.*"
files = [
@@ -1241,6 +1293,7 @@ files = [
name = "entrypoints"
version = "0.4"
description = "Discover and load entry points from installed packages."
+category = "dev"
optional = false
python-versions = ">=3.6"
files = [
@@ -1252,6 +1305,7 @@ files = [
name = "evaluate"
version = "0.4.1"
description = "HuggingFace community-driven open-source library of evaluation"
+category = "main"
optional = true
python-versions = ">=3.7.0"
files = [
@@ -1288,6 +1342,7 @@ torch = ["torch"]
name = "exceptiongroup"
version = "1.2.0"
description = "Backport of PEP 654 (exception groups)"
+category = "main"
optional = false
python-versions = ">=3.7"
files = [
@@ -1302,6 +1357,7 @@ test = ["pytest (>=6)"]
name = "executing"
version = "2.0.1"
description = "Get the currently executing AST node of a frame, and other information"
+category = "dev"
optional = false
python-versions = ">=3.5"
files = [
@@ -1316,6 +1372,7 @@ tests = ["asttokens (>=2.1.0)", "coverage", "coverage-enable-subprocess", "ipyth
name = "fastjsonschema"
version = "2.19.1"
description = "Fastest Python implementation of JSON schema"
+category = "dev"
optional = false
python-versions = "*"
files = [
@@ -1330,6 +1387,7 @@ devel = ["colorama", "json-spec", "jsonschema", "pylint", "pytest", "pytest-benc
name = "filelock"
version = "3.13.1"
description = "A platform independent file lock."
+category = "main"
optional = false
python-versions = ">=3.8"
files = [
@@ -1346,6 +1404,7 @@ typing = ["typing-extensions (>=4.8)"]
name = "flatbuffers"
version = "23.5.26"
description = "The FlatBuffers serialization format for Python"
+category = "main"
optional = true
python-versions = "*"
files = [
@@ -1357,6 +1416,7 @@ files = [
name = "fqdn"
version = "1.5.1"
description = "Validates fully-qualified domain names against RFC 1123, so that they are acceptable to modern bowsers"
+category = "dev"
optional = false
python-versions = ">=2.7, !=3.0, !=3.1, !=3.2, !=3.3, !=3.4, <4"
files = [
@@ -1368,6 +1428,7 @@ files = [
name = "frozenlist"
version = "1.4.1"
description = "A list-like structure which implements collections.abc.MutableSequence"
+category = "main"
optional = false
python-versions = ">=3.8"
files = [
@@ -1454,6 +1515,7 @@ files = [
name = "fsspec"
version = "2023.12.2"
description = "File-system specification"
+category = "main"
optional = false
python-versions = ">=3.8"
files = [
@@ -1493,6 +1555,7 @@ tqdm = ["tqdm"]
name = "furo"
version = "2023.3.27"
description = "A clean customisable Sphinx documentation theme."
+category = "dev"
optional = false
python-versions = ">=3.7"
files = [
@@ -1510,6 +1573,7 @@ sphinx-basic-ng = "*"
name = "google-ai-generativelanguage"
version = "0.4.0"
description = "Google Ai Generativelanguage API client library"
+category = "dev"
optional = false
python-versions = ">=3.7"
files = [
@@ -1518,7 +1582,7 @@ files = [
]
[package.dependencies]
-google-api-core = {version = ">=1.34.0,<2.0.dev0 || >=2.11.dev0,<3.0.0dev", extras = ["grpc"]}
+google-api-core = {version = ">=1.34.0,<2.0.0 || >=2.11.0,<3.0.0dev", extras = ["grpc"]}
proto-plus = ">=1.22.3,<2.0.0dev"
protobuf = ">=3.19.5,<3.20.0 || >3.20.0,<3.20.1 || >3.20.1,<4.21.0 || >4.21.0,<4.21.1 || >4.21.1,<4.21.2 || >4.21.2,<4.21.3 || >4.21.3,<4.21.4 || >4.21.4,<4.21.5 || >4.21.5,<5.0.0dev"
@@ -1526,6 +1590,7 @@ protobuf = ">=3.19.5,<3.20.0 || >3.20.0,<3.20.1 || >3.20.1,<4.21.0 || >4.21.0,<4
name = "google-api-core"
version = "2.15.0"
description = "Google API client core library"
+category = "dev"
optional = false
python-versions = ">=3.7"
files = [
@@ -1537,11 +1602,11 @@ files = [
google-auth = ">=2.14.1,<3.0.dev0"
googleapis-common-protos = ">=1.56.2,<2.0.dev0"
grpcio = [
- {version = ">=1.33.2,<2.0dev", optional = true, markers = "python_version < \"3.11\" and extra == \"grpc\""},
+ {version = ">=1.33.2,<2.0dev", optional = true, markers = "extra == \"grpc\""},
{version = ">=1.49.1,<2.0dev", optional = true, markers = "python_version >= \"3.11\" and extra == \"grpc\""},
]
grpcio-status = [
- {version = ">=1.33.2,<2.0.dev0", optional = true, markers = "python_version < \"3.11\" and extra == \"grpc\""},
+ {version = ">=1.33.2,<2.0.dev0", optional = true, markers = "extra == \"grpc\""},
{version = ">=1.49.1,<2.0.dev0", optional = true, markers = "python_version >= \"3.11\" and extra == \"grpc\""},
]
protobuf = ">=3.19.5,<3.20.0 || >3.20.0,<3.20.1 || >3.20.1,<4.21.0 || >4.21.0,<4.21.1 || >4.21.1,<4.21.2 || >4.21.2,<4.21.3 || >4.21.3,<4.21.4 || >4.21.4,<4.21.5 || >4.21.5,<5.0.0.dev0"
@@ -1556,6 +1621,7 @@ grpcio-gcp = ["grpcio-gcp (>=0.2.2,<1.0.dev0)"]
name = "google-auth"
version = "2.26.1"
description = "Google Authentication Library"
+category = "dev"
optional = false
python-versions = ">=3.7"
files = [
@@ -1579,6 +1645,7 @@ requests = ["requests (>=2.20.0,<3.0.0.dev0)"]
name = "googleapis-common-protos"
version = "1.62.0"
description = "Common protobufs used in Google APIs"
+category = "dev"
optional = false
python-versions = ">=3.7"
files = [
@@ -1596,6 +1663,7 @@ grpc = ["grpcio (>=1.44.0,<2.0.0.dev0)"]
name = "gptcache"
version = "0.1.43"
description = "GPTCache, a powerful caching library that can be used to speed up and lower the cost of chat applications that rely on the LLM service. GPTCache works as a memcache for AIGC applications, similar to how Redis works for traditional applications."
+category = "main"
optional = true
python-versions = ">=3.8.1"
files = [
@@ -1612,6 +1680,7 @@ requests = "*"
name = "gradientai"
version = "1.4.0"
description = "Gradient AI API"
+category = "main"
optional = true
python-versions = ">=3.8.1,<4.0.0"
files = [
@@ -1629,6 +1698,7 @@ urllib3 = ">=1.25.3"
name = "greenlet"
version = "3.0.3"
description = "Lightweight in-process concurrent programming"
+category = "main"
optional = false
python-versions = ">=3.7"
files = [
@@ -1700,6 +1770,7 @@ test = ["objgraph", "psutil"]
name = "grpcio"
version = "1.60.0"
description = "HTTP/2-based RPC framework"
+category = "dev"
optional = false
python-versions = ">=3.7"
files = [
@@ -1766,6 +1837,7 @@ protobuf = ["grpcio-tools (>=1.60.0)"]
name = "grpcio-status"
version = "1.60.0"
description = "Status proto mapping for gRPC"
+category = "dev"
optional = false
python-versions = ">=3.6"
files = [
@@ -1782,6 +1854,7 @@ protobuf = ">=4.21.6"
name = "guidance"
version = "0.0.64"
description = "A guidance language for controlling large language models."
+category = "main"
optional = true
python-versions = "*"
files = [
@@ -1811,6 +1884,7 @@ test = ["pytest", "pytest-cov", "torch", "transformers"]
name = "h11"
version = "0.14.0"
description = "A pure-Python, bring-your-own-I/O implementation of HTTP/1.1"
+category = "main"
optional = false
python-versions = ">=3.7"
files = [
@@ -1822,6 +1896,7 @@ files = [
name = "httpcore"
version = "1.0.2"
description = "A minimal low-level HTTP client."
+category = "main"
optional = false
python-versions = ">=3.8"
files = [
@@ -1836,13 +1911,14 @@ h11 = ">=0.13,<0.15"
[package.extras]
asyncio = ["anyio (>=4.0,<5.0)"]
http2 = ["h2 (>=3,<5)"]
-socks = ["socksio (==1.*)"]
+socks = ["socksio (>=1.0.0,<2.0.0)"]
trio = ["trio (>=0.22.0,<0.23.0)"]
[[package]]
name = "httpx"
version = "0.26.0"
description = "The next generation HTTP client."
+category = "main"
optional = false
python-versions = ">=3.8"
files = [
@@ -1853,20 +1929,21 @@ files = [
[package.dependencies]
anyio = "*"
certifi = "*"
-httpcore = "==1.*"
+httpcore = ">=1.0.0,<2.0.0"
idna = "*"
sniffio = "*"
[package.extras]
brotli = ["brotli", "brotlicffi"]
-cli = ["click (==8.*)", "pygments (==2.*)", "rich (>=10,<14)"]
+cli = ["click (>=8.0.0,<9.0.0)", "pygments (>=2.0.0,<3.0.0)", "rich (>=10,<14)"]
http2 = ["h2 (>=3,<5)"]
-socks = ["socksio (==1.*)"]
+socks = ["socksio (>=1.0.0,<2.0.0)"]
[[package]]
name = "huggingface-hub"
version = "0.20.2"
description = "Client library to download and publish models, datasets and other repos on the huggingface.co hub"
+category = "main"
optional = true
python-versions = ">=3.8.0"
files = [
@@ -1899,6 +1976,7 @@ typing = ["types-PyYAML", "types-requests", "types-simplejson", "types-toml", "t
name = "humanfriendly"
version = "10.0"
description = "Human friendly output for text interfaces using Python"
+category = "main"
optional = true
python-versions = ">=2.7, !=3.0.*, !=3.1.*, !=3.2.*, !=3.3.*, !=3.4.*"
files = [
@@ -1913,6 +1991,7 @@ pyreadline3 = {version = "*", markers = "sys_platform == \"win32\" and python_ve
name = "identify"
version = "2.5.33"
description = "File identification library for Python"
+category = "dev"
optional = false
python-versions = ">=3.8"
files = [
@@ -1927,6 +2006,7 @@ license = ["ukkonen"]
name = "idna"
version = "3.6"
description = "Internationalized Domain Names in Applications (IDNA)"
+category = "main"
optional = false
python-versions = ">=3.5"
files = [
@@ -1938,6 +2018,7 @@ files = [
name = "imagesize"
version = "1.4.1"
description = "Getting image size from png/jpeg/jpeg2000/gif file"
+category = "dev"
optional = false
python-versions = ">=2.7, !=3.0.*, !=3.1.*, !=3.2.*, !=3.3.*"
files = [
@@ -1949,6 +2030,7 @@ files = [
name = "importlib-metadata"
version = "7.0.1"
description = "Read metadata from Python packages"
+category = "dev"
optional = false
python-versions = ">=3.8"
files = [
@@ -1968,6 +2050,7 @@ testing = ["flufl.flake8", "importlib-resources (>=1.3)", "packaging", "pyfakefs
name = "importlib-resources"
version = "6.1.1"
description = "Read resources from Python packages"
+category = "dev"
optional = false
python-versions = ">=3.8"
files = [
@@ -1986,6 +2069,7 @@ testing = ["pytest (>=6)", "pytest-black (>=0.3.7)", "pytest-checkdocs (>=2.4)",
name = "iniconfig"
version = "2.0.0"
description = "brain-dead simple config-ini parsing"
+category = "dev"
optional = false
python-versions = ">=3.7"
files = [
@@ -1997,6 +2081,7 @@ files = [
name = "interegular"
version = "0.3.3"
description = "a regex intersection checker"
+category = "main"
optional = true
python-versions = ">=3.7"
files = [
@@ -2008,6 +2093,7 @@ files = [
name = "ipykernel"
version = "6.28.0"
description = "IPython Kernel for Jupyter"
+category = "dev"
optional = false
python-versions = ">=3.8"
files = [
@@ -2021,7 +2107,7 @@ comm = ">=0.1.1"
debugpy = ">=1.6.5"
ipython = ">=7.23.1"
jupyter-client = ">=6.1.12"
-jupyter-core = ">=4.12,<5.0.dev0 || >=5.1.dev0"
+jupyter-core = ">=4.12,<5.0.0 || >=5.1.0"
matplotlib-inline = ">=0.1"
nest-asyncio = "*"
packaging = "*"
@@ -2041,6 +2127,7 @@ test = ["flaky", "ipyparallel", "pre-commit", "pytest (>=7.0)", "pytest-asyncio"
name = "ipython"
version = "8.10.0"
description = "IPython: Productive Interactive Computing"
+category = "dev"
optional = false
python-versions = ">=3.8"
files = [
@@ -2079,6 +2166,7 @@ test-extra = ["curio", "matplotlib (!=3.2.0)", "nbformat", "numpy (>=1.21)", "pa
name = "ipywidgets"
version = "8.1.1"
description = "Jupyter interactive widgets"
+category = "dev"
optional = false
python-versions = ">=3.7"
files = [
@@ -2100,6 +2188,7 @@ test = ["ipykernel", "jsonschema", "pytest (>=3.6.0)", "pytest-cov", "pytz"]
name = "isoduration"
version = "20.11.0"
description = "Operations with ISO 8601 durations"
+category = "dev"
optional = false
python-versions = ">=3.7"
files = [
@@ -2114,6 +2203,7 @@ arrow = ">=0.15.0"
name = "isort"
version = "5.13.2"
description = "A Python utility / library to sort Python imports."
+category = "dev"
optional = false
python-versions = ">=3.8.0"
files = [
@@ -2128,6 +2218,7 @@ colors = ["colorama (>=0.4.6)"]
name = "jedi"
version = "0.19.1"
description = "An autocompletion tool for Python that can be used for text editors."
+category = "dev"
optional = false
python-versions = ">=3.6"
files = [
@@ -2147,6 +2238,7 @@ testing = ["Django", "attrs", "colorama", "docopt", "pytest (<7.0.0)"]
name = "jinja2"
version = "3.1.2"
description = "A very fast and expressive template engine."
+category = "main"
optional = false
python-versions = ">=3.7"
files = [
@@ -2164,6 +2256,7 @@ i18n = ["Babel (>=2.7)"]
name = "jmespath"
version = "1.0.1"
description = "JSON Matching Expressions"
+category = "dev"
optional = false
python-versions = ">=3.7"
files = [
@@ -2175,6 +2268,7 @@ files = [
name = "joblib"
version = "1.3.2"
description = "Lightweight pipelining with Python functions"
+category = "main"
optional = false
python-versions = ">=3.7"
files = [
@@ -2186,6 +2280,7 @@ files = [
name = "json5"
version = "0.9.14"
description = "A Python implementation of the JSON5 data format."
+category = "dev"
optional = false
python-versions = "*"
files = [
@@ -2200,6 +2295,7 @@ dev = ["hypothesis"]
name = "jsonpatch"
version = "1.33"
description = "Apply JSON-Patches (RFC 6902)"
+category = "main"
optional = true
python-versions = ">=2.7, !=3.0.*, !=3.1.*, !=3.2.*, !=3.3.*, !=3.4.*, !=3.5.*, !=3.6.*"
files = [
@@ -2214,6 +2310,7 @@ jsonpointer = ">=1.9"
name = "jsonpath-ng"
version = "1.6.0"
description = "A final implementation of JSONPath for Python that aims to be standard compliant, including arithmetic and binary comparison operators and providing clear AST for metaprogramming."
+category = "main"
optional = true
python-versions = "*"
files = [
@@ -2228,6 +2325,7 @@ ply = "*"
name = "jsonpointer"
version = "2.4"
description = "Identify specific nodes in a JSON document (RFC 6901)"
+category = "main"
optional = false
python-versions = ">=2.7, !=3.0.*, !=3.1.*, !=3.2.*, !=3.3.*, !=3.4.*, !=3.5.*, !=3.6.*"
files = [
@@ -2239,6 +2337,7 @@ files = [
name = "jsonschema"
version = "4.20.0"
description = "An implementation of JSON Schema validation for Python"
+category = "dev"
optional = false
python-versions = ">=3.8"
files = [
@@ -2270,6 +2369,7 @@ format-nongpl = ["fqdn", "idna", "isoduration", "jsonpointer (>1.13)", "rfc3339-
name = "jsonschema-specifications"
version = "2023.12.1"
description = "The JSON Schema meta-schemas and vocabularies, exposed as a Registry"
+category = "dev"
optional = false
python-versions = ">=3.8"
files = [
@@ -2285,6 +2385,7 @@ referencing = ">=0.31.0"
name = "jupyter"
version = "1.0.0"
description = "Jupyter metapackage. Install all the Jupyter components in one go."
+category = "dev"
optional = false
python-versions = "*"
files = [
@@ -2305,6 +2406,7 @@ qtconsole = "*"
name = "jupyter-cache"
version = "0.6.1"
description = "A defined interface for working with a cache of jupyter notebooks."
+category = "dev"
optional = false
python-versions = "~=3.8"
files = [
@@ -2332,6 +2434,7 @@ testing = ["coverage", "ipykernel", "jupytext", "matplotlib", "nbdime", "nbforma
name = "jupyter-client"
version = "8.6.0"
description = "Jupyter protocol implementation and client libraries"
+category = "dev"
optional = false
python-versions = ">=3.8"
files = [
@@ -2341,7 +2444,7 @@ files = [
[package.dependencies]
importlib-metadata = {version = ">=4.8.3", markers = "python_version < \"3.10\""}
-jupyter-core = ">=4.12,<5.0.dev0 || >=5.1.dev0"
+jupyter-core = ">=4.12,<5.0.0 || >=5.1.0"
python-dateutil = ">=2.8.2"
pyzmq = ">=23.0"
tornado = ">=6.2"
@@ -2355,6 +2458,7 @@ test = ["coverage", "ipykernel (>=6.14)", "mypy", "paramiko", "pre-commit", "pyt
name = "jupyter-console"
version = "6.6.3"
description = "Jupyter terminal console"
+category = "dev"
optional = false
python-versions = ">=3.7"
files = [
@@ -2366,7 +2470,7 @@ files = [
ipykernel = ">=6.14"
ipython = "*"
jupyter-client = ">=7.0.0"
-jupyter-core = ">=4.12,<5.0.dev0 || >=5.1.dev0"
+jupyter-core = ">=4.12,<5.0.0 || >=5.1.0"
prompt-toolkit = ">=3.0.30"
pygments = "*"
pyzmq = ">=17"
@@ -2379,6 +2483,7 @@ test = ["flaky", "pexpect", "pytest"]
name = "jupyter-core"
version = "5.7.1"
description = "Jupyter core package. A base package on which Jupyter projects rely."
+category = "dev"
optional = false
python-versions = ">=3.8"
files = [
@@ -2399,6 +2504,7 @@ test = ["ipykernel", "pre-commit", "pytest", "pytest-cov", "pytest-timeout"]
name = "jupyter-events"
version = "0.9.0"
description = "Jupyter Event System library"
+category = "dev"
optional = false
python-versions = ">=3.8"
files = [
@@ -2424,6 +2530,7 @@ test = ["click", "pre-commit", "pytest (>=7.0)", "pytest-asyncio (>=0.19.0)", "p
name = "jupyter-lsp"
version = "2.2.1"
description = "Multi-Language Server WebSocket proxy for Jupyter Notebook/Lab server"
+category = "dev"
optional = false
python-versions = ">=3.8"
files = [
@@ -2439,6 +2546,7 @@ jupyter-server = ">=1.1.2"
name = "jupyter-server"
version = "2.12.2"
description = "The backend—i.e. core services, APIs, and REST endpoints—to Jupyter web applications."
+category = "dev"
optional = false
python-versions = ">=3.8"
files = [
@@ -2451,7 +2559,7 @@ anyio = ">=3.1.0"
argon2-cffi = "*"
jinja2 = "*"
jupyter-client = ">=7.4.4"
-jupyter-core = ">=4.12,<5.0.dev0 || >=5.1.dev0"
+jupyter-core = ">=4.12,<5.0.0 || >=5.1.0"
jupyter-events = ">=0.9.0"
jupyter-server-terminals = "*"
nbconvert = ">=6.4.4"
@@ -2475,6 +2583,7 @@ test = ["flaky", "ipykernel", "pre-commit", "pytest (>=7.0)", "pytest-console-sc
name = "jupyter-server-terminals"
version = "0.5.1"
description = "A Jupyter Server Extension Providing Terminals."
+category = "dev"
optional = false
python-versions = ">=3.8"
files = [
@@ -2494,6 +2603,7 @@ test = ["jupyter-server (>=2.0.0)", "pytest (>=7.0)", "pytest-jupyter[server] (>
name = "jupyterlab"
version = "4.0.10"
description = "JupyterLab computational environment"
+category = "dev"
optional = false
python-versions = ">=3.8"
files = [
@@ -2527,6 +2637,7 @@ test = ["coverage", "pytest (>=7.0)", "pytest-check-links (>=0.7)", "pytest-cons
name = "jupyterlab-pygments"
version = "0.3.0"
description = "Pygments theme using JupyterLab CSS variables"
+category = "dev"
optional = false
python-versions = ">=3.8"
files = [
@@ -2538,6 +2649,7 @@ files = [
name = "jupyterlab-server"
version = "2.25.2"
description = "A set of server components for JupyterLab and JupyterLab like applications."
+category = "dev"
optional = false
python-versions = ">=3.8"
files = [
@@ -2564,6 +2676,7 @@ test = ["hatch", "ipykernel", "openapi-core (>=0.18.0,<0.19.0)", "openapi-spec-v
name = "jupyterlab-widgets"
version = "3.0.9"
description = "Jupyter interactive widgets for JupyterLab"
+category = "dev"
optional = false
python-versions = ">=3.7"
files = [
@@ -2575,6 +2688,7 @@ files = [
name = "langchain"
version = "0.1.0"
description = "Building applications with LLMs through composability"
+category = "main"
optional = true
python-versions = ">=3.8.1,<4.0"
files = [
@@ -2615,6 +2729,7 @@ text-helpers = ["chardet (>=5.1.0,<6.0.0)"]
name = "langchain-community"
version = "0.0.10"
description = "Community contributed LangChain integrations."
+category = "main"
optional = true
python-versions = ">=3.8.1,<4.0"
files = [
@@ -2641,6 +2756,7 @@ extended-testing = ["aiosqlite (>=0.19.0,<0.20.0)", "aleph-alpha-client (>=2.15.
name = "langchain-core"
version = "0.1.8"
description = "Building applications with LLMs through composability"
+category = "main"
optional = true
python-versions = ">=3.8.1,<4.0"
files = [
@@ -2665,6 +2781,7 @@ extended-testing = ["jinja2 (>=3,<4)"]
name = "langcodes"
version = "3.3.0"
description = "Tools for labeling human languages with IETF language tags"
+category = "main"
optional = true
python-versions = ">=3.6"
files = [
@@ -2679,6 +2796,7 @@ data = ["language-data (>=1.1,<2.0)"]
name = "langsmith"
version = "0.0.77"
description = "Client library to connect to the LangSmith LLM Tracing and Evaluation Platform."
+category = "main"
optional = true
python-versions = ">=3.8.1,<4.0"
files = [
@@ -2694,6 +2812,7 @@ requests = ">=2,<3"
name = "lazy-object-proxy"
version = "1.10.0"
description = "A fast and thorough lazy object proxy."
+category = "dev"
optional = false
python-versions = ">=3.8"
files = [
@@ -2740,6 +2859,7 @@ files = [
name = "livereload"
version = "2.6.3"
description = "Python LiveReload is an awesome tool for web developers"
+category = "dev"
optional = false
python-versions = "*"
files = [
@@ -2755,6 +2875,7 @@ tornado = {version = "*", markers = "python_version > \"2.7\""}
name = "lm-format-enforcer"
version = "0.4.3"
description = "Enforce the output format (JSON Schema, Regex etc) of a language model"
+category = "main"
optional = true
python-versions = ">=3.8,<4.0"
files = [
@@ -2770,6 +2891,7 @@ pydantic = ">=1.10.8"
name = "lxml"
version = "5.1.0"
description = "Powerful and Pythonic XML processing library combining libxml2/libxslt with the ElementTree API."
+category = "dev"
optional = false
python-versions = ">=3.6"
files = [
@@ -2859,6 +2981,7 @@ source = ["Cython (>=3.0.7)"]
name = "m2r2"
version = "0.3.2"
description = "Markdown and reStructuredText in a single file."
+category = "dev"
optional = false
python-versions = "*"
files = [
@@ -2874,6 +2997,7 @@ mistune = "0.8.4"
name = "markdown-it-py"
version = "2.2.0"
description = "Python port of markdown-it. Markdown parsing, done right!"
+category = "dev"
optional = false
python-versions = ">=3.7"
files = [
@@ -2898,6 +3022,7 @@ testing = ["coverage", "pytest", "pytest-cov", "pytest-regressions"]
name = "markupsafe"
version = "2.1.3"
description = "Safely add untrusted strings to HTML/XML markup."
+category = "main"
optional = false
python-versions = ">=3.7"
files = [
@@ -2921,16 +3046,6 @@ files = [
{file = "MarkupSafe-2.1.3-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:5bbe06f8eeafd38e5d0a4894ffec89378b6c6a625ff57e3028921f8ff59318ac"},
{file = "MarkupSafe-2.1.3-cp311-cp311-win32.whl", hash = "sha256:dd15ff04ffd7e05ffcb7fe79f1b98041b8ea30ae9234aed2a9168b5797c3effb"},
{file = "MarkupSafe-2.1.3-cp311-cp311-win_amd64.whl", hash = "sha256:134da1eca9ec0ae528110ccc9e48041e0828d79f24121a1a146161103c76e686"},
- {file = "MarkupSafe-2.1.3-cp312-cp312-macosx_10_9_universal2.whl", hash = "sha256:f698de3fd0c4e6972b92290a45bd9b1536bffe8c6759c62471efaa8acb4c37bc"},
- {file = "MarkupSafe-2.1.3-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:aa57bd9cf8ae831a362185ee444e15a93ecb2e344c8e52e4d721ea3ab6ef1823"},
- {file = "MarkupSafe-2.1.3-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:ffcc3f7c66b5f5b7931a5aa68fc9cecc51e685ef90282f4a82f0f5e9b704ad11"},
- {file = "MarkupSafe-2.1.3-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:47d4f1c5f80fc62fdd7777d0d40a2e9dda0a05883ab11374334f6c4de38adffd"},
- {file = "MarkupSafe-2.1.3-cp312-cp312-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:1f67c7038d560d92149c060157d623c542173016c4babc0c1913cca0564b9939"},
- {file = "MarkupSafe-2.1.3-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:9aad3c1755095ce347e26488214ef77e0485a3c34a50c5a5e2471dff60b9dd9c"},
- {file = "MarkupSafe-2.1.3-cp312-cp312-musllinux_1_1_i686.whl", hash = "sha256:14ff806850827afd6b07a5f32bd917fb7f45b046ba40c57abdb636674a8b559c"},
- {file = "MarkupSafe-2.1.3-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:8f9293864fe09b8149f0cc42ce56e3f0e54de883a9de90cd427f191c346eb2e1"},
- {file = "MarkupSafe-2.1.3-cp312-cp312-win32.whl", hash = "sha256:715d3562f79d540f251b99ebd6d8baa547118974341db04f5ad06d5ea3eb8007"},
- {file = "MarkupSafe-2.1.3-cp312-cp312-win_amd64.whl", hash = "sha256:1b8dd8c3fd14349433c79fa8abeb573a55fc0fdd769133baac1f5e07abf54aeb"},
{file = "MarkupSafe-2.1.3-cp37-cp37m-macosx_10_9_x86_64.whl", hash = "sha256:8e254ae696c88d98da6555f5ace2279cf7cd5b3f52be2b5cf97feafe883b58d2"},
{file = "MarkupSafe-2.1.3-cp37-cp37m-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:cb0932dc158471523c9637e807d9bfb93e06a95cbf010f1a38b98623b929ef2b"},
{file = "MarkupSafe-2.1.3-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:9402b03f1a1b4dc4c19845e5c749e3ab82d5078d16a2a4c2cd2df62d57bb0707"},
@@ -2967,6 +3082,7 @@ files = [
name = "marshmallow"
version = "3.20.1"
description = "A lightweight library for converting complex datatypes to and from native Python datatypes."
+category = "main"
optional = false
python-versions = ">=3.8"
files = [
@@ -2987,6 +3103,7 @@ tests = ["pytest", "pytz", "simplejson"]
name = "matplotlib-inline"
version = "0.1.6"
description = "Inline Matplotlib backend for Jupyter"
+category = "dev"
optional = false
python-versions = ">=3.5"
files = [
@@ -3001,6 +3118,7 @@ traitlets = "*"
name = "mccabe"
version = "0.7.0"
description = "McCabe checker, plugin for flake8"
+category = "dev"
optional = false
python-versions = ">=3.6"
files = [
@@ -3012,6 +3130,7 @@ files = [
name = "mdit-py-plugins"
version = "0.3.5"
description = "Collection of plugins for markdown-it-py"
+category = "dev"
optional = false
python-versions = ">=3.7"
files = [
@@ -3031,6 +3150,7 @@ testing = ["coverage", "pytest", "pytest-cov", "pytest-regressions"]
name = "mdurl"
version = "0.1.2"
description = "Markdown URL utilities"
+category = "dev"
optional = false
python-versions = ">=3.7"
files = [
@@ -3042,6 +3162,7 @@ files = [
name = "mistune"
version = "0.8.4"
description = "The fastest markdown parser in pure Python"
+category = "dev"
optional = false
python-versions = "*"
files = [
@@ -3053,6 +3174,7 @@ files = [
name = "mpmath"
version = "1.3.0"
description = "Python library for arbitrary-precision floating-point arithmetic"
+category = "main"
optional = true
python-versions = "*"
files = [
@@ -3070,6 +3192,7 @@ tests = ["pytest (>=4.6)"]
name = "msal"
version = "1.26.0"
description = "The Microsoft Authentication Library (MSAL) for Python library enables your app to access the Microsoft Cloud by supporting authentication of users with Microsoft Azure Active Directory accounts (AAD) and Microsoft Accounts (MSA) using industry standard OAuth2 and OpenID Connect."
+category = "main"
optional = true
python-versions = ">=2.7"
files = [
@@ -3089,6 +3212,7 @@ broker = ["pymsalruntime (>=0.13.2,<0.14)"]
name = "multidict"
version = "6.0.4"
description = "multidict implementation"
+category = "main"
optional = false
python-versions = ">=3.7"
files = [
@@ -3172,6 +3296,7 @@ files = [
name = "multiprocess"
version = "0.70.15"
description = "better multiprocessing and multithreading in Python"
+category = "main"
optional = true
python-versions = ">=3.7"
files = [
@@ -3200,6 +3325,7 @@ dill = ">=0.3.7"
name = "murmurhash"
version = "1.0.10"
description = "Cython bindings for MurmurHash"
+category = "main"
optional = true
python-versions = ">=3.6"
files = [
@@ -3242,6 +3368,7 @@ files = [
name = "mypy"
version = "0.991"
description = "Optional static typing for Python"
+category = "dev"
optional = false
python-versions = ">=3.7"
files = [
@@ -3292,6 +3419,7 @@ reports = ["lxml"]
name = "mypy-extensions"
version = "1.0.0"
description = "Type system extensions for programs checked with the mypy type checker."
+category = "main"
optional = false
python-versions = ">=3.5"
files = [
@@ -3303,6 +3431,7 @@ files = [
name = "myst-nb"
version = "0.17.2"
description = "A Jupyter Notebook Sphinx reader built on top of the MyST markdown parser."
+category = "dev"
optional = false
python-versions = ">=3.7"
files = [
@@ -3331,6 +3460,7 @@ testing = ["beautifulsoup4", "coverage (>=6.4,<8.0)", "ipykernel (>=5.5,<6.0)",
name = "myst-parser"
version = "0.18.1"
description = "An extended commonmark compliant parser, with bridges to docutils & sphinx."
+category = "dev"
optional = false
python-versions = ">=3.7"
files = [
@@ -3357,6 +3487,7 @@ testing = ["beautifulsoup4", "coverage[toml]", "pytest (>=6,<7)", "pytest-cov",
name = "nbclient"
version = "0.7.4"
description = "A client library for executing notebooks. Formerly nbconvert's ExecutePreprocessor."
+category = "dev"
optional = false
python-versions = ">=3.7.0"
files = [
@@ -3366,7 +3497,7 @@ files = [
[package.dependencies]
jupyter-client = ">=6.1.12"
-jupyter-core = ">=4.12,<5.0.dev0 || >=5.1.dev0"
+jupyter-core = ">=4.12,<5.0.0 || >=5.1.0"
nbformat = ">=5.1"
traitlets = ">=5.3"
@@ -3379,6 +3510,7 @@ test = ["flaky", "ipykernel", "ipython", "ipywidgets", "nbconvert (>=7.0.0)", "p
name = "nbconvert"
version = "6.5.4"
description = "Converting Jupyter Notebooks"
+category = "dev"
optional = false
python-versions = ">=3.7"
files = [
@@ -3416,6 +3548,7 @@ webpdf = ["pyppeteer (>=1,<1.1)"]
name = "nbformat"
version = "5.9.2"
description = "The Jupyter Notebook format"
+category = "dev"
optional = false
python-versions = ">=3.8"
files = [
@@ -3437,6 +3570,7 @@ test = ["pep440", "pre-commit", "pytest", "testpath"]
name = "nest-asyncio"
version = "1.5.8"
description = "Patch asyncio to allow nested event loops"
+category = "main"
optional = false
python-versions = ">=3.5"
files = [
@@ -3448,7 +3582,8 @@ files = [
name = "networkx"
version = "3.1"
description = "Python package for creating and manipulating graphs and networks"
-optional = true
+category = "main"
+optional = false
python-versions = ">=3.8"
files = [
{file = "networkx-3.1-py3-none-any.whl", hash = "sha256:4f33f68cb2afcf86f28a45f43efc27a9386b535d567d2127f8f61d51dec58d36"},
@@ -3466,6 +3601,7 @@ test = ["codecov (>=2.1)", "pytest (>=7.2)", "pytest-cov (>=4.0)"]
name = "nltk"
version = "3.8.1"
description = "Natural Language Toolkit"
+category = "main"
optional = false
python-versions = ">=3.7"
files = [
@@ -3491,6 +3627,7 @@ twitter = ["twython"]
name = "nodeenv"
version = "1.8.0"
description = "Node.js virtual environment builder"
+category = "dev"
optional = false
python-versions = ">=2.7,!=3.0.*,!=3.1.*,!=3.2.*,!=3.3.*,!=3.4.*,!=3.5.*,!=3.6.*"
files = [
@@ -3505,6 +3642,7 @@ setuptools = "*"
name = "notebook"
version = "7.0.6"
description = "Jupyter Notebook - A web-based notebook environment for interactive computing"
+category = "dev"
optional = false
python-versions = ">=3.8"
files = [
@@ -3528,6 +3666,7 @@ test = ["importlib-resources (>=5.0)", "ipykernel", "jupyter-server[test] (>=2.4
name = "notebook-shim"
version = "0.2.3"
description = "A shim layer for notebook traits and config"
+category = "dev"
optional = false
python-versions = ">=3.7"
files = [
@@ -3545,6 +3684,7 @@ test = ["pytest", "pytest-console-scripts", "pytest-jupyter", "pytest-tornasync"
name = "numpy"
version = "1.24.4"
description = "Fundamental package for array computing in Python"
+category = "main"
optional = false
python-versions = ">=3.8"
files = [
@@ -3582,6 +3722,7 @@ files = [
name = "nvidia-cublas-cu12"
version = "12.1.3.1"
description = "CUBLAS native runtime libraries"
+category = "main"
optional = true
python-versions = ">=3"
files = [
@@ -3593,6 +3734,7 @@ files = [
name = "nvidia-cuda-cupti-cu12"
version = "12.1.105"
description = "CUDA profiling tools runtime libs."
+category = "main"
optional = true
python-versions = ">=3"
files = [
@@ -3604,6 +3746,7 @@ files = [
name = "nvidia-cuda-nvrtc-cu12"
version = "12.1.105"
description = "NVRTC native runtime libraries"
+category = "main"
optional = true
python-versions = ">=3"
files = [
@@ -3615,6 +3758,7 @@ files = [
name = "nvidia-cuda-runtime-cu12"
version = "12.1.105"
description = "CUDA Runtime native Libraries"
+category = "main"
optional = true
python-versions = ">=3"
files = [
@@ -3626,6 +3770,7 @@ files = [
name = "nvidia-cudnn-cu12"
version = "8.9.2.26"
description = "cuDNN runtime libraries"
+category = "main"
optional = true
python-versions = ">=3"
files = [
@@ -3639,6 +3784,7 @@ nvidia-cublas-cu12 = "*"
name = "nvidia-cufft-cu12"
version = "11.0.2.54"
description = "CUFFT native runtime libraries"
+category = "main"
optional = true
python-versions = ">=3"
files = [
@@ -3650,6 +3796,7 @@ files = [
name = "nvidia-curand-cu12"
version = "10.3.2.106"
description = "CURAND native runtime libraries"
+category = "main"
optional = true
python-versions = ">=3"
files = [
@@ -3661,6 +3808,7 @@ files = [
name = "nvidia-cusolver-cu12"
version = "11.4.5.107"
description = "CUDA solver native runtime libraries"
+category = "main"
optional = true
python-versions = ">=3"
files = [
@@ -3677,6 +3825,7 @@ nvidia-nvjitlink-cu12 = "*"
name = "nvidia-cusparse-cu12"
version = "12.1.0.106"
description = "CUSPARSE native runtime libraries"
+category = "main"
optional = true
python-versions = ">=3"
files = [
@@ -3691,6 +3840,7 @@ nvidia-nvjitlink-cu12 = "*"
name = "nvidia-nccl-cu12"
version = "2.18.1"
description = "NVIDIA Collective Communication Library (NCCL) Runtime"
+category = "main"
optional = true
python-versions = ">=3"
files = [
@@ -3701,6 +3851,7 @@ files = [
name = "nvidia-nvjitlink-cu12"
version = "12.3.101"
description = "Nvidia JIT LTO Library"
+category = "main"
optional = true
python-versions = ">=3"
files = [
@@ -3712,6 +3863,7 @@ files = [
name = "nvidia-nvtx-cu12"
version = "12.1.105"
description = "NVIDIA Tools Extension"
+category = "main"
optional = true
python-versions = ">=3"
files = [
@@ -3723,6 +3875,7 @@ files = [
name = "onnx"
version = "1.15.0"
description = "Open Neural Network Exchange"
+category = "main"
optional = true
python-versions = ">=3.8"
files = [
@@ -3764,6 +3917,7 @@ reference = ["Pillow", "google-re2"]
name = "onnxruntime"
version = "1.16.3"
description = "ONNX Runtime is a runtime accelerator for Machine Learning models"
+category = "main"
optional = true
python-versions = "*"
files = [
@@ -3805,6 +3959,7 @@ sympy = "*"
name = "openai"
version = "1.6.1"
description = "The official Python library for the openai API"
+category = "main"
optional = false
python-versions = ">=3.7.1"
files = [
@@ -3828,6 +3983,7 @@ datalib = ["numpy (>=1)", "pandas (>=1.2.3)", "pandas-stubs (>=1.1.0.11)"]
name = "optimum"
version = "1.16.1"
description = "Optimum Library is an extension of the Hugging Face Transformers library, providing a framework to integrate third-party libraries from Hardware Partners and interface with their specific functionality."
+category = "main"
optional = true
python-versions = ">=3.7.0"
files = [
@@ -3879,6 +4035,7 @@ tests = ["Pillow", "accelerate", "diffusers (>=0.17.0)", "einops", "invisible-wa
name = "overrides"
version = "7.4.0"
description = "A decorator to automatically detect mismatch when overriding a method."
+category = "dev"
optional = false
python-versions = ">=3.6"
files = [
@@ -3890,6 +4047,7 @@ files = [
name = "packaging"
version = "23.2"
description = "Core utilities for Python packages"
+category = "main"
optional = false
python-versions = ">=3.7"
files = [
@@ -3901,6 +4059,7 @@ files = [
name = "pandas"
version = "2.0.3"
description = "Powerful data structures for data analysis, time series, and statistics"
+category = "main"
optional = false
python-versions = ">=3.8"
files = [
@@ -3934,7 +4093,7 @@ files = [
[package.dependencies]
numpy = [
{version = ">=1.20.3", markers = "python_version < \"3.10\""},
- {version = ">=1.21.0", markers = "python_version >= \"3.10\" and python_version < \"3.11\""},
+ {version = ">=1.21.0", markers = "python_version >= \"3.10\""},
{version = ">=1.23.2", markers = "python_version >= \"3.11\""},
]
python-dateutil = ">=2.8.2"
@@ -3968,6 +4127,7 @@ xml = ["lxml (>=4.6.3)"]
name = "pandocfilters"
version = "1.5.0"
description = "Utilities for writing pandoc filters in python"
+category = "dev"
optional = false
python-versions = ">=2.7, !=3.0.*, !=3.1.*, !=3.2.*, !=3.3.*"
files = [
@@ -3979,6 +4139,7 @@ files = [
name = "parso"
version = "0.8.3"
description = "A Python Parser"
+category = "dev"
optional = false
python-versions = ">=3.6"
files = [
@@ -3994,6 +4155,7 @@ testing = ["docopt", "pytest (<6.0.0)"]
name = "pathspec"
version = "0.12.1"
description = "Utility library for gitignore style pattern matching of file paths."
+category = "dev"
optional = false
python-versions = ">=3.8"
files = [
@@ -4005,6 +4167,7 @@ files = [
name = "pexpect"
version = "4.9.0"
description = "Pexpect allows easy control of interactive console applications."
+category = "dev"
optional = false
python-versions = "*"
files = [
@@ -4019,6 +4182,7 @@ ptyprocess = ">=0.5"
name = "pgvector"
version = "0.1.8"
description = "pgvector support for Python"
+category = "main"
optional = true
python-versions = ">=3.6"
files = [
@@ -4032,6 +4196,7 @@ numpy = "*"
name = "pickleshare"
version = "0.7.5"
description = "Tiny 'shelve'-like database with concurrency support"
+category = "dev"
optional = false
python-versions = "*"
files = [
@@ -4043,6 +4208,7 @@ files = [
name = "pkgutil-resolve-name"
version = "1.3.10"
description = "Resolve a name to an object."
+category = "dev"
optional = false
python-versions = ">=3.6"
files = [
@@ -4054,6 +4220,7 @@ files = [
name = "platformdirs"
version = "4.1.0"
description = "A small Python package for determining appropriate platform-specific dirs, e.g. a \"user data dir\"."
+category = "main"
optional = false
python-versions = ">=3.8"
files = [
@@ -4069,6 +4236,7 @@ test = ["appdirs (==1.4.4)", "covdefaults (>=2.3)", "pytest (>=7.4)", "pytest-co
name = "pluggy"
version = "1.3.0"
description = "plugin and hook calling mechanisms for python"
+category = "dev"
optional = false
python-versions = ">=3.8"
files = [
@@ -4084,6 +4252,7 @@ testing = ["pytest", "pytest-benchmark"]
name = "ply"
version = "3.11"
description = "Python Lex & Yacc"
+category = "main"
optional = true
python-versions = "*"
files = [
@@ -4095,6 +4264,7 @@ files = [
name = "pre-commit"
version = "3.2.0"
description = "A framework for managing and maintaining multi-language pre-commit hooks."
+category = "dev"
optional = false
python-versions = ">=3.8"
files = [
@@ -4113,6 +4283,7 @@ virtualenv = ">=20.10.0"
name = "preshed"
version = "3.0.9"
description = "Cython hash table that trusts the keys are pre-hashed"
+category = "main"
optional = true
python-versions = ">=3.6"
files = [
@@ -4159,6 +4330,7 @@ murmurhash = ">=0.28.0,<1.1.0"
name = "prometheus-client"
version = "0.19.0"
description = "Python client for the Prometheus monitoring system."
+category = "dev"
optional = false
python-versions = ">=3.8"
files = [
@@ -4173,6 +4345,7 @@ twisted = ["twisted"]
name = "prompt-toolkit"
version = "3.0.43"
description = "Library for building powerful interactive command lines in Python"
+category = "dev"
optional = false
python-versions = ">=3.7.0"
files = [
@@ -4187,6 +4360,7 @@ wcwidth = "*"
name = "proto-plus"
version = "1.23.0"
description = "Beautiful, Pythonic protocol buffers."
+category = "dev"
optional = false
python-versions = ">=3.6"
files = [
@@ -4204,6 +4378,7 @@ testing = ["google-api-core[grpc] (>=1.31.5)"]
name = "protobuf"
version = "4.25.1"
description = ""
+category = "main"
optional = false
python-versions = ">=3.8"
files = [
@@ -4224,6 +4399,7 @@ files = [
name = "psutil"
version = "5.9.7"
description = "Cross-platform lib for process and system monitoring in Python."
+category = "main"
optional = false
python-versions = ">=2.7, !=3.0.*, !=3.1.*, !=3.2.*, !=3.3.*, !=3.4.*, !=3.5.*"
files = [
@@ -4252,6 +4428,7 @@ test = ["enum34", "ipaddress", "mock", "pywin32", "wmi"]
name = "psycopg-binary"
version = "3.1.17"
description = "PostgreSQL database adapter for Python -- C optimisation distribution"
+category = "main"
optional = true
python-versions = ">=3.7"
files = [
@@ -4326,6 +4503,7 @@ files = [
name = "psycopg2"
version = "2.9.9"
description = "psycopg2 - Python-PostgreSQL Database Adapter"
+category = "main"
optional = true
python-versions = ">=3.7"
files = [
@@ -4348,6 +4526,7 @@ files = [
name = "ptyprocess"
version = "0.7.0"
description = "Run a subprocess in a pseudo terminal"
+category = "dev"
optional = false
python-versions = "*"
files = [
@@ -4359,6 +4538,7 @@ files = [
name = "pure-eval"
version = "0.2.2"
description = "Safely evaluate AST nodes without side effects"
+category = "dev"
optional = false
python-versions = "*"
files = [
@@ -4373,6 +4553,7 @@ tests = ["pytest"]
name = "pyarrow"
version = "14.0.2"
description = "Python library for Apache Arrow"
+category = "main"
optional = true
python-versions = ">=3.8"
files = [
@@ -4421,6 +4602,7 @@ numpy = ">=1.16.6"
name = "pyasn1"
version = "0.5.1"
description = "Pure-Python implementation of ASN.1 types and DER/BER/CER codecs (X.208)"
+category = "dev"
optional = false
python-versions = "!=3.0.*,!=3.1.*,!=3.2.*,!=3.3.*,!=3.4.*,!=3.5.*,>=2.7"
files = [
@@ -4432,6 +4614,7 @@ files = [
name = "pyasn1-modules"
version = "0.3.0"
description = "A collection of ASN.1-based protocols modules"
+category = "dev"
optional = false
python-versions = "!=3.0.*,!=3.1.*,!=3.2.*,!=3.3.*,!=3.4.*,!=3.5.*,>=2.7"
files = [
@@ -4446,6 +4629,7 @@ pyasn1 = ">=0.4.6,<0.6.0"
name = "pycparser"
version = "2.21"
description = "C parser in Python"
+category = "main"
optional = false
python-versions = ">=2.7, !=3.0.*, !=3.1.*, !=3.2.*, !=3.3.*"
files = [
@@ -4457,6 +4641,7 @@ files = [
name = "pydantic"
version = "1.10.13"
description = "Data validation and settings management using python type hints"
+category = "main"
optional = false
python-versions = ">=3.7"
files = [
@@ -4509,6 +4694,7 @@ email = ["email-validator (>=1.0.3)"]
name = "pygments"
version = "2.17.2"
description = "Pygments is a syntax highlighting package written in Python."
+category = "dev"
optional = false
python-versions = ">=3.7"
files = [
@@ -4524,6 +4710,7 @@ windows-terminal = ["colorama (>=0.4.6)"]
name = "pygtrie"
version = "2.5.0"
description = "A pure Python trie data structure implementation."
+category = "main"
optional = true
python-versions = "*"
files = [
@@ -4535,6 +4722,7 @@ files = [
name = "pyjwt"
version = "2.8.0"
description = "JSON Web Token implementation in Python"
+category = "main"
optional = true
python-versions = ">=3.7"
files = [
@@ -4555,6 +4743,7 @@ tests = ["coverage[toml] (==5.0.4)", "pytest (>=6.0.0,<7.0.0)"]
name = "pylint"
version = "2.15.10"
description = "python code static checker"
+category = "dev"
optional = false
python-versions = ">=3.7.2"
files = [
@@ -4584,6 +4773,7 @@ testutils = ["gitpython (>3)"]
name = "pymongo"
version = "4.6.1"
description = "Python driver for MongoDB <http://www.mongodb.org>"
+category = "dev"
optional = false
python-versions = ">=3.7"
files = [
@@ -4687,6 +4877,7 @@ zstd = ["zstandard"]
name = "pyparsing"
version = "3.1.1"
description = "pyparsing module - Classes and methods to define and execute parsing grammars"
+category = "main"
optional = true
python-versions = ">=3.6.8"
files = [
@@ -4701,6 +4892,7 @@ diagrams = ["jinja2", "railroad-diagrams"]
name = "pyreadline3"
version = "3.4.1"
description = "A python implementation of GNU readline."
+category = "main"
optional = true
python-versions = "*"
files = [
@@ -4712,6 +4904,7 @@ files = [
name = "pytest"
version = "7.2.1"
description = "pytest: simple powerful testing with Python"
+category = "dev"
optional = false
python-versions = ">=3.7"
files = [
@@ -4735,6 +4928,7 @@ testing = ["argcomplete", "hypothesis (>=3.56)", "mock", "nose", "pygments (>=2.
name = "pytest-asyncio"
version = "0.21.0"
description = "Pytest support for asyncio"
+category = "dev"
optional = false
python-versions = ">=3.7"
files = [
@@ -4753,6 +4947,7 @@ testing = ["coverage (>=6.2)", "flaky (>=3.5.0)", "hypothesis (>=5.7.1)", "mypy
name = "pytest-dotenv"
version = "0.5.2"
description = "A py.test plugin that parses environment files before running tests"
+category = "dev"
optional = false
python-versions = "*"
files = [
@@ -4768,6 +4963,7 @@ python-dotenv = ">=0.9.1"
name = "pytest-mock"
version = "3.11.1"
description = "Thin-wrapper around the mock package for easier use with pytest"
+category = "dev"
optional = false
python-versions = ">=3.7"
files = [
@@ -4785,6 +4981,7 @@ dev = ["pre-commit", "pytest-asyncio", "tox"]
name = "python-dateutil"
version = "2.8.2"
description = "Extensions to the standard Python datetime module"
+category = "main"
optional = false
python-versions = "!=3.0.*,!=3.1.*,!=3.2.*,>=2.7"
files = [
@@ -4799,6 +4996,7 @@ six = ">=1.5"
name = "python-dotenv"
version = "1.0.0"
description = "Read key-value pairs from a .env file and set them as environment variables"
+category = "dev"
optional = false
python-versions = ">=3.8"
files = [
@@ -4813,6 +5011,7 @@ cli = ["click (>=5.0)"]
name = "python-json-logger"
version = "2.0.7"
description = "A python library adding a json log formatter"
+category = "dev"
optional = false
python-versions = ">=3.6"
files = [
@@ -4824,6 +5023,7 @@ files = [
name = "pytz"
version = "2023.3.post1"
description = "World timezone definitions, modern and historical"
+category = "main"
optional = false
python-versions = "*"
files = [
@@ -4835,6 +5035,7 @@ files = [
name = "pywin32"
version = "306"
description = "Python for Window Extensions"
+category = "dev"
optional = false
python-versions = "*"
files = [
@@ -4858,6 +5059,7 @@ files = [
name = "pywinpty"
version = "2.0.12"
description = "Pseudo terminal support for Windows from Python."
+category = "dev"
optional = false
python-versions = ">=3.8"
files = [
@@ -4873,6 +5075,7 @@ files = [
name = "pyyaml"
version = "6.0.1"
description = "YAML parser and emitter for Python"
+category = "main"
optional = false
python-versions = ">=3.6"
files = [
@@ -4881,7 +5084,6 @@ files = [
{file = "PyYAML-6.0.1-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:69b023b2b4daa7548bcfbd4aa3da05b3a74b772db9e23b982788168117739938"},
{file = "PyYAML-6.0.1-cp310-cp310-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:81e0b275a9ecc9c0c0c07b4b90ba548307583c125f54d5b6946cfee6360c733d"},
{file = "PyYAML-6.0.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ba336e390cd8e4d1739f42dfe9bb83a3cc2e80f567d8805e11b46f4a943f5515"},
- {file = "PyYAML-6.0.1-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:326c013efe8048858a6d312ddd31d56e468118ad4cdeda36c719bf5bb6192290"},
{file = "PyYAML-6.0.1-cp310-cp310-win32.whl", hash = "sha256:bd4af7373a854424dabd882decdc5579653d7868b8fb26dc7d0e99f823aa5924"},
{file = "PyYAML-6.0.1-cp310-cp310-win_amd64.whl", hash = "sha256:fd1592b3fdf65fff2ad0004b5e363300ef59ced41c2e6b3a99d4089fa8c5435d"},
{file = "PyYAML-6.0.1-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:6965a7bc3cf88e5a1c3bd2e0b5c22f8d677dc88a455344035f03399034eb3007"},
@@ -4889,15 +5091,8 @@ files = [
{file = "PyYAML-6.0.1-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:42f8152b8dbc4fe7d96729ec2b99c7097d656dc1213a3229ca5383f973a5ed6d"},
{file = "PyYAML-6.0.1-cp311-cp311-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:062582fca9fabdd2c8b54a3ef1c978d786e0f6b3a1510e0ac93ef59e0ddae2bc"},
{file = "PyYAML-6.0.1-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:d2b04aac4d386b172d5b9692e2d2da8de7bfb6c387fa4f801fbf6fb2e6ba4673"},
- {file = "PyYAML-6.0.1-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:e7d73685e87afe9f3b36c799222440d6cf362062f78be1013661b00c5c6f678b"},
{file = "PyYAML-6.0.1-cp311-cp311-win32.whl", hash = "sha256:1635fd110e8d85d55237ab316b5b011de701ea0f29d07611174a1b42f1444741"},
{file = "PyYAML-6.0.1-cp311-cp311-win_amd64.whl", hash = "sha256:bf07ee2fef7014951eeb99f56f39c9bb4af143d8aa3c21b1677805985307da34"},
- {file = "PyYAML-6.0.1-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:855fb52b0dc35af121542a76b9a84f8d1cd886ea97c84703eaa6d88e37a2ad28"},
- {file = "PyYAML-6.0.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:40df9b996c2b73138957fe23a16a4f0ba614f4c0efce1e9406a184b6d07fa3a9"},
- {file = "PyYAML-6.0.1-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:6c22bec3fbe2524cde73d7ada88f6566758a8f7227bfbf93a408a9d86bcc12a0"},
- {file = "PyYAML-6.0.1-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:8d4e9c88387b0f5c7d5f281e55304de64cf7f9c0021a3525bd3b1c542da3b0e4"},
- {file = "PyYAML-6.0.1-cp312-cp312-win32.whl", hash = "sha256:d483d2cdf104e7c9fa60c544d92981f12ad66a457afae824d146093b8c294c54"},
- {file = "PyYAML-6.0.1-cp312-cp312-win_amd64.whl", hash = "sha256:0d3304d8c0adc42be59c5f8a4d9e3d7379e6955ad754aa9d6ab7a398b59dd1df"},
{file = "PyYAML-6.0.1-cp36-cp36m-macosx_10_9_x86_64.whl", hash = "sha256:50550eb667afee136e9a77d6dc71ae76a44df8b3e51e41b77f6de2932bfe0f47"},
{file = "PyYAML-6.0.1-cp36-cp36m-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:1fe35611261b29bd1de0070f0b2f47cb6ff71fa6595c077e42bd0c419fa27b98"},
{file = "PyYAML-6.0.1-cp36-cp36m-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:704219a11b772aea0d8ecd7058d0082713c3562b4e271b849ad7dc4a5c90c13c"},
@@ -4914,7 +5109,6 @@ files = [
{file = "PyYAML-6.0.1-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a0cd17c15d3bb3fa06978b4e8958dcdc6e0174ccea823003a106c7d4d7899ac5"},
{file = "PyYAML-6.0.1-cp38-cp38-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:28c119d996beec18c05208a8bd78cbe4007878c6dd15091efb73a30e90539696"},
{file = "PyYAML-6.0.1-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:7e07cbde391ba96ab58e532ff4803f79c4129397514e1413a7dc761ccd755735"},
- {file = "PyYAML-6.0.1-cp38-cp38-musllinux_1_1_x86_64.whl", hash = "sha256:49a183be227561de579b4a36efbb21b3eab9651dd81b1858589f796549873dd6"},
{file = "PyYAML-6.0.1-cp38-cp38-win32.whl", hash = "sha256:184c5108a2aca3c5b3d3bf9395d50893a7ab82a38004c8f61c258d4428e80206"},
{file = "PyYAML-6.0.1-cp38-cp38-win_amd64.whl", hash = "sha256:1e2722cc9fbb45d9b87631ac70924c11d3a401b2d7f410cc0e3bbf249f2dca62"},
{file = "PyYAML-6.0.1-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:9eb6caa9a297fc2c2fb8862bc5370d0303ddba53ba97e71f08023b6cd73d16a8"},
@@ -4922,7 +5116,6 @@ files = [
{file = "PyYAML-6.0.1-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:5773183b6446b2c99bb77e77595dd486303b4faab2b086e7b17bc6bef28865f6"},
{file = "PyYAML-6.0.1-cp39-cp39-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:b786eecbdf8499b9ca1d697215862083bd6d2a99965554781d0d8d1ad31e13a0"},
{file = "PyYAML-6.0.1-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:bc1bf2925a1ecd43da378f4db9e4f799775d6367bdb94671027b73b393a7c42c"},
- {file = "PyYAML-6.0.1-cp39-cp39-musllinux_1_1_x86_64.whl", hash = "sha256:04ac92ad1925b2cff1db0cfebffb6ffc43457495c9b3c39d3fcae417d7125dc5"},
{file = "PyYAML-6.0.1-cp39-cp39-win32.whl", hash = "sha256:faca3bdcf85b2fc05d06ff3fbc1f83e1391b3e724afa3feba7d13eeab355484c"},
{file = "PyYAML-6.0.1-cp39-cp39-win_amd64.whl", hash = "sha256:510c9deebc5c0225e8c96813043e62b680ba2f9c50a08d3724c7f28a747d1486"},
{file = "PyYAML-6.0.1.tar.gz", hash = "sha256:bfdf460b1736c775f2ba9f6a92bca30bc2095067b8a9d77876d1fad6cc3b4a43"},
@@ -4932,6 +5125,7 @@ files = [
name = "pyzmq"
version = "25.1.2"
description = "Python bindings for 0MQ"
+category = "dev"
optional = false
python-versions = ">=3.6"
files = [
@@ -5037,6 +5231,7 @@ cffi = {version = "*", markers = "implementation_name == \"pypy\""}
name = "qtconsole"
version = "5.5.1"
description = "Jupyter Qt console"
+category = "dev"
optional = false
python-versions = ">= 3.8"
files = [
@@ -5062,6 +5257,7 @@ test = ["flaky", "pytest", "pytest-qt"]
name = "qtpy"
version = "2.4.1"
description = "Provides an abstraction layer on top of the various Qt bindings (PyQt5/6 and PySide2/6)."
+category = "dev"
optional = false
python-versions = ">=3.7"
files = [
@@ -5079,6 +5275,7 @@ test = ["pytest (>=6,!=7.0.0,!=7.0.1)", "pytest-cov (>=3.0.0)", "pytest-qt"]
name = "rake-nltk"
version = "1.0.6"
description = "RAKE short for Rapid Automatic Keyword Extraction algorithm, is a domain independent keyword extraction algorithm which tries to determine key phrases in a body of text by analyzing the frequency of word appearance and its co-occurance with other words in the text."
+category = "dev"
optional = false
python-versions = ">=3.6,<4.0"
files = [
@@ -5093,6 +5290,7 @@ nltk = ">=3.6.2,<4.0.0"
name = "rank-bm25"
version = "0.2.2"
description = "Various BM25 algorithms for document ranking"
+category = "main"
optional = true
python-versions = "*"
files = [
@@ -5110,6 +5308,7 @@ dev = ["pytest"]
name = "referencing"
version = "0.32.1"
description = "JSON Referencing + Python"
+category = "dev"
optional = false
python-versions = ">=3.8"
files = [
@@ -5125,6 +5324,7 @@ rpds-py = ">=0.7.0"
name = "regex"
version = "2023.12.25"
description = "Alternative regular expression module, to replace re."
+category = "main"
optional = false
python-versions = ">=3.7"
files = [
@@ -5227,6 +5427,7 @@ files = [
name = "requests"
version = "2.31.0"
description = "Python HTTP for Humans."
+category = "main"
optional = false
python-versions = ">=3.7"
files = [
@@ -5248,6 +5449,7 @@ use-chardet-on-py3 = ["chardet (>=3.0.2,<6)"]
name = "responses"
version = "0.18.0"
description = "A utility library for mocking out the `requests` Python library."
+category = "main"
optional = true
python-versions = ">=3.7"
files = [
@@ -5266,6 +5468,7 @@ tests = ["coverage (>=6.0.0)", "flake8", "mypy", "pytest (>=4.6)", "pytest-cov",
name = "rfc3339-validator"
version = "0.1.4"
description = "A pure python RFC3339 validator"
+category = "dev"
optional = false
python-versions = ">=2.7, !=3.0.*, !=3.1.*, !=3.2.*, !=3.3.*, !=3.4.*"
files = [
@@ -5280,6 +5483,7 @@ six = "*"
name = "rfc3986-validator"
version = "0.1.1"
description = "Pure python rfc3986 validator"
+category = "dev"
optional = false
python-versions = ">=2.7, !=3.0.*, !=3.1.*, !=3.2.*, !=3.3.*, !=3.4.*"
files = [
@@ -5291,6 +5495,7 @@ files = [
name = "rpds-py"
version = "0.16.2"
description = "Python bindings to Rust's persistent data structures (rpds)"
+category = "dev"
optional = false
python-versions = ">=3.8"
files = [
@@ -5399,6 +5604,7 @@ files = [
name = "rsa"
version = "4.9"
description = "Pure-Python RSA implementation"
+category = "dev"
optional = false
python-versions = ">=3.6,<4"
files = [
@@ -5413,6 +5619,7 @@ pyasn1 = ">=0.1.3"
name = "ruff"
version = "0.0.292"
description = "An extremely fast Python linter, written in Rust."
+category = "dev"
optional = false
python-versions = ">=3.7"
files = [
@@ -5439,6 +5646,7 @@ files = [
name = "s3transfer"
version = "0.8.2"
description = "An Amazon S3 Transfer Manager"
+category = "dev"
optional = false
python-versions = ">= 3.7"
files = [
@@ -5456,6 +5664,7 @@ crt = ["botocore[crt] (>=1.33.2,<2.0a.0)"]
name = "safetensors"
version = "0.4.1"
description = ""
+category = "main"
optional = true
python-versions = ">=3.7"
files = [
@@ -5575,6 +5784,7 @@ torch = ["safetensors[numpy]", "torch (>=1.10)"]
name = "scikit-learn"
version = "1.3.2"
description = "A set of python modules for machine learning and data mining"
+category = "main"
optional = true
python-versions = ">=3.8"
files = [
@@ -5622,6 +5832,7 @@ tests = ["black (>=23.3.0)", "matplotlib (>=3.1.3)", "mypy (>=1.3)", "numpydoc (
name = "scipy"
version = "1.9.3"
description = "Fundamental algorithms for scientific computing in Python"
+category = "main"
optional = true
python-versions = ">=3.8"
files = [
@@ -5660,6 +5871,7 @@ test = ["asv", "gmpy2", "mpmath", "pytest", "pytest-cov", "pytest-xdist", "sciki
name = "send2trash"
version = "1.8.2"
description = "Send file to trash natively under Mac OS X, Windows and Linux"
+category = "dev"
optional = false
python-versions = "!=3.0.*,!=3.1.*,!=3.2.*,!=3.3.*,!=3.4.*,>=2.7"
files = [
@@ -5676,6 +5888,7 @@ win32 = ["pywin32"]
name = "sentencepiece"
version = "0.1.99"
description = "SentencePiece python wrapper"
+category = "main"
optional = true
python-versions = "*"
files = [
@@ -5730,6 +5943,7 @@ files = [
name = "setuptools"
version = "69.0.3"
description = "Easily download, build, install, upgrade, and uninstall Python packages"
+category = "main"
optional = false
python-versions = ">=3.8"
files = [
@@ -5746,6 +5960,7 @@ testing-integration = ["build[virtualenv] (>=1.0.3)", "filelock (>=3.4.0)", "jar
name = "six"
version = "1.16.0"
description = "Python 2 and 3 compatibility utilities"
+category = "main"
optional = false
python-versions = ">=2.7, !=3.0.*, !=3.1.*, !=3.2.*"
files = [
@@ -5757,6 +5972,7 @@ files = [
name = "smart-open"
version = "6.4.0"
description = "Utils for streaming large files (S3, HDFS, GCS, Azure Blob Storage, gzip, bz2...)"
+category = "main"
optional = true
python-versions = ">=3.6,<4.0"
files = [
@@ -5778,6 +5994,7 @@ webhdfs = ["requests"]
name = "sniffio"
version = "1.3.0"
description = "Sniff out which async library your code is running under"
+category = "main"
optional = false
python-versions = ">=3.7"
files = [
@@ -5789,6 +6006,7 @@ files = [
name = "snowballstemmer"
version = "2.2.0"
description = "This package provides 29 stemmers for 28 languages generated from Snowball algorithms."
+category = "dev"
optional = false
python-versions = "*"
files = [
@@ -5800,6 +6018,7 @@ files = [
name = "soupsieve"
version = "2.5"
description = "A modern CSS selector implementation for Beautiful Soup."
+category = "main"
optional = false
python-versions = ">=3.8"
files = [
@@ -5811,6 +6030,7 @@ files = [
name = "spacy"
version = "3.7.2"
description = "Industrial-strength Natural Language Processing (NLP) in Python"
+category = "main"
optional = true
python-versions = ">=3.7"
files = [
@@ -5902,6 +6122,7 @@ transformers = ["spacy-transformers (>=1.1.2,<1.4.0)"]
name = "spacy-legacy"
version = "3.0.12"
description = "Legacy registered functions for spaCy backwards compatibility"
+category = "main"
optional = true
python-versions = ">=3.6"
files = [
@@ -5913,6 +6134,7 @@ files = [
name = "spacy-loggers"
version = "1.0.5"
description = "Logging utilities for SpaCy"
+category = "main"
optional = true
python-versions = ">=3.6"
files = [
@@ -5924,6 +6146,7 @@ files = [
name = "sphinx"
version = "5.3.0"
description = "Python documentation generator"
+category = "dev"
optional = false
python-versions = ">=3.6"
files = [
@@ -5959,6 +6182,7 @@ test = ["cython", "html5lib", "pytest (>=4.6)", "typed_ast"]
name = "sphinx-autobuild"
version = "2021.3.14"
description = "Rebuild Sphinx documentation on changes, with live-reload in the browser."
+category = "dev"
optional = false
python-versions = ">=3.6"
files = [
@@ -5978,6 +6202,7 @@ test = ["pytest", "pytest-cov"]
name = "sphinx-automodapi"
version = "0.16.0"
description = "Sphinx extension for auto-generating API documentation for entire modules"
+category = "dev"
optional = false
python-versions = ">=3.7"
files = [
@@ -5996,6 +6221,7 @@ test = ["coverage", "cython", "pytest", "pytest-cov", "setuptools"]
name = "sphinx-basic-ng"
version = "1.0.0b2"
description = "A modern skeleton for Sphinx themes."
+category = "dev"
optional = false
python-versions = ">=3.7"
files = [
@@ -6013,6 +6239,7 @@ docs = ["furo", "ipython", "myst-parser", "sphinx-copybutton", "sphinx-inline-ta
name = "sphinx-reredirects"
version = "0.1.3"
description = "Handles redirects for moved pages in Sphinx documentation projects"
+category = "dev"
optional = false
python-versions = ">=3.5"
files = [
@@ -6027,6 +6254,7 @@ sphinx = "*"
name = "sphinx-rtd-theme"
version = "1.3.0"
description = "Read the Docs theme for Sphinx"
+category = "dev"
optional = false
python-versions = "!=3.0.*,!=3.1.*,!=3.2.*,!=3.3.*,!=3.4.*,!=3.5.*,>=2.7"
files = [
@@ -6046,6 +6274,7 @@ dev = ["bump2version", "sphinxcontrib-httpdomain", "transifex-client", "wheel"]
name = "sphinxcontrib-applehelp"
version = "1.0.4"
description = "sphinxcontrib-applehelp is a Sphinx extension which outputs Apple help books"
+category = "dev"
optional = false
python-versions = ">=3.8"
files = [
@@ -6061,6 +6290,7 @@ test = ["pytest"]
name = "sphinxcontrib-devhelp"
version = "1.0.2"
description = "sphinxcontrib-devhelp is a sphinx extension which outputs Devhelp document."
+category = "dev"
optional = false
python-versions = ">=3.5"
files = [
@@ -6076,6 +6306,7 @@ test = ["pytest"]
name = "sphinxcontrib-gtagjs"
version = "0.2.1"
description = "Sphinx extension to render global site tag of Google."
+category = "dev"
optional = false
python-versions = "*"
files = [
@@ -6093,6 +6324,7 @@ packaging = ["bumpversion"]
name = "sphinxcontrib-htmlhelp"
version = "2.0.1"
description = "sphinxcontrib-htmlhelp is a sphinx extension which renders HTML help files"
+category = "dev"
optional = false
python-versions = ">=3.8"
files = [
@@ -6108,6 +6340,7 @@ test = ["html5lib", "pytest"]
name = "sphinxcontrib-jquery"
version = "4.1"
description = "Extension to include jQuery on newer Sphinx releases"
+category = "dev"
optional = false
python-versions = ">=2.7"
files = [
@@ -6122,6 +6355,7 @@ Sphinx = ">=1.8"
name = "sphinxcontrib-jsmath"
version = "1.0.1"
description = "A sphinx extension which renders display math in HTML via JavaScript"
+category = "dev"
optional = false
python-versions = ">=3.5"
files = [
@@ -6136,6 +6370,7 @@ test = ["flake8", "mypy", "pytest"]
name = "sphinxcontrib-qthelp"
version = "1.0.3"
description = "sphinxcontrib-qthelp is a sphinx extension which outputs QtHelp document."
+category = "dev"
optional = false
python-versions = ">=3.5"
files = [
@@ -6151,6 +6386,7 @@ test = ["pytest"]
name = "sphinxcontrib-serializinghtml"
version = "1.1.5"
description = "sphinxcontrib-serializinghtml is a sphinx extension which outputs \"serialized\" HTML files (json and pickle)."
+category = "dev"
optional = false
python-versions = ">=3.5"
files = [
@@ -6166,6 +6402,7 @@ test = ["pytest"]
name = "sqlalchemy"
version = "2.0.25"
description = "Database Abstraction Library"
+category = "main"
optional = false
python-versions = ">=3.7"
files = [
@@ -6253,6 +6490,7 @@ sqlcipher = ["sqlcipher3_binary"]
name = "srsly"
version = "2.4.8"
description = "Modern high-performance serialization utilities for Python"
+category = "main"
optional = true
python-versions = ">=3.6"
files = [
@@ -6299,6 +6537,7 @@ catalogue = ">=2.0.3,<2.1.0"
name = "stack-data"
version = "0.6.3"
description = "Extract data from python stack frames and tracebacks for informative displays"
+category = "dev"
optional = false
python-versions = "*"
files = [
@@ -6318,6 +6557,7 @@ tests = ["cython", "littleutils", "pygments", "pytest", "typeguard"]
name = "sympy"
version = "1.12"
description = "Computer algebra system (CAS) in Python"
+category = "main"
optional = true
python-versions = ">=3.8"
files = [
@@ -6332,6 +6572,7 @@ mpmath = ">=0.19"
name = "tabulate"
version = "0.9.0"
description = "Pretty-print tabular data"
+category = "dev"
optional = false
python-versions = ">=3.7"
files = [
@@ -6346,6 +6587,7 @@ widechars = ["wcwidth"]
name = "tenacity"
version = "8.2.3"
description = "Retry code until it succeeds"
+category = "main"
optional = false
python-versions = ">=3.7"
files = [
@@ -6360,6 +6602,7 @@ doc = ["reno", "sphinx", "tornado (>=4.5)"]
name = "terminado"
version = "0.18.0"
description = "Tornado websocket backend for the Xterm.js Javascript terminal emulator library."
+category = "dev"
optional = false
python-versions = ">=3.8"
files = [
@@ -6381,6 +6624,7 @@ typing = ["mypy (>=1.6,<2.0)", "traitlets (>=5.11.1)"]
name = "thinc"
version = "8.2.2"
description = "A refreshing functional take on deep learning, compatible with your favorite libraries"
+category = "main"
optional = true
python-versions = ">=3.6"
files = [
@@ -6466,6 +6710,7 @@ torch = ["torch (>=1.6.0)"]
name = "threadpoolctl"
version = "3.2.0"
description = "threadpoolctl"
+category = "main"
optional = true
python-versions = ">=3.8"
files = [
@@ -6477,6 +6722,7 @@ files = [
name = "tiktoken"
version = "0.5.2"
description = "tiktoken is a fast BPE tokeniser for use with OpenAI's models"
+category = "main"
optional = false
python-versions = ">=3.8"
files = [
@@ -6529,6 +6775,7 @@ blobfile = ["blobfile (>=2)"]
name = "tinycss2"
version = "1.2.1"
description = "A tiny CSS parser"
+category = "dev"
optional = false
python-versions = ">=3.7"
files = [
@@ -6547,6 +6794,7 @@ test = ["flake8", "isort", "pytest"]
name = "tokenize-rt"
version = "5.2.0"
description = "A wrapper around the stdlib `tokenize` which roundtrips."
+category = "dev"
optional = false
python-versions = ">=3.8"
files = [
@@ -6558,6 +6806,7 @@ files = [
name = "tokenizers"
version = "0.15.0"
description = ""
+category = "main"
optional = true
python-versions = ">=3.7"
files = [
@@ -6673,6 +6922,7 @@ testing = ["black (==22.3)", "datasets", "numpy", "pytest", "requests"]
name = "tomli"
version = "2.0.1"
description = "A lil' TOML parser"
+category = "dev"
optional = false
python-versions = ">=3.7"
files = [
@@ -6684,6 +6934,7 @@ files = [
name = "tomlkit"
version = "0.12.3"
description = "Style preserving TOML library"
+category = "dev"
optional = false
python-versions = ">=3.7"
files = [
@@ -6695,6 +6946,7 @@ files = [
name = "torch"
version = "2.1.2"
description = "Tensors and Dynamic neural networks in Python with strong GPU acceleration"
+category = "main"
optional = true
python-versions = ">=3.8.0"
files = [
@@ -6748,6 +7000,7 @@ opt-einsum = ["opt-einsum (>=3.3)"]
name = "tornado"
version = "6.4"
description = "Tornado is a Python web framework and asynchronous networking library, originally developed at FriendFeed."
+category = "dev"
optional = false
python-versions = ">= 3.8"
files = [
@@ -6768,6 +7021,7 @@ files = [
name = "tqdm"
version = "4.66.1"
description = "Fast, Extensible Progress Meter"
+category = "main"
optional = false
python-versions = ">=3.7"
files = [
@@ -6788,6 +7042,7 @@ telegram = ["requests"]
name = "traitlets"
version = "5.14.1"
description = "Traitlets Python configuration system"
+category = "dev"
optional = false
python-versions = ">=3.8"
files = [
@@ -6803,6 +7058,7 @@ test = ["argcomplete (>=3.0.3)", "mypy (>=1.7.0)", "pre-commit", "pytest (>=7.0,
name = "transformers"
version = "4.36.2"
description = "State-of-the-art Machine Learning for JAX, PyTorch and TensorFlow"
+category = "main"
optional = true
python-versions = ">=3.8.0"
files = [
@@ -6875,6 +7131,7 @@ vision = ["Pillow (>=10.0.1,<=15.0)"]
name = "tree-sitter"
version = "0.20.4"
description = "Python bindings for the Tree-Sitter parsing library"
+category = "dev"
optional = false
python-versions = ">=3.3"
files = [
@@ -6963,6 +7220,7 @@ setuptools = {version = ">=60.0.0", markers = "python_version >= \"3.12\""}
name = "tree-sitter-languages"
version = "1.9.1"
description = "Binary Python wheels for all tree sitter languages."
+category = "dev"
optional = false
python-versions = "*"
files = [
@@ -7043,6 +7301,7 @@ tree-sitter = "*"
name = "triton"
version = "2.1.0"
description = "A language and compiler for custom Deep Learning operations"
+category = "main"
optional = true
python-versions = "*"
files = [
@@ -7068,6 +7327,7 @@ tutorials = ["matplotlib", "pandas", "tabulate"]
name = "typer"
version = "0.9.0"
description = "Typer, build great CLIs. Easy to code. Based on Python type hints."
+category = "main"
optional = true
python-versions = ">=3.6"
files = [
@@ -7089,6 +7349,7 @@ test = ["black (>=22.3.0,<23.0.0)", "coverage (>=6.2,<7.0)", "isort (>=5.0.6,<6.
name = "types-deprecated"
version = "1.2.9.20240106"
description = "Typing stubs for Deprecated"
+category = "dev"
optional = false
python-versions = ">=3.8"
files = [
@@ -7100,6 +7361,7 @@ files = [
name = "types-docutils"
version = "0.20.0.20240106"
description = "Typing stubs for docutils"
+category = "dev"
optional = false
python-versions = ">=3.8"
files = [
@@ -7111,6 +7373,7 @@ files = [
name = "types-protobuf"
version = "4.24.0.20240106"
description = "Typing stubs for protobuf"
+category = "dev"
optional = false
python-versions = ">=3.8"
files = [
@@ -7122,6 +7385,7 @@ files = [
name = "types-pyopenssl"
version = "23.3.0.20240106"
description = "Typing stubs for pyOpenSSL"
+category = "dev"
optional = false
python-versions = ">=3.8"
files = [
@@ -7136,6 +7400,7 @@ cryptography = ">=35.0.0"
name = "types-python-dateutil"
version = "2.8.19.20240106"
description = "Typing stubs for python-dateutil"
+category = "dev"
optional = false
python-versions = ">=3.8"
files = [
@@ -7147,6 +7412,7 @@ files = [
name = "types-pyyaml"
version = "6.0.12.12"
description = "Typing stubs for PyYAML"
+category = "dev"
optional = false
python-versions = "*"
files = [
@@ -7158,6 +7424,7 @@ files = [
name = "types-redis"
version = "4.5.5.0"
description = "Typing stubs for redis"
+category = "dev"
optional = false
python-versions = "*"
files = [
@@ -7173,6 +7440,7 @@ types-pyOpenSSL = "*"
name = "types-requests"
version = "2.28.11.8"
description = "Typing stubs for requests"
+category = "dev"
optional = false
python-versions = "*"
files = [
@@ -7187,6 +7455,7 @@ types-urllib3 = "<1.27"
name = "types-setuptools"
version = "67.1.0.0"
description = "Typing stubs for setuptools"
+category = "dev"
optional = false
python-versions = "*"
files = [
@@ -7201,6 +7470,7 @@ types-docutils = "*"
name = "types-urllib3"
version = "1.26.25.14"
description = "Typing stubs for urllib3"
+category = "dev"
optional = false
python-versions = "*"
files = [
@@ -7212,6 +7482,7 @@ files = [
name = "typing-extensions"
version = "4.9.0"
description = "Backported and Experimental Type Hints for Python 3.8+"
+category = "main"
optional = false
python-versions = ">=3.8"
files = [
@@ -7223,6 +7494,7 @@ files = [
name = "typing-inspect"
version = "0.9.0"
description = "Runtime inspection utilities for typing module."
+category = "main"
optional = false
python-versions = "*"
files = [
@@ -7238,6 +7510,7 @@ typing-extensions = ">=3.7.4"
name = "tzdata"
version = "2023.4"
description = "Provider of IANA time zone data"
+category = "main"
optional = false
python-versions = ">=2"
files = [
@@ -7249,6 +7522,7 @@ files = [
name = "uri-template"
version = "1.3.0"
description = "RFC 6570 URI Template Processor"
+category = "dev"
optional = false
python-versions = ">=3.7"
files = [
@@ -7263,6 +7537,7 @@ dev = ["flake8", "flake8-annotations", "flake8-bandit", "flake8-bugbear", "flake
name = "urllib3"
version = "1.26.18"
description = "HTTP library with thread-safe connection pooling, file post, and more."
+category = "main"
optional = false
python-versions = ">=2.7, !=3.0.*, !=3.1.*, !=3.2.*, !=3.3.*, !=3.4.*, !=3.5.*"
files = [
@@ -7279,6 +7554,7 @@ socks = ["PySocks (>=1.5.6,!=1.5.7,<2.0)"]
name = "urllib3"
version = "2.0.7"
description = "HTTP library with thread-safe connection pooling, file post, and more."
+category = "main"
optional = false
python-versions = ">=3.7"
files = [
@@ -7296,6 +7572,7 @@ zstd = ["zstandard (>=0.18.0)"]
name = "vellum-ai"
version = "0.0.42"
description = ""
+category = "dev"
optional = false
python-versions = ">=3.7,<4.0"
files = [
@@ -7311,6 +7588,7 @@ pydantic = ">=1.9.2,<2.0.0"
name = "virtualenv"
version = "20.25.0"
description = "Virtual Python Environment builder"
+category = "dev"
optional = false
python-versions = ">=3.7"
files = [
@@ -7331,6 +7609,7 @@ test = ["covdefaults (>=2.3)", "coverage (>=7.2.7)", "coverage-enable-subprocess
name = "wasabi"
version = "1.1.2"
description = "A lightweight console printing and formatting toolkit"
+category = "main"
optional = true
python-versions = ">=3.6"
files = [
@@ -7345,6 +7624,7 @@ colorama = {version = ">=0.4.6", markers = "sys_platform == \"win32\" and python
name = "wcwidth"
version = "0.2.13"
description = "Measures the displayed width of unicode strings in a terminal"
+category = "dev"
optional = false
python-versions = "*"
files = [
@@ -7356,6 +7636,7 @@ files = [
name = "weasel"
version = "0.3.4"
description = "Weasel: A small and easy workflow system"
+category = "main"
optional = true
python-versions = ">=3.6"
files = [
@@ -7378,6 +7659,7 @@ wasabi = ">=0.9.1,<1.2.0"
name = "webcolors"
version = "1.13"
description = "A library for working with the color formats defined by HTML and CSS."
+category = "dev"
optional = false
python-versions = ">=3.7"
files = [
@@ -7393,6 +7675,7 @@ tests = ["pytest", "pytest-cov"]
name = "webencodings"
version = "0.5.1"
description = "Character encoding aliases for legacy web content"
+category = "dev"
optional = false
python-versions = "*"
files = [
@@ -7404,6 +7687,7 @@ files = [
name = "websocket-client"
version = "1.7.0"
description = "WebSocket client for Python with low level API options"
+category = "dev"
optional = false
python-versions = ">=3.8"
files = [
@@ -7420,6 +7704,7 @@ test = ["websockets"]
name = "widgetsnbextension"
version = "4.0.9"
description = "Jupyter interactive widgets for Jupyter Notebook"
+category = "dev"
optional = false
python-versions = ">=3.7"
files = [
@@ -7431,6 +7716,7 @@ files = [
name = "wrapt"
version = "1.16.0"
description = "Module for decorators, wrappers and monkey patching."
+category = "main"
optional = false
python-versions = ">=3.6"
files = [
@@ -7510,6 +7796,7 @@ files = [
name = "xxhash"
version = "3.4.1"
description = "Python binding for xxHash"
+category = "main"
optional = true
python-versions = ">=3.7"
files = [
@@ -7627,6 +7914,7 @@ files = [
name = "yarl"
version = "1.9.4"
description = "Yet another URL library"
+category = "main"
optional = false
python-versions = ">=3.7"
files = [
@@ -7730,6 +8018,7 @@ multidict = ">=4.0"
name = "zipp"
version = "3.17.0"
description = "Backport of pathlib-compatible object wrapper for zip files"
+category = "dev"
optional = false
python-versions = ">=3.8"
files = [
@@ -7751,4 +8040,4 @@ query-tools = ["guidance", "jsonpath-ng", "lm-format-enforcer", "rank-bm25", "sc
[metadata]
lock-version = "2.0"
python-versions = ">=3.8.1,<4.0"
-content-hash = "bef2e7c1b33082e0dee9030ac812d1bb82f7f0a53345688c549a45bd4164288e"
+content-hash = "c70579267131c20843f8364e3e380cdb52272873bc6c1aae3b43815ec4d4bea1"
diff --git a/pyproject.toml b/pyproject.toml
index 845d869ce8..5234d88636 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -77,6 +77,7 @@ rank-bm25 = {optional = true, version = "^0.2.2"}
scikit-learn = {optional = true, version = "*"}
spacy = {optional = true, version = "^3.7.1"}
aiohttp = "^3.8.6"
+networkx = ">=3.0"
psycopg2 = {optional = true, version = "^2.9.9"}
[tool.poetry.extras]
diff --git a/tests/llms/test_custom.py b/tests/llms/test_custom.py
index 90e874e3d5..7535d56922 100644
--- a/tests/llms/test_custom.py
+++ b/tests/llms/test_custom.py
@@ -19,7 +19,9 @@ class TestLLM(CustomLLM):
def metadata(self) -> LLMMetadata:
return LLMMetadata()
- def complete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
+ def complete(
+ self, prompt: str, formatted: bool = False, **kwargs: Any
+ ) -> CompletionResponse:
return CompletionResponse(
text="test output",
additional_kwargs={
@@ -27,7 +29,9 @@ class TestLLM(CustomLLM):
},
)
- def stream_complete(self, prompt: str, **kwargs: Any) -> CompletionResponseGen:
+ def stream_complete(
+ self, prompt: str, formatted: bool = False, **kwargs: Any
+ ) -> CompletionResponseGen:
def gen() -> CompletionResponseGen:
text = "test output"
text_so_far = ""
diff --git a/tests/query_pipeline/__init__.py b/tests/query_pipeline/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/tests/query_pipeline/query.py b/tests/query_pipeline/query.py
new file mode 100644
index 0000000000..d6c2a2a973
--- /dev/null
+++ b/tests/query_pipeline/query.py
@@ -0,0 +1,266 @@
+"""Query pipeline."""
+
+from typing import Any, Dict
+
+import pytest
+from llama_index.core.query_pipeline.query_component import (
+ ChainableMixin,
+ InputComponent,
+ InputKeys,
+ OutputKeys,
+ QueryComponent,
+)
+from llama_index.query_pipeline.query import QueryPipeline
+
+
+class QueryComponent1(QueryComponent):
+ """Query component 1.
+
+ Adds two numbers together.
+
+ """
+
+ def set_callback_manager(self, callback_manager: Any) -> None:
+ """Set callback manager."""
+
+ def _validate_component_inputs(self, input: Dict[str, Any]) -> Dict[str, Any]:
+ """Validate component inputs during run_component."""
+ if "input1" not in input:
+ raise ValueError("input1 not in input")
+ if "input2" not in input:
+ raise ValueError("input2 not in input")
+ return input
+
+ def _run_component(self, **kwargs: Any) -> Any:
+ """Run component."""
+ return {"output": kwargs["input1"] + kwargs["input2"]}
+
+ async def _arun_component(self, **kwargs: Any) -> Any:
+ """Run component."""
+ return self._run_component(**kwargs)
+
+ @property
+ def input_keys(self) -> InputKeys:
+ """Input keys."""
+ return InputKeys.from_keys({"input1", "input2"})
+
+ @property
+ def output_keys(self) -> OutputKeys:
+ """Output keys."""
+ return OutputKeys.from_keys({"output"})
+
+
+class QueryComponent2(QueryComponent):
+ """Query component 1.
+
+ Joins two strings together with ':'
+
+ """
+
+ def set_callback_manager(self, callback_manager: Any) -> None:
+ """Set callback manager."""
+
+ def _validate_component_inputs(self, input: Dict[str, Any]) -> Dict[str, Any]:
+ """Validate component inputs during run_component."""
+ if "input1" not in input:
+ raise ValueError("input1 not in input")
+ if "input2" not in input:
+ raise ValueError("input2 not in input")
+ return input
+
+ def _run_component(self, **kwargs: Any) -> Any:
+ """Run component."""
+ return {"output": f"{kwargs['input1']}:{kwargs['input2']}"}
+
+ async def _arun_component(self, **kwargs: Any) -> Any:
+ """Run component."""
+ return self._run_component(**kwargs)
+
+ @property
+ def input_keys(self) -> InputKeys:
+ """Input keys."""
+ return InputKeys.from_keys({"input1", "input2"})
+
+ @property
+ def output_keys(self) -> OutputKeys:
+ """Output keys."""
+ return OutputKeys.from_keys({"output"})
+
+
+class QueryComponent3(QueryComponent):
+ """Query component 3.
+
+ Takes one input and doubles it.
+
+ """
+
+ def set_callback_manager(self, callback_manager: Any) -> None:
+ """Set callback manager."""
+
+ def _validate_component_inputs(self, input: Dict[str, Any]) -> Dict[str, Any]:
+ """Validate component inputs during run_component."""
+ if "input" not in input:
+ raise ValueError("input not in input")
+ return input
+
+ def _run_component(self, **kwargs: Any) -> Dict:
+ """Run component."""
+ return {"output": kwargs["input"] + kwargs["input"]}
+
+ async def _arun_component(self, **kwargs: Any) -> Any:
+ """Run component."""
+ return self._run_component(**kwargs)
+
+ @property
+ def input_keys(self) -> InputKeys:
+ """Input keys."""
+ return InputKeys.from_keys({"input"})
+
+ @property
+ def output_keys(self) -> OutputKeys:
+ """Output keys."""
+ return OutputKeys.from_keys({"output"})
+
+
+class Chainable2(ChainableMixin):
+ """Chainable mixin."""
+
+ def _as_query_component(self, **kwargs: Any) -> "QueryComponent":
+ """Get query component."""
+ return QueryComponent2()
+
+
+def test_query_pipeline_chain() -> None:
+ """Test query pipeline."""
+ # test qc1 by itself with chain syntax
+ p = QueryPipeline(chain=[QueryComponent1()])
+ output = p.run(input1=1, input2=2)
+ # since there's one output, output is just the value
+ assert output == 3
+
+
+def test_query_pipeline_single_arg_inp() -> None:
+ """Test query pipeline with single arg input (no kwargs)."""
+ # should work if input is a single arg
+ p = QueryPipeline(chain=[QueryComponent3(), QueryComponent3()])
+ # since there's one output, output is just the value
+ output = p.run(3)
+ assert output == 12
+
+
+def test_query_pipeline_input_component() -> None:
+ """Test query pipeline input component."""
+ # test connecting different inputs to different components
+ qc1 = QueryComponent1()
+ qc2 = QueryComponent2()
+ inp = InputComponent()
+ p = QueryPipeline()
+
+ p.add_modules({"qc1": qc1, "qc2": qc2, "inp": inp})
+ # add inp.inp1 to both qc1.input1 and qc2.input2
+ p.add_link("inp", "qc1", src_key="inp1", dest_key="input1")
+ p.add_link("inp", "qc2", src_key="inp1", dest_key="input2")
+ # add inp.inp2 to qc1.input2
+ p.add_link("inp", "qc1", src_key="inp2", dest_key="input2")
+ # add qc1 to qc2.input1
+ p.add_link("qc1", "qc2", dest_key="input1")
+
+ output = p.run(inp1=1, inp2=2)
+ assert output == "3:1"
+
+
+def test_query_pipeline_partial() -> None:
+ """Test query pipeline."""
+ # test qc1 with qc2 with one partial, with chain syntax
+ qc1 = QueryComponent1()
+ qc2 = QueryComponent2()
+ qc2.partial(input2="hello")
+ p = QueryPipeline(chain=[qc1, qc2])
+ output = p.run(input1=1, input2=2)
+ assert output == "3:hello"
+
+ # test qc1 with qc2 with one partial with full syntax
+ qc1 = QueryComponent1()
+ qc2 = QueryComponent2()
+ p = QueryPipeline()
+ p.add_modules({"qc1": qc1, "qc2": qc2})
+ qc2.partial(input2="foo")
+ p.add_link("qc1", "qc2", dest_key="input1")
+ output = p.run(input1=2, input2=2)
+ assert output == "4:foo"
+
+ # test partial with ChainableMixin
+ c2_0 = Chainable2().as_query_component(partial={"input2": "hello"})
+ c2_1 = Chainable2().as_query_component(partial={"input2": "world"})
+ # you can now define a chain because input2 has been defined
+ p = QueryPipeline(chain=[c2_0, c2_1])
+ output = p.run(input1=1)
+ assert output == "1:hello:world"
+
+
+def test_query_pipeline_sub() -> None:
+ """Test query pipeline."""
+ # test qc2 with subpipelines of qc3 w/ full syntax
+ qc2 = QueryComponent2()
+ qc3 = QueryComponent3()
+ p1 = QueryPipeline(chain=[qc3, qc3])
+ p = QueryPipeline()
+ p.add_modules({"qc2": qc2, "p1": p1})
+ # link output of p1 to input1 and input2 of qc2
+ p.add_link("p1", "qc2", dest_key="input1")
+ p.add_link("p1", "qc2", dest_key="input2")
+ output = p.run(input=2)
+ assert output == "8:8"
+
+
+def test_query_pipeline_multi() -> None:
+ """Test query pipeline."""
+ # try run run_multi
+ # link both qc1_0 and qc1_1 to qc2
+ qc1_0 = QueryComponent1()
+ qc1_1 = QueryComponent1()
+ qc2 = QueryComponent2()
+ p = QueryPipeline()
+ p.add_modules({"qc1_0": qc1_0, "qc1_1": qc1_1, "qc2": qc2})
+ p.add_link("qc1_0", "qc2", dest_key="input1")
+ p.add_link("qc1_1", "qc2", dest_key="input2")
+ output = p.run_multi(
+ {"qc1_0": {"input1": 1, "input2": 2}, "qc1_1": {"input1": 3, "input2": 4}}
+ )
+ assert output == {"qc2": {"output": "3:7"}}
+
+
+@pytest.mark.asyncio()
+async def test_query_pipeline_async() -> None:
+ """Test query pipeline in async fashion."""
+ # run some synchronous tests above
+
+ # should work if input is a single arg
+ p = QueryPipeline(chain=[QueryComponent3(), QueryComponent3()])
+ # since there's one output, output is just the value
+ output = await p.arun(3)
+ assert output == 12
+
+ # test qc1 with qc2 with one partial with full syntax
+ qc1 = QueryComponent1()
+ qc2 = QueryComponent2()
+ p = QueryPipeline()
+ p.add_modules({"qc1": qc1, "qc2": qc2})
+ qc2.partial(input2="foo")
+ p.add_link("qc1", "qc2", dest_key="input1")
+ output = await p.arun(input1=2, input2=2)
+ assert output == "4:foo"
+
+ # try run run_multi
+ # link both qc1_0 and qc1_1 to qc2
+ qc1_0 = QueryComponent1()
+ qc1_1 = QueryComponent1()
+ qc2 = QueryComponent2()
+ p = QueryPipeline()
+ p.add_modules({"qc1_0": qc1_0, "qc1_1": qc1_1, "qc2": qc2})
+ p.add_link("qc1_0", "qc2", dest_key="input1")
+ p.add_link("qc1_1", "qc2", dest_key="input2")
+ output = await p.arun_multi(
+ {"qc1_0": {"input1": 1, "input2": 2}, "qc1_1": {"input1": 3, "input2": 4}}
+ )
+ assert output == {"qc2": {"output": "3:7"}}
--
GitLab