diff --git a/llama-index-integrations/graph_rag/llama-index-graph-rag-cognee/.gitignore b/llama-index-integrations/graph_rag/llama-index-graph-rag-cognee/.gitignore index 73c1a9c5f67c4..1200d265cee5b 100644 --- a/llama-index-integrations/graph_rag/llama-index-graph-rag-cognee/.gitignore +++ b/llama-index-integrations/graph_rag/llama-index-graph-rag-cognee/.gitignore @@ -174,4 +174,4 @@ env/ venv/ ENV/ env.bak/ -venv.bak/ \ No newline at end of file +venv.bak/ diff --git a/llama-index-integrations/graph_rag/llama-index-graph-rag-cognee/README.md b/llama-index-integrations/graph_rag/llama-index-graph-rag-cognee/README.md index f7591cb42b069..a0481c9b84ca6 100644 --- a/llama-index-integrations/graph_rag/llama-index-graph-rag-cognee/README.md +++ b/llama-index-integrations/graph_rag/llama-index-graph-rag-cognee/README.md @@ -4,7 +4,6 @@ Cognee assists developers in introducing greater predictability and management i For more information, visit [Cognee documentation](https://docs.cognee.ai/) - ## Installation ```shell @@ -21,6 +20,7 @@ import asyncio from llama_index.core import Document from llama_index.graph_rag.cognee import CogneeGraphRAG + async def example_graph_rag_cognee(): # Gather documents to add to GraphRAG news = pd.read_csv( @@ -28,7 +28,8 @@ async def example_graph_rag_cognee(): )[:5] news.head() documents = [ - Document(text=f"{row['title']}: {row['text']}") for i, row in news.iterrows() + Document(text=f"{row['title']}: {row['text']}") + for i, row in news.iterrows() ] # Instantiate cognee GraphRAG @@ -42,7 +43,7 @@ async def example_graph_rag_cognee(): # Add data to cognee await cogneeRAG.add(documents, "test") - + # Process data into a knowledge graph await cogneeRAG.process_data("test") @@ -61,4 +62,4 @@ async def example_graph_rag_cognee(): if __name__ == "__main__": asyncio.run(example_graph_rag_cognee()) -``` \ No newline at end of file +``` diff --git a/llama-index-integrations/graph_rag/llama-index-graph-rag-cognee/llama_index/graph_rag/cognee/__init__.py b/llama-index-integrations/graph_rag/llama-index-graph-rag-cognee/llama_index/graph_rag/cognee/__init__.py index 48384a2ebb495..03a993e559537 100644 --- a/llama-index-integrations/graph_rag/llama-index-graph-rag-cognee/llama_index/graph_rag/cognee/__init__.py +++ b/llama-index-integrations/graph_rag/llama-index-graph-rag-cognee/llama_index/graph_rag/cognee/__init__.py @@ -1,3 +1,3 @@ from llama_index.graph_rag.cognee.graph_rag import CogneeGraphRAG -__all__ = ['CogneeGraphRAG'] \ No newline at end of file +__all__ = ["CogneeGraphRAG"] diff --git a/llama-index-integrations/graph_rag/llama-index-graph-rag-cognee/llama_index/graph_rag/cognee/base.py b/llama-index-integrations/graph_rag/llama-index-graph-rag-cognee/llama_index/graph_rag/cognee/base.py index 043061a42f804..68b56e5b17e23 100644 --- a/llama-index-integrations/graph_rag/llama-index-graph-rag-cognee/llama_index/graph_rag/cognee/base.py +++ b/llama-index-integrations/graph_rag/llama-index-graph-rag-cognee/llama_index/graph_rag/cognee/base.py @@ -1,32 +1,33 @@ from abc import abstractmethod from typing import Protocol + # NOTE: This is a bare-bone suggestion for an abstract protocol to define GraphRAG for llama-index # This should be expanded upon and integrated to llama-index-core to support multiple different GraphRAG # libraries in the future class GraphRAG(Protocol): """Abstract graph RAG protocol. - This protocol defines the interface for a graphRAG, which is responsible - for adding, storing, processing and retrieving information from knowledge graphs. + This protocol defines the interface for a graphRAG, which is responsible + for adding, storing, processing and retrieving information from knowledge graphs. + + Attributes: + llm_api_key: str: Api key for desired llm. + graph_db_provider: str: The graph database provider. + vector_db_provider: str: The vector database provider. + relational_db_provider: str: The relational database provider. + db_name: str: The name of the databases. + """ - Attributes: - llm_api_key: str: Api key for desired llm. - graph_db_provider: str: The graph database provider. - vector_db_provider: str: The vector database provider. - relational_db_provider: str: The relational database provider. - db_name: str: The name of the databases. - """ @abstractmethod async def add(self, data, dataset_name): """Add data to the specified dataset. This data will later be processed and made into a knowledge graph. - Args: + Args: data (Any): The data to be added to the graph. dataset_name (str): Name of the dataset or node set where the data will be added. """ - pass @abstractmethod async def process_data(self, dataset_name: str): @@ -35,7 +36,6 @@ async def process_data(self, dataset_name: str): Args: dataset_name (str): The dataset name to process. """ - pass @abstractmethod async def search(self, query: str): @@ -44,7 +44,6 @@ async def search(self, query: str): Args: query (str): The query string to match against data from the graph. """ - pass @abstractmethod async def get_related_nodes(self, node_id: str): @@ -53,4 +52,3 @@ async def get_related_nodes(self, node_id: str): Args: node_id (str): The name of the node to match against nodes in the graph. """ - pass diff --git a/llama-index-integrations/graph_rag/llama-index-graph-rag-cognee/llama_index/graph_rag/cognee/graph_rag.py b/llama-index-integrations/graph_rag/llama-index-graph-rag-cognee/llama_index/graph_rag/cognee/graph_rag.py index 9eb26017e9019..0ab8732eaa0ab 100644 --- a/llama-index-integrations/graph_rag/llama-index-graph-rag-cognee/llama_index/graph_rag/cognee/graph_rag.py +++ b/llama-index-integrations/graph_rag/llama-index-graph-rag-cognee/llama_index/graph_rag/cognee/graph_rag.py @@ -2,15 +2,13 @@ import pathlib from typing import List -from cognee.modules.users.methods import get_default_user -from cognee.modules.data.methods import get_datasets_by_name -from cognee.api.v1.search import SearchType import cognee from llama_index.core import Document from .base import GraphRAG + class CogneeGraphRAG(GraphRAG): def __init__( self, @@ -18,14 +16,11 @@ def __init__( graph_db_provider, vector_db_provider, relational_db_provider, - db_name - ): - + db_name, + ) -> None: cognee.config.set_llm_config({"llm_api_key": llm_api_key}) - cognee.config.set_vector_db_config( - {"vector_db_provider": vector_db_provider} - ) + cognee.config.set_vector_db_config({"vector_db_provider": vector_db_provider}) cognee.config.set_relational_db_config( {"db_provider": relational_db_provider, "db_name": db_name} ) @@ -49,7 +44,7 @@ async def add(self, data, dataset_name: str): """Add data to the specified dataset. This data will later be processed and made into a knowledge graph. - Args: + Args: data (Any): The data to be added to the graph. dataset_name (str): Name of the dataset or node set where the data will be added. """ @@ -68,7 +63,9 @@ async def process_data(self, dataset_names: str): dataset_name (str): The dataset name to process. """ user = await cognee.modules.users.methods.get_default_user() - datasets = await cognee.modules.data.methods.get_datasets_by_name(dataset_names, user.id) + datasets = await cognee.modules.data.methods.get_datasets_by_name( + dataset_names, user.id + ) await cognee.cognify(datasets, user) async def get_graph_url(self, graphistry_password, graphistry_username) -> str: @@ -103,7 +100,9 @@ async def search(self, query: str): query (str): The query string to match against data from the graph. """ user = await cognee.modules.users.methods.get_default_user() - return await cognee.search(cognee.api.v1.search.SearchType.SUMMARIES, query, user) + return await cognee.search( + cognee.api.v1.search.SearchType.SUMMARIES, query, user + ) async def get_related_nodes(self, node_id: str): """Search the graph for relevant nodes or relationships based on node id. @@ -112,4 +111,6 @@ async def get_related_nodes(self, node_id: str): node_id (str): The name of the node to match against nodes in the graph. """ user = await cognee.modules.users.methods.get_default_user() - return await cognee.search(cognee.api.v1.search.SearchType.INSIGHTS, node_id, user) + return await cognee.search( + cognee.api.v1.search.SearchType.INSIGHTS, node_id, user + ) diff --git a/llama-index-integrations/graph_rag/llama-index-graph-rag-cognee/tests/test_graph_rag_cognee.py b/llama-index-integrations/graph_rag/llama-index-graph-rag-cognee/tests/test_graph_rag_cognee.py index 0a7babc338827..17ac24dc3d95c 100644 --- a/llama-index-integrations/graph_rag/llama-index-graph-rag-cognee/tests/test_graph_rag_cognee.py +++ b/llama-index-integrations/graph_rag/llama-index-graph-rag-cognee/tests/test_graph_rag_cognee.py @@ -49,4 +49,4 @@ async def test_graph_rag_cognee(): if __name__ == "__main__": - asyncio.run(test_graph_rag_cognee()) \ No newline at end of file + asyncio.run(test_graph_rag_cognee())