Create app.py

app.py

@@ -0,0 +1,204 @@
+import os
+import gradio as gr
+from dotenv import load_dotenv
+import traceback # For detailed error logging
+import torch # Required for Hugging Face transformers
+
+# --- LangChain and Hugging Face Transformers Imports ---
+from langchain_neo4j import Neo4jGraph
+# from langchain_openai import ChatOpenAI # We will replace this
+from langchain_community.llms import HuggingFacePipeline # For using HuggingFace models
+from langchain_community.chains.graph_qa.cypher import GraphCypherQAChain
+from langchain_core.prompts import PromptTemplate
+
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+
+# --- Environment Variable Loading ---
+load_dotenv()
+print("Environment variables loaded:")
+print(f"NEO4J_URI: {'Set' if os.getenv('NEO4J_URI') else 'Not Set'}")
+print(f"NEO4J_USER: {'Set' if os.getenv('NEO4J_USER') else 'Not Set'}")
+print(f"NEO4J_PASSWORD: {'Set' if os.getenv('NEO4J_PASSWORD') else 'Not Set'}")
+# OPENAI_API_KEY is no longer the primary concern if using local/HF models
+# print(f"OPENAI_API_KEY: {'Set' if os.getenv('OPENAI_API_KEY') else 'Not Set'}")
+print(f"HUGGINGFACE_HUB_TOKEN: {'Set' if os.getenv('HUGGINGFACE_HUB_TOKEN') else 'Not Set (may be needed for certain models)'}")
+
+
+# --- Global LangChain chain variable ---
+chain = None
+graph_connection_error = None # To store graph connection error
+llm_initialization_error = None # To store LLM setup error
+
+# --- Neo4j, Hugging Face LLM, and LangChain Setup ---
+try:
+    print("Attempting to connect to Neo4j...")
+    graph = Neo4jGraph(
+        url=os.getenv("NEO4J_URI"),
+        username=os.getenv("NEO4J_USER"),
+        password=os.getenv("NEO4J_PASSWORD"),
+    )
+    print("Successfully connected to Neo4j.")
+
+    # --- Hugging Face LLM Setup ---
+    print("Initializing Hugging Face LLM...")
+    # IMPORTANT: Replace "gpt2" with your desired Hugging Face model.
+    # For larger models like Llama-2, ensure you have enough resources (VRAM/RAM)
+    # and handle authentication if it's a gated model (e.g., using huggingface-cli login
+    # or by passing use_auth_token=os.getenv("HUGGINGFACE_HUB_TOKEN") if supported and necessary).
+    model_id = "gpt2" # REPLACE THIS with your chosen model, e.g., "NousResearch/Llama-2-7b-chat-hf"
+    # model_id = "meta-llama/Llama-2-7b-chat-hf" # Example from the prompt, requires auth and resources
+
+    try:
+        tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True) # trust_remote_code might be needed for some models
+        # For large models, device_map='auto' and torch_dtype are crucial.
+        # For smaller models like gpt2, they might not be strictly necessary or could be simplified.
+        hf_model = AutoModelForCausalLM.from_pretrained(
+            model_id,
+            trust_remote_code=True,
+            device_map='auto', # Automatically distributes model layers across available devices (CPU/GPU)
+            torch_dtype=torch.float16, # Use float16 for memory efficiency if GPU supports it
+            # use_auth_token=os.getenv("HUGGINGFACE_HUB_TOKEN") # If your model requires a token
+        )
+        hf_model.eval() # Set the model to evaluation mode
+
+        # Create a text-generation pipeline
+        # Adjust max_new_tokens, do_sample, top_k as needed for your model and task
+        pipe = pipeline(
+            "text-generation",
+            model=hf_model,
+            tokenizer=tokenizer,
+            # torch_dtype=torch.bfloat16, # Alternative dtype
+            # device_map="auto", # Already set in model loading
+            max_new_tokens=512,  # Max tokens for the generated Cypher query + answer synthesis
+            do_sample=True,
+            top_k=30,
+            num_return_sequences=1,
+            eos_token_id=tokenizer.eos_token_id,
+            pad_token_id=tokenizer.eos_token_id # Often good to set for open-ended generation
+        )
+
+        # Wrap the pipeline in LangChain's HuggingFacePipeline
+        llm = HuggingFacePipeline(
+            pipeline=pipe,
+            # model_kwargs can be used to pass additional arguments to the pipeline's __call__ method
+            # or to the model's generate method.
+            model_kwargs={'temperature': 0.1, 'max_length': 2000} # max_length here includes prompt
+        )
+        print(f"Hugging Face LLM ({model_id}) initialized successfully.")
+
+    except Exception as e_llm:
+        llm_initialization_error_message = f"Error initializing Hugging Face LLM ({model_id}): {str(e_llm)}\n"
+        llm_initialization_error_message += "Full Traceback:\n" + traceback.format_exc()
+        print(llm_initialization_error_message)
+        llm_initialization_error = llm_initialization_error_message
+        llm = None
+
+
+    if llm: # Proceed only if LLM initialized successfully
+        # --- GraphCypherQAChain Setup ---
+        print("Initializing GraphCypherQAChain...")
+        CYPHER_GENERATION_TEMPLATE = """You are an expert Neo4j Cypher translator.
+        Task: Convert the natural language question into a Cypher query that can retrieve relevant information from a Neo4j graph.
+        Instructions:
+        1. Use only the provided schema details. Do not use any other node labels or relationship types.
+        2. Understand the question and identify the key entities and relationships.
+        3. Construct a Cypher query that accurately reflects the question's intent.
+        4. Output ONLY the Cypher query. No explanations, no introductory text, no markdown. Just the query.
+
+        Schema:
+        {schema}
+
+        Question: {question}
+        Cypher Query:"""
+        cypher_prompt = PromptTemplate.from_template(CYPHER_GENERATION_TEMPLATE)
+
+        # For the QA part, the default prompt is often okay, but you might want to customize it too.
+        # Here's an example if you choose to:
+        # QA_TEMPLATE = """You are an assistant that answers questions based on query results from a graph database.
+        # Use the provided query result to answer the question.
+        # If the result is empty or does not contain the answer, say so.
+        # Do not make up information.
+        # Question: {question}
+        # Cypher Query Result: {context}
+        # Answer:"""
+        # qa_prompt = PromptTemplate.from_template(QA_TEMPLATE)
+
+        chain = GraphCypherQAChain.from_llm(
+            llm=llm,
+            graph=graph,
+            verbose=True,
+            return_intermediate_steps=True,
+            cypher_prompt=cypher_prompt,
+            # qa_prompt=qa_prompt # Uncomment if you want to use a custom QA prompt
+        )
+        print("LangChain integration with GraphCypherQAChain initialized successfully.")
+    else:
+        # This case is now handled by the llm_initialization_error check in process_query
+        pass
+
+
+except Exception as e_graph:
+    graph_connection_error_message = f"Error setting up Neo4j connection: {str(e_graph)}\n"
+    graph_connection_error_message += "Full Traceback:\n" + traceback.format_exc()
+    print(graph_connection_error_message)
+    graph_connection_error = graph_connection_error_message
+    chain = None
+
+# --- Gradio Interface Function ---
+def process_query(message: str, history: list):
+    if graph_connection_error:
+        return f"Application Initialization Error (Neo4j): {graph_connection_error}"
+    if llm_initialization_error:
+        return f"Application Initialization Error (LLM): {llm_initialization_error}"
+    if not chain:
+        return "Error: LangChain QA Chain is not available. Please check server logs for initialization issues."
+
+    print(f"Processing message: {message}")
+    try:
+        result = chain.invoke({"query": message})
+        print(f"Chain result: {result}")
+
+        answer = result.get("result", "No answer found or an error occurred in processing.")
+        intermediate_steps = result.get("intermediate_steps", [])
+        generated_cypher = "Could not extract Cypher query from intermediate steps."
+
+        if intermediate_steps and isinstance(intermediate_steps, list) and len(intermediate_steps) > 0:
+            if isinstance(intermediate_steps[0], dict) and "query" in intermediate_steps[0]:
+                generated_cypher = intermediate_steps[0]["query"]
+            # Sometimes the Cypher query might be in a different structure or a later step
+            # depending on the chain's verbosity and internal structure.
+            # You might need to inspect intermediate_steps more closely if the above doesn't work.
+
+        return f"📝 Generated Cypher:\n```cypher\n{generated_cypher}\n```\n\n💬 Answer:\n{answer}"
+
+    except Exception as e:
+        error_message = f"Error processing query: {str(e)}"
+        print(error_message)
+        print(traceback.format_exc())
+        # Specific error check for Hugging Face model issues (e.g. out of memory)
+        if "CUDA out of memory" in str(e):
+            return "LLM Error: CUDA out of memory. The model may be too large for your GPU. Try a smaller model or reduce batch size if applicable."
+        return error_message
+
+# --- Gradio Interface Definition ---
+print("Setting up Gradio interface...")
+demo = gr.ChatInterface(
+    fn=process_query,
+    chatbot=gr.Chatbot(height=600, type="messages"),
+    title="Neo4j Graph Database Assistant (with Hugging Face LLM)",
+    description="Ask questions about your Neo4j database. Model responses depend on the chosen Hugging Face LLM.",
+    examples=[
+        "How many nodes are in the database?",
+        "What types of nodes exist?",
+        "List all relationship types.",
+    ],
+    theme=gr.themes.Soft(),
+    cache_examples=False
+)
+
+# --- Main Execution ---
+if __name__ == "__main__":
+    print("Launching Gradio interface...")
+    # To make accessible on the network (e.g., in Docker):
+    # demo.launch(server_name="0.0.0.0")
+    demo.launch()