From 83995fa548869ae0d03d05a79eedac7bbc405669 Mon Sep 17 00:00:00 2001 From: Leon Luithlen Date: Tue, 12 Nov 2024 16:28:22 +0100 Subject: [PATCH] Try old version of classify_documents --- cognee/tasks/documents/classify_documents.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/cognee/tasks/documents/classify_documents.py b/cognee/tasks/documents/classify_documents.py index 472249582..9ab0726ab 100644 --- a/cognee/tasks/documents/classify_documents.py +++ b/cognee/tasks/documents/classify_documents.py @@ -10,8 +10,10 @@ EXTENSION_TO_DOCUMENT_CLASS = { def classify_documents(data_documents: list[Data]) -> list[Document]: documents = [ - EXTENSION_TO_DOCUMENT_CLASS[data_item.extension](id = data_item.id, title=f"{data_item.name}.{data_item.extension}", raw_data_location=data_item.raw_data_location) + PdfDocument(id = data_item.id, name=f"{data_item.name}.{data_item.extension}", raw_data_location=data_item.raw_data_location) if data_item.extension == "pdf" else + AudioDocument(id = data_item.id, name=f"{data_item.name}.{data_item.extension}", raw_data_location=data_item.raw_data_location) if data_item.extension == "audio" else + ImageDocument(id = data_item.id, name=f"{data_item.name}.{data_item.extension}", raw_data_location=data_item.raw_data_location) if data_item.extension == "image" else + TextDocument(id = data_item.id, name=f"{data_item.name}.{data_item.extension}", raw_data_location=data_item.raw_data_location) for data_item in data_documents ] - return documents