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

cognee/modules/search/methods/get_search_type_retriever_instance.py

@@ -0,0 +1,252 @@
+import os
+from typing import Callable, List, Optional, Type, Tuple
+
+from cognee.modules.retrieval.base_retriever import BaseRetriever
+
+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_retriever_instance(
+    query_type: SearchType,
+    query_text: str,
+    **kwargs,
+) -> BaseRetriever:
+    """
+    Factory method to get the appropriate retriever instance based on the search type.
+
+    Args:
+        query_type: SearchType enum indicating the type of search.
+        query_text: query string.
+        retriever_specific_config: Retriever specific configuration dictionary.
+        **kwargs: General keyword arguments for retriever initialization.
+
+    Returns:
+
+    """
+    # Transform retriever specific config if empty to avoid None checks later
+    retriever_specific_config = kwargs.get("retriever_specific_config")
+    if retriever_specific_config is None:
+        retriever_specific_config = {}
+
+    # Extract common defaults with fallback values from kwargs
+    top_k = kwargs.get("top_k", 10)
+    system_prompt_path = kwargs.get("system_prompt_path", "answer_simple_question.txt")
+    system_prompt = kwargs.get("system_prompt")
+    node_type = kwargs.get("node_type", NodeSet)
+    node_name = kwargs.get("node_name")
+    save_interaction = kwargs.get("save_interaction", False)
+    wide_search_top_k = kwargs.get("wide_search_top_k", 100)
+    triplet_distance_penalty = kwargs.get("triplet_distance_penalty", 3.5)
+    session_id = kwargs.get("session_id")
+
+    # Registry mapping search types to their corresponding retriever classes and input parameters
+    search_core_registry: dict[SearchType, Tuple[BaseRetriever, dict]] = {
+        SearchType.SUMMARIES: (SummariesRetriever, {"top_k": top_k, "session_id": session_id}),
+        SearchType.CHUNKS: (
+            ChunksRetriever,
+            {"top_k": top_k},
+        ),
+        SearchType.RAG_COMPLETION: (
+            CompletionRetriever,
+            {
+                "system_prompt_path": system_prompt_path,
+                "top_k": top_k,
+                "system_prompt": system_prompt,
+                "session_id": session_id,
+                "response_model": retriever_specific_config.get("response_model", str),
+            },
+        ),
+        SearchType.TRIPLET_COMPLETION: (
+            TripletRetriever,
+            {
+                "system_prompt_path": system_prompt_path,
+                "top_k": top_k,
+                "system_prompt": system_prompt,
+                "session_id": session_id,
+                "response_model": retriever_specific_config.get("response_model", str),
+            },
+        ),
+        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,
+                "session_id": session_id,
+                "response_model": retriever_specific_config.get("response_model", str),
+            },
+        ),
+        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,
+                "max_iter": retriever_specific_config.get("max_iter", 4),
+                "validation_system_prompt_path": retriever_specific_config.get(
+                    "validation_system_prompt_path", "cot_validation_system_prompt.txt"
+                ),
+                "validation_user_prompt_path": retriever_specific_config.get(
+                    "validation_user_prompt_path", "cot_validation_user_prompt.txt"
+                ),
+                "followup_system_prompt_path": retriever_specific_config.get(
+                    "followup_system_prompt_path", "cot_followup_system_prompt.txt"
+                ),
+                "followup_user_prompt_path": retriever_specific_config.get(
+                    "followup_user_prompt_path", "cot_followup_user_prompt.txt"
+                ),
+                "session_id": session_id,
+                "response_model": retriever_specific_config.get("response_model", str),
+            },
+        ),
+        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,
+                "context_extension_rounds": retriever_specific_config.get(
+                    "context_extension_rounds", 4
+                ),
+                "session_id": session_id,
+                "response_model": retriever_specific_config.get("response_model", str),
+            },
+        ),
+        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,
+                "session_id": session_id,
+                "summarize_prompt_path": retriever_specific_config.get(
+                    "summarize_prompt_path", "summarize_search_results.txt"
+                ),
+            },
+        ),
+        SearchType.CYPHER: (
+            CypherSearchRetriever,
+            {
+                "user_prompt_path": retriever_specific_config.get(
+                    "user_prompt_path", "context_for_question.txt"
+                ),
+                "system_prompt_path": retriever_specific_config.get(
+                    "system_prompt_path", "answer_simple_question.txt"
+                ),
+                "session_id": session_id,
+            },
+        ),
+        SearchType.NATURAL_LANGUAGE: (
+            NaturalLanguageRetriever,
+            {
+                "session_id": session_id,
+                "system_prompt_path": retriever_specific_config.get(
+                    "system_prompt_path", "natural_language_retriever_system.txt"
+                ),
+                "max_attempts": retriever_specific_config.get("max_attempts", 3),
+            },
+        ),
+        SearchType.TEMPORAL: (
+            TemporalRetriever,
+            {
+                "top_k": top_k,
+                "wide_search_top_k": wide_search_top_k,
+                "triplet_distance_penalty": triplet_distance_penalty,
+                "session_id": session_id,
+                "response_model": retriever_specific_config.get("response_model", str),
+                "user_prompt_path": retriever_specific_config.get(
+                    "user_prompt_path", "graph_context_for_question.txt"
+                ),
+                "system_prompt_path": retriever_specific_config.get(
+                    "system_prompt_path", "answer_simple_question.txt"
+                ),
+                "time_extraction_prompt_path": retriever_specific_config.get(
+                    "time_extraction_prompt_path", "extract_query_time.txt"
+                ),
+                "node_type": node_type,
+                "node_name": node_name,
+            },
+        ),
+        SearchType.CHUNKS_LEXICAL: (JaccardChunksRetriever, {"top_k": top_k}),
+        SearchType.CODING_RULES: (
+            CodingRulesRetriever,
+            {"rules_nodeset_name": node_name},
+        ),
+    }
+
+    # 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.get(query_type)
+
+        if not retriever:
+            raise UnsupportedSearchTypeError(str(query_type))
+        # TODO: Fix community retrievers on the community side so they get all input parameters properly
+        retriever_instance = retriever(**kwargs)
+    else:
+        retriever_info = search_core_registry.get(query_type)
+        # Check if retriever info is found for the given query type
+        if not retriever_info:
+            raise UnsupportedSearchTypeError(str(query_type))
+
+        # If it exists unpack the retriever class and its initialization arguments
+        retriever_cls, retriever_args = retriever_info
+
+        retriever_instance = retriever_cls(**retriever_args)
+
+    return retriever_instance

cognee/modules/search/methods/search.py

@@ -14,8 +14,6 @@ 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.search.operations import log_query, log_result
@@ -25,9 +23,7 @@ from cognee.modules.data.methods.get_authorized_existing_datasets import (
     get_authorized_existing_datasets,
 )
 from cognee import __version__ as cognee_version
-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
+from cognee.modules.search.methods.get_retriever_output import get_retriever_output
 
 logger = get_logger()
 
@@ -45,11 +41,12 @@ async def search(
     save_interaction: bool = False,
     last_k: Optional[int] = None,
     only_context: bool = False,
-    use_combined_context: bool = False,
     session_id: Optional[str] = None,
     wide_search_top_k: Optional[int] = 100,
     triplet_distance_penalty: Optional[float] = 3.5,
-) -> Union[CombinedSearchResult, List[SearchResult]]:
+    verbose=False,
+    retriever_specific_config: Optional[dict] = None,
+) -> List[SearchResult]:
     """
 
     Args:
@@ -75,44 +72,24 @@ async def search(
         },
     )
 
-    # Use search function filtered by permissions if access control is enabled
-    if backend_access_control_enabled():
-        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,
-            session_id=session_id,
-            wide_search_top_k=wide_search_top_k,
-            triplet_distance_penalty=triplet_distance_penalty,
-        )
-    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,
-                session_id=session_id,
-                wide_search_top_k=wide_search_top_k,
-                triplet_distance_penalty=triplet_distance_penalty,
-            )
-        ]
+    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,
+        session_id=session_id,
+        wide_search_top_k=wide_search_top_k,
+        triplet_distance_penalty=triplet_distance_penalty,
+        retriever_specific_config=retriever_specific_config,
+    )
 
     send_telemetry(
         "cognee.search EXECUTION COMPLETED",
@@ -125,89 +102,11 @@ async def search(
 
     await log_result(
         query.id,
-        json.dumps(
-            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
-                ]
-            )
-        ),
+        json.dumps(jsonable_encoder(search_results)),
         user.id,
     )
 
-    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:
-        # This is for maintaining backwards compatibility
-        if backend_access_control_enabled():
-            return_value = []
-            for search_result in search_results:
-                prepared_search_results = await prepare_search_result(search_result)
-
-                result = prepared_search_results["result"]
-                graphs = prepared_search_results["graphs"]
-                context = prepared_search_results["context"]
-                datasets = prepared_search_results["datasets"]
-
-                if only_context:
-                    return_value.append(
-                        {
-                            "search_result": [context] if context else None,
-                            "dataset_id": datasets[0].id,
-                            "dataset_name": datasets[0].name,
-                            "dataset_tenant_id": datasets[0].tenant_id,
-                            "graphs": graphs,
-                        }
-                    )
-                else:
-                    return_value.append(
-                        {
-                            "search_result": [result] if result else None,
-                            "dataset_id": datasets[0].id,
-                            "dataset_name": datasets[0].name,
-                            "dataset_tenant_id": datasets[0].tenant_id,
-                            "graphs": graphs,
-                        }
-                    )
-            return return_value
-        else:
-            return_value = []
-            if only_context:
-                for search_result in search_results:
-                    prepared_search_results = await prepare_search_result(search_result)
-                    return_value.append(prepared_search_results["context"])
-            else:
-                for search_result in search_results:
-                    result, context, datasets = search_result
-                    return_value.append(result)
-            # For maintaining backwards compatibility
-            if len(return_value) == 1 and isinstance(return_value[0], list):
-                return return_value[0]
-            else:
-                return return_value
+    return _backwards_compatible_search_results(search_results, verbose)
 
 
 async def authorized_search(
@@ -223,14 +122,11 @@ async def authorized_search(
     save_interaction: bool = False,
     last_k: Optional[int] = None,
     only_context: bool = False,
-    use_combined_context: bool = False,
     session_id: Optional[str] = None,
     wide_search_top_k: Optional[int] = 100,
     triplet_distance_penalty: Optional[float] = 3.5,
-) -> Union[
-    Tuple[Any, Union[List[Edge], str], List[Dataset]],
-    List[Tuple[Any, Union[List[Edge], str], List[Dataset]]],
-]:
+    retriever_specific_config: Optional[dict] = None,
+) -> 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.
@@ -240,70 +136,6 @@ async def authorized_search(
         datasets=dataset_ids, permission_type="read", user=user
     )
 
-    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,
-            last_k=last_k,
-            only_context=True,
-            session_id=session_id,
-            wide_search_top_k=wide_search_top_k,
-            triplet_distance_penalty=triplet_distance_penalty,
-        )
-
-        context = {}
-        datasets: List[Dataset] = []
-
-        for _, search_context, search_datasets in search_responses:
-            for dataset in search_datasets:
-                context[str(dataset.id)] = search_context
-
-            datasets.extend(search_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,
-            last_k=last_k,
-            wide_search_top_k=wide_search_top_k,
-            triplet_distance_penalty=triplet_distance_penalty,
-        )
-        search_tools = specific_search_tools
-        if len(search_tools) == 2:
-            [get_completion, _] = search_tools
-        else:
-            get_completion = search_tools[0]
-
-        def prepare_combined_context(
-            context,
-        ) -> Union[List[Edge], str]:
-            combined_context = []
-
-            for dataset_context in context.values():
-                combined_context += dataset_context
-
-            if combined_context and isinstance(combined_context[0], str):
-                return "\n".join(combined_context)
-
-            return combined_context
-
-        combined_context = prepare_combined_context(context)
-        completion = await get_completion(query_text, combined_context, session_id=session_id)
-
-        return completion, combined_context, datasets
-
     # Searches all provided datasets and handles setting up of appropriate database context based on permissions
     search_results = await search_in_datasets_context(
         search_datasets=search_datasets,
@@ -319,6 +151,8 @@ async def authorized_search(
         only_context=only_context,
         session_id=session_id,
         wide_search_top_k=wide_search_top_k,
+        triplet_distance_penalty=triplet_distance_penalty,
+        retriever_specific_config=retriever_specific_config,
     )
 
     return search_results
@@ -336,10 +170,10 @@ async def search_in_datasets_context(
     save_interaction: bool = False,
     last_k: Optional[int] = None,
     only_context: bool = False,
-    context: Optional[Any] = None,
     session_id: Optional[str] = None,
     wide_search_top_k: Optional[int] = 100,
     triplet_distance_penalty: Optional[float] = 3.5,
+    retriever_specific_config: Optional[dict] = 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.
@@ -358,17 +192,17 @@ async def search_in_datasets_context(
         save_interaction: bool = False,
         last_k: Optional[int] = None,
         only_context: bool = False,
-        context: Optional[Any] = None,
         session_id: Optional[str] = None,
         wide_search_top_k: Optional[int] = 100,
         triplet_distance_penalty: Optional[float] = 3.5,
+        retriever_specific_config: Optional[dict] = 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)
 
+        # Check if graph for dataset is empty and log warnings if necessary
         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
             from cognee.modules.data.methods import get_dataset_data
@@ -382,12 +216,14 @@ async def search_in_datasets_context(
                 )
             else:
                 logger.warning(
-                    "Search attempt on an empty knowledge graph - no data has been added to this dataset"
+                    f"Search attempt on an empty knowledge graph - no data has been added to this dataset: {dataset.name}"
                 )
 
-        specific_search_tools = await get_search_type_tools(
+        # Get retriever output in the context of the current dataset
+        return await get_retriever_output(
             query_type=query_type,
             query_text=query_text,
+            dataset=dataset,
             system_prompt_path=system_prompt_path,
             system_prompt=system_prompt,
             top_k=top_k,
@@ -395,24 +231,12 @@ async def search_in_datasets_context(
             node_name=node_name,
             save_interaction=save_interaction,
             last_k=last_k,
+            only_context=only_context,
+            session_id=session_id,
             wide_search_top_k=wide_search_top_k,
             triplet_distance_penalty=triplet_distance_penalty,
+            retriever_specific_config=retriever_specific_config,
         )
-        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, session_id=session_id)
-
-            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 = []
@@ -430,11 +254,62 @@ async def search_in_datasets_context(
                 save_interaction=save_interaction,
                 last_k=last_k,
                 only_context=only_context,
-                context=context,
                 session_id=session_id,
                 wide_search_top_k=wide_search_top_k,
                 triplet_distance_penalty=triplet_distance_penalty,
+                retriever_specific_config=retriever_specific_config,
             )
         )
 
     return await asyncio.gather(*tasks)
+
+
+def _backwards_compatible_search_results(search_results, verbose: bool):
+    """
+    Prepares search results in a format compatible with previous versions of the API.
+    """
+    # This is for maintaining backwards compatibility
+    if backend_access_control_enabled():
+        return_value = []
+        for search_result in search_results:
+            # Dataset info needs to be always included
+            search_result_dict = {
+                "dataset_id": search_result.dataset_id,
+                "dataset_name": search_result.dataset_name,
+                "dataset_tenant_id": search_result.dataset_tenant_id,
+            }
+            if verbose:
+                # Include all different types of results only in verbose mode
+                search_result_dict["text_result"] = search_result.completion
+                search_result_dict["context_result"] = search_result.context
+                search_result_dict["objects_result"] = search_result.result_object
+            else:
+                # Result attribute handles returning appropriate result based on set flags and outputs
+                search_result_dict["search_result"] = search_result.result
+
+            return_value.append(search_result_dict)
+        return return_value
+    else:
+        return_value = []
+        if verbose:
+            for search_result in search_results:
+                # Include all different types of results only in verbose mode
+                search_result_dict = {
+                    "text_result": search_result.completion,
+                    "context_result": search_result.context,
+                    "objects_result": search_result.result_object,
+                }
+                return_value.append(search_result_dict)
+            return return_value
+        else:
+            for search_result in search_results:
+                # Result attribute handles returning appropriate result based on set flags and outputs
+                return_value.append(search_result.result)
+
+            # For maintaining backwards compatibility
+            if len(return_value) == 1 and isinstance(return_value[0], list):
+                # If a single element list return the element directly
+                return return_value[0]
+            else:
+                # Otherwise return the list of results
+                return return_value