Skip to main content

CogneeRetriever

This will help you get started with the Cognee retriever. For detailed documentation of all CogneeRetriever features and configurations head to the API reference.

Integration details

Bring-your-own data (i.e., index and search a custom corpus of documents):
RetrieverSelf-hostCloud offeringPackage
CogneeRetrieverlangchain-cognee

Setup

For cognee default setup, only thing you need is your OpenAI API key. If you want to get automated tracing from individual queries, you can also set your LangSmith API key by uncommenting below:
# os.environ["LANGSMITH_API_KEY"] = getpass.getpass("Enter your LangSmith API key: ")
# os.environ["LANGSMITH_TRACING"] = "true"

Installation

This retriever lives in the langchain-cognee package:
%pip install -qU langchain-cognee
import nest_asyncio

nest_asyncio.apply()

Instantiation

Now we can instantiate our retriever:
from langchain_cognee import CogneeRetriever

retriever = CogneeRetriever(
    llm_api_key="sk-",  # OpenAI API Key
    dataset_name="my_dataset",
    k=3,
)

Usage

Add some documents, process them, and then run queries. Cognee retrieves relevant knowledge to your queries and generates final answers.
# Example of adding and processing documents
from langchain_core.documents import Document

docs = [
    Document(page_content="Elon Musk is the CEO of SpaceX."),
    Document(page_content="SpaceX focuses on rockets and space travel."),
]

retriever.add_documents(docs)
retriever.process_data()

# Now let's query the retriever
query = "Tell me about Elon Musk"
results = retriever.invoke(query)

for idx, doc in enumerate(results, start=1):
    print(f"Doc {idx}: {doc.page_content}")

API reference

TODO: add link to API reference.
I