Generating Multi-Faceted Queries

Consider a RAG system that needs to search through a large database of code examples and tutorials. A user might ask for “Python examples using the chat endpoint” or “JavaScript tutorials for text summarization”.

In a basic RAG setup, these queries would be passed as-is to a search function, potentially missing important context or failing to leverage the structured nature of the data. For example, the code examples database might consist of metadata such as the programming language, the created time, the tech stack used, and so on.

It would be great if we could design a system that could leverage this metadata as a filter to retrieve only the relevant results.

We can achieve this using a tool use approach. Here, we can build a system that generates multi-faceted queries to capture the full intent of a user’s request. This allows for more precise and relevant results by utilizing the semi-structured nature of the data.

Here are some examples of how this approach can be applied:

1. E-commerce product searches: Filtering by price range, category, brand, customer ratings, and availability.
2. Academic research databases: Narrowing results by publication year, field of study, citation count, and peer-review status.
3. Job search platforms: Refining job listings by location, experience level, salary range, and required skills.

In this tutorial, we’ll cover:

• Defining the function for data querying
• Creating the tool for generating multi-faceted queries
• Building an agent for performing multi-faceted queries
• Running the agent

We’ll build an agent that helps developers find relevant code examples and tutorials for using Cohere.

Setup

To get started, first we need to install the cohere library and create a Cohere client.

1
import json
2
import os
3
import cohere
4
5
co = cohere.ClientV2("COHERE_API_KEY") # Get your free API key: https://dashboard.cohere.com/api-keys

Important: the source code for tool definitions can be found here. Make sure to have the tool_def.py file in the same directory as this notebook for the imports to work correctly.

Defining the function for data querying

We’ll remove the other tools from Part 1 and just use one – search_code_examples.

Now, instead of just the query parameter, we’ll add two more parameters: programming_language and endpoints:

• programming_language: The programming language of the code example or tutorial.
• endpoints: The Cohere endpoints used in the code example or tutorial.

We’ll use these parameters as the metadata to filter the code examples and tutorials.

Let’s rename the function to search_code_examples_detailed to reflect this change.

And as in Part 1, for simplicity, we create query as just a mock parameter and no actual search logic will be performed based on it.

1
from tool_def import (
2
    search_code_examples_detailed,
3
    search_code_examples_detailed_tool,
4
)

1
functions_map = {
2
    "search_code_examples_detailed": search_code_examples_detailed,
3
}

Creating the tool for generating multi-faceted queries

With the search_code_examples modified, we now need to modify the tool definition as well. Here, we are adding two new properties to the tool definition:

• programming_language: This is a string property which we provide a list of options for the model to choose. We do this by adding “Possible enum values” to the description, which in our case is py, js.
• endpoints: We want the model to be able to choose from more than one endpoint, and so here we define an array property. When defining an array property, we need to specify the type of the items in the array using the items key, which in our case is string. We also provide a list of endpoint options for the model to choose from, which is chat, embed, rerank, classify.

We set only the query parameter as required, while the other two parameters are optional.

1
search_code_examples_detailed_tool = {
2
    "type": "function",
3
    "function": {
4
        "name": "search_code_examples_detailed",
5
        "description": "Searches code examples or tutorials of using Cohere.",
6
        "parameters": {
7
            "type": "object",
8
            "properties": {
9
                "query": {
10
                    "type": "string",
11
                    "description": "The search query."
12
                },
13
                "programming_language": {
14
                    "type": "string",
15
                    "description": "The programming language of the code example or tutorial. Only use this property when specified by the user. Possible enum values: py, js.",
16
                },
17
                "endpoints": {
18
                    "type": "array",
19
                    "items": {"type": "string"},
20
                    "description": "The Cohere endpoints used in the code example or tutorial. Only use this property when asked by the user. Possible enum values: chat, embed, rerank, classify.",
21
                }
22
            },
23
            "required": ["query"]
24
        }
25
    }
26
}

1
tools = [search_code_examples_detailed_tool]

Building an agent for performing multi-faceted queries

Next, let’s create a run_agent function to run the agentic RAG workflow, the same as in Part 1.

The only change we are making here is to make the system message simpler and more specific since the agent now only has one tool.

1
system_message="""## Task and Context
2
You are an assistant who helps developers find code examples and tutorials on using Cohere."""

1
model = "command-r-plus-08-2024"
2
3
def run_agent(query, messages=None):
4
    if messages is None:
5
        messages = []
6
        
7
    if "system" not in {m.get("role") for m in messages}:
8
        messages.append({"role": "system", "content": system_message})
9
    
10
    # Step 1: get user message
11
    print(f"QUESTION:\n{query}")
12
    print("="*50)
13
    
14
    messages.append({"role": "user", "content": query})
15
16
    # Step 2: Generate tool calls (if any)
17
    response = co.chat(
18
        model=model,
19
        messages=messages,
20
        tools=tools,
21
        temperature=0.1
22
    )
23
24
    while response.message.tool_calls:
25
        
26
        print("TOOL PLAN:")
27
        print(response.message.tool_plan,"\n")
28
        print("TOOL CALLS:")
29
        for tc in response.message.tool_calls:
30
            print(f"Tool name: {tc.function.name} | Parameters: {tc.function.arguments}")
31
        print("="*50)
32
33
        messages.append({"role": "assistant", "tool_calls": response.message.tool_calls, "tool_plan": response.message.tool_plan})        
34
        
35
        # Step 3: Get tool results
36
        for tc in response.message.tool_calls:
37
            tool_result = functions_map[tc.function.name](**json.loads(tc.function.arguments))
38
            tool_content = []
39
            for data in tool_result:
40
                tool_content.append({"type": "document", "document": {"data": json.dumps(data)}})
41
                # Optional: add an "id" field in the "document" object, otherwise IDs are auto-generated
42
            messages.append({"role": "tool", "tool_call_id": tc.id, "content": tool_content})
43
        
44
        # Step 4: Generate response and citations 
45
        response = co.chat(
46
            model=model,
47
            messages=messages,
48
            tools=tools,
49
            temperature=0.1
50
        )
51
    
52
    messages.append({"role": "assistant", "content": response.message.content[0].text})
53
        
54
    # Print final response
55
    print("RESPONSE:")
56
    print(response.message.content[0].text)
57
    print("="*50)
58
    
59
    # Print citations (if any)
60
    verbose_source = False # Change to True to display the contents of a source
61
    if response.message.citations:
62
        print("CITATIONS:\n")
63
        for citation in response.message.citations:
64
            print(f"Start: {citation.start}| End:{citation.end}| Text:'{citation.text}' ")
65
            print("Sources:")
66
            for idx, source in enumerate(citation.sources):
67
                print(f"{idx+1}. {source.id}")
68
                if verbose_source:
69
                    print(f"{source.tool_output}")
70
            print("\n")
71
                    
72
    return messages

Running the agent

Let’s start with a broad query about “RAG code examples”.

Since it’s broad, this query shouldn’t require any metadata filtering.

And this is shown by the agent’s response, which provides only one parameter, query, in its tool call.

1
messages = run_agent("Do you have any RAG code examples")
2
# Tool name: search_code_examples | Parameters: {"query":"RAG code examples"}

QUESTION:
Do you have any RAG code examples
==================================================
TOOL PLAN:
I will search for RAG code examples. 
TOOL CALLS:
Tool name: search_code_examples_detailed | Parameters: {"query":"RAG"}
==================================================
RESPONSE:
Yes, we have a code example for RAG with Chat Embed and Rerank via Pinecone.
==================================================
CITATIONS:
Start: 32| End:76| Text:'RAG with Chat Embed and Rerank via Pinecone.' 
Sources:
1. search_code_examples_detailed_bqm9ww9cp4nw:2

Let’s try a more specific query about “javascript tutorials on text summarization”.

This time, the agent uses the programming_language parameter and passed the value js to it.

1
messages = run_agent("Javascript tutorials on summarization")
2
# Tool name: search_code_examples | Parameters: {"programming_language":"js","query":"..."}

QUESTION:
Javascript tutorials on summarization
==================================================
TOOL PLAN:
I will search for 'summarization' in the 'js' programming language. 
TOOL CALLS:
Tool name: search_code_examples_detailed | Parameters: {"programming_language":"js","query":"summarization"}
==================================================
RESPONSE:
Here is a tutorial on building a Chrome extension to summarize web pages using JavaScript.
==================================================
CITATIONS:
Start: 22| End:72| Text:'building a Chrome extension to summarize web pages' 
Sources:
1. search_code_examples_detailed_hf9rv18987vp:0

Let’s now try a query that involves filtering based on the endpoints. Here, the user asks for “code examples of using embed and rerank endpoints”.

And since we have set up the endpoints parameter to be an array, the agent is able to call the tool with a list of endpoints as its argument.

1
messages = run_agent("Code examples of using embed and rerank endpoints.")
2
3
# Tool name: search_code_examples | Parameters: {"endpoints":["embed","rerank"],"query":"..."}

QUESTION:
Code examples of using embed and rerank endpoints.
==================================================
TOOL PLAN:
I will search for code examples of using the embed and rerank endpoints. 
TOOL CALLS:
Tool name: search_code_examples_detailed | Parameters: {"endpoints":["embed","rerank"],"query":"embed and rerank"}
==================================================
RESPONSE:
Here are some code examples of using the embed and rerank endpoints:
- Wikipedia Semantic Search with Cohere Embedding Archives
- RAG With Chat Embed and Rerank via Pinecone
- Build Chatbots That Know Your Business with MongoDB and Cohere
==================================================
CITATIONS:
Start: 71| End:127| Text:'Wikipedia Semantic Search with Cohere Embedding Archives' 
Sources:
1. search_code_examples_detailed_d3br7bwd05bp:0
Start: 130| End:173| Text:'RAG With Chat Embed and Rerank via Pinecone' 
Sources:
1. search_code_examples_detailed_d3br7bwd05bp:1
Start: 176| End:238| Text:'Build Chatbots That Know Your Business with MongoDB and Cohere' 
Sources:
1. search_code_examples_detailed_d3br7bwd05bp:2

Finally, let’s try a query that involves filtering based on both the programming language and the endpoints. Here, the user asks for “Python examples of using the chat endpoint”.

And the agent correctly uses both parameters to query the code examples.

1
messages = run_agent("Python examples of using the chat endpoint.")
2
3
#Tool name: search_code_examples | Parameters: {"endpoints":["chat"],"programming_language":"py","query":"..."}

QUESTION:
Python examples of using the chat endpoint.
==================================================
TOOL PLAN:
I will search for Python examples of using the chat endpoint. 
TOOL CALLS:
Tool name: search_code_examples_detailed | Parameters: {"endpoints":["chat"],"programming_language":"py","query":"Python examples of using the chat endpoint"}
==================================================
RESPONSE:
Here are some Python examples of using the chat endpoint:
- Calendar Agent with Native Multi Step Tool
- RAG With Chat Embed and Rerank via Pinecone
- Build Chatbots That Know Your Business with MongoDB and Cohere
==================================================
CITATIONS:
Start: 60| End:102| Text:'Calendar Agent with Native Multi Step Tool' 
Sources:
1. search_code_examples_detailed_kp25g4k49ej5:0
Start: 105| End:148| Text:'RAG With Chat Embed and Rerank via Pinecone' 
Sources:
1. search_code_examples_detailed_kp25g4k49ej5:2
Start: 151| End:213| Text:'Build Chatbots That Know Your Business with MongoDB and Cohere' 
Sources:
1. search_code_examples_detailed_kp25g4k49ej5:3

Summary

In this tutorial, we learned about:

• How to define the function for data querying
• How to create the tool for generating multi-faceted queries
• How to build an agent for performing multi-faceted queries
• How to run the agent

By implementing multi-faceted queries over semi-structured data, we’ve enhanced our RAG system to handle more specific and targeted searches. This approach allows for better utilization of metadata and more precise filtering of results, which is particularly useful when dealing with large collections of code examples and tutorials.

While this tutorial demonstrates how to work with semi-structured data, the agentic RAG approach can be applied to structured data as well. That means we can build agents that can translate natural language queries into queries for tables or relational databases.

In Part 5, we’ll learn how to perform RAG over structured data (tables).