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

cognee/modules/users/roles/methods/create_role.py

@@ -16,6 +16,16 @@ async def create_role(
     role_name: str,
     owner_id: UUID,
 ):
+    """
+        Create a new role with the given name, if the request owner with the given id
+        has the necessary permission.
+    Args:
+        role_name: Name of the new role.
+        owner_id: Id of the request owner.
+
+    Returns:
+        None
+    """
     db_engine = get_relational_engine()
     async with db_engine.get_async_session() as session:
         user = await get_user(owner_id)

cognee/modules/users/tenants/methods/add_user_to_tenant.py

@@ -13,6 +13,18 @@ from cognee.modules.users.exceptions import (
 
 
 async def add_user_to_tenant(user_id: UUID, tenant_id: UUID, owner_id: UUID):
+    """
+        Add a user with the given id to the tenant with the given id.
+        This can only be successful if the request owner with the given id is the tenant owner.
+    Args:
+        user_id: Id of the user.
+        tenant_id: Id of the tenant.
+        owner_id: Id of the request owner.
+
+    Returns:
+        None
+
+    """
     db_engine = get_relational_engine()
     async with db_engine.get_async_session() as session:
         user = await get_user(user_id)

cognee/modules/users/tenants/methods/create_tenant.py

@@ -8,6 +8,16 @@ from cognee.modules.users.methods import get_user
 
 
 async def create_tenant(tenant_name: str, user_id: UUID):
+    """
+        Create a new tenant with the given name, for the user with the given id.
+        This user is the owner of the tenant.
+    Args:
+        tenant_name: Name of the new tenant.
+        user_id: Id of the user.
+
+    Returns:
+        None
+    """
     db_engine = get_relational_engine()
     async with db_engine.get_async_session() as session:
         try:

cognee/root_dir.py

@@ -20,6 +20,11 @@ def ensure_absolute_path(path: str) -> str:
     """
     if path is None:
         raise ValueError("Path cannot be None")
+
+    # Check if it's an S3 URL - S3 URLs are absolute by definition
+    if path.startswith("s3://"):
+        return path
+
     path_obj = Path(path).expanduser()
     if path_obj.is_absolute():
         return str(path_obj.resolve())

cognee/shared/cache.py

@@ -0,0 +1,346 @@
+"""
+Storage-aware cache management utilities for Cognee.
+
+This module provides cache functionality that works with both local and cloud storage
+backends (like S3) through the StorageManager abstraction.
+"""
+
+import hashlib
+import zipfile
+import asyncio
+from typing import Optional, Tuple
+import aiohttp
+import logging
+from io import BytesIO
+
+from cognee.base_config import get_base_config
+from cognee.infrastructure.files.storage.get_file_storage import get_file_storage
+from cognee.infrastructure.files.storage.StorageManager import StorageManager
+from cognee.shared.utils import create_secure_ssl_context
+
+logger = logging.getLogger(__name__)
+
+
+class StorageAwareCache:
+    """
+    A cache manager that works with different storage backends (local, S3, etc.)
+    """
+
+    def __init__(self, cache_subdir: str = "cache"):
+        """
+        Initialize the cache manager.
+
+        Args:
+            cache_subdir: Subdirectory name within the system root for caching
+        """
+        self.base_config = get_base_config()
+        # Since we're using cache_root_directory, don't add extra cache prefix
+        self.cache_base_path = ""
+        self.storage_manager: StorageManager = get_file_storage(
+            self.base_config.cache_root_directory
+        )
+
+        # Print absolute path
+        storage_path = self.storage_manager.storage.storage_path
+        if storage_path.startswith("s3://"):
+            absolute_path = storage_path  # S3 paths are already absolute
+        else:
+            import os
+
+            absolute_path = os.path.abspath(storage_path)
+        logger.info(f"Storage manager absolute path: {absolute_path}")
+
+    async def get_cache_dir(self) -> str:
+        """Get the base cache directory path."""
+        cache_path = self.cache_base_path or "."  # Use "." for root when cache_base_path is empty
+        await self.storage_manager.ensure_directory_exists(cache_path)
+        return cache_path
+
+    async def get_cache_subdir(self, name: str) -> str:
+        """Get a specific cache subdirectory."""
+        if self.cache_base_path:
+            cache_path = f"{self.cache_base_path}/{name}"
+        else:
+            cache_path = name
+        await self.storage_manager.ensure_directory_exists(cache_path)
+
+        # Return the absolute path based on storage system
+        if self.storage_manager.storage.storage_path.startswith("s3://"):
+            return cache_path
+        elif hasattr(self.storage_manager.storage, "storage_path"):
+            return f"{self.storage_manager.storage.storage_path}/{cache_path}"
+        else:
+            # Fallback for other storage types
+            return cache_path
+
+    async def delete_cache(self):
+        """Delete the entire cache directory."""
+        logger.info("Deleting cache...")
+        try:
+            await self.storage_manager.remove_all(self.cache_base_path)
+            logger.info("✓ Cache deleted successfully!")
+        except Exception as e:
+            logger.error(f"Error deleting cache: {e}")
+            raise
+
+    async def _is_cache_valid(self, cache_dir: str, version_or_hash: str) -> bool:
+        """Check if cached content is valid for the given version/hash."""
+        version_file = f"{cache_dir}/version.txt"
+
+        if not await self.storage_manager.file_exists(version_file):
+            return False
+
+        try:
+            async with self.storage_manager.open(version_file, "r") as f:
+                cached_version = (await asyncio.to_thread(f.read)).strip()
+                return cached_version == version_or_hash
+        except Exception as e:
+            logger.debug(f"Error checking cache validity: {e}")
+            return False
+
+    async def _clear_cache(self, cache_dir: str) -> None:
+        """Clear a cache directory."""
+        try:
+            await self.storage_manager.remove_all(cache_dir)
+        except Exception as e:
+            logger.debug(f"Error clearing cache directory {cache_dir}: {e}")
+
+    async def _check_remote_content_freshness(
+        self, url: str, cache_dir: str
+    ) -> Tuple[bool, Optional[str]]:
+        """
+        Check if remote content is fresher than cached version using HTTP headers.
+
+        Returns:
+            Tuple of (is_fresh: bool, new_identifier: Optional[str])
+        """
+        try:
+            # Make a HEAD request to check headers without downloading
+            ssl_context = create_secure_ssl_context()
+            connector = aiohttp.TCPConnector(ssl=ssl_context)
+            async with aiohttp.ClientSession(connector=connector) as session:
+                async with session.head(url, timeout=aiohttp.ClientTimeout(total=30)) as response:
+                    response.raise_for_status()
+
+                    # Try ETag first (most reliable)
+                    etag = response.headers.get("ETag", "").strip('"')
+                    last_modified = response.headers.get("Last-Modified", "")
+
+            # Use ETag if available, otherwise Last-Modified
+            remote_identifier = etag if etag else last_modified
+
+            if not remote_identifier:
+                logger.debug("No freshness headers available, cannot check for updates")
+                return True, None  # Assume fresh if no headers
+
+            # Check cached identifier
+            identifier_file = f"{cache_dir}/content_id.txt"
+            if await self.storage_manager.file_exists(identifier_file):
+                async with self.storage_manager.open(identifier_file, "r") as f:
+                    cached_identifier = (await asyncio.to_thread(f.read)).strip()
+                    if cached_identifier == remote_identifier:
+                        logger.debug(f"Content is fresh (identifier: {remote_identifier[:20]}...)")
+                        return True, None
+                    else:
+                        logger.info(
+                            f"Content has changed (old: {cached_identifier[:20]}..., new: {remote_identifier[:20]}...)"
+                        )
+                        return False, remote_identifier
+            else:
+                # No cached identifier, treat as stale
+                return False, remote_identifier
+
+        except Exception as e:
+            logger.debug(f"Could not check remote freshness: {e}")
+            return True, None  # Assume fresh if we can't check
+
+    async def download_and_extract_zip(
+        self, url: str, cache_subdir_name: str, version_or_hash: str, force: bool = False
+    ) -> str:
+        """
+        Download a zip file and extract it to cache directory with content freshness checking.
+
+        Args:
+            url: URL to download zip file from
+            cache_subdir_name: Name of the cache subdirectory
+            version_or_hash: Version string or content hash for cache validation
+            force: If True, re-download even if already cached
+
+        Returns:
+            Path to the cached directory
+        """
+        cache_dir = await self.get_cache_subdir(cache_subdir_name)
+
+        # Check if already cached and valid
+        if not force and await self._is_cache_valid(cache_dir, version_or_hash):
+            # Also check if remote content has changed
+            is_fresh, new_identifier = await self._check_remote_content_freshness(url, cache_dir)
+            if is_fresh:
+                logger.debug(f"Content already cached and fresh for version {version_or_hash}")
+                return cache_dir
+            else:
+                logger.info("Cached content is stale, updating...")
+
+        # Clear old cache if it exists
+        await self._clear_cache(cache_dir)
+
+        logger.info(f"Downloading content from {url}...")
+
+        # Download the zip file
+        zip_content = BytesIO()
+        etag = ""
+        last_modified = ""
+        ssl_context = create_secure_ssl_context()
+        connector = aiohttp.TCPConnector(ssl=ssl_context)
+        async with aiohttp.ClientSession(connector=connector) as session:
+            async with session.get(url, timeout=aiohttp.ClientTimeout(total=60)) as response:
+                response.raise_for_status()
+
+                # Extract headers before consuming response
+                etag = response.headers.get("ETag", "").strip('"')
+                last_modified = response.headers.get("Last-Modified", "")
+
+                # Read the response content
+                async for chunk in response.content.iter_chunked(8192):
+                    zip_content.write(chunk)
+        zip_content.seek(0)
+
+        # Extract the archive
+        await self.storage_manager.ensure_directory_exists(cache_dir)
+
+        # Extract files and store them using StorageManager
+        with zipfile.ZipFile(zip_content, "r") as zip_file:
+            for file_info in zip_file.infolist():
+                if file_info.is_dir():
+                    # Create directory
+                    dir_path = f"{cache_dir}/{file_info.filename}"
+                    await self.storage_manager.ensure_directory_exists(dir_path)
+                else:
+                    # Extract and store file
+                    file_data = zip_file.read(file_info.filename)
+                    file_path = f"{cache_dir}/{file_info.filename}"
+                    await self.storage_manager.store(file_path, BytesIO(file_data), overwrite=True)
+
+        # Write version info for future cache validation
+        version_file = f"{cache_dir}/version.txt"
+        await self.storage_manager.store(version_file, version_or_hash, overwrite=True)
+
+        # Store content identifier from response headers for freshness checking
+        content_identifier = etag if etag else last_modified
+
+        if content_identifier:
+            identifier_file = f"{cache_dir}/content_id.txt"
+            await self.storage_manager.store(identifier_file, content_identifier, overwrite=True)
+            logger.debug(f"Stored content identifier: {content_identifier[:20]}...")
+
+        logger.info("✓ Content downloaded and cached successfully!")
+        return cache_dir
+
+    async def file_exists(self, file_path: str) -> bool:
+        """Check if a file exists in cache storage."""
+        return await self.storage_manager.file_exists(file_path)
+
+    async def read_file(self, file_path: str, encoding: str = "utf-8"):
+        """Read a file from cache storage."""
+        return self.storage_manager.open(file_path, encoding=encoding)
+
+    async def list_files(self, directory_path: str):
+        """List files in a cache directory."""
+        try:
+            file_list = await self.storage_manager.list_files(directory_path)
+
+            # For S3 storage, convert relative paths to full S3 URLs
+            if self.storage_manager.storage.storage_path.startswith("s3://"):
+                full_paths = []
+                for file_path in file_list:
+                    full_s3_path = f"{self.storage_manager.storage.storage_path}/{file_path}"
+                    full_paths.append(full_s3_path)
+                return full_paths
+            else:
+                # For local storage, return absolute paths
+                storage_path = self.storage_manager.storage.storage_path
+                if not storage_path.startswith("/"):
+                    import os
+
+                    storage_path = os.path.abspath(storage_path)
+
+                full_paths = []
+                for file_path in file_list:
+                    if file_path.startswith("/"):
+                        full_paths.append(file_path)  # Already absolute
+                    else:
+                        full_paths.append(f"{storage_path}/{file_path}")
+                return full_paths
+
+        except Exception as e:
+            logger.debug(f"Error listing files in {directory_path}: {e}")
+            return []
+
+
+# Convenience functions that maintain API compatibility
+_cache_manager = None
+
+
+def get_cache_manager() -> StorageAwareCache:
+    """Get a singleton cache manager instance."""
+    global _cache_manager
+    if _cache_manager is None:
+        _cache_manager = StorageAwareCache()
+    return _cache_manager
+
+
+def generate_content_hash(url: str, additional_data: str = "") -> str:
+    """Generate a content hash from URL and optional additional data."""
+    content = f"{url}:{additional_data}"
+    return hashlib.md5(content.encode()).hexdigest()[:12]  # Short hash for readability
+
+
+# Async wrapper functions for backward compatibility
+async def delete_cache():
+    """Delete the Cognee cache directory."""
+    cache_manager = get_cache_manager()
+    await cache_manager.delete_cache()
+
+
+async def get_cognee_cache_dir() -> str:
+    """Get the base Cognee cache directory."""
+    cache_manager = get_cache_manager()
+    return await cache_manager.get_cache_dir()
+
+
+async def get_cache_subdir(name: str) -> str:
+    """Get a specific cache subdirectory."""
+    cache_manager = get_cache_manager()
+    return await cache_manager.get_cache_subdir(name)
+
+
+async def download_and_extract_zip(
+    url: str, cache_dir_name: str, version_or_hash: str, force: bool = False
+) -> str:
+    """Download a zip file and extract it to cache directory."""
+    cache_manager = get_cache_manager()
+    return await cache_manager.download_and_extract_zip(url, cache_dir_name, version_or_hash, force)
+
+
+async def get_tutorial_data_dir() -> str:
+    """Get the tutorial data cache directory."""
+    return await get_cache_subdir("tutorial_data")
+
+
+# Cache file operations
+async def cache_file_exists(file_path: str) -> bool:
+    """Check if a file exists in cache storage."""
+    cache_manager = get_cache_manager()
+    return await cache_manager.file_exists(file_path)
+
+
+async def read_cache_file(file_path: str, encoding: str = "utf-8"):
+    """Read a file from cache storage."""
+    cache_manager = get_cache_manager()
+    return await cache_manager.read_file(file_path, encoding)
+
+
+async def list_cache_files(directory_path: str):
+    """List files in a cache directory."""
+    cache_manager = get_cache_manager()
+    return await cache_manager.list_files(directory_path)

cognee/shared/utils.py

@@ -1,6 +1,7 @@
 """This module contains utility functions for the cognee."""
 
 import os
+import ssl
 import requests
 from datetime import datetime, timezone
 import matplotlib.pyplot as plt
@@ -18,6 +19,17 @@ from cognee.infrastructure.databases.graph import get_graph_engine
 proxy_url = "https://test.prometh.ai"
 
 
+def create_secure_ssl_context() -> ssl.SSLContext:
+    """
+    Create a secure SSL context.
+
+    By default, use the system's certificate store.
+    If users report SSL issues, I'm keeping this open in case we need to switch to:
+        ssl.create_default_context(cafile=certifi.where())
+    """
+    return ssl.create_default_context()
+
+
 def get_entities(tagged_tokens):
     import nltk
 

cognee/tasks/graph/extract_graph_from_data.py

@@ -3,8 +3,14 @@ from typing import Type, List, Optional
 from pydantic import BaseModel
 
 from cognee.infrastructure.databases.graph import get_graph_engine
+from cognee.modules.ontology.ontology_env_config import get_ontology_env_config
 from cognee.tasks.storage.add_data_points import add_data_points
-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.ontology.base_ontology_resolver import BaseOntologyResolver
 from cognee.modules.chunking.models.DocumentChunk import DocumentChunk
 from cognee.modules.graph.utils import (
     expand_with_nodes_and_edges,
@@ -24,9 +30,28 @@ async def integrate_chunk_graphs(
     data_chunks: list[DocumentChunk],
     chunk_graphs: list,
     graph_model: Type[BaseModel],
-    ontology_adapter: OntologyResolver,
+    ontology_resolver: BaseOntologyResolver,
 ) -> List[DocumentChunk]:
-    """Updates DocumentChunk objects, integrates data points and edges into databases."""
+    """Integrate chunk graphs with ontology validation and store in databases.
+
+    This function processes document chunks and their associated knowledge graphs,
+    validates entities against an ontology resolver, and stores the integrated
+    data points and edges in the configured databases.
+
+    Args:
+        data_chunks: List of document chunks containing source data
+        chunk_graphs: List of knowledge graphs corresponding to each chunk
+        graph_model: Pydantic model class for graph data validation
+        ontology_resolver: Resolver for validating entities against ontology
+
+    Returns:
+        List of updated DocumentChunk objects with integrated data
+
+    Raises:
+        InvalidChunkGraphInputError: If input validation fails
+        InvalidGraphModelError: If graph model validation fails
+        InvalidOntologyAdapterError: If ontology resolver validation fails
+    """
 
     if not isinstance(data_chunks, list) or not isinstance(chunk_graphs, list):
         raise InvalidChunkGraphInputError("data_chunks and chunk_graphs must be lists.")
@@ -36,9 +61,9 @@ async def integrate_chunk_graphs(
         )
     if not isinstance(graph_model, type) or not issubclass(graph_model, BaseModel):
         raise InvalidGraphModelError(graph_model)
-    if ontology_adapter is None or not hasattr(ontology_adapter, "get_subgraph"):
+    if ontology_resolver is None or not hasattr(ontology_resolver, "get_subgraph"):
         raise InvalidOntologyAdapterError(
-            type(ontology_adapter).__name__ if ontology_adapter else "None"
+            type(ontology_resolver).__name__ if ontology_resolver else "None"
         )
 
     graph_engine = await get_graph_engine()
@@ -55,7 +80,7 @@ async def integrate_chunk_graphs(
     )
 
     graph_nodes, graph_edges = expand_with_nodes_and_edges(
-        data_chunks, chunk_graphs, ontology_adapter, existing_edges_map
+        data_chunks, chunk_graphs, ontology_resolver, existing_edges_map
     )
 
     if len(graph_nodes) > 0:
@@ -70,7 +95,7 @@ async def integrate_chunk_graphs(
 async def extract_graph_from_data(
     data_chunks: List[DocumentChunk],
     graph_model: Type[BaseModel],
-    ontology_adapter: OntologyResolver = None,
+    config: Config = None,
     custom_prompt: Optional[str] = None,
 ) -> List[DocumentChunk]:
     """
@@ -101,6 +126,24 @@ async def extract_graph_from_data(
                 if edge.source_node_id in valid_node_ids and edge.target_node_id in valid_node_ids
             ]
 
-    return await integrate_chunk_graphs(
-        data_chunks, chunk_graphs, graph_model, ontology_adapter or OntologyResolver()
-    )
+    # Extract resolver from config if provided, otherwise get default
+    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()}
+            }
+
+    ontology_resolver = config["ontology_config"]["ontology_resolver"]
+
+    return await integrate_chunk_graphs(data_chunks, chunk_graphs, graph_model, ontology_resolver)

cognee/tasks/graph/extract_graph_from_data_v2.py

@@ -3,7 +3,7 @@ from typing import List
 
 from cognee.modules.chunking.models.DocumentChunk import DocumentChunk
 from cognee.shared.data_models import KnowledgeGraph
-from cognee.modules.ontology.rdf_xml.OntologyResolver import OntologyResolver
+from cognee.modules.ontology.base_ontology_resolver import BaseOntologyResolver
 from cognee.tasks.graph.cascade_extract.utils.extract_nodes import extract_nodes
 from cognee.tasks.graph.cascade_extract.utils.extract_content_nodes_and_relationship_names import (
     extract_content_nodes_and_relationship_names,
@@ -17,9 +17,21 @@ from cognee.tasks.graph.extract_graph_from_data import integrate_chunk_graphs
 async def extract_graph_from_data(
     data_chunks: List[DocumentChunk],
     n_rounds: int = 2,
-    ontology_adapter: OntologyResolver = None,
+    ontology_adapter: BaseOntologyResolver = None,
 ) -> List[DocumentChunk]:
-    """Extract and update graph data from document chunks in multiple steps."""
+    """Extract and update graph data from document chunks using cascade extraction.
+
+    This function performs multi-step graph extraction from document chunks,
+    using cascade extraction techniques to build comprehensive knowledge graphs.
+
+    Args:
+        data_chunks: List of document chunks to process
+        n_rounds: Number of extraction rounds to perform (default: 2)
+        ontology_adapter: Resolver for validating entities against ontology
+
+    Returns:
+        List of updated DocumentChunk objects with extracted graph data
+    """
     chunk_nodes = await asyncio.gather(
         *[extract_nodes(chunk.text, n_rounds) for chunk in data_chunks]
     )
@@ -44,5 +56,5 @@ async def extract_graph_from_data(
         data_chunks=data_chunks,
         chunk_graphs=chunk_graphs,
         graph_model=KnowledgeGraph,
-        ontology_adapter=ontology_adapter or OntologyResolver(),
+        ontology_adapter=ontology_adapter,
     )

cognee/tasks/ingestion/save_data_item_to_storage.py

@@ -41,6 +41,7 @@ async def save_data_item_to_storage(data_item: Union[BinaryIO, str, Any]) -> str
             abs_path.is_file()
         except (OSError, ValueError):
             # In case file path is too long it's most likely not a relative path
+            abs_path = data_item
             logger.debug(f"Data item was too long to be a possible file path: {abs_path}")
             abs_path = Path("")
 

cognee/tasks/temporal_graph/models.py

@@ -3,12 +3,17 @@ from pydantic import BaseModel, Field
 
 
 class Timestamp(BaseModel):
-    year: int = Field(..., ge=1, le=9999)
-    month: int = Field(..., ge=1, le=12)
-    day: int = Field(..., ge=1, le=31)
-    hour: int = Field(..., ge=0, le=23)
-    minute: int = Field(..., ge=0, le=59)
-    second: int = Field(..., ge=0, le=59)
+    year: int = Field(
+        ...,
+        ge=1,
+        le=9999,
+        description="Always required. If only a year is known, use it.",
+    )
+    month: int = Field(1, ge=1, le=12, description="If unknown, default to 1")
+    day: int = Field(1, ge=1, le=31, description="If unknown, default to 1")
+    hour: int = Field(0, ge=0, le=23, description="If unknown, default to 0")
+    minute: int = Field(0, ge=0, le=59, description="If unknown, default to 0")
+    second: int = Field(0, ge=0, le=59, description="If unknown, default to 0")
 
 
 class Interval(BaseModel):

cognee/tests/cli_tests/cli_unit_tests/test_cli_main.py

@@ -49,7 +49,7 @@ class TestCliMain:
     def test_main_no_command(self, mock_create_parser):
         """Test main function when no command is provided"""
         mock_parser = MagicMock()
-        mock_parser.parse_args.return_value = MagicMock(command=None)
+        mock_parser.parse_args.return_value = MagicMock(command=None, spec={})
         mock_create_parser.return_value = (mock_parser, {})
 
         result = main()
@@ -64,7 +64,7 @@ class TestCliMain:
         mock_command.execute.return_value = None
 
         mock_parser = MagicMock()
-        mock_args = MagicMock(command="test")
+        mock_args = MagicMock(command="test", spec={})
         mock_parser.parse_args.return_value = mock_args
 
         mock_create_parser.return_value = (mock_parser, {"test": mock_command})
@@ -84,7 +84,7 @@ class TestCliMain:
         mock_command.execute.side_effect = CliCommandException("Test error", error_code=2)
 
         mock_parser = MagicMock()
-        mock_args = MagicMock(command="test")
+        mock_args = MagicMock(command="test", spec={})
         mock_parser.parse_args.return_value = mock_args
 
         mock_create_parser.return_value = (mock_parser, {"test": mock_command})
@@ -103,7 +103,7 @@ class TestCliMain:
         mock_command.execute.side_effect = Exception("Generic error")
 
         mock_parser = MagicMock()
-        mock_args = MagicMock(command="test")
+        mock_args = MagicMock(command="test", spec={})
         mock_parser.parse_args.return_value = mock_args
 
         mock_create_parser.return_value = (mock_parser, {"test": mock_command})
@@ -126,7 +126,7 @@ class TestCliMain:
         mock_command.execute.side_effect = test_exception
 
         mock_parser = MagicMock()
-        mock_args = MagicMock(command="test")
+        mock_args = MagicMock(command="test", spec={})
         mock_parser.parse_args.return_value = mock_args
 
         mock_create_parser.return_value = (mock_parser, {"test": mock_command})