RAG with Cohere

1
# pip install cohere
2
3
import cohere
4
import numpy as np
5
import json
6
from typing import List
7
8
co = cohere.ClientV2(api_key="COHERE_API_KEY") # Get your free API key: https://dashboard.cohere.com/api-keys

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
]

1
# Add the user query
2
query = "Are there health benefits?"
3
4
# Generate the response
5
response = co.chat(model="command-r-plus-08-2024",
6
                   messages=[{'role': 'user', 'content': query}],
7
                   documents=documents)
8
9
# Display the response
10
print(response.message.content[0].text)
11
12
# Display the citations and source documents
13
if response.message.citations:
14
    print("\nCITATIONS:")
15
    for citation in response.message.citations:
16
        print(citation, "\n")

1
def generate_search_queries(message: str) -> List[str]:
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
26
    # Define a preamble to optimize search query generation
27
    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."
28
29
    # Generate search queries (if any)
30
    search_queries = []
31
    
32
    res = co.chat(
33
        model="command-r-08-2024",
34
        messages=[
35
            {"role": "system", "content": instructions},
36
            {"role": "user", "content": message},
37
        ],
38
        tools=query_gen_tool
39
    )
40
    
41
    if res.message.tool_calls:
42
        for tc in res.message.tool_calls:
43
            queries = json.loads(tc.function.arguments)["queries"]
44
            search_queries.extend(queries)
45
46
    return search_queries

1
query = "How to stay connected with the company, and do you organize team events?"
2
queries_for_search = generate_search_queries(query)
3
print(queries_for_search)

1
query = "How flexible are the working hours"
2
queries_for_search = generate_search_queries(query)
3
print(queries_for_search)

1
query = "What is 2 + 2"
2
queries_for_search = generate_search_queries(query)
3
print(queries_for_search)

1
# Define the documents
2
faqs_long = [
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 faqs_long],
60
            embedding_types=["float"]).embeddings.float

1
# Add the user query
2
query = "How to get to know my teammates"
3
4
# Generate the search query
5
# Note: For simplicity, we are assuming only one query generated. For actual implementations, you will need to perform search for each query.
6
queries_for_search = generate_search_queries(query)[0]
7
print("Search query: ", queries_for_search)
8
9
# Embed the search query
10
query_emb = co.embed(
11
    model="embed-english-v3.0",
12
    input_type="search_query",
13
    texts=[queries_for_search],
14
    embedding_types=["float"]).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 = [faqs_long[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(query=queries_for_search,
3
                    documents=[doc['data']['text'] for doc in retrieved_documents],
4
                    top_n=2,
5
                    model='rerank-english-v3.0')
6
7
# Display the reranking results
8
for idx, result in enumerate(results.results):
9
    print(f"Rank: {idx+1}") 
10
    print(f"Score: {result.relevance_score}")
11
    print(f"Document: {retrieved_documents[result.index]}\n")
12
    
13
reranked_documents = [retrieved_documents[result.index] for result in results.results]

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