Python code for extracting entities from audio with AssemblyAI, then uploading that data into a Neo4j Graph database

app.py

import assemblyai as aai
from neo4j import GraphDatabase

aai.settings.api_key = "<ASSEMBLYAI_API_KEY>"

audio_file = "https://assembly.ai/wildfires.mp3"
# audio_file = "./local_file.mp3" #syntax for using a local file

config = aai.TranscriptionConfig(entity_detection=True)
transcript = aai.Transcriber().transcribe(audio_file, config)
print(f"Transcript ID:", transcript.id)

def upload_entities_to_neo4j(source, entities):
    # Replace with your Neo4j connection details
    NEO4J_URI = "bolt://localhost:7687"
    NEO4J_USER = "neo4j"
    NEO4J_PASSWORD = "<your_password>"

    # Create a list of all entities to process
    entity_data = [
        {
            "text": entity.text,
            "labels": [entity.entity_type, "_entity"],
            "start": entity.start,
            "end": entity.end
        }
        for entity in entities
    ]

    with GraphDatabase.driver(NEO4J_URI, auth=(NEO4J_USER, NEO4J_PASSWORD)) as driver:

        # Can optionally make sure each Entity node is unique by name
        # drive.execute_query("CREATE CONSTRAINT IF NOT EXISTS ON (e:_entity) ASSERT e.name IS UNIQUE")

        driver.execute_query(
            """
            MERGE (d:_source {name: $source})
            WITH d, $entities AS entities
            UNWIND entities AS entity
            MERGE (e:$(entity.labels) {name: entity.text})
            MERGE (d)-[:MENTIONS]->(e)
            MERGE (s:_timestamp {name: entity.start})
            MERGE (e)-[:STARTS_AT]->(s)
            MERGE (e)-[:ENDS_AT]->(e2:_timestamp {name: entity.end})
            """,
            source=source,
            entities=entity_data
        )


upload_entities_to_neo4j(audio_file, transcript.entities)