feat: Add Cognee integration

[hande-k]

Closes https://github.com/deepset-ai/haystack-private/issues/240

Summary

• Adds Cognee integration with 4 components: CogneeWriter, CogneeCognifier, CogneeRetriever, and CogneeMemoryStore
• CogneeWriter ingests Haystack Documents into Cognee memory via cognee.add() + optional cognee.cognify()
• CogneeRetriever searches Cognee's memory and returns Haystack Documents
• CogneeCognifier wraps cognee.cognify() as a standalone pipeline step
• CogneeMemoryStore implements the MemoryStore protocol from haystack-experimental for use with Haystack's experimental Agent

Test plan

• Unit tests pass via hatch run test:unit
• Linting passes via hatch run fmt-check
• Type checking passes via hatch run test:types
• Demo scripts tested manually (demo_pipeline.py, demo_memory_agent.py)

[CLAassistant]

All committers have signed the CLA.

[davidsbatista]

@hande-k thank you for this contribution!

I left some initial comments/suggestions for improvements.

[hande-k]

Thanks for the review @davidsbatista & @sjrl! I've addressed all the comments. A couple of notes:

• py.typed: Since the components are now split across components/retrievers/, components/writers/,
  and memory_stores/, I added py.typed markers at the parent level for each (matching the pattern used
  by other multi-package integrations).
• dataset_name on CogneeCognifier: Added str | list[str] | None support so it can target one or multiple
  datasets

Let me know if anything needs further adjustment!

[sjrl]

Would it be possible to make the type for item more strict? Looking at the code below it looks like it could be one of three things str, dict and some Cognee specific object.

[hande-k]

Cognee's search API returns Any internally. Results can be str (LLM completions), dict (structured outputs), or cognee model objects (DataPoint subclasses) depending on the search type. Since these internal types aren't part of cognee's API, I've kept Any but expanded the docstring to document the three expected categories. Would that work?

[sjrl]

Yeah expanding the docstring with the three supported types and the note that they are internal types would be great.

[sjrl]

Can you help me understand Cognee a bit better. It looks like from your code that dataset_name is required which makes sense.

My question is what level of scope is dataset_name meant to have? For example, is it normal to have a new dataset_name per user or is the intention for dataset to be scoped to multiple users (e.g. like at an org level)?

[sjrl]

To provide more context on why I'm asking this question is that our existing Mem0 memory store, most memories are scoped to users which is why our experimental protocol typically requires a user_id in all of its methods since the idea is that we are always scoping the request to a specific user. For example our protocol for add_memories looks like

def add_memories(self, *, messages: list[ChatMessage], user_id: str | None = None, **kwargs: Any) -> None:

with the idea being that user_id is set at run time to scope the request to a specific user.

This way we don't need to create a new MemoryStore for every user which is useful when we usually want to easily reuse a Haystack Agent or Pipeline for many users.

So my hope is that Cognee can also fit into this pattern.

[hande-k]

Cognee supports user-level scoping. Its add(), search(), and cognify() APIs all accept a user parameter for ACL-based access control. I've added user_id support to CogneeMemoryStore to match your existing pattern: a single store instance can serve multiple users via user_id at runtime. When user_id is provided, it's resolved to a cognee User and forwarded to all API calls.
dataset_name groups data logically (like a collection), while user_id controls who can access it. Search is scoped to the store's dataset_name for the given user, and shared datasets (where another user granted read permission) are automatically resolved. I've updated the examples/demo_memory_agent.py for a demo of isolation + sharing.
Does the current logic make sense for the integration?

[sjrl]

Yup that logic makes a lot of sense! Thanks for the changes.

I'd only say it would be good to expose the user resolution to the Retriever and Writer components as well since those are how pipeline builders will typically interact with the Cognee Memory Store.

[sjrl]

It would be great to also add some integration tests that require access to a Cognee account. We then mark those test with a header like

@pytest.mark.skipif(
    not os.environ.get("COGNEE_API_KEY", None),
    reason="Export an env var called COGNEE_API_KEY containing the Cognee API key to run this test.",
)
@pytest.mark.integration

[hande-k]

Added test_integration.py gated on LLM_API_KEY. This integration uses Cognee SDK locally as a library, so the only external dependency is the LLM provider (in default config it is OpenAI API key).
We can extend the integration in the future to use our Cloud as well, which would require using a cognee account

[sjrl]

Oh okay good to know. When using the Cognee SDK locally like this where do the memories get stored? Is the cognee library running a local db in the background or storing them in memory?

[sjrl]

Could we also add tests for CogneeCognifier and for the connectors/cognee/_utils.py file?

[sjrl]

Thanks for the review @davidsbatista & @sjrl! I've addressed all the comments. A couple of notes:

• py.typed: Since the components are now split across components/retrievers/, components/writers/,
  and memory_stores/, I added py.typed markers at the parent level for each (matching the pattern used
  by other multi-package integrations).
• dataset_name on CogneeCognifier: Added str | list[str] | None support so it can target one or multiple
  datasets

Let me know if anything needs further adjustment!

@hande-k thanks for your patience with us! This is going to be a great addition to Haystack. Since it’s the first version of a new abstraction, we really appreciate you working through all the comments with us as we refine how it fits in.

[sjrl]

Hi @hande-k, just checking in, do you have time to continue working on this? Happy to help or take over some of the remaining changes if that’s useful.

[hande-k]

hi @sjrl thanks for the new comments and offering your help to handle the requested changes, appreciate it!
I'll try to address them all early this week, hope that is fine. Let me know if there's urgency so that I can try to prioritize it.

[sjrl]

@hande-k you're welcome and that sounds good to me!

[sjrl]

What happens if the user is None?

[sjrl]

Let's make sure to only only use single backticks to wrap code sections otherwise our doc build complains.

Suggested change

	Search is restricted to the store's ``dataset_name``. If the user owns the
	dataset it is resolved by name; otherwise the store checks whether the user
	has been granted read access (e.g. via shared permissions) and searches by
	dataset UUID.
	When ``None``, cognee's default user is used.
	Search is restricted to the store's `dataset_name`. If the user owns the
	dataset it is resolved by name; otherwise the store checks whether the user
	has been granted read access (e.g. via shared permissions) and searches by
	dataset UUID.
	When `None`, cognee's default user is used.

[sjrl]

Following the pattern for our DocumentStores and their respective retrievers I'd expect the init method to take in a CogneeMemoryStore as init param. Check out our OpenSearchBM25Retriever as an example

haystack-core-integrations/integrations/opensearch/src/haystack_integrations/components/retrievers/opensearch/bm25_retriever.py

Lines 20 to 39 in c40817b

	class OpenSearchBM25Retriever:
	"""
	Fetches documents from OpenSearchDocumentStore using the keyword-based BM25 algorithm.
	BM25 computes a weighted word overlap between the query string and a document to determine its similarity.
	"""
	def __init__(
	self,
	*,
	document_store: OpenSearchDocumentStore,
	filters: dict[str, Any] | None = None,
	fuzziness: int | str = "AUTO",
	top_k: int = 10,
	scale_score: bool = False,
	all_terms_must_match: bool = False,
	filter_policy: str | FilterPolicy = FilterPolicy.REPLACE,
	custom_query: dict[str, Any] | None = None,
	raise_on_failure: bool = True,
	) -> None:

The idea is for the retriever to call the method directly from the Store. E.g. this is how the bm25 retrieval is run in the linked component docs = doc_store._bm25_retrieval(**bm25_args).

So it would be great if we could follow that pattern here as well.

[sjrl]

Related to this comment https://github.com/deepset-ai/haystack-core-integrations/pull/2979/changes#r3027766776 the run method should more or less accept all of the same arguments as the CogneeMemoryStore.search_memories function