End-to-end example of RAG with Chat, Embed, and Rerank

1
message = "How to get to know my teammates"
2
3
# Define the query generation tool
4
query_gen_tool = [
5
    {
6
        "type": "function",
7
        "function": {
8
            "name": "internet_search",
9
            "description": "Returns a list of relevant document snippets for a textual query retrieved from the internet",
10
            "parameters": {
11
                "type": "object",
12
                "properties": {
13
                    "queries": {
14
                        "type": "array",
15
                        "items": {"type": "string"},
16
                        "description": "a list of queries to search the internet with.",
17
                    }
18
                },
19
                "required": ["queries"],
20
            },
21
        },
22
    }
23
]
24
25
# Define a system message to optimize search query generation
26
instructions = "Write a search query that will find helpful information for answering the user's question accurately. If you need more than one search query, write a list of search queries. If you decide that a search is very unlikely to find information that would be useful in constructing a response to the user, you should instead directly answer."
27
28
# Generate search queries (if any)
29
search_queries = []
30
31
res = co.chat(
32
    model="command-a-03-2025",
33
    messages=[
34
        {"role": "system", "content": instructions},
35
        {"role": "user", "content": message},
36
    ],
37
    tools=query_gen_tool,
38
)
39
40
if res.message.tool_calls:
41
    for tc in res.message.tool_calls:
42
        queries = json.loads(tc.function.arguments)["queries"]
43
        search_queries.extend(queries)
44
45
print(search_queries)

1
# Define the documents
2
documents = [
3
    {
4
        "data": {
5
            "text": "Joining Slack Channels: You will receive an invite via email. Be sure to join relevant channels to stay informed and engaged."
6
        }
7
    },
8
    {
9
        "data": {
10
            "text": "Finding Coffee Spots: For your caffeine fix, head to the break room's coffee machine or cross the street to the café for artisan coffee."
11
        }
12
    },
13
    {
14
        "data": {
15
            "text": "Team-Building Activities: We foster team spirit with monthly outings and weekly game nights. Feel free to suggest new activity ideas anytime!"
16
        }
17
    },
18
    {
19
        "data": {
20
            "text": "Working Hours Flexibility: We prioritize work-life balance. While our core hours are 9 AM to 5 PM, we offer flexibility to adjust as needed."
21
        }
22
    },
23
    {
24
        "data": {
25
            "text": "Side Projects Policy: We encourage you to pursue your passions. Just be mindful of any potential conflicts of interest with our business."
26
        }
27
    },
28
    {
29
        "data": {
30
            "text": "Reimbursing Travel Expenses: Easily manage your travel expenses by submitting them through our finance tool. Approvals are prompt and straightforward."
31
        }
32
    },
33
    {
34
        "data": {
35
            "text": "Working from Abroad: Working remotely from another country is possible. Simply coordinate with your manager and ensure your availability during core hours."
36
        }
37
    },
38
    {
39
        "data": {
40
            "text": "Health and Wellness Benefits: We care about your well-being and offer gym memberships, on-site yoga classes, and comprehensive health insurance."
41
        }
42
    },
43
    {
44
        "data": {
45
            "text": "Performance Reviews Frequency: We conduct informal check-ins every quarter and formal performance reviews twice a year."
46
        }
47
    },
48
    {
49
        "data": {
50
            "text": "Proposing New Ideas: Innovation is welcomed! Share your brilliant ideas at our weekly team meetings or directly with your team lead."
51
        }
52
    },
53
]
54
55
# Embed the documents
56
doc_emb = co.embed(
57
    model="embed-english-v3.0",
58
    input_type="search_document",
59
    texts=[doc["data"]["text"] for doc in documents],
60
    embedding_types=["float"],
61
).embeddings.float

1
# Embed the search query
2
query_emb = co.embed(
3
    model="embed-english-v3.0",
4
    input_type="search_query",
5
    texts=search_queries,
6
    embedding_types=["float"],
7
).embeddings.float

1
# Compute dot product similarity and display results
2
n = 5
3
scores = np.dot(query_emb, np.transpose(doc_emb))[0]
4
max_idx = np.argsort(-scores)[:n]
5
6
retrieved_documents = [documents[item] for item in max_idx]
7
8
for rank, idx in enumerate(max_idx):
9
    print(f"Rank: {rank+1}")
10
    print(f"Score: {scores[idx]}")
11
    print(f"Document: {retrieved_documents[rank]}\n")

1
# Rerank the documents
2
results = co.rerank(
3
    model="rerank-v3.5",
4
    query=search_queries[0],
5
    documents=[doc["data"]["text"] for doc in retrieved_documents],
6
    top_n=2,
7
)
8
9
# Display the reranking results
10
for idx, result in enumerate(results.results):
11
    print(f"Rank: {idx+1}")
12
    print(f"Score: {result.relevance_score}")
13
    print(f"Document: {retrieved_documents[result.index]}\n")
14
15
reranked_documents = [
16
    retrieved_documents[result.index] for result in results.results
17
]

1
messages = [{"role": "user", "content": message}]
2
3
# Generate the response
4
response = co.chat(
5
    model="command-a-03-2025",
6
    messages=messages,
7
    documents=reranked_documents,
8
)
9
10
# Display the response
11
print(response.message.content[0].text)
12
13
# Display the citations and source documents
14
if response.message.citations:
15
    print("\nCITATIONS:")
16
    for citation in response.message.citations:
17
        print(citation, "\n")