wehub-resource-sync
a0c8464e58
Build Package / build (ubuntu-latest) (push) Failing after 1s
CodeQL / Analyze (python) (push) Failing after 1s
Core Typecheck / core-typecheck (push) Failing after 1s
Linting / lint (push) Failing after 1s
llama-dev tests / test-llama-dev (push) Failing after 1s
Publish Sub-Package to PyPI if Needed / publish_subpackage_if_needed (push) Has been skipped
Sync Docs to Developer Hub / sync-docs (push) Failing after 0s
Build Package / build (windows-latest) (push) Has been cancelled
557 lines
19 KiB
Plaintext
557 lines
19 KiB
Plaintext
{
|
|
"cells": [
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"# LongRAG Workflow"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"This notebook shows how to implement LongRAG using LlamaIndex workflows."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"import nest_asyncio\n",
|
|
"\n",
|
|
"nest_asyncio.apply()"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"%pip install -U llama-index"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"import os\n",
|
|
"\n",
|
|
"os.environ[\"OPENAI_API_KEY\"] = \"sk-proj-...\""
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"!wget https://github.com/user-attachments/files/16474262/data.zip -O data.zip\n",
|
|
"!unzip -o data.zip\n",
|
|
"!rm data.zip"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"Since workflows are async first, this all runs fine in a notebook. If you were running in your own code, you would want to use `asyncio.run()` to start an async event loop if one isn't already running.\n",
|
|
"\n",
|
|
"```python\n",
|
|
"async def main():\n",
|
|
" <async code>\n",
|
|
"\n",
|
|
"if __name__ == \"__main__\":\n",
|
|
" import asyncio\n",
|
|
" asyncio.run(main())\n",
|
|
"```"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Helper Functions\n",
|
|
"\n",
|
|
"These helper functions will help us split documents into smaller pieces and group nodes based on their relationships."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"from typing import List, Dict, Optional, Set, FrozenSet\n",
|
|
"\n",
|
|
"from llama_index.core.schema import BaseNode, TextNode\n",
|
|
"from llama_index.core.node_parser import SentenceSplitter\n",
|
|
"\n",
|
|
"# constants\n",
|
|
"DEFAULT_CHUNK_SIZE = 4096 # optionally splits documents into CHUNK_SIZE, then regroups them to demonstrate grouping algorithm\n",
|
|
"DEFAULT_MAX_GROUP_SIZE = 20 # maximum number of documents in a group\n",
|
|
"DEFAULT_SMALL_CHUNK_SIZE = 512 # small chunk size for generating embeddings\n",
|
|
"DEFAULT_TOP_K = 8 # top k for retrieving\n",
|
|
"\n",
|
|
"\n",
|
|
"def split_doc(chunk_size: int, documents: List[BaseNode]) -> List[TextNode]:\n",
|
|
" \"\"\"Splits documents into smaller pieces.\n",
|
|
"\n",
|
|
" Args:\n",
|
|
" chunk_size (int): Chunk size\n",
|
|
" documents (List[BaseNode]): Documents\n",
|
|
"\n",
|
|
" Returns:\n",
|
|
" List[TextNode]: Smaller chunks\n",
|
|
" \"\"\"\n",
|
|
" # split docs into tokens\n",
|
|
" text_parser = SentenceSplitter(chunk_size=chunk_size)\n",
|
|
" return text_parser.get_nodes_from_documents(documents)\n",
|
|
"\n",
|
|
"\n",
|
|
"def group_docs(\n",
|
|
" nodes: List[str],\n",
|
|
" adj: Dict[str, List[str]],\n",
|
|
" max_group_size: Optional[int] = DEFAULT_MAX_GROUP_SIZE,\n",
|
|
") -> Set[FrozenSet[str]]:\n",
|
|
" \"\"\"Groups documents.\n",
|
|
"\n",
|
|
" Args:\n",
|
|
" nodes (List[str]): documents IDs\n",
|
|
" adj (Dict[str, List[str]]): related documents for each document; id -> list of doc strings\n",
|
|
" max_group_size (Optional[int], optional): max group size, None if no max group size. Defaults to DEFAULT_MAX_GROUP_SIZE.\n",
|
|
" \"\"\"\n",
|
|
" docs = sorted(nodes, key=lambda node: len(adj[node]))\n",
|
|
" groups = set() # set of set of IDs\n",
|
|
" for d in docs:\n",
|
|
" related_groups = set()\n",
|
|
" for r in adj[d]:\n",
|
|
" for g in groups:\n",
|
|
" if r in g:\n",
|
|
" related_groups = related_groups.union(frozenset([g]))\n",
|
|
"\n",
|
|
" gnew = {d}\n",
|
|
" related_groupsl = sorted(related_groups, key=lambda el: len(el))\n",
|
|
" for g in related_groupsl:\n",
|
|
" if max_group_size is None or len(gnew) + len(g) <= max_group_size:\n",
|
|
" gnew = gnew.union(g)\n",
|
|
" if g in groups:\n",
|
|
" groups.remove(g)\n",
|
|
"\n",
|
|
" groups.add(frozenset(gnew))\n",
|
|
"\n",
|
|
" return groups\n",
|
|
"\n",
|
|
"\n",
|
|
"def get_grouped_docs(\n",
|
|
" nodes: List[TextNode],\n",
|
|
" max_group_size: Optional[int] = DEFAULT_MAX_GROUP_SIZE,\n",
|
|
") -> List[TextNode]:\n",
|
|
" \"\"\"Gets list of documents that are grouped.\n",
|
|
"\n",
|
|
" Args:\n",
|
|
" nodes (t.List[TextNode]): Input list\n",
|
|
" max_group_size (Optional[int], optional): max group size, None if no max group size. Defaults to DEFAULT_MAX_GROUP_SIZE.\n",
|
|
"\n",
|
|
" Returns:\n",
|
|
" t.List[TextNode]: Output list\n",
|
|
" \"\"\"\n",
|
|
" # node IDs\n",
|
|
" nodes_str = [node.id_ for node in nodes]\n",
|
|
" # maps node ID -> related node IDs based on that node's relationships\n",
|
|
" adj: Dict[str, List[str]] = {\n",
|
|
" node.id_: [val.node_id for val in node.relationships.values()]\n",
|
|
" for node in nodes\n",
|
|
" }\n",
|
|
" # node ID -> node\n",
|
|
" nodes_dict = {node.id_: node for node in nodes}\n",
|
|
"\n",
|
|
" res = group_docs(nodes_str, adj, max_group_size)\n",
|
|
"\n",
|
|
" ret_nodes = []\n",
|
|
" for g in res:\n",
|
|
" cur_node = TextNode()\n",
|
|
"\n",
|
|
" for node_id in g:\n",
|
|
" cur_node.text += nodes_dict[node_id].text + \"\\n\\n\"\n",
|
|
" cur_node.metadata.update(nodes_dict[node_id].metadata)\n",
|
|
"\n",
|
|
" ret_nodes.append(cur_node)\n",
|
|
"\n",
|
|
" return ret_nodes"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Making the Retriever\n",
|
|
"\n",
|
|
"LongRAG needs a custom retriever, which is shown below:"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"from llama_index.core.retrievers import BaseRetriever\n",
|
|
"from llama_index.core.vector_stores.simple import BasePydanticVectorStore\n",
|
|
"from llama_index.core.schema import QueryBundle, NodeWithScore\n",
|
|
"from llama_index.core.vector_stores.types import VectorStoreQuery\n",
|
|
"from llama_index.core.settings import Settings\n",
|
|
"\n",
|
|
"\n",
|
|
"class LongRAGRetriever(BaseRetriever):\n",
|
|
" \"\"\"Long RAG Retriever.\"\"\"\n",
|
|
"\n",
|
|
" def __init__(\n",
|
|
" self,\n",
|
|
" grouped_nodes: List[TextNode],\n",
|
|
" small_toks: List[TextNode],\n",
|
|
" vector_store: BasePydanticVectorStore,\n",
|
|
" similarity_top_k: int = DEFAULT_TOP_K,\n",
|
|
" ) -> None:\n",
|
|
" \"\"\"Constructor.\n",
|
|
"\n",
|
|
" Args:\n",
|
|
" grouped_nodes (List[TextNode]): Long retrieval units, nodes with docs grouped together based on relationships\n",
|
|
" small_toks (List[TextNode]): Smaller tokens\n",
|
|
" embed_model (BaseEmbedding, optional): Embed model. Defaults to None.\n",
|
|
" similarity_top_k (int, optional): Similarity top k. Defaults to 8.\n",
|
|
" \"\"\"\n",
|
|
" self._grouped_nodes = grouped_nodes\n",
|
|
" self._grouped_nodes_dict = {node.id_: node for node in grouped_nodes}\n",
|
|
" self._small_toks = small_toks\n",
|
|
" self._small_toks_dict = {node.id_: node for node in self._small_toks}\n",
|
|
"\n",
|
|
" self._similarity_top_k = similarity_top_k\n",
|
|
" self._vec_store = vector_store\n",
|
|
" self._embed_model = Settings.embed_model\n",
|
|
"\n",
|
|
" def _retrieve(self, query_bundle: QueryBundle) -> List[NodeWithScore]:\n",
|
|
" \"\"\"Retrieves.\n",
|
|
"\n",
|
|
" Args:\n",
|
|
" query_bundle (QueryBundle): query bundle\n",
|
|
"\n",
|
|
" Returns:\n",
|
|
" List[NodeWithScore]: nodes with scores\n",
|
|
" \"\"\"\n",
|
|
" # make query\n",
|
|
" query_embedding = self._embed_model.get_query_embedding(\n",
|
|
" query_bundle.query_str\n",
|
|
" )\n",
|
|
" vector_store_query = VectorStoreQuery(\n",
|
|
" query_embedding=query_embedding, similarity_top_k=500\n",
|
|
" )\n",
|
|
"\n",
|
|
" # query for answer\n",
|
|
" query_res = self._vec_store.query(vector_store_query)\n",
|
|
"\n",
|
|
" # determine top parents of most similar children (these are long retrieval units)\n",
|
|
" top_parents_set: Set[str] = set()\n",
|
|
" top_parents: List[NodeWithScore] = []\n",
|
|
" for id_, similarity in zip(query_res.ids, query_res.similarities):\n",
|
|
" cur_node = self._small_toks_dict[id_]\n",
|
|
" parent_id = cur_node.ref_doc_id\n",
|
|
" if parent_id not in top_parents_set:\n",
|
|
" top_parents_set.add(parent_id)\n",
|
|
"\n",
|
|
" parent_node = self._grouped_nodes_dict[parent_id]\n",
|
|
" node_with_score = NodeWithScore(\n",
|
|
" node=parent_node, score=similarity\n",
|
|
" )\n",
|
|
" top_parents.append(node_with_score)\n",
|
|
"\n",
|
|
" if len(top_parents_set) >= self._similarity_top_k:\n",
|
|
" break\n",
|
|
"\n",
|
|
" assert len(top_parents) == min(\n",
|
|
" self._similarity_top_k, len(self._grouped_nodes)\n",
|
|
" )\n",
|
|
"\n",
|
|
" return top_parents"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Designing the Workflow"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"LongRAG consists of the following steps:\n",
|
|
"\n",
|
|
"1. Ingesting the data — grouping documents and putting them in long retrieval units, splitting the long retrieval units into smaller tokens to generate embeddings, and indexing the small nodes.\n",
|
|
"2. Constructing the retriever and query engine.\n",
|
|
"3. Querying over the data given a string."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"We define an event that passes the long and small retrieval units into the retriever and query engine."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"from typing import Iterable\n",
|
|
"\n",
|
|
"from llama_index.core import VectorStoreIndex\n",
|
|
"from llama_index.core.llms import LLM\n",
|
|
"from llama_index.core.workflow import Event\n",
|
|
"\n",
|
|
"\n",
|
|
"class LoadNodeEvent(Event):\n",
|
|
" \"\"\"Event for loading nodes.\"\"\"\n",
|
|
"\n",
|
|
" small_nodes: Iterable[TextNode]\n",
|
|
" grouped_nodes: list[TextNode]\n",
|
|
" index: VectorStoreIndex\n",
|
|
" similarity_top_k: int\n",
|
|
" llm: LLM"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"After defining our events, we can write our workflow and steps:"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"from llama_index.core.workflow import (\n",
|
|
" Workflow,\n",
|
|
" step,\n",
|
|
" StartEvent,\n",
|
|
" StopEvent,\n",
|
|
" Context,\n",
|
|
")\n",
|
|
"from llama_index.core import SimpleDirectoryReader\n",
|
|
"from llama_index.core.query_engine import RetrieverQueryEngine\n",
|
|
"\n",
|
|
"\n",
|
|
"class LongRAGWorkflow(Workflow):\n",
|
|
" \"\"\"Long RAG Workflow.\"\"\"\n",
|
|
"\n",
|
|
" @step\n",
|
|
" async def ingest(self, ev: StartEvent) -> LoadNodeEvent | None:\n",
|
|
" \"\"\"Ingestion step.\n",
|
|
"\n",
|
|
" Args:\n",
|
|
" ctx (Context): Context\n",
|
|
" ev (StartEvent): start event\n",
|
|
"\n",
|
|
" Returns:\n",
|
|
" StopEvent | None: stop event with result\n",
|
|
" \"\"\"\n",
|
|
" data_dir: str = ev.get(\"data_dir\")\n",
|
|
" llm: LLM = ev.get(\"llm\")\n",
|
|
" chunk_size: int | None = ev.get(\"chunk_size\")\n",
|
|
" similarity_top_k: int = ev.get(\"similarity_top_k\")\n",
|
|
" small_chunk_size: int = ev.get(\"small_chunk_size\")\n",
|
|
" index: VectorStoreIndex | None = ev.get(\"index\")\n",
|
|
" index_kwargs: dict[str, t.Any] | None = ev.get(\"index_kwargs\")\n",
|
|
"\n",
|
|
" if any(\n",
|
|
" i is None\n",
|
|
" for i in [data_dir, llm, similarity_top_k, small_chunk_size]\n",
|
|
" ):\n",
|
|
" return None\n",
|
|
"\n",
|
|
" if not index:\n",
|
|
" docs = SimpleDirectoryReader(data_dir).load_data()\n",
|
|
" if chunk_size is not None:\n",
|
|
" nodes = split_doc(\n",
|
|
" chunk_size, docs\n",
|
|
" ) # split documents into chunks of chunk_size\n",
|
|
" grouped_nodes = get_grouped_docs(\n",
|
|
" nodes\n",
|
|
" ) # get list of nodes after grouping (groups are combined into one node), these are long retrieval units\n",
|
|
" else:\n",
|
|
" grouped_nodes = docs\n",
|
|
"\n",
|
|
" # split large retrieval units into smaller nodes\n",
|
|
" small_nodes = split_doc(small_chunk_size, grouped_nodes)\n",
|
|
"\n",
|
|
" index_kwargs = index_kwargs or {}\n",
|
|
" index = VectorStoreIndex(small_nodes, **index_kwargs)\n",
|
|
" else:\n",
|
|
" # get smaller nodes from index and form large retrieval units from these nodes\n",
|
|
" small_nodes = index.docstore.docs.values()\n",
|
|
" grouped_nodes = get_grouped_docs(small_nodes, None)\n",
|
|
"\n",
|
|
" return LoadNodeEvent(\n",
|
|
" small_nodes=small_nodes,\n",
|
|
" grouped_nodes=grouped_nodes,\n",
|
|
" index=index,\n",
|
|
" similarity_top_k=similarity_top_k,\n",
|
|
" llm=llm,\n",
|
|
" )\n",
|
|
"\n",
|
|
" @step\n",
|
|
" async def make_query_engine(\n",
|
|
" self, ctx: Context, ev: LoadNodeEvent\n",
|
|
" ) -> StopEvent:\n",
|
|
" \"\"\"Query engine construction step.\n",
|
|
"\n",
|
|
" Args:\n",
|
|
" ctx (Context): context\n",
|
|
" ev (LoadNodeEvent): event\n",
|
|
"\n",
|
|
" Returns:\n",
|
|
" StopEvent: stop event\n",
|
|
" \"\"\"\n",
|
|
" # make retriever and query engine\n",
|
|
" retriever = LongRAGRetriever(\n",
|
|
" grouped_nodes=ev.grouped_nodes,\n",
|
|
" small_toks=ev.small_nodes,\n",
|
|
" similarity_top_k=ev.similarity_top_k,\n",
|
|
" vector_store=ev.index.vector_store,\n",
|
|
" )\n",
|
|
" query_eng = RetrieverQueryEngine.from_args(retriever, ev.llm)\n",
|
|
"\n",
|
|
" return StopEvent(\n",
|
|
" result={\n",
|
|
" \"retriever\": retriever,\n",
|
|
" \"query_engine\": query_eng,\n",
|
|
" \"index\": ev.index,\n",
|
|
" }\n",
|
|
" )\n",
|
|
"\n",
|
|
" @step\n",
|
|
" async def query(self, ctx: Context, ev: StartEvent) -> StopEvent | None:\n",
|
|
" \"\"\"Query step.\n",
|
|
"\n",
|
|
" Args:\n",
|
|
" ctx (Context): context\n",
|
|
" ev (StartEvent): start event\n",
|
|
"\n",
|
|
" Returns:\n",
|
|
" StopEvent | None: stop event with result\n",
|
|
" \"\"\"\n",
|
|
" query_str: str | None = ev.get(\"query_str\")\n",
|
|
" query_eng = ev.get(\"query_eng\")\n",
|
|
"\n",
|
|
" if query_str is None:\n",
|
|
" return None\n",
|
|
"\n",
|
|
" result = query_eng.query(query_str)\n",
|
|
" return StopEvent(result=result)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"Walkthrough:\n",
|
|
"- There are 2 entry points: one for ingesting and indexing and another for querying.\n",
|
|
"- When ingesting, it first reads the documents, splits them into smaller nodes, and indexes them. After that, it sends a `LoadNodeEvent` which triggers the execution of `make_query_engine`, constructing a retriever and query engine from the nodes. It returns a result of the retriever, the query engine, and the index.\n",
|
|
"- When querying, it takes in the query from the `StartEvent`, feeds it into the query engine in the context, and returns the result of the query.\n",
|
|
"- The context is used to store the query engine."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Running the Workflow"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"from llama_index.llms.openai import OpenAI\n",
|
|
"\n",
|
|
"wf = LongRAGWorkflow(timeout=60)\n",
|
|
"llm = OpenAI(\"gpt-4o\")\n",
|
|
"data_dir = \"data\"\n",
|
|
"\n",
|
|
"# initialize the workflow\n",
|
|
"result = await wf.run(\n",
|
|
" data_dir=data_dir,\n",
|
|
" llm=llm,\n",
|
|
" chunk_size=DEFAULT_CHUNK_SIZE,\n",
|
|
" similarity_top_k=DEFAULT_TOP_K,\n",
|
|
" small_chunk_size=DEFAULT_SMALL_CHUNK_SIZE,\n",
|
|
")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"data": {
|
|
"text/markdown": [
|
|
"Pittsburgh can become a startup hub by leveraging its increasing population of young people, particularly those aged 25 to 29, who are crucial for the startup ecosystem. The city should encourage the youth-driven food boom, preserve historic buildings, and make the city more bicycle and pedestrian-friendly. Additionally, Carnegie Mellon University (CMU) should focus on being an even better research university to attract ambitious talent. The city should also foster a culture of tolerance and gradually build an investor community.\n",
|
|
"\n",
|
|
"There are two types of moderates: intentional moderates and accidental moderates. Intentional moderates deliberately choose positions midway between the extremes of right and left, while accidental moderates form their opinions independently on each issue, resulting in a broad range of views that average to a moderate position."
|
|
],
|
|
"text/plain": [
|
|
"<IPython.core.display.Markdown object>"
|
|
]
|
|
},
|
|
"metadata": {},
|
|
"output_type": "display_data"
|
|
}
|
|
],
|
|
"source": [
|
|
"from IPython.display import display, Markdown\n",
|
|
"\n",
|
|
"# run a query\n",
|
|
"res = await wf.run(\n",
|
|
" query_str=\"How can Pittsburgh become a startup hub, and what are the two types of moderates?\",\n",
|
|
" query_eng=result[\"query_engine\"],\n",
|
|
")\n",
|
|
"display(Markdown(str(res)))"
|
|
]
|
|
}
|
|
],
|
|
"metadata": {
|
|
"kernelspec": {
|
|
"display_name": "llama-index-RvIdVF4_-py3.11",
|
|
"language": "python",
|
|
"name": "python3"
|
|
},
|
|
"language_info": {
|
|
"codemirror_mode": {
|
|
"name": "ipython",
|
|
"version": 3
|
|
},
|
|
"file_extension": ".py",
|
|
"mimetype": "text/x-python",
|
|
"name": "python",
|
|
"nbconvert_exporter": "python",
|
|
"pygments_lexer": "ipython3"
|
|
}
|
|
},
|
|
"nbformat": 4,
|
|
"nbformat_minor": 2
|
|
}
|