From 322ef156cb5efce2d75bc7e0df0ebec9484903c9 Mon Sep 17 00:00:00 2001 From: Daulet Amirkhanov Date: Tue, 21 Oct 2025 17:10:45 +0100 Subject: [PATCH] redefine `preferred_loaders` param to allow for args per loader --- cognee/api/v1/add/add.py | 2 +- cognee/api/v1/update/update.py | 4 ++-- cognee/infrastructure/loaders/LoaderEngine.py | 6 ++++-- cognee/tasks/ingestion/data_item_to_text_file.py | 4 ++-- cognee/tasks/ingestion/ingest_data.py | 4 ++-- 5 files changed, 11 insertions(+), 9 deletions(-) diff --git a/cognee/api/v1/add/add.py b/cognee/api/v1/add/add.py index 216911ec0..73a3081be 100644 --- a/cognee/api/v1/add/add.py +++ b/cognee/api/v1/add/add.py @@ -23,7 +23,7 @@ async def add( vector_db_config: dict = None, graph_db_config: dict = None, dataset_id: Optional[UUID] = None, - preferred_loaders: List[str] = None, + preferred_loaders: dict[str, dict[str, Any]] = None, incremental_loading: bool = True, data_per_batch: Optional[int] = 20, ): diff --git a/cognee/api/v1/update/update.py b/cognee/api/v1/update/update.py index a421b3dc0..83b92c50f 100644 --- a/cognee/api/v1/update/update.py +++ b/cognee/api/v1/update/update.py @@ -1,5 +1,5 @@ from uuid import UUID -from typing import Union, BinaryIO, List, Optional +from typing import Union, BinaryIO, List, Optional, Any from cognee.modules.users.models import User from cognee.api.v1.delete import delete @@ -15,7 +15,7 @@ async def update( node_set: Optional[List[str]] = None, vector_db_config: dict = None, graph_db_config: dict = None, - preferred_loaders: List[str] = None, + preferred_loaders: dict[str, dict[str, Any]] = None, incremental_loading: bool = True, ): """ diff --git a/cognee/infrastructure/loaders/LoaderEngine.py b/cognee/infrastructure/loaders/LoaderEngine.py index 6b62f7641..84ecee0de 100644 --- a/cognee/infrastructure/loaders/LoaderEngine.py +++ b/cognee/infrastructure/loaders/LoaderEngine.py @@ -64,7 +64,9 @@ class LoaderEngine: return True def get_loader( - self, file_path: str, preferred_loaders: List[str] = None + self, + file_path: str, + preferred_loaders: dict[str, dict[str, Any]], ) -> Optional[LoaderInterface]: """ Get appropriate loader for a file. @@ -105,7 +107,7 @@ class LoaderEngine: async def load_file( self, file_path: str, - preferred_loaders: Optional[List[str]] = None, + preferred_loaders: dict[str, dict[str, Any]] = None, **kwargs, ): """ diff --git a/cognee/tasks/ingestion/data_item_to_text_file.py b/cognee/tasks/ingestion/data_item_to_text_file.py index 4b9e4bb23..0303f6c92 100644 --- a/cognee/tasks/ingestion/data_item_to_text_file.py +++ b/cognee/tasks/ingestion/data_item_to_text_file.py @@ -1,6 +1,6 @@ import os from urllib.parse import urlparse -from typing import List, Tuple +from typing import Any, List, Tuple from pathlib import Path import tempfile @@ -35,7 +35,7 @@ async def pull_from_s3(file_path, destination_file) -> None: async def data_item_to_text_file( data_item_path: str, - preferred_loaders: List[str], + preferred_loaders: dict[str, dict[str, Any]] = None, ) -> Tuple[str, LoaderInterface]: if isinstance(data_item_path, str): parsed_url = urlparse(data_item_path) diff --git a/cognee/tasks/ingestion/ingest_data.py b/cognee/tasks/ingestion/ingest_data.py index 02987b893..7b081cc34 100644 --- a/cognee/tasks/ingestion/ingest_data.py +++ b/cognee/tasks/ingestion/ingest_data.py @@ -27,7 +27,7 @@ async def ingest_data( user: User, node_set: Optional[List[str]] = None, dataset_id: UUID = None, - preferred_loaders: List[str] = None, + preferred_loaders: dict[str, dict[str, Any]] = None, ): if not user: user = await get_default_user() @@ -44,7 +44,7 @@ async def ingest_data( user: User, node_set: Optional[List[str]] = None, dataset_id: UUID = None, - preferred_loaders: List[str] = None, + preferred_loaders: dict[str, dict[str, Any]] = None, ): new_datapoints = [] existing_data_points = []