DocumentationAPI Reference📓 Tutorials🧑‍🍳 Cookbook🤝 Integrations💜 Discord

CohereGenerator

CohereGenerator enables text generation using Cohere's large language models (LLMs).

NameCohereGenerator
Pathhttps://github.com/deepset-ai/haystack-core-integrations/tree/main/integrations/cohere
Most common Position in a PipelineAfter a PromptBuilder
Mandatory Input variables“prompt”: a string containing the prompt for the LLM
Output variables“replies”: a list of strings with all the replies generated by the LLM

”meta”: a list of dictionaries with the metadata associated with each reply, such as token count, finish reason, and so on

This integration supports Cohere generate models such as command and command-light. Check out the most recent full list in Cohere documentation.

Parameters Overview

CohereGenerator needs a Cohere API key to work. You can write this key in:

  • The api_key init parameter
  • The COHERE_API_KEY environment variable (recommended)

Then, the component needs a prompt to operate, but you can pass any text generation parameters valid for the Co.generate method directly to this component using the generation_kwargs parameter, both at initialization and to run() method. For more details on the parameters supported by the Cohere API, refer to the Cohere documentation.

Streaming

CohereGenerator supports streaming the tokens from the LLM directly in output. To do so, pass a function to the streaming_callback init parameter.

Usage

You need to install cohere-haystack package to use the CohereGenerator:

pip install cohere-haystack

On its own

Basic usage:

from haystack_integrations.components.generators.cohere import CohereGenerator


client = CohereGenerator(api_key=api_key)
response = client.run("Briefly explain what NLP is in one sentence.")
print(response)

>>> {'replies': ["Natural Language Processing (NLP) is a subfield of artificial intelligence and computational linguistics that focuses on the interaction between computers and human languages..."],
 'meta': [{'finish_reason': 'COMPLETE'}]}

With streaming:

from haystack_integrations.components.generators.cohere import CohereGenerator


client = CohereGenerator(streaming_callback=lambda chunk: print(chunk.text, end="", flush=True))
response = client.run("Briefly explain what NLP is in one sentence.")
print(response)

>>> Natural Language Processing (NLP) is the study of natural language and how it can be used to solve problems through computational methods, enabling machines to understand, interpret, and generate human language. 

>>>{'replies': [' Natural Language Processing (NLP) is the study of natural language and how it can be used to solve problems through computational methods, enabling machines to understand, interpret, and generate human language.'], 'meta': [{'index': 0, 'finish_reason': 'COMPLETE'}]}

In a pipeline

In a RAG pipeline:

from haystack import Pipeline
from haystack.components.retrievers.in_memory import InMemoryBM25Retriever
from haystack.components.builders.prompt_builder import PromptBuilder
from haystack.document_stores.in_memory import InMemoryDocumentStore
from haystack_integrations.components.generators.cohere import CohereGenerator
from haystack import Document

docstore = InMemoryDocumentStore()
docstore.write_documents([Document(content="Rome is the capital of Italy"), Document(content="Paris is the capital of France")])

query = "What is the capital of France?"

template = """
Given the following information, answer the question.

Context: 
{% for document in documents %}
    {{ document.content }}
{% endfor %}

Question: {{ query }}?
"""
pipe = Pipeline()

pipe.add_component("retriever", InMemoryBM25Retriever(document_store=docstore))
pipe.add_component("prompt_builder", PromptBuilder(template=template))
pipe.add_component("llm", CohereGenerator(api_key="YOUR_API_KEY"))
pipe.connect("retriever", "prompt_builder.documents")
pipe.connect("prompt_builder", "llm")

res=pipe.run({
    "prompt_builder": {
        "query": query
    },
    "retriever": {
        "query": query
    }
})

print(res)

Related Links

Check out the API reference in the GitHub repo or in our docs: