Refactor SdlcNode for improved prompt handling and error logging; streamline user story and design document generation processes.

src/langgraphagenticai/nodes/sdlc_node.py

@@ -6,8 +6,8 @@ import json
 from datetime import datetime
 from typing import List
 from src.langgraphagenticai.logging.logging_utils import logger, log_entry_exit
-from src.langgraphagenticai.prompt_library import prompt
-import json
+from src.langgraphagenticai.prompt_library import prompt # Assuming your prompts are here
+# import json # Duplicate import, already imported
 from tenacity import retry, stop_after_attempt, wait_exponential
 import functools
 import time
@@ -16,7 +16,7 @@ import re
 
 class SdlcNode:
     def __init__(self, model):
-        """Initialize the BlogGenerationNode with an LLM."""
+        """Initialize the SdlcNode with an LLM."""
         self.llm = model
 
     @log_entry_exit
@@ -51,21 +51,17 @@ class SdlcNode:
                 "project_scope": state.project_scope if state.project_scope is not None else "No project scope provided",
                 "project_objectives": state.project_objectives if state.project_objectives is not None else "No project objectives provided",
             }
-            # Prepare dynamic input
             requirements_input_json = json.dumps(requirements_input, indent=2)
 
-            # Render prompt string with runtime data
             prompt_string = prompt.REQUIREMENTS_PROMPT_STRING.format(requirements_input=requirements_input_json)
-            sys_prompt = prompt.REQUIREMENTS_sys_prompt
-                                
-              
+            # Assuming REQUIREMENTS_sys_prompt does not need .format() or is formatted elsewhere if needed
+            sys_prompt_content = prompt.REQUIREMENTS_sys_prompt 
             
             messages = [
-                SystemMessage(content=sys_prompt),
+                SystemMessage(content=sys_prompt_content),
                 HumanMessage(content=prompt_string)
             ]
            
-            
             response = self.llm.invoke(messages)
             state.generated_requirements = response.content if hasattr(response, 'content') else str(response)  
             return {"generated_requirements": state.generated_requirements}
@@ -77,113 +73,173 @@ class SdlcNode:
     @log_entry_exit
     def generate_user_stories(self, state: State) -> dict:
         """Generate user stories based on the requirements."""
-        
         logger.info("Generating user stories")
         if not state.generated_requirements:
             state.user_stories = "No requirements generated yet."
             logger.warning("Cannot generate user stories without requirements.")
             return {"user_stories": state.user_stories}
+
+        formatted_requirements = str(state.generated_requirements).replace('{', '{{').replace('}', '}}')
+
         feedback = state.get_last_feedback_for_stage(SDLCStages.PLANNING)
         logger.info(f"Feedback for user stories: {feedback}")
+
+        project_name_val = state.project_name or 'N/A'
+        project_name_formatted = str(project_name_val).replace('{', '{{').replace('}', '}}')
+
+
         if feedback:
-            prompt_string = prompt.USER_STORIES_FEEDBACK_PROMPT_STRING.format(
-                                generated_requirements=state.generated_requirements,
-                                # project_name=state.project_name or 'N/A',
-                                feedback=feedback
-            )
-            
-            sys_prompt = prompt.USER_STORIES_FEEDBACK_SYS_PROMPT.format(
-                                generated_requirements=state.generated_requirements,
-                                project_name=state.project_name or 'N/A',
-                                feedback=feedback
-                            )
-            
+            formatted_feedback = str(feedback).replace('{', '{{').replace('}', '}}')
+            try:
+                prompt_string = prompt.USER_STORIES_FEEDBACK_PROMPT_STRING.format(
+                                    generated_requirements=formatted_requirements,
+                                    feedback=formatted_feedback
+                )
+                sys_prompt_content = prompt.USER_STORIES_FEEDBACK_SYS_PROMPT.format(
+                                    generated_requirements=formatted_requirements,
+                                    project_name=project_name_formatted,
+                                    feedback=formatted_feedback
+                                )
+            except KeyError as e:
+                logger.error(f"KeyError during formatting USER_STORIES_FEEDBACK prompts: {e}")
+                logger.error(f"Available keys: generated_requirements, project_name, feedback")
+                state.user_stories = f"Error formatting user story prompt: {e}"
+                return {"user_stories": state.user_stories}
+        else:
+            try:
+                prompt_string = prompt.USER_STORIES_NO_FEEDBACK_PROMPT_STRING.format(
+                                    generated_requirements=formatted_requirements,
+                                    project_name=project_name_formatted
+                                )
+                sys_prompt_content = prompt.USER_STORIES_NO_FEEDBACK_SYS_PROMPT.format(
+                                    generated_requirements=formatted_requirements,
+                                    project_name=project_name_formatted
+                                )
+            except KeyError as e:
+                logger.error(f"KeyError during formatting USER_STORIES_NO_FEEDBACK prompts: {e}")
+                logger.error(f"Available keys: generated_requirements, project_name")
+                state.user_stories = f"Error formatting user story prompt: {e}"
+                return {"user_stories": state.user_stories}
+        
+        try:
             messages = [
-                SystemMessage(content=sys_prompt),
+                SystemMessage(content=sys_prompt_content),
                 HumanMessage(content=prompt_string)
             ]
             response = self.llm.invoke(messages)
             state.user_stories = response.content if hasattr(response, 'content') else str(response)
+            st.session_state["user_stories"] = state.user_stories # Update session state
+            logger.info(f"--- RAW state.user_stories after generation ---")
+            logger.info(state.user_stories)
+            logger.info(f"--- END RAW state.user_stories ---")
+
+            return {"user_stories": state.user_stories}
+        except Exception as e:
+            logger.error(f"Error generating user stories: {e}")
+            state.user_stories = f"Error generating user stories: {str(e)}"
             return {"user_stories": state.user_stories}
-        else:
-            try:
-                if state.generated_requirements:
-                    prompt_string = prompt.USER_STORIES_NO_FEEDBACK_PROMPT_STRING.format(
-                                generated_requirements=state.generated_requirements,
-                                project_name=state.project_name or 'N/A'
-                            )
-                            
-                    sys_prompt = prompt.USER_STORIES_NO_FEEDBACK_SYS_PROMPT.format(
-                                                generated_requirements=state.generated_requirements,
-                                                project_name=state.project_name or 'N/A'
-                                            )
-                    
-                    messages = [
-                        SystemMessage(content=sys_prompt),
-                        HumanMessage(content=prompt_string)
-                    ]
-                    response = self.llm.invoke(messages)
-                    state.user_stories = response.content if hasattr(response, 'content') else str(response)
-                    return {"user_stories": state.user_stories}
-                else:
-                    state.user_stories = "No requirements generated yet."
-                    return {"user_stories": state.user_stories}
-            except Exception as e:
-                logger.error(f"Error generating user stories: {e}")
-                state.user_stories = f"Error generating user stories: {str(e)}"
-                return {"user_stories": state.user_stories}
-    
 
     @log_entry_exit
-    def design_documents(self, state: State) -> dict:
+    def design_documents(self, state: State) -> dict[str, str]:
         """Generate design documents based on user stories with robust validation."""
         state.feedback_decision = None
         logger.info("Generating design documents")
-        if not state.user_stories:
-            state.design_documents = "No user stories generated yet."
+
+        if not state.user_stories or not str(state.user_stories).strip():
+            state.design_documents = "No user stories provided for design document generation."
             logger.warning("Cannot generate design documents without user stories.")
             return {"design_documents": state.design_documents}
+
+        
+        user_stories_for_prompt = str(state.user_stories).replace('{', '{{').replace('}', '}}')
+        
+        project_name_val = state.project_name or 'N/A'
+        project_name_for_prompt = str(project_name_val).replace('{', '{{').replace('}', '}}')
+
+        logger.info(f"--- User Stories for TDD Prompt (escaped) ---")
+        logger.info(user_stories_for_prompt[:1000] + "..." if len(user_stories_for_prompt) > 1000 else user_stories_for_prompt) # Log a snippet
+        logger.info(f"--- END User Stories for TDD Prompt ---")
+
+
+        # Get feedback for design stage (currently not used in these prompts but good to have)
+        # feedback = state.get_last_feedback_for_stage(SDLCStages.DESIGN)
+        # logger.info(f"Feedback for design documents: {feedback or 'None'}")
+
         try:
-            prompt_string = f"Your task is to generate design documents based on the following user stories:\n\n{state.user_stories}\n\nPlease provide detailed design documents that cover the architecture, components, and interactions required to implement these user stories."
-            sys_prompt = f" You are an expert software architect. Your task is to generate design documents based on the user stories provided. The design documents should be comprehensive and cover all aspects of the architecture, components, and interactions required to implement the user stories.\n\nProject Name: {state.project_name or 'N/A'}"
+            # Use the detailed prompts from prompt_library
+            # Ensure the placeholders in your prompts match what you provide here.
+            # The prompts DESIGN_DOCUMENTS_NO_FEEDBACK_... expect {project_name} and {user_stories}
+
+            sys_prompt_content = prompt.DESIGN_DOCUMENTS_NO_FEEDBACK_SYS_PROMPT.format(
+                user_stories=user_stories_for_prompt,
+                project_name=project_name_for_prompt
+            )
+            prompt_string_content = prompt.DESIGN_DOCUMENTS_NO_FEEDBACK_PROMPT_STRING.format(
+                user_stories=user_stories_for_prompt,
+                project_name=project_name_for_prompt # Assuming this is also needed here
+            )
+            
+            logger.debug(f"--- Formatted System Prompt for TDD (first 500 chars) ---")
+            logger.debug(sys_prompt_content[:500] + "...")
+            logger.debug(f"--- Formatted Human Prompt for TDD (first 500 chars) ---")
+            logger.debug(prompt_string_content[:500] + "...")
+
             messages = [
-                SystemMessage(content=sys_prompt),
-                HumanMessage(content=prompt_string)
+                SystemMessage(content=sys_prompt_content),
+                HumanMessage(content=prompt_string_content)
             ]
+           
             response = self.llm.invoke(messages)
             state.design_documents = response.content if hasattr(response, 'content') else str(response)
+            logger.info("Design documents generated successfully.")
+            return {"design_documents": state.design_documents}
+
+        except KeyError as ke:
+            error_msg = f"KeyError during TDD prompt formatting: {str(ke)}. This means a placeholder in your DESIGN_DOCUMENTS_... prompts was not found in the .format() call."
+            logger.error(error_msg)
+            logger.error(f"User stories (snippet): {user_stories_for_prompt[:200]}")
+            logger.error(f"Project name: {project_name_for_prompt}")
+            state.design_documents = error_msg
             return {"design_documents": state.design_documents}
         except Exception as e:
-            logger.error(f"Error generating design documents: {e}")
-            state.design_documents = f"Error generating design documents: {str(e)}"
+            error_msg = f"Error generating design documents: {type(e).__name__} - {str(e)}"
+            state.design_documents = error_msg
+            logger.error(f"{error_msg}, Input User Stories (snippet): {user_stories_for_prompt[:200]}...")
             return {"design_documents": state.design_documents}
-        
-
-      
       
     @log_entry_exit
     def development_artifact(self, state: State) -> dict:
         """Generate development artifacts based on design documents."""
-        # state.development_artifact = "DUMMY DEVELOPMENT ARTIFACT: This is a dummy development artifact for testing."
-        # return {"development_artifact": state.development_artifact}
         logger.info("Generating development artifacts")
         if not state.design_documents:
             state.development_artifact = "No design documents generated yet."
             logger.warning("Cannot generate development artifacts without design documents.")
             return {"development_artifact": state.development_artifact}
+        
+        # Escape design_documents and project_name for .format()
+        design_documents_for_prompt = str(state.design_documents).replace('{', '{{').replace('}', '}}')
+        project_name_val = state.project_name or 'N/A'
+        project_name_for_prompt = str(project_name_val).replace('{', '{{').replace('}', '}}')
+
         try:
             prompt_string = prompt.DEVELOPMENT_ARTIFACT_PROMPT_STRING.format(
-                design_documents=state.design_documents
+                design_documents=design_documents_for_prompt
+            )
+            sys_prompt_content = prompt.DEVELOPMENT_ARTIFACT_SYS_PROMPT.format(
+                project_name=project_name_for_prompt
             )
             messages = [
-                SystemMessage(content=prompt.DEVELOPMENT_ARTIFACT_SYS_PROMPT.format(
-                    project_name=state.project_name or 'N/A'
-                )),
+                SystemMessage(content=sys_prompt_content),
                 HumanMessage(content=prompt_string)
             ]
             response = self.llm.invoke(messages)
             state.development_artifact = response.content if hasattr(response, 'content') else str(response)
             return {"development_artifact": state.development_artifact}
+        except KeyError as ke:
+            error_msg = f"KeyError during DEVELOPMENT_ARTIFACT prompt formatting: {str(ke)}."
+            logger.error(error_msg)
+            state.development_artifact = error_msg
+            return {"development_artifact": state.development_artifact}
         except Exception as e:
             logger.error(f"Error generating development artifacts: {e}")
             state.development_artifact = f"Error generating development artifacts: {str(e)}"
@@ -192,27 +248,38 @@ class SdlcNode:
     @log_entry_exit
     def testing_artifact(self, state: State) -> dict:
         """Generate testing artifacts based on development artifacts."""
-        # state.testing_artifact = "DUMMY TESTING ARTIFACT: This is a dummy testing artifact for testing."
-        # return {"testing_artifact": state.testing_artifact}
         logger.info("Generating testing artifacts")
         if not state.development_artifact:
             state.testing_artifact = "No development artifacts generated yet."
             logger.warning("Cannot generate testing artifacts without development artifacts.")
             return {"testing_artifact": state.testing_artifact}
+
+        # Escape inputs for .format()
+        user_stories_for_prompt = str(state.user_stories).replace('{', '{{').replace('}', '}}')
+        development_artifact_for_prompt = str(state.development_artifact).replace('{', '{{').replace('}', '}}')
+        project_name_val = state.project_name or 'N/A'
+        project_name_for_prompt = str(project_name_val).replace('{', '{{').replace('}', '}}')
+
         try:
             prompt_string = prompt.TESTING_ARTIFACT_PROMPT_STRING.format(
-                user_stories=state.user_stories,
-                development_artifact=state.development_artifact
+                user_stories=user_stories_for_prompt,
+                development_artifact=development_artifact_for_prompt
+            )
+            sys_prompt_content = prompt.TESTING_ARTIFACT_SYS_PROMPT.format(
+                project_name=project_name_for_prompt
             )
             messages = [
-                SystemMessage(content=prompt.TESTING_ARTIFACT_SYS_PROMPT.format(
-                    project_name=state.project_name or 'N/A'
-                )),
+                SystemMessage(content=sys_prompt_content),
                 HumanMessage(content=prompt_string)
             ]
             response = self.llm.invoke(messages)
             state.testing_artifact = response.content if hasattr(response, 'content') else str(response)
             return {"testing_artifact": state.testing_artifact}
+        except KeyError as ke:
+            error_msg = f"KeyError during TESTING_ARTIFACT prompt formatting: {str(ke)}."
+            logger.error(error_msg)
+            state.testing_artifact = error_msg
+            return {"testing_artifact": state.testing_artifact}
         except Exception as e:
             logger.error(f"Error generating testing artifacts: {e}")
             state.testing_artifact = f"Error generating testing artifacts: {str(e)}"
@@ -221,19 +288,64 @@ class SdlcNode:
     @log_entry_exit
     def deployment_artifact(self, state: State) -> dict:
         """Generate deployment artifacts based on testing artifacts."""
-        # state.deployment_artifact = "DUMMY DEPLOYMENT ARTIFACT: This is a dummy deployment artifact for testing."
-        # return {"deployment_artifact": state.deployment_artifact}
         logger.info("Generating deployment artifacts")
         if not state.testing_artifact:
             state.deployment_artifact = "No testing artifacts generated yet."
             logger.warning("Cannot generate deployment artifacts without testing artifacts.")
             return {"deployment_artifact": state.deployment_artifact}
+
+        # Escape inputs for .format()
+        # The DEPLOYMENT_ARTIFACT_PROMPT_STRING uses {state}, which is not standard for .format()
+        # It should be specific fields like {state.testing_artifact} or {testing_artifact}
+        # For now, assuming it wants the string representation of the testing_artifact and project_name
+        testing_artifact_for_prompt = str(state.testing_artifact).replace('{', '{{').replace('}', '}}')
+        project_name_val = state.project_name or 'N/A'
+        project_name_for_prompt = str(project_name_val).replace('{', '{{').replace('}', '}}')
+        
         try:
-            prompt_string = prompt.DEPLOYMENT_ARTIFACT_PROMPT_STRING.format(state=state)
-            messages = [SystemMessage(content=prompt.DEPLOYMENT_ARTIFACT_SYS_PROMPT), HumanMessage(content=prompt_string)]
+            # --- IMPORTANT: Review prompt.DEPLOYMENT_ARTIFACT_PROMPT_STRING ---
+            # It currently has .format(state=state). This is unusual.
+            # It should likely be .format(testing_artifact=testing_artifact_for_prompt, project_name=project_name_for_prompt)
+            # or similar, depending on its actual placeholders.
+            # For now, I'll assume it expects 'testing_artifact' and 'project_name' as keys.
+            # You MUST adjust this if your prompt uses different placeholder names.
+            
+            # Tentative formatting based on common patterns, ADJUST IF YOUR PROMPT IS DIFFERENT
+            try:
+                prompt_string_content = prompt.DEPLOYMENT_ARTIFACT_PROMPT_STRING.format(
+                    testing_artifact=testing_artifact_for_prompt 
+                    # Add other fields from 'state' if your prompt actually uses them like {state.project_name}
+                )
+                # Assuming DEPLOYMENT_ARTIFACT_SYS_PROMPT expects project_name
+                sys_prompt_content = prompt.DEPLOYMENT_ARTIFACT_SYS_PROMPT.format(
+                    project_name=project_name_for_prompt
+                )
+            except KeyError as ke:
+                 # Fallback if the prompt string uses {state.testing_artifact} directly (less common for general prompts)
+                if 'state.testing_artifact' in str(ke): # Check if the error is about 'state.testing_artifact'
+                    prompt_string_content = prompt.DEPLOYMENT_ARTIFACT_PROMPT_STRING.format(
+                        state=state # Pass the whole state object if the prompt needs it this way
+                    )
+                    sys_prompt_content = prompt.DEPLOYMENT_ARTIFACT_SYS_PROMPT.format(
+                        project_name=project_name_for_prompt # This part is likely fine
+                    )
+                else: # Re-raise if it's a different KeyError
+                    raise ke
+
+
+            messages = [
+                SystemMessage(content=sys_prompt_content), 
+                HumanMessage(content=prompt_string_content)
+            ]
             response = self.llm.invoke(messages)
             state.deployment_artifact = response.content if hasattr(response, 'content') else str(response)
             return {"deployment_artifact": state.deployment_artifact}
+        except KeyError as ke:
+            error_msg = f"KeyError during DEPLOYMENT_ARTIFACT prompt formatting: {str(ke)}. Review prompt placeholders."
+            logger.error(error_msg)
+            logger.error(f"Prompt string might be: {prompt.DEPLOYMENT_ARTIFACT_PROMPT_STRING}")
+            state.deployment_artifact = error_msg
+            return {"deployment_artifact": state.deployment_artifact}
         except Exception as e:
             logger.error(f"Error generating deployment artifacts: {e}")
             state.deployment_artifact = f"Error generating deployment artifacts: {str(e)}"
@@ -243,31 +355,30 @@ class SdlcNode:
     def process_feedback(self, state: State) -> dict:
         """
         Process user feedback and update state with decision.
-        Only { "current_stage": ["accept"] } ends the process.
         """
-        
-        # logger.info(f"Input state: {state.to_dict()}")
-        logger.info(f"Feedback: {state.feedback}")
-
+        logger.info(f"Processing feedback. Current feedback state: {state.feedback}")
         current_stage = state.current_stage
-        feedback_for_stage = state.get_last_feedback_for_stage(current_stage)
-
-        if feedback_for_stage is None:
-            logger.warning(f"No feedback found for stage {current_stage}. Available feedback: {state.feedback}")
-
-        state.add_feedback(current_stage, str(feedback_for_stage))
+        # feedback_decision is set by the UI/controller before calling this node
+        feedback_text_from_ui = state.feedback.get(current_stage.value, [None])[-1] # Get latest feedback for current stage
 
-        logger.info(f"Processing feedback for stage: {current_stage}")
-        logger.info(f"Feedback received: {feedback_for_stage}")
+        logger.info(f"Processing feedback for stage: {current_stage}, Decision from UI: {state.feedback_decision}, Text: {feedback_text_from_ui}")
 
-        if feedback_for_stage and feedback_for_stage.strip().lower() == "accept":
-            logger.info(f"Feedback for stage '{current_stage}' is ACCEPT. Ending flow.")
-            state.feedback_decision = "accept"
-        else:
-            logger.info(f"Feedback for stage '{current_stage}' is not accept. Looping back. Feedback is: {feedback_for_stage}")
+        if state.feedback_decision == "accept":
+            logger.info(f"Feedback for stage '{current_stage}' is ACCEPT based on UI decision.")
+            # Optionally, store "accept" as a feedback entry if needed for history, though decision is primary
+            state.add_feedback(current_stage, "User accepted.") 
+        elif state.feedback_decision == "reject":
+            logger.info(f"Feedback for stage '{current_stage}' is REJECT based on UI decision. Feedback text: {feedback_text_from_ui}")
+            if feedback_text_from_ui:
+                 state.add_feedback(current_stage, str(feedback_text_from_ui)) # Add the actual feedback text
+            else:
+                 state.add_feedback(current_stage, "User rejected, no specific feedback provided.")
+        else: # Should not happen if UI sets feedback_decision correctly
+            logger.warning(f"Unknown feedback_decision '{state.feedback_decision}' for stage {current_stage}. Defaulting to reject.")
             state.feedback_decision = "reject"
-        logger.info(f"Feedback: {state.feedback}")
+            state.add_feedback(current_stage, f"System default to reject due to unknown decision: {state.feedback_decision}")
        
+        logger.info(f"Updated feedback state: {state.feedback}")
         return {"feedback_decision": state.feedback_decision}   
 
     @log_entry_exit
@@ -277,26 +388,31 @@ class SdlcNode:
 
         if not isinstance(state, State):
             logger.error(f"Invalid state type: {type(state)}. Routing to reject.")
-            return "reject"
+            return "reject" # Or handle as an error state
 
-        if state.feedback_decision == "accept":
+        decision = state.feedback_decision # Already processed in process_feedback
+
+        if decision == "accept":
             logger.info("Feedback accepted. Routing to next stage.")
             next_stage = state.get_next_stage()
             if next_stage:
-                state.update_stage(next_stage)
-                logger.info(f"Updated state to next stage: {next_stage}")
+                state.update_stage(next_stage) # This should update current_stage
+                st.session_state["current_stage"] = state.current_stage # Reflect in Streamlit
+                logger.info(f"Updated state to next stage: {state.current_stage}")
             else:
                 logger.info("No next stage available. Ending process.")
-                
-            state.clear_feedback_decision()
-            st.session_state["feedback_decision"] = None
-            logger.info(f"Feedback: {state.feedback}")
-            logger.info(f"Feedback decision: {state.feedback_decision}")
+                # No need to update current_stage if ending
+            
+            state.clear_feedback_decision() # Clear for the next cycle
+            st.session_state["feedback_decision"] = None # Reflect in Streamlit
+            # Clear specific feedback text from UI for the accepted stage if desired
+            # st.session_state["feedback"] = {} # Or selectively clear
+            logger.info(f"Feedback decision cleared. Current state feedback: {state.feedback}")
             return "accept"
-        else:
-            logger.info("Feedback rejected. Routing back for revision.")
+        else: # Covers 'reject' and any other case defaulting to reject
+            logger.info(f"Feedback decision is '{decision}'. Routing back for revision of stage {state.current_stage}.")
+            # The current stage remains the same for revision.
             state.clear_feedback_decision()
             st.session_state["feedback_decision"] = None
-            return "reject"
-
-
+            logger.info(f"Feedback decision cleared for revision. Current state feedback: {state.feedback}")
+            return "reject"