LlamaIndex and Cohere's Models ↗

from llama_index.llms.cohere import Cohere
from llama_index.core.llms import ChatMessage

cohere_model = Cohere(
    api_key="COHERE_API_KEY", model="command-a-03-2025"
)

message = ChatMessage(role="user", content="What is 2 + 3?")

response = cohere_model.chat([message])
print(response)

from llama_index.embeddings.cohere import CohereEmbedding

embed_model = CohereEmbedding(
    api_key="COHERE_API_KEY",
    model_name="embed-english-v3.0",
    input_type="search_document",  # Use search_query for queries, search_document for documents
    max_tokens=8000,
    embedding_types=["float"],
)

embeddings = embed_model.get_text_embedding("Welcome to Cohere!")

# Print embeddings
print(len(embeddings))
print(embeddings[:5])

from llama_index.postprocessor.cohere_rerank import CohereRerank
from llama_index.readers.web import (
    SimpleWebPageReader,
)  # first, run `pip install llama-index-readers-web`

# create index (we are using an example page from Cohere's docs)
documents = SimpleWebPageReader(html_to_text=True).load_data(
    ["https://docs.cohere.com/v2/docs/cohere-embed"]
)  # you can replace this with any other reader or documents
index = VectorStoreIndex.from_documents(documents=documents)

# create reranker
cohere_rerank = CohereRerank(
    api_key="COHERE_API_KEY", model="rerank-english-v3.0", top_n=2
)

# query the index
query_engine = index.as_query_engine(
    similarity_top_k=10,
    node_postprocessors=[cohere_rerank],
)

print(query_engine)

# generate a response
response = query_engine.query(
    "What is Cohere's Embed Model?",
)

print(response)

# To view the source documents
from llama_index.core.response.pprint_utils import pprint_response

pprint_response(response, show_source=True)

from llama_index.llms.cohere import Cohere
from llama_index.embeddings.cohere import CohereEmbedding
from llama_index.postprocessor.cohere_rerank import CohereRerank
from llama_index.core import Settings
from llama_index.core import VectorStoreIndex
from llama_index.readers.web import (
    SimpleWebPageReader,
)  # first, run `pip install llama-index-readers-web`

# Create the embedding model
embed_model = CohereEmbedding(
    api_key="COHERE_API_KEY",
    model="embed-english-v3.0",
    input_type="search_document",
    max_tokens=8000,
    embedding_types=["float"],
)

# Create the service context with the cohere model for generation and embedding model
Settings.llm = Cohere(
    api_key="COHERE_API_KEY", model="command-a-03-2025"
)
Settings.embed_model = embed_model

# create index (we are using an example page from Cohere's docs)
documents = SimpleWebPageReader(html_to_text=True).load_data(
    ["https://docs.cohere.com/v2/docs/cohere-embed"]
)  # you can replace this with any other reader or documents
index = VectorStoreIndex.from_documents(documents=documents)

# Create a cohere reranker
cohere_rerank = CohereRerank(
    api_key="COHERE_API_KEY", model="rerank-english-v3.0", top_n=2
)

# Create the query engine
query_engine = index.as_query_engine(
    node_postprocessors=[cohere_rerank]
)

# Generate the response
response = query_engine.query("What is Cohere's Embed model?")
print(response)

from llama_index.llms.cohere import Cohere
from llama_index.core.tools import FunctionTool
from llama_index.core.agent import FunctionCallingAgent


# Define tools
def multiply(a: int, b: int) -> int:
    """Multiple two integers and returns the result integer"""
    return a * b


multiply_tool = FunctionTool.from_defaults(fn=multiply)


def add(a: int, b: int) -> int:
    """Add two integers and returns the result integer"""
    return a + b


add_tool = FunctionTool.from_defaults(fn=add)

# Define LLM
llm = Cohere(api_key="COHERE_API_KEY", model="command-a-03-2025")

# Create agent
agent = FunctionCallingAgent.from_tools(
    [multiply_tool, add_tool],
    llm=llm,
    verbose=True,
    allow_parallel_tool_calls=True,
)

# Run agent
response = await agent.achat("What is (121 * 3) + (5 * 8)?")