From f1144abc543212f2178ea31d21d635fc617381e6 Mon Sep 17 00:00:00 2001 From: Igor Ilic Date: Tue, 9 Sep 2025 13:50:16 +0200 Subject: [PATCH] refactor: remove LLMGateway usage where not needed --- .../evaluation/direct_llm_eval_adapter.py | 5 +-- cognee/infrastructure/llm/config.py | 6 ++-- .../llm/anthropic/adapter.py | 32 ------------------- .../litellm_instructor/llm/gemini/adapter.py | 31 ------------------ cognee/modules/retrieval/code_retriever.py | 3 +- .../graph_completion_cot_retriever.py | 11 ++++--- .../retrieval/natural_language_retriever.py | 3 +- .../modules/retrieval/temporal_retriever.py | 3 +- cognee/modules/retrieval/utils/completion.py | 11 +++---- .../chunk_naive_llm_classifier.py | 4 +-- .../codingagents/coding_rule_associations.py | 5 +-- .../entity_extractors/llm_entity_extractor.py | 5 +-- ...ct_content_nodes_and_relationship_names.py | 5 +-- .../utils/extract_edge_triplets.py | 5 +-- .../cascade_extract/utils/extract_nodes.py | 5 +-- cognee/tasks/graph/extract_graph_from_code.py | 4 +-- cognee/tasks/summarization/summarize_code.py | 4 +-- examples/python/graphiti_example.py | 5 +-- 18 files changed, 46 insertions(+), 101 deletions(-) diff --git a/cognee/eval_framework/evaluation/direct_llm_eval_adapter.py b/cognee/eval_framework/evaluation/direct_llm_eval_adapter.py index 971f38d45..12d099b79 100644 --- a/cognee/eval_framework/evaluation/direct_llm_eval_adapter.py +++ b/cognee/eval_framework/evaluation/direct_llm_eval_adapter.py @@ -3,6 +3,7 @@ from pydantic import BaseModel from cognee.eval_framework.evaluation.base_eval_adapter import BaseEvalAdapter from cognee.eval_framework.eval_config import EvalConfig +from cognee.infrastructure.llm.prompts import render_prompt, read_query_prompt from cognee.infrastructure.llm import LLMGateway @@ -25,8 +26,8 @@ class DirectLLMEvalAdapter(BaseEvalAdapter): ) -> Dict[str, Any]: args = {"question": question, "answer": answer, "golden_answer": golden_answer} - user_prompt = LLMGateway.render_prompt(self.eval_prompt_path, args) - system_prompt = LLMGateway.read_query_prompt(self.system_prompt_path) + user_prompt = render_prompt(self.eval_prompt_path, args) + system_prompt = read_query_prompt(self.system_prompt_path) evaluation = await LLMGateway.acreate_structured_output( text_input=user_prompt, diff --git a/cognee/infrastructure/llm/config.py b/cognee/infrastructure/llm/config.py index 7aa8f33f7..51ba0f210 100644 --- a/cognee/infrastructure/llm/config.py +++ b/cognee/infrastructure/llm/config.py @@ -76,10 +76,10 @@ class LLMConfig(BaseSettings): provider=self.baml_llm_provider, options={ "model": self.baml_llm_model, - "temperature": self.baml_llm_temperature, + # "temperature": self.baml_llm_temperature, "api_key": self.baml_llm_api_key, - "base_url": self.baml_llm_endpoint, - "api_version": self.baml_llm_api_version, + # "base_url": self.baml_llm_endpoint, + # "api_version": self.baml_llm_api_version, }, ) # Sets the primary client diff --git a/cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/anthropic/adapter.py b/cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/anthropic/adapter.py index 103ac5e08..7f34fb84b 100644 --- a/cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/anthropic/adapter.py +++ b/cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/anthropic/adapter.py @@ -68,35 +68,3 @@ class AnthropicAdapter(LLMInterface): ], response_model=response_model, ) - - def show_prompt(self, text_input: str, system_prompt: str) -> str: - """ - Format and display the prompt for a user query. - - Parameters: - ----------- - - - text_input (str): The input text from the user, defaults to a placeholder if - empty. - - system_prompt (str): The path to the system prompt to be read and formatted. - - Returns: - -------- - - - str: A formatted string displaying the system prompt and user input. - """ - - if not text_input: - text_input = "No user input provided." - if not system_prompt: - raise MissingSystemPromptPathError() - - system_prompt = LLMGateway.read_query_prompt(system_prompt) - - formatted_prompt = ( - f"""System Prompt:\n{system_prompt}\n\nUser Input:\n{text_input}\n""" - if system_prompt - else None - ) - - return formatted_prompt diff --git a/cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/gemini/adapter.py b/cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/gemini/adapter.py index 2ba94e06f..2c7b73eef 100644 --- a/cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/gemini/adapter.py +++ b/cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/gemini/adapter.py @@ -113,34 +113,3 @@ class GeminiAdapter(LLMInterface): logger.error(f"Schema validation failed: {str(e)}") logger.debug(f"Raw response: {e.raw_response}") raise ValueError(f"Response failed schema validation: {str(e)}") - - def show_prompt(self, text_input: str, system_prompt: str) -> str: - """ - Format and display the prompt for a user query. - - Raises an MissingQueryParameterError if no system prompt is provided. - - Parameters: - ----------- - - - text_input (str): The user input text to display along with the system prompt. - - system_prompt (str): The path or content of the system prompt to be read and - displayed. - - Returns: - -------- - - - str: Returns a formatted string containing the system prompt and user input. - """ - if not text_input: - text_input = "No user input provided." - if not system_prompt: - raise MissingSystemPromptPathError() - system_prompt = LLMGateway.read_query_prompt(system_prompt) - - formatted_prompt = ( - f"""System Prompt:\n{system_prompt}\n\nUser Input:\n{text_input}\n""" - if system_prompt - else None - ) - return formatted_prompt diff --git a/cognee/modules/retrieval/code_retriever.py b/cognee/modules/retrieval/code_retriever.py index 76b5e758c..d178e7ddd 100644 --- a/cognee/modules/retrieval/code_retriever.py +++ b/cognee/modules/retrieval/code_retriever.py @@ -7,6 +7,7 @@ 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") @@ -41,7 +42,7 @@ class CodeRetriever(BaseRetriever): f"Processing query with LLM: '{query[:100]}{'...' if len(query) > 100 else ''}'" ) - system_prompt = LLMGateway.read_query_prompt("codegraph_retriever_system.txt") + system_prompt = read_query_prompt("codegraph_retriever_system.txt") try: result = await LLMGateway.acreate_structured_output( diff --git a/cognee/modules/retrieval/graph_completion_cot_retriever.py b/cognee/modules/retrieval/graph_completion_cot_retriever.py index 7e14078e4..59f9af359 100644 --- a/cognee/modules/retrieval/graph_completion_cot_retriever.py +++ b/cognee/modules/retrieval/graph_completion_cot_retriever.py @@ -1,9 +1,10 @@ -from typing import Any, Optional, List, Tuple, Type +from typing import Any, Optional, List, Type from cognee.shared.logging_utils import get_logger from cognee.modules.retrieval.graph_completion_retriever import GraphCompletionRetriever from cognee.modules.retrieval.utils.completion import generate_completion from cognee.infrastructure.llm.LLMGateway import LLMGateway +from cognee.infrastructure.llm.prompts import render_prompt, read_query_prompt logger = get_logger() @@ -104,10 +105,10 @@ class GraphCompletionCotRetriever(GraphCompletionRetriever): logger.info(f"Chain-of-thought: round {round_idx} - answer: {completion}") if round_idx < max_iter: valid_args = {"query": query, "answer": completion, "context": context} - valid_user_prompt = LLMGateway.render_prompt( + valid_user_prompt = render_prompt( filename=self.validation_user_prompt_path, context=valid_args ) - valid_system_prompt = LLMGateway.read_query_prompt( + valid_system_prompt = read_query_prompt( prompt_file_name=self.validation_system_prompt_path ) @@ -117,10 +118,10 @@ class GraphCompletionCotRetriever(GraphCompletionRetriever): response_model=str, ) followup_args = {"query": query, "answer": completion, "reasoning": reasoning} - followup_prompt = LLMGateway.render_prompt( + followup_prompt = render_prompt( filename=self.followup_user_prompt_path, context=followup_args ) - followup_system = LLMGateway.read_query_prompt( + followup_system = read_query_prompt( prompt_file_name=self.followup_system_prompt_path ) diff --git a/cognee/modules/retrieval/natural_language_retriever.py b/cognee/modules/retrieval/natural_language_retriever.py index 9797c28a0..35c9f6666 100644 --- a/cognee/modules/retrieval/natural_language_retriever.py +++ b/cognee/modules/retrieval/natural_language_retriever.py @@ -2,6 +2,7 @@ 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.llm.LLMGateway import LLMGateway +from cognee.infrastructure.llm.prompts import render_prompt 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 @@ -49,7 +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.""" - system_prompt = LLMGateway.render_prompt( + system_prompt = render_prompt( self.system_prompt_path, context={ "edge_schemas": edge_schemas, diff --git a/cognee/modules/retrieval/temporal_retriever.py b/cognee/modules/retrieval/temporal_retriever.py index 0ab8e2ecf..9d2474264 100644 --- a/cognee/modules/retrieval/temporal_retriever.py +++ b/cognee/modules/retrieval/temporal_retriever.py @@ -6,6 +6,7 @@ from operator import itemgetter from cognee.infrastructure.databases.vector import get_vector_engine from cognee.modules.retrieval.utils.completion import generate_completion from cognee.infrastructure.databases.graph import get_graph_engine +from cognee.infrastructure.llm.prompts import render_prompt from cognee.infrastructure.llm import LLMGateway from cognee.modules.retrieval.graph_completion_retriever import GraphCompletionRetriever from cognee.shared.logging_utils import get_logger @@ -72,7 +73,7 @@ class TemporalRetriever(GraphCompletionRetriever): else: base_directory = None - system_prompt = LLMGateway.render_prompt(prompt_path, {}, base_directory=base_directory) + system_prompt = render_prompt(prompt_path, {}, base_directory=base_directory) interval = await LLMGateway.acreate_structured_output(query, system_prompt, QueryInterval) diff --git a/cognee/modules/retrieval/utils/completion.py b/cognee/modules/retrieval/utils/completion.py index 81e636aad..e02301bb3 100644 --- a/cognee/modules/retrieval/utils/completion.py +++ b/cognee/modules/retrieval/utils/completion.py @@ -1,5 +1,6 @@ from typing import Optional from cognee.infrastructure.llm.LLMGateway import LLMGateway +from cognee.infrastructure.llm.prompts import render_prompt, read_query_prompt async def generate_completion( @@ -12,10 +13,8 @@ async def generate_completion( ) -> str: """Generates a completion using LLM with given context and prompts.""" args = {"question": query, "context": context} - user_prompt = LLMGateway.render_prompt(user_prompt_path, args) - system_prompt = ( - system_prompt if system_prompt else LLMGateway.read_query_prompt(system_prompt_path) - ) + user_prompt = render_prompt(user_prompt_path, args) + system_prompt = system_prompt if system_prompt else read_query_prompt(system_prompt_path) if only_context: return context @@ -33,9 +32,7 @@ async def summarize_text( system_prompt: str = None, ) -> str: """Summarizes text using LLM with the specified prompt.""" - system_prompt = ( - system_prompt if system_prompt else LLMGateway.read_query_prompt(system_prompt_path) - ) + system_prompt = system_prompt if system_prompt else read_query_prompt(system_prompt_path) return await LLMGateway.acreate_structured_output( text_input=text, diff --git a/cognee/tasks/chunk_naive_llm_classifier/chunk_naive_llm_classifier.py b/cognee/tasks/chunk_naive_llm_classifier/chunk_naive_llm_classifier.py index 5de81d4f0..f2c9f11ef 100644 --- a/cognee/tasks/chunk_naive_llm_classifier/chunk_naive_llm_classifier.py +++ b/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.infrastructure.llm.LLMGateway import LLMGateway +from cognee.infrastructure.llm.extraction import extract_categories 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( - *[LLMGateway.extract_categories(chunk.text, classification_model) for chunk in data_chunks], + *[extract_categories(chunk.text, classification_model) for chunk in data_chunks], ) classification_data_points = [] diff --git a/cognee/tasks/codingagents/coding_rule_associations.py b/cognee/tasks/codingagents/coding_rule_associations.py index c809bc68f..9264343ba 100644 --- a/cognee/tasks/codingagents/coding_rule_associations.py +++ b/cognee/tasks/codingagents/coding_rule_associations.py @@ -4,6 +4,7 @@ from cognee.infrastructure.databases.graph import get_graph_engine from cognee.infrastructure.databases.vector import get_vector_engine from cognee.low_level import DataPoint +from cognee.infrastructure.llm.prompts import render_prompt from cognee.infrastructure.llm import LLMGateway from cognee.shared.logging_utils import get_logger from cognee.modules.engine.models import NodeSet @@ -106,8 +107,8 @@ async def add_rule_associations( user_context = {"chat": data, "rules": existing_rules} - user_prompt = LLMGateway.render_prompt(user_prompt_location, context=user_context) - system_prompt = LLMGateway.render_prompt(system_prompt_location, context={}) + user_prompt = render_prompt(user_prompt_location, context=user_context) + system_prompt = render_prompt(system_prompt_location, context={}) rule_list = await LLMGateway.acreate_structured_output( text_input=user_prompt, system_prompt=system_prompt, response_model=RuleSet diff --git a/cognee/tasks/entity_completion/entity_extractors/llm_entity_extractor.py b/cognee/tasks/entity_completion/entity_extractors/llm_entity_extractor.py index 2b5274b9e..9d9fd1912 100644 --- a/cognee/tasks/entity_completion/entity_extractors/llm_entity_extractor.py +++ b/cognee/tasks/entity_completion/entity_extractors/llm_entity_extractor.py @@ -3,6 +3,7 @@ from typing import List from pydantic import BaseModel +from cognee.infrastructure.llm.prompts import render_prompt, read_query_prompt from cognee.infrastructure.entities.BaseEntityExtractor import BaseEntityExtractor from cognee.modules.engine.models import Entity from cognee.modules.engine.models.EntityType import EntityType @@ -50,8 +51,8 @@ class LLMEntityExtractor(BaseEntityExtractor): try: logger.info(f"Extracting entities from text: {text[:100]}...") - user_prompt = LLMGateway.render_prompt(self.user_prompt_template, {"text": text}) - system_prompt = LLMGateway.read_query_prompt(self.system_prompt_template) + user_prompt = render_prompt(self.user_prompt_template, {"text": text}) + system_prompt = read_query_prompt(self.system_prompt_template) response = await LLMGateway.acreate_structured_output( text_input=user_prompt, diff --git a/cognee/tasks/graph/cascade_extract/utils/extract_content_nodes_and_relationship_names.py b/cognee/tasks/graph/cascade_extract/utils/extract_content_nodes_and_relationship_names.py index ef09375ea..00678caa5 100644 --- a/cognee/tasks/graph/cascade_extract/utils/extract_content_nodes_and_relationship_names.py +++ b/cognee/tasks/graph/cascade_extract/utils/extract_content_nodes_and_relationship_names.py @@ -1,6 +1,7 @@ from typing import List, Tuple from pydantic import BaseModel +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 @@ -32,12 +33,12 @@ async def extract_content_nodes_and_relationship_names( } base_directory = get_absolute_path("./tasks/graph/cascade_extract/prompts") - text_input = LLMGateway.render_prompt( + text_input = render_prompt( "extract_graph_relationship_names_prompt_input.txt", context, base_directory=base_directory, ) - system_prompt = LLMGateway.read_query_prompt( + system_prompt = read_query_prompt( "extract_graph_relationship_names_prompt_system.txt", base_directory=base_directory ) response = await LLMGateway.acreate_structured_output( diff --git a/cognee/tasks/graph/cascade_extract/utils/extract_edge_triplets.py b/cognee/tasks/graph/cascade_extract/utils/extract_edge_triplets.py index 8aadb57c7..7d846802e 100644 --- a/cognee/tasks/graph/cascade_extract/utils/extract_edge_triplets.py +++ b/cognee/tasks/graph/cascade_extract/utils/extract_edge_triplets.py @@ -1,5 +1,6 @@ from typing import List +from cognee.infrastructure.llm.prompts import render_prompt, read_query_prompt from cognee.infrastructure.llm.LLMGateway import LLMGateway from cognee.shared.data_models import KnowledgeGraph from cognee.root_dir import get_absolute_path @@ -26,10 +27,10 @@ async def extract_edge_triplets( } base_directory = get_absolute_path("./tasks/graph/cascade_extract/prompts") - text_input = LLMGateway.render_prompt( + text_input = render_prompt( "extract_graph_edge_triplets_prompt_input.txt", context, base_directory=base_directory ) - system_prompt = LLMGateway.read_query_prompt( + system_prompt = read_query_prompt( "extract_graph_edge_triplets_prompt_system.txt", base_directory=base_directory ) extracted_graph = await LLMGateway.acreate_structured_output( diff --git a/cognee/tasks/graph/cascade_extract/utils/extract_nodes.py b/cognee/tasks/graph/cascade_extract/utils/extract_nodes.py index e1fac4721..7b7218a1c 100644 --- a/cognee/tasks/graph/cascade_extract/utils/extract_nodes.py +++ b/cognee/tasks/graph/cascade_extract/utils/extract_nodes.py @@ -1,6 +1,7 @@ from typing import List from pydantic import BaseModel +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 @@ -24,10 +25,10 @@ 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 = LLMGateway.render_prompt( + text_input = render_prompt( "extract_graph_nodes_prompt_input.txt", context, base_directory=base_directory ) - system_prompt = LLMGateway.read_query_prompt( + system_prompt = read_query_prompt( "extract_graph_nodes_prompt_system.txt", base_directory=base_directory ) response = await LLMGateway.acreate_structured_output( diff --git a/cognee/tasks/graph/extract_graph_from_code.py b/cognee/tasks/graph/extract_graph_from_code.py index 4d0da2e1b..555cb59a3 100644 --- a/cognee/tasks/graph/extract_graph_from_code.py +++ b/cognee/tasks/graph/extract_graph_from_code.py @@ -2,7 +2,7 @@ import asyncio from typing import Type, List from pydantic import BaseModel -from cognee.infrastructure.llm.LLMGateway import LLMGateway +from cognee.infrastructure.llm.extraction import extract_content_graph from cognee.modules.chunking.models.DocumentChunk import DocumentChunk from cognee.tasks.storage import add_data_points @@ -18,7 +18,7 @@ async def extract_graph_from_code( - Graph nodes are stored using the `add_data_points` function for later retrieval or analysis. """ chunk_graphs = await asyncio.gather( - *[LLMGateway.extract_content_graph(chunk.text, graph_model) for chunk in data_chunks] + *[extract_content_graph(chunk.text, graph_model) for chunk in data_chunks] ) for chunk_index, chunk in enumerate(data_chunks): diff --git a/cognee/tasks/summarization/summarize_code.py b/cognee/tasks/summarization/summarize_code.py index 75b1c42c3..c7ad69c2a 100644 --- a/cognee/tasks/summarization/summarize_code.py +++ b/cognee/tasks/summarization/summarize_code.py @@ -3,7 +3,7 @@ from typing import AsyncGenerator, Union from uuid import uuid5 from cognee.infrastructure.engine import DataPoint -from cognee.infrastructure.llm.LLMGateway import LLMGateway +from cognee.infrastructure.llm.extraction import extract_code_summary from .models import CodeSummary @@ -16,7 +16,7 @@ async def summarize_code( code_data_points = [file for file in code_graph_nodes if hasattr(file, "source_code")] file_summaries = await asyncio.gather( - *[LLMGateway.extract_code_summary(file.source_code) for file in code_data_points] + *[extract_code_summary(file.source_code) for file in code_data_points] ) file_summaries_map = { diff --git a/examples/python/graphiti_example.py b/examples/python/graphiti_example.py index 471fd1bb7..ece9c452b 100644 --- a/examples/python/graphiti_example.py +++ b/examples/python/graphiti_example.py @@ -12,6 +12,7 @@ from cognee.tasks.temporal_awareness.index_graphiti_objects import ( ) from cognee.modules.retrieval.utils.brute_force_triplet_search import brute_force_triplet_search from cognee.modules.retrieval.graph_completion_retriever import GraphCompletionRetriever +from cognee.infrastructure.llm.prompts import render_prompt, read_query_prompt from cognee.infrastructure.llm.LLMGateway import LLMGateway from cognee.modules.users.methods import get_default_user @@ -59,8 +60,8 @@ async def main(): "context": context, } - user_prompt = LLMGateway.render_prompt("graph_context_for_question.txt", args) - system_prompt = LLMGateway.read_query_prompt("answer_simple_question_restricted.txt") + user_prompt = render_prompt("graph_context_for_question.txt", args) + system_prompt = read_query_prompt("answer_simple_question_restricted.txt") computed_answer = await LLMGateway.acreate_structured_output( text_input=user_prompt,