Bagatur/self query doc update (#12461)

This commit is contained in:
Bagatur 2023-10-28 14:37:14 -07:00 committed by GitHub
parent 689853902e
commit e130680d74
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 582 additions and 213 deletions

View File

@ -7,9 +7,33 @@
"source": [
"# Building hotel room search with self-querying retrieval\n",
"\n",
"In this example we'll walk through how to build and iterate on a hotel room search service that leverages an LLM to generate structured filter queries that can then be passed to a vector store.\n",
"\n",
"For an introduction to self-querying retrieval [check out the docs](https://python.langchain.com/docs/modules/data_connection/retrievers/self_query)."
]
},
{
"cell_type": "markdown",
"id": "d621de99-d993-4f4b-b94a-d02b2c7ad4e0",
"metadata": {},
"source": [
"## Imports and data prep\n",
"\n",
"In this example we use `ChatOpenAI` for the model and `ElasticsearchStore` for the vector store, but these can be swapped out with an LLM/ChatModel and [any VectorStore that support self-querying](https://python.langchain.com/docs/integrations/retrievers/self_query/).\n",
"\n",
"Download data from: https://www.kaggle.com/datasets/keshavramaiah/hotel-recommendation"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "8ecd1fbb-bdba-420b-bcc7-5ea8a232ab11",
"metadata": {},
"outputs": [],
"source": [
"!pip install langchain lark openai elasticsearch pandas"
]
},
{
"cell_type": "code",
"execution_count": 1,
@ -1142,7 +1166,9 @@
}
],
"source": [
"results = retriever.get_relevant_documents(\"I want to stay somewhere highly rated along the coast. I want a room with a patio and a fireplace.\")\n",
"results = retriever.get_relevant_documents(\n",
" \"I want to stay somewhere highly rated along the coast. I want a room with a patio and a fireplace.\"\n",
")\n",
"for res in results:\n",
" print(res.page_content)\n",
" print(\"\\n\" + \"-\" * 20 + \"\\n\")"

View File

@ -0,0 +1,555 @@
{
"cells": [
{
"cell_type": "markdown",
"id": "c0bc3390-4bed-49d3-96ce-072badb4110b",
"metadata": {},
"source": [
"# Self-querying\n",
"\n",
":::info\n",
"Head to [Integrations](/docs/integrations/retrievers/self_query) for documentation on vector stores with built-in support for self-querying.\n",
":::\n",
"\n",
"A self-querying retriever is one that, as the name suggests, has the ability to query itself. Specifically, given any natural language query, the retriever uses a query-constructing LLM chain to write a structured query and then applies that structured query to its underlying VectorStore. This allows the retriever to not only use the user-input query for semantic similarity comparison with the contents of stored documents but to also extract filters from the user query on the metadata of stored documents and to execute those filters.\n",
"\n",
"![](https://drive.google.com/uc?id=1OQUN-0MJcDUxmPXofgS7MqReEs720pqS)\n",
"\n",
"## Get started\n",
"For demonstration purposes we'll use a `Chroma` vector store. We've created a small demo set of documents that contain summaries of movies.\n",
"\n",
"**Note:** The self-query retriever requires you to have `lark` package installed."
]
},
{
"cell_type": "code",
"execution_count": 1,
"id": "e1486ca4-9785-4107-90bd-923505542167",
"metadata": {},
"outputs": [],
"source": [
"# !pip install lark chromadb"
]
},
{
"cell_type": "code",
"execution_count": 2,
"id": "beec3e35-3750-408c-9f2a-d92cf0a9a321",
"metadata": {},
"outputs": [],
"source": [
"from langchain.embeddings import OpenAIEmbeddings\n",
"from langchain.schema import Document\n",
"from langchain.vectorstores import Chroma\n",
"\n",
"\n",
"docs = [\n",
" Document(\n",
" page_content=\"A bunch of scientists bring back dinosaurs and mayhem breaks loose\", \n",
" metadata={\"year\": 1993, \"rating\": 7.7, \"genre\": \"science fiction\"})\n",
" ,\n",
" Document(\n",
" page_content=\"Leo DiCaprio gets lost in a dream within a dream within a dream within a ...\", \n",
" metadata={\"year\": 2010, \"director\": \"Christopher Nolan\", \"rating\": 8.2}\n",
" ),\n",
" Document(\n",
" page_content=\"A psychologist / detective gets lost in a series of dreams within dreams within dreams and Inception reused the idea\", \n",
" metadata={\"year\": 2006, \"director\": \"Satoshi Kon\", \"rating\": 8.6}\n",
" ),\n",
" Document(\n",
" page_content=\"A bunch of normal-sized women are supremely wholesome and some men pine after them\", \n",
" metadata={\"year\": 2019, \"director\": \"Greta Gerwig\", \"rating\": 8.3}\n",
" ),\n",
" Document(\n",
" page_content=\"Toys come alive and have a blast doing so\", \n",
" metadata={\"year\": 1995, \"genre\": \"animated\"}\n",
" ),\n",
" Document(\n",
" page_content=\"Three men walk into the Zone, three men walk out of the Zone\", \n",
" metadata={\"year\": 1979, \"rating\": 9.9, \"director\": \"Andrei Tarkovsky\", \"genre\": \"thriller\", \"rating\": 9.9}\n",
" )\n",
"]\n",
"vectorstore = Chroma.from_documents(docs, OpenAIEmbeddings())"
]
},
{
"cell_type": "markdown",
"id": "99771131-1efb-42e2-95f8-2aaa95f37677",
"metadata": {},
"source": [
"### Creating our self-querying retriever\n",
"\n",
"Now we can instantiate our retriever. To do this we'll need to provide some information upfront about the metadata fields that our documents support and a short description of the document contents."
]
},
{
"cell_type": "code",
"execution_count": 28,
"id": "7832ca43-cc17-4375-bf4e-679b99584568",
"metadata": {},
"outputs": [],
"source": [
"from langchain.chat_models import ChatOpenAI\n",
"from langchain.chains.query_constructor.base import AttributeInfo\n",
"from langchain.retrievers.self_query.base import SelfQueryRetriever\n",
"\n",
"\n",
"metadata_field_info=[\n",
" AttributeInfo(\n",
" name=\"genre\",\n",
" description=\"The genre of the movie. One of ['science fiction', 'comedy', 'drama', 'thriller', 'romance', 'action', 'animated']\",\n",
" type=\"string\",\n",
" ),\n",
" AttributeInfo(\n",
" name=\"year\",\n",
" description=\"The year the movie was released\",\n",
" type=\"integer\",\n",
" ),\n",
" AttributeInfo(\n",
" name=\"director\",\n",
" description=\"The name of the movie director\",\n",
" type=\"string\",\n",
" ),\n",
" AttributeInfo(\n",
" name=\"rating\",\n",
" description=\"A 1-10 rating for the movie\",\n",
" type=\"float\"\n",
" ),\n",
"]\n",
"document_content_description = \"Brief summary of a movie\"\n",
"llm = ChatOpenAI(temperature=0)\n",
"retriever = SelfQueryRetriever.from_llm(\n",
" llm, \n",
" vectorstore, \n",
" document_content_description, \n",
" metadata_field_info, \n",
")"
]
},
{
"cell_type": "markdown",
"id": "9c66f4c8-3682-46ac-8f17-0839194888a3",
"metadata": {},
"source": [
"### Testing it out\n",
"\n",
"And now we can actually try using our retriever!"
]
},
{
"cell_type": "code",
"execution_count": 5,
"id": "21c5df28-ea78-4f4e-99d6-489c864d1a04",
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"[Document(page_content='Three men walk into the Zone, three men walk out of the Zone', metadata={'director': 'Andrei Tarkovsky', 'genre': 'thriller', 'rating': 9.9, 'year': 1979}),\n",
" Document(page_content='A psychologist / detective gets lost in a series of dreams within dreams within dreams and Inception reused the idea', metadata={'director': 'Satoshi Kon', 'rating': 8.6, 'year': 2006})]"
]
},
"execution_count": 5,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"# This example only specifies a filter\n",
"retriever.invoke(\"I want to watch a movie rated higher than 8.5\")"
]
},
{
"cell_type": "code",
"execution_count": 6,
"id": "228e5d70-d4cf-43bb-bc8e-3d6f11e784f2",
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"[Document(page_content='A bunch of normal-sized women are supremely wholesome and some men pine after them', metadata={'director': 'Greta Gerwig', 'rating': 8.3, 'year': 2019})]"
]
},
"execution_count": 6,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"# This example specifies a query and a filter\n",
"retriever.invoke(\"Has Greta Gerwig directed any movies about women\")"
]
},
{
"cell_type": "code",
"execution_count": 7,
"id": "8244591e-97b5-4aba-b1e5-fe5e1996cb99",
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"[Document(page_content='A psychologist / detective gets lost in a series of dreams within dreams within dreams and Inception reused the idea', metadata={'director': 'Satoshi Kon', 'rating': 8.6, 'year': 2006}),\n",
" Document(page_content='Three men walk into the Zone, three men walk out of the Zone', metadata={'director': 'Andrei Tarkovsky', 'genre': 'thriller', 'rating': 9.9, 'year': 1979})]"
]
},
"execution_count": 7,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"# This example specifies a composite filter\n",
"retriever.invoke(\"What's a highly rated (above 8.5) science fiction film?\")"
]
},
{
"cell_type": "code",
"execution_count": 8,
"id": "420a6906-66fb-449f-8626-2e399ae5e6a8",
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"[Document(page_content='Toys come alive and have a blast doing so', metadata={'genre': 'animated', 'year': 1995})]"
]
},
"execution_count": 8,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"# This example specifies a query and composite filter\n",
"retriever.invoke(\"What's a movie after 1990 but before 2005 that's all about toys, and preferably is animated\")"
]
},
{
"cell_type": "markdown",
"id": "4f25a751-f1d2-405e-84d6-fe9e4f60ce95",
"metadata": {},
"source": [
"### Filter k\n",
"\n",
"We can also use the self query retriever to specify `k`: the number of documents to fetch.\n",
"\n",
"We can do this by passing `enable_limit=True` to the constructor."
]
},
{
"cell_type": "code",
"execution_count": 9,
"id": "ab56595f-0fb4-4b7f-8fc1-e85eff13255a",
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"[Document(page_content='A bunch of scientists bring back dinosaurs and mayhem breaks loose', metadata={'genre': 'science fiction', 'rating': 7.7, 'year': 1993}),\n",
" Document(page_content='Toys come alive and have a blast doing so', metadata={'genre': 'animated', 'year': 1995})]"
]
},
"execution_count": 9,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"retriever = SelfQueryRetriever.from_llm(\n",
" llm,\n",
" vectorstore,\n",
" document_content_description,\n",
" metadata_field_info,\n",
" enable_limit=True,\n",
")\n",
"\n",
"# This example only specifies a relevant query\n",
"retriever.invoke(\"What are two movies about dinosaurs\")"
]
},
{
"cell_type": "markdown",
"id": "51e144c4-cbf4-4540-92e7-9a68e05f2480",
"metadata": {},
"source": [
"## Constructing from scratch with LCEL\n",
"\n",
"To see what's going on under the hood, and to have more custom control, we can reconstruct our retriever from scratch.\n",
"\n",
"First, we need to create a query-construction chain. This chain will take a |user query and generated a `StructuredQuery` object which captures the filters specified by the user. We provide some helper functions for creating a prompt and output parser. These have a number of tunable params that we'll ignore here for simplicity."
]
},
{
"cell_type": "code",
"execution_count": 29,
"id": "c5f501ac-46c1-4a54-9d23-c0530e8c88f0",
"metadata": {},
"outputs": [],
"source": [
"from langchain.chains.query_constructor.base import (\n",
" get_query_constructor_prompt,\n",
" StructuredQueryOutputParser,\n",
")\n",
"\n",
"prompt = get_query_constructor_prompt(document_content_description, metadata_field_info,)\n",
"output_parser = StructuredQueryOutputParser.from_components()\n",
"query_constructor = prompt | llm | output_parser"
]
},
{
"cell_type": "markdown",
"id": "8deb5d44-632f-4f41-b139-fc811979e6e8",
"metadata": {},
"source": [
"Let's look at our prompt:"
]
},
{
"cell_type": "code",
"execution_count": 30,
"id": "eed553cb-8575-486b-8349-0806b7817a8c",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Your goal is to structure the user's query to match the request schema provided below.\n",
"\n",
"<< Structured Request Schema >>\n",
"When responding use a markdown code snippet with a JSON object formatted in the following schema:\n",
"\n",
"```json\n",
"{\n",
" \"query\": string \\ text string to compare to document contents\n",
" \"filter\": string \\ logical condition statement for filtering documents\n",
"}\n",
"```\n",
"\n",
"The query string should contain only text that is expected to match the contents of documents. Any conditions in the filter should not be mentioned in the query as well.\n",
"\n",
"A logical condition statement is composed of one or more comparison and logical operation statements.\n",
"\n",
"A comparison statement takes the form: `comp(attr, val)`:\n",
"- `comp` (eq | ne | gt | gte | lt | lte | contain | like | in | nin): comparator\n",
"- `attr` (string): name of attribute to apply the comparison to\n",
"- `val` (string): is the comparison value\n",
"\n",
"A logical operation statement takes the form `op(statement1, statement2, ...)`:\n",
"- `op` (and | or | not): logical operator\n",
"- `statement1`, `statement2`, ... (comparison statements or logical operation statements): one or more statements to apply the operation to\n",
"\n",
"Make sure that you only use the comparators and logical operators listed above and no others.\n",
"Make sure that filters only refer to attributes that exist in the data source.\n",
"Make sure that filters only use the attributed names with its function names if there are functions applied on them.\n",
"Make sure that filters only use format `YYYY-MM-DD` when handling timestamp data typed values.\n",
"Make sure that filters take into account the descriptions of attributes and only make comparisons that are feasible given the type of data being stored.\n",
"Make sure that filters are only used as needed. If there are no filters that should be applied return \"NO_FILTER\" for the filter value.\n",
"\n",
"<< Example 1. >>\n",
"Data Source:\n",
"```json\n",
"{\n",
" \"content\": \"Lyrics of a song\",\n",
" \"attributes\": {\n",
" \"artist\": {\n",
" \"type\": \"string\",\n",
" \"description\": \"Name of the song artist\"\n",
" },\n",
" \"length\": {\n",
" \"type\": \"integer\",\n",
" \"description\": \"Length of the song in seconds\"\n",
" },\n",
" \"genre\": {\n",
" \"type\": \"string\",\n",
" \"description\": \"The song genre, one of \"pop\", \"rock\" or \"rap\"\"\n",
" }\n",
" }\n",
"}\n",
"```\n",
"\n",
"User Query:\n",
"What are songs by Taylor Swift or Katy Perry about teenage romance under 3 minutes long in the dance pop genre\n",
"\n",
"Structured Request:\n",
"```json\n",
"{\n",
" \"query\": \"teenager love\",\n",
" \"filter\": \"and(or(eq(\\\"artist\\\", \\\"Taylor Swift\\\"), eq(\\\"artist\\\", \\\"Katy Perry\\\")), lt(\\\"length\\\", 180), eq(\\\"genre\\\", \\\"pop\\\"))\"\n",
"}\n",
"```\n",
"\n",
"\n",
"<< Example 2. >>\n",
"Data Source:\n",
"```json\n",
"{\n",
" \"content\": \"Lyrics of a song\",\n",
" \"attributes\": {\n",
" \"artist\": {\n",
" \"type\": \"string\",\n",
" \"description\": \"Name of the song artist\"\n",
" },\n",
" \"length\": {\n",
" \"type\": \"integer\",\n",
" \"description\": \"Length of the song in seconds\"\n",
" },\n",
" \"genre\": {\n",
" \"type\": \"string\",\n",
" \"description\": \"The song genre, one of \"pop\", \"rock\" or \"rap\"\"\n",
" }\n",
" }\n",
"}\n",
"```\n",
"\n",
"User Query:\n",
"What are songs that were not published on Spotify\n",
"\n",
"Structured Request:\n",
"```json\n",
"{\n",
" \"query\": \"\",\n",
" \"filter\": \"NO_FILTER\"\n",
"}\n",
"```\n",
"\n",
"\n",
"<< Example 3. >>\n",
"Data Source:\n",
"```json\n",
"{\n",
" \"content\": \"Brief summary of a movie\",\n",
" \"attributes\": {\n",
" \"genre\": {\n",
" \"description\": \"The genre of the movie. One of ['science fiction', 'comedy', 'drama', 'thriller', 'romance', 'action', 'animated']\",\n",
" \"type\": \"string\"\n",
" },\n",
" \"year\": {\n",
" \"description\": \"The year the movie was released\",\n",
" \"type\": \"integer\"\n",
" },\n",
" \"director\": {\n",
" \"description\": \"The name of the movie director\",\n",
" \"type\": \"string\"\n",
" },\n",
" \"rating\": {\n",
" \"description\": \"A 1-10 rating for the movie\",\n",
" \"type\": \"float\"\n",
" }\n",
"}\n",
"}\n",
"```\n",
"\n",
"User Query:\n",
"dummy question\n",
"\n",
"Structured Request:\n",
"\n"
]
}
],
"source": [
"print(prompt.format(query=\"dummy question\"))"
]
},
{
"cell_type": "markdown",
"id": "00420512-c395-4661-8d07-c7f6f1b45793",
"metadata": {},
"source": [
"And what our full chain produces:"
]
},
{
"cell_type": "code",
"execution_count": 32,
"id": "139cce01-ca75-452b-8de2-033ceec27158",
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"StructuredQuery(query='taxi driver', filter=Operation(operator=<Operator.AND: 'and'>, arguments=[Comparison(comparator=<Comparator.EQ: 'eq'>, attribute='genre', value='science fiction'), Operation(operator=<Operator.AND: 'and'>, arguments=[Comparison(comparator=<Comparator.GTE: 'gte'>, attribute='year', value=1990), Comparison(comparator=<Comparator.LT: 'lt'>, attribute='year', value=2000)]), Comparison(comparator=<Comparator.EQ: 'eq'>, attribute='director', value='Luc Besson')]), limit=None)"
]
},
"execution_count": 32,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"query_constructor.invoke({\"query\": \"What are some sci-fi movies from the 90's directed by Luc Besson about taxi drivers\"})"
]
},
{
"cell_type": "markdown",
"id": "9582a5fa-ffed-4d50-ad74-9b12d7d94b2a",
"metadata": {},
"source": [
"The query constructor is the key element of the self-query retriever. To make a great retrieval system you'll need to make sure your query constructor works well. Often this requires adjusting the prompt, the examples in the prompt, the attribute descriptions, etc. For an example that walks through refining a query constructor on some hotel inventory data, [check out this cookbook](https://github.com/langchain-ai/langchain/blob/master/cookbook/self_query_hotel_search.ipynb).\n",
"\n",
"The next key element is the structured query translator. This is the object responsible for translating the generic `StructuredQuery` object into a metadata filter in the syntax of the vector store you're using. LangChain comes with a number of built-in translators. To see them all head to the [Integrations section](/docs/integrations/retrievers/self_query)."
]
},
{
"cell_type": "code",
"execution_count": 33,
"id": "05f07ead-9aac-4079-9dde-784cb7aa1a8a",
"metadata": {},
"outputs": [],
"source": [
"from langchain.retrievers.self_query.chroma import ChromaTranslator\n",
"\n",
"retriever = SelfQueryRetriever(\n",
" query_constructor=query_constructor, \n",
" vectorstore=vectorstore, \n",
" structured_query_translator=ChromaTranslator(),\n",
")"
]
},
{
"cell_type": "code",
"execution_count": 34,
"id": "0ee155c9-7b02-4fe9-8de3-e37385c465af",
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"[Document(page_content='Toys come alive and have a blast doing so', metadata={'genre': 'animated', 'year': 1995})]"
]
},
"execution_count": 34,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"retriever.invoke(\"What's a movie after 1990 but before 2005 that's all about toys, and preferably is animated\")"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "poetry-venv",
"language": "python",
"name": "poetry-venv"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.9.1"
}
},
"nbformat": 4,
"nbformat_minor": 5
}

View File

@ -1,212 +0,0 @@
# Self-querying
:::info
Head to [Integrations](/docs/integrations/retrievers/self_query) for documentation on vector stores with built-in support for self-querying.
:::
A self-querying retriever is one that, as the name suggests, has the ability to query itself. Specifically, given any natural language query, the retriever uses a query-constructing LLM chain to write a structured query and then applies that structured query to its underlying VectorStore. This allows the retriever to not only use the user-input query for semantic similarity comparison with the contents of stored documents but to also extract filters from the user query on the metadata of stored documents and to execute those filters.
![](https://drive.google.com/uc?id=1OQUN-0MJcDUxmPXofgS7MqReEs720pqS)
## Get started
For demonstration purposes we'll use a `Pinecone` vector store.
First we'll want to create a `Pinecone` vector store and seed it with some data. We've created a small demo set of documents that contain summaries of movies.
To use Pinecone, you need to have `pinecone` package installed and you must have an API key and an environment. Here are the [installation instructions](https://docs.pinecone.io/docs/quickstart).
**Note:** The self-query retriever requires you to have `lark` package installed.
```python
# !pip install lark pinecone-client
```
```python
import os
import pinecone
pinecone.init(api_key=os.environ["PINECONE_API_KEY"], environment=os.environ["PINECONE_ENV"])
```
```python
from langchain.schema import Document
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.vectorstores import Pinecone
embeddings = OpenAIEmbeddings()
# create new index
pinecone.create_index("langchain-self-retriever-demo", dimension=1536)
```
```python
docs = [
Document(page_content="A bunch of scientists bring back dinosaurs and mayhem breaks loose", metadata={"year": 1993, "rating": 7.7, "genre": ["action", "science fiction"]}),
Document(page_content="Leo DiCaprio gets lost in a dream within a dream within a dream within a ...", metadata={"year": 2010, "director": "Christopher Nolan", "rating": 8.2}),
Document(page_content="A psychologist / detective gets lost in a series of dreams within dreams within dreams and Inception reused the idea", metadata={"year": 2006, "director": "Satoshi Kon", "rating": 8.6}),
Document(page_content="A bunch of normal-sized women are supremely wholesome and some men pine after them", metadata={"year": 2019, "director": "Greta Gerwig", "rating": 8.3}),
Document(page_content="Toys come alive and have a blast doing so", metadata={"year": 1995, "genre": "animated"}),
Document(page_content="Three men walk into the Zone, three men walk out of the Zone", metadata={"year": 1979, "rating": 9.9, "director": "Andrei Tarkovsky", "genre": ["science fiction", "thriller"], "rating": 9.9})
]
vectorstore = Pinecone.from_documents(
docs, embeddings, index_name="langchain-self-retriever-demo"
)
```
## Creating our self-querying retriever
Now we can instantiate our retriever. To do this we'll need to provide some information upfront about the metadata fields that our documents support and a short description of the document contents.
```python
from langchain.llms import OpenAI
from langchain.retrievers.self_query.base import SelfQueryRetriever
from langchain.chains.query_constructor.base import AttributeInfo
metadata_field_info=[
AttributeInfo(
name="genre",
description="The genre of the movie",
type="string or list[string]",
),
AttributeInfo(
name="year",
description="The year the movie was released",
type="integer",
),
AttributeInfo(
name="director",
description="The name of the movie director",
type="string",
),
AttributeInfo(
name="rating",
description="A 1-10 rating for the movie",
type="float"
),
]
document_content_description = "Brief summary of a movie"
llm = OpenAI(temperature=0)
retriever = SelfQueryRetriever.from_llm(llm, vectorstore, document_content_description, metadata_field_info, verbose=True)
```
## Testing it out
And now we can actually try using our retriever!
```python
# This example only specifies a relevant query
retriever.get_relevant_documents("What are some movies about dinosaurs")
```
<CodeOutputBlock lang="python">
```
query='dinosaur' filter=None
[Document(page_content='A bunch of scientists bring back dinosaurs and mayhem breaks loose', metadata={'genre': ['action', 'science fiction'], 'rating': 7.7, 'year': 1993.0}),
Document(page_content='Toys come alive and have a blast doing so', metadata={'genre': 'animated', 'year': 1995.0}),
Document(page_content='A psychologist / detective gets lost in a series of dreams within dreams within dreams and Inception reused the idea', metadata={'director': 'Satoshi Kon', 'rating': 8.6, 'year': 2006.0}),
Document(page_content='Leo DiCaprio gets lost in a dream within a dream within a dream within a ...', metadata={'director': 'Christopher Nolan', 'rating': 8.2, 'year': 2010.0})]
```
</CodeOutputBlock>
```python
# This example only specifies a filter
retriever.get_relevant_documents("I want to watch a movie rated higher than 8.5")
```
<CodeOutputBlock lang="python">
```
query=' ' filter=Comparison(comparator=<Comparator.GT: 'gt'>, attribute='rating', value=8.5)
[Document(page_content='A psychologist / detective gets lost in a series of dreams within dreams within dreams and Inception reused the idea', metadata={'director': 'Satoshi Kon', 'rating': 8.6, 'year': 2006.0}),
Document(page_content='Three men walk into the Zone, three men walk out of the Zone', metadata={'director': 'Andrei Tarkovsky', 'genre': ['science fiction', 'thriller'], 'rating': 9.9, 'year': 1979.0})]
```
</CodeOutputBlock>
```python
# This example specifies a query and a filter
retriever.get_relevant_documents("Has Greta Gerwig directed any movies about women")
```
<CodeOutputBlock lang="python">
```
query='women' filter=Comparison(comparator=<Comparator.EQ: 'eq'>, attribute='director', value='Greta Gerwig')
[Document(page_content='A bunch of normal-sized women are supremely wholesome and some men pine after them', metadata={'director': 'Greta Gerwig', 'rating': 8.3, 'year': 2019.0})]
```
</CodeOutputBlock>
```python
# This example specifies a composite filter
retriever.get_relevant_documents("What's a highly rated (above 8.5) science fiction film?")
```
<CodeOutputBlock lang="python">
```
query=' ' filter=Operation(operator=<Operator.AND: 'and'>, arguments=[Comparison(comparator=<Comparator.EQ: 'eq'>, attribute='genre', value='science fiction'), Comparison(comparator=<Comparator.GT: 'gt'>, attribute='rating', value=8.5)])
[Document(page_content='Three men walk into the Zone, three men walk out of the Zone', metadata={'director': 'Andrei Tarkovsky', 'genre': ['science fiction', 'thriller'], 'rating': 9.9, 'year': 1979.0})]
```
</CodeOutputBlock>
```python
# This example specifies a query and composite filter
retriever.get_relevant_documents("What's a movie after 1990 but before 2005 that's all about toys, and preferably is animated")
```
<CodeOutputBlock lang="python">
```
query='toys' filter=Operation(operator=<Operator.AND: 'and'>, arguments=[Comparison(comparator=<Comparator.GT: 'gt'>, attribute='year', value=1990.0), Comparison(comparator=<Comparator.LT: 'lt'>, attribute='year', value=2005.0), Comparison(comparator=<Comparator.EQ: 'eq'>, attribute='genre', value='animated')])
[Document(page_content='Toys come alive and have a blast doing so', metadata={'genre': 'animated', 'year': 1995.0})]
```
</CodeOutputBlock>
## Filter k
We can also use the self query retriever to specify `k`: the number of documents to fetch.
We can do this by passing `enable_limit=True` to the constructor.
```python
retriever = SelfQueryRetriever.from_llm(
llm,
vectorstore,
document_content_description,
metadata_field_info,
enable_limit=True,
verbose=True
)
```
```python
# This example only specifies a relevant query
retriever.get_relevant_documents("What are two movies about dinosaurs")
```