""" requirements: neo4j, pydantic """ from neo4j import GraphDatabase from pydantic import BaseModel, Field class Tools: def __init__(self): self.valves = self.Valves() class Valves(BaseModel): neo4j_uri: str = Field("bolt://host.docker.internal:7687") neo4j_user: str = Field("neo4j", description="Neo4j username") neo4j_password: str = Field("", description="Neo4j password") def query_neo4j(self, __user__, query: str) -> str: """ Execute a Cypher query on a neo4j graph database and return the result :param query: the query to execute """ def classify_cypher_query(query: str) -> str: """ Classify a Cypher query as 'read-only', 'writing', or 'deleting'. Args: query (str): The Cypher query. Returns: str: 'read-only', 'writing', or 'deleting'. """ # Normalize the query by stripping extra whitespace and converting to lowercase normalized_query = query.strip().lower() # Tokens that indicate writing operations writing_keywords = { "create", "merge", "set", "replace", "on create set", "on match set", } # Tokens that indicate deleting operations deleting_keywords = {"delete", "detach delete", "remove"} # Tokens that indicate read-only operations read_only_keywords = {"match", "return", "with", "call", "unwind", "union"} # Check for writing keywords if any(keyword in normalized_query for keyword in writing_keywords): return "writing" # Check for deleting keywords if any(keyword in normalized_query for keyword in deleting_keywords): return "deleting" # If it's not writing or deleting, assume it's read-only if any(keyword in normalized_query for keyword in read_only_keywords): return "read-only" # Default to unknown if keywords are not found return "unknown" try: lines = query.split("\n") for line in lines: if ( not __user__.get("role", "x") == "admin" and not classify_cypher_query(line) == "read-only" ): raise Exception( "Access denied, you can only read from the database" ) driver = GraphDatabase.driver( self.valves.neo4j_uri, auth=(self.valves.neo4j_user, self.valves.neo4j_password), ) with driver.session() as session: result = session.run(query) return f"{query}: {str([record for record in result])}" except Exception as e: print(e) return f"{query}: Invalid Cypher query" def close(driver): """Close the Neo4j driver connection.""" driver.close()