diff --git a/.data/multimedia/example.png b/.data/multimedia/example.png new file mode 100644 index 000000000..4d406cafd Binary files /dev/null and b/.data/multimedia/example.png differ diff --git a/.data/multimedia/text_to_speech.mp3 b/.data/multimedia/text_to_speech.mp3 new file mode 100644 index 000000000..e84aea505 Binary files /dev/null and b/.data/multimedia/text_to_speech.mp3 differ diff --git a/cognee/tasks/documents/classify_documents.py b/cognee/tasks/documents/classify_documents.py index d881514a2..8ee87bcad 100644 --- a/cognee/tasks/documents/classify_documents.py +++ b/cognee/tasks/documents/classify_documents.py @@ -1,16 +1,51 @@ from cognee.modules.data.models import Data -from cognee.modules.data.processing.document_types import Document, PdfDocument, AudioDocument, ImageDocument, TextDocument +from cognee.modules.data.processing.document_types import ( + Document, + PdfDocument, + AudioDocument, + ImageDocument, + TextDocument, +) EXTENSION_TO_DOCUMENT_CLASS = { - "pdf": PdfDocument, - "audio": AudioDocument, - "image": ImageDocument, - "txt": TextDocument + "pdf": PdfDocument, # Text documents + "txt": TextDocument, + "png": ImageDocument, # Image documents + "dwg": ImageDocument, + "xcf": ImageDocument, + "jpg": ImageDocument, + "jpx": ImageDocument, + "apng": ImageDocument, + "gif": ImageDocument, + "webp": ImageDocument, + "cr2": ImageDocument, + "tif": ImageDocument, + "bmp": ImageDocument, + "jxr": ImageDocument, + "psd": ImageDocument, + "ico": ImageDocument, + "heic": ImageDocument, + "avif": ImageDocument, + "aac": AudioDocument, # Audio documents + "mid": AudioDocument, + "mp3": AudioDocument, + "m4a": AudioDocument, + "ogg": AudioDocument, + "flac": AudioDocument, + "wav": AudioDocument, + "amr": AudioDocument, + "aiff": AudioDocument, } + 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, name=data_item.name) + 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, + name=data_item.name, + ) for data_item in data_documents ] return documents diff --git a/examples/python/multimedia_example.py b/examples/python/multimedia_example.py new file mode 100644 index 000000000..6c8bc5995 --- /dev/null +++ b/examples/python/multimedia_example.py @@ -0,0 +1,48 @@ +import os +import asyncio +import pathlib + +import cognee +from cognee.api.v1.search import SearchType + +# Prerequisites: +# 1. Copy `.env.template` and rename it to `.env`. +# 2. Add your OpenAI API key to the `.env` file in the `LLM_API_KEY` field: +# LLM_API_KEY = "your_key_here" + + +async def main(): + # Create a clean slate for cognee -- reset data and system state + await cognee.prune.prune_data() + await cognee.prune.prune_system(metadata=True) + + # cognee knowledge graph will be created based on the text + # and description of these files + mp3_file_path = os.path.join( + pathlib.Path(__file__).parent.parent.parent, + ".data/multimedia/text_to_speech.mp3", + ) + png_file_path = os.path.join( + pathlib.Path(__file__).parent.parent.parent, + ".data/multimedia/example.png", + ) + + # Add the files, and make it available for cognify + await cognee.add([mp3_file_path, png_file_path]) + + # Use LLMs and cognee to create knowledge graph + await cognee.cognify() + + # Query cognee for summaries of the data in the multimedia files + search_results = await cognee.search( + SearchType.SUMMARIES, + query_text="What is in the multimedia files?", + ) + + # Display search results + for result_text in search_results: + print(result_text) + + +if __name__ == "__main__": + asyncio.run(main())