cognee 0.2.4__py3-none-any.whl → 0.3.0__py3-none-any.whl

cognee/modules/search/methods/get_search_type_tools.py

@@ -0,0 +1,168 @@
+from typing import Callable, List, Optional, Type
+
+from cognee.modules.engine.models.node_set import NodeSet
+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.insights_retriever import InsightsRetriever
+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.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.code_retriever import CodeRetriever
+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,
+) -> list:
+    search_tasks: dict[SearchType, List[Callable]] = {
+        SearchType.SUMMARIES: [
+            SummariesRetriever(top_k=top_k).get_completion,
+            SummariesRetriever(top_k=top_k).get_context,
+        ],
+        SearchType.INSIGHTS: [
+            InsightsRetriever(top_k=top_k).get_completion,
+            InsightsRetriever(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.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,
+            ).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,
+            ).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,
+            ).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,
+            ).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,
+            ).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,
+            ).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,
+            ).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,
+            ).get_context,
+        ],
+        SearchType.CODE: [
+            CodeRetriever(top_k=top_k).get_completion,
+            CodeRetriever(top_k=top_k).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).get_completion,
+            TemporalRetriever(top_k=top_k).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)
+
+    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

@@ -0,0 +1,47 @@
+from typing import Any, List, Optional, Tuple, Type, Union
+
+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 .get_search_type_tools import get_search_type_tools
+
+
+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,
+) -> 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,
+    )
+    if len(search_tools) == 2:
+        [get_completion, get_context] = search_tools
+
+        if only_context:
+            return await get_context(query_text)
+
+        context = await get_context(query_text)
+        result = await get_completion(query_text, context)
+    else:
+        unknown_tool = search_tools[0]
+        result = await unknown_tool(query_text)
+        context = ""
+
+    return result, context, []

cognee/modules/search/methods/search.py

@@ -2,33 +2,28 @@ import os
 import json
 import asyncio
 from uuid import UUID
-from typing import Callable, List, Optional, Type, Union
+from fastapi.encoders import jsonable_encoder
+from typing import Any, List, Optional, Tuple, Type, Union
 
-from cognee.modules.retrieval.user_qa_feedback import UserQAFeedback
-from cognee.modules.search.exceptions import UnsupportedSearchTypeError
+from cognee.shared.utils import send_telemetry
 from cognee.context_global_variables import set_database_global_context_variables
-from cognee.modules.retrieval.chunks_retriever import ChunksRetriever
-from cognee.modules.retrieval.insights_retriever import InsightsRetriever
-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.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.engine.models.node_set import NodeSet
+from cognee.modules.graph.cognee_graph.CogneeGraphElements import Edge
+from cognee.modules.search.types import (
+    SearchResult,
+    CombinedSearchResult,
+    SearchResultDataset,
+    SearchType,
 )
-from cognee.modules.retrieval.code_retriever import CodeRetriever
-from cognee.modules.retrieval.cypher_search_retriever import CypherSearchRetriever
-from cognee.modules.retrieval.natural_language_retriever import NaturalLanguageRetriever
-from cognee.modules.search.types import SearchType
-from cognee.modules.storage.utils import JSONEncoder
+from cognee.modules.search.operations import log_query, log_result
 from cognee.modules.users.models import User
 from cognee.modules.data.models import Dataset
-from cognee.shared.utils import send_telemetry
 from cognee.modules.users.permissions.methods import get_specific_user_permission_datasets
-from cognee.modules.search.operations import log_query, log_result, select_search_type
+
+from .get_search_type_tools import get_search_type_tools
+from .no_access_control_search import no_access_control_search
+from ..utils.prepare_search_result import prepare_search_result
 
 
 async def search(
@@ -37,12 +32,15 @@ async def search(
     dataset_ids: Union[list[UUID], None],
     user: User,
     system_prompt_path="answer_simple_question.txt",
+    system_prompt: Optional[str] = None,
     top_k: int = 10,
-    node_type: Optional[Type] = None,
+    node_type: Optional[Type] = NodeSet,
     node_name: Optional[List[str]] = None,
-    save_interaction: Optional[bool] = False,
+    save_interaction: bool = False,
     last_k: Optional[int] = None,
-):
+    only_context: bool = False,
+    use_combined_context: bool = False,
+) -> Union[CombinedSearchResult, List[SearchResult]]:
     """
 
     Args:
@@ -58,192 +56,274 @@ async def search(
     Notes:
         Searching by dataset is only available in ENABLE_BACKEND_ACCESS_CONTROL mode
     """
+    query = await log_query(query_text, query_type.value, user.id)
+    send_telemetry("cognee.search EXECUTION STARTED", user.id)
+
     # Use search function filtered by permissions if access control is enabled
     if os.getenv("ENABLE_BACKEND_ACCESS_CONTROL", "false").lower() == "true":
-        return await authorized_search(
-            query_text=query_text,
+        search_results = await authorized_search(
             query_type=query_type,
+            query_text=query_text,
             user=user,
             dataset_ids=dataset_ids,
             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,
+            only_context=only_context,
+            use_combined_context=use_combined_context,
         )
+    else:
+        search_results = [
+            await no_access_control_search(
+                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,
+                only_context=only_context,
+            )
+        ]
 
-    query = await log_query(query_text, query_type.value, user.id)
-
-    search_results = await specific_search(
-        query_type,
-        query_text,
-        user,
-        system_prompt_path=system_prompt_path,
-        top_k=top_k,
-        node_type=node_type,
-        node_name=node_name,
-        save_interaction=save_interaction,
-        last_k=last_k,
-    )
+    send_telemetry("cognee.search EXECUTION COMPLETED", user.id)
 
     await log_result(
         query.id,
         json.dumps(
-            search_results if len(search_results) > 1 else search_results[0], cls=JSONEncoder
+            jsonable_encoder(
+                await prepare_search_result(
+                    search_results[0] if isinstance(search_results, list) else search_results
+                )
+                if use_combined_context
+                else [
+                    await prepare_search_result(search_result) for search_result in search_results
+                ]
+            )
         ),
         user.id,
     )
 
-    return search_results
+    if use_combined_context:
+        prepared_search_results = await prepare_search_result(
+            search_results[0] if isinstance(search_results, list) else search_results
+        )
+        result = prepared_search_results["result"]
+        graphs = prepared_search_results["graphs"]
+        context = prepared_search_results["context"]
+        datasets = prepared_search_results["datasets"]
 
+        return CombinedSearchResult(
+            result=result,
+            graphs=graphs,
+            context=context,
+            datasets=[
+                SearchResultDataset(
+                    id=dataset.id,
+                    name=dataset.name,
+                )
+                for dataset in datasets
+            ],
+        )
+    else:
+        return [
+            SearchResult(
+                search_result=result,
+                dataset_id=datasets[min(index, len(datasets) - 1)].id if datasets else None,
+                dataset_name=datasets[min(index, len(datasets) - 1)].name if datasets else None,
+            )
+            for index, (result, _, datasets) in enumerate(search_results)
+        ]
 
-async def specific_search(
+
+async def authorized_search(
     query_type: SearchType,
-    query: str,
+    query_text: str,
     user: User,
-    system_prompt_path="answer_simple_question.txt",
+    dataset_ids: Optional[list[UUID]] = None,
+    system_prompt_path: str = "answer_simple_question.txt",
+    system_prompt: Optional[str] = None,
     top_k: int = 10,
-    node_type: Optional[Type] = None,
+    node_type: Optional[Type] = NodeSet,
     node_name: Optional[List[str]] = None,
-    save_interaction: Optional[bool] = False,
+    save_interaction: bool = False,
     last_k: Optional[int] = None,
-) -> list:
-    search_tasks: dict[SearchType, Callable] = {
-        SearchType.SUMMARIES: SummariesRetriever(top_k=top_k).get_completion,
-        SearchType.INSIGHTS: InsightsRetriever(top_k=top_k).get_completion,
-        SearchType.CHUNKS: ChunksRetriever(top_k=top_k).get_completion,
-        SearchType.RAG_COMPLETION: CompletionRetriever(
-            system_prompt_path=system_prompt_path, top_k=top_k
-        ).get_completion,
-        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,
-        ).get_completion,
-        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,
-        ).get_completion,
-        SearchType.GRAPH_COMPLETION_CONTEXT_EXTENSION: GraphCompletionContextExtensionRetriever(
+    only_context: bool = False,
+    use_combined_context: bool = False,
+) -> Union[
+    Tuple[Any, Union[List[Edge], str], List[Dataset]],
+    List[Tuple[Any, Union[List[Edge], str], List[Dataset]]],
+]:
+    """
+    Verifies access for provided datasets or uses all datasets user has read access for and performs search per dataset.
+    Not to be used outside of active access control mode.
+    """
+    # Find datasets user has read access for (if datasets are provided only return them. Provided user has read access)
+    search_datasets = await get_specific_user_permission_datasets(user.id, "read", dataset_ids)
+
+    if use_combined_context:
+        search_responses = await search_in_datasets_context(
+            search_datasets=search_datasets,
+            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,
-        ).get_completion,
-        SearchType.GRAPH_SUMMARY_COMPLETION: GraphSummaryCompletionRetriever(
+            last_k=last_k,
+            only_context=True,
+        )
+
+        context = {}
+        datasets: List[Dataset] = []
+
+        for _, search_context, datasets in search_responses:
+            for dataset in datasets:
+                context[str(dataset.id)] = search_context
+
+            datasets.extend(datasets)
+
+        specific_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,
-        ).get_completion,
-        SearchType.CODE: CodeRetriever(top_k=top_k).get_completion,
-        SearchType.CYPHER: CypherSearchRetriever().get_completion,
-        SearchType.NATURAL_LANGUAGE: NaturalLanguageRetriever().get_completion,
-        SearchType.FEEDBACK: UserQAFeedback(last_k=last_k).add_feedback,
-    }
-
-    # 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)
-
-    search_task = search_tasks.get(query_type)
-
-    if search_task is None:
-        raise UnsupportedSearchTypeError(str(query_type))
-
-    send_telemetry("cognee.search EXECUTION STARTED", user.id)
-
-    results = await search_task(query)
+            last_k=last_k,
+        )
+        search_tools = specific_search_tools
+        if len(search_tools) == 2:
+            [get_completion, _] = search_tools
+        else:
+            get_completion = search_tools[0]
 
-    send_telemetry("cognee.search EXECUTION COMPLETED", user.id)
+        def prepare_combined_context(
+            context,
+        ) -> Union[List[Edge], str]:
+            combined_context = []
 
-    return results
+            for dataset_context in context.values():
+                combined_context += dataset_context
 
+            if combined_context and isinstance(combined_context[0], str):
+                return "\n".join(combined_context)
 
-async def authorized_search(
-    query_text: str,
-    query_type: SearchType,
-    user: User = None,
-    dataset_ids: Optional[list[UUID]] = None,
-    system_prompt_path: str = "answer_simple_question.txt",
-    top_k: int = 10,
-    save_interaction: bool = False,
-    last_k: Optional[int] = None,
-) -> list:
-    """
-    Verifies access for provided datasets or uses all datasets user has read access for and performs search per dataset.
-    Not to be used outside of active access control mode.
-    """
+            return combined_context
 
-    query = await log_query(query_text, query_type.value, user.id)
+        combined_context = prepare_combined_context(context)
+        completion = await get_completion(query_text, combined_context)
 
-    # Find datasets user has read access for (if datasets are provided only return them. Provided user has read access)
-    search_datasets = await get_specific_user_permission_datasets(user.id, "read", dataset_ids)
+        return completion, combined_context, datasets
 
     # Searches all provided datasets and handles setting up of appropriate database context based on permissions
-    search_results = await specific_search_by_context(
-        search_datasets,
-        query_text,
-        query_type,
-        user,
-        system_prompt_path,
-        top_k,
-        save_interaction,
+    search_results = await search_in_datasets_context(
+        search_datasets=search_datasets,
+        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,
+        only_context=only_context,
     )
 
-    await log_result(query.id, json.dumps(search_results, cls=JSONEncoder), user.id)
-
     return search_results
 
 
-async def specific_search_by_context(
+async def search_in_datasets_context(
     search_datasets: list[Dataset],
-    query_text: str,
     query_type: SearchType,
-    user: User,
-    system_prompt_path: str,
-    top_k: int,
+    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,
+    context: Optional[Any] = None,
+) -> List[Tuple[Any, Union[str, List[Edge]], List[Dataset]]]:
     """
     Searches all provided datasets and handles setting up of appropriate database context based on permissions.
     Not to be used outside of active access control mode.
     """
 
-    async def _search_by_context(
-        dataset, user, query_type, query_text, system_prompt_path, top_k, last_k
-    ):
+    async def _search_in_dataset_context(
+        dataset: Dataset,
+        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,
+        context: Optional[Any] = None,
+    ) -> Tuple[Any, Union[str, List[Edge]], List[Dataset]]:
         # Set database configuration in async context for each dataset user has access for
         await set_database_global_context_variables(dataset.id, dataset.owner_id)
-        search_results = await specific_search(
-            query_type,
-            query_text,
-            user,
+
+        specific_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,
         )
-        return {
-            "search_result": search_results,
-            "dataset_id": dataset.id,
-            "dataset_name": dataset.name,
-        }
+        search_tools = specific_search_tools
+        if len(search_tools) == 2:
+            [get_completion, get_context] = search_tools
+
+            if only_context:
+                return None, await get_context(query_text), [dataset]
+
+            search_context = context or await get_context(query_text)
+            search_result = await get_completion(query_text, search_context)
+
+            return search_result, search_context, [dataset]
+        else:
+            unknown_tool = search_tools[0]
+
+            return await unknown_tool(query_text), "", [dataset]
 
     # Search every dataset async based on query and appropriate database configuration
     tasks = []
     for dataset in search_datasets:
         tasks.append(
-            _search_by_context(
-                dataset, user, query_type, query_text, system_prompt_path, top_k, last_k
+            _search_in_dataset_context(
+                dataset=dataset,
+                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,
+                only_context=only_context,
+                context=context,
             )
         )