cognee 0.3.3__py3-none-any.whl → 0.3.4__py3-none-any.whl

cognee/api/v1/cloud/routers/get_checks_router.py

@@ -16,7 +16,7 @@ def get_checks_router():
         api_token = request.headers.get("X-Api-Key")
 
         if api_token is None:
-            return CloudApiKeyMissingError()
+            raise CloudApiKeyMissingError()
 
         return await check_api_key(api_token)
 

cognee/api/v1/cognify/cognify.py

@@ -3,6 +3,7 @@ from pydantic import BaseModel
 from typing import Union, Optional
 from uuid import UUID
 
+from cognee.modules.ontology.ontology_env_config import get_ontology_env_config
 from cognee.shared.logging_utils import get_logger
 from cognee.shared.data_models import KnowledgeGraph
 from cognee.infrastructure.llm import get_max_chunk_tokens
@@ -10,7 +11,11 @@ from cognee.infrastructure.llm import get_max_chunk_tokens
 from cognee.modules.pipelines import run_pipeline
 from cognee.modules.pipelines.tasks.task import Task
 from cognee.modules.chunking.TextChunker import TextChunker
-from cognee.modules.ontology.rdf_xml.OntologyResolver import OntologyResolver
+from cognee.modules.ontology.ontology_config import Config
+from cognee.modules.ontology.get_default_ontology_resolver import (
+    get_default_ontology_resolver,
+    get_ontology_resolver_from_env,
+)
 from cognee.modules.users.models import User
 
 from cognee.tasks.documents import (
@@ -39,7 +44,7 @@ async def cognify(
     graph_model: BaseModel = KnowledgeGraph,
     chunker=TextChunker,
     chunk_size: int = None,
-    ontology_file_path: Optional[str] = None,
+    config: Config = None,
     vector_db_config: dict = None,
     graph_db_config: dict = None,
     run_in_background: bool = False,
@@ -100,8 +105,6 @@ async def cognify(
                    Formula: min(embedding_max_completion_tokens, llm_max_completion_tokens // 2)
                    Default limits: ~512-8192 tokens depending on models.
                    Smaller chunks = more granular but potentially fragmented knowledge.
-        ontology_file_path: Path to RDF/OWL ontology file for domain-specific entity types.
-                          Useful for specialized fields like medical or legal documents.
         vector_db_config: Custom vector database configuration for embeddings storage.
         graph_db_config: Custom graph database configuration for relationship storage.
         run_in_background: If True, starts processing asynchronously and returns immediately.
@@ -188,11 +191,28 @@ async def cognify(
         - LLM_RATE_LIMIT_ENABLED: Enable rate limiting (default: False)
         - LLM_RATE_LIMIT_REQUESTS: Max requests per interval (default: 60)
     """
+    if config is None:
+        ontology_config = get_ontology_env_config()
+        if (
+            ontology_config.ontology_file_path
+            and ontology_config.ontology_resolver
+            and ontology_config.matching_strategy
+        ):
+            config: Config = {
+                "ontology_config": {
+                    "ontology_resolver": get_ontology_resolver_from_env(**ontology_config.to_dict())
+                }
+            }
+        else:
+            config: Config = {
+                "ontology_config": {"ontology_resolver": get_default_ontology_resolver()}
+            }
+
     if temporal_cognify:
         tasks = await get_temporal_tasks(user, chunker, chunk_size)
     else:
         tasks = await get_default_tasks(
-            user, graph_model, chunker, chunk_size, ontology_file_path, custom_prompt
+            user, graph_model, chunker, chunk_size, config, custom_prompt
         )
 
     # By calling get pipeline executor we get a function that will have the run_pipeline run in the background or a function that we will need to wait for
@@ -216,9 +236,26 @@ async def get_default_tasks(  # TODO: Find out a better way to do this (Boris's
     graph_model: BaseModel = KnowledgeGraph,
     chunker=TextChunker,
     chunk_size: int = None,
-    ontology_file_path: Optional[str] = None,
+    config: Config = None,
     custom_prompt: Optional[str] = None,
 ) -> list[Task]:
+    if config is None:
+        ontology_config = get_ontology_env_config()
+        if (
+            ontology_config.ontology_file_path
+            and ontology_config.ontology_resolver
+            and ontology_config.matching_strategy
+        ):
+            config: Config = {
+                "ontology_config": {
+                    "ontology_resolver": get_ontology_resolver_from_env(**ontology_config.to_dict())
+                }
+            }
+        else:
+            config: Config = {
+                "ontology_config": {"ontology_resolver": get_default_ontology_resolver()}
+            }
+
     default_tasks = [
         Task(classify_documents),
         Task(check_permissions_on_dataset, user=user, permissions=["write"]),
@@ -230,7 +267,7 @@ async def get_default_tasks(  # TODO: Find out a better way to do this (Boris's
         Task(
             extract_graph_from_data,
             graph_model=graph_model,
-            ontology_adapter=OntologyResolver(ontology_file=ontology_file_path),
+            config=config,
             custom_prompt=custom_prompt,
             task_config={"batch_size": 10},
         ),  # Generate knowledge graphs from the document chunks.

cognee/api/v1/cognify/routers/get_cognify_router.py

@@ -3,6 +3,7 @@ import asyncio
 from uuid import UUID
 from pydantic import Field
 from typing import List, Optional
+from fastapi.encoders import jsonable_encoder
 from fastapi.responses import JSONResponse
 from fastapi import APIRouter, WebSocket, Depends, WebSocketDisconnect
 from starlette.status import WS_1000_NORMAL_CLOSURE, WS_1008_POLICY_VIOLATION
@@ -119,7 +120,7 @@ def get_cognify_router() -> APIRouter:
 
             # If any cognify run errored return JSONResponse with proper error status code
             if any(isinstance(v, PipelineRunErrored) for v in cognify_run.values()):
-                return JSONResponse(status_code=420, content=cognify_run)
+                return JSONResponse(status_code=420, content=jsonable_encoder(cognify_run))
             return cognify_run
         except Exception as error:
             return JSONResponse(status_code=409, content={"error": str(error)})

cognee/api/v1/prune/prune.py

@@ -7,8 +7,8 @@ class prune:
         await _prune_data()
 
     @staticmethod
-    async def prune_system(graph=True, vector=True, metadata=False):
-        await _prune_system(graph, vector, metadata)
+    async def prune_system(graph=True, vector=True, metadata=False, cache=True):
+        await _prune_system(graph, vector, metadata, cache)
 
 
 if __name__ == "__main__":

cognee/api/v1/search/search.py

@@ -22,7 +22,7 @@ async def search(
     node_type: Optional[Type] = NodeSet,
     node_name: Optional[List[str]] = None,
     save_interaction: bool = False,
-    last_k: Optional[int] = None,
+    last_k: Optional[int] = 1,
     only_context: bool = False,
     use_combined_context: bool = False,
 ) -> Union[List[SearchResult], CombinedSearchResult]:

cognee/api/v1/sync/sync.py

@@ -23,6 +23,7 @@ from cognee.modules.sync.methods import (
     mark_sync_completed,
     mark_sync_failed,
 )
+from cognee.shared.utils import create_secure_ssl_context
 
 logger = get_logger("sync")
 
@@ -583,7 +584,9 @@ async def _check_hashes_diff(
     logger.info(f"Checking missing hashes on cloud for dataset {dataset.id}")
 
     try:
-        async with aiohttp.ClientSession() as session:
+        ssl_context = create_secure_ssl_context()
+        connector = aiohttp.TCPConnector(ssl=ssl_context)
+        async with aiohttp.ClientSession(connector=connector) as session:
             async with session.post(url, json=payload.dict(), headers=headers) as response:
                 if response.status == 200:
                     data = await response.json()
@@ -630,7 +633,9 @@ async def _download_missing_files(
 
     headers = {"X-Api-Key": auth_token}
 
-    async with aiohttp.ClientSession() as session:
+    ssl_context = create_secure_ssl_context()
+    connector = aiohttp.TCPConnector(ssl=ssl_context)
+    async with aiohttp.ClientSession(connector=connector) as session:
         for file_hash in hashes_missing_on_local:
             try:
                 # Download file from cloud by hash
@@ -749,7 +754,9 @@ async def _upload_missing_files(
 
     headers = {"X-Api-Key": auth_token}
 
-    async with aiohttp.ClientSession() as session:
+    ssl_context = create_secure_ssl_context()
+    connector = aiohttp.TCPConnector(ssl=ssl_context)
+    async with aiohttp.ClientSession(connector=connector) as session:
         for file_info in files_to_upload:
             try:
                 file_dir = os.path.dirname(file_info.raw_data_location)
@@ -809,7 +816,9 @@ async def _prune_cloud_dataset(
     logger.info("Pruning cloud dataset to match local state")
 
     try:
-        async with aiohttp.ClientSession() as session:
+        ssl_context = create_secure_ssl_context()
+        connector = aiohttp.TCPConnector(ssl=ssl_context)
+        async with aiohttp.ClientSession(connector=connector) as session:
             async with session.put(url, json=payload.dict(), headers=headers) as response:
                 if response.status == 200:
                     data = await response.json()
@@ -852,7 +861,9 @@ async def _trigger_remote_cognify(
     logger.info(f"Triggering cognify processing for dataset {dataset_id}")
 
     try:
-        async with aiohttp.ClientSession() as session:
+        ssl_context = create_secure_ssl_context()
+        connector = aiohttp.TCPConnector(ssl=ssl_context)
+        async with aiohttp.ClientSession(connector=connector) as session:
             async with session.post(url, json=payload, headers=headers) as response:
                 if response.status == 200:
                     data = await response.json()

cognee/base_config.py

@@ -10,13 +10,30 @@ import pydantic
 class BaseConfig(BaseSettings):
     data_root_directory: str = get_absolute_path(".data_storage")
     system_root_directory: str = get_absolute_path(".cognee_system")
-    monitoring_tool: object = Observer.LANGFUSE
+    cache_root_directory: str = get_absolute_path(".cognee_cache")
+    monitoring_tool: object = Observer.NONE
 
     @pydantic.model_validator(mode="after")
     def validate_paths(self):
+        # Adding this here temporarily to ensure that the cache root directory is set correctly for S3 storage automatically
+        # I'll remove this after we update documentation for S3 storage
+        # Auto-configure cache root directory for S3 storage if not explicitly set
+        storage_backend = os.getenv("STORAGE_BACKEND", "").lower()
+        cache_root_env = os.getenv("CACHE_ROOT_DIRECTORY")
+
+        if storage_backend == "s3" and not cache_root_env:
+            # Auto-generate S3 cache path when using S3 storage
+            bucket_name = os.getenv("STORAGE_BUCKET_NAME")
+            if bucket_name:
+                self.cache_root_directory = f"s3://{bucket_name}/cognee/cache"
+
         # Require absolute paths for root directories
         self.data_root_directory = ensure_absolute_path(self.data_root_directory)
         self.system_root_directory = ensure_absolute_path(self.system_root_directory)
+        # Set monitoring tool based on available keys
+        if self.langfuse_public_key and self.langfuse_secret_key:
+            self.monitoring_tool = Observer.LANGFUSE
+
         return self
 
     langfuse_public_key: Optional[str] = os.getenv("LANGFUSE_PUBLIC_KEY")
@@ -31,6 +48,7 @@ class BaseConfig(BaseSettings):
             "data_root_directory": self.data_root_directory,
             "system_root_directory": self.system_root_directory,
             "monitoring_tool": self.monitoring_tool,
+            "cache_root_directory": self.cache_root_directory,
         }
 
 

cognee/eval_framework/corpus_builder/task_getters/get_default_tasks_by_indices.py

@@ -5,7 +5,7 @@ from cognee.modules.chunking.TextChunker import TextChunker
 from cognee.tasks.graph import extract_graph_from_data
 from cognee.tasks.storage import add_data_points
 from cognee.shared.data_models import KnowledgeGraph
-from cognee.modules.ontology.rdf_xml.OntologyResolver import OntologyResolver
+from cognee.modules.ontology.rdf_xml.RDFLibOntologyResolver import RDFLibOntologyResolver
 
 
 async def get_default_tasks_by_indices(
@@ -33,7 +33,7 @@ async def get_no_summary_tasks(
     # Get base tasks (0=classify, 1=check_permissions, 2=extract_chunks)
     base_tasks = await get_default_tasks_by_indices([0, 1, 2], chunk_size, chunker)
 
-    ontology_adapter = OntologyResolver(ontology_file=ontology_file_path)
+    ontology_adapter = RDFLibOntologyResolver(ontology_file=ontology_file_path)
 
     graph_task = Task(
         extract_graph_from_data,

cognee/infrastructure/databases/graph/kuzu/remote_kuzu_adapter.py

@@ -7,6 +7,7 @@ import aiohttp
 from uuid import UUID
 
 from cognee.infrastructure.databases.graph.kuzu.adapter import KuzuAdapter
+from cognee.shared.utils import create_secure_ssl_context
 
 logger = get_logger()
 
@@ -42,7 +43,9 @@ class RemoteKuzuAdapter(KuzuAdapter):
     async def _get_session(self) -> aiohttp.ClientSession:
         """Get or create an aiohttp session."""
         if self._session is None or self._session.closed:
-            self._session = aiohttp.ClientSession()
+            ssl_context = create_secure_ssl_context()
+            connector = aiohttp.TCPConnector(ssl=ssl_context)
+            self._session = aiohttp.ClientSession(connector=connector)
         return self._session
 
     async def close(self):

cognee/infrastructure/databases/relational/ModelBase.py

@@ -1,7 +1,8 @@
 from sqlalchemy.orm import DeclarativeBase
+from sqlalchemy.ext.asyncio import AsyncAttrs
 
 
-class Base(DeclarativeBase):
+class Base(AsyncAttrs, DeclarativeBase):
     """
     Represents a base class for declarative models using SQLAlchemy.
 

cognee/infrastructure/databases/vector/chromadb/ChromaDBAdapter.py

@@ -83,7 +83,7 @@ def process_data_for_chroma(data):
         elif isinstance(value, list):
             # Store lists as JSON strings with special prefix
             processed_data[f"{key}__list"] = json.dumps(value)
-        elif isinstance(value, (str, int, float, bool)) or value is None:
+        elif isinstance(value, (str, int, float, bool)):
             processed_data[key] = value
         else:
             processed_data[key] = str(value)
@@ -553,8 +553,4 @@ class ChromaDBAdapter(VectorDBInterface):
             Returns a list of collection names.
         """
         client = await self.get_connection()
-        collections = await client.list_collections()
-        return [
-            collection.name if hasattr(collection, "name") else collection["name"]
-            for collection in collections
-        ]
+        return await client.list_collections()

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

@@ -14,6 +14,7 @@ from cognee.infrastructure.databases.vector.embeddings.embedding_rate_limiter im
     embedding_rate_limit_async,
     embedding_sleep_and_retry_async,
 )
+from cognee.shared.utils import create_secure_ssl_context
 
 logger = get_logger("OllamaEmbeddingEngine")
 
@@ -94,16 +95,16 @@ class OllamaEmbeddingEngine(EmbeddingEngine):
         """
         Internal method to call the Ollama embeddings endpoint for a single prompt.
         """
-        payload = {
-            "model": self.model,
-            "prompt": prompt,
-        }
+        payload = {"model": self.model, "prompt": prompt, "input": prompt}
+
         headers = {}
         api_key = os.getenv("LLM_API_KEY")
         if api_key:
             headers["Authorization"] = f"Bearer {api_key}"
 
-        async with aiohttp.ClientSession() as session:
+        ssl_context = create_secure_ssl_context()
+        connector = aiohttp.TCPConnector(ssl=ssl_context)
+        async with aiohttp.ClientSession(connector=connector) as session:
             async with session.post(
                 self.endpoint, json=payload, headers=headers, timeout=60.0
             ) as response:

cognee/infrastructure/files/storage/LocalFileStorage.py

@@ -253,6 +253,56 @@ class LocalFileStorage(Storage):
         if os.path.exists(full_file_path):
             os.remove(full_file_path)
 
+    def list_files(self, directory_path: str, recursive: bool = False) -> list[str]:
+        """
+        List all files in the specified directory.
+
+        Parameters:
+        -----------
+            - directory_path (str): The directory path to list files from
+            - recursive (bool): If True, list files recursively in subdirectories
+
+        Returns:
+        --------
+            - list[str]: List of file paths relative to the storage root
+        """
+        from pathlib import Path
+
+        parsed_storage_path = get_parsed_path(self.storage_path)
+
+        if directory_path:
+            full_directory_path = os.path.join(parsed_storage_path, directory_path)
+        else:
+            full_directory_path = parsed_storage_path
+
+        directory_pathlib = Path(full_directory_path)
+
+        if not directory_pathlib.exists() or not directory_pathlib.is_dir():
+            return []
+
+        files = []
+
+        if recursive:
+            # Use rglob for recursive search
+            for file_path in directory_pathlib.rglob("*"):
+                if file_path.is_file():
+                    # Get relative path from storage root
+                    relative_path = os.path.relpath(str(file_path), parsed_storage_path)
+                    # Normalize path separators for consistency
+                    relative_path = relative_path.replace(os.sep, "/")
+                    files.append(relative_path)
+        else:
+            # Use iterdir for just immediate directory
+            for file_path in directory_pathlib.iterdir():
+                if file_path.is_file():
+                    # Get relative path from storage root
+                    relative_path = os.path.relpath(str(file_path), parsed_storage_path)
+                    # Normalize path separators for consistency
+                    relative_path = relative_path.replace(os.sep, "/")
+                    files.append(relative_path)
+
+        return files
+
     def remove_all(self, tree_path: str = None):
         """
         Remove an entire directory tree at the specified path, including all files and

cognee/infrastructure/files/storage/S3FileStorage.py

@@ -155,21 +155,19 @@ class S3FileStorage(Storage):
         """
         Ensure that the specified directory exists, creating it if necessary.
 
-        If the directory already exists, no action is taken.
+        For S3 storage, this is a no-op since directories are created implicitly
+        when files are written to paths. S3 doesn't have actual directories,
+        just object keys with prefixes that appear as directories.
 
         Parameters:
         -----------
 
             - directory_path (str): The path of the directory to check or create.
         """
-        if not directory_path.strip():
-            directory_path = self.storage_path.replace("s3://", "")
-
-        def ensure_directory():
-            if not self.s3.exists(directory_path):
-                self.s3.makedirs(directory_path, exist_ok=True)
-
-        await run_async(ensure_directory)
+        # In S3, directories don't exist as separate entities - they're just prefixes
+        # When you write a file to s3://bucket/path/to/file.txt, the "directories"
+        # path/ and path/to/ are implicitly created. No explicit action needed.
+        pass
 
     async def copy_file(self, source_file_path: str, destination_file_path: str):
         """
@@ -213,6 +211,55 @@ class S3FileStorage(Storage):
 
         await run_async(remove_file)
 
+    async def list_files(self, directory_path: str, recursive: bool = False) -> list[str]:
+        """
+        List all files in the specified directory.
+
+        Parameters:
+        -----------
+            - directory_path (str): The directory path to list files from
+            - recursive (bool): If True, list files recursively in subdirectories
+
+        Returns:
+        --------
+            - list[str]: List of file paths relative to the storage root
+        """
+
+        def list_files_sync():
+            if directory_path:
+                # Combine storage path with directory path
+                full_path = os.path.join(self.storage_path.replace("s3://", ""), directory_path)
+            else:
+                full_path = self.storage_path.replace("s3://", "")
+
+            if recursive:
+                # Use ** for recursive search
+                pattern = f"{full_path}/**"
+            else:
+                # Just files in the immediate directory
+                pattern = f"{full_path}/*"
+
+            # Use s3fs glob to find files
+            try:
+                all_paths = self.s3.glob(pattern)
+                # Filter to only files (not directories)
+                files = [path for path in all_paths if self.s3.isfile(path)]
+
+                # Convert back to relative paths from storage root
+                storage_prefix = self.storage_path.replace("s3://", "")
+                relative_files = []
+                for file_path in files:
+                    if file_path.startswith(storage_prefix):
+                        relative_path = file_path[len(storage_prefix) :].lstrip("/")
+                        relative_files.append(relative_path)
+
+                return relative_files
+            except Exception:
+                # If directory doesn't exist or other error, return empty list
+                return []
+
+        return await run_async(list_files_sync)
+
     async def remove_all(self, tree_path: str):
         """
         Remove an entire directory tree at the specified path, including all files and

cognee/infrastructure/files/storage/StorageManager.py

@@ -135,6 +135,24 @@ class StorageManager:
         else:
             return self.storage.remove(file_path)
 
+    async def list_files(self, directory_path: str, recursive: bool = False) -> list[str]:
+        """
+        List all files in the specified directory.
+
+        Parameters:
+        -----------
+            - directory_path (str): The directory path to list files from
+            - recursive (bool): If True, list files recursively in subdirectories
+
+        Returns:
+        --------
+            - list[str]: List of file paths relative to the storage root
+        """
+        if inspect.iscoroutinefunction(self.storage.list_files):
+            return await self.storage.list_files(directory_path, recursive)
+        else:
+            return self.storage.list_files(directory_path, recursive)
+
     async def remove_all(self, tree_path: str = None):
         """
         Remove an entire directory tree at the specified path, including all files and

cognee/infrastructure/files/utils/get_file_metadata.py

@@ -56,7 +56,12 @@ 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 = Path(file_path).stem if file_path else None
+
+    if isinstance(file_path, str):
+        file_name = Path(file_path).stem if file_path else None
+    else:
+        # In case file_path does not exist or is a integer return None
+        file_name = None
 
     # Get file size
     pos = file.tell()  # remember current pointer

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

@@ -12,6 +12,7 @@ from cognee.infrastructure.llm.structured_output_framework.litellm_instructor.ll
 )
 
 from cognee.infrastructure.llm.LLMGateway import LLMGateway
+from cognee.infrastructure.llm.config import get_llm_config
 
 
 class AnthropicAdapter(LLMInterface):
@@ -27,7 +28,8 @@ class AnthropicAdapter(LLMInterface):
         import anthropic
 
         self.aclient = instructor.patch(
-            create=anthropic.AsyncAnthropic().messages.create, mode=instructor.Mode.ANTHROPIC_TOOLS
+            create=anthropic.AsyncAnthropic(api_key=get_llm_config().llm_api_key).messages.create,
+            mode=instructor.Mode.ANTHROPIC_TOOLS,
         )
 
         self.model = model
@@ -57,7 +59,7 @@ class AnthropicAdapter(LLMInterface):
 
         return await self.aclient(
             model=self.model,
-            max_completion_tokens=4096,
+            max_tokens=4096,
             max_retries=5,
             messages=[
                 {

cognee/modules/cloud/operations/check_api_key.py

@@ -1,6 +1,7 @@
 import aiohttp
 
 from cognee.modules.cloud.exceptions import CloudConnectionError
+from cognee.shared.utils import create_secure_ssl_context
 
 
 async def check_api_key(auth_token: str):
@@ -10,7 +11,9 @@ async def check_api_key(auth_token: str):
     headers = {"X-Api-Key": auth_token}
 
     try:
-        async with aiohttp.ClientSession() as session:
+        ssl_context = create_secure_ssl_context()
+        connector = aiohttp.TCPConnector(ssl=ssl_context)
+        async with aiohttp.ClientSession(connector=connector) as session:
             async with session.post(url, headers=headers) as response:
                 if response.status == 200:
                     return

cognee/modules/data/deletion/prune_system.py

@@ -1,9 +1,10 @@
 from cognee.infrastructure.databases.vector import get_vector_engine
 from cognee.infrastructure.databases.graph.get_graph_engine import get_graph_engine
 from cognee.infrastructure.databases.relational import get_relational_engine
+from cognee.shared.cache import delete_cache
 
 
-async def prune_system(graph=True, vector=True, metadata=True):
+async def prune_system(graph=True, vector=True, metadata=True, cache=True):
     if graph:
         graph_engine = await get_graph_engine()
         await graph_engine.delete_graph()
@@ -15,3 +16,6 @@ async def prune_system(graph=True, vector=True, metadata=True):
     if metadata:
         db_engine = get_relational_engine()
         await db_engine.delete_database()
+
+    if cache:
+        await delete_cache()

cognee/modules/data/methods/create_authorized_dataset.py

@@ -6,6 +6,15 @@ from .create_dataset import create_dataset
 
 
 async def create_authorized_dataset(dataset_name: str, user: User) -> Dataset:
+    """
+        Create a new dataset and give all permissions on this dataset to the given user.
+    Args:
+        dataset_name: Name of the dataset.
+        user: The user object.
+
+    Returns:
+        Dataset: The new authorized dataset.
+    """
     db_engine = get_relational_engine()
 
     async with db_engine.get_async_session() as session:

cognee/modules/data/methods/get_authorized_dataset.py

@@ -15,7 +15,7 @@ async def get_authorized_dataset(
     Get a specific dataset with permissions for a user.
 
     Args:
-        user_id (UUID): user id
+        user: User object
         dataset_id (UUID): dataset id
         permission_type (str): permission type(read, write, delete, share), default is read
 

cognee/modules/data/methods/get_authorized_dataset_by_name.py

@@ -11,6 +11,17 @@ from ..models import Dataset
 async def get_authorized_dataset_by_name(
     dataset_name: str, user: User, permission_type: str
 ) -> Optional[Dataset]:
+    """
+    Get a specific dataset with the given name, with permissions for a given user.
+
+    Args:
+        dataset_name: Name of the dataset.
+        user: User object.
+        permission_type (str): permission type(read, write, delete, share), default is read
+
+    Returns:
+        Optional[Dataset]: dataset with permissions
+    """
     authorized_datasets = await get_authorized_existing_datasets([], permission_type, user)
 
     return next((dataset for dataset in authorized_datasets if dataset.name == dataset_name), None)