cognee 0.2.3.dev1__py3-none-any.whl → 0.3.0__py3-none-any.whl

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/extraction/knowledge_graph/extract_event_graph.py

@@ -0,0 +1,46 @@
+import os
+from pydantic import BaseModel
+from typing import Type
+from cognee.infrastructure.llm.LLMGateway import LLMGateway
+
+from cognee.infrastructure.llm.config import (
+    get_llm_config,
+)
+
+
+async def extract_event_graph(content: str, response_model: Type[BaseModel]):
+    """
+    Extracts an event graph from the given content using an LLM with a structured output format.
+
+    This function loads a temporal graph extraction prompt from the LLM configuration,
+    renders it as a system prompt, and queries the LLM to produce a structured event
+    graph matching the specified response model.
+
+    Args:
+        content (str): The input text from which to extract the event graph.
+        response_model (Type[BaseModel]): A Pydantic model defining the structure of the expected output.
+
+    Returns:
+        BaseModel: An instance of the response_model populated with the extracted event graph.
+    """
+
+    llm_config = get_llm_config()
+
+    prompt_path = llm_config.temporal_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
+    )
+
+    return content_graph

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

@@ -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=[
                 {

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

@@ -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,

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

@@ -54,11 +54,15 @@ 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:
@@ -74,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,
@@ -94,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:
@@ -102,7 +106,9 @@ 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:
@@ -117,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,
@@ -134,7 +140,7 @@ 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,

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

@@ -23,9 +23,12 @@ from cognee.infrastructure.llm.structured_output_framework.litellm_instructor.ll
     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()
+
 
 class OpenAIAdapter(LLMInterface):
     """
@@ -64,7 +67,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,
@@ -77,7 +80,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
@@ -129,6 +132,7 @@ class OpenAIAdapter(LLMInterface):
                 api_version=self.api_version,
                 response_model=response_model,
                 max_retries=self.MAX_RETRIES,
+                extra_body={"reasoning_effort": "minimal"},
             )
         except (
             ContentFilterFinishReasonError,
@@ -139,7 +143,27 @@ class OpenAIAdapter(LLMInterface):
                 isinstance(error, InstructorRetryException)
                 and "content management policy" not in str(error).lower()
             ):
-                raise error
+                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,
+                )
 
             if not (self.fallback_model and self.fallback_api_key):
                 raise ContentPolicyFilterError(
@@ -301,7 +325,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,
         )
 

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

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

@@ -16,17 +16,17 @@ class MistralTokenizer(TokenizerInterface):
 
     Instance variables include:
     - model: str
-    - max_tokens: int
+    - max_completion_tokens: int
     - tokenizer: MistralTokenizer
     """
 
     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
 
         # Import here to make it an optional dependency
         from mistral_common.tokens.tokenizers.mistral import MistralTokenizer

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

@@ -13,10 +13,10 @@ class TikTokenTokenizer(TokenizerInterface):
     def __init__(
         self,
         model: Optional[str] = None,
-        max_tokens: int = 8191,
+        max_completion_tokens: int = 8191,
     ):
         self.model = model
-        self.max_tokens = max_tokens
+        self.max_completion_tokens = max_completion_tokens
         # Initialize TikToken for GPT based on model
         if model:
             self.tokenizer = tiktoken.encoding_for_model(self.model)
@@ -93,9 +93,9 @@ class TikTokenTokenizer(TokenizerInterface):
         num_tokens = len(self.tokenizer.encode(text))
         return num_tokens
 
-    def trim_text_to_max_tokens(self, text: str) -> str:
+    def trim_text_to_max_completion_tokens(self, text: str) -> str:
         """
-        Trim the text so that the number of tokens does not exceed max_tokens.
+        Trim the text so that the number of tokens does not exceed max_completion_tokens.
 
         Parameters:
         -----------
@@ -111,13 +111,13 @@ class TikTokenTokenizer(TokenizerInterface):
         num_tokens = self.count_tokens(text)
 
         # If the number of tokens is within the limit, return the text as is
-        if num_tokens <= self.max_tokens:
+        if num_tokens <= self.max_completion_tokens:
             return text
 
         # If the number exceeds the limit, trim the text
         # This is a simple trim, it may cut words in half; consider using word boundaries for a cleaner cut
         encoded_text = self.tokenizer.encode(text)
-        trimmed_encoded_text = encoded_text[: self.max_tokens]
+        trimmed_encoded_text = encoded_text[: self.max_completion_tokens]
         # Decoding the trimmed text
         trimmed_text = self.tokenizer.decode(trimmed_encoded_text)
         return trimmed_text

cognee/infrastructure/llm/utils.py

@@ -32,13 +32,13 @@ def get_max_chunk_tokens():
 
     # 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
-    llm_cutoff_point = llm_client.max_tokens // 2  # Round down the division
-    max_chunk_tokens = min(embedding_engine.max_tokens, llm_cutoff_point)
+    llm_cutoff_point = llm_client.max_completion_tokens // 2  # Round down the division
+    max_chunk_tokens = min(embedding_engine.max_completion_tokens, llm_cutoff_point)
 
     return max_chunk_tokens
 
 
-def get_model_max_tokens(model_name: str):
+def get_model_max_completion_tokens(model_name: str):
     """
     Retrieve the maximum token limit for a specified model name if it exists.
 
@@ -56,15 +56,15 @@ def get_model_max_tokens(model_name: str):
 
         Number of max tokens of model, or None if model is unknown
     """
-    max_tokens = None
+    max_completion_tokens = None
 
     if model_name in litellm.model_cost:
-        max_tokens = litellm.model_cost[model_name]["max_tokens"]
-        logger.debug(f"Max input tokens for {model_name}: {max_tokens}")
+        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.")
 
-    return max_tokens
+    return max_completion_tokens
 
 
 async def test_llm_connection():

cognee/infrastructure/utils/run_sync.py

@@ -8,8 +8,15 @@ def run_sync(coro, timeout=None):
 
     def runner():
         nonlocal result, exception
+
         try:
-            result = asyncio.run(coro)
+            try:
+                running_loop = asyncio.get_running_loop()
+
+                result = asyncio.run_coroutine_threadsafe(coro, running_loop).result(timeout)
+            except RuntimeError:
+                result = asyncio.run(coro)
+
         except Exception as e:
             exception = e
 

cognee/modules/chunking/models/DocumentChunk.py

@@ -1,8 +1,9 @@
-from typing import List
+from typing import List, Union
 
 from cognee.infrastructure.engine import DataPoint
 from cognee.modules.data.processing.document_types import Document
 from cognee.modules.engine.models import Entity
+from cognee.tasks.temporal_graph.models import Event
 
 
 class DocumentChunk(DataPoint):
@@ -20,7 +21,7 @@ class DocumentChunk(DataPoint):
     - chunk_index: The index of the chunk in the original document.
     - cut_type: The type of cut that defined this chunk.
     - is_part_of: The document to which this chunk belongs.
-    - contains: A list of entities contained within the chunk (default is None).
+    - contains: A list of entities or events contained within the chunk (default is None).
     - metadata: A dictionary to hold meta information related to the chunk, including index
     fields.
     """
@@ -30,6 +31,6 @@ class DocumentChunk(DataPoint):
     chunk_index: int
     cut_type: str
     is_part_of: Document
-    contains: List[Entity] = None
+    contains: List[Union[Entity, Event]] = None
 
     metadata: dict = {"index_fields": ["text"]}

cognee/modules/cloud/exceptions/CloudApiKeyMissingError.py

@@ -0,0 +1,15 @@
+from fastapi import status
+
+from cognee.exceptions.exceptions import CogneeConfigurationError
+
+
+class CloudApiKeyMissingError(CogneeConfigurationError):
+    """Raised when the API key for the cloud service is not provided."""
+
+    def __init__(
+        self,
+        message: str = "Failed to connect to the cloud service. Please add your API key to local instance.",
+        name: str = "CloudApiKeyMissingError",
+        status_code=status.HTTP_400_BAD_REQUEST,
+    ):
+        super().__init__(message, name, status_code)

cognee/modules/cloud/exceptions/CloudConnectionError.py

@@ -0,0 +1,15 @@
+from fastapi import status
+
+from cognee.exceptions.exceptions import CogneeConfigurationError
+
+
+class CloudConnectionError(CogneeConfigurationError):
+    """Raised when the connection to the cloud service fails."""
+
+    def __init__(
+        self,
+        message: str = "Failed to connect to the cloud service. Please check your cloud API key in local instance.",
+        name: str = "CloudConnnectionError",
+        status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
+    ):
+        super().__init__(message, name, status_code)

cognee/modules/cloud/exceptions/__init__.py

@@ -0,0 +1,2 @@
+from .CloudConnectionError import CloudConnectionError
+from .CloudApiKeyMissingError import CloudApiKeyMissingError

cognee/modules/cloud/operations/__init__.py

@@ -0,0 +1 @@
+from .check_api_key import check_api_key

cognee/modules/cloud/operations/check_api_key.py

@@ -0,0 +1,25 @@
+import aiohttp
+
+from cognee.modules.cloud.exceptions import CloudConnectionError
+
+
+async def check_api_key(auth_token: str):
+    cloud_base_url = "http://localhost:8001"
+
+    url = f"{cloud_base_url}/api/api-keys/check"
+    headers = {"X-Api-Key": auth_token}
+
+    try:
+        async with aiohttp.ClientSession() as session:
+            async with session.post(url, headers=headers) as response:
+                if response.status == 200:
+                    return
+                else:
+                    error_text = await response.text()
+
+                    raise CloudConnectionError(
+                        f"Failed to connect to cloud instance: {response.status} - {error_text}"
+                    )
+
+    except Exception as e:
+        raise CloudConnectionError(f"Failed to connect to cloud instance: {str(e)}")

cognee/modules/data/deletion/prune_system.py

@@ -3,7 +3,7 @@ from cognee.infrastructure.databases.graph.get_graph_engine import get_graph_eng
 from cognee.infrastructure.databases.relational import get_relational_engine
 
 
-async def prune_system(graph=True, vector=True, metadata=False):
+async def prune_system(graph=True, vector=True, metadata=True):
     if graph:
         graph_engine = await get_graph_engine()
         await graph_engine.delete_graph()

cognee/modules/data/methods/__init__.py

@@ -7,6 +7,7 @@ from .get_datasets import get_datasets
 from .get_datasets_by_name import get_datasets_by_name
 from .get_dataset_data import get_dataset_data
 from .get_authorized_dataset import get_authorized_dataset
+from .get_authorized_dataset_by_name import get_authorized_dataset_by_name
 from .get_data import get_data
 from .get_unique_dataset_id import get_unique_dataset_id
 from .get_authorized_existing_datasets import get_authorized_existing_datasets
@@ -18,6 +19,7 @@ from .delete_data import delete_data
 
 # Create
 from .load_or_create_datasets import load_or_create_datasets
+from .create_authorized_dataset import create_authorized_dataset
 
 # Check
 from .check_dataset_name import check_dataset_name

cognee/modules/data/methods/check_dataset_name.py

@@ -1,3 +1,3 @@
 def check_dataset_name(dataset_name: str):
     if "." in dataset_name or " " in dataset_name:
-        raise ValueError("Dataset name cannot contain spaces or underscores")
+        raise ValueError("Dataset name cannot contain spaces or dots.")

cognee/modules/data/methods/create_authorized_dataset.py

@@ -0,0 +1,19 @@
+from cognee.infrastructure.databases.relational import get_relational_engine
+from cognee.modules.users.models import User
+from cognee.modules.data.models import Dataset
+from cognee.modules.users.permissions.methods import give_permission_on_dataset
+from .create_dataset import create_dataset
+
+
+async def create_authorized_dataset(dataset_name: str, user: User) -> Dataset:
+    db_engine = get_relational_engine()
+
+    async with db_engine.get_async_session() as session:
+        new_dataset = await create_dataset(dataset_name, user, session)
+
+    await give_permission_on_dataset(user, new_dataset.id, "read")
+    await give_permission_on_dataset(user, new_dataset.id, "write")
+    await give_permission_on_dataset(user, new_dataset.id, "delete")
+    await give_permission_on_dataset(user, new_dataset.id, "share")
+
+    return new_dataset