cognee 0.2.1.dev7__py3-none-any.whl → 0.2.2.dev1__py3-none-any.whl

cognee/modules/retrieval/graph_completion_cot_retriever.py

@@ -1,9 +1,9 @@
 from typing import Any, Optional, List, Type
 from cognee.shared.logging_utils import get_logger
-from cognee.infrastructure.llm.get_llm_client import get_llm_client
+
 from cognee.modules.retrieval.graph_completion_retriever import GraphCompletionRetriever
 from cognee.modules.retrieval.utils.completion import generate_completion
-from cognee.infrastructure.llm.prompts import read_query_prompt, render_prompt
+from cognee.infrastructure.llm.LLMGateway import LLMGateway
 
 logger = get_logger()
 
@@ -73,7 +73,6 @@ class GraphCompletionCotRetriever(GraphCompletionRetriever):
 
             - List[str]: A list containing the generated answer to the user's query.
         """
-        llm_client = get_llm_client()
         followup_question = ""
         triplets = []
         answer = [""]
@@ -95,27 +94,27 @@ class GraphCompletionCotRetriever(GraphCompletionRetriever):
             logger.info(f"Chain-of-thought: round {round_idx} - answer: {answer}")
             if round_idx < max_iter:
                 valid_args = {"query": query, "answer": answer, "context": context}
-                valid_user_prompt = render_prompt(
+                valid_user_prompt = LLMGateway.render_prompt(
                     filename=self.validation_user_prompt_path, context=valid_args
                 )
-                valid_system_prompt = read_query_prompt(
+                valid_system_prompt = LLMGateway.read_query_prompt(
                     prompt_file_name=self.validation_system_prompt_path
                 )
 
-                reasoning = await llm_client.acreate_structured_output(
+                reasoning = await LLMGateway.acreate_structured_output(
                     text_input=valid_user_prompt,
                     system_prompt=valid_system_prompt,
                     response_model=str,
                 )
                 followup_args = {"query": query, "answer": answer, "reasoning": reasoning}
-                followup_prompt = render_prompt(
+                followup_prompt = LLMGateway.render_prompt(
                     filename=self.followup_user_prompt_path, context=followup_args
                 )
-                followup_system = read_query_prompt(
+                followup_system = LLMGateway.read_query_prompt(
                     prompt_file_name=self.followup_system_prompt_path
                 )
 
-                followup_question = await llm_client.acreate_structured_output(
+                followup_question = await LLMGateway.acreate_structured_output(
                     text_input=followup_prompt, system_prompt=followup_system, response_model=str
                 )
                 logger.info(

cognee/modules/retrieval/graph_completion_retriever.py

@@ -10,7 +10,7 @@ from cognee.modules.retrieval.utils.completion import generate_completion
 from cognee.modules.retrieval.utils.stop_words import DEFAULT_STOP_WORDS
 from cognee.shared.logging_utils import get_logger
 
-logger = get_logger()
+logger = get_logger("GraphCompletionRetriever")
 
 
 class GraphCompletionRetriever(BaseRetriever):

cognee/modules/retrieval/insights_retriever.py

@@ -1,12 +1,15 @@
 import asyncio
 from typing import Any, Optional
 
+from cognee.shared.logging_utils import get_logger
 from cognee.infrastructure.databases.graph import get_graph_engine
 from cognee.infrastructure.databases.vector import get_vector_engine
 from cognee.modules.retrieval.base_retriever import BaseRetriever
 from cognee.modules.retrieval.exceptions.exceptions import NoDataError
 from cognee.infrastructure.databases.vector.exceptions.exceptions import CollectionNotFoundError
 
+logger = get_logger("InsightsRetriever")
+
 
 class InsightsRetriever(BaseRetriever):
     """
@@ -63,6 +66,7 @@ class InsightsRetriever(BaseRetriever):
                     vector_engine.search("EntityType_name", query_text=query, limit=self.top_k),
                 )
             except CollectionNotFoundError as error:
+                logger.error("Entity collections not found")
                 raise NoDataError("No data found in the system, please add data first.") from error
 
             results = [*results[0], *results[1]]

cognee/modules/retrieval/natural_language_retriever.py

@@ -1,14 +1,13 @@
 from typing import Any, Optional
-import logging
+from cognee.shared.logging_utils import get_logger
 from cognee.infrastructure.databases.graph import get_graph_engine
 from cognee.infrastructure.databases.graph.networkx.adapter import NetworkXAdapter
-from cognee.infrastructure.llm.get_llm_client import get_llm_client
-from cognee.infrastructure.llm.prompts import render_prompt
+from cognee.infrastructure.llm.LLMGateway import LLMGateway
 from cognee.modules.retrieval.base_retriever import BaseRetriever
 from cognee.modules.retrieval.exceptions import SearchTypeNotSupported
 from cognee.infrastructure.databases.graph.graph_db_interface import GraphDBInterface
 
-logger = logging.getLogger("NaturalLanguageRetriever")
+logger = get_logger("NaturalLanguageRetriever")
 
 
 class NaturalLanguageRetriever(BaseRetriever):
@@ -51,8 +50,7 @@ class NaturalLanguageRetriever(BaseRetriever):
 
     async def _generate_cypher_query(self, query: str, edge_schemas, previous_attempts=None) -> str:
         """Generate a Cypher query using LLM based on natural language query and schema information."""
-        llm_client = get_llm_client()
-        system_prompt = render_prompt(
+        system_prompt = LLMGateway.render_prompt(
             self.system_prompt_path,
             context={
                 "edge_schemas": edge_schemas,
@@ -60,7 +58,7 @@ class NaturalLanguageRetriever(BaseRetriever):
             },
         )
 
-        return await llm_client.acreate_structured_output(
+        return await LLMGateway.acreate_structured_output(
             text_input=query,
             system_prompt=system_prompt,
             response_model=str,
@@ -123,16 +121,12 @@ class NaturalLanguageRetriever(BaseRetriever):
             - Optional[Any]: Returns the context retrieved from the graph database based on the
               query.
         """
-        try:
-            graph_engine = await get_graph_engine()
+        graph_engine = await get_graph_engine()
 
-            if isinstance(graph_engine, (NetworkXAdapter)):
-                raise SearchTypeNotSupported("Natural language search type not supported.")
+        if isinstance(graph_engine, (NetworkXAdapter)):
+            raise SearchTypeNotSupported("Natural language search type not supported.")
 
-            return await self._execute_cypher_query(query, graph_engine)
-        except Exception as e:
-            logger.error("Failed to execute natural language search retrieval: %s", str(e))
-            raise e
+        return await self._execute_cypher_query(query, graph_engine)
 
     async def get_completion(self, query: str, context: Optional[Any] = None) -> Any:
         """

cognee/modules/retrieval/summaries_retriever.py

@@ -1,10 +1,13 @@
 from typing import Any, Optional
 
+from cognee.shared.logging_utils import get_logger
 from cognee.infrastructure.databases.vector import get_vector_engine
 from cognee.modules.retrieval.base_retriever import BaseRetriever
 from cognee.modules.retrieval.exceptions.exceptions import NoDataError
 from cognee.infrastructure.databases.vector.exceptions.exceptions import CollectionNotFoundError
 
+logger = get_logger("SummariesRetriever")
+
 
 class SummariesRetriever(BaseRetriever):
     """
@@ -40,16 +43,24 @@ class SummariesRetriever(BaseRetriever):
 
             - Any: A list of payloads from the retrieved summaries.
         """
+        logger.info(
+            f"Starting summary retrieval for query: '{query[:100]}{'...' if len(query) > 100 else ''}'"
+        )
+
         vector_engine = get_vector_engine()
 
         try:
             summaries_results = await vector_engine.search(
                 "TextSummary_text", query, limit=self.top_k
             )
+            logger.info(f"Found {len(summaries_results)} summaries from vector search")
         except CollectionNotFoundError as error:
+            logger.error("TextSummary_text collection not found in vector database")
             raise NoDataError("No data found in the system, please add data first.") from error
 
-        return [summary.payload for summary in summaries_results]
+        summary_payloads = [summary.payload for summary in summaries_results]
+        logger.info(f"Returning {len(summary_payloads)} summary payloads")
+        return summary_payloads
 
     async def get_completion(self, query: str, context: Optional[Any] = None) -> Any:
         """
@@ -70,6 +81,17 @@ class SummariesRetriever(BaseRetriever):
 
             - Any: The generated completion context, which is either provided or retrieved.
         """
+        logger.info(
+            f"Starting completion generation for query: '{query[:100]}{'...' if len(query) > 100 else ''}'"
+        )
+
         if context is None:
+            logger.debug("No context provided, retrieving context from vector database")
             context = await self.get_context(query)
+        else:
+            logger.debug("Using provided context")
+
+        logger.info(
+            f"Returning context with {len(context) if isinstance(context, list) else 1} item(s)"
+        )
         return context

cognee/modules/retrieval/utils/brute_force_triplet_search.py

@@ -1,4 +1,5 @@
 import asyncio
+import time
 from typing import List, Optional, Type
 
 from cognee.shared.logging_utils import get_logger, ERROR
@@ -59,13 +60,13 @@ async def get_memory_fragment(
     node_name: Optional[List[str]] = None,
 ) -> CogneeGraph:
     """Creates and initializes a CogneeGraph memory fragment with optional property projections."""
-    graph_engine = await get_graph_engine()
-    memory_fragment = CogneeGraph()
-
     if properties_to_project is None:
         properties_to_project = ["id", "description", "name", "type", "text"]
 
     try:
+        graph_engine = await get_graph_engine()
+        memory_fragment = CogneeGraph()
+
         await memory_fragment.project_graph_from_db(
             graph_engine,
             node_properties_to_project=properties_to_project,
@@ -73,7 +74,13 @@ async def get_memory_fragment(
             node_type=node_type,
             node_name=node_name,
         )
+
     except EntityNotFoundError:
+        # This is expected behavior - continue with empty fragment
+        pass
+    except Exception as e:
+        logger.error(f"Error during memory fragment creation: {str(e)}")
+        # Still return the fragment even if projection failed
         pass
 
     return memory_fragment
@@ -168,6 +175,8 @@ async def brute_force_search(
             return []
 
     try:
+        start_time = time.time()
+
         results = await asyncio.gather(
             *[search_in_collection(collection_name) for collection_name in collections]
         )
@@ -175,10 +184,20 @@ async def brute_force_search(
         if all(not item for item in results):
             return []
 
+        # Final statistics
+        projection_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"
+        )
+
         node_distances = {collection: result for collection, result in zip(collections, results)}
 
+        edge_distances = node_distances.get("EdgeType_relationship_name", None)
+
         await memory_fragment.map_vector_distances_to_graph_nodes(node_distances=node_distances)
-        await memory_fragment.map_vector_distances_to_graph_edges(vector_engine, query)
+        await memory_fragment.map_vector_distances_to_graph_edges(
+            vector_engine=vector_engine, query_vector=query_vector, edge_distances=edge_distances
+        )
 
         results = await memory_fragment.calculate_top_triplet_importances(k=top_k)
 

cognee/modules/retrieval/utils/completion.py

@@ -1,5 +1,4 @@
-from cognee.infrastructure.llm.get_llm_client import get_llm_client
-from cognee.infrastructure.llm.prompts import read_query_prompt, render_prompt
+from cognee.infrastructure.llm.LLMGateway import LLMGateway
 
 
 async def generate_completion(
@@ -10,11 +9,10 @@ async def generate_completion(
 ) -> str:
     """Generates a completion using LLM with given context and prompts."""
     args = {"question": query, "context": context}
-    user_prompt = render_prompt(user_prompt_path, args)
-    system_prompt = read_query_prompt(system_prompt_path)
+    user_prompt = LLMGateway.render_prompt(user_prompt_path, args)
+    system_prompt = LLMGateway.read_query_prompt(system_prompt_path)
 
-    llm_client = get_llm_client()
-    return await llm_client.acreate_structured_output(
+    return await LLMGateway.acreate_structured_output(
         text_input=user_prompt,
         system_prompt=system_prompt,
         response_model=str,
@@ -26,10 +24,9 @@ async def summarize_text(
     prompt_path: str = "summarize_search_results.txt",
 ) -> str:
     """Summarizes text using LLM with the specified prompt."""
-    system_prompt = read_query_prompt(prompt_path)
-    llm_client = get_llm_client()
+    system_prompt = LLMGateway.read_query_prompt(prompt_path)
 
-    return await llm_client.acreate_structured_output(
+    return await LLMGateway.acreate_structured_output(
         text_input=text,
         system_prompt=system_prompt,
         response_model=str,

cognee/modules/retrieval/utils/description_to_codepart_search.py

@@ -9,7 +9,7 @@ from cognee.modules.users.methods import get_default_user
 from cognee.modules.users.models import User
 from cognee.shared.utils import send_telemetry
 from cognee.modules.search.methods import search
-from cognee.infrastructure.llm.get_llm_client import get_llm_client
+from cognee.infrastructure.llm.LLMGateway import LLMGateway
 
 logger = get_logger(level=ERROR)
 
@@ -71,8 +71,7 @@ async def code_description_to_code_part(
                 if isinstance(obj, dict) and "description" in obj
             )
 
-            llm_client = get_llm_client()
-            context_from_documents = await llm_client.acreate_structured_output(
+            context_from_documents = await LLMGateway.acreate_structured_output(
                 text_input=f"The retrieved context from documents is {concatenated_descriptions}.",
                 system_prompt="You are a Senior Software Engineer, summarize the context from documents"
                 f" in a way that it is gonna be provided next to codeparts as context"

cognee/modules/search/methods/search.py

@@ -27,7 +27,7 @@ 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
+from cognee.modules.search.operations import log_query, log_result, select_search_type
 
 
 async def search(
@@ -129,6 +129,10 @@ async def specific_search(
         SearchType.NATURAL_LANGUAGE: NaturalLanguageRetriever().get_completion,
     }
 
+    # 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:

cognee/modules/search/operations/__init__.py

@@ -1,3 +1,4 @@
 from .log_query import log_query
 from .log_result import log_result
 from .get_history import get_history
+from .select_search_type import select_search_type

cognee/modules/search/operations/select_search_type.py

@@ -0,0 +1,42 @@
+from cognee.infrastructure.llm.prompts import read_query_prompt
+from cognee.modules.search.types import SearchType
+from cognee.shared.logging_utils import get_logger
+from cognee.infrastructure.llm.LLMGateway import LLMGateway
+
+logger = get_logger("SearchTypeSelector")
+
+
+async def select_search_type(
+    query: str,
+    system_prompt_path: str = "search_type_selector_prompt.txt",
+) -> SearchType:
+    """
+    Analyzes the query and Selects the best search type.
+
+    Args:
+        query: The query to analyze.
+        system_prompt_path: The path to the system prompt.
+
+    Returns:
+        The best search type given by the LLM.
+    """
+    default_search_type = SearchType.RAG_COMPLETION
+    system_prompt = read_query_prompt(system_prompt_path)
+
+    try:
+        response = await LLMGateway.acreate_structured_output(
+            text_input=query,
+            system_prompt=system_prompt,
+            response_model=str,
+        )
+
+        if response.upper() in SearchType.__members__:
+            logger.info(f"Selected lucky search type: {response.upper()}")
+            return SearchType(response.upper())
+
+        # If the response is not a valid search type, return the default search type
+        logger.info(f"LLM gives an invalid search type: {response.upper()}")
+        return default_search_type
+    except Exception as e:
+        logger.error(f"Failed to select search type intelligently from LLM: {str(e)}")
+        return default_search_type

cognee/modules/search/types/SearchType.py

@@ -13,3 +13,4 @@ class SearchType(Enum):
     NATURAL_LANGUAGE = "NATURAL_LANGUAGE"
     GRAPH_COMPLETION_COT = "GRAPH_COMPLETION_COT"
     GRAPH_COMPLETION_CONTEXT_EXTENSION = "GRAPH_COMPLETION_CONTEXT_EXTENSION"
+    FEELING_LUCKY = "FEELING_LUCKY"

cognee/modules/settings/get_settings.py

@@ -43,14 +43,6 @@ def get_settings() -> SettingsDict:
     llm_config = get_llm_config()
 
     vector_dbs = [
-        {
-            "value": "weaviate",
-            "label": "Weaviate",
-        },
-        {
-            "value": "qdrant",
-            "label": "Qdrant",
-        },
         {
             "value": "lancedb",
             "label": "LanceDB",

cognee/modules/settings/save_vector_db_config.py

@@ -6,7 +6,7 @@ from cognee.infrastructure.databases.vector import get_vectordb_config
 class VectorDBConfig(BaseModel):
     url: str
     api_key: str
-    provider: Union[Literal["lancedb"], Literal["qdrant"], Literal["weaviate"], Literal["pgvector"]]
+    provider: Union[Literal["lancedb"], Literal["pgvector"]]
 
 
 async def save_vector_db_config(vector_db_config: VectorDBConfig):

cognee/shared/data_models.py

@@ -4,7 +4,9 @@ from enum import Enum, auto
 from typing import Any, Dict, List, Optional, Union
 
 from pydantic import BaseModel, Field
-from cognee.infrastructure.llm.config import get_llm_config
+from cognee.infrastructure.llm.config import (
+    get_llm_config,
+)
 
 if get_llm_config().llm_provider.lower() == "gemini":
     """

cognee/shared/logging_utils.py

@@ -175,17 +175,13 @@ def log_database_configuration(logger):
     try:
         # Log relational database configuration
         relational_config = get_relational_config()
-        logger.info(f"Relational database: {relational_config.db_provider}")
         if relational_config.db_provider == "postgres":
             logger.info(f"Postgres host: {relational_config.db_host}:{relational_config.db_port}")
-            logger.info(f"Postgres database: {relational_config.db_name}")
         elif relational_config.db_provider == "sqlite":
             logger.info(f"SQLite path: {relational_config.db_path}")
-            logger.info(f"SQLite database: {relational_config.db_name}")
 
         # Log vector database configuration
         vector_config = get_vectordb_config()
-        logger.info(f"Vector database: {vector_config.vector_db_provider}")
         if vector_config.vector_db_provider == "lancedb":
             logger.info(f"Vector database path: {vector_config.vector_db_url}")
         else:
@@ -193,7 +189,6 @@ def log_database_configuration(logger):
 
         # Log graph database configuration
         graph_config = get_graph_config()
-        logger.info(f"Graph database: {graph_config.graph_database_provider}")
         if graph_config.graph_database_provider == "kuzu":
             logger.info(f"Graph database path: {graph_config.graph_file_path}")
         else:

cognee/tasks/chunk_naive_llm_classifier/chunk_naive_llm_classifier.py

@@ -7,7 +7,7 @@ from pydantic import BaseModel
 from cognee.infrastructure.databases.graph import get_graph_engine
 from cognee.infrastructure.databases.vector import get_vector_engine
 from cognee.infrastructure.engine.models import DataPoint
-from cognee.modules.data.extraction.extract_categories import extract_categories
+from cognee.infrastructure.llm.LLMGateway import LLMGateway
 from cognee.modules.chunking.models.DocumentChunk import DocumentChunk
 
 
@@ -40,7 +40,7 @@ async def chunk_naive_llm_classifier(
         return data_chunks
 
     chunk_classifications = await asyncio.gather(
-        *[extract_categories(chunk.text, classification_model) for chunk in data_chunks],
+        *[LLMGateway.extract_categories(chunk.text, classification_model) for chunk in data_chunks],
     )
 
     classification_data_points = []

cognee/tasks/documents/extract_chunks_from_documents.py

@@ -8,7 +8,6 @@ from cognee.modules.data.models import Data
 from cognee.infrastructure.databases.relational import get_relational_engine
 from cognee.modules.chunking.TextChunker import TextChunker
 from cognee.modules.chunking.Chunker import Chunker
-from cognee.modules.data.processing.document_types.exceptions.exceptions import PyPdfInternalError
 
 
 async def update_document_token_count(document_id: UUID, token_count: int) -> None:
@@ -40,15 +39,14 @@ async def extract_chunks_from_documents(
     """
     for document in documents:
         document_token_count = 0
-        try:
-            async for document_chunk in document.read(
-                max_chunk_size=max_chunk_size, chunker_cls=chunker
-            ):
-                document_token_count += document_chunk.chunk_size
-                document_chunk.belongs_to_set = document.belongs_to_set
-                yield document_chunk
-
-            await update_document_token_count(document.id, document_token_count)
-        except PyPdfInternalError:
-            pass
+
+        async for document_chunk in document.read(
+            max_chunk_size=max_chunk_size, chunker_cls=chunker
+        ):
+            document_token_count += document_chunk.chunk_size
+            document_chunk.belongs_to_set = document.belongs_to_set
+            yield document_chunk
+
+        await update_document_token_count(document.id, document_token_count)
+
         # todo rita

cognee/tasks/entity_completion/entity_extractors/llm_entity_extractor.py

@@ -6,8 +6,7 @@ from pydantic import BaseModel
 from cognee.infrastructure.entities.BaseEntityExtractor import BaseEntityExtractor
 from cognee.modules.engine.models import Entity
 from cognee.modules.engine.models.EntityType import EntityType
-from cognee.infrastructure.llm.prompts import read_query_prompt, render_prompt
-from cognee.infrastructure.llm.get_llm_client import get_llm_client
+from cognee.infrastructure.llm.LLMGateway import LLMGateway
 
 logger = get_logger("llm_entity_extractor")
 
@@ -51,11 +50,10 @@ class LLMEntityExtractor(BaseEntityExtractor):
         try:
             logger.info(f"Extracting entities from text: {text[:100]}...")
 
-            llm_client = get_llm_client()
-            user_prompt = render_prompt(self.user_prompt_template, {"text": text})
-            system_prompt = read_query_prompt(self.system_prompt_template)
+            user_prompt = LLMGateway.render_prompt(self.user_prompt_template, {"text": text})
+            system_prompt = LLMGateway.read_query_prompt(self.system_prompt_template)
 
-            response = await llm_client.acreate_structured_output(
+            response = await LLMGateway.acreate_structured_output(
                 text_input=user_prompt,
                 system_prompt=system_prompt,
                 response_model=EntityList,

cognee/tasks/graph/cascade_extract/utils/extract_content_nodes_and_relationship_names.py

@@ -1,8 +1,7 @@
 from typing import List, Tuple
 from pydantic import BaseModel
 
-from cognee.infrastructure.llm.get_llm_client import get_llm_client
-from cognee.infrastructure.llm.prompts import render_prompt, read_query_prompt
+from cognee.infrastructure.llm.LLMGateway import LLMGateway
 from cognee.root_dir import get_absolute_path
 
 
@@ -17,7 +16,6 @@ async def extract_content_nodes_and_relationship_names(
     content: str, existing_nodes: List[str], n_rounds: int = 2
 ) -> Tuple[List[str], List[str]]:
     """Extracts node names and relationship_names from content through multiple rounds of analysis."""
-    llm_client = get_llm_client()
     all_nodes: List[str] = existing_nodes.copy()
     all_relationship_names: List[str] = []
     existing_node_set = {node.lower() for node in all_nodes}
@@ -34,15 +32,15 @@ async def extract_content_nodes_and_relationship_names(
         }
 
         base_directory = get_absolute_path("./tasks/graph/cascade_extract/prompts")
-        text_input = render_prompt(
+        text_input = LLMGateway.render_prompt(
             "extract_graph_relationship_names_prompt_input.txt",
             context,
             base_directory=base_directory,
         )
-        system_prompt = read_query_prompt(
+        system_prompt = LLMGateway.read_query_prompt(
             "extract_graph_relationship_names_prompt_system.txt", base_directory=base_directory
         )
-        response = await llm_client.acreate_structured_output(
+        response = await LLMGateway.acreate_structured_output(
             text_input=text_input,
             system_prompt=system_prompt,
             response_model=PotentialNodesAndRelationshipNames,

cognee/tasks/graph/cascade_extract/utils/extract_edge_triplets.py

@@ -1,6 +1,6 @@
-from typing import List, Tuple
-from cognee.infrastructure.llm.get_llm_client import get_llm_client
-from cognee.infrastructure.llm.prompts import render_prompt, read_query_prompt
+from typing import List
+
+from cognee.infrastructure.llm.LLMGateway import LLMGateway
 from cognee.shared.data_models import KnowledgeGraph
 from cognee.root_dir import get_absolute_path
 
@@ -9,7 +9,6 @@ async def extract_edge_triplets(
     content: str, nodes: List[str], relationship_names: List[str], n_rounds: int = 2
 ) -> KnowledgeGraph:
     """Creates a knowledge graph by identifying relationships between the provided nodes."""
-    llm_client = get_llm_client()
     final_graph = KnowledgeGraph(nodes=[], edges=[])
     existing_nodes = set()
     existing_node_ids = set()
@@ -27,13 +26,13 @@ async def extract_edge_triplets(
         }
 
         base_directory = get_absolute_path("./tasks/graph/cascade_extract/prompts")
-        text_input = render_prompt(
+        text_input = LLMGateway.render_prompt(
             "extract_graph_edge_triplets_prompt_input.txt", context, base_directory=base_directory
         )
-        system_prompt = read_query_prompt(
+        system_prompt = LLMGateway.read_query_prompt(
             "extract_graph_edge_triplets_prompt_system.txt", base_directory=base_directory
         )
-        extracted_graph = await llm_client.acreate_structured_output(
+        extracted_graph = await LLMGateway.acreate_structured_output(
             text_input=text_input, system_prompt=system_prompt, response_model=KnowledgeGraph
         )
 

cognee/tasks/graph/cascade_extract/utils/extract_nodes.py

@@ -1,9 +1,7 @@
 from typing import List
 from pydantic import BaseModel
 
-from cognee.modules.chunking.models.DocumentChunk import DocumentChunk
-from cognee.infrastructure.llm.get_llm_client import get_llm_client
-from cognee.infrastructure.llm.prompts import render_prompt, read_query_prompt
+from cognee.infrastructure.llm.LLMGateway import LLMGateway
 from cognee.root_dir import get_absolute_path
 
 
@@ -15,7 +13,6 @@ class PotentialNodes(BaseModel):
 
 async def extract_nodes(text: str, n_rounds: int = 2) -> List[str]:
     """Extracts node names from content through multiple rounds of analysis."""
-    llm_client = get_llm_client()
     all_nodes: List[str] = []
     existing_nodes = set()
 
@@ -27,13 +24,13 @@ async def extract_nodes(text: str, n_rounds: int = 2) -> List[str]:
             "text": text,
         }
         base_directory = get_absolute_path("./tasks/graph/cascade_extract/prompts")
-        text_input = render_prompt(
+        text_input = LLMGateway.render_prompt(
             "extract_graph_nodes_prompt_input.txt", context, base_directory=base_directory
         )
-        system_prompt = read_query_prompt(
+        system_prompt = LLMGateway.read_query_prompt(
             "extract_graph_nodes_prompt_system.txt", base_directory=base_directory
         )
-        response = await llm_client.acreate_structured_output(
+        response = await LLMGateway.acreate_structured_output(
             text_input=text_input, system_prompt=system_prompt, response_model=PotentialNodes
         )