Spaces:
Running
Running
| import streamlit as st | |
| import pandas as pd | |
| import logging | |
| import time | |
| import json # For displaying dicts/lists nicely | |
| # Import core components from the refactored library | |
| from kig_core.config import settings # Loads config on import | |
| from kig_core.schemas import PlannerState, KeyIssue, GraphConfig | |
| from kig_core.planner import build_graph | |
| from kig_core.utils import key_issues_to_dataframe, dataframe_to_excel_bytes | |
| from kig_core.graph_client import neo4j_client # Import the initialized client instance | |
| from langchain_core.messages import HumanMessage, AIMessage, SystemMessage | |
| # Configure logging for Streamlit app | |
| logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s') | |
| logger = logging.getLogger(__name__) | |
| # --- Streamlit Page Configuration --- | |
| st.set_page_config(page_title="Key Issue Generator (KIG)", layout="wide") | |
| st.title(" KIG - Key Issue Generator ") | |
| st.write("Generate structured Key Issues from knowledge graph context.") | |
| # --- Sidebar --- | |
| with st.sidebar: | |
| st.header(" Status & Info ") | |
| # Check Neo4j Connectivity on startup | |
| neo4j_status = st.empty() | |
| try: | |
| # Accessing the client instance will trigger verification if not already done | |
| neo4j_client._get_driver() # Ensure connection is attempted | |
| neo4j_status.success("Neo4j Connection Verified") | |
| can_run = True | |
| except ConnectionError as e: | |
| neo4j_status.error(f"Neo4j Error: {e}") | |
| can_run = False | |
| except Exception as e: | |
| neo4j_status.error(f"Neo4j Init Error: {e}") | |
| can_run = False | |
| st.header("Configuration") | |
| # Display some key settings (be careful with secrets) | |
| st.text(f"Main LLM: {settings.main_llm_model}") | |
| st.text(f"Neo4j URI: {settings.neo4j_uri}") | |
| st.text(f"Plan Method: {settings.plan_method}") | |
| st.text(f"Max Docs: {settings.max_docs}") | |
| st.header("About") | |
| st.info(""" | |
| This app uses LLMs and a Neo4j graph to: | |
| 1. Plan an approach based on your query. | |
| 2. Execute the plan, retrieving & processing graph data. | |
| 3. Generate structured Key Issues. | |
| 4. Output results to an Excel file. | |
| """) | |
| # --- Main Application Logic --- | |
| st.header("Enter Your Query") | |
| user_query = st.text_area( | |
| "Describe the technical requirement or area you want to explore for Key Issues:", | |
| "What are the main challenges and potential key issues in deploying edge computing for real-time AI-driven traffic management systems in smart cities?", | |
| height=150 | |
| ) | |
| # Session state to store results across reruns if needed | |
| if 'key_issues_result' not in st.session_state: | |
| st.session_state.key_issues_result = None | |
| if 'log_messages' not in st.session_state: | |
| st.session_state.log_messages = [] | |
| # Placeholder for status updates | |
| status_placeholder = st.empty() | |
| results_placeholder = st.container() | |
| log_placeholder = st.expander("Show Execution Log") | |
| if st.button("Generate Key Issues", type="primary", disabled=not can_run): | |
| if not user_query: | |
| st.error("Please enter a query.") | |
| else: | |
| st.session_state.key_issues_result = None # Clear previous results | |
| st.session_state.log_messages = ["Starting Key Issue generation..."] | |
| with st.spinner("Processing... Building graph and executing workflow..."): | |
| start_time = time.time() | |
| try: | |
| # Build the graph | |
| status_placeholder.info("Building workflow graph...") | |
| app_graph = build_graph() | |
| st.session_state.log_messages.append("Workflow graph built.") | |
| # Define the initial state | |
| initial_state: PlannerState = { | |
| "user_query": user_query, | |
| "messages": [HumanMessage(content=user_query)], | |
| "plan": [], | |
| "current_plan_step_index": -1, # Will be set by start_planning | |
| "step_outputs": {}, | |
| "key_issues": [], | |
| "error": None | |
| } | |
| # Configuration for the graph run (e.g., thread_id for memory) | |
| # Using user query hash as a simple thread identifier for memory (if used) | |
| import hashlib | |
| thread_id = hashlib.sha256(user_query.encode()).hexdigest()[:8] | |
| config: GraphConfig = {"configurable": {"thread_id": thread_id}} | |
| status_placeholder.info("Executing workflow... (This may take a while)") | |
| st.session_state.log_messages.append("Invoking graph stream...") | |
| final_state = None | |
| # Stream events for logging/updates | |
| for i, step_state in enumerate(app_graph.stream(initial_state, config=config)): | |
| # step_state is a dictionary where keys are node names | |
| node_name = list(step_state.keys())[0] | |
| node_output = step_state[node_name] | |
| log_msg = f"Step {i+1}: Node '{node_name}' executed." | |
| st.session_state.log_messages.append(log_msg) | |
| # logger.info(log_msg) # Log to console as well | |
| # logger.debug(f"Node output: {node_output}") | |
| # You could update the status placeholder more dynamically here | |
| # status_placeholder.info(f"Executing: {node_name}...") | |
| final_state = node_output # Keep track of the latest state | |
| end_time = time.time() | |
| st.session_state.log_messages.append(f"Workflow finished in {end_time - start_time:.2f} seconds.") | |
| status_placeholder.success(f"Processing Complete! ({end_time - start_time:.2f}s)") | |
| # --- Process Final Results --- | |
| if final_state and not final_state.get("error"): | |
| generated_issues = final_state.get("key_issues", []) | |
| st.session_state.key_issues_result = generated_issues | |
| st.session_state.log_messages.append(f"Successfully extracted {len(generated_issues)} key issues.") | |
| elif final_state and final_state.get("error"): | |
| error_msg = final_state.get("error", "Unknown error") | |
| st.session_state.log_messages.append(f"Workflow failed: {error_msg}") | |
| status_placeholder.error(f"Workflow failed: {error_msg}") | |
| else: | |
| st.session_state.log_messages.append("Workflow finished, but no final state or key issues found.") | |
| status_placeholder.warning("Workflow finished, but no key issues were generated.") | |
| except Exception as e: | |
| end_time = time.time() | |
| logger.error(f"An error occurred during graph execution: {e}", exc_info=True) | |
| status_placeholder.error(f"An unexpected error occurred: {e}") | |
| st.session_state.log_messages.append(f"FATAL ERROR: {e}") | |
| # --- Display Results --- | |
| if st.session_state.key_issues_result: | |
| issues = st.session_state.key_issues_result | |
| results_placeholder.subheader(f"Generated Key Issues ({len(issues)})") | |
| df = key_issues_to_dataframe(issues) | |
| if not df.empty: | |
| # Display as DataFrame | |
| results_placeholder.dataframe(df, use_container_width=True) | |
| # Provide download button | |
| excel_bytes = dataframe_to_excel_bytes(df) | |
| results_placeholder.download_button( | |
| label="📥 Download Key Issues as Excel", | |
| data=excel_bytes, | |
| file_name="key_issues_output.xlsx", | |
| mime="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet" | |
| ) | |
| else: | |
| results_placeholder.info("No key issues were generated or parsed correctly.") | |
| # Display logs | |
| with log_placeholder: | |
| st.code("\n".join(st.session_state.log_messages), language="text") |