import os
os.environ["OPENAI_API_KEY"] = "sk-*"
from langchain.prompts import PromptTemplate
from langchain_neo4j import Neo4jGraph
from langchain_neo4j import GraphCypherQAChain
from langchain_openai import ChatOpenAI
 
os.environ["NEO4J_URI"] = os.getenv("NEO4J_URI", "bolt://localhost:7687")
os.environ["NEO4J_USERNAME"] = os.getenv("NEO4J_USERNAME", "neo4j")
os.environ["NEO4J_PASSWORD"] = os.getenv("NEO4J_PASSWORD", "password")
 
graph = Neo4jGraph()
schema = graph.schema
 
sample_query = """
MATCH (a:Author)-[:WROTE]->(b:Book)
RETURN b.title AS book_title, b.publisher AS publisher, b.language AS language, b.rating AS rating, a.name AS author_name
"""
result = graph.query(sample_query)
 
llm = ChatOpenAI(model="gpt-4o-mini", temperature=0.1)  # Adjusted temperature for better responses
 
chain = GraphCypherQAChain.from_llm(
    graph=graph,
    llm=llm,
    verbose=True,
    validate_cypher=True,
    allow_dangerous_requests=True,  # Use cautiously in production
    return_intermediate_steps=True,  # Useful for debugging
)
 
def get_response(query_input):
    response = chain.invoke({
        "query": query_input
    })
    return response
 
prompt_template = """You are an expert in querying Neo4j graphs and liberating AI. Your task is to answer questions about a graph database containing data about Goodreads books and authors. The nodes in the graph represent Author and Book, and the edges between them represent the relationship WROTE.
 
For the given query input, follow these steps:
 
- use contains() function to match the query and format the response
- use the context output, cypher query and intermidiate steps to answer the query
 
 
1. Identify the nodes
  - Break down the input query to identify the relevant Author and Book entities (nodes) mentioned.
  - Look for keywords or entities that could correspond to nodes in the graph.
 
2. Match the identified nodes in the graph
  - Using the identified nodes, construct a search that matches them within the graph database.
  - Ensure that the node relationships (i.e., `WROTE`) are properly accounted for when linking authors and books.
 
3. Create and execute the Cypher query
  - Based on the extracted entities, formulate an appropriate Cypher query to retrieve the data from Neo4j.
  - Execute the query to fetch the relevant results from the graph database.
 
4. Format the response:
  - Return the result in a human-readable format (not the raw Cypher query), summarizing the data from the query.
 
Query: {query_input}
Schema: {schema}
 
Please follow the steps and return a human-readable answer based on the graph data, not the Cypher query.
 
Cypher Query:
"""
 
search_results = ""
while True:
    query_input = input("Enter the query: ")
    if query_input == "exit":
        break
    else:
        query_input += search_results
 
    template = PromptTemplate(input_variables=["query_input", "schema"],
                              template=prompt_template)
    formatted_prompt = template.format(query_input=query_input, schema=schema)
 
    response = get_response(formatted_prompt)
 
    result = response["result"]
    print("Query: ", query_input)
    print("Result: ", result)
 
# share all books published by the Penguin Books?
# what is the author name of The Complete Verse and Other ?