Integrations

LlamaIndex

Prerequisite

To use LlamaIndex and Cohere, you will need:

  • LlamaIndex Package. To install it, run:
    • pip install llama-index
    • pip install llama-index-llms-cohere (to use the Command models)
    • pip install llama-index-embeddings-cohere (to use the Embed models)
    • pip install llama-index-postprocessor-cohere-rerank (to use the Rerank models)
  • Cohere’s SDK. To install it, run pip install cohere. If you run into any issues or want more details on Cohere’s SDK, see this wiki.
  • A Cohere API Key. For more details on pricing see this page. When you create an account with Cohere, we automatically create a trial API key for you. This key will be available on the dashboard where you can copy it, and it’s in the dashboard section called “API Keys” as well.

Cohere Chat with LlamaIndex

To use Cohere’s chat functionality with LlamaIndex create a Cohere model object and call the chat function.

PYTHON
1from llama_index.llms.cohere import Cohere
2from llama_index.core.llms import ChatMessage
3
4cohere_model = Cohere(api_key="COHERE_API_KEY",
5 model="command-r-plus")
6
7message = ChatMessage(role="user", content= "What is 2 + 3?")
8
9response = cohere_model.chat([message])
10print(response)

Cohere Embeddings with LlamaIndex

To use Cohere’s embeddings with LlamaIndex create a Cohere Embeddings object with an embedding model from this list and call get_text_embedding.

PYTHON
1from llama_index.embeddings.cohere import CohereEmbedding
2
3embed_model = CohereEmbedding(
4 api_key="COHERE_API_KEY",
5 model_name="embed-english-v3.0",
6 input_type="search_document", # Use search_query for queries, search_document for documents
7)
8
9# Generate Embeddings
10embeddings = embed_model.get_text_embedding("Welcome to Cohere!")
11
12# Print embeddings
13print(len(embeddings))
14print(embeddings[:5])

Cohere Rerank with LlamaIndex

To use Cohere’s rerank functionality with LlamaIndex create a Cohere Rerank object and use as a node post processor.

PYTHON
1from llama_index.postprocessor.cohere_rerank import CohereRerank
2from llama_index.readers.web import SimpleWebPageReader # first, run `pip install llama-index-readers-web`
3
4# create index (we are using an example page from Cohere's docs)
5documents = SimpleWebPageReader(html_to_text=True).load_data(
6 ["https://docs.cohere.com/v2/docs/prompt-tuner"]
7) # you can replace this with any other reader or documents
8index = VectorStoreIndex.from_documents(documents=documents)
9
10# create reranker
11cohere_rerank = CohereRerank(api_key="COHERE_API_KEY",
12 model="rerank-english-v3.0",
13 top_n=2)
14
15# query the index
16query_engine = index.as_query_engine(
17 similarity_top_k=10,
18 node_postprocessors=[cohere_rerank],
19)
20
21print(query_engine)
22
23# generate a response
24response = query_engine.query(
25 "What is Cohere Prompt Tuner?",
26)
27
28print(response)
29
30# To view the source documents
31from llama_index.core.response.pprint_utils import pprint_response
32pprint_response(response, show_source=True)

Cohere RAG with LlamaIndex

The following example uses Cohere’s chat model, embeddings and rerank functionality to generate a response based on your data.

PYTHON
1from llama_index.llms.cohere import Cohere
2from llama_index.embeddings.cohere import CohereEmbedding
3from llama_index.postprocessor.cohere_rerank import CohereRerank
4from llama_index.core import Settings
5from llama_index.core import VectorStoreIndex
6from llama_index.readers.web import SimpleWebPageReader # first, run `pip install llama-index-readers-web`
7
8# Create the embedding model
9embed_model = CohereEmbedding(
10 api_key="COHERE_API_KEY",
11 model_name="embed-english-v3.0",
12 input_type="search_query",
13)
14
15# Create the service context with the cohere model for generation and embedding model
16Settings.llm = Cohere(api_key="COHERE_API_KEY",
17 model="command-r-plus")
18Settings.embed_model = embed_model
19
20# create index (we are using an example page from Cohere's docs)
21documents = SimpleWebPageReader(html_to_text=True).load_data(
22 ["https://docs.cohere.com/v2/docs/prompt-tuner"]
23) # you can replace this with any other reader or documents
24index = VectorStoreIndex.from_documents(documents=documents)
25
26# Create a cohere reranker
27cohere_rerank = CohereRerank(api_key="COHERE_API_KEY",
28 model="rerank-english-v3.0",
29 top_n=2)
30
31# Create the query engine
32query_engine = index.as_query_engine(node_postprocessors=[cohere_rerank])
33
34# Generate the response
35response = query_engine.query("What is Cohere Prompt Tuner?")
36print(response)

Cohere Tool Use (Function Calling) with LlamaIndex

To use Cohere’s tool use functionality with LlamaIndex, you can use the FunctionTool class to create a tool that uses Cohere’s API.

PYTHON
1from llama_index.llms.cohere import Cohere
2from llama_index.core.tools import FunctionTool
3from llama_index.core.agent import FunctionCallingAgent
4
5# Define tools
6def multiply(a: int, b: int) -> int:
7 """Multiple two integers and returns the result integer"""
8 return a * b
9
10multiply_tool = FunctionTool.from_defaults(fn=multiply)
11
12def add(a: int, b: int) -> int:
13 """Add two integers and returns the result integer"""
14 return a + b
15
16add_tool = FunctionTool.from_defaults(fn=add)
17
18# Define LLM
19llm = Cohere(api_key="COHERE_API_KEY",
20 model="command-r-plus")
21
22# Create agent
23agent = FunctionCallingAgent.from_tools(
24 [multiply_tool, add_tool],
25 llm=llm,
26 verbose=True,
27 allow_parallel_tool_calls=True,
28)
29
30# Run agent
31response = await agent.achat("What is (121 * 3) + (5 * 8)?")