cognee 0.3.6__py3-none-any.whl → 0.3.7.dev1__py3-none-any.whl

cognee/infrastructure/databases/vector/embeddings/OllamaEmbeddingEngine.py

@@ -3,8 +3,16 @@ from cognee.shared.logging_utils import get_logger
 import aiohttp
 from typing import List, Optional
 import os
-
+import litellm
+import logging
 import aiohttp.http_exceptions
+from tenacity import (
+    retry,
+    stop_after_delay,
+    wait_exponential_jitter,
+    retry_if_not_exception_type,
+    before_sleep_log,
+)
 
 from cognee.infrastructure.databases.vector.embeddings.EmbeddingEngine import EmbeddingEngine
 from cognee.infrastructure.llm.tokenizer.HuggingFace import (
@@ -69,7 +77,6 @@ class OllamaEmbeddingEngine(EmbeddingEngine):
             enable_mocking = str(enable_mocking).lower()
         self.mock = enable_mocking in ("true", "1", "yes")
 
-    @embedding_rate_limit_async
     async def embed_text(self, text: List[str]) -> List[List[float]]:
         """
         Generate embedding vectors for a list of text prompts.
@@ -92,7 +99,13 @@ class OllamaEmbeddingEngine(EmbeddingEngine):
         embeddings = await asyncio.gather(*[self._get_embedding(prompt) for prompt in text])
         return embeddings
 
-    @embedding_sleep_and_retry_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 _get_embedding(self, prompt: str) -> List[float]:
         """
         Internal method to call the Ollama embeddings endpoint for a single prompt.
@@ -111,7 +124,7 @@ class OllamaEmbeddingEngine(EmbeddingEngine):
                 self.endpoint, json=payload, headers=headers, timeout=60.0
             ) as response:
                 data = await response.json()
-                return data["embedding"]
+                return data["embeddings"][0]
 
     def get_vector_size(self) -> int:
         """

cognee/infrastructure/databases/vector/embeddings/config.py

@@ -24,11 +24,10 @@ class EmbeddingConfig(BaseSettings):
     model_config = SettingsConfigDict(env_file=".env", extra="allow")
 
     def model_post_init(self, __context) -> None:
-        # If embedding batch size is not defined use 2048 as default for OpenAI and 100 for all other embedding models
         if not self.embedding_batch_size and self.embedding_provider.lower() == "openai":
-            self.embedding_batch_size = 2048
+            self.embedding_batch_size = 36
         elif not self.embedding_batch_size:
-            self.embedding_batch_size = 100
+            self.embedding_batch_size = 36
 
     def to_dict(self) -> dict:
         """

cognee/infrastructure/databases/vector/exceptions/exceptions.py

@@ -15,7 +15,7 @@ class CollectionNotFoundError(CogneeValidationError):
         self,
         message,
         name: str = "CollectionNotFoundError",
-        status_code: int = status.HTTP_422_UNPROCESSABLE_ENTITY,
+        status_code: int = status.HTTP_422_UNPROCESSABLE_CONTENT,
         log=True,
         log_level="DEBUG",
     ):

cognee/infrastructure/databases/vector/lancedb/LanceDBAdapter.py

@@ -324,7 +324,6 @@ class LanceDBAdapter(VectorDBInterface):
 
     def get_data_point_schema(self, model_type: BaseModel):
         related_models_fields = []
-
         for field_name, field_config in model_type.model_fields.items():
             if hasattr(field_config, "model_fields"):
                 related_models_fields.append(field_name)

cognee/infrastructure/files/exceptions.py

@@ -8,6 +8,6 @@ class FileContentHashingError(Exception):
         self,
         message: str = "Failed to hash content of the file.",
         name: str = "FileContentHashingError",
-        status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+        status_code=status.HTTP_422_UNPROCESSABLE_CONTENT,
     ):
         super().__init__(message, name, status_code)

cognee/infrastructure/files/storage/LocalFileStorage.py

@@ -82,16 +82,16 @@ class LocalFileStorage(Storage):
         self.ensure_directory_exists(file_dir_path)
 
         if overwrite or not os.path.exists(full_file_path):
-            with open(
-                full_file_path,
-                mode="w" if isinstance(data, str) else "wb",
-                encoding="utf-8" if isinstance(data, str) else None,
-            ) as file:
-                if hasattr(data, "read"):
-                    data.seek(0)
-                    file.write(data.read())
-                else:
+            if isinstance(data, str):
+                with open(full_file_path, mode="w", encoding="utf-8", newline="\n") as file:
                     file.write(data)
+            else:
+                with open(full_file_path, mode="wb") as file:
+                    if hasattr(data, "read"):
+                        data.seek(0)
+                        file.write(data.read())
+                    else:
+                        file.write(data)
 
                 file.close()
 

cognee/infrastructure/files/storage/S3FileStorage.py

@@ -70,18 +70,18 @@ class S3FileStorage(Storage):
         if overwrite or not await self.file_exists(file_path):
 
             def save_data_to_file():
-                with self.s3.open(
-                    full_file_path,
-                    mode="w" if isinstance(data, str) else "wb",
-                    encoding="utf-8" if isinstance(data, str) else None,
-                ) as file:
-                    if hasattr(data, "read"):
-                        data.seek(0)
-                        file.write(data.read())
-                    else:
+                if isinstance(data, str):
+                    with self.s3.open(
+                        full_file_path, mode="w", encoding="utf-8", newline="\n"
+                    ) as file:
                         file.write(data)
-
-                    file.close()
+                else:
+                    with self.s3.open(full_file_path, mode="wb") as file:
+                        if hasattr(data, "read"):
+                            data.seek(0)
+                            file.write(data.read())
+                        else:
+                            file.write(data)
 
             await run_async(save_data_to_file)
 

cognee/infrastructure/files/utils/guess_file_type.py

@@ -124,6 +124,12 @@ def guess_file_type(file: BinaryIO) -> filetype.Type:
     """
     file_type = filetype.guess(file)
 
+    # If file type could not be determined consider it a plain text file as they don't have magic number encoding
+    if file_type is None:
+        from filetype.types.base import Type
+
+        file_type = Type("text/plain", "txt")
+
     if file_type is None:
         raise FileTypeException(f"Unknown file detected: {file.name}.")
 

cognee/infrastructure/llm/prompts/feedback_reaction_prompt.txt

@@ -0,0 +1,14 @@
+A question was previously answered, but the answer received negative feedback.
+Please reconsider and improve the response.
+
+Question: {question}
+Context originally used: {context}
+Previous answer: {wrong_answer}
+Feedback on that answer: {negative_feedback}
+
+Task: Provide a better response. The new answer should be short and direct.
+Then explain briefly why this answer is better.
+
+Format your reply as:
+Answer: <improved answer>
+Explanation: <short explanation>

cognee/infrastructure/llm/prompts/feedback_report_prompt.txt

@@ -0,0 +1,13 @@
+Write a concise, stand-alone paragraph that explains the correct answer to the question below.
+The paragraph should read naturally on its own, providing all necessary context and reasoning
+so the answer is clear and well-supported.
+
+Question: {question}
+Correct answer: {improved_answer}
+Supporting context: {new_context}
+
+Your paragraph should:
+- First sentence clearly states the correct answer as a full sentence
+- Remainder flows from first sentence and provides explanation based on context
+- Use simple, direct language that is easy to follow
+- Use shorter sentences, no long-winded explanations

cognee/infrastructure/llm/prompts/feedback_user_context_prompt.txt

@@ -0,0 +1,5 @@
+Question: {question}
+Context: {context}
+
+Provide a one paragraph human readable summary of this interaction context,
+listing all the relevant facts and information in a simple and direct way.

cognee/infrastructure/llm/prompts/search_type_selector_prompt.txt

@@ -10,8 +10,6 @@ Here are the available `SearchType` tools and their specific functions:
   - Summarizing large amounts of information 
   - Quick understanding of complex subjects
 
-* **`INSIGHTS`**: The `INSIGHTS` search type discovers connections and relationships between entities in the knowledge graph.
-
   **Best for:**
 
   - Discovering how entities are connected
@@ -95,9 +93,6 @@ Here are the available `SearchType` tools and their specific functions:
 Query: "Summarize the key findings from these research papers"
 Response: `SUMMARIES`
 
-Query: "What is the relationship between the methodologies used in these papers?"
-Response: `INSIGHTS`
-
 Query: "When was Einstein born?"
 Response: `CHUNKS`
 

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

@@ -1,19 +1,24 @@
+import logging
 from typing import Type
 from pydantic import BaseModel
+import litellm
 import instructor
+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,
+)
 
-from cognee.infrastructure.llm.exceptions import MissingSystemPromptPathError
 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.llm.LLMGateway import LLMGateway
 from cognee.infrastructure.llm.config import get_llm_config
 
+logger = get_logger()
+
 
 class AnthropicAdapter(LLMInterface):
     """
@@ -35,8 +40,13 @@ class AnthropicAdapter(LLMInterface):
         self.model = model
         self.max_completion_tokens = max_completion_tokens
 
-    @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:

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

@@ -12,11 +12,18 @@ 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.structured_output_framework.litellm_instructor.llm.rate_limiter import (
-    rate_limit_async,
-    sleep_and_retry_async,
+import logging
+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 GeminiAdapter(LLMInterface):
     """
@@ -58,8 +65,13 @@ class GeminiAdapter(LLMInterface):
 
         self.aclient = instructor.from_litellm(litellm.acompletion, 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:

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

@@ -12,11 +12,18 @@ 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.structured_output_framework.litellm_instructor.llm.rate_limiter import (
-    rate_limit_async,
-    sleep_and_retry_async,
+import logging
+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 GenericAPIAdapter(LLMInterface):
     """
@@ -58,8 +65,13 @@ class GenericAPIAdapter(LLMInterface):
 
         self.aclient = instructor.from_litellm(litellm.acompletion, 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:

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.