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

cognee/modules/cognify/config.py

@@ -8,14 +8,12 @@ import os
 class CognifyConfig(BaseSettings):
     classification_model: object = DefaultContentPrediction
     summarization_model: object = SummarizedContent
-    triplet_embedding: bool = False
     model_config = SettingsConfigDict(env_file=".env", extra="allow")
 
     def to_dict(self) -> dict:
         return {
             "classification_model": self.classification_model,
             "summarization_model": self.summarization_model,
-            "triplet_embedding": self.triplet_embedding,
         }
 
 

cognee/modules/data/deletion/prune_system.py

@@ -1,67 +1,17 @@
-from sqlalchemy.exc import OperationalError
-
-from cognee.infrastructure.databases.exceptions import EntityNotFoundError
-from cognee.context_global_variables import backend_access_control_enabled
 from cognee.infrastructure.databases.vector import get_vector_engine
 from cognee.infrastructure.databases.graph.get_graph_engine import get_graph_engine
 from cognee.infrastructure.databases.relational import get_relational_engine
-from cognee.infrastructure.databases.utils import (
-    get_graph_dataset_database_handler,
-    get_vector_dataset_database_handler,
-)
 from cognee.shared.cache import delete_cache
-from cognee.modules.users.models import DatasetDatabase
-from cognee.shared.logging_utils import get_logger
-
-logger = get_logger()
-
-
-async def prune_graph_databases():
-    db_engine = get_relational_engine()
-    try:
-        dataset_databases = await db_engine.get_all_data_from_table("dataset_database")
-        # Go through each dataset database and delete the graph database
-        for dataset_database in dataset_databases:
-            handler = get_graph_dataset_database_handler(dataset_database)
-            await handler["handler_instance"].delete_dataset(dataset_database)
-    except (OperationalError, EntityNotFoundError) as e:
-        logger.debug(
-            "Skipping pruning of graph DB. Error when accessing dataset_database table: %s",
-            e,
-        )
-        return
-
-
-async def prune_vector_databases():
-    db_engine = get_relational_engine()
-    try:
-        dataset_databases = await db_engine.get_all_data_from_table("dataset_database")
-        # Go through each dataset database and delete the vector database
-        for dataset_database in dataset_databases:
-            handler = get_vector_dataset_database_handler(dataset_database)
-            await handler["handler_instance"].delete_dataset(dataset_database)
-    except (OperationalError, EntityNotFoundError) as e:
-        logger.debug(
-            "Skipping pruning of vector DB. Error when accessing dataset_database table: %s",
-            e,
-        )
-        return
 
 
 async def prune_system(graph=True, vector=True, metadata=True, cache=True):
-    # Note: prune system should not be available through the API, it has no permission checks and will
-    #       delete all graph and vector databases if called. It should only be used in development or testing environments.
-    if graph and not backend_access_control_enabled():
+    if graph:
         graph_engine = await get_graph_engine()
         await graph_engine.delete_graph()
-    elif graph and backend_access_control_enabled():
-        await prune_graph_databases()
 
-    if vector and not backend_access_control_enabled():
+    if vector:
         vector_engine = get_vector_engine()
         await vector_engine.prune()
-    elif vector and backend_access_control_enabled():
-        await prune_vector_databases()
 
     if metadata:
         db_engine = get_relational_engine()

cognee/modules/data/methods/delete_dataset.py

@@ -1,34 +1,8 @@
-from cognee.modules.users.models import DatasetDatabase
-from sqlalchemy import select
-
 from cognee.modules.data.models import Dataset
-from cognee.infrastructure.databases.utils.get_vector_dataset_database_handler import (
-    get_vector_dataset_database_handler,
-)
-from cognee.infrastructure.databases.utils.get_graph_dataset_database_handler import (
-    get_graph_dataset_database_handler,
-)
 from cognee.infrastructure.databases.relational import get_relational_engine
 
 
 async def delete_dataset(dataset: Dataset):
     db_engine = get_relational_engine()
 
-    async with db_engine.get_async_session() as session:
-        stmt = select(DatasetDatabase).where(
-            DatasetDatabase.dataset_id == dataset.id,
-        )
-        dataset_database: DatasetDatabase = await session.scalar(stmt)
-        if dataset_database:
-            graph_dataset_database_handler = get_graph_dataset_database_handler(dataset_database)
-            vector_dataset_database_handler = get_vector_dataset_database_handler(dataset_database)
-            await graph_dataset_database_handler["handler_instance"].delete_dataset(
-                dataset_database
-            )
-            await vector_dataset_database_handler["handler_instance"].delete_dataset(
-                dataset_database
-            )
-    # TODO: Remove dataset from pipeline_run_status in Data objects related to dataset as well
-    #       This blocks recreation of the dataset with the same name and data after deletion as
-    #       it's marked as completed and will be just skipped even though it's empty.
     return await db_engine.delete_entity_by_id(dataset.__tablename__, dataset.id)

cognee/modules/engine/models/__init__.py

@@ -7,4 +7,3 @@ from .ColumnValue import ColumnValue
 from .Timestamp import Timestamp
 from .Interval import Interval
 from .Event import Event
-from .Triplet import Triplet

cognee/modules/graph/cognee_graph/CogneeGraph.py

@@ -56,68 +56,6 @@ class CogneeGraph(CogneeAbstractGraph):
     def get_edges(self) -> List[Edge]:
         return self.edges
 
-    async def _get_nodeset_subgraph(
-        self,
-        adapter,
-        node_type,
-        node_name,
-    ):
-        """Retrieve subgraph based on node type and name."""
-        logger.info("Retrieving graph filtered by node type and node name (NodeSet).")
-        nodes_data, edges_data = await adapter.get_nodeset_subgraph(
-            node_type=node_type, node_name=node_name
-        )
-        if not nodes_data or not edges_data:
-            raise EntityNotFoundError(
-                message="Nodeset does not exist, or empty nodeset projected from the database."
-            )
-        return nodes_data, edges_data
-
-    async def _get_full_or_id_filtered_graph(
-        self,
-        adapter,
-        relevant_ids_to_filter,
-    ):
-        """Retrieve full or ID-filtered graph with fallback."""
-        if relevant_ids_to_filter is None:
-            logger.info("Retrieving full graph.")
-            nodes_data, edges_data = await adapter.get_graph_data()
-            if not nodes_data or not edges_data:
-                raise EntityNotFoundError(message="Empty graph projected from the database.")
-            return nodes_data, edges_data
-
-        get_graph_data_fn = getattr(adapter, "get_id_filtered_graph_data", adapter.get_graph_data)
-        if getattr(adapter.__class__, "get_id_filtered_graph_data", None):
-            logger.info("Retrieving ID-filtered graph from database.")
-            nodes_data, edges_data = await get_graph_data_fn(target_ids=relevant_ids_to_filter)
-        else:
-            logger.info("Retrieving full graph from database.")
-            nodes_data, edges_data = await get_graph_data_fn()
-        if hasattr(adapter, "get_id_filtered_graph_data") and (not nodes_data or not edges_data):
-            logger.warning(
-                "Id filtered graph returned empty, falling back to full graph retrieval."
-            )
-            logger.info("Retrieving full graph")
-            nodes_data, edges_data = await adapter.get_graph_data()
-
-        if not nodes_data or not edges_data:
-            raise EntityNotFoundError("Empty graph projected from the database.")
-        return nodes_data, edges_data
-
-    async def _get_filtered_graph(
-        self,
-        adapter,
-        memory_fragment_filter,
-    ):
-        """Retrieve graph filtered by attributes."""
-        logger.info("Retrieving graph filtered by memory fragment")
-        nodes_data, edges_data = await adapter.get_filtered_graph_data(
-            attribute_filters=memory_fragment_filter
-        )
-        if not nodes_data or not edges_data:
-            raise EntityNotFoundError(message="Empty filtered graph projected from the database.")
-        return nodes_data, edges_data
-
     async def project_graph_from_db(
         self,
         adapter: Union[GraphDBInterface],
@@ -129,39 +67,40 @@ class CogneeGraph(CogneeAbstractGraph):
         memory_fragment_filter=[],
         node_type: Optional[Type] = None,
         node_name: Optional[List[str]] = None,
-        relevant_ids_to_filter: Optional[List[str]] = None,
-        triplet_distance_penalty: float = 3.5,
     ) -> None:
         if node_dimension < 1 or edge_dimension < 1:
             raise InvalidDimensionsError()
         try:
+            import time
+
+            start_time = time.time()
+
+            # Determine projection strategy
             if node_type is not None and node_name not in [None, [], ""]:
-                nodes_data, edges_data = await self._get_nodeset_subgraph(
-                    adapter, node_type, node_name
+                nodes_data, edges_data = await adapter.get_nodeset_subgraph(
+                    node_type=node_type, node_name=node_name
                 )
+                if not nodes_data or not edges_data:
+                    raise EntityNotFoundError(
+                        message="Nodeset does not exist, or empty nodetes projected from the database."
+                    )
             elif len(memory_fragment_filter) == 0:
-                nodes_data, edges_data = await self._get_full_or_id_filtered_graph(
-                    adapter, relevant_ids_to_filter
-                )
+                nodes_data, edges_data = await adapter.get_graph_data()
+                if not nodes_data or not edges_data:
+                    raise EntityNotFoundError(message="Empty graph projected from the database.")
             else:
-                nodes_data, edges_data = await self._get_filtered_graph(
-                    adapter, memory_fragment_filter
+                nodes_data, edges_data = await adapter.get_filtered_graph_data(
+                    attribute_filters=memory_fragment_filter
                 )
+                if not nodes_data or not edges_data:
+                    raise EntityNotFoundError(
+                        message="Empty filtered graph projected from the database."
+                    )
 
-            import time
-
-            start_time = time.time()
             # Process nodes
             for node_id, properties in nodes_data:
                 node_attributes = {key: properties.get(key) for key in node_properties_to_project}
-                self.add_node(
-                    Node(
-                        str(node_id),
-                        node_attributes,
-                        dimension=node_dimension,
-                        node_penalty=triplet_distance_penalty,
-                    )
-                )
+                self.add_node(Node(str(node_id), node_attributes, dimension=node_dimension))
 
             # Process edges
             for source_id, target_id, relationship_type, properties in edges_data:
@@ -179,7 +118,6 @@ class CogneeGraph(CogneeAbstractGraph):
                         attributes=edge_attributes,
                         directed=directed,
                         dimension=edge_dimension,
-                        edge_penalty=triplet_distance_penalty,
                     )
                     self.add_edge(edge)
 
@@ -211,10 +149,24 @@ class CogneeGraph(CogneeAbstractGraph):
                     node.add_attribute("vector_distance", score)
                     mapped_nodes += 1
 
-    async def map_vector_distances_to_graph_edges(self, edge_distances) -> None:
+    async def map_vector_distances_to_graph_edges(
+        self, vector_engine, query_vector, edge_distances
+    ) -> None:
         try:
+            if query_vector is None or len(query_vector) == 0:
+                raise ValueError("Failed to generate query embedding.")
+
             if edge_distances is None:
-                return
+                start_time = time.time()
+                edge_distances = await vector_engine.search(
+                    collection_name="EdgeType_relationship_name",
+                    query_vector=query_vector,
+                    limit=None,
+                )
+                projection_time = time.time() - start_time
+                logger.info(
+                    f"Edge collection distances were calculated separately from nodes in {projection_time:.2f}s"
+                )
 
             embedding_map = {result.payload["text"]: result.score for result in edge_distances}
 

cognee/modules/graph/cognee_graph/CogneeGraphElements.py

@@ -20,17 +20,13 @@ class Node:
     status: np.ndarray
 
     def __init__(
-        self,
-        node_id: str,
-        attributes: Optional[Dict[str, Any]] = None,
-        dimension: int = 1,
-        node_penalty: float = 3.5,
+        self, node_id: str, attributes: Optional[Dict[str, Any]] = None, dimension: int = 1
     ):
         if dimension <= 0:
             raise InvalidDimensionsError()
         self.id = node_id
         self.attributes = attributes if attributes is not None else {}
-        self.attributes["vector_distance"] = node_penalty
+        self.attributes["vector_distance"] = float("inf")
         self.skeleton_neighbours = []
         self.skeleton_edges = []
         self.status = np.ones(dimension, dtype=int)
@@ -109,14 +105,13 @@ class Edge:
         attributes: Optional[Dict[str, Any]] = None,
         directed: bool = True,
         dimension: int = 1,
-        edge_penalty: float = 3.5,
     ):
         if dimension <= 0:
             raise InvalidDimensionsError()
         self.node1 = node1
         self.node2 = node2
         self.attributes = attributes if attributes is not None else {}
-        self.attributes["vector_distance"] = edge_penalty
+        self.attributes["vector_distance"] = float("inf")
         self.directed = directed
         self.status = np.ones(dimension, dtype=int)
 

cognee/modules/memify/memify.py

@@ -12,6 +12,9 @@ from cognee.modules.users.models import User
 from cognee.modules.pipelines.layers.resolve_authorized_user_datasets import (
     resolve_authorized_user_datasets,
 )
+from cognee.modules.pipelines.layers.reset_dataset_pipeline_run_status import (
+    reset_dataset_pipeline_run_status,
+)
 from cognee.modules.engine.operations.setup import setup
 from cognee.modules.pipelines.layers.pipeline_execution_mode import get_pipeline_executor
 from cognee.tasks.memify.extract_subgraph_chunks import extract_subgraph_chunks
@@ -94,6 +97,10 @@ async def memify(
         *enrichment_tasks,
     ]
 
+    await reset_dataset_pipeline_run_status(
+        authorized_dataset.id, user, pipeline_names=["memify_pipeline"]
+    )
+
     # By calling get pipeline executor we get a function that will have the run_pipeline run in the background or a function that we will need to wait for
     pipeline_executor_func = get_pipeline_executor(run_in_background=run_in_background)
 
@@ -106,7 +113,6 @@ async def memify(
         datasets=authorized_dataset.id,
         vector_db_config=vector_db_config,
         graph_db_config=graph_db_config,
-        use_pipeline_cache=False,
         incremental_loading=False,
         pipeline_name="memify_pipeline",
     )

cognee/modules/pipelines/operations/pipeline.py

@@ -20,9 +20,6 @@ from cognee.modules.pipelines.layers.resolve_authorized_user_datasets import (
 from cognee.modules.pipelines.layers.check_pipeline_run_qualification import (
     check_pipeline_run_qualification,
 )
-from cognee.modules.pipelines.models.PipelineRunInfo import (
-    PipelineRunStarted,
-)
 from typing import Any
 
 logger = get_logger("cognee.pipeline")
@@ -38,7 +35,6 @@ async def run_pipeline(
     pipeline_name: str = "custom_pipeline",
     vector_db_config: dict = None,
     graph_db_config: dict = None,
-    use_pipeline_cache: bool = False,
     incremental_loading: bool = False,
     data_per_batch: int = 20,
 ):
@@ -55,7 +51,6 @@ async def run_pipeline(
             data=data,
             pipeline_name=pipeline_name,
             context={"dataset": dataset},
-            use_pipeline_cache=use_pipeline_cache,
             incremental_loading=incremental_loading,
             data_per_batch=data_per_batch,
         ):
@@ -69,7 +64,6 @@ async def run_pipeline_per_dataset(
     data=None,
     pipeline_name: str = "custom_pipeline",
     context: dict = None,
-    use_pipeline_cache=False,
     incremental_loading=False,
     data_per_batch: int = 20,
 ):
@@ -83,18 +77,8 @@ async def run_pipeline_per_dataset(
     if process_pipeline_status:
         # If pipeline was already processed or is currently being processed
         # return status information to async generator and finish execution
-        if use_pipeline_cache:
-            # If pipeline caching is enabled we do not proceed with re-processing
-            yield process_pipeline_status
-            return
-        else:
-            # If pipeline caching is disabled we always return pipeline started information and proceed with re-processing
-            yield PipelineRunStarted(
-                pipeline_run_id=process_pipeline_status.pipeline_run_id,
-                dataset_id=dataset.id,
-                dataset_name=dataset.name,
-                payload=data,
-            )
+        yield process_pipeline_status
+        return
 
     pipeline_run = run_tasks(
         tasks,

cognee/modules/retrieval/__init__.py

@@ -1 +1 @@
-
+from cognee.modules.retrieval.code_retriever import CodeRetriever

cognee/modules/retrieval/code_retriever.py

@@ -0,0 +1,232 @@
+from typing import Any, Optional, List
+import asyncio
+import aiofiles
+from pydantic import BaseModel
+
+from cognee.shared.logging_utils import get_logger
+from cognee.modules.retrieval.base_retriever import BaseRetriever
+from cognee.infrastructure.databases.graph import get_graph_engine
+from cognee.infrastructure.databases.vector import get_vector_engine
+from cognee.infrastructure.llm.prompts import read_query_prompt
+from cognee.infrastructure.llm.LLMGateway import LLMGateway
+
+logger = get_logger("CodeRetriever")
+
+
+class CodeRetriever(BaseRetriever):
+    """Retriever for handling code-based searches."""
+
+    class CodeQueryInfo(BaseModel):
+        """
+        Model for representing the result of a query related to code files.
+
+        This class holds a list of filenames and the corresponding source code extracted from a
+        query. It is used to encapsulate response data in a structured format.
+        """
+
+        filenames: List[str] = []
+        sourcecode: str
+
+    def __init__(self, top_k: int = 3):
+        """Initialize retriever with search parameters."""
+        self.top_k = top_k
+        self.file_name_collections = ["CodeFile_name"]
+        self.classes_and_functions_collections = [
+            "ClassDefinition_source_code",
+            "FunctionDefinition_source_code",
+        ]
+
+    async def _process_query(self, query: str) -> "CodeRetriever.CodeQueryInfo":
+        """Process the query using LLM to extract file names and source code parts."""
+        logger.debug(
+            f"Processing query with LLM: '{query[:100]}{'...' if len(query) > 100 else ''}'"
+        )
+
+        system_prompt = read_query_prompt("codegraph_retriever_system.txt")
+
+        try:
+            result = await LLMGateway.acreate_structured_output(
+                text_input=query,
+                system_prompt=system_prompt,
+                response_model=self.CodeQueryInfo,
+            )
+            logger.info(
+                f"LLM extracted {len(result.filenames)} filenames and {len(result.sourcecode)} chars of source code"
+            )
+            return result
+        except Exception as e:
+            logger.error(f"Failed to retrieve structured output from LLM: {str(e)}")
+            raise RuntimeError("Failed to retrieve structured output from LLM") from e
+
+    async def get_context(self, query: str) -> Any:
+        """Find relevant code files based on the query."""
+        logger.info(
+            f"Starting code retrieval for query: '{query[:100]}{'...' if len(query) > 100 else ''}'"
+        )
+
+        if not query or not isinstance(query, str):
+            logger.error("Invalid query: must be a non-empty string")
+            raise ValueError("The query must be a non-empty string.")
+
+        try:
+            vector_engine = get_vector_engine()
+            graph_engine = await get_graph_engine()
+            logger.debug("Successfully initialized vector and graph engines")
+        except Exception as e:
+            logger.error(f"Database initialization error: {str(e)}")
+            raise RuntimeError("Database initialization error in code_graph_retriever, ") from e
+
+        files_and_codeparts = await self._process_query(query)
+
+        similar_filenames = []
+        similar_codepieces = []
+
+        if not files_and_codeparts.filenames or not files_and_codeparts.sourcecode:
+            logger.info("No specific files/code extracted from query, performing general search")
+
+            for collection in self.file_name_collections:
+                logger.debug(f"Searching {collection} collection with general query")
+                search_results_file = await vector_engine.search(
+                    collection, query, limit=self.top_k
+                )
+                logger.debug(f"Found {len(search_results_file)} results in {collection}")
+                for res in search_results_file:
+                    similar_filenames.append(
+                        {"id": res.id, "score": res.score, "payload": res.payload}
+                    )
+
+            existing_collection = []
+            for collection in self.classes_and_functions_collections:
+                if await vector_engine.has_collection(collection):
+                    existing_collection.append(collection)
+
+            if not existing_collection:
+                raise RuntimeError("No collection found for code retriever")
+
+            for collection in existing_collection:
+                logger.debug(f"Searching {collection} collection with general query")
+                search_results_code = await vector_engine.search(
+                    collection, query, limit=self.top_k
+                )
+                logger.debug(f"Found {len(search_results_code)} results in {collection}")
+                for res in search_results_code:
+                    similar_codepieces.append(
+                        {"id": res.id, "score": res.score, "payload": res.payload}
+                    )
+        else:
+            logger.info(
+                f"Using extracted filenames ({len(files_and_codeparts.filenames)}) and source code for targeted search"
+            )
+
+            for collection in self.file_name_collections:
+                for file_from_query in files_and_codeparts.filenames:
+                    logger.debug(f"Searching {collection} for specific file: {file_from_query}")
+                    search_results_file = await vector_engine.search(
+                        collection, file_from_query, limit=self.top_k
+                    )
+                    logger.debug(
+                        f"Found {len(search_results_file)} results for file {file_from_query}"
+                    )
+                    for res in search_results_file:
+                        similar_filenames.append(
+                            {"id": res.id, "score": res.score, "payload": res.payload}
+                        )
+
+            for collection in self.classes_and_functions_collections:
+                logger.debug(f"Searching {collection} with extracted source code")
+                search_results_code = await vector_engine.search(
+                    collection, files_and_codeparts.sourcecode, limit=self.top_k
+                )
+                logger.debug(f"Found {len(search_results_code)} results for source code search")
+                for res in search_results_code:
+                    similar_codepieces.append(
+                        {"id": res.id, "score": res.score, "payload": res.payload}
+                    )
+
+        total_items = len(similar_filenames) + len(similar_codepieces)
+        logger.info(
+            f"Total search results: {total_items} items ({len(similar_filenames)} filenames, {len(similar_codepieces)} code pieces)"
+        )
+
+        if total_items == 0:
+            logger.warning("No search results found, returning empty list")
+            return []
+
+        logger.debug("Getting graph connections for all search results")
+        relevant_triplets = await asyncio.gather(
+            *[
+                graph_engine.get_connections(similar_piece["id"])
+                for similar_piece in similar_filenames + similar_codepieces
+            ]
+        )
+        logger.info(f"Retrieved graph connections for {len(relevant_triplets)} items")
+
+        paths = set()
+        for i, sublist in enumerate(relevant_triplets):
+            logger.debug(f"Processing connections for item {i}: {len(sublist)} connections")
+            for tpl in sublist:
+                if isinstance(tpl, tuple) and len(tpl) >= 3:
+                    if "file_path" in tpl[0]:
+                        paths.add(tpl[0]["file_path"])
+                    if "file_path" in tpl[2]:
+                        paths.add(tpl[2]["file_path"])
+
+        logger.info(f"Found {len(paths)} unique file paths to read")
+
+        retrieved_files = {}
+        read_tasks = []
+        for file_path in paths:
+
+            async def read_file(fp):
+                try:
+                    logger.debug(f"Reading file: {fp}")
+                    async with aiofiles.open(fp, "r", encoding="utf-8") as f:
+                        content = await f.read()
+                        retrieved_files[fp] = content
+                        logger.debug(f"Successfully read {len(content)} characters from {fp}")
+                except Exception as e:
+                    logger.error(f"Error reading {fp}: {e}")
+                    retrieved_files[fp] = ""
+
+            read_tasks.append(read_file(file_path))
+
+        await asyncio.gather(*read_tasks)
+        logger.info(
+            f"Successfully read {len([f for f in retrieved_files.values() if f])} files (out of {len(paths)} total)"
+        )
+
+        result = [
+            {
+                "name": file_path,
+                "description": file_path,
+                "content": retrieved_files[file_path],
+            }
+            for file_path in paths
+        ]
+
+        logger.info(f"Returning {len(result)} code file contexts")
+        return result
+
+    async def get_completion(
+        self, query: str, context: Optional[Any] = None, session_id: Optional[str] = None
+    ) -> Any:
+        """
+        Returns the code files context.
+
+        Parameters:
+        -----------
+
+            - query (str): The query string to retrieve code context for.
+            - context (Optional[Any]): Optional pre-fetched context; 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)
+
+        Returns:
+        --------
+
+            - Any: The code files context, either provided or retrieved.
+        """
+        if context is None:
+            context = await self.get_context(query)
+        return context

cognee/modules/retrieval/graph_completion_context_extension_retriever.py

@@ -39,8 +39,6 @@ class GraphCompletionContextExtensionRetriever(GraphCompletionRetriever):
         node_type: Optional[Type] = None,
         node_name: Optional[List[str]] = None,
         save_interaction: bool = False,
-        wide_search_top_k: Optional[int] = 100,
-        triplet_distance_penalty: Optional[float] = 3.5,
     ):
         super().__init__(
             user_prompt_path=user_prompt_path,
@@ -50,8 +48,6 @@ class GraphCompletionContextExtensionRetriever(GraphCompletionRetriever):
             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,
         )
 
     async def get_completion(