cognee 0.3.6__py3-none-any.whl → 0.3.7__py3-none-any.whl

cognee/modules/retrieval/insights_retriever.py

@@ -1,133 +0,0 @@
-import asyncio
-from typing import Any, Optional
-
-from cognee.modules.graph.cognee_graph.CogneeGraphElements import Edge, Node
-from cognee.modules.retrieval.base_graph_retriever import BaseGraphRetriever
-from cognee.shared.logging_utils import get_logger
-from cognee.infrastructure.databases.graph import get_graph_engine
-from cognee.infrastructure.databases.vector import get_vector_engine
-from cognee.modules.retrieval.exceptions.exceptions import NoDataError
-from cognee.infrastructure.databases.vector.exceptions.exceptions import CollectionNotFoundError
-
-logger = get_logger("InsightsRetriever")
-
-
-class InsightsRetriever(BaseGraphRetriever):
-    """
-    Retriever for handling graph connection-based insights.
-
-    Public methods include:
-    - get_context
-    - get_completion
-
-    Instance variables include:
-    - exploration_levels
-    - top_k
-    """
-
-    def __init__(self, exploration_levels: int = 1, top_k: Optional[int] = 5):
-        """Initialize retriever with exploration levels and search parameters."""
-        self.exploration_levels = exploration_levels
-        self.top_k = top_k
-
-    async def get_context(self, query: str) -> list:
-        """
-        Find neighbours of a given node in the graph.
-
-        If the provided query does not correspond to an existing node,
-        search for similar entities and retrieve their connections.
-        Reraises NoDataError if there is no data found in the system.
-
-        Parameters:
-        -----------
-
-            - query (str): A string identifier for the node whose neighbours are to be
-              retrieved.
-
-        Returns:
-        --------
-
-            - list: A list of unique connections found for the queried node.
-        """
-        if query is None:
-            return []
-
-        node_id = query
-        graph_engine = await get_graph_engine()
-        exact_node = await graph_engine.extract_node(node_id)
-
-        if exact_node is not None and "id" in exact_node:
-            node_connections = await graph_engine.get_connections(str(exact_node["id"]))
-        else:
-            vector_engine = get_vector_engine()
-
-            try:
-                results = await asyncio.gather(
-                    vector_engine.search("Entity_name", query_text=query, limit=self.top_k),
-                    vector_engine.search("EntityType_name", query_text=query, limit=self.top_k),
-                )
-            except CollectionNotFoundError as error:
-                logger.error("Entity collections not found")
-                raise NoDataError("No data found in the system, please add data first.") from error
-
-            results = [*results[0], *results[1]]
-            relevant_results = [result for result in results if result.score < 0.5][: self.top_k]
-
-            if len(relevant_results) == 0:
-                return []
-
-            node_connections_results = await asyncio.gather(
-                *[graph_engine.get_connections(result.id) for result in relevant_results]
-            )
-
-            node_connections = []
-            for neighbours in node_connections_results:
-                node_connections.extend(neighbours)
-
-        unique_node_connections_map = {}
-        unique_node_connections = []
-
-        for node_connection in node_connections:
-            if "id" not in node_connection[0] or "id" not in node_connection[2]:
-                continue
-
-            unique_id = f"{node_connection[0]['id']} {node_connection[1]['relationship_name']} {node_connection[2]['id']}"
-            if unique_id not in unique_node_connections_map:
-                unique_node_connections_map[unique_id] = True
-                unique_node_connections.append(node_connection)
-
-        return unique_node_connections
-        # return [
-        #     Edge(
-        #         node1=Node(node_id=connection[0]["id"], attributes=connection[0]),
-        #         node2=Node(node_id=connection[2]["id"], attributes=connection[2]),
-        #         attributes={
-        #             **connection[1],
-        #             "relationship_type": connection[1]["relationship_name"],
-        #         },
-        #     )
-        #     for connection in unique_node_connections
-        # ]
-
-    async def get_completion(self, query: str, context: Optional[Any] = None) -> Any:
-        """
-        Returns the graph connections context.
-
-        If a context is not provided, it fetches the context using the query provided.
-
-        Parameters:
-        -----------
-
-            - query (str): A string identifier used to fetch the context.
-            - context (Optional[Any]): An optional context to use for the completion; if None,
-              it fetches the context based on the query. (default None)
-
-        Returns:
-        --------
-
-            - Any: The context used for the completion, which is either provided or fetched
-              based on the query.
-        """
-        if context is None:
-            context = await self.get_context(query)
-        return context

cognee/tests/test_memgraph.py

@@ -1,109 +0,0 @@
-import os
-
-import pathlib
-import cognee
-from cognee.infrastructure.files.storage import get_storage_config
-from cognee.modules.search.operations import get_history
-from cognee.modules.users.methods import get_default_user
-from cognee.shared.logging_utils import get_logger
-from cognee.modules.search.types import SearchType
-
-
-logger = get_logger()
-
-
-async def main():
-    cognee.config.set_graph_database_provider("memgraph")
-    data_directory_path = str(
-        pathlib.Path(
-            os.path.join(pathlib.Path(__file__).parent, ".data_storage/test_memgraph")
-        ).resolve()
-    )
-    cognee.config.data_root_directory(data_directory_path)
-    cognee_directory_path = str(
-        pathlib.Path(
-            os.path.join(pathlib.Path(__file__).parent, ".cognee_system/test_memgraph")
-        ).resolve()
-    )
-    cognee.config.system_root_directory(cognee_directory_path)
-
-    await cognee.prune.prune_data()
-    await cognee.prune.prune_system(metadata=True)
-
-    dataset_name = "cs_explanations"
-
-    explanation_file_path = os.path.join(
-        pathlib.Path(__file__).parent, "test_data/Natural_language_processing.txt"
-    )
-    await cognee.add([explanation_file_path], dataset_name)
-
-    text = """A quantum computer is a computer that takes advantage of quantum mechanical phenomena.
-    At small scales, physical matter exhibits properties of both particles and waves, and quantum computing leverages this behavior, specifically quantum superposition and entanglement, using specialized hardware that supports the preparation and manipulation of quantum states.
-    Classical physics cannot explain the operation of these quantum devices, and a scalable quantum computer could perform some calculations exponentially faster (with respect to input size scaling) than any modern "classical" computer. In particular, a large-scale quantum computer could break widely used encryption schemes and aid physicists in performing physical simulations; however, the current state of the technology is largely experimental and impractical, with several obstacles to useful applications. Moreover, scalable quantum computers do not hold promise for many practical tasks, and for many important tasks quantum speedups are proven impossible.
-    The basic unit of information in quantum computing is the qubit, similar to the bit in traditional digital electronics. Unlike a classical bit, a qubit can exist in a superposition of its two "basis" states. When measuring a qubit, the result is a probabilistic output of a classical bit, therefore making quantum computers nondeterministic in general. If a quantum computer manipulates the qubit in a particular way, wave interference effects can amplify the desired measurement results. The design of quantum algorithms involves creating procedures that allow a quantum computer to perform calculations efficiently and quickly.
-    Physically engineering high-quality qubits has proven challenging. If a physical qubit is not sufficiently isolated from its environment, it suffers from quantum decoherence, introducing noise into calculations. Paradoxically, perfectly isolating qubits is also undesirable because quantum computations typically need to initialize qubits, perform controlled qubit interactions, and measure the resulting quantum states. Each of those operations introduces errors and suffers from noise, and such inaccuracies accumulate.
-    In principle, a non-quantum (classical) computer can solve the same computational problems as a quantum computer, given enough time. Quantum advantage comes in the form of time complexity rather than computability, and quantum complexity theory shows that some quantum algorithms for carefully selected tasks require exponentially fewer computational steps than the best known non-quantum algorithms. Such tasks can in theory be solved on a large-scale quantum computer whereas classical computers would not finish computations in any reasonable amount of time. However, quantum speedup is not universal or even typical across computational tasks, since basic tasks such as sorting are proven to not allow any asymptotic quantum speedup. Claims of quantum supremacy have drawn significant attention to the discipline, but are demonstrated on contrived tasks, while near-term practical use cases remain limited.
-    """
-
-    await cognee.add([text], dataset_name)
-
-    await cognee.cognify([dataset_name])
-
-    from cognee.infrastructure.databases.vector import get_vector_engine
-
-    vector_engine = get_vector_engine()
-    random_node = (await vector_engine.search("Entity_name", "Quantum computer"))[0]
-    random_node_name = random_node.payload["text"]
-
-    search_results = await cognee.search(
-        query_type=SearchType.INSIGHTS, query_text=random_node_name
-    )
-    assert len(search_results) != 0, "The search results list is empty."
-    print("\n\nExtracted sentences are:\n")
-    for result in search_results:
-        print(f"{result}\n")
-
-    search_results = await cognee.search(query_type=SearchType.CHUNKS, query_text=random_node_name)
-    assert len(search_results) != 0, "The search results list is empty."
-    print("\n\nExtracted chunks are:\n")
-    for result in search_results:
-        print(f"{result}\n")
-
-    search_results = await cognee.search(
-        query_type=SearchType.SUMMARIES, query_text=random_node_name
-    )
-    assert len(search_results) != 0, "Query related summaries don't exist."
-    print("\nExtracted results are:\n")
-    for result in search_results:
-        print(f"{result}\n")
-
-    search_results = await cognee.search(
-        query_type=SearchType.NATURAL_LANGUAGE,
-        query_text=f"Find nodes connected to node with name {random_node_name}",
-    )
-    assert len(search_results) != 0, "Query related natural language don't exist."
-    print("\nExtracted results are:\n")
-    for result in search_results:
-        print(f"{result}\n")
-
-    user = await get_default_user()
-    history = await get_history(user.id)
-
-    assert len(history) == 8, "Search history is not correct."
-
-    await cognee.prune.prune_data()
-    data_root_directory = get_storage_config()["data_root_directory"]
-    assert not os.path.isdir(data_root_directory), "Local data files are not deleted"
-
-    await cognee.prune.prune_system(metadata=True)
-    from cognee.infrastructure.databases.graph import get_graph_engine
-
-    graph_engine = await get_graph_engine()
-    nodes, edges = await graph_engine.get_graph_data()
-    assert len(nodes) == 0 and len(edges) == 0, "Memgraph graph database is not empty"
-
-
-if __name__ == "__main__":
-    import asyncio
-
-    asyncio.run(main())

cognee/tests/unit/modules/retrieval/insights_retriever_test.py

@@ -1,251 +0,0 @@
-import os
-import pytest
-import pathlib
-
-import cognee
-from cognee.low_level import setup
-from cognee.tasks.storage import add_data_points
-from cognee.modules.engine.models import Entity, EntityType
-from cognee.infrastructure.databases.graph import get_graph_engine
-from cognee.infrastructure.databases.vector import get_vector_engine
-from cognee.modules.retrieval.exceptions.exceptions import NoDataError
-from cognee.modules.retrieval.insights_retriever import InsightsRetriever
-
-
-class TestInsightsRetriever:
-    @pytest.mark.asyncio
-    async def test_insights_context_simple(self):
-        system_directory_path = os.path.join(
-            pathlib.Path(__file__).parent, ".cognee_system/test_insights_context_simple"
-        )
-        cognee.config.system_root_directory(system_directory_path)
-        data_directory_path = os.path.join(
-            pathlib.Path(__file__).parent, ".data_storage/test_insights_context_simple"
-        )
-        cognee.config.data_root_directory(data_directory_path)
-
-        await cognee.prune.prune_data()
-        await cognee.prune.prune_system(metadata=True)
-        await setup()
-
-        entityTypePerson = EntityType(
-            name="Person",
-            description="An individual",
-        )
-
-        person1 = Entity(
-            name="Steve Rodger",
-            is_a=entityTypePerson,
-            description="An American actor, comedian, and filmmaker",
-        )
-
-        person2 = Entity(
-            name="Mike Broski",
-            is_a=entityTypePerson,
-            description="Financial advisor and philanthropist",
-        )
-
-        person3 = Entity(
-            name="Christina Mayer",
-            is_a=entityTypePerson,
-            description="Maker of next generation of iconic American music videos",
-        )
-
-        entityTypeCompany = EntityType(
-            name="Company",
-            description="An organization that operates on an annual basis",
-        )
-
-        company1 = Entity(
-            name="Apple",
-            is_a=entityTypeCompany,
-            description="An American multinational technology company headquartered in Cupertino, California",
-        )
-
-        company2 = Entity(
-            name="Google",
-            is_a=entityTypeCompany,
-            description="An American multinational technology company that specializes in Internet-related services and products",
-        )
-
-        company3 = Entity(
-            name="Facebook",
-            is_a=entityTypeCompany,
-            description="An American social media, messaging, and online platform",
-        )
-
-        entities = [person1, person2, person3, company1, company2, company3]
-
-        await add_data_points(entities)
-
-        retriever = InsightsRetriever()
-
-        context = await retriever.get_context("Mike")
-
-        assert context[0][0]["name"] == "Mike Broski", "Failed to get Mike Broski"
-
-    @pytest.mark.asyncio
-    async def test_insights_context_complex(self):
-        system_directory_path = os.path.join(
-            pathlib.Path(__file__).parent, ".cognee_system/test_insights_context_complex"
-        )
-        cognee.config.system_root_directory(system_directory_path)
-        data_directory_path = os.path.join(
-            pathlib.Path(__file__).parent, ".data_storage/test_insights_context_complex"
-        )
-        cognee.config.data_root_directory(data_directory_path)
-
-        await cognee.prune.prune_data()
-        await cognee.prune.prune_system(metadata=True)
-        await setup()
-
-        entityTypePerson = EntityType(
-            name="Person",
-            description="An individual",
-        )
-
-        person1 = Entity(
-            name="Steve Rodger",
-            is_a=entityTypePerson,
-            description="An American actor, comedian, and filmmaker",
-        )
-
-        person2 = Entity(
-            name="Mike Broski",
-            is_a=entityTypePerson,
-            description="Financial advisor and philanthropist",
-        )
-
-        person3 = Entity(
-            name="Christina Mayer",
-            is_a=entityTypePerson,
-            description="Maker of next generation of iconic American music videos",
-        )
-
-        person4 = Entity(
-            name="Jason Statham",
-            is_a=entityTypePerson,
-            description="An American actor",
-        )
-
-        person5 = Entity(
-            name="Mike Tyson",
-            is_a=entityTypePerson,
-            description="A former professional boxer from the United States",
-        )
-
-        entityTypeCompany = EntityType(
-            name="Company",
-            description="An organization that operates on an annual basis",
-        )
-
-        company1 = Entity(
-            name="Apple",
-            is_a=entityTypeCompany,
-            description="An American multinational technology company headquartered in Cupertino, California",
-        )
-
-        company2 = Entity(
-            name="Google",
-            is_a=entityTypeCompany,
-            description="An American multinational technology company that specializes in Internet-related services and products",
-        )
-
-        company3 = Entity(
-            name="Facebook",
-            is_a=entityTypeCompany,
-            description="An American social media, messaging, and online platform",
-        )
-
-        entities = [person1, person2, person3, company1, company2, company3]
-
-        await add_data_points(entities)
-
-        graph_engine = await get_graph_engine()
-
-        await graph_engine.add_edges(
-            [
-                (
-                    (str)(person1.id),
-                    (str)(company1.id),
-                    "works_for",
-                    dict(
-                        relationship_name="works_for",
-                        source_node_id=person1.id,
-                        target_node_id=company1.id,
-                    ),
-                ),
-                (
-                    (str)(person2.id),
-                    (str)(company2.id),
-                    "works_for",
-                    dict(
-                        relationship_name="works_for",
-                        source_node_id=person2.id,
-                        target_node_id=company2.id,
-                    ),
-                ),
-                (
-                    (str)(person3.id),
-                    (str)(company3.id),
-                    "works_for",
-                    dict(
-                        relationship_name="works_for",
-                        source_node_id=person3.id,
-                        target_node_id=company3.id,
-                    ),
-                ),
-                (
-                    (str)(person4.id),
-                    (str)(company1.id),
-                    "works_for",
-                    dict(
-                        relationship_name="works_for",
-                        source_node_id=person4.id,
-                        target_node_id=company1.id,
-                    ),
-                ),
-                (
-                    (str)(person5.id),
-                    (str)(company1.id),
-                    "works_for",
-                    dict(
-                        relationship_name="works_for",
-                        source_node_id=person5.id,
-                        target_node_id=company1.id,
-                    ),
-                ),
-            ]
-        )
-
-        retriever = InsightsRetriever(top_k=20)
-
-        context = await retriever.get_context("Christina")
-
-        assert context[0][0]["name"] == "Christina Mayer", "Failed to get Christina Mayer"
-
-    @pytest.mark.asyncio
-    async def test_insights_context_on_empty_graph(self):
-        system_directory_path = os.path.join(
-            pathlib.Path(__file__).parent, ".cognee_system/test_insights_context_on_empty_graph"
-        )
-        cognee.config.system_root_directory(system_directory_path)
-        data_directory_path = os.path.join(
-            pathlib.Path(__file__).parent, ".data_storage/test_insights_context_on_empty_graph"
-        )
-        cognee.config.data_root_directory(data_directory_path)
-
-        await cognee.prune.prune_data()
-        await cognee.prune.prune_system(metadata=True)
-
-        retriever = InsightsRetriever()
-
-        with pytest.raises(NoDataError):
-            await retriever.get_context("Christina Mayer")
-
-        vector_engine = get_vector_engine()
-        await vector_engine.create_collection("Entity_name", payload_schema=Entity)
-        await vector_engine.create_collection("EntityType_name", payload_schema=EntityType)
-
-        context = await retriever.get_context("Christina Mayer")
-        assert context == [], "Returned context should be empty on an empty graph"