CohereGenerator
CohereGenerator
enables text generation using Cohere's large language models (LLMs).
Most common position in a pipeline | After a PromptBuilder |
Mandatory init variables | "api_key": The Cohere API key. Can be set with COHERE_API_KEY or CO_API_KEY env var. |
Mandatory run 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 |
API reference | Cohere |
GitHub link | https://github.com/deepset-ai/haystack-core-integrations/tree/main/integrations/cohere |
This integration supports Cohere models such as command
, command-r
and comman-r-plus
. Check out the most recent full list in Cohere documentation.
Overview
CohereGenerator
needs a Cohere API key to work. You can write this key in:
- The
api_key
init parameter using Secret API - The
COHERE_API_KEY
environment variable (recommended)
Then, the component needs a prompt to operate, but you can pass any text generation parameters directly to this component using the generation_kwargs
parameter at initialization. For more details on the parameters supported by the Cohere API, refer to the Cohere documentation.
Streaming
This Generator 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()
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.content, 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())
pipe.connect("retriever", "prompt_builder.documents")
pipe.connect("prompt_builder", "llm")
res=pipe.run({
"prompt_builder": {
"query": query
},
"retriever": {
"query": query
}
})
print(res)
Updated 3 months ago
Check out the API reference in the GitHub repo or in our docs: