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
435 lines
11 KiB
Plaintext
435 lines
11 KiB
Plaintext
{
|
|
"cells": [
|
|
{
|
|
"attachments": {},
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"<a href=\"https://colab.research.google.com/github/run-llama/llama_index/blob/main/docs/examples/query_engine/pydantic_query_engine.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"# Query Engine with Pydantic Outputs\n",
|
|
"\n",
|
|
"Every query engine has support for integrated structured responses using the following `response_mode`s in `RetrieverQueryEngine`:\n",
|
|
"- `refine`\n",
|
|
"- `compact`\n",
|
|
"- `tree_summarize`\n",
|
|
"- `accumulate` (beta, requires extra parsing to convert to objects)\n",
|
|
"- `compact_accumulate` (beta, requires extra parsing to convert to objects)\n",
|
|
"\n",
|
|
"In this notebook, we walk through a small example demonstrating the usage.\n",
|
|
"\n",
|
|
"Under the hood, every LLM response will be a pydantic object. If that response needs to be refined or summarized, it is converted into a JSON string for the next response. Then, the final response is returned as a pydantic object.\n",
|
|
"\n",
|
|
"**NOTE:** This can technically work with any LLM, but non-openai is support is still in development and considered beta."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Setup"
|
|
]
|
|
},
|
|
{
|
|
"attachments": {},
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"If you're opening this Notebook on colab, you will probably need to install LlamaIndex 🦙."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"%pip install llama-index-llms-anthropic\n",
|
|
"%pip install llama-index-llms-openai"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"!pip install llama-index"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"import os\n",
|
|
"import openai\n",
|
|
"\n",
|
|
"os.environ[\"OPENAI_API_KEY\"] = \"sk-...\"\n",
|
|
"openai.api_key = os.environ[\"OPENAI_API_KEY\"]"
|
|
]
|
|
},
|
|
{
|
|
"attachments": {},
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"Download Data"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"!mkdir -p 'data/paul_graham/'\n",
|
|
"!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/examples/data/paul_graham/paul_graham_essay.txt' -O 'data/paul_graham/paul_graham_essay.txt'"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"from llama_index.core import SimpleDirectoryReader\n",
|
|
"\n",
|
|
"documents = SimpleDirectoryReader(\"./data/paul_graham\").load_data()"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"### Create our Pydanitc Output Object"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"from typing import List\n",
|
|
"from pydantic import BaseModel\n",
|
|
"\n",
|
|
"\n",
|
|
"class Biography(BaseModel):\n",
|
|
" \"\"\"Data model for a biography.\"\"\"\n",
|
|
"\n",
|
|
" name: str\n",
|
|
" best_known_for: List[str]\n",
|
|
" extra_info: str"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Create the Index + Query Engine (OpenAI)\n",
|
|
"\n",
|
|
"When using OpenAI, the function calling API will be leveraged for reliable structured outputs."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"from llama_index.core import VectorStoreIndex\n",
|
|
"from llama_index.llms.openai import OpenAI\n",
|
|
"\n",
|
|
"llm = OpenAI(model=\"gpt-3.5-turbo\", temperature=0.1)\n",
|
|
"\n",
|
|
"index = VectorStoreIndex.from_documents(\n",
|
|
" documents,\n",
|
|
")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"query_engine = index.as_query_engine(\n",
|
|
" output_cls=Biography, response_mode=\"compact\", llm=llm\n",
|
|
")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"response = query_engine.query(\"Who is Paul Graham?\")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"name": "stdout",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"Paul Graham\n",
|
|
"['working on Bel', 'co-founding Viaweb', 'creating the programming language Arc']\n",
|
|
"Paul Graham is a computer scientist, entrepreneur, and writer. He is best known for his work on Bel, a programming language, and for co-founding Viaweb, an early web application company that was later acquired by Yahoo. Graham also created the programming language Arc. He has written numerous essays on topics such as startups, programming, and life.\n"
|
|
]
|
|
}
|
|
],
|
|
"source": [
|
|
"print(response.name)\n",
|
|
"print(response.best_known_for)\n",
|
|
"print(response.extra_info)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"name": "stdout",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"<class '__main__.Biography'>\n"
|
|
]
|
|
}
|
|
],
|
|
"source": [
|
|
"# get the full pydanitc object\n",
|
|
"print(type(response.response))"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Create the Index + Query Engine (Non-OpenAI, Beta)\n",
|
|
"\n",
|
|
"When using an LLM that does not support function calling, we rely on the LLM to write the JSON itself, and we parse the JSON into the proper pydantic object."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"import os\n",
|
|
"\n",
|
|
"os.environ[\"ANTHROPIC_API_KEY\"] = \"sk-...\""
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"from llama_index.core import VectorStoreIndex\n",
|
|
"from llama_index.llms.anthropic import Anthropic\n",
|
|
"\n",
|
|
"llm = Anthropic(model=\"claude-instant-1.2\", temperature=0.1)\n",
|
|
"\n",
|
|
"index = VectorStoreIndex.from_documents(\n",
|
|
" documents,\n",
|
|
")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"query_engine = index.as_query_engine(\n",
|
|
" output_cls=Biography, response_mode=\"tree_summarize\", llm=llm\n",
|
|
")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"response = query_engine.query(\"Who is Paul Graham?\")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"name": "stdout",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"Paul Graham\n",
|
|
"['Co-founder of Y Combinator', 'Essayist and programmer']\n",
|
|
"He is known for creating Viaweb, one of the first web application builders, and for founding Y Combinator, one of the world's top startup accelerators. Graham has also written extensively about technology, investing, and philosophy.\n"
|
|
]
|
|
}
|
|
],
|
|
"source": [
|
|
"print(response.name)\n",
|
|
"print(response.best_known_for)\n",
|
|
"print(response.extra_info)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"name": "stdout",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"<class '__main__.Biography'>\n"
|
|
]
|
|
}
|
|
],
|
|
"source": [
|
|
"# get the full pydanitc object\n",
|
|
"print(type(response.response))"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Accumulate Examples (Beta)\n",
|
|
"\n",
|
|
"Accumulate with pydantic objects requires some extra parsing. This is still a beta feature, but it's still possible to get accumulate pydantic objects."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"from typing import List\n",
|
|
"from pydantic import BaseModel\n",
|
|
"\n",
|
|
"\n",
|
|
"class Company(BaseModel):\n",
|
|
" \"\"\"Data model for a companies mentioned.\"\"\"\n",
|
|
"\n",
|
|
" company_name: str\n",
|
|
" context_info: str"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"from llama_index.core import VectorStoreIndex,\n",
|
|
"from llama_index.llms.openai import OpenAI\n",
|
|
"\n",
|
|
"llm = OpenAI(model=\"gpt-3.5-turbo\", temperature=0.1)\n",
|
|
"\n",
|
|
"index = VectorStoreIndex.from_documents(\n",
|
|
" documents,\n",
|
|
")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"query_engine = index.as_query_engine(\n",
|
|
" output_cls=Company, response_mode=\"accumulate\", llm=llm\n",
|
|
")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"response = query_engine.query(\"What companies are mentioned in the text?\")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"In accumulate, responses are separated by a default separator, and prepended with a prefix."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"companies = []\n",
|
|
"\n",
|
|
"# split by the default separator\n",
|
|
"for response_str in str(response).split(\"\\n---------------------\\n\"):\n",
|
|
" # remove the prefix -- every response starts like `Response 1: {...}`\n",
|
|
" # so, we find the first bracket and remove everything before it\n",
|
|
" response_str = response_str[response_str.find(\"{\") :]\n",
|
|
" companies.append(Company.parse_raw(response_str))"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"name": "stdout",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"[Company(company_name='Yahoo', context_info='Yahoo bought us'), Company(company_name='Yahoo', context_info=\"I'd been meaning to since Yahoo bought us\")]\n"
|
|
]
|
|
}
|
|
],
|
|
"source": [
|
|
"print(companies)"
|
|
]
|
|
}
|
|
],
|
|
"metadata": {
|
|
"kernelspec": {
|
|
"display_name": "llama-index",
|
|
"language": "python",
|
|
"name": "llama-index"
|
|
},
|
|
"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
|
|
}
|