cognee 0.5.0.dev0__py3-none-any.whl → 0.5.1__py3-none-any.whl

cognee/modules/retrieval/triplet_retriever.py

@@ -0,0 +1,182 @@
+import asyncio
+from typing import Any, Optional, Type, List
+
+from cognee.shared.logging_utils import get_logger
+from cognee.infrastructure.databases.vector import get_vector_engine
+from cognee.modules.retrieval.utils.completion import generate_completion, summarize_text
+from cognee.modules.retrieval.utils.session_cache import (
+    save_conversation_history,
+    get_conversation_history,
+)
+from cognee.modules.retrieval.base_retriever import BaseRetriever
+from cognee.modules.retrieval.exceptions.exceptions import NoDataError
+from cognee.infrastructure.databases.vector.exceptions import CollectionNotFoundError
+from cognee.context_global_variables import session_user
+from cognee.infrastructure.databases.cache.config import CacheConfig
+
+logger = get_logger("TripletRetriever")
+
+
+class TripletRetriever(BaseRetriever):
+    """
+    Retriever for handling LLM-based completion searches using triplets.
+
+    Public methods:
+    - get_context(query: str) -> str
+    - get_completion(query: str, context: Optional[Any] = None) -> Any
+    """
+
+    def __init__(
+        self,
+        user_prompt_path: str = "context_for_question.txt",
+        system_prompt_path: str = "answer_simple_question.txt",
+        system_prompt: Optional[str] = None,
+        top_k: Optional[int] = 5,
+    ):
+        """Initialize retriever with optional custom prompt paths."""
+        self.user_prompt_path = user_prompt_path
+        self.system_prompt_path = system_prompt_path
+        self.top_k = top_k if top_k is not None else 1
+        self.system_prompt = system_prompt
+
+    async def get_context(self, query: str) -> str:
+        """
+        Retrieves relevant triplets as context.
+
+        Fetches triplets based on a query from a vector engine and combines their text.
+        Returns empty string if no triplets are found. Raises NoDataError if the collection is not
+        found.
+
+        Parameters:
+        -----------
+
+            - query (str): The query string used to search for relevant triplets.
+
+        Returns:
+        --------
+
+            - str: A string containing the combined text of the retrieved triplets, or an
+              empty string if none are found.
+        """
+        vector_engine = get_vector_engine()
+
+        try:
+            if not await vector_engine.has_collection(collection_name="Triplet_text"):
+                logger.error("Triplet_text collection not found")
+                raise NoDataError(
+                    "In order to use TRIPLET_COMPLETION first use the create_triplet_embeddings memify pipeline. "
+                )
+
+            found_triplets = await vector_engine.search("Triplet_text", query, limit=self.top_k)
+
+            if len(found_triplets) == 0:
+                return ""
+
+            triplets_payload = [found_triplet.payload["text"] for found_triplet in found_triplets]
+            combined_context = "\n".join(triplets_payload)
+            return combined_context
+        except CollectionNotFoundError as error:
+            logger.error("Triplet_text collection not found")
+            raise NoDataError("No data found in the system, please add data first.") from error
+
+    async def get_completion(
+        self,
+        query: str,
+        context: Optional[Any] = None,
+        session_id: Optional[str] = None,
+        response_model: Type = str,
+    ) -> List[Any]:
+        """
+        Generates an LLM completion using the context.
+
+        Retrieves context if not provided and generates a completion based on the query and
+        context using an external completion generator.
+
+        Parameters:
+        -----------
+
+            - query (str): The query string to be used for generating a completion.
+            - context (Optional[Any]): Optional pre-fetched context to use for generating the
+              completion; if None, it retrieves the context for the query. (default None)
+            - session_id (Optional[str]): Optional session identifier for caching. If None,
+              defaults to 'default_session'. (default None)
+            - response_model (Type): The Pydantic model type for structured output. (default str)
+
+        Returns:
+        --------
+
+            - Any: The generated completion based on the provided query and context.
+        """
+        if context is None:
+            context = await self.get_context(query)
+
+        cache_config = CacheConfig()
+        user = session_user.get()
+        user_id = getattr(user, "id", None)
+        session_save = user_id and cache_config.caching
+
+        if session_save:
+            completion = await self._get_completion_with_session(
+                query=query,
+                context=context,
+                session_id=session_id,
+                response_model=response_model,
+            )
+        else:
+            completion = await self._get_completion_without_session(
+                query=query,
+                context=context,
+                response_model=response_model,
+            )
+
+        return [completion]
+
+    async def _get_completion_with_session(
+        self,
+        query: str,
+        context: str,
+        session_id: Optional[str],
+        response_model: Type,
+    ) -> Any:
+        """Generate completion with session history and caching."""
+        conversation_history = await get_conversation_history(session_id=session_id)
+
+        context_summary, completion = await asyncio.gather(
+            summarize_text(context),
+            generate_completion(
+                query=query,
+                context=context,
+                user_prompt_path=self.user_prompt_path,
+                system_prompt_path=self.system_prompt_path,
+                system_prompt=self.system_prompt,
+                conversation_history=conversation_history,
+                response_model=response_model,
+            ),
+        )
+
+        await save_conversation_history(
+            query=query,
+            context_summary=context_summary,
+            answer=completion,
+            session_id=session_id,
+        )
+
+        return completion
+
+    async def _get_completion_without_session(
+        self,
+        query: str,
+        context: str,
+        response_model: Type,
+    ) -> Any:
+        """Generate completion without session history."""
+        completion = await generate_completion(
+            query=query,
+            context=context,
+            user_prompt_path=self.user_prompt_path,
+            system_prompt_path=self.system_prompt_path,
+            system_prompt=self.system_prompt,
+            response_model=response_model,
+        )
+
+        return completion

cognee/modules/retrieval/utils/brute_force_triplet_search.py

@@ -58,6 +58,8 @@ async def get_memory_fragment(
     properties_to_project: Optional[List[str]] = None,
     node_type: Optional[Type] = None,
     node_name: Optional[List[str]] = None,
+    relevant_ids_to_filter: Optional[List[str]] = None,
+    triplet_distance_penalty: Optional[float] = 3.5,
 ) -> CogneeGraph:
     """Creates and initializes a CogneeGraph memory fragment with optional property projections."""
     if properties_to_project is None:
@@ -74,6 +76,8 @@ async def get_memory_fragment(
             edge_properties_to_project=["relationship_name", "edge_text"],
             node_type=node_type,
             node_name=node_name,
+            relevant_ids_to_filter=relevant_ids_to_filter,
+            triplet_distance_penalty=triplet_distance_penalty,
         )
 
     except EntityNotFoundError:
@@ -95,6 +99,8 @@ async def brute_force_triplet_search(
     memory_fragment: Optional[CogneeGraph] = None,
     node_type: Optional[Type] = None,
     node_name: Optional[List[str]] = None,
+    wide_search_top_k: Optional[int] = 100,
+    triplet_distance_penalty: Optional[float] = 3.5,
 ) -> List[Edge]:
     """
     Performs a brute force search to retrieve the top triplets from the graph.
@@ -107,6 +113,8 @@ async def brute_force_triplet_search(
         memory_fragment (Optional[CogneeGraph]): Existing memory fragment to reuse.
         node_type: node type to filter
         node_name: node name to filter
+        wide_search_top_k (Optional[int]): Number of initial elements to retrieve from collections
+        triplet_distance_penalty (Optional[float]): Default distance penalty in graph projection
 
     Returns:
         list: The top triplet results.
@@ -116,10 +124,10 @@ async def brute_force_triplet_search(
     if top_k <= 0:
         raise ValueError("top_k must be a positive integer.")
 
-    if memory_fragment is None:
-        memory_fragment = await get_memory_fragment(
-            properties_to_project, node_type=node_type, node_name=node_name
-        )
+    # Setting wide search limit based on the parameters
+    non_global_search = node_name is None
+
+    wide_search_limit = wide_search_top_k if non_global_search else None
 
     if collections is None:
         collections = [
@@ -129,6 +137,9 @@ async def brute_force_triplet_search(
             "DocumentChunk_text",
         ]
 
+    if "EdgeType_relationship_name" not in collections:
+        collections.append("EdgeType_relationship_name")
+
     try:
         vector_engine = get_vector_engine()
     except Exception as e:
@@ -140,7 +151,7 @@ async def brute_force_triplet_search(
     async def search_in_collection(collection_name: str):
         try:
             return await vector_engine.search(
-                collection_name=collection_name, query_vector=query_vector, limit=None
+                collection_name=collection_name, query_vector=query_vector, limit=wide_search_limit
             )
         except CollectionNotFoundError:
             return []
@@ -156,19 +167,40 @@ async def brute_force_triplet_search(
             return []
 
         # Final statistics
-        projection_time = time.time() - start_time
+        vector_collection_search_time = time.time() - start_time
         logger.info(
-            f"Vector collection retrieval completed: Retrieved distances from {sum(1 for res in results if res)} collections in {projection_time:.2f}s"
+            f"Vector collection retrieval completed: Retrieved distances from {sum(1 for res in results if res)} collections in {vector_collection_search_time:.2f}s"
         )
 
         node_distances = {collection: result for collection, result in zip(collections, results)}
 
         edge_distances = node_distances.get("EdgeType_relationship_name", None)
 
+        if wide_search_limit is not None:
+            relevant_ids_to_filter = list(
+                {
+                    str(getattr(scored_node, "id"))
+                    for collection_name, score_collection in node_distances.items()
+                    if collection_name != "EdgeType_relationship_name"
+                    and isinstance(score_collection, (list, tuple))
+                    for scored_node in score_collection
+                    if getattr(scored_node, "id", None)
+                }
+            )
+        else:
+            relevant_ids_to_filter = None
+
+        if memory_fragment is None:
+            memory_fragment = await get_memory_fragment(
+                properties_to_project=properties_to_project,
+                node_type=node_type,
+                node_name=node_name,
+                relevant_ids_to_filter=relevant_ids_to_filter,
+                triplet_distance_penalty=triplet_distance_penalty,
+            )
+
         await memory_fragment.map_vector_distances_to_graph_nodes(node_distances=node_distances)
-        await memory_fragment.map_vector_distances_to_graph_edges(
-            vector_engine=vector_engine, query_vector=query_vector, edge_distances=edge_distances
-        )
+        await memory_fragment.map_vector_distances_to_graph_edges(edge_distances=edge_distances)
 
         results = await memory_fragment.calculate_top_triplet_importances(k=top_k)
 

cognee/modules/run_custom_pipeline/run_custom_pipeline.py

@@ -18,6 +18,8 @@ async def run_custom_pipeline(
     user: User = None,
     vector_db_config: Optional[dict] = None,
     graph_db_config: Optional[dict] = None,
+    use_pipeline_cache: bool = False,
+    incremental_loading: bool = False,
     data_per_batch: int = 20,
     run_in_background: bool = False,
     pipeline_name: str = "custom_pipeline",
@@ -40,6 +42,10 @@ async def run_custom_pipeline(
         user: User context for authentication and data access. Uses default if None.
         vector_db_config: Custom vector database configuration for embeddings storage.
         graph_db_config: Custom graph database configuration for relationship storage.
+        use_pipeline_cache: If True, pipelines with the same ID that are currently executing and pipelines with the same ID that were completed won't process data again.
+                        Pipelines ID is created based on the generate_pipeline_id function. Pipeline status can be manually reset with the reset_dataset_pipeline_run_status function.
+        incremental_loading: If True, only new or modified data will be processed to avoid duplication. (Only works if data is used with the Cognee python Data model).
+                            The incremental system stores and compares hashes of processed data in the Data model and skips data with the same content hash.
         data_per_batch: Number of data items to be processed in parallel.
         run_in_background: If True, starts processing asynchronously and returns immediately.
                           If False, waits for completion before returning.
@@ -63,7 +69,8 @@ async def run_custom_pipeline(
         datasets=dataset,
         vector_db_config=vector_db_config,
         graph_db_config=graph_db_config,
-        incremental_loading=False,
+        use_pipeline_cache=use_pipeline_cache,
+        incremental_loading=incremental_loading,
         data_per_batch=data_per_batch,
         pipeline_name=pipeline_name,
     )

cognee/modules/search/methods/get_search_type_tools.py

@@ -2,6 +2,7 @@ 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
@@ -22,7 +23,6 @@ from cognee.modules.retrieval.graph_completion_cot_retriever import GraphComplet
 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
 
@@ -37,6 +37,8 @@ async def get_search_type_tools(
     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: [
@@ -59,6 +61,18 @@ async def get_search_type_tools(
                 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,
@@ -67,6 +81,8 @@ async def get_search_type_tools(
                 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,
@@ -75,6 +91,8 @@ async def get_search_type_tools(
                 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: [
@@ -85,6 +103,8 @@ async def get_search_type_tools(
                 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,
@@ -93,6 +113,8 @@ async def get_search_type_tools(
                 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: [
@@ -103,6 +125,8 @@ async def get_search_type_tools(
                 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,
@@ -111,6 +135,8 @@ async def get_search_type_tools(
                 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: [
@@ -121,6 +147,8 @@ async def get_search_type_tools(
                 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,
@@ -129,12 +157,10 @@ async def get_search_type_tools(
                 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.CODE: [
-            CodeRetriever(top_k=top_k).get_completion,
-            CodeRetriever(top_k=top_k).get_context,
-        ],
         SearchType.CYPHER: [
             CypherSearchRetriever().get_completion,
             CypherSearchRetriever().get_context,
@@ -145,8 +171,16 @@ async def get_search_type_tools(
         ],
         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,
+            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): [
@@ -169,7 +203,19 @@ async def get_search_type_tools(
     ):
         raise UnsupportedSearchTypeError("Cypher query search types are disabled.")
 
-    search_type_tools = search_tasks.get(query_type)
+    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))

cognee/modules/search/methods/no_access_control_search.py

@@ -24,6 +24,8 @@ async def no_access_control_search(
     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,
@@ -35,6 +37,8 @@ async def no_access_control_search(
         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()

cognee/modules/search/methods/search.py

@@ -47,6 +47,9 @@ async def search(
     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,
+    verbose: bool = False,
 ) -> Union[CombinedSearchResult, List[SearchResult]]:
     """
 
@@ -90,6 +93,8 @@ async def search(
             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 = [
@@ -105,6 +110,8 @@ async def search(
                 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,
             )
         ]
 
@@ -134,6 +141,7 @@ async def search(
     )
 
     if use_combined_context:
+        # Note: combined context search must always be verbose and return a CombinedSearchResult with graphs info
         prepared_search_results = await prepare_search_result(
             search_results[0] if isinstance(search_results, list) else search_results
         )
@@ -167,25 +175,30 @@ async def search(
                 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,
-                        }
-                    )
+                    search_result_dict = {
+                        "search_result": [context] if context else None,
+                        "dataset_id": datasets[0].id,
+                        "dataset_name": datasets[0].name,
+                        "dataset_tenant_id": datasets[0].tenant_id,
+                    }
+                    if verbose:
+                        # Include graphs only in verbose mode
+                        search_result_dict["graphs"] = graphs
+
+                    return_value.append(search_result_dict)
                 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,
-                        }
-                    )
+                    search_result_dict = {
+                        "search_result": [result] if result else None,
+                        "dataset_id": datasets[0].id,
+                        "dataset_name": datasets[0].name,
+                        "dataset_tenant_id": datasets[0].tenant_id,
+                    }
+                    if verbose:
+                        # Include graphs only in verbose mode
+                        search_result_dict["graphs"] = graphs
+
+                    return_value.append(search_result_dict)
+
             return return_value
         else:
             return_value = []
@@ -219,6 +232,8 @@ async def authorized_search(
     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]]],
@@ -246,6 +261,8 @@ async def authorized_search(
             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 = {}
@@ -267,6 +284,8 @@ async def authorized_search(
             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:
@@ -306,6 +325,7 @@ async def authorized_search(
         last_k=last_k,
         only_context=only_context,
         session_id=session_id,
+        wide_search_top_k=wide_search_top_k,
     )
 
     return search_results
@@ -325,6 +345,8 @@ async def search_in_datasets_context(
     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,
 ) -> List[Tuple[Any, Union[str, List[Edge]], List[Dataset]]]:
     """
     Searches all provided datasets and handles setting up of appropriate database context based on permissions.
@@ -345,6 +367,8 @@ async def search_in_datasets_context(
         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,
     ) -> 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)
@@ -378,6 +402,8 @@ async def search_in_datasets_context(
             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:
@@ -413,6 +439,8 @@ async def search_in_datasets_context(
                 only_context=only_context,
                 context=context,
                 session_id=session_id,
+                wide_search_top_k=wide_search_top_k,
+                triplet_distance_penalty=triplet_distance_penalty,
             )
         )
 

cognee/modules/search/types/SearchType.py

@@ -5,9 +5,9 @@ class SearchType(Enum):
     SUMMARIES = "SUMMARIES"
     CHUNKS = "CHUNKS"
     RAG_COMPLETION = "RAG_COMPLETION"
+    TRIPLET_COMPLETION = "TRIPLET_COMPLETION"
     GRAPH_COMPLETION = "GRAPH_COMPLETION"
     GRAPH_SUMMARY_COMPLETION = "GRAPH_SUMMARY_COMPLETION"
-    CODE = "CODE"
     CYPHER = "CYPHER"
     NATURAL_LANGUAGE = "NATURAL_LANGUAGE"
     GRAPH_COMPLETION_COT = "GRAPH_COMPLETION_COT"