cognee 0.3.4.dev4__py3-none-any.whl → 0.3.5__py3-none-any.whl

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

@@ -1,146 +1,155 @@
+"""Adapter for Generic API LLM provider API"""
+
 import litellm
+import instructor
+from typing import Type
 from pydantic import BaseModel
-from typing import Type, Optional
-from litellm import acompletion, JSONSchemaValidationError
+from openai import ContentFilterFinishReasonError
+from litellm.exceptions import ContentPolicyViolationError
+from instructor.core import InstructorRetryException
 
-from cognee.shared.logging_utils import get_logger
-from cognee.modules.observability.get_observe import get_observe
-from cognee.infrastructure.llm.exceptions import MissingSystemPromptPathError
+from cognee.infrastructure.llm.exceptions import ContentPolicyFilterError
 from cognee.infrastructure.llm.structured_output_framework.litellm_instructor.llm.llm_interface import (
     LLMInterface,
 )
-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,
 )
 
-logger = get_logger()
-observe = get_observe()
-
 
 class GeminiAdapter(LLMInterface):
     """
-    Handles interactions with a language model API.
+    Adapter for Gemini API LLM provider.
 
-    Public methods include:
-    - acreate_structured_output
-    - show_prompt
+    This class initializes the API adapter with necessary credentials and configurations for
+    interacting with the gemini LLM models. It provides methods for creating structured outputs
+    based on user input and system prompts.
+
+    Public methods:
+    - acreate_structured_output(text_input: str, system_prompt: str, response_model:
+    Type[BaseModel]) -> BaseModel
     """
 
-    MAX_RETRIES = 5
+    name: str
+    model: str
+    api_key: str
 
     def __init__(
         self,
+        endpoint,
         api_key: str,
         model: str,
+        api_version: str,
         max_completion_tokens: int,
-        endpoint: Optional[str] = None,
-        api_version: Optional[str] = None,
-        streaming: bool = False,
-    ) -> None:
-        self.api_key = api_key
+        fallback_model: str = None,
+        fallback_api_key: str = None,
+        fallback_endpoint: str = None,
+    ):
         self.model = model
+        self.api_key = api_key
         self.endpoint = endpoint
         self.api_version = api_version
-        self.streaming = streaming
         self.max_completion_tokens = max_completion_tokens
 
-    @observe(as_type="generation")
+        self.fallback_model = fallback_model
+        self.fallback_api_key = fallback_api_key
+        self.fallback_endpoint = fallback_endpoint
+
+        self.aclient = instructor.from_litellm(litellm.acompletion, mode=instructor.Mode.JSON)
+
     @sleep_and_retry_async()
     @rate_limit_async
     async def acreate_structured_output(
         self, text_input: str, system_prompt: str, response_model: Type[BaseModel]
     ) -> BaseModel:
         """
-        Generate structured output from the language model based on the provided input and
-        system prompt.
+        Generate a response from a user query.
 
-        This method handles retries and raises a ValueError if the request fails or the response
-        does not conform to the expected schema, logging errors accordingly.
+        This asynchronous method sends a user query and a system prompt to a language model and
+        retrieves the generated response. It handles API communication and retries up to a
+        specified limit in case of request failures.
 
         Parameters:
         -----------
 
-            - text_input (str): The user input text to generate a response for.
-            - system_prompt (str): The system's prompt or context to influence the language
-              model's generation.
-            - response_model (Type[BaseModel]): A model type indicating the expected format of
-              the response.
+            - text_input (str): The input text from the user to generate a response for.
+            - system_prompt (str): A prompt that provides context or instructions for the
+              response generation.
+            - response_model (Type[BaseModel]): A Pydantic model that defines the structure of
+              the expected response.
 
         Returns:
         --------
 
-            - BaseModel: Returns the generated response as an instance of the specified response
-              model.
+            - BaseModel: An instance of the specified response model containing the structured
+              output from the language model.
         """
-        try:
-            if response_model is str:
-                response_schema = {"type": "string"}
-            else:
-                response_schema = response_model
 
-            messages = [
-                {"role": "system", "content": system_prompt},
-                {"role": "user", "content": text_input},
-            ]
+        try:
+            return await self.aclient.chat.completions.create(
+                model=self.model,
+                messages=[
+                    {
+                        "role": "user",
+                        "content": f"""{text_input}""",
+                    },
+                    {
+                        "role": "system",
+                        "content": system_prompt,
+                    },
+                ],
+                api_key=self.api_key,
+                max_retries=5,
+                api_base=self.endpoint,
+                api_version=self.api_version,
+                response_model=response_model,
+            )
+        except (
+            ContentFilterFinishReasonError,
+            ContentPolicyViolationError,
+            InstructorRetryException,
+        ) as error:
+            if (
+                isinstance(error, InstructorRetryException)
+                and "content management policy" not in str(error).lower()
+            ):
+                raise error
+
+            if not (self.fallback_model and self.fallback_api_key and self.fallback_endpoint):
+                raise ContentPolicyFilterError(
+                    f"The provided input contains content that is not aligned with our content policy: {text_input}"
+                )
 
             try:
-                response = await acompletion(
-                    model=f"{self.model}",
-                    messages=messages,
-                    api_key=self.api_key,
-                    max_completion_tokens=self.max_completion_tokens,
-                    temperature=0.1,
-                    response_format=response_schema,
-                    timeout=100,
-                    num_retries=self.MAX_RETRIES,
+                return await self.aclient.chat.completions.create(
+                    model=self.fallback_model,
+                    messages=[
+                        {
+                            "role": "user",
+                            "content": f"""{text_input}""",
+                        },
+                        {
+                            "role": "system",
+                            "content": system_prompt,
+                        },
+                    ],
+                    max_retries=5,
+                    api_key=self.fallback_api_key,
+                    api_base=self.fallback_endpoint,
+                    response_model=response_model,
                 )
-
-                if response.choices and response.choices[0].message.content:
-                    content = response.choices[0].message.content
-                    if response_model is str:
-                        return content
-                    return response_model.model_validate_json(content)
-
-            except litellm.exceptions.BadRequestError as e:
-                logger.error(f"Bad request error: {str(e)}")
-                raise ValueError(f"Invalid request: {str(e)}")
-
-            raise ValueError("Failed to get valid response after retries")
-
-        except JSONSchemaValidationError as e:
-            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
+            except (
+                ContentFilterFinishReasonError,
+                ContentPolicyViolationError,
+                InstructorRetryException,
+            ) as error:
+                if (
+                    isinstance(error, InstructorRetryException)
+                    and "content management policy" not in str(error).lower()
+                ):
+                    raise error
+                else:
+                    raise ContentPolicyFilterError(
+                        f"The provided input contains content that is not aligned with our content policy: {text_input}"
+                    )

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

@@ -6,7 +6,7 @@ from typing import Type
 from pydantic import BaseModel
 from openai import ContentFilterFinishReasonError
 from litellm.exceptions import ContentPolicyViolationError
-from instructor.exceptions import InstructorRetryException
+from instructor.core import InstructorRetryException
 
 from cognee.infrastructure.llm.exceptions import ContentPolicyFilterError
 from cognee.infrastructure.llm.structured_output_framework.litellm_instructor.llm.llm_interface import (
@@ -56,9 +56,7 @@ class GenericAPIAdapter(LLMInterface):
         self.fallback_api_key = fallback_api_key
         self.fallback_endpoint = fallback_endpoint
 
-        self.aclient = instructor.from_litellm(
-            litellm.acompletion, mode=instructor.Mode.JSON, api_key=api_key
-        )
+        self.aclient = instructor.from_litellm(litellm.acompletion, mode=instructor.Mode.JSON)
 
     @sleep_and_retry_async()
     @rate_limit_async
@@ -102,6 +100,7 @@ class GenericAPIAdapter(LLMInterface):
                     },
                 ],
                 max_retries=5,
+                api_key=self.api_key,
                 api_base=self.endpoint,
                 response_model=response_model,
             )
@@ -119,7 +118,7 @@ class GenericAPIAdapter(LLMInterface):
             if not (self.fallback_model and self.fallback_api_key and self.fallback_endpoint):
                 raise ContentPolicyFilterError(
                     f"The provided input contains content that is not aligned with our content policy: {text_input}"
-                )
+                ) from error
 
             try:
                 return await self.aclient.chat.completions.create(
@@ -152,4 +151,4 @@ class GenericAPIAdapter(LLMInterface):
                 else:
                     raise ContentPolicyFilterError(
                         f"The provided input contains content that is not aligned with our content policy: {text_input}"
-                    )
+                    ) from error

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

@@ -32,7 +32,7 @@ class LLMProvider(Enum):
     GEMINI = "gemini"
 
 
-def get_llm_client():
+def get_llm_client(raise_api_key_error: bool = True):
     """
     Get the LLM client based on the configuration using Enums.
 
@@ -65,7 +65,7 @@ def get_llm_client():
     )
 
     if provider == LLMProvider.OPENAI:
-        if llm_config.llm_api_key is None:
+        if llm_config.llm_api_key is None and raise_api_key_error:
             raise LLMAPIKeyNotSetError()
 
         from cognee.infrastructure.llm.structured_output_framework.litellm_instructor.llm.openai.adapter import (
@@ -86,7 +86,7 @@ def get_llm_client():
         )
 
     elif provider == LLMProvider.OLLAMA:
-        if llm_config.llm_api_key is None:
+        if llm_config.llm_api_key is None and raise_api_key_error:
             raise LLMAPIKeyNotSetError()
 
         from cognee.infrastructure.llm.structured_output_framework.litellm_instructor.llm.generic_llm_api.adapter import (
@@ -111,7 +111,7 @@ def get_llm_client():
         )
 
     elif provider == LLMProvider.CUSTOM:
-        if llm_config.llm_api_key is None:
+        if llm_config.llm_api_key is None and raise_api_key_error:
             raise LLMAPIKeyNotSetError()
 
         from cognee.infrastructure.llm.structured_output_framework.litellm_instructor.llm.generic_llm_api.adapter import (
@@ -130,7 +130,7 @@ def get_llm_client():
         )
 
     elif provider == LLMProvider.GEMINI:
-        if llm_config.llm_api_key is None:
+        if llm_config.llm_api_key is None and raise_api_key_error:
             raise LLMAPIKeyNotSetError()
 
         from cognee.infrastructure.llm.structured_output_framework.litellm_instructor.llm.gemini.adapter import (
@@ -143,7 +143,6 @@ def get_llm_client():
             max_completion_tokens=max_completion_tokens,
             endpoint=llm_config.llm_endpoint,
             api_version=llm_config.llm_api_version,
-            streaming=llm_config.llm_streaming,
         )
 
     else:

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

@@ -36,29 +36,3 @@ class LLMInterface(Protocol):
               output.
         """
         raise NotImplementedError
-
-    def show_prompt(self, text_input: str, system_prompt: str) -> str:
-        """
-        Format and display the prompt for a user query.
-
-        Parameters:
-        -----------
-
-            - text_input (str): Input text from the user to be included in the prompt.
-            - system_prompt (str): The system prompt that will be shown alongside the user
-              input.
-
-        Returns:
-        --------
-
-            - str: The formatted prompt string combining system prompt and user input.
-        """
-        if not text_input:
-            text_input = "No user input provided."
-        if not system_prompt:
-            raise ValueError("No system prompt path provided.")
-        system_prompt = LLMGateway.read_query_prompt(system_prompt)
-
-        formatted_prompt = f"""System Prompt:\n{system_prompt}\n\nUser Input:\n{text_input}\n"""
-
-        return formatted_prompt

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

@@ -5,15 +5,13 @@ from typing import Type
 from pydantic import BaseModel
 from openai import ContentFilterFinishReasonError
 from litellm.exceptions import ContentPolicyViolationError
-from instructor.exceptions import InstructorRetryException
+from instructor.core import InstructorRetryException
 
-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,
-    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 (
@@ -73,8 +71,19 @@ class OpenAIAdapter(LLMInterface):
         fallback_api_key: str = None,
         fallback_endpoint: str = None,
     ):
-        self.aclient = instructor.from_litellm(litellm.acompletion)
-        self.client = instructor.from_litellm(litellm.completion)
+        # TODO: With gpt5 series models OpenAI expects JSON_SCHEMA as a mode for structured outputs.
+        #       Make sure all new gpt models will work with this mode as well.
+        if "gpt-5" in model:
+            self.aclient = instructor.from_litellm(
+                litellm.acompletion, mode=instructor.Mode.JSON_SCHEMA
+            )
+            self.client = instructor.from_litellm(
+                litellm.completion, mode=instructor.Mode.JSON_SCHEMA
+            )
+        else:
+            self.aclient = instructor.from_litellm(litellm.acompletion)
+            self.client = instructor.from_litellm(litellm.completion)
+
         self.transcription_model = transcription_model
         self.model = model
         self.api_key = api_key
@@ -132,43 +141,16 @@ class OpenAIAdapter(LLMInterface):
                 api_version=self.api_version,
                 response_model=response_model,
                 max_retries=self.MAX_RETRIES,
-                extra_body={"reasoning_effort": "minimal"},
             )
         except (
             ContentFilterFinishReasonError,
             ContentPolicyViolationError,
             InstructorRetryException,
-        ) as error:
-            if (
-                isinstance(error, InstructorRetryException)
-                and "content management policy" not in str(error).lower()
-            ):
-                logger.debug(
-                    "LLM Model does not support reasoning_effort parameter, trying call without the parameter."
-                )
-                return await self.aclient.chat.completions.create(
-                    model=self.model,
-                    messages=[
-                        {
-                            "role": "user",
-                            "content": f"""{text_input}""",
-                        },
-                        {
-                            "role": "system",
-                            "content": system_prompt,
-                        },
-                    ],
-                    api_key=self.api_key,
-                    api_base=self.endpoint,
-                    api_version=self.api_version,
-                    response_model=response_model,
-                    max_retries=self.MAX_RETRIES,
-                )
-
+        ) as e:
             if not (self.fallback_model and self.fallback_api_key):
                 raise ContentPolicyFilterError(
                     f"The provided input contains content that is not aligned with our content policy: {text_input}"
-                )
+                ) from e
 
             try:
                 return await self.aclient.chat.completions.create(
@@ -201,7 +183,7 @@ class OpenAIAdapter(LLMInterface):
                 else:
                     raise ContentPolicyFilterError(
                         f"The provided input contains content that is not aligned with our content policy: {text_input}"
-                    )
+                    ) from error
 
     @observe
     @sleep_and_retry_sync()
@@ -328,35 +310,3 @@ class OpenAIAdapter(LLMInterface):
             max_completion_tokens=300,
             max_retries=self.MAX_RETRIES,
         )
-
-    def show_prompt(self, text_input: str, system_prompt: str) -> str:
-        """
-        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 MissingSystemPromptPathError if the system prompt is not
-        provided.
-
-        Parameters:
-        -----------
-
-            - text_input (str): The input text provided by the user.
-            - system_prompt (str): The system's prompt to guide the model's response.
-
-        Returns:
-        --------
-
-            - str: A formatted string representing the user input and system prompt.
-        """
-        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

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

@@ -3,6 +3,7 @@ from typing import List, Any
 from ..tokenizer_interface import TokenizerInterface
 
 
+# NOTE: DEPRECATED as to count tokens you need to send an API request to Google it is too slow to use with Cognee
 class GeminiTokenizer(TokenizerInterface):
     """
     Implements a tokenizer interface for the Gemini model, managing token extraction and
@@ -16,10 +17,10 @@ class GeminiTokenizer(TokenizerInterface):
 
     def __init__(
         self,
-        model: str,
+        llm_model: str,
         max_completion_tokens: int = 3072,
     ):
-        self.model = model
+        self.llm_model = llm_model
         self.max_completion_tokens = max_completion_tokens
 
         # Get LLM API key from config
@@ -28,12 +29,11 @@ class GeminiTokenizer(TokenizerInterface):
             get_llm_config,
         )
 
-        config = get_embedding_config()
         llm_config = get_llm_config()
 
-        import google.generativeai as genai
+        from google import genai
 
-        genai.configure(api_key=config.embedding_api_key or llm_config.llm_api_key)
+        self.client = genai.Client(api_key=llm_config.llm_api_key)
 
     def extract_tokens(self, text: str) -> List[Any]:
         """
@@ -77,6 +77,7 @@ class GeminiTokenizer(TokenizerInterface):
 
             - int: The number of tokens in the given text.
         """
-        import google.generativeai as genai
 
-        return len(genai.embed_content(model=f"models/{self.model}", content=text))
+        tokens_response = self.client.models.count_tokens(model=self.llm_model, contents=text)
+
+        return tokens_response.total_tokens

cognee/infrastructure/llm/utils.py

@@ -3,6 +3,7 @@ import litellm
 from cognee.infrastructure.llm.structured_output_framework.litellm_instructor.llm.get_llm_client import (
     get_llm_client,
 )
+from cognee.infrastructure.llm.LLMGateway import LLMGateway
 from cognee.shared.logging_utils import get_logger
 
 logger = get_logger()
@@ -28,7 +29,7 @@ def get_max_chunk_tokens():
 
     # Calculate max chunk size based on the following formula
     embedding_engine = get_vector_engine().embedding_engine
-    llm_client = get_llm_client()
+    llm_client = get_llm_client(raise_api_key_error=False)
 
     # We need to make sure chunk size won't take more than half of LLM max context token size
     # but it also can't be bigger than the embedding engine max token size
@@ -62,7 +63,7 @@ def get_model_max_completion_tokens(model_name: str):
         max_completion_tokens = litellm.model_cost[model_name]["max_tokens"]
         logger.debug(f"Max input tokens for {model_name}: {max_completion_tokens}")
     else:
-        logger.info("Model not found in LiteLLM's model_cost.")
+        logger.debug("Model not found in LiteLLM's model_cost.")
 
     return max_completion_tokens
 
@@ -76,8 +77,7 @@ async def test_llm_connection():
     the connection attempt and re-raise the exception for further handling.
     """
     try:
-        llm_adapter = get_llm_client()
-        await llm_adapter.acreate_structured_output(
+        await LLMGateway.acreate_structured_output(
             text_input="test",
             system_prompt='Respond to me with the following string: "test"',
             response_model=str,

cognee/infrastructure/loaders/LoaderEngine.py

@@ -27,6 +27,7 @@ class LoaderEngine:
 
         self.default_loader_priority = [
             "text_loader",
+            "advanced_pdf_loader",
             "pypdf_loader",
             "image_loader",
             "audio_loader",
@@ -86,7 +87,7 @@ class LoaderEngine:
                     if loader.can_handle(extension=file_info.extension, mime_type=file_info.mime):
                         return loader
                 else:
-                    raise ValueError(f"Loader does not exist: {loader_name}")
+                    logger.info(f"Skipping {loader_name}: Preferred Loader not registered")
 
         # Try default priority order
         for loader_name in self.default_loader_priority:
@@ -95,7 +96,9 @@ class LoaderEngine:
                 if loader.can_handle(extension=file_info.extension, mime_type=file_info.mime):
                     return loader
             else:
-                raise ValueError(f"Loader does not exist: {loader_name}")
+                logger.info(
+                    f"Skipping {loader_name}: Loader not registered (in default priority list)."
+                )
 
         return None
 

cognee/infrastructure/loaders/external/__init__.py

@@ -20,3 +20,10 @@ try:
     __all__.append("UnstructuredLoader")
 except ImportError:
     pass
+
+try:
+    from .advanced_pdf_loader import AdvancedPdfLoader
+
+    __all__.append("AdvancedPdfLoader")
+except ImportError:
+    pass