From a9ec51691ed3c81b9a84d891489f887185a55498 Mon Sep 17 00:00:00 2001 From: Igor Ilic Date: Tue, 5 Aug 2025 21:12:09 +0200 Subject: [PATCH] refactor: Rename LLMAdapter to LLMGateway --- .../src/codingagents/coding_rule_associations.py | 8 ++++---- .../evaluation/direct_llm_eval_adapter.py | 8 ++++---- .../llm/{LLMAdapter.py => LLMGateway.py} | 2 +- cognee/infrastructure/llm/__init__.py | 2 +- .../extraction/extract_categories.py | 6 +++--- .../extraction/extract_summary.py | 6 +++--- .../knowledge_graph/extract_content_graph.py | 6 +++--- .../litellm_instructor/llm/anthropic/adapter.py | 4 ++-- .../litellm_instructor/llm/gemini/adapter.py | 4 ++-- .../litellm_instructor/llm/llm_interface.py | 4 ++-- .../litellm_instructor/llm/openai/adapter.py | 4 ++-- .../processing/document_types/AudioDocument.py | 4 ++-- .../processing/document_types/ImageDocument.py | 4 ++-- cognee/modules/retrieval/code_retriever.py | 6 +++--- .../retrieval/graph_completion_cot_retriever.py | 14 +++++++------- .../retrieval/natural_language_retriever.py | 6 +++--- cognee/modules/retrieval/utils/completion.py | 12 ++++++------ .../utils/description_to_codepart_search.py | 4 ++-- .../search/operations/select_search_type.py | 4 ++-- .../chunk_naive_llm_classifier.py | 4 ++-- .../entity_extractors/llm_entity_extractor.py | 8 ++++---- ...extract_content_nodes_and_relationship_names.py | 8 ++++---- .../cascade_extract/utils/extract_edge_triplets.py | 8 ++++---- .../graph/cascade_extract/utils/extract_nodes.py | 8 ++++---- cognee/tasks/graph/extract_graph_from_code.py | 4 ++-- cognee/tasks/graph/extract_graph_from_data.py | 4 ++-- cognee/tasks/graph/infer_data_ontology.py | 6 +++--- cognee/tasks/summarization/summarize_code.py | 4 ++-- cognee/tasks/summarization/summarize_text.py | 4 ++-- .../agentic_reasoning_procurement_example.py | 4 ++-- examples/python/graphiti_example.py | 8 ++++---- 31 files changed, 89 insertions(+), 89 deletions(-) rename cognee/infrastructure/llm/{LLMAdapter.py => LLMGateway.py} (99%) diff --git a/cognee-mcp/src/codingagents/coding_rule_associations.py b/cognee-mcp/src/codingagents/coding_rule_associations.py index 19d94b9f9..a34c9233b 100644 --- a/cognee-mcp/src/codingagents/coding_rule_associations.py +++ b/cognee-mcp/src/codingagents/coding_rule_associations.py @@ -4,7 +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 import LLMAdapter +from cognee.infrastructure.llm import LLMGateway from cognee.shared.logging_utils import get_logger from cognee.modules.engine.models import NodeSet from cognee.tasks.storage import add_data_points, index_graph_edges @@ -96,12 +96,12 @@ async def add_rule_associations(data: str, rules_nodeset_name: str): user_context = {"chat": data, "rules": existing_rules} - user_prompt = LLMAdapter.render_prompt( + user_prompt = LLMGateway.render_prompt( "coding_rule_association_agent_user.txt", context=user_context ) - system_prompt = LLMAdapter.render_prompt("coding_rule_association_agent_system.txt", context={}) + system_prompt = LLMGateway.render_prompt("coding_rule_association_agent_system.txt", context={}) - rule_list = await LLMAdapter.acreate_structured_output( + rule_list = await LLMGateway.acreate_structured_output( text_input=user_prompt, system_prompt=system_prompt, response_model=RuleSet ) diff --git a/cognee/eval_framework/evaluation/direct_llm_eval_adapter.py b/cognee/eval_framework/evaluation/direct_llm_eval_adapter.py index d9c7b9851..971f38d45 100644 --- a/cognee/eval_framework/evaluation/direct_llm_eval_adapter.py +++ b/cognee/eval_framework/evaluation/direct_llm_eval_adapter.py @@ -3,7 +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 import LLMAdapter +from cognee.infrastructure.llm import LLMGateway class CorrectnessEvaluation(BaseModel): @@ -25,10 +25,10 @@ class DirectLLMEvalAdapter(BaseEvalAdapter): ) -> Dict[str, Any]: args = {"question": question, "answer": answer, "golden_answer": golden_answer} - user_prompt = LLMAdapter.render_prompt(self.eval_prompt_path, args) - system_prompt = LLMAdapter.read_query_prompt(self.system_prompt_path) + user_prompt = LLMGateway.render_prompt(self.eval_prompt_path, args) + system_prompt = LLMGateway.read_query_prompt(self.system_prompt_path) - evaluation = await LLMAdapter.acreate_structured_output( + evaluation = await LLMGateway.acreate_structured_output( text_input=user_prompt, system_prompt=system_prompt, response_model=CorrectnessEvaluation, diff --git a/cognee/infrastructure/llm/LLMAdapter.py b/cognee/infrastructure/llm/LLMGateway.py similarity index 99% rename from cognee/infrastructure/llm/LLMAdapter.py rename to cognee/infrastructure/llm/LLMGateway.py index 24c67068f..e7941dcce 100644 --- a/cognee/infrastructure/llm/LLMAdapter.py +++ b/cognee/infrastructure/llm/LLMGateway.py @@ -4,7 +4,7 @@ from typing import Coroutine from cognee.infrastructure.llm import get_llm_config -class LLMAdapter: +class LLMGateway: """ Class handles selection of structured output frameworks and LLM functions. Class used as a namespace for LLM related functions, should not be instantiated, all methods are static. diff --git a/cognee/infrastructure/llm/__init__.py b/cognee/infrastructure/llm/__init__.py index 51b5c816e..3f866d9f8 100644 --- a/cognee/infrastructure/llm/__init__.py +++ b/cognee/infrastructure/llm/__init__.py @@ -11,4 +11,4 @@ from cognee.infrastructure.llm.utils import ( test_embedding_connection, ) -from cognee.infrastructure.llm.LLMAdapter import LLMAdapter +from cognee.infrastructure.llm.LLMGateway import LLMGateway diff --git a/cognee/infrastructure/llm/structured_output_framework/litellm_instructor/extraction/extract_categories.py b/cognee/infrastructure/llm/structured_output_framework/litellm_instructor/extraction/extract_categories.py index cd9646a72..b723ca0c1 100644 --- a/cognee/infrastructure/llm/structured_output_framework/litellm_instructor/extraction/extract_categories.py +++ b/cognee/infrastructure/llm/structured_output_framework/litellm_instructor/extraction/extract_categories.py @@ -1,12 +1,12 @@ from typing import Type from pydantic import BaseModel -from cognee.infrastructure.llm.LLMAdapter import LLMAdapter +from cognee.infrastructure.llm.LLMGateway import LLMGateway async def extract_categories(content: str, response_model: Type[BaseModel]): - system_prompt = LLMAdapter.read_query_prompt("classify_content.txt") + system_prompt = LLMGateway.read_query_prompt("classify_content.txt") - llm_output = await LLMAdapter.acreate_structured_output(content, system_prompt, response_model) + llm_output = await LLMGateway.acreate_structured_output(content, system_prompt, response_model) return llm_output diff --git a/cognee/infrastructure/llm/structured_output_framework/litellm_instructor/extraction/extract_summary.py b/cognee/infrastructure/llm/structured_output_framework/litellm_instructor/extraction/extract_summary.py index 672a3a78a..341f73736 100644 --- a/cognee/infrastructure/llm/structured_output_framework/litellm_instructor/extraction/extract_summary.py +++ b/cognee/infrastructure/llm/structured_output_framework/litellm_instructor/extraction/extract_summary.py @@ -5,7 +5,7 @@ from typing import Type from instructor.exceptions import InstructorRetryException from pydantic import BaseModel -from cognee.infrastructure.llm.LLMAdapter import LLMAdapter +from cognee.infrastructure.llm.LLMGateway import LLMGateway from cognee.shared.data_models import SummarizedCode logger = get_logger("extract_summary") @@ -25,9 +25,9 @@ def get_mock_summarized_code(): async def extract_summary(content: str, response_model: Type[BaseModel]): - system_prompt = LLMAdapter.read_query_prompt("summarize_content.txt") + system_prompt = LLMGateway.read_query_prompt("summarize_content.txt") - llm_output = await LLMAdapter.acreate_structured_output(content, system_prompt, response_model) + llm_output = await LLMGateway.acreate_structured_output(content, system_prompt, response_model) return llm_output diff --git a/cognee/infrastructure/llm/structured_output_framework/litellm_instructor/extraction/knowledge_graph/extract_content_graph.py b/cognee/infrastructure/llm/structured_output_framework/litellm_instructor/extraction/knowledge_graph/extract_content_graph.py index 3eabd2873..9b945d167 100644 --- a/cognee/infrastructure/llm/structured_output_framework/litellm_instructor/extraction/knowledge_graph/extract_content_graph.py +++ b/cognee/infrastructure/llm/structured_output_framework/litellm_instructor/extraction/knowledge_graph/extract_content_graph.py @@ -2,7 +2,7 @@ import os from typing import Type from pydantic import BaseModel -from cognee.infrastructure.llm.LLMAdapter import LLMAdapter +from cognee.infrastructure.llm.LLMGateway import LLMGateway from cognee.infrastructure.llm.config import ( get_llm_config, ) @@ -22,9 +22,9 @@ async def extract_content_graph(content: str, response_model: Type[BaseModel]): else: base_directory = None - system_prompt = LLMAdapter.render_prompt(prompt_path, {}, base_directory=base_directory) + system_prompt = LLMGateway.render_prompt(prompt_path, {}, base_directory=base_directory) - content_graph = await LLMAdapter.acreate_structured_output( + content_graph = await LLMGateway.acreate_structured_output( content, system_prompt, response_model ) 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 5a3b2d5f1..636e6c0f2 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 @@ -11,7 +11,7 @@ from cognee.infrastructure.llm.structured_output_framework.litellm_instructor.ll sleep_and_retry_async, ) -from cognee.infrastructure.llm.LLMAdapter import LLMAdapter +from cognee.infrastructure.llm.LLMGateway import LLMGateway class AnthropicAdapter(LLMInterface): @@ -91,7 +91,7 @@ class AnthropicAdapter(LLMInterface): if not system_prompt: raise InvalidValueError(message="No system prompt path provided.") - system_prompt = LLMAdapter.read_query_prompt(system_prompt) + system_prompt = LLMGateway.read_query_prompt(system_prompt) formatted_prompt = ( f"""System Prompt:\n{system_prompt}\n\nUser Input:\n{text_input}\n""" 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 139d92980..61d42ff5f 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 @@ -9,7 +9,7 @@ from cognee.exceptions import InvalidValueError from cognee.infrastructure.llm.structured_output_framework.litellm_instructor.llm.llm_interface import ( LLMInterface, ) -from cognee.infrastructure.llm.LLMAdapter import LLMAdapter +from cognee.infrastructure.llm.LLMGateway import LLMGateway from cognee.infrastructure.llm.structured_output_framework.litellm_instructor.llm.rate_limiter import ( rate_limit_async, sleep_and_retry_async, @@ -136,7 +136,7 @@ class GeminiAdapter(LLMInterface): text_input = "No user input provided." if not system_prompt: raise InvalidValueError(message="No system prompt path provided.") - system_prompt = LLMAdapter.read_query_prompt(system_prompt) + system_prompt = LLMGateway.read_query_prompt(system_prompt) formatted_prompt = ( f"""System Prompt:\n{system_prompt}\n\nUser Input:\n{text_input}\n""" diff --git a/cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/llm_interface.py b/cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/llm_interface.py index 7d6780d19..853716a93 100644 --- a/cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/llm_interface.py +++ b/cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/llm_interface.py @@ -3,7 +3,7 @@ from typing import Type, Protocol from abc import abstractmethod from pydantic import BaseModel -from cognee.infrastructure.llm.LLMAdapter import LLMAdapter +from cognee.infrastructure.llm.LLMGateway import LLMGateway class LLMInterface(Protocol): @@ -57,7 +57,7 @@ class LLMInterface(Protocol): text_input = "No user input provided." if not system_prompt: raise ValueError("No system prompt path provided.") - system_prompt = LLMAdapter.read_query_prompt(system_prompt) + system_prompt = LLMGateway.read_query_prompt(system_prompt) formatted_prompt = f"""System Prompt:\n{system_prompt}\n\nUser Input:\n{text_input}\n""" diff --git a/cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/openai/adapter.py b/cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/openai/adapter.py index 7fd0e0d95..69c1bac1c 100644 --- a/cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/openai/adapter.py +++ b/cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/openai/adapter.py @@ -8,7 +8,7 @@ from litellm.exceptions import ContentPolicyViolationError from instructor.exceptions import InstructorRetryException from cognee.exceptions import InvalidValueError -from cognee.infrastructure.llm.LLMAdapter import LLMAdapter +from cognee.infrastructure.llm.LLMGateway import LLMGateway from cognee.infrastructure.llm.structured_output_framework.litellm_instructor.llm.llm_interface import ( LLMInterface, ) @@ -326,7 +326,7 @@ class OpenAIAdapter(LLMInterface): text_input = "No user input provided." if not system_prompt: raise InvalidValueError(message="No system prompt path provided.") - system_prompt = LLMAdapter.read_query_prompt(system_prompt) + system_prompt = LLMGateway.read_query_prompt(system_prompt) formatted_prompt = ( f"""System Prompt:\n{system_prompt}\n\nUser Input:\n{text_input}\n""" diff --git a/cognee/modules/data/processing/document_types/AudioDocument.py b/cognee/modules/data/processing/document_types/AudioDocument.py index 8d3971cdd..1cc77f19e 100644 --- a/cognee/modules/data/processing/document_types/AudioDocument.py +++ b/cognee/modules/data/processing/document_types/AudioDocument.py @@ -1,5 +1,5 @@ from cognee.modules.chunking.Chunker import Chunker -from cognee.infrastructure.llm.LLMAdapter import LLMAdapter +from cognee.infrastructure.llm.LLMGateway import LLMGateway from .Document import Document @@ -8,7 +8,7 @@ class AudioDocument(Document): type: str = "audio" async def create_transcript(self): - result = await LLMAdapter.create_transcript(self.raw_data_location) + result = await LLMGateway.create_transcript(self.raw_data_location) return result.text async def read(self, chunker_cls: Chunker, max_chunk_size: int): diff --git a/cognee/modules/data/processing/document_types/ImageDocument.py b/cognee/modules/data/processing/document_types/ImageDocument.py index b99ecb8b7..4a0642ba4 100644 --- a/cognee/modules/data/processing/document_types/ImageDocument.py +++ b/cognee/modules/data/processing/document_types/ImageDocument.py @@ -1,4 +1,4 @@ -from cognee.infrastructure.llm.LLMAdapter import LLMAdapter +from cognee.infrastructure.llm.LLMGateway import LLMGateway from cognee.modules.chunking.Chunker import Chunker from .Document import Document @@ -8,7 +8,7 @@ class ImageDocument(Document): type: str = "image" async def transcribe_image(self): - result = await LLMAdapter.transcribe_image(self.raw_data_location) + result = await LLMGateway.transcribe_image(self.raw_data_location) return result.choices[0].message.content async def read(self, chunker_cls: Chunker, max_chunk_size: int): diff --git a/cognee/modules/retrieval/code_retriever.py b/cognee/modules/retrieval/code_retriever.py index a177bc5dd..6e819d8a7 100644 --- a/cognee/modules/retrieval/code_retriever.py +++ b/cognee/modules/retrieval/code_retriever.py @@ -7,7 +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.LLMAdapter import LLMAdapter +from cognee.infrastructure.llm.LLMGateway import LLMGateway logger = get_logger("CodeRetriever") @@ -41,10 +41,10 @@ class CodeRetriever(BaseRetriever): f"Processing query with LLM: '{query[:100]}{'...' if len(query) > 100 else ''}'" ) - system_prompt = LLMAdapter.read_query_prompt("codegraph_retriever_system.txt") + system_prompt = LLMGateway.read_query_prompt("codegraph_retriever_system.txt") try: - result = await LLMAdapter.acreate_structured_output( + result = await LLMGateway.acreate_structured_output( text_input=query, system_prompt=system_prompt, response_model=self.CodeQueryInfo, diff --git a/cognee/modules/retrieval/graph_completion_cot_retriever.py b/cognee/modules/retrieval/graph_completion_cot_retriever.py index 52f62e50b..b3e3bfbd4 100644 --- a/cognee/modules/retrieval/graph_completion_cot_retriever.py +++ b/cognee/modules/retrieval/graph_completion_cot_retriever.py @@ -3,7 +3,7 @@ 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.LLMAdapter import LLMAdapter +from cognee.infrastructure.llm.LLMGateway import LLMGateway logger = get_logger() @@ -94,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 = LLMAdapter.render_prompt( + valid_user_prompt = LLMGateway.render_prompt( filename=self.validation_user_prompt_path, context=valid_args ) - valid_system_prompt = LLMAdapter.read_query_prompt( + valid_system_prompt = LLMGateway.read_query_prompt( prompt_file_name=self.validation_system_prompt_path ) - reasoning = await LLMAdapter.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 = LLMAdapter.render_prompt( + followup_prompt = LLMGateway.render_prompt( filename=self.followup_user_prompt_path, context=followup_args ) - followup_system = LLMAdapter.read_query_prompt( + followup_system = LLMGateway.read_query_prompt( prompt_file_name=self.followup_system_prompt_path ) - followup_question = await LLMAdapter.acreate_structured_output( + followup_question = await LLMGateway.acreate_structured_output( text_input=followup_prompt, system_prompt=followup_system, response_model=str ) logger.info( diff --git a/cognee/modules/retrieval/natural_language_retriever.py b/cognee/modules/retrieval/natural_language_retriever.py index 67760eb99..d9a77bd7e 100644 --- a/cognee/modules/retrieval/natural_language_retriever.py +++ b/cognee/modules/retrieval/natural_language_retriever.py @@ -2,7 +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.databases.graph.networkx.adapter import NetworkXAdapter -from cognee.infrastructure.llm.LLMAdapter import LLMAdapter +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 @@ -50,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 = LLMAdapter.render_prompt( + system_prompt = LLMGateway.render_prompt( self.system_prompt_path, context={ "edge_schemas": edge_schemas, @@ -58,7 +58,7 @@ class NaturalLanguageRetriever(BaseRetriever): }, ) - return await LLMAdapter.acreate_structured_output( + return await LLMGateway.acreate_structured_output( text_input=query, system_prompt=system_prompt, response_model=str, diff --git a/cognee/modules/retrieval/utils/completion.py b/cognee/modules/retrieval/utils/completion.py index 5a87e9a89..ca0b30c18 100644 --- a/cognee/modules/retrieval/utils/completion.py +++ b/cognee/modules/retrieval/utils/completion.py @@ -1,4 +1,4 @@ -from cognee.infrastructure.llm.LLMAdapter import LLMAdapter +from cognee.infrastructure.llm.LLMGateway import LLMGateway async def generate_completion( @@ -9,10 +9,10 @@ async def generate_completion( ) -> str: """Generates a completion using LLM with given context and prompts.""" args = {"question": query, "context": context} - user_prompt = LLMAdapter.render_prompt(user_prompt_path, args) - system_prompt = LLMAdapter.read_query_prompt(system_prompt_path) + user_prompt = LLMGateway.render_prompt(user_prompt_path, args) + system_prompt = LLMGateway.read_query_prompt(system_prompt_path) - return await LLMAdapter.acreate_structured_output( + return await LLMGateway.acreate_structured_output( text_input=user_prompt, system_prompt=system_prompt, response_model=str, @@ -24,9 +24,9 @@ async def summarize_text( prompt_path: str = "summarize_search_results.txt", ) -> str: """Summarizes text using LLM with the specified prompt.""" - system_prompt = LLMAdapter.read_query_prompt(prompt_path) + system_prompt = LLMGateway.read_query_prompt(prompt_path) - return await LLMAdapter.acreate_structured_output( + return await LLMGateway.acreate_structured_output( text_input=text, system_prompt=system_prompt, response_model=str, diff --git a/cognee/modules/retrieval/utils/description_to_codepart_search.py b/cognee/modules/retrieval/utils/description_to_codepart_search.py index 7c285011c..a61feb574 100644 --- a/cognee/modules/retrieval/utils/description_to_codepart_search.py +++ b/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.LLMAdapter import LLMAdapter +from cognee.infrastructure.llm.LLMGateway import LLMGateway logger = get_logger(level=ERROR) @@ -71,7 +71,7 @@ async def code_description_to_code_part( if isinstance(obj, dict) and "description" in obj ) - context_from_documents = await LLMAdapter.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" diff --git a/cognee/modules/search/operations/select_search_type.py b/cognee/modules/search/operations/select_search_type.py index a3013fec1..58e9fe57b 100644 --- a/cognee/modules/search/operations/select_search_type.py +++ b/cognee/modules/search/operations/select_search_type.py @@ -1,7 +1,7 @@ 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.LLMAdapter import LLMAdapter +from cognee.infrastructure.llm.LLMGateway import LLMGateway logger = get_logger("SearchTypeSelector") @@ -24,7 +24,7 @@ async def select_search_type( system_prompt = read_query_prompt(system_prompt_path) try: - response = await LLMAdapter.acreate_structured_output( + response = await LLMGateway.acreate_structured_output( text_input=query, system_prompt=system_prompt, response_model=str, 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 ef0eaf6f8..5de81d4f0 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.LLMAdapter import LLMAdapter +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( - *[LLMAdapter.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 = [] 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 d52e7bf4b..2b5274b9e 100644 --- a/cognee/tasks/entity_completion/entity_extractors/llm_entity_extractor.py +++ b/cognee/tasks/entity_completion/entity_extractors/llm_entity_extractor.py @@ -6,7 +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.LLMAdapter import LLMAdapter +from cognee.infrastructure.llm.LLMGateway import LLMGateway logger = get_logger("llm_entity_extractor") @@ -50,10 +50,10 @@ class LLMEntityExtractor(BaseEntityExtractor): try: logger.info(f"Extracting entities from text: {text[:100]}...") - user_prompt = LLMAdapter.render_prompt(self.user_prompt_template, {"text": text}) - system_prompt = LLMAdapter.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 LLMAdapter.acreate_structured_output( + response = await LLMGateway.acreate_structured_output( text_input=user_prompt, system_prompt=system_prompt, response_model=EntityList, 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 359ace6f3..ef09375ea 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,7 +1,7 @@ from typing import List, Tuple from pydantic import BaseModel -from cognee.infrastructure.llm.LLMAdapter import LLMAdapter +from cognee.infrastructure.llm.LLMGateway import LLMGateway from cognee.root_dir import get_absolute_path @@ -32,15 +32,15 @@ async def extract_content_nodes_and_relationship_names( } base_directory = get_absolute_path("./tasks/graph/cascade_extract/prompts") - text_input = LLMAdapter.render_prompt( + text_input = LLMGateway.render_prompt( "extract_graph_relationship_names_prompt_input.txt", context, base_directory=base_directory, ) - system_prompt = LLMAdapter.read_query_prompt( + system_prompt = LLMGateway.read_query_prompt( "extract_graph_relationship_names_prompt_system.txt", base_directory=base_directory ) - response = await LLMAdapter.acreate_structured_output( + response = await LLMGateway.acreate_structured_output( text_input=text_input, system_prompt=system_prompt, response_model=PotentialNodesAndRelationshipNames, 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 fe8ce6696..8aadb57c7 100644 --- a/cognee/tasks/graph/cascade_extract/utils/extract_edge_triplets.py +++ b/cognee/tasks/graph/cascade_extract/utils/extract_edge_triplets.py @@ -1,6 +1,6 @@ from typing import List -from cognee.infrastructure.llm.LLMAdapter import LLMAdapter +from cognee.infrastructure.llm.LLMGateway import LLMGateway from cognee.shared.data_models import KnowledgeGraph from cognee.root_dir import get_absolute_path @@ -26,13 +26,13 @@ async def extract_edge_triplets( } base_directory = get_absolute_path("./tasks/graph/cascade_extract/prompts") - text_input = LLMAdapter.render_prompt( + text_input = LLMGateway.render_prompt( "extract_graph_edge_triplets_prompt_input.txt", context, base_directory=base_directory ) - system_prompt = LLMAdapter.read_query_prompt( + system_prompt = LLMGateway.read_query_prompt( "extract_graph_edge_triplets_prompt_system.txt", base_directory=base_directory ) - extracted_graph = await LLMAdapter.acreate_structured_output( + extracted_graph = await LLMGateway.acreate_structured_output( text_input=text_input, system_prompt=system_prompt, response_model=KnowledgeGraph ) diff --git a/cognee/tasks/graph/cascade_extract/utils/extract_nodes.py b/cognee/tasks/graph/cascade_extract/utils/extract_nodes.py index 1cda91924..e1fac4721 100644 --- a/cognee/tasks/graph/cascade_extract/utils/extract_nodes.py +++ b/cognee/tasks/graph/cascade_extract/utils/extract_nodes.py @@ -1,7 +1,7 @@ from typing import List from pydantic import BaseModel -from cognee.infrastructure.llm.LLMAdapter import LLMAdapter +from cognee.infrastructure.llm.LLMGateway import LLMGateway from cognee.root_dir import get_absolute_path @@ -24,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 = LLMAdapter.render_prompt( + text_input = LLMGateway.render_prompt( "extract_graph_nodes_prompt_input.txt", context, base_directory=base_directory ) - system_prompt = LLMAdapter.read_query_prompt( + system_prompt = LLMGateway.read_query_prompt( "extract_graph_nodes_prompt_system.txt", base_directory=base_directory ) - response = await LLMAdapter.acreate_structured_output( + response = await LLMGateway.acreate_structured_output( text_input=text_input, system_prompt=system_prompt, response_model=PotentialNodes ) diff --git a/cognee/tasks/graph/extract_graph_from_code.py b/cognee/tasks/graph/extract_graph_from_code.py index 72a6df2bb..4d0da2e1b 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.LLMAdapter import LLMAdapter +from cognee.infrastructure.llm.LLMGateway import LLMGateway 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( - *[LLMAdapter.extract_content_graph(chunk.text, graph_model) for chunk in data_chunks] + *[LLMGateway.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/graph/extract_graph_from_data.py b/cognee/tasks/graph/extract_graph_from_data.py index 581e548b5..aa415d504 100644 --- a/cognee/tasks/graph/extract_graph_from_data.py +++ b/cognee/tasks/graph/extract_graph_from_data.py @@ -11,7 +11,7 @@ from cognee.modules.graph.utils import ( retrieve_existing_edges, ) from cognee.shared.data_models import KnowledgeGraph -from cognee.infrastructure.llm.LLMAdapter import LLMAdapter +from cognee.infrastructure.llm.LLMGateway import LLMGateway async def integrate_chunk_graphs( @@ -56,7 +56,7 @@ async def extract_graph_from_data( Extracts and integrates a knowledge graph from the text content of document chunks using a specified graph model. """ chunk_graphs = await asyncio.gather( - *[LLMAdapter.extract_content_graph(chunk.text, graph_model) for chunk in data_chunks] + *[LLMGateway.extract_content_graph(chunk.text, graph_model) for chunk in data_chunks] ) # Note: Filter edges with missing source or target nodes diff --git a/cognee/tasks/graph/infer_data_ontology.py b/cognee/tasks/graph/infer_data_ontology.py index d19225b04..93b02db9f 100644 --- a/cognee/tasks/graph/infer_data_ontology.py +++ b/cognee/tasks/graph/infer_data_ontology.py @@ -27,7 +27,7 @@ from cognee.modules.data.methods.add_model_class_to_graph import ( from cognee.tasks.graph.models import NodeModel, GraphOntology from cognee.shared.data_models import KnowledgeGraph from cognee.modules.engine.utils import generate_node_id, generate_node_name -from cognee.infrastructure.llm.LLMAdapter import LLMAdapter +from cognee.infrastructure.llm.LLMGateway import LLMGateway logger = get_logger("task:infer_data_ontology") @@ -53,9 +53,9 @@ async def extract_ontology(content: str, response_model: Type[BaseModel]): The structured ontology extracted from the content. """ - system_prompt = LLMAdapter.read_query_prompt("extract_ontology.txt") + system_prompt = LLMGateway.read_query_prompt("extract_ontology.txt") - ontology = await LLMAdapter.acreate_structured_output(content, system_prompt, response_model) + ontology = await LLMGateway.acreate_structured_output(content, system_prompt, response_model) return ontology diff --git a/cognee/tasks/summarization/summarize_code.py b/cognee/tasks/summarization/summarize_code.py index ce48fbad2..75b1c42c3 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.LLMAdapter import LLMAdapter +from cognee.infrastructure.llm.LLMGateway import LLMGateway 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( - *[LLMAdapter.extract_code_summary(file.source_code) for file in code_data_points] + *[LLMGateway.extract_code_summary(file.source_code) for file in code_data_points] ) file_summaries_map = { diff --git a/cognee/tasks/summarization/summarize_text.py b/cognee/tasks/summarization/summarize_text.py index 30a475ed5..ca6964f83 100644 --- a/cognee/tasks/summarization/summarize_text.py +++ b/cognee/tasks/summarization/summarize_text.py @@ -4,7 +4,7 @@ from uuid import uuid5 from pydantic import BaseModel from cognee.modules.chunking.models.DocumentChunk import DocumentChunk -from cognee.infrastructure.llm.LLMAdapter import LLMAdapter +from cognee.infrastructure.llm.LLMGateway import LLMGateway from cognee.modules.cognify.config import get_cognify_config from .models import TextSummary @@ -43,7 +43,7 @@ async def summarize_text( summarization_model = cognee_config.summarization_model chunk_summaries = await asyncio.gather( - *[LLMAdapter.extract_summary(chunk.text, summarization_model) for chunk in data_chunks] + *[LLMGateway.extract_summary(chunk.text, summarization_model) for chunk in data_chunks] ) summaries = [ diff --git a/examples/python/agentic_reasoning_procurement_example.py b/examples/python/agentic_reasoning_procurement_example.py index 6e39d0515..5aa3caa70 100644 --- a/examples/python/agentic_reasoning_procurement_example.py +++ b/examples/python/agentic_reasoning_procurement_example.py @@ -3,7 +3,7 @@ import logging import cognee import asyncio -from cognee.infrastructure.llm.LLMAdapter import LLMAdapter +from cognee.infrastructure.llm.LLMGateway import LLMGateway from dotenv import load_dotenv from cognee.api.v1.search import SearchType from cognee.modules.engine.models import NodeSet @@ -185,7 +185,7 @@ async def run_procurement_example(): print(research_information) print("\nPassing research to LLM for final procurement recommendation...\n") - final_decision = await LLMAdapter.acreate_structured_output( + final_decision = await LLMGateway.acreate_structured_output( text_input=research_information, system_prompt="""You are a procurement decision assistant. Use the provided QA pairs that were collected through a research phase. Recommend the best vendor, based on pricing, delivery, warranty, policy fit, and past performance. Be concise and justify your choice with evidence. diff --git a/examples/python/graphiti_example.py b/examples/python/graphiti_example.py index 6c0ff5f54..471fd1bb7 100644 --- a/examples/python/graphiti_example.py +++ b/examples/python/graphiti_example.py @@ -12,7 +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.LLMAdapter import LLMAdapter +from cognee.infrastructure.llm.LLMGateway import LLMGateway from cognee.modules.users.methods import get_default_user text_list = [ @@ -59,10 +59,10 @@ async def main(): "context": context, } - user_prompt = LLMAdapter.render_prompt("graph_context_for_question.txt", args) - system_prompt = LLMAdapter.read_query_prompt("answer_simple_question_restricted.txt") + user_prompt = LLMGateway.render_prompt("graph_context_for_question.txt", args) + system_prompt = LLMGateway.read_query_prompt("answer_simple_question_restricted.txt") - computed_answer = await LLMAdapter.acreate_structured_output( + computed_answer = await LLMGateway.acreate_structured_output( text_input=user_prompt, system_prompt=system_prompt, response_model=str,