from flask import Flask, jsonify, request
from neo4j import GraphDatabase
from pymilvus import connections, Collection

app = Flask(__name__)

# Neo4j Connection
neo4j_driver = GraphDatabase.driver("bolt://neo4j:7687", auth=("neo4j", "testpassword"))

# Milvus Connection
connections.connect("default", host="milvus", port="19530")

@app.route('/')
def home():
    return jsonify({'message': 'Hello, World!'})

@app.route('/neo4j_test')
def neo4j_test():
    with neo4j_driver.session() as session:
        result = session.run("MATCH (n) RETURN count(n) AS count")
        count = result.single()["count"]
    return jsonify({'neo4j_node_count': count})

@app.route('/milvus_test')
def milvus_test():
    collections = Collection.list()
    return jsonify({'milvus_collections': collections})

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000)