feat: Add support for multiple audio and image formats

Added support for multiple audio and image formats with example

Feature COG-507

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

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())