Skip to content

Commit

Permalink
AutoRetrievel Filter using OpenAI agent and multi-filters (run-llama#…
Browse files Browse the repository at this point in the history
…9281)

* AutoRetrievel Filter using OpenAI agent and multi-filters

* cr
  • Loading branch information
hatianzhang authored Dec 4, 2023
1 parent 067de92 commit 91647d1
Showing 1 changed file with 196 additions and 29 deletions.
225 changes: 196 additions & 29 deletions docs/examples/agent/openai_agent_query_cookbook.ipynb
Original file line number Diff line number Diff line change
Expand Up @@ -67,7 +67,23 @@
"import os\n",
"\n",
"api_key = os.environ[\"PINECONE_API_KEY\"]\n",
"pinecone.init(api_key=api_key, environment=\"us-west1-gcp\")"
"pinecone.init(api_key=api_key, environment=\"us-west4-gcp-free\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "6639b685",
"metadata": {},
"outputs": [],
"source": [
"import os\n",
"import getpass\n",
"\n",
"# os.environ[\"OPENAI_API_KEY\"] = getpass.getpass(\"OpenAI API Key:\")\n",
"import openai\n",
"\n",
"openai.api_key = \"sk-<your-key>\""
]
},
{
Expand All @@ -80,7 +96,7 @@
"# dimensions are for text-embedding-ada-002\n",
"try:\n",
" pinecone.create_index(\n",
" \"quickstart\", dimension=1536, metric=\"euclidean\", pod_type=\"p1\"\n",
" \"quickstart-index\", dimension=1536, metric=\"euclidean\", pod_type=\"p1\"\n",
" )\n",
"except Exception:\n",
" # most likely index already exists\n",
Expand All @@ -94,7 +110,7 @@
"metadata": {},
"outputs": [],
"source": [
"pinecone_index = pinecone.Index(\"quickstart\")"
"pinecone_index = pinecone.Index(\"quickstart-index\")"
]
},
{
Expand Down Expand Up @@ -149,6 +165,8 @@
" metadata={\n",
" \"category\": \"Sports\",\n",
" \"country\": \"United States\",\n",
" \"gender\": \"male\",\n",
" \"born\": 1963,\n",
" },\n",
" ),\n",
" TextNode(\n",
Expand All @@ -160,6 +178,8 @@
" metadata={\n",
" \"category\": \"Entertainment\",\n",
" \"country\": \"United States\",\n",
" \"gender\": \"female\",\n",
" \"born\": 1975,\n",
" },\n",
" ),\n",
" TextNode(\n",
Expand All @@ -171,6 +191,8 @@
" metadata={\n",
" \"category\": \"Business\",\n",
" \"country\": \"United States\",\n",
" \"gender\": \"male\",\n",
" \"born\": 1971,\n",
" },\n",
" ),\n",
" TextNode(\n",
Expand All @@ -182,6 +204,8 @@
" metadata={\n",
" \"category\": \"Music\",\n",
" \"country\": \"Barbados\",\n",
" \"gender\": \"female\",\n",
" \"born\": 1988,\n",
" },\n",
" ),\n",
" TextNode(\n",
Expand All @@ -194,6 +218,8 @@
" metadata={\n",
" \"category\": \"Sports\",\n",
" \"country\": \"Portugal\",\n",
" \"gender\": \"male\",\n",
" \"born\": 1985,\n",
" },\n",
" ),\n",
"]"
Expand Down Expand Up @@ -222,7 +248,7 @@
"name": "stderr",
"output_type": "stream",
"text": [
"Upserted vectors: 100%|███████████████████████████████████████████████████| 5/5 [00:00<00:00, 9.61it/s]\n"
"Upserted vectors: 100%|██████████| 5/5 [00:00<00:00, 5.79it/s]\n"
]
}
],
Expand Down Expand Up @@ -256,8 +282,10 @@
"from llama_index.vector_stores.types import (\n",
" VectorStoreInfo,\n",
" MetadataInfo,\n",
" ExactMatchFilter,\n",
" MetadataFilter,\n",
" MetadataFilters,\n",
" FilterCondition,\n",
" FilterOperator,\n",
")\n",
"from llama_index.retrievers import VectorIndexRetriever\n",
"from llama_index.query_engine import RetrieverQueryEngine\n",
Expand Down Expand Up @@ -288,27 +316,80 @@
" \" Portugal]\"\n",
" ),\n",
" ),\n",
" MetadataInfo(\n",
" name=\"gender\",\n",
" type=\"str\",\n",
" description=(\"Gender of the celebrity, one of [male, female]\"),\n",
" ),\n",
" MetadataInfo(\n",
" name=\"born\",\n",
" type=\"int\",\n",
" description=(\"Born year of the celebrity, could be any integer\"),\n",
" ),\n",
" ],\n",
")\n",
"\n",
"\n",
")"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "0594ca0f",
"metadata": {},
"outputs": [],
"source": [
"# define pydantic model for auto-retrieval function\n",
"class AutoRetrieveModel(BaseModel):\n",
" query: str = Field(..., description=\"natural language query string\")\n",
" filter_key_list: List[str] = Field(\n",
" ..., description=\"List of metadata filter field names\"\n",
" )\n",
" filter_value_list: List[str] = Field(\n",
" filter_value_list: List[Any] = Field(\n",
" ...,\n",
" description=(\n",
" \"List of metadata filter field values (corresponding to names\"\n",
" \" specified in filter_key_list)\"\n",
" ),\n",
" )\n",
" filter_operator_list: List[str] = Field(\n",
" ...,\n",
" description=(\n",
" \"Metadata filters conditions (could be one of <, <=, >, >=, ==, !=)\"\n",
" ),\n",
" )\n",
" filter_condition: str = Field(\n",
" ...,\n",
" description=(\"Metadata filters condition values (could be AND or OR)\"),\n",
" )\n",
"\n",
"\n",
"description = f\"\"\"\\\n",
"Use this tool to look up biographical information about celebrities.\n",
"The vector database schema is given below:\n",
"{vector_store_info.json()}\n",
"\"\"\""
]
},
{
"cell_type": "markdown",
"id": "a759ea66",
"metadata": {},
"source": [
"Define AutoRetrieve Functions"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "a8951dd4",
"metadata": {},
"outputs": [],
"source": [
"def auto_retrieve_fn(\n",
" query: str, filter_key_list: List[str], filter_value_list: List[str]\n",
" query: str,\n",
" filter_key_list: List[str],\n",
" filter_value_list: List[any],\n",
" filter_operator_list: List[str],\n",
" filter_condition: str,\n",
"):\n",
" \"\"\"Auto retrieval function.\n",
"\n",
Expand All @@ -317,13 +398,17 @@
" \"\"\"\n",
" query = query or \"Query\"\n",
"\n",
" exact_match_filters = [\n",
" ExactMatchFilter(key=k, value=v)\n",
" for k, v in zip(filter_key_list, filter_value_list)\n",
" metadata_filters = [\n",
" MetadataFilter(key=k, value=v, operator=op)\n",
" for k, v, op in zip(\n",
" filter_key_list, filter_value_list, filter_operator_list\n",
" )\n",
" ]\n",
" retriever = VectorIndexRetriever(\n",
" index,\n",
" filters=MetadataFilters(filters=exact_match_filters),\n",
" filters=MetadataFilters(\n",
" filters=metadata_filters, condition=filter_condition\n",
" ),\n",
" top_k=top_k,\n",
" )\n",
" query_engine = RetrieverQueryEngine.from_args(retriever)\n",
Expand All @@ -332,12 +417,6 @@
" return str(response)\n",
"\n",
"\n",
"description = f\"\"\"\\\n",
"Use this tool to look up biographical information about celebrities.\n",
"The vector database schema is given below:\n",
"{vector_store_info.json()}\n",
"\"\"\"\n",
"\n",
"auto_retrieve_tool = FunctionTool.from_defaults(\n",
" fn=auto_retrieve_fn,\n",
" name=\"celebrity_bios\",\n",
Expand Down Expand Up @@ -382,18 +461,28 @@
"name": "stdout",
"output_type": "stream",
"text": [
"STARTING TURN 1\n",
"---------------\n",
"\n",
"=== Calling Function ===\n",
"Calling function: celebrity_bios with args: {\n",
" \"query\": \"celebrities\",\n",
" \"filter_key_list\": [\"country\"],\n",
" \"filter_value_list\": [\"United States\"]\n",
"\"query\": \"celebrities from the United States\",\n",
"\"filter_key_list\": [\"country\"],\n",
"\"filter_value_list\": [\"United States\"],\n",
"\"filter_operator_list\": [\"==\"],\n",
"\"filter_condition\": \"and\"\n",
"}\n",
"Got output: \n",
"Celebrities in the United States who are associated with Entertainment and Sports include Angelina Jolie and Michael Jordan.\n",
"Got output: Angelina Jolie and Michael Jordan are both celebrities from the United States.\n",
"========================\n",
"Angelina Jolie is an American actress, filmmaker, and humanitarian. She has received an Academy Award and three Golden Globe Awards, and has been cited as Hollywood's highest-paid actress. Jolie made her screen debut as a child alongside her father, Jon Voight, in Lookin' to Get Out (1982), and her film career began in earnest a decade later with the low-budget production Cyborg 2 (1993), followed by her first leading role in a major film, Hackers (1995).\n",
"\n",
"Michael Jordan is a retired professional basketball player from the United States. He is widely regarded as one of the greatest basketball players in history. Jordan was one of the most effectively marketed athletes of his generation and was instrumental in popularizing the NBA around the world in the 1980s and 1990s. He played 15 seasons in the NBA, winning six championships with the Chicago Bulls.\n"
"STARTING TURN 2\n",
"---------------\n",
"\n",
"Here are two celebrities from the United States:\n",
"\n",
"1. **Angelina Jolie**: She is an American actress, filmmaker, and humanitarian. The recipient of numerous accolities, including an Academy Award and three Golden Globe Awards, she has been named Hollywood's highest-paid actress multiple times.\n",
"\n",
"2. **Michael Jordan**: He is a former professional basketball player and the principal owner of the Charlotte Hornets of the National Basketball Association (NBA). He played 15 seasons in the NBA, winning six championships with the Chicago Bulls. He is considered one of the greatest players in the history of the NBA.\n"
]
}
],
Expand All @@ -402,6 +491,84 @@
"print(str(response))"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "b389d555",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"STARTING TURN 1\n",
"---------------\n",
"\n",
"=== Calling Function ===\n",
"Calling function: celebrity_bios with args: {\n",
"\"query\": \"celebrities born after 1980\",\n",
"\"filter_key_list\": [\"born\"],\n",
"\"filter_value_list\": [1980],\n",
"\"filter_operator_list\": [\">\"],\n",
"\"filter_condition\": \"and\"\n",
"}\n",
"Got output: Rihanna and Cristiano Ronaldo are both celebrities who were born after 1980.\n",
"========================\n",
"\n",
"STARTING TURN 2\n",
"---------------\n",
"\n",
"Here are two celebrities who were born after 1980:\n",
"\n",
"1. **Rihanna**: She is a Barbadian singer, actress, and businesswoman. Born in Saint Michael and raised in Bridgetown, Barbados, Rihanna was discovered by American record producer Evan Rogers who invited her to the United States to record demo tapes. She rose to fame with her debut album \"Music of the Sun\" and its follow-up \"A Girl like Me\".\n",
"\n",
"2. **Cristiano Ronaldo**: He is a Portuguese professional footballer who plays as a forward for Serie A club Juventus and captains the Portugal\n"
]
}
],
"source": [
"response = agent.chat(\"Tell me about two celebrities born after 1980. \")\n",
"print(str(response))"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "4558faa2",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"STARTING TURN 1\n",
"---------------\n",
"\n",
"=== Calling Function ===\n",
"Calling function: celebrity_bios with args: {\n",
"\"query\": \"business celebrities born after 1950\",\n",
"\"filter_key_list\": [\"category\", \"born\"],\n",
"\"filter_value_list\": [\"Business\", 1950],\n",
"\"filter_operator_list\": [\"==\", \">\"],\n",
"\"filter_condition\": \"and\"\n",
"}\n",
"Got output: Elon Musk is a notable business celebrity who was born in 1971.\n",
"========================\n",
"\n",
"STARTING TURN 2\n",
"---------------\n",
"\n",
"Elon Musk is a business celebrity who was born after 1950. He is a business magnate and investor. He is the founder, CEO, CTO, and chief designer of SpaceX; early investor, CEO and product architect of Tesla, Inc.; founder of The\n"
]
}
],
"source": [
"response = agent.chat(\n",
" \"Tell me about few celebrities under category business and born after 1950. \"\n",
")\n",
"print(str(response))"
]
},
{
"attachments": {},
"cell_type": "markdown",
Expand Down Expand Up @@ -954,9 +1121,9 @@
],
"metadata": {
"kernelspec": {
"display_name": "llama_index_v2",
"display_name": "venv",
"language": "python",
"name": "llama_index_v2"
"name": "python3"
},
"language_info": {
"codemirror_mode": {
Expand Down

0 comments on commit 91647d1

Please sign in to comment.