cognee 0.2.1.dev4__tar.gz → 0.2.1.dev6__tar.gz

{cognee-0.2.1.dev4 → cognee-0.2.1.dev6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cognee
-Version: 0.2.1.dev4
+Version: 0.2.1.dev6
 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

{cognee-0.2.1.dev4 → cognee-0.2.1.dev6}/cognee/api/v1/cognify/cognify.py

@@ -237,8 +237,12 @@ async def run_cognify_as_background_process(
     graph_db_config: dict = None,
     vector_db_config: dict = False,
 ):
+    # Convert dataset to list if it's a string
+    if isinstance(datasets, str):
+        datasets = [datasets]
+
     # Store pipeline status for all pipelines
-    pipeline_run_started_info = []
+    pipeline_run_started_info = {}
 
     async def handle_rest_of_the_run(pipeline_list):
         # Execute all provided pipelines one by one to avoid database write conflicts
@@ -269,7 +273,14 @@ async def run_cognify_as_background_process(
         )
 
         # Save dataset Pipeline run started info
-        pipeline_run_started_info.append(await anext(pipeline_run))
+        run_info = await anext(pipeline_run)
+        pipeline_run_started_info[run_info.dataset_id] = run_info
+
+        if pipeline_run_started_info[run_info.dataset_id].payload:
+            # Remove payload info to avoid serialization
+            # TODO: Handle payload serialization
+            pipeline_run_started_info[run_info.dataset_id].payload = []
+
         pipeline_list.append(pipeline_run)
 
     # Send all started pipelines to execute one by one in background

{cognee-0.2.1.dev4 → cognee-0.2.1.dev6}/cognee/api/v1/cognify/routers/get_cognify_router.py

@@ -32,7 +32,6 @@ logger = get_logger("api.cognify")
 class CognifyPayloadDTO(InDTO):
     datasets: Optional[List[str]] = None
     dataset_ids: Optional[List[UUID]] = None
-    graph_model: Optional[BaseModel] = KnowledgeGraph
     run_in_background: Optional[bool] = False
 
 
@@ -96,7 +95,7 @@ def get_cognify_router() -> APIRouter:
             datasets = payload.dataset_ids if payload.dataset_ids else payload.datasets
 
             cognify_run = await cognee_cognify(
-                datasets, user, payload.graph_model, run_in_background=payload.run_in_background
+                datasets, user, run_in_background=payload.run_in_background
             )
 
             return cognify_run

{cognee-0.2.1.dev4 → cognee-0.2.1.dev6}/cognee/api/v1/datasets/routers/get_datasets_router.py

@@ -331,11 +331,19 @@ def get_datasets_router() -> APIRouter:
         ## Error Codes
         - **500 Internal Server Error**: Error retrieving status information
         """
-        from cognee.modules.data.methods import get_dataset_status
+        from cognee.api.v1.datasets.datasets import datasets as cognee_datasets
 
-        dataset_status = await get_dataset_status(datasets, user.id)
+        try:
+            # Verify user has permission to read dataset
+            authorized_datasets = await get_authorized_existing_datasets(datasets, "read", user)
+
+            datasets_statuses = await cognee_datasets.get_status(
+                [dataset.id for dataset in authorized_datasets]
+            )
 
-        return dataset_status
+            return datasets_statuses
+        except Exception as error:
+            return JSONResponse(status_code=409, content={"error": str(error)})
 
     @router.get("/{dataset_id}/data/{data_id}/raw", response_class=FileResponse)
     async def get_raw_data(

{cognee-0.2.1.dev4 → cognee-0.2.1.dev6}/cognee/infrastructure/files/storage/S3FileStorage.py

@@ -26,7 +26,8 @@ class S3FileStorage(Storage):
                 key=s3_config.aws_access_key_id,
                 secret=s3_config.aws_secret_access_key,
                 anon=False,
-                endpoint_url="https://s3-eu-west-1.amazonaws.com",
+                endpoint_url=s3_config.aws_endpoint_url,
+                client_kwargs={"region_name": s3_config.aws_region},
             )
         else:
             raise ValueError("S3 credentials are not set in the configuration.")

{cognee-0.2.1.dev4 → cognee-0.2.1.dev6}/cognee/infrastructure/files/storage/s3_config.py

@@ -4,6 +4,8 @@ from pydantic_settings import BaseSettings, SettingsConfigDict
 
 
 class S3Config(BaseSettings):
+    aws_region: Optional[str] = None
+    aws_endpoint_url: Optional[str] = None
     aws_access_key_id: Optional[str] = None
     aws_secret_access_key: Optional[str] = None
     model_config = SettingsConfigDict(env_file=".env", extra="allow")

{cognee-0.2.1.dev4 → cognee-0.2.1.dev6}/cognee/shared/utils.py

@@ -54,6 +54,11 @@ def extract_pos_tags(sentence):
 
 def get_anonymous_id():
     """Creates or reads a anonymous user id"""
+    tracking_id = os.getenv("TRACKING_ID", None)
+
+    if tracking_id:
+        return tracking_id
+
     home_dir = str(pathlib.Path(pathlib.Path(__file__).parent.parent.parent.resolve()))
 
     if not os.path.isdir(home_dir):

{cognee-0.2.1.dev4 → cognee-0.2.1.dev6}/cognee/tests/test_s3_file_storage.py

@@ -13,10 +13,11 @@ logger = get_logger()
 
 
 async def main():
+    bucket_name = os.getenv("STORAGE_BUCKET_NAME")
     test_run_id = uuid4()
-    data_directory_path = f"s3://cognee-storage-dev/{test_run_id}/data"
+    data_directory_path = f"s3://{bucket_name}/{test_run_id}/data"
     cognee.config.data_root_directory(data_directory_path)
-    cognee_directory_path = f"s3://cognee-storage-dev/{test_run_id}/system"
+    cognee_directory_path = f"s3://{bucket_name}/{test_run_id}/system"
     cognee.config.system_root_directory(cognee_directory_path)
 
     await cognee.prune.prune_data()

{cognee-0.2.1.dev4 → cognee-0.2.1.dev6}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "cognee"
-version = "0.2.1-dev4"
+version = "0.2.1-dev6"
 description = "Cognee - is a library for enriching LLM context with a semantic layer for better understanding and reasoning."
 authors = [
     { name = "Vasilije Markovic" },

{cognee-0.2.1.dev4 → cognee-0.2.1.dev6}/uv.lock

@@ -878,7 +878,7 @@ wheels = [
 
 [[package]]
 name = "cognee"
-version = "0.2.1.dev4"
+version = "0.2.1.dev6"
 source = { editable = "." }
 dependencies = [
     { name = "aiofiles" },