app.py

import os
import streamlit as st
from langchain_community.graphs import Neo4jGraph
import pandas as pd
import json
import time

from ki_gen.planner import build_planner_graph
from ki_gen.utils import init_app, memory
from ki_gen.prompts import get_initial_prompt

from neo4j import GraphDatabase

# Set page config
st.set_page_config(page_title="Key Issue Generator", layout="wide")

# Neo4j Database Configuration
NEO4J_URI = "neo4j+s://4985272f.databases.neo4j.io"
NEO4J_USERNAME = "neo4j"
NEO4J_PASSWORD = os.getenv("neo4j_password")

# API Keys for LLM services
OPENAI_API_KEY = os.getenv("openai_api_key")
GROQ_API_KEY = os.getenv("groq_api_key")
LANGSMITH_API_KEY = os.getenv("langsmith_api_key")

def verify_neo4j_connectivity():
    """Verify connection to Neo4j database"""
    try:
        with GraphDatabase.driver(NEO4J_URI, auth=(NEO4J_USERNAME, NEO4J_PASSWORD)) as driver:
            return driver.verify_connectivity()
    except Exception as e:
        return f"Error: {str(e)}"

def load_config():
    """Load configuration with custom parameters"""
    # Custom configuration based on provided parameters
    custom_config = {
        "main_llm": "deepseek-r1-distill-llama-70b",
        "plan_method": "generation",
        "use_detailed_query": False,
        "cypher_gen_method": "guided",
        "validate_cypher": False,
        "summarize_model": "deepseek-r1-distill-llama-70b",
        "eval_method": "binary",
        "eval_threshold": 0.7,
        "max_docs": 15,
        "compression_method": "llm_lingua",
        "compress_rate": 0.33,
        "force_tokens": ["."],  # Converting to list format as expected by the application
        "eval_model": "deepseek-r1-distill-llama-70b",
        "thread_id": "3"
    }
    
    # Add Neo4j graph object to config
    try:
        neo_graph = Neo4jGraph(
            url=NEO4J_URI, 
            username=NEO4J_USERNAME, 
            password=NEO4J_PASSWORD
        )
        custom_config["graph"] = neo_graph
    except Exception as e:
        st.error(f"Error connecting to Neo4j: {e}")
        return None
        
    return {"configurable": custom_config}

def generate_key_issues(user_query):
    """Main function to generate key issues from Neo4j data"""
    # Initialize application with API keys
    init_app(
        openai_key=OPENAI_API_KEY,
        groq_key=GROQ_API_KEY,
        langsmith_key=LANGSMITH_API_KEY
    )
    
    # Load configuration with custom parameters
    config = load_config()
    if not config:
        return None
    
    # Create status containers
    plan_status = st.empty()
    plan_display = st.empty()
    retrieval_status = st.empty()
    processing_status = st.empty()
    
    # Build planner graph
    plan_status.info("Building planner graph...")
    graph = build_planner_graph(memory, config["configurable"])
    
    # Execute initial prompt generation
    plan_status.info(f"Generating plan for query: {user_query}")
    
    messages_content = []
    for event in graph.stream(get_initial_prompt(config, user_query), config, stream_mode="values"):
        if "messages" in event:
            event["messages"][-1].pretty_print()
            messages_content.append(event["messages"][-1].content)
    
    # Get the state with the generated plan
    state = graph.get_state(config)
    steps = [i for i in range(1, len(state.values['store_plan'])+1)]
    plan_df = pd.DataFrame({'Plan steps': steps, 'Description': state.values['store_plan']})
    
    # Display the plan
    plan_status.success("Plan generation complete!")
    plan_display.dataframe(plan_df, use_container_width=True)
    
    # Continue with plan execution for document retrieval
    retrieval_status.info("Retrieving documents...")
    for event in graph.stream(None, config, stream_mode="values"):
        if "messages" in event:
            event["messages"][-1].pretty_print()
            messages_content.append(event["messages"][-1].content)
    
    # Get updated state after document retrieval
    snapshot = graph.get_state(config)
    doc_count = len(snapshot.values.get('valid_docs', []))
    retrieval_status.success(f"Retrieved {doc_count} documents")
    
    # Proceed to document processing
    processing_status.info("Processing documents...")
    process_steps = ["summarize"]  # Using summarize as default processing step
    
    # Update state to indicate human validation is complete and specify processing steps
    graph.update_state(config, {'human_validated': True, 'process_steps': process_steps}, as_node="human_validation")
    
    # Continue execution with document processing
    for event in graph.stream(None, config, stream_mode="values"):
        if "messages" in event:
            event["messages"][-1].pretty_print()
            messages_content.append(event["messages"][-1].content)
    
    # Get final state after processing
    final_snapshot = graph.get_state(config)
    processing_status.success("Document processing complete!")
    
    if "messages" in final_snapshot.values:
        final_result = final_snapshot.values["messages"][-1].content
        return final_result, final_snapshot.values.get('valid_docs', [])
    
    return None, []

# App header
st.title("Key Issue Generator")
st.write("Generate key issues from a Neo4j knowledge graph using advanced language models.")

# Check database connectivity
connectivity_status = verify_neo4j_connectivity()
st.sidebar.header("Database Status")
if "Error" not in str(connectivity_status):
    st.sidebar.success("Connected to Neo4j database")
else:
    st.sidebar.error(f"Database connection issue: {connectivity_status}")

# User input section
st.header("Enter Your Query")
user_query = st.text_area("What would you like to explore?", 
                         "What are the main challenges in AI adoption for healthcare systems?", 
                         height=100)

# Process button
if st.button("Generate Key Issues", type="primary"):
    if not OPENAI_API_KEY or not GROQ_API_KEY or not LANGSMITH_API_KEY or not NEO4J_PASSWORD:
        st.error("Required API keys or database credentials are missing. Please check your environment variables.")
    else:
        with st.spinner("Processing your query..."):
            start_time = time.time()
            final_result, valid_docs = generate_key_issues(user_query)
            end_time = time.time()
            
            if final_result:
                # Display execution time
                st.sidebar.info(f"Total execution time: {round(end_time - start_time, 2)} seconds")
                
                # Display final result
                st.header("Generated Key Issues")
                st.markdown(final_result)
                
                # Option to download results
                st.download_button(
                    label="Download Results",
                    data=final_result,
                    file_name="key_issues_results.txt",
                    mime="text/plain"
                )
                
                # Display retrieved documents in expandable section
                if valid_docs:
                    with st.expander("View Retrieved Documents"):
                        for i, doc in enumerate(valid_docs):
                            st.markdown(f"### Document {i+1}")
                            for key in doc:
                                st.markdown(f"**{key}**: {doc[key]}")
                            st.divider()
            else:
                st.error("An error occurred during processing. Please check the logs for details.")

# Help information in sidebar
with st.sidebar:
    st.header("About")
    st.info("""
    This application uses advanced language models to analyze a Neo4j knowledge graph and generate key issues 
    based on your query. The process involves:
    
    1. Creating a plan based on your query
    2. Retrieving relevant documents from the database
    3. Processing and summarizing the information
    4. Generating a comprehensive response
    """)