c56bef871b
CodeQL / Analyze (python) (push) Has been cancelled
Update Platform Components Table / update (push) Has been cancelled
Docker image release / Build base image (push) Has been cancelled
Sync docs with Docusaurus / sync (push) Has been cancelled
Tests / Check if changed (push) Has been cancelled
Tests / format (push) Has been cancelled
Tests / check-imports (push) Has been cancelled
Tests / Unit / macos-latest (push) Has been cancelled
Tests / Unit / ubuntu-latest (push) Has been cancelled
Tests / Unit / windows-latest (push) Has been cancelled
Tests / mypy (push) Has been cancelled
Tests / Integration / ubuntu-latest (push) Has been cancelled
Tests / Integration / macos-latest (push) Has been cancelled
Tests / Integration / windows-latest (push) Has been cancelled
Tests / notify-slack-on-failure (push) Has been cancelled
Tests / Mark tests as completed (push) Has been cancelled
140 lines
4.4 KiB
Plaintext
140 lines
4.4 KiB
Plaintext
---
|
|
title: "CogneeRetriever"
|
|
id: cogneeretriever
|
|
slug: "/cogneeretriever"
|
|
description: "Retrieves memories from a CogneeMemoryStore and returns them as system ChatMessage objects."
|
|
---
|
|
|
|
# CogneeRetriever
|
|
|
|
Retrieves memories from a `CogneeMemoryStore` and returns them as system `ChatMessage` objects.
|
|
|
|
<div className="key-value-table">
|
|
|
|
| | |
|
|
| --- | --- |
|
|
| **Most common position in a pipeline** | Before an [`Agent`](../agents-1/agent.mdx) or Chat Generator in memory-augmented pipelines |
|
|
| **Mandatory init variables** | `memory_store`: A `CogneeMemoryStore` instance |
|
|
| **Optional init variables** | `top_k`: Maximum number of memories to return (defaults to the store's `top_k`) |
|
|
| **Mandatory run variables** | `query`: A text query to search memories |
|
|
| **Optional run variables** | `user_id`: Cognee user ID to scope the retrieval; pass `None` to use Cognee's default user |
|
|
| **Output variables** | `messages`: A list of system `ChatMessage` objects |
|
|
| **API reference** | [Cognee](/reference/integrations-cognee#cogneeretriever) |
|
|
| **GitHub link** | https://github.com/deepset-ai/haystack-core-integrations/tree/main/integrations/cognee |
|
|
| **Package name** | `cognee-haystack` |
|
|
|
|
</div>
|
|
|
|
## Overview
|
|
|
|
`CogneeRetriever` retrieves memories from a `CogneeMemoryStore` and returns them as system `ChatMessage` objects.
|
|
|
|
Use it to inject long-term memory into an Agent or a chat generation pipeline before the model produces a response.
|
|
|
|
Search behavior — including the search strategy (`search_type`), dataset, and session tier — is configured on the `CogneeMemoryStore`. The retriever is a thin pipeline adapter over `search_memories`.
|
|
|
|
The `user_id` parameter scopes the retrieval to a specific Cognee user. Pass `None` to use Cognee's default user.
|
|
|
|
## Installation
|
|
|
|
Install the Cognee integration:
|
|
|
|
```bash
|
|
pip install cognee-haystack
|
|
```
|
|
|
|
Set your LLM API key (used by Cognee for graph extraction and queries):
|
|
|
|
```bash
|
|
export LLM_API_KEY="your-llm-api-key"
|
|
```
|
|
|
|
Optionally, set a separate embedding API key (defaults to `LLM_API_KEY` when unset):
|
|
|
|
```bash
|
|
export EMBEDDING_API_KEY="your-embedding-api-key"
|
|
```
|
|
|
|
## Usage
|
|
|
|
### On its own
|
|
|
|
```python
|
|
from haystack.dataclasses import ChatMessage
|
|
|
|
from haystack_integrations.components.retrievers.cognee import CogneeRetriever
|
|
from haystack_integrations.memory_stores.cognee import CogneeMemoryStore
|
|
|
|
store = CogneeMemoryStore(search_type="GRAPH_COMPLETION", top_k=5)
|
|
|
|
# Write some memories first
|
|
store.add_memories(
|
|
messages=[ChatMessage.from_user("Alice prefers concise Python examples.")],
|
|
user_id="a1b2c3d4-e5f6-7890-abcd-ef1234567890",
|
|
)
|
|
|
|
retriever = CogneeRetriever(memory_store=store, top_k=3)
|
|
|
|
result = retriever.run(
|
|
query="What does Alice prefer?",
|
|
user_id="a1b2c3d4-e5f6-7890-abcd-ef1234567890",
|
|
)
|
|
memories = result["messages"]
|
|
print([message.text for message in memories])
|
|
```
|
|
|
|
### In a Pipeline
|
|
|
|
This example retrieves memories, prepends them to the current user message, and passes the combined message list to an Agent.
|
|
|
|
```python
|
|
from haystack import Pipeline
|
|
from haystack.components.agents import Agent
|
|
from haystack.components.converters import OutputAdapter
|
|
from haystack.components.generators.chat import OpenAIChatGenerator
|
|
from haystack.dataclasses import ChatMessage
|
|
|
|
from haystack_integrations.components.retrievers.cognee import CogneeRetriever
|
|
from haystack_integrations.memory_stores.cognee import CogneeMemoryStore
|
|
|
|
store = CogneeMemoryStore(dataset_name="my_agent_memory", session_id="alice_session_1")
|
|
|
|
pipeline = Pipeline()
|
|
pipeline.add_component("retriever", CogneeRetriever(memory_store=store, top_k=5))
|
|
pipeline.add_component(
|
|
"memory_context",
|
|
OutputAdapter(
|
|
template="{{ memories + user_messages }}",
|
|
output_type=list[ChatMessage],
|
|
unsafe=True,
|
|
),
|
|
)
|
|
pipeline.add_component(
|
|
"agent",
|
|
Agent(
|
|
chat_generator=OpenAIChatGenerator(model="gpt-4o-mini"),
|
|
system_prompt=(
|
|
"Use any system messages at the start of the conversation as long-term memory. "
|
|
"Answer concisely."
|
|
),
|
|
),
|
|
)
|
|
|
|
pipeline.connect("retriever.messages", "memory_context.memories")
|
|
pipeline.connect("memory_context.output", "agent.messages")
|
|
|
|
query = "Give me a short implementation tip."
|
|
|
|
pipeline.run(
|
|
{
|
|
"retriever": {
|
|
"query": query,
|
|
"user_id": "a1b2c3d4-e5f6-7890-abcd-ef1234567890",
|
|
},
|
|
"memory_context": {
|
|
"user_messages": [ChatMessage.from_user(query)],
|
|
},
|
|
}
|
|
)
|
|
```
|