import os import json import gradio as gr import uuid from typing import List, Any, Tuple, Dict # Added Dict from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder from langchain_core.messages import HumanMessage, AIMessage from langchain_google_genai import ChatGoogleGenerativeAI from langchain_core.output_parsers import StrOutputParser # Define custom CSS for Gradio UI (can be empty or contain styles) custom_css = """ .gradio-container { max-width: none !important; } """ # --- Secrets and API Key Configuration --- os.environ["GOOGLE_API_KEY"] = os.environ.get("GOOGLE_API_KEY") if not os.environ["GOOGLE_API_KEY"]: print("FATAL: GOOGLE_API_KEY is not set. Please add it in the Hugging Face secrets.") exit() # --- 1) LLM and Prompt Configuration --- llm = ChatGoogleGenerativeAI(model="gemini-2.5-flash", temperature=0.5) JSON_ONLY_INSTRUCTION = "Your output MUST BE ONLY the raw JSON object, without any markdown formatting like ```json or any other text." # Define the comprehensive list of BRD sections and their descriptions BRD_SECTIONS = { "Project Scope": "What are the specific, measurable features for the MVP (Minimum Viable Product)? What is explicitly out of scope?", "Target Audience": "Detailed personas (e.g., \"Busy Professional\" needs more detail: age, tech-savviness, cooking skill).", "Business Goals": "What are the key business metrics for success (e.g., user acquisition targets, revenue goals, churn rate)?", "Functional Requirements": "How should AI personalization work? What are the steps in the grocery ordering and delivery process? How does inventory tracking function?", "Non-Functional Requirements": "What are the expected performance benchmarks (e.g., app load time)? What are the security and data privacy considerations?", "Data Sources & Constraints": "What specific APIs will be used for recipes or grocery data? Are there any budget or technology stack limitations?", "Timeline": "What is the desired deadline or a rough project timeline?" } # Clarifier system template now includes a placeholder for the dynamic checklist clarifier_system_template = """You are a BRD-building assistant. Your primary goal is to gather details by asking clarifying questions. When you have enough information, you will signal you are ready to report. You will be given a conversation history. Analyze it. - If you have enough information to write a great BRD, set "status" to "REPORT_READY". - If you still need more information, set "status" to "ASK" and provide a list of questions you need answers to. **Required Information Checklist:** {dynamic_checklist_content} Return ONLY valid JSON with this exact schema: {{{{ "status": "ASK" | "REPORT_READY", "questions": string[], "reason": string }}}} """ # Removed {JSON_ONLY_INSTRUCTION} from here # clarifier_prompt will be created dynamically in chat_logic report_system = "You are a meticulous BRD writer. Create a polished, structured BRD based on all information collected in the conversation." report_prompt = ChatPromptTemplate.from_messages([ ("system", report_system), MessagesPlaceholder("history"), ("human", "Excellent. Please compile all the information into a comprehensive Business Requirement Document.") ]) # NEW: Prompt for generating a summary of current understanding summarizer_system_template = """You are a helpful assistant. Based on the provided conversation history, summarize the key information gathered so far for each of the following BRD sections. If a section has no information yet, state that. Be concise and use bullet points. **BRD Sections to Summarize:** {dynamic_checklist_content} Provide a summary for each section. """ # --- 2) Chat History & Session Management --- HISTORY_DIR = "chat_histories" if not os.path.exists(HISTORY_DIR): os.makedirs(HISTORY_DIR) def get_session_filepath(session_id: str) -> str: return os.path.join(HISTORY_DIR, f"session_{session_id}.json") def save_session_state(session_id: str, history: List[Any], pending_questions: List[str], selected_sections: List[str]): if not session_id: return filepath = get_session_filepath(session_id) state = { "history": [{"type": msg.type, "content": msg.content} for msg in history], "pending_questions": pending_questions, "selected_sections": selected_sections # NEW: Save selected sections } with open(filepath, "w") as f: json.dump(state, f, indent=2) def load_session_state(session_id: str) -> Tuple[List[Any], List[str], List[str]]: # NEW: Return type includes selected_sections filepath = get_session_filepath(session_id) if not os.path.exists(filepath): return [], [], list(BRD_SECTIONS.keys()) # Default all sections if new with open(filepath, "r") as f: try: state = json.load(f) history = [ HumanMessage(content=item["content"]) if item["type"] == "human" else AIMessage(content=item["content"]) for item in state.get("history", []) ] pending_questions = state.get("pending_questions", []) selected_sections = state.get("selected_sections", list(BRD_SECTIONS.keys())) # NEW: Load selected sections, default all return history, pending_questions, selected_sections except (json.JSONDecodeError, TypeError): return [], [], list(BRD_SECTIONS.keys()) # Default all sections on error def get_saved_sessions() -> List[Tuple[str, str]]: sessions = [] for f in os.listdir(HISTORY_DIR): if f.startswith("session_") and f.endswith(".json"): session_id = f.replace("session_", "").replace(".json", "") history, _, _ = load_session_state(session_id) label = "New Chat" # Safely get the content of the first message for the label if history and isinstance(history[0], (HumanMessage, AIMessage)) and history[0].content: first_message_content = history[0].content label = first_message_content[:40] + ("..." if len(first_message_content) > 40 else "") sessions.append((label, session_id)) # Fix: Use x[1] to get the session_id from the tuple for get_session_filepath sorted_sessions = sorted(sessions, key=lambda x: os.path.getmtime(get_session_filepath(x[1])), reverse=True) return [("➕ New Chat", "NEW_CHAT_SESSION")] + sorted_sessions # Add a fixed "New Chat" option # --- 3) Professional Output Formatting --- def format_ai_output(result: dict) -> str: """ Transforms the AI's JSON output (clarifier questions) into a polished, professional message. """ reason = result.get("reason", "No reason provided.") questions = result.get("questions", []) response = ( "Thank you. To build a comprehensive document, please provide details on the following points.\n\n" f"**Reasoning:** *{reason}*\n\n" "--- \n\n" ) categorized = {} for q in questions: # Try to extract category from question, e.g., "Project Scope: What are..." parts = q.split(":", 1) if len(parts) == 2: category, question_text = parts category = category.replace("Regarding ", "").strip() if category not in categorized: categorized[category] = [] categorized[category].append(question_text.strip()) else: # If no clear category, put in 'General' if "General" not in categorized: categorized["General"] = [] categorized["General"].append(q) q_counter = 1 for category, q_list in categorized.items(): response += f"**{category}**\n" for q_text in q_list: response += f"- {q_counter}. {q_text}\n" # Add numbering q_counter += 1 response += "\n" return response.strip() # --- 4) Core Gradio Application Logic (One-by-One Questions) --- def chat_logic(user_message: str, history: List[Any], pending_questions: List[str], session_id: str, selected_sections: List[str]) -> Tuple: """Main logic for a single chat turn with a question queue.""" # If this is a "NEW_CHAT_SESSION", generate a real UUID and update the state if session_id == "NEW_CHAT_SESSION": session_id = str(uuid.uuid4()) is_first_message_in_new_session = True else: is_first_message_in_new_session = len(history) == 0 # Check if history is empty for existing session # Determine if checkboxes should be interactive (only interactive for truly new sessions before first message) checkbox_interactive_state = gr.update(interactive=False) if history else gr.update(interactive=True) # 1. Add user's latest message to the history (this history is the one saved and passed to LLM) history.append(HumanMessage(content=user_message)) # If this is the first message in a new session, save it now if is_first_message_in_new_session: save_session_state(session_id, history, pending_questions, selected_sections) updated_session_list_choices = gr.update(choices=get_saved_sessions(), value=session_id) else: updated_session_list_choices = gr.update() # Initialize content for the BRD display brd_display_content = "Thank you for your input. Processing...\n\n" # 2. Dynamically build the checklist content for prompts dynamic_checklist_content = "\n".join([f"- {section}: {BRD_SECTIONS[section]}" for section in selected_sections if section in BRD_SECTIONS]) # 3. Generate Progress Summary (for BRD Preview) current_summarizer_system = summarizer_system_template.format(dynamic_checklist_content=dynamic_checklist_content) summarizer_prompt_instance = ChatPromptTemplate.from_messages([ ("system", current_summarizer_system), MessagesPlaceholder("history") # Pass the current history (user message added) ]) summary_chain = summarizer_prompt_instance | llm | StrOutputParser() progress_summary = summary_chain.invoke({"history": history}) brd_display_content += f"**Current Understanding:**\n{progress_summary}\n\n" # 4. Determine next AI action (clarification or report) ai_response_for_chat = "" # This will be the message shown in the chatbot if pending_questions: # User was answering a pending question, so pop the next one ai_response_for_chat = pending_questions.pop(0) else: # No pending questions, call clarifier to get new questions or report current_clarifier_system = clarifier_system_template.format(dynamic_checklist_content=dynamic_checklist_content) + "\n" + JSON_ONLY_INSTRUCTION clarifier_prompt_instance = ChatPromptTemplate.from_messages([("system", current_clarifier_system), MessagesPlaceholder("history")]) clarifier_chain = clarifier_prompt_instance | llm | StrOutputParser() raw_response = clarifier_chain.invoke({"history": history}) try: result = json.loads(raw_response) status = result.get("status") if status == "ASK" and result.get("questions"): pending_questions.extend(result["questions"]) # Add all new questions to the queue reason = result.get("reason", "To gather more details, I have some questions.") if pending_questions: first_question = pending_questions.pop(0) # Pop only the first one for current turn ai_response_for_chat = f"Thank you for your input. {reason}\n\nHere is my first question:\n- {first_question}" else: ai_response_for_chat = "I need more information, but I couldn't generate a specific question. Can you please elaborate further?" elif status == "REPORT_READY": report_intro = "Excellent, I have all the information required. Generating the final report now..." final_chain = report_prompt | llm | StrOutputParser() report = final_chain.invoke({"history": history}) ai_response_for_chat = f"{report_intro}\n\n---\n\n{report}" brd_display_content += f"\n\n=== FINAL BRD ===\n\n{report}" # Add final report to BRD display else: ai_response_for_chat = "I received an unexpected status. Let's try again. Can you rephrase?" except (json.JSONDecodeError, ValueError): ai_response_for_chat = "I'm having a little trouble processing that. Could you please clarify or rephrase your last message?" # 5. Append the AI's direct response (question or report) to history for LLM context # This is the only AI message that should be part of the conversational history for the LLM. history.append(AIMessage(content=ai_response_for_chat)) # Append the *formatted* message for consistency # 6. Save state and update display save_session_state(session_id, history, pending_questions, selected_sections) # Format history for chatbot display (now simpler as history only contains relevant messages) chatbot_display = _format_history_for_chatbot(history) # Return values: user_input, history_state, pending_questions_state, chatbot, brd_display, selected_sections_state, brd_sections_checkboxes, session_id_state, session_list return ( "", # user_input history, # history_state pending_questions, # pending_questions_state chatbot_display, # chatbot brd_display_content, # brd_display (now includes summary and potentially final report) selected_sections, # selected_sections_state checkbox_interactive_state, # brd_sections_checkboxes session_id, # session_id_state updated_session_list_choices # session_list ) def _format_history_for_chatbot(history: List[Any]) -> List[Dict[str, str]]: """Converts message history to Gradio chatbot 'messages' format.""" # Since history now only contains human messages and the *single* relevant AI response per turn, # this function becomes straightforward. chatbot_display = [] for msg in history: role = "user" if msg.type == "human" else "assistant" chatbot_display.append({"role": role, "content": msg.content}) return chatbot_display # --- 4) UI Functions --- def start_new_chat(): """Resets the state for a new, unsaved chat.""" default_sections = list(BRD_SECTIONS.keys()) # Do NOT save here. Session will be saved on first user message. # Return values: history_state, pending_questions_state, session_id_state, chatbot, user_input, session_list, selected_sections_state, brd_sections_checkboxes, brd_display return ( [], # history_state (empty list for chatbot) [], # pending_questions_state (empty list) "NEW_CHAT_SESSION", # session_id_state [], # chatbot (empty list for display) "", # user_input (empty string) gr.update(choices=get_saved_sessions(), value="NEW_CHAT_SESSION"), # session_list default_sections, # selected_sections_state (list of strings) gr.update(value=default_sections, interactive=True), # brd_sections_checkboxes "BRD will appear here once generated." # brd_display reset ) def load_chat_session(session_id: str): """Loads a selected chat session from the sidebar or starts a new one.""" if session_id == "NEW_CHAT_SESSION": return start_new_chat() history, pending_questions, selected_sections = load_session_state(session_id) chatbot_display = _format_history_for_chatbot(history) # Determine if checkboxes should be interactive based on whether history exists checkbox_interactive_state = gr.update(interactive=False) if history else gr.update(interactive=True) # Re-generate the last summary for the BRD display if history exists current_brd_display_content = "BRD will appear here once generated." if history: dynamic_checklist_content = "\n".join([f"- {section}: {BRD_SECTIONS[section]}" for section in selected_sections if section in BRD_SECTIONS]) current_summarizer_system = summarizer_system_template.format(dynamic_checklist_content=dynamic_checklist_content) summarizer_prompt_instance = ChatPromptTemplate.from_messages([ ("system", current_summarizer_system), MessagesPlaceholder("history") ]) summary_chain = summarizer_prompt_instance | llm | StrOutputParser() current_brd_display_content = f"**Current Understanding:**\n{summary_chain.invoke({'history': history})}" # If the last message in history was a final report, append it to the BRD display if history and history[-1].type == "ai" and "Excellent, I have all the information required. Generating the final report now..." in history[-1].content: current_brd_display_content += f"\n\n---\n\n{history[-1].content}" # Append the full report text # Return values: history_state, pending_questions_state, session_id_state, chatbot, user_input, session_list, selected_sections_state, brd_sections_checkboxes, brd_display return ( history, pending_questions, session_id, chatbot_display, "", # user_input (empty string) gr.update(choices=get_saved_sessions(), value=session_id), # session_list selected_sections, gr.update(value=selected_sections, interactive=checkbox_interactive_state['interactive']), # brd_sections_checkboxes current_brd_display_content # brd_display updated ) def on_ui_load(): """Called when the UI is first loaded in the browser.""" return start_new_chat() # --- 5) Gradio UI Layout --- with gr.Blocks(theme=gr.themes.Soft(), title="BRD Assistant", css=custom_css) as app: # Added custom_css history_state = gr.State([]) pending_questions_state = gr.State([]) session_id_state = gr.State("") brd_text_state = gr.State("") selected_sections_state = gr.State(list(BRD_SECTIONS.keys())) # NEW STATE: Initialize with all sections with gr.Row(): with gr.Column(scale=1, min_width=250): gr.Markdown("### Chat History") # Changed to ### for consistency new_chat_button = gr.Button("➕ New Chat", variant="primary") session_list = gr.Radio( label="Past Conversations", choices=get_saved_sessions(), interactive=True, type="value" ) with gr.Column(scale=3): gr.Markdown("### BRD Sections") # NEW HEADER - MOVED HERE brd_sections_checkboxes = gr.CheckboxGroup( # NEW COMPONENT - MOVED HERE label="Select sections for your BRD", choices=list(BRD_SECTIONS.keys()), value=list(BRD_SECTIONS.keys()), # Default all selected interactive=True ) chatbot = gr.Chatbot( label="Conversation", # Removed bubble_full_width=False (deprecated) height=500, show_copy_button=True, type='messages' # Set chatbot type to 'messages' ) with gr.Row(): user_input = gr.Textbox( show_label=False, placeholder="Enter your project idea or answers here...", scale=5, container=False ) submit_button = gr.Button("Send", variant="primary", scale=1, min_width=150) # ✅ Add the BRD Preview here: with gr.Column(scale=3): # Adjusted column scale for better layout gr.Markdown("### 📄 Current BRD Preview (Auto-Updated)") brd_display = gr.Markdown("BRD will appear here once generated.") # --- Event Handlers --- app.load( on_ui_load, None, [history_state, pending_questions_state, session_id_state, chatbot, user_input, session_list, selected_sections_state, brd_sections_checkboxes, brd_display] # Added brd_display to outputs ) submit_triggers = [user_input.submit, submit_button.click] for trigger in submit_triggers: trigger( chat_logic, [user_input, history_state, pending_questions_state, session_id_state, selected_sections_state], [user_input, history_state, pending_questions_state, chatbot, brd_display, selected_sections_state, brd_sections_checkboxes, session_id_state, session_list] ) new_chat_button.click( start_new_chat, [], [history_state, pending_questions_state, session_id_state, chatbot, user_input, session_list, selected_sections_state, brd_sections_checkboxes, brd_display] # Added brd_display to outputs ) session_list.change( load_chat_session, [session_list], [history_state, pending_questions_state, session_id_state, chatbot, user_input, session_list, selected_sections_state, brd_sections_checkboxes, brd_display] # Added brd_display to outputs ) # NEW HANDLER: Update selected_sections_state when checkboxes change brd_sections_checkboxes.change( lambda x: x, # Simple passthrough function to update the state [brd_sections_checkboxes], [selected_sections_state] ) if __name__ == "__main__": app.launch() #debug=True, share=True) # TODO: show sections # Show outputs (aknowledgment, questions, report) # How to canvas.