cognee 0.5.1__py3-none-any.whl → 0.5.2__py3-none-any.whl

cognee/api/.env.example

@@ -1,5 +0,0 @@
-# Comma-separated list of allowed origins for CORS (leave empty to block all cross-origin requests)
-# Example:
-# CORS_ALLOWED_ORIGINS="https://yourdomain.com,https://another.com"
-# For local development, you might use:
-# CORS_ALLOWED_ORIGINS="http://localhost:3000"

cognee/modules/retrieval/base_graph_retriever.py

@@ -1,24 +0,0 @@
-from typing import Any, List, Optional, Type
-from abc import ABC, abstractmethod
-
-from cognee.modules.graph.cognee_graph.CogneeGraphElements import Edge
-
-
-class BaseGraphRetriever(ABC):
-    """Base class for all graph based retrievers."""
-
-    @abstractmethod
-    async def get_context(self, query: str) -> List[Edge]:
-        """Retrieves triplets based on the query."""
-        pass
-
-    @abstractmethod
-    async def get_completion(
-        self,
-        query: str,
-        context: Optional[List[Edge]] = None,
-        session_id: Optional[str] = None,
-        response_model: Type = str,
-    ) -> List[Any]:
-        """Generates a response using the query and optional context (triplets)."""
-        pass

cognee/modules/search/methods/get_search_type_tools.py

@@ -1,223 +0,0 @@
-import os
-from typing import Callable, List, Optional, Type
-
-from cognee.modules.engine.models.node_set import NodeSet
-from cognee.modules.retrieval.triplet_retriever import TripletRetriever
-from cognee.modules.search.types import SearchType
-from cognee.modules.search.operations import select_search_type
-from cognee.modules.search.exceptions import UnsupportedSearchTypeError
-
-# Retrievers
-from cognee.modules.retrieval.user_qa_feedback import UserQAFeedback
-from cognee.modules.retrieval.chunks_retriever import ChunksRetriever
-from cognee.modules.retrieval.summaries_retriever import SummariesRetriever
-from cognee.modules.retrieval.completion_retriever import CompletionRetriever
-from cognee.modules.retrieval.graph_completion_retriever import GraphCompletionRetriever
-from cognee.modules.retrieval.temporal_retriever import TemporalRetriever
-from cognee.modules.retrieval.coding_rules_retriever import CodingRulesRetriever
-from cognee.modules.retrieval.jaccard_retrival import JaccardChunksRetriever
-from cognee.modules.retrieval.graph_summary_completion_retriever import (
-    GraphSummaryCompletionRetriever,
-)
-from cognee.modules.retrieval.graph_completion_cot_retriever import GraphCompletionCotRetriever
-from cognee.modules.retrieval.graph_completion_context_extension_retriever import (
-    GraphCompletionContextExtensionRetriever,
-)
-from cognee.modules.retrieval.cypher_search_retriever import CypherSearchRetriever
-from cognee.modules.retrieval.natural_language_retriever import NaturalLanguageRetriever
-
-
-async def get_search_type_tools(
-    query_type: SearchType,
-    query_text: str,
-    system_prompt_path: str = "answer_simple_question.txt",
-    system_prompt: Optional[str] = None,
-    top_k: int = 10,
-    node_type: Optional[Type] = NodeSet,
-    node_name: Optional[List[str]] = None,
-    save_interaction: bool = False,
-    last_k: Optional[int] = None,
-    wide_search_top_k: Optional[int] = 100,
-    triplet_distance_penalty: Optional[float] = 3.5,
-) -> list:
-    search_tasks: dict[SearchType, List[Callable]] = {
-        SearchType.SUMMARIES: [
-            SummariesRetriever(top_k=top_k).get_completion,
-            SummariesRetriever(top_k=top_k).get_context,
-        ],
-        SearchType.CHUNKS: [
-            ChunksRetriever(top_k=top_k).get_completion,
-            ChunksRetriever(top_k=top_k).get_context,
-        ],
-        SearchType.RAG_COMPLETION: [
-            CompletionRetriever(
-                system_prompt_path=system_prompt_path,
-                top_k=top_k,
-                system_prompt=system_prompt,
-            ).get_completion,
-            CompletionRetriever(
-                system_prompt_path=system_prompt_path,
-                top_k=top_k,
-                system_prompt=system_prompt,
-            ).get_context,
-        ],
-        SearchType.TRIPLET_COMPLETION: [
-            TripletRetriever(
-                system_prompt_path=system_prompt_path,
-                top_k=top_k,
-                system_prompt=system_prompt,
-            ).get_completion,
-            TripletRetriever(
-                system_prompt_path=system_prompt_path,
-                top_k=top_k,
-                system_prompt=system_prompt,
-            ).get_context,
-        ],
-        SearchType.GRAPH_COMPLETION: [
-            GraphCompletionRetriever(
-                system_prompt_path=system_prompt_path,
-                top_k=top_k,
-                node_type=node_type,
-                node_name=node_name,
-                save_interaction=save_interaction,
-                system_prompt=system_prompt,
-                wide_search_top_k=wide_search_top_k,
-                triplet_distance_penalty=triplet_distance_penalty,
-            ).get_completion,
-            GraphCompletionRetriever(
-                system_prompt_path=system_prompt_path,
-                top_k=top_k,
-                node_type=node_type,
-                node_name=node_name,
-                save_interaction=save_interaction,
-                system_prompt=system_prompt,
-                wide_search_top_k=wide_search_top_k,
-                triplet_distance_penalty=triplet_distance_penalty,
-            ).get_context,
-        ],
-        SearchType.GRAPH_COMPLETION_COT: [
-            GraphCompletionCotRetriever(
-                system_prompt_path=system_prompt_path,
-                top_k=top_k,
-                node_type=node_type,
-                node_name=node_name,
-                save_interaction=save_interaction,
-                system_prompt=system_prompt,
-                wide_search_top_k=wide_search_top_k,
-                triplet_distance_penalty=triplet_distance_penalty,
-            ).get_completion,
-            GraphCompletionCotRetriever(
-                system_prompt_path=system_prompt_path,
-                top_k=top_k,
-                node_type=node_type,
-                node_name=node_name,
-                save_interaction=save_interaction,
-                system_prompt=system_prompt,
-                wide_search_top_k=wide_search_top_k,
-                triplet_distance_penalty=triplet_distance_penalty,
-            ).get_context,
-        ],
-        SearchType.GRAPH_COMPLETION_CONTEXT_EXTENSION: [
-            GraphCompletionContextExtensionRetriever(
-                system_prompt_path=system_prompt_path,
-                top_k=top_k,
-                node_type=node_type,
-                node_name=node_name,
-                save_interaction=save_interaction,
-                system_prompt=system_prompt,
-                wide_search_top_k=wide_search_top_k,
-                triplet_distance_penalty=triplet_distance_penalty,
-            ).get_completion,
-            GraphCompletionContextExtensionRetriever(
-                system_prompt_path=system_prompt_path,
-                top_k=top_k,
-                node_type=node_type,
-                node_name=node_name,
-                save_interaction=save_interaction,
-                system_prompt=system_prompt,
-                wide_search_top_k=wide_search_top_k,
-                triplet_distance_penalty=triplet_distance_penalty,
-            ).get_context,
-        ],
-        SearchType.GRAPH_SUMMARY_COMPLETION: [
-            GraphSummaryCompletionRetriever(
-                system_prompt_path=system_prompt_path,
-                top_k=top_k,
-                node_type=node_type,
-                node_name=node_name,
-                save_interaction=save_interaction,
-                system_prompt=system_prompt,
-                wide_search_top_k=wide_search_top_k,
-                triplet_distance_penalty=triplet_distance_penalty,
-            ).get_completion,
-            GraphSummaryCompletionRetriever(
-                system_prompt_path=system_prompt_path,
-                top_k=top_k,
-                node_type=node_type,
-                node_name=node_name,
-                save_interaction=save_interaction,
-                system_prompt=system_prompt,
-                wide_search_top_k=wide_search_top_k,
-                triplet_distance_penalty=triplet_distance_penalty,
-            ).get_context,
-        ],
-        SearchType.CYPHER: [
-            CypherSearchRetriever().get_completion,
-            CypherSearchRetriever().get_context,
-        ],
-        SearchType.NATURAL_LANGUAGE: [
-            NaturalLanguageRetriever().get_completion,
-            NaturalLanguageRetriever().get_context,
-        ],
-        SearchType.FEEDBACK: [UserQAFeedback(last_k=last_k).add_feedback],
-        SearchType.TEMPORAL: [
-            TemporalRetriever(
-                top_k=top_k,
-                wide_search_top_k=wide_search_top_k,
-                triplet_distance_penalty=triplet_distance_penalty,
-            ).get_completion,
-            TemporalRetriever(
-                top_k=top_k,
-                wide_search_top_k=wide_search_top_k,
-                triplet_distance_penalty=triplet_distance_penalty,
-            ).get_context,
-        ],
-        SearchType.CHUNKS_LEXICAL: (
-            lambda _r=JaccardChunksRetriever(top_k=top_k): [
-                _r.get_completion,
-                _r.get_context,
-            ]
-        )(),
-        SearchType.CODING_RULES: [
-            CodingRulesRetriever(rules_nodeset_name=node_name).get_existing_rules,
-        ],
-    }
-
-    # If the query type is FEELING_LUCKY, select the search type intelligently
-    if query_type is SearchType.FEELING_LUCKY:
-        query_type = await select_search_type(query_text)
-
-    if (
-        query_type in [SearchType.CYPHER, SearchType.NATURAL_LANGUAGE]
-        and os.getenv("ALLOW_CYPHER_QUERY", "true").lower() == "false"
-    ):
-        raise UnsupportedSearchTypeError("Cypher query search types are disabled.")
-
-    from cognee.modules.retrieval.registered_community_retrievers import (
-        registered_community_retrievers,
-    )
-
-    if query_type in registered_community_retrievers:
-        retriever = registered_community_retrievers[query_type]
-        retriever_instance = retriever(top_k=top_k)
-        search_type_tools = [
-            retriever_instance.get_completion,
-            retriever_instance.get_context,
-        ]
-    else:
-        search_type_tools = search_tasks.get(query_type)
-
-    if not search_type_tools:
-        raise UnsupportedSearchTypeError(str(query_type))
-
-    return search_type_tools

cognee/modules/search/methods/no_access_control_search.py

@@ -1,62 +0,0 @@
-from typing import Any, List, Optional, Tuple, Type, Union
-
-from cognee.infrastructure.databases.graph import get_graph_engine
-from cognee.modules.data.models.Dataset import Dataset
-from cognee.modules.engine.models.node_set import NodeSet
-from cognee.modules.graph.cognee_graph.CogneeGraphElements import Edge
-from cognee.modules.search.types import SearchType
-from cognee.shared.logging_utils import get_logger
-
-from .get_search_type_tools import get_search_type_tools
-
-logger = get_logger()
-
-
-async def no_access_control_search(
-    query_type: SearchType,
-    query_text: str,
-    system_prompt_path: str = "answer_simple_question.txt",
-    system_prompt: Optional[str] = None,
-    top_k: int = 10,
-    node_type: Optional[Type] = NodeSet,
-    node_name: Optional[List[str]] = None,
-    save_interaction: bool = False,
-    last_k: Optional[int] = None,
-    only_context: bool = False,
-    session_id: Optional[str] = None,
-    wide_search_top_k: Optional[int] = 100,
-    triplet_distance_penalty: Optional[float] = 3.5,
-) -> Tuple[Any, Union[str, List[Edge]], List[Dataset]]:
-    search_tools = await get_search_type_tools(
-        query_type=query_type,
-        query_text=query_text,
-        system_prompt_path=system_prompt_path,
-        system_prompt=system_prompt,
-        top_k=top_k,
-        node_type=node_type,
-        node_name=node_name,
-        save_interaction=save_interaction,
-        last_k=last_k,
-        wide_search_top_k=wide_search_top_k,
-        triplet_distance_penalty=triplet_distance_penalty,
-    )
-    graph_engine = await get_graph_engine()
-    is_empty = await graph_engine.is_empty()
-
-    if is_empty:
-        # TODO: we can log here, but not all search types use graph. Still keeping this here for reviewer input
-        logger.warning("Search attempt on an empty knowledge graph")
-    if len(search_tools) == 2:
-        [get_completion, get_context] = search_tools
-
-        if only_context:
-            return None, await get_context(query_text), []
-
-        context = await get_context(query_text)
-        result = await get_completion(query_text, context, session_id=session_id)
-    else:
-        unknown_tool = search_tools[0]
-        result = await unknown_tool(query_text)
-        context = ""
-
-    return result, context, []

cognee/modules/search/utils/prepare_search_result.py

@@ -1,63 +0,0 @@
-from typing import List, cast
-from uuid import uuid5, NAMESPACE_OID
-
-from cognee.modules.graph.utils import resolve_edges_to_text
-from cognee.modules.graph.cognee_graph.CogneeGraphElements import Edge
-from cognee.modules.search.types.SearchResult import SearchResultDataset
-from cognee.modules.search.utils.transform_context_to_graph import transform_context_to_graph
-from cognee.modules.search.utils.transform_insights_to_graph import transform_insights_to_graph
-
-
-async def prepare_search_result(search_result):
-    results, context, datasets = search_result
-
-    graphs = None
-    result_graph = None
-    context_texts = {}
-
-    if isinstance(datasets, list) and len(datasets) == 0:
-        datasets = [
-            SearchResultDataset(
-                id=uuid5(NAMESPACE_OID, "*"),
-                name="all available datasets",
-            )
-        ]
-
-    if (
-        isinstance(context, List)
-        and len(context) > 0
-        and isinstance(context[0], tuple)
-        and context[0][1].get("relationship_name")
-    ):
-        context_graph = transform_insights_to_graph(context)
-        graphs = {
-            ", ".join([dataset.name for dataset in datasets]): context_graph,
-        }
-        results = None
-    elif isinstance(context, List) and len(context) > 0 and isinstance(context[0], Edge):
-        context_graph = transform_context_to_graph(context)
-
-        graphs = {
-            ", ".join([dataset.name for dataset in datasets]): context_graph,
-        }
-        context_texts = {
-            ", ".join([dataset.name for dataset in datasets]): await resolve_edges_to_text(context),
-        }
-    elif isinstance(context, str):
-        context_texts = {
-            ", ".join([dataset.name for dataset in datasets]): context,
-        }
-    elif isinstance(context, List) and len(context) > 0 and isinstance(context[0], str):
-        context_texts = {
-            ", ".join([dataset.name for dataset in datasets]): "\n".join(cast(List[str], context)),
-        }
-
-    if isinstance(results, List) and len(results) > 0 and isinstance(results[0], Edge):
-        result_graph = transform_context_to_graph(results)
-
-    return {
-        "result": result_graph or results[0] if results and len(results) == 1 else results,
-        "graphs": graphs,
-        "context": context_texts,
-        "datasets": datasets,
-    }

cognee/tests/test_feedback_enrichment.py

@@ -1,174 +0,0 @@
-"""
-End-to-end integration test for feedback enrichment feature.
-
-Tests the complete feedback enrichment pipeline:
-1. Add data and cognify
-2. Run search with save_interaction=True to create CogneeUserInteraction nodes
-3. Submit feedback to create CogneeUserFeedback nodes
-4. Run memify with feedback enrichment tasks to create FeedbackEnrichment nodes
-5. Verify all nodes and edges are properly created and linked in the graph
-"""
-
-import os
-import pathlib
-from collections import Counter
-
-import cognee
-from cognee.infrastructure.databases.graph import get_graph_engine
-from cognee.modules.pipelines.tasks.task import Task
-from cognee.modules.search.types import SearchType
-from cognee.shared.data_models import KnowledgeGraph
-from cognee.shared.logging_utils import get_logger
-from cognee.tasks.feedback.create_enrichments import create_enrichments
-from cognee.tasks.feedback.extract_feedback_interactions import (
-    extract_feedback_interactions,
-)
-from cognee.tasks.feedback.generate_improved_answers import generate_improved_answers
-from cognee.tasks.feedback.link_enrichments_to_feedback import (
-    link_enrichments_to_feedback,
-)
-from cognee.tasks.graph import extract_graph_from_data
-from cognee.tasks.storage import add_data_points
-
-logger = get_logger()
-
-
-async def main():
-    data_directory_path = str(
-        pathlib.Path(
-            os.path.join(
-                pathlib.Path(__file__).parent,
-                ".data_storage/test_feedback_enrichment",
-            )
-        ).resolve()
-    )
-    cognee_directory_path = str(
-        pathlib.Path(
-            os.path.join(
-                pathlib.Path(__file__).parent,
-                ".cognee_system/test_feedback_enrichment",
-            )
-        ).resolve()
-    )
-
-    cognee.config.data_root_directory(data_directory_path)
-    cognee.config.system_root_directory(cognee_directory_path)
-
-    await cognee.prune.prune_data()
-    await cognee.prune.prune_system(metadata=True)
-
-    dataset_name = "feedback_enrichment_test"
-
-    await cognee.add("Cognee turns documents into AI memory.", dataset_name)
-    await cognee.cognify([dataset_name])
-
-    question_text = "Say something."
-    result = await cognee.search(
-        query_type=SearchType.GRAPH_COMPLETION,
-        query_text=question_text,
-        save_interaction=True,
-    )
-
-    assert len(result) > 0, "Search should return non-empty results"
-
-    feedback_text = "This answer was completely useless, my feedback is definitely negative."
-    await cognee.search(
-        query_type=SearchType.FEEDBACK,
-        query_text=feedback_text,
-        last_k=1,
-    )
-
-    graph_engine = await get_graph_engine()
-    nodes_before, edges_before = await graph_engine.get_graph_data()
-
-    interaction_nodes_before = [
-        (node_id, props)
-        for node_id, props in nodes_before
-        if props.get("type") == "CogneeUserInteraction"
-    ]
-    feedback_nodes_before = [
-        (node_id, props)
-        for node_id, props in nodes_before
-        if props.get("type") == "CogneeUserFeedback"
-    ]
-
-    edge_types_before = Counter(edge[2] for edge in edges_before)
-
-    assert len(interaction_nodes_before) >= 1, (
-        f"Expected at least 1 CogneeUserInteraction node, found {len(interaction_nodes_before)}"
-    )
-    assert len(feedback_nodes_before) >= 1, (
-        f"Expected at least 1 CogneeUserFeedback node, found {len(feedback_nodes_before)}"
-    )
-
-    for node_id, props in feedback_nodes_before:
-        sentiment = props.get("sentiment", "")
-        score = props.get("score", 0)
-        feedback_text = props.get("feedback", "")
-        logger.info(
-            "Feedback node created",
-            feedback=feedback_text,
-            sentiment=sentiment,
-            score=score,
-        )
-
-    assert edge_types_before.get("gives_feedback_to", 0) >= 1, (
-        f"Expected at least 1 'gives_feedback_to' edge, found {edge_types_before.get('gives_feedback_to', 0)}"
-    )
-
-    extraction_tasks = [Task(extract_feedback_interactions, last_n=5)]
-    enrichment_tasks = [
-        Task(generate_improved_answers, top_k=20),
-        Task(create_enrichments),
-        Task(
-            extract_graph_from_data,
-            graph_model=KnowledgeGraph,
-            task_config={"batch_size": 10},
-        ),
-        Task(add_data_points, task_config={"batch_size": 10}),
-        Task(link_enrichments_to_feedback),
-    ]
-
-    await cognee.memify(
-        extraction_tasks=extraction_tasks,
-        enrichment_tasks=enrichment_tasks,
-        data=[{}],
-        dataset=dataset_name,
-    )
-
-    nodes_after, edges_after = await graph_engine.get_graph_data()
-
-    enrichment_nodes = [
-        (node_id, props)
-        for node_id, props in nodes_after
-        if props.get("type") == "FeedbackEnrichment"
-    ]
-
-    assert len(enrichment_nodes) >= 1, (
-        f"Expected at least 1 FeedbackEnrichment node, found {len(enrichment_nodes)}"
-    )
-
-    for node_id, props in enrichment_nodes:
-        assert "text" in props, f"FeedbackEnrichment node {node_id} missing 'text' property"
-
-    enrichment_node_ids = {node_id for node_id, _ in enrichment_nodes}
-    edges_with_enrichments = [
-        edge
-        for edge in edges_after
-        if edge[0] in enrichment_node_ids or edge[1] in enrichment_node_ids
-    ]
-
-    assert len(edges_with_enrichments) >= 1, (
-        f"Expected enrichment nodes to have at least 1 edge, found {len(edges_with_enrichments)}"
-    )
-
-    await cognee.prune.prune_data()
-    await cognee.prune.prune_system(metadata=True)
-
-    logger.info("All feedback enrichment tests passed successfully")
-
-
-if __name__ == "__main__":
-    import asyncio
-
-    asyncio.run(main())