Retrieval augmented generation (RAG) - quickstart

Setup

First, install the Cohere Python SDK with the following command.

$
pip install -U cohere

Next, import the library and create a client.

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

Documents

First, define the documents that will passed as the context for RAG. These documents are typically retrieved from sources such as vector databases via semantic search, or any system that can retrieve unstructured data given a user query.

Each document is a data object that can take any number of fields e.g. title, url, text, etc.

1
documents = [
2
    {
3
        "data": {
4
            "text": "Reimbursing Travel Expenses: Easily manage your travel expenses by submitting them through our finance tool. Approvals are prompt and straightforward."
5
        }
6
    },
7
    {
8
        "data": {
9
            "text": "Working from Abroad: Working remotely from another country is possible. Simply coordinate with your manager and ensure your availability during core hours."
10
        }
11
    },
12
    {
13
        "data": {
14
            "text": "Health and Wellness Benefits: We care about your well-being and offer gym memberships, on-site yoga classes, and comprehensive health insurance."
15
        }
16
    },
17
]

Response Generation

Next, call the Chat API by passing the documents in the documents parameter. This tells the model to run in RAG-mode and use these documents as the context in its response.

1
# Add the user query
2
query = "Are there health benefits?"
3
4
# Generate the response
5
response = co.chat(
6
    model="command-r-plus-08-2024",
7
    messages=[{"role": "user", "content": query}],
8
    documents=documents,
9
)
10
11
# Display the response
12
print(response.message.content[0].text)

1
Yes, there are health benefits. We offer gym memberships, on-site yoga classes, and comprehensive health insurance.

Citation Generation

The response object contains a citations field, which contains specific text spans from the documents on which the response is grounded.

1
if response.message.citations:
2
    for citation in response.message.citations:
3
        print(citation, "\n")

1
start=14 end=88 text='gym memberships, on-site yoga classes, and comprehensive health insurance.' document_ids=['doc_1']