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
756 lines
22 KiB
Plaintext
756 lines
22 KiB
Plaintext
{
|
|
"cells": [
|
|
{
|
|
"attachments": {},
|
|
"cell_type": "markdown",
|
|
"id": "0d7688a7",
|
|
"metadata": {},
|
|
"source": [
|
|
"<a href=\"https://colab.research.google.com/github/run-llama/llama_index/blob/main/docs/examples/output_parsing/openai_pydantic_program.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "530c973e-916d-4c9e-9365-e2d5306d7e3d",
|
|
"metadata": {},
|
|
"source": [
|
|
"# OpenAI Pydantic Program"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "18461ba1-6978-4b5b-861e-6dceec36857b",
|
|
"metadata": {},
|
|
"source": [
|
|
"This guide shows you how to generate structured data with [new OpenAI API](https://openai.com/blog/function-calling-and-other-api-updates) via LlamaIndex. The user just needs to specify a Pydantic object.\n",
|
|
"\n",
|
|
"We demonstrate two settings:\n",
|
|
"- Extraction into an `Album` object (which can contain a list of Song objects)\n",
|
|
"- Extraction into a `DirectoryTree` object (which can contain recursive Node objects)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "8fcefc79-68b4-4481-b1ef-a902fc12e4c8",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Extraction into `Album`\n",
|
|
"\n",
|
|
"This is a simple example of parsing an output into an `Album` schema, which can contain multiple songs."
|
|
]
|
|
},
|
|
{
|
|
"attachments": {},
|
|
"cell_type": "markdown",
|
|
"id": "81e5dde0",
|
|
"metadata": {},
|
|
"source": [
|
|
"If you're opening this Notebook on colab, you will probably need to install LlamaIndex 🦙."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "511a8171",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"%pip install llama-index-llms-openai\n",
|
|
"%pip install llama-index-program-openai"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "b2833cea",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"%pip install llama-index"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "f7a83b49-5c34-45d5-8cf4-62f348fb1299",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"from pydantic import BaseModel\n",
|
|
"from typing import List\n",
|
|
"\n",
|
|
"from llama_index.program.openai import OpenAIPydanticProgram"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "6311f7ae",
|
|
"metadata": {},
|
|
"source": [
|
|
"### Without docstring in Model"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "0563f1ba-8086-4dcc-ba35-bfda31c45ae4",
|
|
"metadata": {},
|
|
"source": [
|
|
"Define output schema (without docstring)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "42053ea8-2580-4639-9dcf-566e8427c44e",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"class Song(BaseModel):\n",
|
|
" title: str\n",
|
|
" length_seconds: int\n",
|
|
"\n",
|
|
"\n",
|
|
"class Album(BaseModel):\n",
|
|
" name: str\n",
|
|
" artist: str\n",
|
|
" songs: List[Song]"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "4afff44e-a746-4b9f-85a9-72058bcdd29f",
|
|
"metadata": {},
|
|
"source": [
|
|
"Define openai pydantic program"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "fe756697-c299-4f9a-a108-944b6693f824",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"prompt_template_str = \"\"\"\\\n",
|
|
"Generate an example album, with an artist and a list of songs. \\\n",
|
|
"Using the movie {movie_name} as inspiration.\\\n",
|
|
"\"\"\"\n",
|
|
"program = OpenAIPydanticProgram.from_defaults(\n",
|
|
" output_cls=Album, prompt_template_str=prompt_template_str, verbose=True\n",
|
|
")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "b7be01dc-433e-4485-bab0-36a04c3afbcb",
|
|
"metadata": {},
|
|
"source": [
|
|
"Run program to get structured output. "
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "25d02228-2907-4810-932e-83ec9fc71f6b",
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"name": "stdout",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"Function call: Album with args: {\n",
|
|
" \"name\": \"The Shining\",\n",
|
|
" \"artist\": \"Various Artists\",\n",
|
|
" \"songs\": [\n",
|
|
" {\n",
|
|
" \"title\": \"Main Title\",\n",
|
|
" \"length_seconds\": 180\n",
|
|
" },\n",
|
|
" {\n",
|
|
" \"title\": \"Opening Credits\",\n",
|
|
" \"length_seconds\": 120\n",
|
|
" },\n",
|
|
" {\n",
|
|
" \"title\": \"The Overlook Hotel\",\n",
|
|
" \"length_seconds\": 240\n",
|
|
" },\n",
|
|
" {\n",
|
|
" \"title\": \"Redrum\",\n",
|
|
" \"length_seconds\": 150\n",
|
|
" },\n",
|
|
" {\n",
|
|
" \"title\": \"Here's Johnny!\",\n",
|
|
" \"length_seconds\": 200\n",
|
|
" }\n",
|
|
" ]\n",
|
|
"}\n"
|
|
]
|
|
}
|
|
],
|
|
"source": [
|
|
"output = program(\n",
|
|
" movie_name=\"The Shining\", description=\"Data model for an album.\"\n",
|
|
")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "4c2af9a5",
|
|
"metadata": {},
|
|
"source": [
|
|
"### With docstring in Model"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "35c01bec",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"class Song(BaseModel):\n",
|
|
" \"\"\"Data model for a song.\"\"\"\n",
|
|
"\n",
|
|
" title: str\n",
|
|
" length_seconds: int\n",
|
|
"\n",
|
|
"\n",
|
|
"class Album(BaseModel):\n",
|
|
" \"\"\"Data model for an album.\"\"\"\n",
|
|
"\n",
|
|
" name: str\n",
|
|
" artist: str\n",
|
|
" songs: List[Song]"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "22268e2a",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"prompt_template_str = \"\"\"\\\n",
|
|
"Generate an example album, with an artist and a list of songs. \\\n",
|
|
"Using the movie {movie_name} as inspiration.\\\n",
|
|
"\"\"\"\n",
|
|
"program = OpenAIPydanticProgram.from_defaults(\n",
|
|
" output_cls=Album, prompt_template_str=prompt_template_str, verbose=True\n",
|
|
")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "9411d0f1",
|
|
"metadata": {},
|
|
"source": [
|
|
"Run program to get structured output. "
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "066e9c2d",
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"name": "stdout",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"Function call: Album with args: {\n",
|
|
" \"name\": \"The Shining\",\n",
|
|
" \"artist\": \"Various Artists\",\n",
|
|
" \"songs\": [\n",
|
|
" {\n",
|
|
" \"title\": \"Main Title\",\n",
|
|
" \"length_seconds\": 180\n",
|
|
" },\n",
|
|
" {\n",
|
|
" \"title\": \"Opening Credits\",\n",
|
|
" \"length_seconds\": 120\n",
|
|
" },\n",
|
|
" {\n",
|
|
" \"title\": \"The Overlook Hotel\",\n",
|
|
" \"length_seconds\": 240\n",
|
|
" },\n",
|
|
" {\n",
|
|
" \"title\": \"Redrum\",\n",
|
|
" \"length_seconds\": 150\n",
|
|
" },\n",
|
|
" {\n",
|
|
" \"title\": \"Here's Johnny\",\n",
|
|
" \"length_seconds\": 200\n",
|
|
" }\n",
|
|
" ]\n",
|
|
"}\n"
|
|
]
|
|
}
|
|
],
|
|
"source": [
|
|
"output = program(movie_name=\"The Shining\")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "27ec0777-28d5-494b-b419-daf6bce2b20e",
|
|
"metadata": {},
|
|
"source": [
|
|
"The output is a valid Pydantic object that we can then use to call functions/APIs. "
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "3e51bcf4-e7df-47b9-b380-8e5b900a31e1",
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"data": {
|
|
"text/plain": [
|
|
"Album(name='The Shining', artist='Various Artists', songs=[Song(title='Main Title', length_seconds=180), Song(title='Opening Credits', length_seconds=120), Song(title='The Overlook Hotel', length_seconds=240), Song(title='Redrum', length_seconds=150), Song(title=\"Here's Johnny\", length_seconds=200)])"
|
|
]
|
|
},
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"output_type": "execute_result"
|
|
}
|
|
],
|
|
"source": [
|
|
"output"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "cbc6cd54",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Stream partial intermediate Pydantic Objects"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "9f237a81",
|
|
"metadata": {},
|
|
"source": [
|
|
"Instead of waiting for the Function Call to generate the entire JSON, we can use the `stream_partial_objects()` method of the `program` to stream valid intermediate instances of the Pydantic Output class as soon as they're available 🔥 "
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "06312c87",
|
|
"metadata": {},
|
|
"source": [
|
|
"First let's define the Output Pydantic class"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "213794ef",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"from pydantic import BaseModel, Field\n",
|
|
"\n",
|
|
"\n",
|
|
"class CharacterInfo(BaseModel):\n",
|
|
" \"\"\"Information about a character.\"\"\"\n",
|
|
"\n",
|
|
" character_name: str\n",
|
|
" name: str = Field(..., description=\"Name of the actor/actress\")\n",
|
|
" hometown: str\n",
|
|
"\n",
|
|
"\n",
|
|
"class Characters(BaseModel):\n",
|
|
" \"\"\"List of characters.\"\"\"\n",
|
|
"\n",
|
|
" characters: list[CharacterInfo] = Field(default_factory=list)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "254b4254",
|
|
"metadata": {},
|
|
"source": [
|
|
"Now we'll initialilze the program with prompt template"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "4718559f",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"from llama_index.program.openai import OpenAIPydanticProgram\n",
|
|
"\n",
|
|
"prompt_template_str = \"Information about 3 characters from the movie: {movie}\"\n",
|
|
"\n",
|
|
"program = OpenAIPydanticProgram.from_defaults(\n",
|
|
" output_cls=Characters, prompt_template_str=prompt_template_str\n",
|
|
")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "935cae78",
|
|
"metadata": {},
|
|
"source": [
|
|
"Finally we stream the partial objects using the `stream_partial_objects()` method"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "221df824",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"for partial_object in program.stream_partial_objects(movie=\"Harry Potter\"):\n",
|
|
" # send the partial object to the frontend for better user experience\n",
|
|
" print(partial_object)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "025c567c",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Extracting List of `Album` (with Parallel Function Calling)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "273d1716",
|
|
"metadata": {},
|
|
"source": [
|
|
"With the latest [parallel function calling](https://platform.openai.com/docs/guides/function-calling/parallel-function-calling) feature from OpenAI, we can simultaneously extract multiple structured data from a single prompt!"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "a3956a6f",
|
|
"metadata": {},
|
|
"source": [
|
|
"To do this, we need to:\n",
|
|
"1. pick one of the latest models (e.g. `gpt-3.5-turbo-1106`), and \n",
|
|
"2. set `allow_multiple` to True in our `OpenAIPydanticProgram` (if not, it will only return the first object, and raise a warning)."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "74b49a50",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"from llama_index.llms.openai import OpenAI\n",
|
|
"\n",
|
|
"prompt_template_str = \"\"\"\\\n",
|
|
"Generate 4 albums about spring, summer, fall, and winter.\n",
|
|
"\"\"\"\n",
|
|
"program = OpenAIPydanticProgram.from_defaults(\n",
|
|
" output_cls=Album,\n",
|
|
" llm=OpenAI(model=\"gpt-3.5-turbo-1106\"),\n",
|
|
" prompt_template_str=prompt_template_str,\n",
|
|
" allow_multiple=True,\n",
|
|
" verbose=True,\n",
|
|
")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "23428a9a",
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"name": "stdout",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"Function call: Album with args: {\"name\": \"Spring\", \"artist\": \"Various Artists\", \"songs\": [{\"title\": \"Blossom\", \"length_seconds\": 180}, {\"title\": \"Sunshine\", \"length_seconds\": 240}, {\"title\": \"Renewal\", \"length_seconds\": 200}]}\n",
|
|
"Function call: Album with args: {\"name\": \"Summer\", \"artist\": \"Beach Boys\", \"songs\": [{\"title\": \"Beach Party\", \"length_seconds\": 220}, {\"title\": \"Heatwave\", \"length_seconds\": 260}, {\"title\": \"Vacation\", \"length_seconds\": 180}]}\n",
|
|
"Function call: Album with args: {\"name\": \"Fall\", \"artist\": \"Autumn Leaves\", \"songs\": [{\"title\": \"Golden Days\", \"length_seconds\": 210}, {\"title\": \"Harvest Moon\", \"length_seconds\": 240}, {\"title\": \"Crisp Air\", \"length_seconds\": 190}]}\n",
|
|
"Function call: Album with args: {\"name\": \"Winter\", \"artist\": \"Snowflakes\", \"songs\": [{\"title\": \"Frosty Morning\", \"length_seconds\": 190}, {\"title\": \"Snowfall\", \"length_seconds\": 220}, {\"title\": \"Cozy Nights\", \"length_seconds\": 250}]}\n"
|
|
]
|
|
}
|
|
],
|
|
"source": [
|
|
"output = program()"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "a4c3ad95",
|
|
"metadata": {},
|
|
"source": [
|
|
"The output is a list of valid Pydantic object."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "8a20fdce",
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"data": {
|
|
"text/plain": [
|
|
"[Album(name='Spring', artist='Various Artists', songs=[Song(title='Blossom', length_seconds=180), Song(title='Sunshine', length_seconds=240), Song(title='Renewal', length_seconds=200)]),\n",
|
|
" Album(name='Summer', artist='Beach Boys', songs=[Song(title='Beach Party', length_seconds=220), Song(title='Heatwave', length_seconds=260), Song(title='Vacation', length_seconds=180)]),\n",
|
|
" Album(name='Fall', artist='Autumn Leaves', songs=[Song(title='Golden Days', length_seconds=210), Song(title='Harvest Moon', length_seconds=240), Song(title='Crisp Air', length_seconds=190)]),\n",
|
|
" Album(name='Winter', artist='Snowflakes', songs=[Song(title='Frosty Morning', length_seconds=190), Song(title='Snowfall', length_seconds=220), Song(title='Cozy Nights', length_seconds=250)])]"
|
|
]
|
|
},
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"output_type": "execute_result"
|
|
}
|
|
],
|
|
"source": [
|
|
"output"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "58eb69c1-bafe-49e4-9d14-f676662e74ac",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Extraction into `Album` (Streaming)\n",
|
|
"\n",
|
|
"We also support streaming a list of objects through our `stream_list` function.\n",
|
|
"\n",
|
|
"Full credits to this idea go to `openai_function_call` repo: https://github.com/jxnl/openai_function_call/tree/main/examples/streaming_multitask"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "0a155e9c-a7e0-46b8-addb-7fbf844c0107",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"prompt_template_str = \"{input_str}\"\n",
|
|
"program = OpenAIPydanticProgram.from_defaults(\n",
|
|
" output_cls=Album,\n",
|
|
" prompt_template_str=prompt_template_str,\n",
|
|
" verbose=False,\n",
|
|
")\n",
|
|
"\n",
|
|
"output = program.stream_list(\n",
|
|
" input_str=\"make up 5 random albums\",\n",
|
|
")\n",
|
|
"for obj in output:\n",
|
|
" print(obj.json(indent=2))"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "2d6159d2-c967-4523-ad3c-82bc17009b2c",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Extraction into `DirectoryTree` object\n",
|
|
"\n",
|
|
"This is directly inspired by jxnl's awesome repo here: https://github.com/jxnl/openai_function_call.\n",
|
|
"\n",
|
|
"That repository shows how you can use OpenAI's function API to parse recursive Pydantic objects. The main requirement is that you want to \"wrap\" a recursive Pydantic object with a non-recursive one.\n",
|
|
"\n",
|
|
"Here we show an example in a \"directory\" setting, where a `DirectoryTree` object wraps recursive `Node` objects, to parse a file structure."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "b58f6a12-3f5c-414b-80df-4492f6e18be5",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"# NOTE: defining recursive objects in a notebook causes errors\n",
|
|
"from directory import DirectoryTree, Node"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "bc1c7eeb-f0a3-4d72-86ee-c6b63e01b0ff",
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"data": {
|
|
"text/plain": [
|
|
"{'title': 'DirectoryTree',\n",
|
|
" 'description': 'Container class representing a directory tree.\\n\\nArgs:\\n root (Node): The root node of the tree.',\n",
|
|
" 'type': 'object',\n",
|
|
" 'properties': {'root': {'title': 'Root',\n",
|
|
" 'description': 'Root folder of the directory tree',\n",
|
|
" 'allOf': [{'$ref': '#/definitions/Node'}]}},\n",
|
|
" 'required': ['root'],\n",
|
|
" 'definitions': {'NodeType': {'title': 'NodeType',\n",
|
|
" 'description': 'Enumeration representing the types of nodes in a filesystem.',\n",
|
|
" 'enum': ['file', 'folder'],\n",
|
|
" 'type': 'string'},\n",
|
|
" 'Node': {'title': 'Node',\n",
|
|
" 'description': 'Class representing a single node in a filesystem. Can be either a file or a folder.\\nNote that a file cannot have children, but a folder can.\\n\\nArgs:\\n name (str): The name of the node.\\n children (List[Node]): The list of child nodes (if any).\\n node_type (NodeType): The type of the node, either a file or a folder.',\n",
|
|
" 'type': 'object',\n",
|
|
" 'properties': {'name': {'title': 'Name',\n",
|
|
" 'description': 'Name of the folder',\n",
|
|
" 'type': 'string'},\n",
|
|
" 'children': {'title': 'Children',\n",
|
|
" 'description': 'List of children nodes, only applicable for folders, files cannot have children',\n",
|
|
" 'type': 'array',\n",
|
|
" 'items': {'$ref': '#/definitions/Node'}},\n",
|
|
" 'node_type': {'description': 'Either a file or folder, use the name to determine which it could be',\n",
|
|
" 'default': 'file',\n",
|
|
" 'allOf': [{'$ref': '#/definitions/NodeType'}]}},\n",
|
|
" 'required': ['name']}}}"
|
|
]
|
|
},
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"output_type": "execute_result"
|
|
}
|
|
],
|
|
"source": [
|
|
"DirectoryTree.schema()"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "02c4a7a1-f145-40bc-83b8-4153a531a8eb",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"program = OpenAIPydanticProgram.from_defaults(\n",
|
|
" output_cls=DirectoryTree,\n",
|
|
" prompt_template_str=\"{input_str}\",\n",
|
|
" verbose=True,\n",
|
|
")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "c88cf49f-a52f-4415-bddc-14d70c897629",
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"name": "stdout",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"Function call: DirectoryTree with args: {\n",
|
|
" \"root\": {\n",
|
|
" \"name\": \"root\",\n",
|
|
" \"children\": [\n",
|
|
" {\n",
|
|
" \"name\": \"folder1\",\n",
|
|
" \"children\": [\n",
|
|
" {\n",
|
|
" \"name\": \"file1.txt\",\n",
|
|
" \"children\": [],\n",
|
|
" \"node_type\": \"file\"\n",
|
|
" },\n",
|
|
" {\n",
|
|
" \"name\": \"file2.txt\",\n",
|
|
" \"children\": [],\n",
|
|
" \"node_type\": \"file\"\n",
|
|
" }\n",
|
|
" ],\n",
|
|
" \"node_type\": \"folder\"\n",
|
|
" },\n",
|
|
" {\n",
|
|
" \"name\": \"folder2\",\n",
|
|
" \"children\": [\n",
|
|
" {\n",
|
|
" \"name\": \"file3.txt\",\n",
|
|
" \"children\": [],\n",
|
|
" \"node_type\": \"file\"\n",
|
|
" },\n",
|
|
" {\n",
|
|
" \"name\": \"subfolder1\",\n",
|
|
" \"children\": [\n",
|
|
" {\n",
|
|
" \"name\": \"file4.txt\",\n",
|
|
" \"children\": [],\n",
|
|
" \"node_type\": \"file\"\n",
|
|
" }\n",
|
|
" ],\n",
|
|
" \"node_type\": \"folder\"\n",
|
|
" }\n",
|
|
" ],\n",
|
|
" \"node_type\": \"folder\"\n",
|
|
" }\n",
|
|
" ],\n",
|
|
" \"node_type\": \"folder\"\n",
|
|
" }\n",
|
|
"}\n"
|
|
]
|
|
}
|
|
],
|
|
"source": [
|
|
"input_str = \"\"\"\n",
|
|
"root\n",
|
|
"├── folder1\n",
|
|
"│ ├── file1.txt\n",
|
|
"│ └── file2.txt\n",
|
|
"└── folder2\n",
|
|
" ├── file3.txt\n",
|
|
" └── subfolder1\n",
|
|
" └── file4.txt\n",
|
|
"\"\"\"\n",
|
|
"\n",
|
|
"output = program(input_str=input_str)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "0bf23c34-2cc7-4d0f-9389-75f7cf6bf9a2",
|
|
"metadata": {},
|
|
"source": [
|
|
"The output is a full DirectoryTree structure with recursive `Node` objects."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "3885032f-0f3a-4afb-9157-54851e810843",
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"data": {
|
|
"text/plain": [
|
|
"DirectoryTree(root=Node(name='root', children=[Node(name='folder1', children=[Node(name='file1.txt', children=[], node_type=<NodeType.FILE: 'file'>), Node(name='file2.txt', children=[], node_type=<NodeType.FILE: 'file'>)], node_type=<NodeType.FOLDER: 'folder'>), Node(name='folder2', children=[Node(name='file3.txt', children=[], node_type=<NodeType.FILE: 'file'>), Node(name='subfolder1', children=[Node(name='file4.txt', children=[], node_type=<NodeType.FILE: 'file'>)], node_type=<NodeType.FOLDER: 'folder'>)], node_type=<NodeType.FOLDER: 'folder'>)], node_type=<NodeType.FOLDER: 'folder'>))"
|
|
]
|
|
},
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"output_type": "execute_result"
|
|
}
|
|
],
|
|
"source": [
|
|
"output"
|
|
]
|
|
}
|
|
],
|
|
"metadata": {
|
|
"kernelspec": {
|
|
"display_name": "llama-index-vs8PXMh0-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": 5
|
|
}
|