cognee 0.3.5__py3-none-any.whl → 0.3.7__py3-none-any.whl

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

@@ -23,6 +23,7 @@ class LLMProvider(Enum):
     - ANTHROPIC: Represents the Anthropic provider.
     - CUSTOM: Represents a custom provider option.
     - GEMINI: Represents the Gemini provider.
+    - MISTRAL: Represents the Mistral AI provider.
     """
 
     OPENAI = "openai"
@@ -30,6 +31,7 @@ class LLMProvider(Enum):
     ANTHROPIC = "anthropic"
     CUSTOM = "custom"
     GEMINI = "gemini"
+    MISTRAL = "mistral"
 
 
 def get_llm_client(raise_api_key_error: bool = True):
@@ -145,5 +147,35 @@ def get_llm_client(raise_api_key_error: bool = True):
             api_version=llm_config.llm_api_version,
         )
 
+    elif provider == LLMProvider.MISTRAL:
+        if llm_config.llm_api_key is None:
+            raise LLMAPIKeyNotSetError()
+
+        from cognee.infrastructure.llm.structured_output_framework.litellm_instructor.llm.mistral.adapter import (
+            MistralAdapter,
+        )
+
+        return MistralAdapter(
+            api_key=llm_config.llm_api_key,
+            model=llm_config.llm_model,
+            max_completion_tokens=max_completion_tokens,
+            endpoint=llm_config.llm_endpoint,
+        )
+
+    elif provider == LLMProvider.MISTRAL:
+        if llm_config.llm_api_key is None:
+            raise LLMAPIKeyNotSetError()
+
+        from cognee.infrastructure.llm.structured_output_framework.litellm_instructor.llm.mistral.adapter import (
+            MistralAdapter,
+        )
+
+        return MistralAdapter(
+            api_key=llm_config.llm_api_key,
+            model=llm_config.llm_model,
+            max_completion_tokens=max_completion_tokens,
+            endpoint=llm_config.llm_endpoint,
+        )
+
     else:
         raise UnsupportedLLMProviderError(provider)

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

@@ -0,0 +1,109 @@
+import litellm
+import instructor
+from pydantic import BaseModel
+from typing import Type
+from litellm import JSONSchemaValidationError
+
+from cognee.shared.logging_utils import get_logger
+from cognee.modules.observability.get_observe import get_observe
+from cognee.infrastructure.llm.structured_output_framework.litellm_instructor.llm.llm_interface import (
+    LLMInterface,
+)
+from cognee.infrastructure.llm.config import get_llm_config
+
+import logging
+from tenacity import (
+    retry,
+    stop_after_delay,
+    wait_exponential_jitter,
+    retry_if_not_exception_type,
+    before_sleep_log,
+)
+
+logger = get_logger()
+observe = get_observe()
+
+
+class MistralAdapter(LLMInterface):
+    """
+    Adapter for Mistral AI API, for structured output generation and prompt display.
+
+    Public methods:
+    - acreate_structured_output
+    - show_prompt
+    """
+
+    name = "Mistral"
+    model: str
+    api_key: str
+    max_completion_tokens: int
+
+    def __init__(self, api_key: str, model: str, max_completion_tokens: int, endpoint: str = None):
+        from mistralai import Mistral
+
+        self.model = model
+        self.max_completion_tokens = max_completion_tokens
+
+        self.aclient = instructor.from_litellm(
+            litellm.acompletion,
+            mode=instructor.Mode.MISTRAL_TOOLS,
+            api_key=get_llm_config().llm_api_key,
+        )
+
+    @retry(
+        stop=stop_after_delay(128),
+        wait=wait_exponential_jitter(2, 128),
+        retry=retry_if_not_exception_type(litellm.exceptions.NotFoundError),
+        before_sleep=before_sleep_log(logger, logging.DEBUG),
+        reraise=True,
+    )
+    async def acreate_structured_output(
+        self, text_input: str, system_prompt: str, response_model: Type[BaseModel]
+    ) -> BaseModel:
+        """
+        Generate a response from the user query.
+
+        Parameters:
+        -----------
+            - text_input (str): The input text from the user to be processed.
+            - system_prompt (str): A prompt that sets the context for the query.
+            - response_model (Type[BaseModel]): The model to structure the response according to
+              its format.
+
+        Returns:
+        --------
+            - BaseModel: An instance of BaseModel containing the structured response.
+        """
+        try:
+            messages = [
+                {
+                    "role": "system",
+                    "content": system_prompt,
+                },
+                {
+                    "role": "user",
+                    "content": f"""Use the given format to extract information
+                from the following input: {text_input}""",
+                },
+            ]
+            try:
+                response = await self.aclient.chat.completions.create(
+                    model=self.model,
+                    max_tokens=self.max_completion_tokens,
+                    max_retries=5,
+                    messages=messages,
+                    response_model=response_model,
+                )
+                if response.choices and response.choices[0].message.content:
+                    content = response.choices[0].message.content
+                    return response_model.model_validate_json(content)
+                else:
+                    raise ValueError("Failed to get valid response after retries")
+            except litellm.exceptions.BadRequestError as e:
+                logger.error(f"Bad request error: {str(e)}")
+                raise ValueError(f"Invalid request: {str(e)}")
+
+        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)}")

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

@@ -1,4 +1,6 @@
 import base64
+import litellm
+import logging
 import instructor
 from typing import Type
 from openai import OpenAI
@@ -7,11 +9,17 @@ from pydantic import BaseModel
 from cognee.infrastructure.llm.structured_output_framework.litellm_instructor.llm.llm_interface import (
     LLMInterface,
 )
-from cognee.infrastructure.llm.structured_output_framework.litellm_instructor.llm.rate_limiter import (
-    rate_limit_async,
-    sleep_and_retry_async,
-)
 from cognee.infrastructure.files.utils.open_data_file import open_data_file
+from cognee.shared.logging_utils import get_logger
+from tenacity import (
+    retry,
+    stop_after_delay,
+    wait_exponential_jitter,
+    retry_if_not_exception_type,
+    before_sleep_log,
+)
+
+logger = get_logger()
 
 
 class OllamaAPIAdapter(LLMInterface):
@@ -47,8 +55,13 @@ class OllamaAPIAdapter(LLMInterface):
             OpenAI(base_url=self.endpoint, api_key=self.api_key), mode=instructor.Mode.JSON
         )
 
-    @sleep_and_retry_async()
-    @rate_limit_async
+    @retry(
+        stop=stop_after_delay(128),
+        wait=wait_exponential_jitter(2, 128),
+        retry=retry_if_not_exception_type(litellm.exceptions.NotFoundError),
+        before_sleep=before_sleep_log(logger, logging.DEBUG),
+        reraise=True,
+    )
     async def acreate_structured_output(
         self, text_input: str, system_prompt: str, response_model: Type[BaseModel]
     ) -> BaseModel:
@@ -90,7 +103,13 @@ class OllamaAPIAdapter(LLMInterface):
 
         return response
 
-    @rate_limit_async
+    @retry(
+        stop=stop_after_delay(128),
+        wait=wait_exponential_jitter(2, 128),
+        retry=retry_if_not_exception_type(litellm.exceptions.NotFoundError),
+        before_sleep=before_sleep_log(logger, logging.DEBUG),
+        reraise=True,
+    )
     async def create_transcript(self, input_file: str) -> str:
         """
         Generate an audio transcript from a user query.
@@ -123,7 +142,13 @@ class OllamaAPIAdapter(LLMInterface):
 
         return transcription.text
 
-    @rate_limit_async
+    @retry(
+        stop=stop_after_delay(128),
+        wait=wait_exponential_jitter(2, 128),
+        retry=retry_if_not_exception_type(litellm.exceptions.NotFoundError),
+        before_sleep=before_sleep_log(logger, logging.DEBUG),
+        reraise=True,
+    )
     async def transcribe_image(self, input_file: str) -> str:
         """
         Transcribe content from an image using base64 encoding.

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

@@ -7,6 +7,15 @@ from openai import ContentFilterFinishReasonError
 from litellm.exceptions import ContentPolicyViolationError
 from instructor.core import InstructorRetryException
 
+import logging
+from tenacity import (
+    retry,
+    stop_after_delay,
+    wait_exponential_jitter,
+    retry_if_not_exception_type,
+    before_sleep_log,
+)
+
 from cognee.infrastructure.llm.structured_output_framework.litellm_instructor.llm.llm_interface import (
     LLMInterface,
 )
@@ -14,19 +23,13 @@ from cognee.infrastructure.llm.exceptions import (
     ContentPolicyFilterError,
 )
 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,
-    rate_limit_sync,
-    sleep_and_retry_async,
-    sleep_and_retry_sync,
-)
 from cognee.modules.observability.get_observe import get_observe
 from cognee.shared.logging_utils import get_logger
 
-observe = get_observe()
-
 logger = get_logger()
 
+observe = get_observe()
+
 
 class OpenAIAdapter(LLMInterface):
     """
@@ -97,8 +100,13 @@ class OpenAIAdapter(LLMInterface):
         self.fallback_endpoint = fallback_endpoint
 
     @observe(as_type="generation")
-    @sleep_and_retry_async()
-    @rate_limit_async
+    @retry(
+        stop=stop_after_delay(128),
+        wait=wait_exponential_jitter(2, 128),
+        retry=retry_if_not_exception_type(litellm.exceptions.NotFoundError),
+        before_sleep=before_sleep_log(logger, logging.DEBUG),
+        reraise=True,
+    )
     async def acreate_structured_output(
         self, text_input: str, system_prompt: str, response_model: Type[BaseModel]
     ) -> BaseModel:
@@ -148,10 +156,7 @@ class OpenAIAdapter(LLMInterface):
             InstructorRetryException,
         ) 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
-
+                raise e
             try:
                 return await self.aclient.chat.completions.create(
                     model=self.fallback_model,
@@ -186,8 +191,13 @@ class OpenAIAdapter(LLMInterface):
                     ) from error
 
     @observe
-    @sleep_and_retry_sync()
-    @rate_limit_sync
+    @retry(
+        stop=stop_after_delay(128),
+        wait=wait_exponential_jitter(2, 128),
+        retry=retry_if_not_exception_type(litellm.exceptions.NotFoundError),
+        before_sleep=before_sleep_log(logger, logging.DEBUG),
+        reraise=True,
+    )
     def create_structured_output(
         self, text_input: str, system_prompt: str, response_model: Type[BaseModel]
     ) -> BaseModel:
@@ -231,7 +241,13 @@ class OpenAIAdapter(LLMInterface):
             max_retries=self.MAX_RETRIES,
         )
 
-    @rate_limit_async
+    @retry(
+        stop=stop_after_delay(128),
+        wait=wait_exponential_jitter(2, 128),
+        retry=retry_if_not_exception_type(litellm.exceptions.NotFoundError),
+        before_sleep=before_sleep_log(logger, logging.DEBUG),
+        reraise=True,
+    )
     async def create_transcript(self, input):
         """
         Generate an audio transcript from a user query.
@@ -263,7 +279,13 @@ class OpenAIAdapter(LLMInterface):
 
         return transcription
 
-    @rate_limit_async
+    @retry(
+        stop=stop_after_delay(128),
+        wait=wait_exponential_jitter(2, 128),
+        retry=retry_if_not_exception_type(litellm.exceptions.NotFoundError),
+        before_sleep=before_sleep_log(logger, logging.DEBUG),
+        reraise=True,
+    )
     async def transcribe_image(self, input) -> BaseModel:
         """
         Generate a transcription of an image from a user query.

cognee/infrastructure/loaders/LoaderEngine.py

@@ -27,11 +27,11 @@ class LoaderEngine:
 
         self.default_loader_priority = [
             "text_loader",
-            "advanced_pdf_loader",
             "pypdf_loader",
             "image_loader",
             "audio_loader",
             "unstructured_loader",
+            "advanced_pdf_loader",
         ]
 
     def register_loader(self, loader: LoaderInterface) -> bool:
@@ -64,7 +64,9 @@ class LoaderEngine:
         return True
 
     def get_loader(
-        self, file_path: str, preferred_loaders: List[str] = None
+        self,
+        file_path: str,
+        preferred_loaders: dict[str, dict[str, Any]],
     ) -> Optional[LoaderInterface]:
         """
         Get appropriate loader for a file.
@@ -76,14 +78,21 @@ class LoaderEngine:
         Returns:
             LoaderInterface that can handle the file, or None if not found
         """
+        from pathlib import Path
 
         file_info = filetype.guess(file_path)
 
+        path_extension = Path(file_path).suffix.lstrip(".")
+
         # Try preferred loaders first
         if preferred_loaders:
             for loader_name in preferred_loaders:
                 if loader_name in self._loaders:
                     loader = self._loaders[loader_name]
+                    # Try with path extension first (for text formats like html)
+                    if loader.can_handle(extension=path_extension, mime_type=file_info.mime):
+                        return loader
+                    # Fall back to content-detected extension
                     if loader.can_handle(extension=file_info.extension, mime_type=file_info.mime):
                         return loader
                 else:
@@ -93,6 +102,10 @@ class LoaderEngine:
         for loader_name in self.default_loader_priority:
             if loader_name in self._loaders:
                 loader = self._loaders[loader_name]
+                # Try with path extension first (for text formats like html)
+                if loader.can_handle(extension=path_extension, mime_type=file_info.mime):
+                    return loader
+                # Fall back to content-detected extension
                 if loader.can_handle(extension=file_info.extension, mime_type=file_info.mime):
                     return loader
             else:
@@ -105,8 +118,7 @@ class LoaderEngine:
     async def load_file(
         self,
         file_path: str,
-        file_stream: Optional[Any],
-        preferred_loaders: Optional[List[str]] = None,
+        preferred_loaders: dict[str, dict[str, Any]] = None,
         **kwargs,
     ):
         """
@@ -114,7 +126,7 @@ class LoaderEngine:
 
         Args:
             file_path: Path to the file to be processed
-            preferred_loaders: List of preferred loader names to try first
+            preferred_loaders: Dict of loader names to their configurations
             **kwargs: Additional loader-specific configuration
 
         Raises:
@@ -126,8 +138,16 @@ class LoaderEngine:
             raise ValueError(f"No loader found for file: {file_path}")
 
         logger.debug(f"Loading {file_path} with {loader.loader_name}")
-        # TODO: loading needs to be reworked to work with both file streams and file locations
-        return await loader.load(file_path, **kwargs)
+
+        # Extract loader-specific config from preferred_loaders
+        loader_config = {}
+        if preferred_loaders and loader.loader_name in preferred_loaders:
+            loader_config = preferred_loaders[loader.loader_name]
+
+        # Merge with any additional kwargs (kwargs take precedence)
+        merged_kwargs = {**loader_config, **kwargs}
+
+        return await loader.load(file_path, **merged_kwargs)
 
     def get_available_loaders(self) -> List[str]:
         """

cognee/infrastructure/loaders/external/__init__.py

@@ -27,3 +27,10 @@ try:
     __all__.append("AdvancedPdfLoader")
 except ImportError:
     pass
+
+try:
+    from .beautiful_soup_loader import BeautifulSoupLoader
+
+    __all__.append("BeautifulSoupLoader")
+except ImportError:
+    pass

cognee/infrastructure/loaders/external/advanced_pdf_loader.py

@@ -14,14 +14,6 @@ from cognee.infrastructure.loaders.external.pypdf_loader import PyPdfLoader
 
 logger = get_logger(__name__)
 
-try:
-    from unstructured.partition.pdf import partition_pdf
-except ImportError as e:
-    logger.info(
-        "unstructured[pdf] not installed, can't use AdvancedPdfLoader, will use PyPdfLoader instead."
-    )
-    raise ImportError from e
-
 
 @dataclass
 class _PageBuffer:
@@ -88,6 +80,8 @@ class AdvancedPdfLoader(LoaderInterface):
                 **kwargs,
             }
             # Use partition to extract elements
+            from unstructured.partition.pdf import partition_pdf
+
             elements = partition_pdf(**partition_kwargs)
 
             # Process elements into text content