forked from Archives/langchain
You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
61 lines
1.7 KiB
Python
61 lines
1.7 KiB
Python
"""Test Graph Database Chain."""
|
|
import os
|
|
|
|
from langchain.chains.graph_qa.cypher import GraphCypherQAChain
|
|
from langchain.graphs import Neo4jGraph
|
|
from langchain.llms.openai import OpenAI
|
|
|
|
|
|
def test_connect_neo4j() -> None:
|
|
"""Test that Neo4j database is correctly instantiated and connected."""
|
|
url = os.environ.get("NEO4J_URL")
|
|
username = os.environ.get("NEO4J_USERNAME")
|
|
password = os.environ.get("NEO4J_PASSWORD")
|
|
assert url is not None
|
|
assert username is not None
|
|
assert password is not None
|
|
|
|
graph = Neo4jGraph(
|
|
url=url,
|
|
username=username,
|
|
password=password,
|
|
)
|
|
|
|
output = graph.query(
|
|
"""
|
|
RETURN "test" AS output
|
|
"""
|
|
)
|
|
expected_output = [{"output": "test"}]
|
|
assert output == expected_output
|
|
|
|
|
|
def test_cypher_generating_run() -> None:
|
|
"""Test that Cypher statement is correctly generated and executed."""
|
|
url = os.environ.get("NEO4J_URL")
|
|
username = os.environ.get("NEO4J_USERNAME")
|
|
password = os.environ.get("NEO4J_PASSWORD")
|
|
assert url is not None
|
|
assert username is not None
|
|
assert password is not None
|
|
|
|
graph = Neo4jGraph(
|
|
url=url,
|
|
username=username,
|
|
password=password,
|
|
)
|
|
# Delete all nodes in the graph
|
|
graph.query("MATCH (n) DETACH DELETE n")
|
|
# Create two nodes and a relationship
|
|
graph.query(
|
|
"CREATE (a:Actor {name:'Bruce Willis'})"
|
|
"-[:ACTED_IN]->(:Movie {title: 'Pulp Fiction'})"
|
|
)
|
|
# Refresh schema information
|
|
graph.refresh_schema()
|
|
|
|
chain = GraphCypherQAChain.from_llm(OpenAI(temperature=0), graph=graph)
|
|
output = chain.run("Who played in Pulp Fiction?")
|
|
expected_output = " Bruce Willis played in Pulp Fiction."
|
|
assert output == expected_output
|