cognee 0.3.2__tar.gz → 0.3.4__tar.gz

{cognee-0.3.2 → cognee-0.3.4}/.env.template

@@ -47,6 +47,28 @@ BAML_LLM_API_VERSION=""
 # DATA_ROOT_DIRECTORY='/Users/<user>/Desktop/cognee/.cognee_data/'
 # SYSTEM_ROOT_DIRECTORY='/Users/<user>/Desktop/cognee/.cognee_system/'
 
+################################################################################
+# ☁️ Storage Backend Settings
+################################################################################
+# Configure storage backend (local filesystem or S3)
+# STORAGE_BACKEND="local"  # Default: uses local filesystem
+# 
+# -- To switch to S3 storage, uncomment and fill these: ---------------------
+# STORAGE_BACKEND="s3"
+# STORAGE_BUCKET_NAME="your-bucket-name"
+# AWS_REGION="us-east-1"
+# AWS_ACCESS_KEY_ID="your-access-key"
+# AWS_SECRET_ACCESS_KEY="your-secret-key"
+# 
+# -- S3 Root Directories (optional) -----------------------------------------
+# DATA_ROOT_DIRECTORY="s3://your-bucket/cognee/data"
+# SYSTEM_ROOT_DIRECTORY="s3://your-bucket/cognee/system"
+# 
+# -- Cache Directory (auto-configured for S3) -------------------------------
+# When STORAGE_BACKEND=s3, cache automatically uses S3: s3://BUCKET/cognee/cache
+# To override the automatic S3 cache location, uncomment:
+# CACHE_ROOT_DIRECTORY="s3://your-bucket/cognee/cache"
+
 ################################################################################
 # 🗄️ Relational database settings
 ################################################################################
@@ -94,7 +116,15 @@ VECTOR_DB_PROVIDER="lancedb"
 VECTOR_DB_URL=
 VECTOR_DB_KEY=
 
+################################################################################
+# 🧩 Ontology resolver settings
+################################################################################
 
+# -- Ontology resolver params --------------------------------------
+# ONTOLOGY_RESOLVER=rdflib  # Default: uses rdflib and owl file to read ontology structures
+# MATCHING_STRATEGY=fuzzy # Default: uses fuzzy matching with 80% similarity threshold
+# ONTOLOGY_FILE_PATH=YOUR_FULL_FULE_PATH # Default: empty
+# To add ontology resolvers, either set them as it is set in ontology_example or add full_path and settings as envs.
 
 ################################################################################
 #  🔄  MIGRATION (RELATIONAL → GRAPH) SETTINGS
@@ -121,6 +151,9 @@ ACCEPT_LOCAL_FILE_PATH=True
 # This protects against Server Side Request Forgery when proper infrastructure is not in place.
 ALLOW_HTTP_REQUESTS=True
 
+# When set to false don't allow cypher search to be used in Cognee.
+ALLOW_CYPHER_QUERY=True
+
 # When set to False errors during data processing will be returned as info but not raised to allow handling of faulty documents
 RAISE_INCREMENTAL_LOADING_ERRORS=True
 

{cognee-0.3.2 → cognee-0.3.4}/.gitignore

@@ -186,6 +186,7 @@ cognee/cache/
 # Default cognee system directory, used in development
 .cognee_system/
 .data_storage/
+.cognee_cache/
 .artifacts/
 .anon_id
 

{cognee-0.3.2 → cognee-0.3.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cognee
-Version: 0.3.2
+Version: 0.3.4
 Summary: Cognee - is a library for enriching LLM context with a semantic layer for better understanding and reasoning.
 Project-URL: Homepage, https://www.cognee.ai
 Project-URL: Repository, https://github.com/topoteretes/cognee
@@ -25,6 +25,7 @@ Requires-Dist: dlt[sqlalchemy]<2,>=1.9.0
 Requires-Dist: fastapi-users[sqlalchemy]<15.0.0,>=14.0.1
 Requires-Dist: fastapi<1.0.0,>=0.115.7
 Requires-Dist: filetype<2.0.0,>=1.2.0
+Requires-Dist: gunicorn<24,>=20.1.0
 Requires-Dist: instructor<2.0.0,>=1.9.1
 Requires-Dist: jinja2<4,>=3.1.3
 Requires-Dist: kuzu==0.11.0
@@ -33,6 +34,7 @@ Requires-Dist: langfuse<3,>=2.32.0
 Requires-Dist: limits<5,>=4.4.1
 Requires-Dist: litellm<2.0.0,>=1.71.0
 Requires-Dist: matplotlib<4,>=3.8.3
+Requires-Dist: nbformat<6.0.0,>=5.7.0
 Requires-Dist: networkx<4,>=3.4.2
 Requires-Dist: nltk<4.0.0,>=3.9.1
 Requires-Dist: numpy<=4.0.0,>=1.26.4
@@ -56,17 +58,16 @@ Requires-Dist: sqlalchemy<3.0.0,>=2.0.39
 Requires-Dist: structlog<26,>=25.2.0
 Requires-Dist: tiktoken<1.0.0,>=0.8.0
 Requires-Dist: typing-extensions<5.0.0,>=4.12.2
+Requires-Dist: uvicorn<1.0.0,>=0.34.0
+Requires-Dist: websockets<16.0.0,>=15.0.1
 Provides-Extra: anthropic
-Requires-Dist: anthropic<0.27,>=0.26.1; extra == 'anthropic'
+Requires-Dist: anthropic>=0.27; extra == 'anthropic'
 Provides-Extra: api
-Requires-Dist: gunicorn<24,>=20.1.0; extra == 'api'
-Requires-Dist: uvicorn<1.0.0,>=0.34.0; extra == 'api'
-Requires-Dist: websockets<16.0.0,>=15.0.1; extra == 'api'
 Provides-Extra: aws
 Requires-Dist: s3fs[boto3]==2025.3.2; extra == 'aws'
 Provides-Extra: chromadb
-Requires-Dist: chromadb<0.7,>=0.3.0; extra == 'chromadb'
-Requires-Dist: pypika==0.48.8; extra == 'chromadb'
+Requires-Dist: chromadb<0.7,>=0.6; extra == 'chromadb'
+Requires-Dist: pypika==0.48.9; extra == 'chromadb'
 Provides-Extra: codegraph
 Requires-Dist: fastembed<=0.6.0; (python_version < '3.13') and extra == 'codegraph'
 Requires-Dist: transformers<5,>=4.46.3; extra == 'codegraph'
@@ -316,7 +317,9 @@ You can also cognify your files and query using cognee UI.
 
 <img src="assets/cognee-new-ui.webp" width="100%" alt="Cognee UI 2"></a>
 
-Try cognee UI by runnning ``` cognee -ui ``` command on your terminal.
+### Running the UI
+
+Try cognee UI by running ``` cognee-cli -ui ``` command on your terminal.
 
 ## Understand our architecture
 

{cognee-0.3.2 → cognee-0.3.4}/README.md

@@ -176,7 +176,9 @@ You can also cognify your files and query using cognee UI.
 
 <img src="assets/cognee-new-ui.webp" width="100%" alt="Cognee UI 2"></a>
 
-Try cognee UI by runnning ``` cognee -ui ``` command on your terminal.
+### Running the UI
+
+Try cognee UI by running ``` cognee-cli -ui ``` command on your terminal.
 
 ## Understand our architecture
 

{cognee-0.3.2 → cognee-0.3.4}/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-0.3.2 → cognee-0.3.4}/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-0.3.2 → cognee-0.3.4}/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-0.3.2 → cognee-0.3.4}/cognee/api/v1/notebooks/routers/get_notebooks_router.py

@@ -31,7 +31,8 @@ def get_notebooks_router():
 
     @router.get("")
     async def get_notebooks_endpoint(user: User = Depends(get_authenticated_user)):
-        return await get_notebooks(user.id)
+        async with get_async_session() as session:
+            return await get_notebooks(user.id, session)
 
     @router.post("")
     async def create_notebook_endpoint(

{cognee-0.3.2 → cognee-0.3.4}/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-0.3.2 → cognee-0.3.4}/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-0.3.2 → cognee-0.3.4}/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-0.3.2 → cognee-0.3.4}/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-0.3.2 → cognee-0.3.4}/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-0.3.2 → cognee-0.3.4}/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-0.3.2 → cognee-0.3.4}/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-0.3.2 → cognee-0.3.4}/cognee/infrastructure/databases/relational/sqlalchemy/SqlAlchemyAdapter.py

@@ -58,8 +58,8 @@ class SQLAlchemyAdapter:
         else:
             self.engine = create_async_engine(
                 connection_string,
-                pool_size=5,
-                max_overflow=10,
+                pool_size=20,
+                max_overflow=20,
                 pool_recycle=280,
                 pool_pre_ping=True,
                 pool_timeout=280,

{cognee-0.3.2 → cognee-0.3.4}/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-0.3.2 → cognee-0.3.4}/cognee/infrastructure/databases/vector/config.py

@@ -39,7 +39,7 @@ class VectorConfig(BaseSettings):
             values.vector_db_url = ensure_absolute_path(
                 values.vector_db_url,
             )
-        else:
+        elif not values.vector_db_url:
             # Default path
             databases_directory_path = os.path.join(base_config.system_root_directory, "databases")
             values.vector_db_url = os.path.join(databases_directory_path, "cognee.lancedb")

{cognee-0.3.2 → cognee-0.3.4}/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-0.3.2 → cognee-0.3.4}/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