Migrating Monolithic Prompts to Command-R with RAG

Back to CookbooksOpen in GitHub

Command-R is a powerful LLM optimized for long context tasks such as retrieval augmented generation (RAG). Migrating a monolithic task such as question-answering or query-focused summarization to RAG can improve the quality of responses due to reduced hallucination and improved conciseness through grounding.

Previously, migrating an existing use case to RAG involved a lot of manual work around indexing documents, implementing at least a basic search strategy, extensive post-processing to introduce proper grounding through citations, and of course fine-tuning an LLM to work well in the RAG paradigm.

This cookbook demonstrates automatic migration of monolithic prompts through two diverse use cases where an original prompt is broken down into two parts: (1) context; and (2) instructions. The former can be done automatically or through simple chunking, while the latter is done automatically by Command-R through single shot prompt optimization.

The two use cases demonstrated here are:

  1. Autobiography Assistant; and
  2. Legal Question Answering
PYTHON
1#!pip install cohere
PYTHON
1import json
2import os
3import re
4
5import cohere
6import getpass
PYTHON
1CO_API_KEY = getpass.getpass('cohere API key:')
Output
cohere API key:··········
PYTHON
1co = cohere.Client(CO_API_KEY)

Autobiography Assistant

This application scenario is a common LLM-as-assistant use case: given some context, help the user to complete a task. In this case, the task is to write a concise autobiographical summary.

PYTHON
1original_prompt = '''## information
2Current Job Title: Senior Software Engineer
3Current Company Name: GlobalSolTech
4Work Experience: Over 15 years of experience in software engineering, specializing in AI and machine learning. Proficient in Python, C++, and Java, with expertise in developing algorithms for natural language processing, computer vision, and recommendation systems.
5Current Department Name: AI Research and Development
6Education: B.Sc. in Physics from Trent University (2004), Ph.D. in Statistics from HEC in Paris (2010)
7Hobbies: I love hiking in the mountains, free diving, and collecting and restoring vintage world war one mechanical watches.
8Family: Married with 4 children and 3 grandchildren.
9
10## instructions
11Your task is to assist a user in writing a short biography for social media.
12The length of the text should be no more than 100 words.
13Write the summary in first person.'''
PYTHON
1response = co.chat(
2    message=original_prompt,
3    model='command-r',
4)
PYTHON
1print(response.text)
    I'm a Senior Software Engineer at GlobalSolTech, with over 15 years of experience in AI and machine learning. My expertise lies in developing innovative algorithms for natural language processing, computer vision, and recommendation systems. I hold a B.Sc. in Physics and a Ph.D. in Statistics and enjoy hiking, free diving, and collecting vintage watches in my spare time. I'm passionate about using my skills to contribute to cutting-edge AI research and development. At GlobalSolTech, I'm proud to be part of a dynamic team driving technological advancement.

Using Command-R, we can automatically upgrade the original prompt to a RAG-style prompt to get more faithful adherence to the instructions, a clearer and more concise prompt, and in-line citations for free. Consider the following meta-prompt:

PYTHON
1meta_prompt = f'''Below is a task for an LLM delimited with ## Original Task. Your task is to split that task into two parts: (1) the context; and (2) the instructions.
2The context should be split into several separate parts and returned as a JSON object where each part has a name describing its contents and the value is the contents itself.
3Make sure to include all of the context contained in the original task description and do not change its meaning.
4The instructions should be re-written so that they are very clear and concise. Do not change the meaning of the instructions or task, just make sure they are very direct and clear.
5Return everything in a JSON object with the following structure:
6
7{{
8  "context": [{{"<description 1="" of="" part="">": "<content 1="" of="" part="">"}}, ...],
9  "instructions": "<the instructions="" re-written="">"
10}}
11
12## Original Task
13{original_prompt}
14'''
PYTHON
1print(meta_prompt)
Output
Below is a task for an LLM delimited with ## Original Task. Your task is to split that task into two parts: (1) the context; and (2) the instructions.
The context should be split into several separate parts and returned as a JSON object where each part has a name describing its contents and the value is the contents itself.
Make sure to include all of the context contained in the original task description and do not change its meaning.
The instructions should be re-written so that they are very clear and concise. Do not change the meaning of the instructions or task, just make sure they are very direct and clear.
Return everything in a JSON object with the following structure:
{
    "context": [{"<description 1="" of="" part="">": "<content 1="" of="" part="">"}, ...],
    "instructions": "<the instructions="" re-written="">"
}
## Original Task
## information
Current Job Title: Senior Software Engineer
Current Company Name: GlobalSolTech
Work Experience: Over 15 years of experience in software engineering, specializing in AI and machine learning. Proficient in Python, C++, and Java, with expertise in developing algorithms for natural language processing, computer vision, and recommendation systems.
Current Department Name: AI Research and Development
Education: B.Sc. in Physics from Trent University (2004), Ph.D. in Statistics from HEC in Paris (2010)
Hobbies: I love hiking in the mountains, free diving, and collecting and restoring vintage world war one mechanical watches.
Family: Married with 4 children and 3 grandchildren.
## instructions
Your task is to assist a user in writing a short biography for social media.
The length of the text should be no more than 100 words.
Write the summary in first person.

Command-R returns with the following:

PYTHON
1upgraded_prompt = co.chat(
2    message=meta_prompt,
3    model='command-r',
4)
PYTHON
1print(upgraded_prompt.text)
Output
Here is the task delved into a JSON object as requested:
```json JSON
{
    "context": [
    {
        "Work Experience": "Over 15 years of AI and machine learning engineering experience. Proficient in Python, C++, and Java, with expertise in developing algorithms for natural language processing, computer vision, and recommendation systems."
    },
    {
        "Education": "B.Sc. in Physics (Trent University, 2004) and Ph.D. in Statistics (HEC Paris, 2010)."
    },
    {
        "Personal Life": "I’m a married senior software engineer with 4 children and 3 grandchildren. I enjoy hiking, free diving, and vintage watch restoration."
    },
    {
        "Current Position": "I work at GlobalSolTech in the AI Research and Development department as a senior software engineer."
    }
    ],
    "instructions": "Using the provided information, write a concise, first-person social media biography of no more than 100 words."
}
```

To extract the returned information, we will write two simple functions to post-process out the JSON and then parse it.

PYTHON
1def get_json(text: str) -> str:
2    matches = [m.group(1) for m in re.finditer("```([\w\W]*?)```", text)]
3    if len(matches):
4        postproced = matches[0]
5        if postproced[:4] == 'json':
6            return postproced[4:]
7        return postproced
8    return text
PYTHON
1def get_prompt_and_docs(text: str) -> tuple:
2    json_obj = json.loads(get_json(text))
3    prompt = json_obj['instructions']
4    docs = []
5    for item in json_obj['context']:
6        for k,v in item.items():
7            docs.append({"title": k, "snippet": v})
8    return prompt, docs
PYTHON
1new_prompt, docs = get_prompt_and_docs(upgraded_prompt.text)
PYTHON
1new_prompt, docs
Output
('Using the provided information, write a concise, first-person social media biography of no more than 100 words.',
    [{'title': 'Work Experience',
    'snippet': 'Over 15 years of AI and machine learning engineering experience. Proficient in Python, C++, and Java, with expertise in developing algorithms for natural language processing, computer vision, and recommendation systems.'},
    {'title': 'Education',
    'snippet': 'B.Sc. in Physics (Trent University, 2004) and Ph.D. in Statistics (HEC Paris, 2010).'},
    {'title': 'Personal Life',
    'snippet': 'I’m a married senior software engineer with 4 children and 3 grandchildren. I enjoy hiking, free diving, and vintage watch restoration.'},
    {'title': 'Current Position',
    'snippet': 'I work at GlobalSolTech in the AI Research and Development department as a senior software engineer.'}])

As we can see above, the new prompt is much more concise and gets right to the point. The context has been split into 4 “documents” that Command-R can ground the information to. Now let’s run the same task with the new prompt while leveraging the documents= parameter. Note that the docs variable is a list of dict objects with title describing the contents of a text and snippet containing the text itself:

PYTHON
1response = co.chat(
2    message=new_prompt,
3    model='command-r',
4    documents=docs,
5)
PYTHON
1print(response.text)
Output
I'm a senior software engineer with a Ph.D. in Statistics and over 15 years of AI and machine learning engineering experience. My current focus at GlobalSolTech's AI R&amp;D department is developing algorithms for natural language processing, computer vision, and recommendation systems. In my free time, I enjoy hiking, freediving, and restoring vintage watches, and I'm a married father of four with three grandchildren.

The response is concise. More importantly, we can ensure that there is no hallucination because the text is automatically grounded in the input documents. Using the simple function below, we can add this grounding information to the text as citations:

PYTHON
1def insert_citations(text: str, citations: list[dict], add_one: bool=False):
2    """
3    A helper function to pretty print citations.
4    """
5    offset = 0
6    # Process citations in the order they were provided
7    for citation in citations:
8        # Adjust start/end with offset
9        start, end = citation.start + offset, citation.end + offset
10        if add_one:
11            cited_docs = [str(int(doc[4:]) + 1) for doc in citation.document_ids]
12        else:
13            cited_docs = [doc[4:] for doc in citation.document_ids]
14        # Shorten citations if they're too long for convenience
15        if len(cited_docs) > 3:
16            placeholder = "[" + ", ".join(cited_docs[:3]) + "...]"
17        else:
18            placeholder = "[" + ", ".join(cited_docs) + "]"
19        # ^ doc[4:] removes the 'doc_' prefix, and leaves the quoted document
20        modification = f'{text[start:end]} {placeholder}'
21        # Replace the cited text with its bolded version + placeholder
22        text = text[:start] + modification + text[end:]
23        # Update the offset for subsequent replacements
24        offset += len(modification) - (end - start)
25
26    return text
PYTHON
1print(insert_citations(response.text, response.citations, True))
Output
I'm a senior software engineer [3, 4] with a Ph.D. in Statistics [2] and over 15 years of AI and machine learning engineering experience. [1] My current focus at GlobalSolTech's AI R&amp;D department [4] is developing algorithms for natural language processing, computer vision, and recommendation systems. [1] In my free time, I enjoy hiking, freediving, and restoring vintage watches [3], and I'm a married father of four with three grandchildren. [3]

Now let’s move on to an arguably more difficult problem.

On March 21st, the DOJ announced that it is suing Apple for anti-competitive practices. The complaint is 88 pages long and consists of about 230 paragraphs of text. To understand what the suit alleges, a common use case would be to ask for a summary. Because Command-R has a context window of 128K, even an 88-page legal complaint fits comfortably within the window.

PYTHON
1apple = open('data/apple_mod.txt').read()
PYTHON
1tokens = co.tokenize(text=apple, model='command-r')
2len(tokens.tokens)
Output
29697

We can set up a prompt template that allows us to ask questions on the original text.

PYTHON
1prompt_template = '''
2{legal_text}
3
4{question}
5'''
PYTHON
1question = '''Please summarize the attached legal complaint succinctly. Focus on answering the question: what does the complaint allege?'''
2rendered_prompt = prompt_template.format(legal_text=apple, question=question)
PYTHON
1response = co.chat(
2    message=rendered_prompt,
3    model='command-r',
4    temperature=0.3,
5)
PYTHON
1print(response.text)
Output
The complaint alleges that Apple has violated antitrust laws by engaging in a pattern of anticompetitive conduct to maintain its monopoly power over the U.S. markets for smartphones and performance smartphones. Apple is accused of using its control over app distribution and access to its operating system to impede competition and innovation. Specifically, the company is said to have restricted developers' ability to create certain apps and limited the functionality of others, making it harder for consumers to switch away from iPhones to rival smartphones. This conduct is alleged to have harmed consumers and developers by reducing choice, increasing prices, and stifling innovation. The plaintiffs seek injunctive relief and potential monetary awards to remedy these illegal practices.

The summary seems clear enough. But we are interested in the specific allegations that the DOJ makes. For example, skimming the full complaint, it looks like the DOJ is alleging that Apple could encrypt text messages sent to Android phones if it wanted to do so. We can amend the rendered prompt and ask:

PYTHON
1question = '''Does the DOJ allege that Apple could encrypt text messages sent to Android phones?'''
2rendered_prompt = prompt_template.format(legal_text=apple, question=question)
PYTHON
1response = co.chat(
2    message=rendered_prompt,
3    model='command-r',
4)
PYTHON
1print(response.text)
Output
Yes, the DOJ alleges that Apple could allow iPhone users to send encrypted messages to Android users while still using iMessage on their iPhones but chooses not to do so. According to the DOJ, this would instantly improve the privacy and security of iPhones and other smartphones.

This is a very interesting allegation that at first glance suggests that the model could be hallucinating. Because RAG has been shown to help reduce hallucinations and grounds its responses in the input text, we should convert this prompt to the RAG style paradigm to gain confidence in its response.

While previously we asked Command-R to chunk the text for us, the legal complaint is highly structured with numbered paragraphs so we can use the following function to break the complaint into input docs ready for RAG:

PYTHON
1def chunk_doc(input_doc: str) -> list:
2    chunks = []
3    current_para = 'Preamble'
4    current_chunk = ''
5    # pattern to find an integer number followed by a dot (finding the explicitly numbered paragraph numbers)
6    pattern = r'^\d+\.$'
7
8    for line in input_doc.splitlines():
9        if re.match(pattern, line):
10            chunks.append((current_para.replace('.', ''), current_chunk))
11            current_chunk = ''
12            current_para = line
13        else:
14            current_chunk += line + '\n'
15
16    docs = []
17    for chunk in chunks:
18        docs.append({"title": chunk[0], "snippet": chunk[1]})
19
20    return docs
PYTHON
1chunks = chunk_doc(apple)
PYTHON
1print(chunks[18])
Output
1    {'title': '18', 'snippet': '\nProtecting competition and the innovation that competition inevitably ushers in\nfor consumers, developers, publishers, content creators, and device manufacturers is why\nPlaintiffs bring this lawsuit under Section 2 of the Sherman Act to challenge Apple’s\nmaintenance of its monopoly over smartphone markets, which affect hundreds of millions of\nAmericans every day. Plaintiffs bring this case to rid smartphone markets of Apple’s\nmonopolization and exclusionary conduct and to ensure that the next generation of innovators\ncan upend the technological world as we know it with new and transformative technologies.\n\n\nII.\n\nDefendant Apple\n\n'}

We can now try the same question but ask it directly to Command-R with the chunks as grounding information.

PYTHON
1response = co.chat(
2    message='''Does the DOJ allege that Apple could encrypt text messages sent to Android phones?''',
3    model='command-r',
4    documents=chunks,
5)
PYTHON
1print(response.text)
Output
Yes, according to the DOJ, Apple could encrypt text messages sent from iPhones to Android phones. The DOJ claims that Apple degrades the security and privacy of its users by impeding cross-platform encryption and preventing developers from fixing the broken cross-platform messaging experience. Apple's conduct makes it harder to switch from iPhone to Android, as messages sent from iPhones to Android phones are unencrypted.

The responses seem similar, but we should add citations and check the citation to get confidence in the response.

PYTHON
1print(insert_citations(response.text, response.citations))
Output
Yes, according to the DOJ, Apple could encrypt text messages sent from iPhones to Android phones. [144] The DOJ claims that Apple degrades the security and privacy [144] of its users by impeding cross-platform encryption [144] and preventing developers from fixing the broken cross-platform messaging experience. [93] Apple's conduct makes it harder to switch from iPhone to Android [144], as messages sent from iPhones to Android phones are unencrypted. [144]

The most important passage seems to be paragraph 144. Paragraph 93 is also cited. Let’s check what they contain.

PYTHON
1print(chunks[144]['snippet'])
Output
Apple is also willing to make the iPhone less secure and less private if that helps
maintain its monopoly power. For example, text messages sent from iPhones to Android phones
are unencrypted as a result of Apple’s conduct. If Apple wanted to, Apple could allow iPhone
users to send encrypted messages to Android users while still using iMessage on their iPhone,
which would instantly improve the privacy and security of iPhone and other smartphone users.
PYTHON
1print(chunks[93]['snippet'])
Output
Recently, Apple blocked a third-party developer from fixing the broken cross-
platform messaging experience in Apple Messages and providing end-to-end encryption for
messages between Apple Messages and Android users. By rejecting solutions that would allow
for cross-platform encryption, Apple continues to make iPhone users’ less secure than they could
otherwise be.
ii.

Paragraph 144 indeed contains the important allegation: If Apple wanted to, Apple could allow iPhone users to send encrypted messages to Android users.

In this cookbook we have shown how one can easily take an existing monolithic prompt and migrate it to the RAG paradigm to get less hallucination, grounded information, and in-line citations. We also demonstrated Command-R’s ability to re-write an instruction prompt in a single shot to make it more concise and potentially lead to higher quality completions.