CogneeRetriever

CogneeMemoryStore에서 메모리를 검색해 시스템 ChatMessage 객체로 반환하는 Retriever예요.

출처: 문서

본문

항목	내용
파이프라인에서의 일반적인 위치	메모리 보강 파이프라인에서 `Agent` 또는 Chat Generator 이전
필수 init 변수	`memory_store` — `CogneeMemoryStore` 인스턴스
선택적 init 변수	`top_k` — 반환할 최대 메모리 수(기본값은 스토어의 `top_k`)
필수 run 변수	`query` — 메모리를 검색할 텍스트 쿼리
선택적 run 변수	`user_id` — 검색 범위를 지정할 Cognee 사용자 ID. `None`을 넘기면 Cognee 기본 사용자 사용
출력 변수	`messages` — 시스템 `ChatMessage` 객체 리스트
API reference	Cognee
GitHub link	https://github.com/deepset-ai/haystack-core-integrations/tree/main/integrations/cognee
Package name	`cognee-haystack`

Overview

CogneeRetriever는 CogneeMemoryStore에서 메모리를 검색해 시스템 ChatMessage 객체로 반환합니다.

모델이 응답을 생성하기 전에 Agent나 채팅 생성 파이프라인에 장기 메모리를 주입할 때 사용해요.

검색 동작(검색 전략 search_type, 데이터셋, 세션 계층 포함)은 CogneeMemoryStore에서 구성합니다. retriever는 search_memories 위의 얇은 파이프라인 어댑터예요.

user_id 파라미터는 검색 범위를 특정 Cognee 사용자로 한정합니다. None을 넘기면 Cognee 기본 사용자를 사용해요.

Installation

Cognee 통합을 설치합니다:

pip install cognee-haystack

LLM API 키를 설정합니다(Cognee가 그래프 추출·쿼리에 사용):

export LLM_API_KEY="your-llm-api-key"

선택적으로 별도의 임베딩 API 키를 설정할 수 있어요(설정하지 않으면 LLM_API_KEY를 사용):

export EMBEDDING_API_KEY="your-embedding-api-key"

Usage

On its own

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

이 예시는 메모리를 검색해 현재 사용자 메시지 앞에 붙이고, 결합된 메시지 리스트를 Agent에 전달합니다.

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)],
        },
    }
)

더 알아보기 (Learn more)

CogneeRetriever — Haystack 공식 문서