From 78e952dcfb67541f59d35344b4ec5ad72c423974 Mon Sep 17 00:00:00 2001
From: jp170na <jan.ptak@student.tuke.sk>
Date: Tue, 29 Sep 2026 18:41:56 +0200
Subject: [PATCH] Add Neo4j infrastructure

---
 docker-compose.yml     | 34 +++++++++++++++++++
 requirements.txt       |  1 +
 scripts/graph_db.py    | 75 ++++++++++++++++++++++++++++++++++++++++++
 scripts/neo4j_smoke.py | 14 ++++++++
 4 files changed, 124 insertions(+)
 create mode 100644 scripts/graph_db.py
 create mode 100644 scripts/neo4j_smoke.py

diff --git a/docker-compose.yml b/docker-compose.yml
index 262a51e..fd7e4fe 100644
--- a/docker-compose.yml
+++ b/docker-compose.yml
@@ -15,6 +15,10 @@ services:
       CHUNK_MIN_TOKENS: "80"
       CHUNK_TOKEN_ENCODING: cl100k_base
       HF_HOME: /cache/huggingface
+      NEO4J_URI: bolt://neo4j:7687
+      NEO4J_USER: ${NEO4J_USER:-neo4j}
+      NEO4J_PASSWORD: ${NEO4J_PASSWORD}
+      NEO4J_DATABASE: ${NEO4J_DATABASE:-neo4j}
 
     volumes:
       - ./data:/app/data
@@ -23,5 +27,35 @@ services:
 
     restart: unless-stopped
 
+  neo4j:
+    image: neo4j:2026.09.0
+    container_name: zp-agent-neo4j
+
+    ports:
+      - "127.0.0.1:7474:7474"
+      - "127.0.0.1:7687:7687"
+
+    environment:
+      NEO4J_AUTH: ${NEO4J_USER:-neo4j}/${NEO4J_PASSWORD}
+
+    volumes:
+      - neo4j-data:/data
+      - neo4j-logs:/logs
+
+    healthcheck:
+      test:
+        [
+          "CMD-SHELL",
+          "u=$${NEO4J_AUTH%%/*}; p=$${NEO4J_AUTH#*/}; cypher-shell -u \"$$u\" -p \"$$p\" 'RETURN 1' >/dev/null 2>&1"
+        ]
+      interval: 10s
+      timeout: 5s
+      retries: 10
+      start_period: 20s
+
+    restart: unless-stopped
+
 volumes:
   hf-cache:
+  neo4j-data:
+  neo4j-logs:
diff --git a/requirements.txt b/requirements.txt
index 83462fc..f6c9441 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -6,3 +6,4 @@ tiktoken>=0.8,<1
 uvicorn[standard]==0.48.0
 numpy==2.2.6
 sentence-transformers==5.7.0
+neo4j==6.3.1
diff --git a/scripts/graph_db.py b/scripts/graph_db.py
new file mode 100644
index 0000000..18c8e0f
--- /dev/null
+++ b/scripts/graph_db.py
@@ -0,0 +1,75 @@
+import os
+from dataclasses import dataclass
+
+from neo4j import GraphDatabase
+
+
+@dataclass(frozen=True)
+class Neo4jSettings:
+    uri: str
+    user: str
+    password: str
+    database: str
+
+
+def get_neo4j_settings() -> Neo4jSettings:
+    password = os.getenv("NEO4J_PASSWORD")
+
+    if not password:
+        raise RuntimeError(
+            "NEO4J_PASSWORD is not configured."
+        )
+
+    return Neo4jSettings(
+        uri=os.getenv(
+            "NEO4J_URI",
+            "bolt://localhost:7687",
+        ),
+        user=os.getenv(
+            "NEO4J_USER",
+            "neo4j",
+        ),
+        password=password,
+        database=os.getenv(
+            "NEO4J_DATABASE",
+            "neo4j",
+        ),
+    )
+
+
+def create_neo4j_driver(
+    settings: Neo4jSettings | None = None,
+):
+    resolved = settings or get_neo4j_settings()
+
+    return GraphDatabase.driver(
+        resolved.uri,
+        auth=(
+            resolved.user,
+            resolved.password,
+        ),
+    )
+
+
+def verify_neo4j_connection() -> Neo4jSettings:
+    settings = get_neo4j_settings()
+
+    with create_neo4j_driver(settings) as driver:
+        driver.verify_connectivity()
+
+        with driver.session(
+            database=settings.database
+        ) as session:
+            record = session.run(
+                "RETURN 1 AS value"
+            ).single()
+
+            if (
+                record is None
+                or record["value"] != 1
+            ):
+                raise RuntimeError(
+                    "Neo4j connectivity check failed."
+                )
+
+    return settings
diff --git a/scripts/neo4j_smoke.py b/scripts/neo4j_smoke.py
new file mode 100644
index 0000000..41ae918
--- /dev/null
+++ b/scripts/neo4j_smoke.py
@@ -0,0 +1,14 @@
+from scripts.graph_db import verify_neo4j_connection
+
+
+def main() -> None:
+    settings = verify_neo4j_connection()
+
+    print("Neo4j connection: OK")
+    print(f"URI: {settings.uri}")
+    print(f"Database: {settings.database}")
+    print(f"User: {settings.user}")
+
+
+if __name__ == "__main__":
+    main()