cognee 0.2.3.dev0__py3-none-any.whl → 0.2.4__py3-none-any.whl

cognee/infrastructure/databases/vector/pgvector/PGVectorAdapter.py

@@ -9,7 +9,7 @@ from sqlalchemy.exc import ProgrammingError
 from tenacity import retry, retry_if_exception_type, stop_after_attempt, wait_exponential
 from asyncpg import DeadlockDetectedError, DuplicateTableError, UniqueViolationError
 
-from cognee.exceptions import InvalidValueError
+
 from cognee.shared.logging_utils import get_logger
 from cognee.infrastructure.engine import DataPoint
 from cognee.infrastructure.engine.utils import parse_id
@@ -17,6 +17,7 @@ from cognee.infrastructure.databases.relational import get_relational_engine
 
 from distributed.utils import override_distributed
 from distributed.tasks.queued_add_data_points import queued_add_data_points
+from cognee.infrastructure.databases.exceptions import MissingQueryParameterError
 
 from ...relational.ModelBase import Base
 from ...relational.sqlalchemy.SqlAlchemyAdapter import SQLAlchemyAdapter
@@ -275,7 +276,7 @@ class PGVectorAdapter(SQLAlchemyAdapter, VectorDBInterface):
                 return metadata.tables[collection_name]
             else:
                 raise CollectionNotFoundError(
-                    f"Collection '{collection_name}' not found!", log_level="DEBUG"
+                    f"Collection '{collection_name}' not found!",
                 )
 
     async def retrieve(self, collection_name: str, data_point_ids: List[str]):
@@ -302,7 +303,7 @@ class PGVectorAdapter(SQLAlchemyAdapter, VectorDBInterface):
         with_vector: bool = False,
     ) -> List[ScoredResult]:
         if query_text is None and query_vector is None:
-            raise InvalidValueError(message="One of query_text or query_vector must be provided!")
+            raise MissingQueryParameterError()
 
         if query_text and not query_vector:
             query_vector = (await self.embedding_engine.embed_text([query_text]))[0]

cognee/infrastructure/files/utils/get_data_file_path.py

@@ -5,19 +5,24 @@ from urllib.parse import urlparse
 def get_data_file_path(file_path: str):
     # Check if this is a file URI BEFORE normalizing (which corrupts URIs)
     if file_path.startswith("file://"):
+        # Remove first occurrence of file:// prefix
+        pure_file_path = file_path.replace("file://", "", 1)
         # Normalize the file URI for Windows - replace backslashes with forward slashes
-        normalized_file_uri = os.path.normpath(file_path)
+        normalized_file_uri = os.path.normpath(pure_file_path)
 
-        parsed_url = urlparse(normalized_file_uri)
-
-        # Convert URI path to file system path
+        # Convert path to proper file system path
         if os.name == "nt":  # Windows
             # Handle Windows drive letters correctly
-            fs_path = parsed_url.path
-            if fs_path.startswith("/") and len(fs_path) > 1 and fs_path[2] == ":":
-                fs_path = fs_path[1:]  # Remove leading slash for Windows drive paths
-        else:  # Unix-like systems
-            fs_path = parsed_url.path
+            fs_path = normalized_file_uri
+            if (
+                (fs_path.startswith("/") or fs_path.startswith("\\"))
+                and len(fs_path) > 1
+                and fs_path[2] == ":"
+            ):
+                fs_path = fs_path[1:]
+        else:
+            # Unix - like systems
+            fs_path = normalized_file_uri
 
         # Now split the actual filesystem path
         actual_fs_path = os.path.normpath(fs_path)

cognee/infrastructure/files/utils/get_file_metadata.py

@@ -1,6 +1,7 @@
 import io
 import os.path
 from typing import BinaryIO, TypedDict
+from pathlib import Path
 
 from cognee.shared.logging_utils import get_logger
 from cognee.infrastructure.files.utils.get_file_content_hash import get_file_content_hash
@@ -55,7 +56,7 @@ async def get_file_metadata(file: BinaryIO) -> FileMetadata:
     file_type = guess_file_type(file)
 
     file_path = getattr(file, "name", None) or getattr(file, "full_name", None)
-    file_name = str(file_path).split("/")[-1].split(".")[0] if file_path else None
+    file_name = Path(file_path).stem if file_path else None
 
     # Get file size
     pos = file.tell()  # remember current pointer

cognee/infrastructure/llm/LLMGateway.py

@@ -1,6 +1,5 @@
-from typing import Type
+from typing import Type, Optional, Coroutine
 from pydantic import BaseModel
-from typing import Coroutine
 from cognee.infrastructure.llm import get_llm_config
 
 
@@ -79,7 +78,10 @@ class LLMGateway:
 
     @staticmethod
     def extract_content_graph(
-        content: str, response_model: Type[BaseModel], mode: str = "simple"
+        content: str,
+        response_model: Type[BaseModel],
+        mode: str = "simple",
+        custom_prompt: Optional[str] = None,
     ) -> Coroutine:
         llm_config = get_llm_config()
         if llm_config.structured_output_framework.upper() == "BAML":
@@ -87,13 +89,20 @@ class LLMGateway:
                 extract_content_graph,
             )
 
-            return extract_content_graph(content=content, response_model=response_model, mode=mode)
+            return extract_content_graph(
+                content=content,
+                response_model=response_model,
+                mode=mode,
+                custom_prompt=custom_prompt,
+            )
         else:
             from cognee.infrastructure.llm.structured_output_framework.litellm_instructor.extraction import (
                 extract_content_graph,
             )
 
-            return extract_content_graph(content=content, response_model=response_model)
+            return extract_content_graph(
+                content=content, response_model=response_model, custom_prompt=custom_prompt
+            )
 
     @staticmethod
     def extract_categories(content: str, response_model: Type[BaseModel]) -> Coroutine:

cognee/infrastructure/llm/config.py

@@ -18,7 +18,7 @@ class LLMConfig(BaseSettings):
     - llm_api_version
     - llm_temperature
     - llm_streaming
-    - llm_max_tokens
+    - llm_max_completion_tokens
     - transcription_model
     - graph_prompt_path
     - llm_rate_limit_enabled
@@ -35,16 +35,16 @@ class LLMConfig(BaseSettings):
 
     structured_output_framework: str = "instructor"
     llm_provider: str = "openai"
-    llm_model: str = "gpt-4o-mini"
+    llm_model: str = "gpt-5-mini"
     llm_endpoint: str = ""
     llm_api_key: Optional[str] = None
     llm_api_version: Optional[str] = None
     llm_temperature: float = 0.0
     llm_streaming: bool = False
-    llm_max_tokens: int = 16384
+    llm_max_completion_tokens: int = 16384
 
     baml_llm_provider: str = "openai"
-    baml_llm_model: str = "gpt-4o-mini"
+    baml_llm_model: str = "gpt-5-mini"
     baml_llm_endpoint: str = ""
     baml_llm_api_key: Optional[str] = None
     baml_llm_temperature: float = 0.0
@@ -171,7 +171,7 @@ class LLMConfig(BaseSettings):
             "api_version": self.llm_api_version,
             "temperature": self.llm_temperature,
             "streaming": self.llm_streaming,
-            "max_tokens": self.llm_max_tokens,
+            "max_completion_tokens": self.llm_max_completion_tokens,
             "transcription_model": self.transcription_model,
             "graph_prompt_path": self.graph_prompt_path,
             "rate_limit_enabled": self.llm_rate_limit_enabled,

cognee/infrastructure/llm/exceptions.py

@@ -1,5 +1,33 @@
-from cognee.exceptions.exceptions import CriticalError
+from cognee.exceptions.exceptions import CogneeValidationError
 
 
-class ContentPolicyFilterError(CriticalError):
+class ContentPolicyFilterError(CogneeValidationError):
     pass
+
+
+class LLMAPIKeyNotSetError(CogneeValidationError):
+    """
+    Raised when the LLM API key is not set in the configuration.
+    """
+
+    def __init__(self, message: str = "LLM API key is not set."):
+        super().__init__(message=message, name="LLMAPIKeyNotSetError")
+
+
+class UnsupportedLLMProviderError(CogneeValidationError):
+    """
+    Raised when an unsupported LLM provider is specified in the configuration.
+    """
+
+    def __init__(self, provider: str):
+        message = f"Unsupported LLM provider: {provider}"
+        super().__init__(message=message, name="UnsupportedLLMProviderError")
+
+
+class MissingSystemPromptPathError(CogneeValidationError):
+    def __init__(
+        self,
+        name: str = "MissingSystemPromptPathError",
+    ):
+        message = "No system prompt path provided."
+        super().__init__(message, name)

cognee/infrastructure/llm/structured_output_framework/baml/baml_src/extraction/knowledge_graph/extract_content_graph.py

@@ -1,4 +1,4 @@
-from typing import Type
+from typing import Type, Optional
 from pydantic import BaseModel
 from cognee.infrastructure.llm.config import get_llm_config
 from cognee.shared.logging_utils import get_logger, setup_logging
@@ -6,7 +6,10 @@ from cognee.infrastructure.llm.structured_output_framework.baml.baml_client.asyn
 
 
 async def extract_content_graph(
-    content: str, response_model: Type[BaseModel], mode: str = "simple"
+    content: str,
+    response_model: Type[BaseModel],
+    mode: str = "simple",
+    custom_prompt: Optional[str] = None,
 ):
     config = get_llm_config()
     setup_logging()
@@ -26,8 +29,16 @@ async def extract_content_graph(
     #     return graph
 
     # else:
-    graph = await b.ExtractContentGraphGeneric(
-        content, mode=mode, baml_options={"client_registry": config.baml_registry}
-    )
+    if custom_prompt:
+        graph = await b.ExtractContentGraphGeneric(
+            content,
+            mode="custom",
+            custom_prompt_content=custom_prompt,
+            baml_options={"client_registry": config.baml_registry},
+        )
+    else:
+        graph = await b.ExtractContentGraphGeneric(
+            content, mode=mode, baml_options={"client_registry": config.baml_registry}
+        )
 
     return graph

cognee/infrastructure/llm/structured_output_framework/litellm_instructor/extraction/knowledge_graph/extract_content_graph.py

@@ -1,5 +1,5 @@
 import os
-from typing import Type
+from typing import Type, Optional
 from pydantic import BaseModel
 
 from cognee.infrastructure.llm.LLMGateway import LLMGateway
@@ -8,21 +8,25 @@ from cognee.infrastructure.llm.config import (
 )
 
 
-async def extract_content_graph(content: str, response_model: Type[BaseModel]):
-    llm_config = get_llm_config()
-
-    prompt_path = llm_config.graph_prompt_path
-
-    # Check if the prompt path is an absolute path or just a filename
-    if os.path.isabs(prompt_path):
-        # directory containing the file
-        base_directory = os.path.dirname(prompt_path)
-        # just the filename itself
-        prompt_path = os.path.basename(prompt_path)
+async def extract_content_graph(
+    content: str, response_model: Type[BaseModel], custom_prompt: Optional[str] = None
+):
+    if custom_prompt:
+        system_prompt = custom_prompt
     else:
-        base_directory = None
-
-    system_prompt = LLMGateway.render_prompt(prompt_path, {}, base_directory=base_directory)
+        llm_config = get_llm_config()
+        prompt_path = llm_config.graph_prompt_path
+
+        # Check if the prompt path is an absolute path or just a filename
+        if os.path.isabs(prompt_path):
+            # directory containing the file
+            base_directory = os.path.dirname(prompt_path)
+            # just the filename itself
+            prompt_path = os.path.basename(prompt_path)
+        else:
+            base_directory = None
+
+        system_prompt = LLMGateway.render_prompt(prompt_path, {}, base_directory=base_directory)
 
     content_graph = await LLMGateway.acreate_structured_output(
         content, system_prompt, response_model

cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/anthropic/adapter.py

@@ -2,7 +2,7 @@ from typing import Type
 from pydantic import BaseModel
 import instructor
 
-from cognee.exceptions import InvalidValueError
+from cognee.infrastructure.llm.exceptions import MissingSystemPromptPathError
 from cognee.infrastructure.llm.structured_output_framework.litellm_instructor.llm.llm_interface import (
     LLMInterface,
 )
@@ -23,7 +23,7 @@ class AnthropicAdapter(LLMInterface):
     name = "Anthropic"
     model: str
 
-    def __init__(self, max_tokens: int, model: str = None):
+    def __init__(self, max_completion_tokens: int, model: str = None):
         import anthropic
 
         self.aclient = instructor.patch(
@@ -31,7 +31,7 @@ class AnthropicAdapter(LLMInterface):
         )
 
         self.model = model
-        self.max_tokens = max_tokens
+        self.max_completion_tokens = max_completion_tokens
 
     @sleep_and_retry_async()
     @rate_limit_async
@@ -57,7 +57,7 @@ class AnthropicAdapter(LLMInterface):
 
         return await self.aclient(
             model=self.model,
-            max_tokens=4096,
+            max_completion_tokens=4096,
             max_retries=5,
             messages=[
                 {
@@ -89,7 +89,7 @@ class AnthropicAdapter(LLMInterface):
         if not text_input:
             text_input = "No user input provided."
         if not system_prompt:
-            raise InvalidValueError(message="No system prompt path provided.")
+            raise MissingSystemPromptPathError()
 
         system_prompt = LLMGateway.read_query_prompt(system_prompt)
 

cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/gemini/adapter.py

@@ -5,7 +5,7 @@ from litellm import acompletion, JSONSchemaValidationError
 
 from cognee.shared.logging_utils import get_logger
 from cognee.modules.observability.get_observe import get_observe
-from cognee.exceptions import InvalidValueError
+from cognee.infrastructure.llm.exceptions import MissingSystemPromptPathError
 from cognee.infrastructure.llm.structured_output_framework.litellm_instructor.llm.llm_interface import (
     LLMInterface,
 )
@@ -34,7 +34,7 @@ class GeminiAdapter(LLMInterface):
         self,
         api_key: str,
         model: str,
-        max_tokens: int,
+        max_completion_tokens: int,
         endpoint: Optional[str] = None,
         api_version: Optional[str] = None,
         streaming: bool = False,
@@ -44,7 +44,7 @@ class GeminiAdapter(LLMInterface):
         self.endpoint = endpoint
         self.api_version = api_version
         self.streaming = streaming
-        self.max_tokens = max_tokens
+        self.max_completion_tokens = max_completion_tokens
 
     @observe(as_type="generation")
     @sleep_and_retry_async()
@@ -90,7 +90,7 @@ class GeminiAdapter(LLMInterface):
                     model=f"{self.model}",
                     messages=messages,
                     api_key=self.api_key,
-                    max_tokens=self.max_tokens,
+                    max_completion_tokens=self.max_completion_tokens,
                     temperature=0.1,
                     response_format=response_schema,
                     timeout=100,
@@ -118,7 +118,7 @@ class GeminiAdapter(LLMInterface):
         """
         Format and display the prompt for a user query.
 
-        Raises an InvalidValueError if no system prompt is provided.
+        Raises an MissingQueryParameterError if no system prompt is provided.
 
         Parameters:
         -----------
@@ -135,7 +135,7 @@ class GeminiAdapter(LLMInterface):
         if not text_input:
             text_input = "No user input provided."
         if not system_prompt:
-            raise InvalidValueError(message="No system prompt path provided.")
+            raise MissingSystemPromptPathError()
         system_prompt = LLMGateway.read_query_prompt(system_prompt)
 
         formatted_prompt = (

cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/generic_llm_api/adapter.py

@@ -41,7 +41,7 @@ class GenericAPIAdapter(LLMInterface):
         api_key: str,
         model: str,
         name: str,
-        max_tokens: int,
+        max_completion_tokens: int,
         fallback_model: str = None,
         fallback_api_key: str = None,
         fallback_endpoint: str = None,
@@ -50,7 +50,7 @@ class GenericAPIAdapter(LLMInterface):
         self.model = model
         self.api_key = api_key
         self.endpoint = endpoint
-        self.max_tokens = max_tokens
+        self.max_completion_tokens = max_completion_tokens
 
         self.fallback_model = fallback_model
         self.fallback_api_key = fallback_api_key

cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/get_llm_client.py

@@ -2,11 +2,14 @@
 
 from enum import Enum
 
-from cognee.exceptions import InvalidValueError
 from cognee.infrastructure.llm import get_llm_config
 from cognee.infrastructure.llm.structured_output_framework.litellm_instructor.llm.ollama.adapter import (
     OllamaAPIAdapter,
 )
+from cognee.infrastructure.llm.exceptions import (
+    LLMAPIKeyNotSetError,
+    UnsupportedLLMProviderError,
+)
 
 
 # Define an Enum for LLM Providers
@@ -35,7 +38,7 @@ def get_llm_client():
 
     This function retrieves the configuration for the LLM provider and model, and
     initializes the appropriate LLM client adapter accordingly. It raises an
-    InvalidValueError if the LLM API key is not set for certain providers or if the provider
+    LLMAPIKeyNotSetError if the LLM API key is not set for certain providers or if the provider
     is unsupported.
 
     Returns:
@@ -51,15 +54,19 @@ def get_llm_client():
     # Check if max_token value is defined in liteLLM for given model
     # if not use value from cognee configuration
     from cognee.infrastructure.llm.utils import (
-        get_model_max_tokens,
+        get_model_max_completion_tokens,
     )  # imported here to avoid circular imports
 
-    model_max_tokens = get_model_max_tokens(llm_config.llm_model)
-    max_tokens = model_max_tokens if model_max_tokens else llm_config.llm_max_tokens
+    model_max_completion_tokens = get_model_max_completion_tokens(llm_config.llm_model)
+    max_completion_tokens = (
+        model_max_completion_tokens
+        if model_max_completion_tokens
+        else llm_config.llm_max_completion_tokens
+    )
 
     if provider == LLMProvider.OPENAI:
         if llm_config.llm_api_key is None:
-            raise InvalidValueError(message="LLM API key is not set.")
+            raise LLMAPIKeyNotSetError()
 
         from cognee.infrastructure.llm.structured_output_framework.litellm_instructor.llm.openai.adapter import (
             OpenAIAdapter,
@@ -71,7 +78,7 @@ def get_llm_client():
             api_version=llm_config.llm_api_version,
             model=llm_config.llm_model,
             transcription_model=llm_config.transcription_model,
-            max_tokens=max_tokens,
+            max_completion_tokens=max_completion_tokens,
             streaming=llm_config.llm_streaming,
             fallback_api_key=llm_config.fallback_api_key,
             fallback_endpoint=llm_config.fallback_endpoint,
@@ -80,7 +87,7 @@ def get_llm_client():
 
     elif provider == LLMProvider.OLLAMA:
         if llm_config.llm_api_key is None:
-            raise InvalidValueError(message="LLM API key is not set.")
+            raise LLMAPIKeyNotSetError()
 
         from cognee.infrastructure.llm.structured_output_framework.litellm_instructor.llm.generic_llm_api.adapter import (
             GenericAPIAdapter,
@@ -91,7 +98,7 @@ def get_llm_client():
             llm_config.llm_api_key,
             llm_config.llm_model,
             "Ollama",
-            max_tokens=max_tokens,
+            max_completion_tokens=max_completion_tokens,
         )
 
     elif provider == LLMProvider.ANTHROPIC:
@@ -99,11 +106,13 @@ def get_llm_client():
             AnthropicAdapter,
         )
 
-        return AnthropicAdapter(max_tokens=max_tokens, model=llm_config.llm_model)
+        return AnthropicAdapter(
+            max_completion_tokens=max_completion_tokens, model=llm_config.llm_model
+        )
 
     elif provider == LLMProvider.CUSTOM:
         if llm_config.llm_api_key is None:
-            raise InvalidValueError(message="LLM API key is not set.")
+            raise LLMAPIKeyNotSetError()
 
         from cognee.infrastructure.llm.structured_output_framework.litellm_instructor.llm.generic_llm_api.adapter import (
             GenericAPIAdapter,
@@ -114,7 +123,7 @@ def get_llm_client():
             llm_config.llm_api_key,
             llm_config.llm_model,
             "Custom",
-            max_tokens=max_tokens,
+            max_completion_tokens=max_completion_tokens,
             fallback_api_key=llm_config.fallback_api_key,
             fallback_endpoint=llm_config.fallback_endpoint,
             fallback_model=llm_config.fallback_model,
@@ -122,7 +131,7 @@ def get_llm_client():
 
     elif provider == LLMProvider.GEMINI:
         if llm_config.llm_api_key is None:
-            raise InvalidValueError(message="LLM API key is not set.")
+            raise LLMAPIKeyNotSetError()
 
         from cognee.infrastructure.llm.structured_output_framework.litellm_instructor.llm.gemini.adapter import (
             GeminiAdapter,
@@ -131,11 +140,11 @@ def get_llm_client():
         return GeminiAdapter(
             api_key=llm_config.llm_api_key,
             model=llm_config.llm_model,
-            max_tokens=max_tokens,
+            max_completion_tokens=max_completion_tokens,
             endpoint=llm_config.llm_endpoint,
             api_version=llm_config.llm_api_version,
             streaming=llm_config.llm_streaming,
         )
 
     else:
-        raise InvalidValueError(message=f"Unsupported LLM provider: {provider}")
+        raise UnsupportedLLMProviderError(provider)

cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/ollama/adapter.py

@@ -30,16 +30,18 @@ class OllamaAPIAdapter(LLMInterface):
     - model
     - api_key
     - endpoint
-    - max_tokens
+    - max_completion_tokens
     - aclient
     """
 
-    def __init__(self, endpoint: str, api_key: str, model: str, name: str, max_tokens: int):
+    def __init__(
+        self, endpoint: str, api_key: str, model: str, name: str, max_completion_tokens: int
+    ):
         self.name = name
         self.model = model
         self.api_key = api_key
         self.endpoint = endpoint
-        self.max_tokens = max_tokens
+        self.max_completion_tokens = max_completion_tokens
 
         self.aclient = instructor.from_openai(
             OpenAI(base_url=self.endpoint, api_key=self.api_key), mode=instructor.Mode.JSON
@@ -159,7 +161,7 @@ class OllamaAPIAdapter(LLMInterface):
                     ],
                 }
             ],
-            max_tokens=300,
+            max_completion_tokens=300,
         )
 
         # Ensure response is valid before accessing .choices[0].message.content

cognee/infrastructure/llm/structured_output_framework/litellm_instructor/llm/openai/adapter.py

@@ -7,12 +7,14 @@ from openai import ContentFilterFinishReasonError
 from litellm.exceptions import ContentPolicyViolationError
 from instructor.exceptions import InstructorRetryException
 
-from cognee.exceptions import InvalidValueError
 from cognee.infrastructure.llm.LLMGateway import LLMGateway
 from cognee.infrastructure.llm.structured_output_framework.litellm_instructor.llm.llm_interface import (
     LLMInterface,
 )
-from cognee.infrastructure.llm.exceptions import ContentPolicyFilterError
+from cognee.infrastructure.llm.exceptions import (
+    ContentPolicyFilterError,
+    MissingSystemPromptPathError,
+)
 from cognee.infrastructure.files.utils.open_data_file import open_data_file
 from cognee.infrastructure.llm.structured_output_framework.litellm_instructor.llm.rate_limiter import (
     rate_limit_async,
@@ -62,7 +64,7 @@ class OpenAIAdapter(LLMInterface):
         api_version: str,
         model: str,
         transcription_model: str,
-        max_tokens: int,
+        max_completion_tokens: int,
         streaming: bool = False,
         fallback_model: str = None,
         fallback_api_key: str = None,
@@ -75,7 +77,7 @@ class OpenAIAdapter(LLMInterface):
         self.api_key = api_key
         self.endpoint = endpoint
         self.api_version = api_version
-        self.max_tokens = max_tokens
+        self.max_completion_tokens = max_completion_tokens
         self.streaming = streaming
 
         self.fallback_model = fallback_model
@@ -299,7 +301,7 @@ class OpenAIAdapter(LLMInterface):
             api_key=self.api_key,
             api_base=self.endpoint,
             api_version=self.api_version,
-            max_tokens=300,
+            max_completion_tokens=300,
             max_retries=self.MAX_RETRIES,
         )
 
@@ -308,7 +310,7 @@ class OpenAIAdapter(LLMInterface):
         Format and display the prompt for a user query.
 
         This method formats the prompt using the provided user input and system prompt,
-        returning a string representation. Raises InvalidValueError if the system prompt is not
+        returning a string representation. Raises MissingSystemPromptPathError if the system prompt is not
         provided.
 
         Parameters:
@@ -325,7 +327,7 @@ class OpenAIAdapter(LLMInterface):
         if not text_input:
             text_input = "No user input provided."
         if not system_prompt:
-            raise InvalidValueError(message="No system prompt path provided.")
+            raise MissingSystemPromptPathError()
         system_prompt = LLMGateway.read_query_prompt(system_prompt)
 
         formatted_prompt = (

cognee/infrastructure/llm/tokenizer/Gemini/adapter.py

@@ -17,10 +17,10 @@ class GeminiTokenizer(TokenizerInterface):
     def __init__(
         self,
         model: str,
-        max_tokens: int = 3072,
+        max_completion_tokens: int = 3072,
     ):
         self.model = model
-        self.max_tokens = max_tokens
+        self.max_completion_tokens = max_completion_tokens
 
         # Get LLM API key from config
         from cognee.infrastructure.databases.vector.embeddings.config import get_embedding_config

cognee/infrastructure/llm/tokenizer/HuggingFace/adapter.py

@@ -14,17 +14,17 @@ class HuggingFaceTokenizer(TokenizerInterface):
 
     Instance variables include:
     - model: str
-    - max_tokens: int
+    - max_completion_tokens: int
     - tokenizer: AutoTokenizer
     """
 
     def __init__(
         self,
         model: str,
-        max_tokens: int = 512,
+        max_completion_tokens: int = 512,
     ):
         self.model = model
-        self.max_tokens = max_tokens
+        self.max_completion_tokens = max_completion_tokens
 
         # Import here to make it an optional dependency
         from transformers import AutoTokenizer