Implement Feature Workflow

prometheus/lang_graph/graphs/issue_graph.py

@@ -11,6 +11,7 @@
 from prometheus.lang_graph.nodes.issue_classification_subgraph_node import (
     IssueClassificationSubgraphNode,
 )
+from prometheus.lang_graph.nodes.issue_feature_subgraph_node import IssueFeatureSubgraphNode
 from prometheus.lang_graph.nodes.issue_question_subgraph_node import IssueQuestionSubgraphNode
 from prometheus.lang_graph.nodes.noop_node import NoopNode
 
@@ -67,13 +68,24 @@ def __init__(
             repository_id=repository_id,
         )
 
+        # Subgraph node for handling feature request issues
+        issue_feature_subgraph_node = IssueFeatureSubgraphNode(
+            advanced_model=advanced_model,
+            base_model=base_model,
+            container=container,
+            kg=kg,
+            git_repo=git_repo,
+            repository_id=repository_id,
+        )
+
         # Create the state graph for the issue handling workflow
         workflow = StateGraph(IssueState)
         # Add nodes to the workflow
         workflow.add_node("issue_type_branch_node", issue_type_branch_node)
         workflow.add_node("issue_classification_subgraph_node", issue_classification_subgraph_node)
         workflow.add_node("issue_bug_subgraph_node", issue_bug_subgraph_node)
         workflow.add_node("issue_question_subgraph_node", issue_question_subgraph_node)
+        workflow.add_node("issue_feature_subgraph_node", issue_feature_subgraph_node)
         # Set the entry point for the workflow
         workflow.set_entry_point("issue_type_branch_node")
         # Define the edges and conditions for the workflow
@@ -84,7 +96,7 @@ def __init__(
             {
                 IssueType.AUTO: "issue_classification_subgraph_node",
                 IssueType.BUG: "issue_bug_subgraph_node",
-                IssueType.FEATURE: END,
+                IssueType.FEATURE: "issue_feature_subgraph_node",
                 IssueType.DOCUMENTATION: END,
                 IssueType.QUESTION: "issue_question_subgraph_node",
             },
@@ -95,14 +107,15 @@ def __init__(
             lambda state: state["issue_type"],
             {
                 IssueType.BUG: "issue_bug_subgraph_node",
-                IssueType.FEATURE: END,
+                IssueType.FEATURE: "issue_feature_subgraph_node",
                 IssueType.DOCUMENTATION: END,
                 IssueType.QUESTION: "issue_question_subgraph_node",
             },
         )
         # Add edges for ending the workflow
         workflow.add_edge("issue_bug_subgraph_node", END)
         workflow.add_edge("issue_question_subgraph_node", END)
+        workflow.add_edge("issue_feature_subgraph_node", END)
 
         self.graph = workflow.compile()
 

prometheus/lang_graph/nodes/edit_message_node.py

@@ -14,15 +14,15 @@ class EditMessageNode:
 {issue_info}
 --- END ISSUE INFO ---
 
-Bug Context Found:
---- BEGIN BUG FIX CONTEXT ---
-{bug_fix_context}
---- END BUG FIX CONTEXT ---
+Context Found:
+--- BEGIN CONTEXT ---
+{context}
+--- END CONTEXT ---
 
-Bug analyzer agent has analyzed the issue and provided instruction on how to fix it:
---- BEGIN BUG ANALYZER MESSAGE ---
-{bug_analyzer_message}
---- END BUG ANALYZER MESSAGE ---
+Analyzer agent has analyzed the issue and provided instruction on the issue:
+--- BEGIN ANALYZER MESSAGE ---
+{analyzer_message}
+--- END ANALYZER MESSAGE ---
 
 Please implement these changes precisely, following the exact specifications from the analyzer.
 """
@@ -33,17 +33,18 @@ class EditMessageNode:
 {edit_error}
 --- END EDIT ERROR ---
 
-Bug analyzer agent has analyzed the issue and provided instruction on how to fix it:
---- BEGIN BUG ANALYZER MESSAGE ---
-{bug_analyzer_message}
---- END BUG ANALYZER MESSAGE ---
+Analyzer agent has analyzed the issue and provided instruction on the issue:
+--- BEGIN ANALYZER MESSAGE ---
+{analyzer_message}
+--- END ANALYZER MESSAGE ---
 
-Please implement these revised changes carefully, ensuring you address the
-specific issues that caused the previous error.
+Please implement these revised changes carefully, ensuring you address the specific issues that caused the previous error.
 """
 
-    def __init__(self):
+    def __init__(self, context_key: str, analyzer_message_key: str):
         self._logger = logging.getLogger(f"thread-{threading.get_ident()}.{__name__}")
+        self.context_key = context_key
+        self.analyzer_message_key = analyzer_message_key
 
     def format_human_message(self, state: Dict):
         edit_error = ""
@@ -58,9 +59,7 @@ def format_human_message(self, state: Dict):
             return HumanMessage(
                 self.FOLLOWUP_HUMAN_PROMPT.format(
                     edit_error=edit_error,
-                    bug_analyzer_message=get_last_message_content(
-                        state["issue_bug_analyzer_messages"]
-                    ),
+                    analyzer_message=get_last_message_content(state[self.analyzer_message_key]),
                 )
             )
 
@@ -69,8 +68,8 @@ def format_human_message(self, state: Dict):
                 issue_info=format_issue_info(
                     state["issue_title"], state["issue_body"], state["issue_comments"]
                 ),
-                bug_fix_context="\n\n".join([str(context) for context in state["bug_fix_context"]]),
-                bug_analyzer_message=get_last_message_content(state["issue_bug_analyzer_messages"]),
+                context="\n\n".join([str(context) for context in state[self.context_key]]),
+                analyzer_message=get_last_message_content(state[self.analyzer_message_key]),
             )
         )
 

prometheus/lang_graph/nodes/final_patch_selection_node.py

@@ -26,7 +26,7 @@ class FinalPatchSelectionNode:
 4. STYLE COHERENCE: The patch should maintain consistent coding style with the surrounding code
 
 Analysis Process:
-1. First, understand the issue from the provided issue_info and bug_context
+1. First, understand the issue from the provided issue_info and context
 2. Examine each patch carefully, considering:
    - Does it fix the root cause of the issue?
    - Does it maintain existing behavior (if appropriate)?
@@ -49,7 +49,7 @@ class FinalPatchSelectionNode:
 Comments: - Occurs in production environment
           - Affects customer-facing API
 
-Bug Context:
+Context:
 ```java
 // File: src/main/java/com/example/service/UserService.java
 public User getUser(String userId) {
@@ -113,8 +113,8 @@ class FinalPatchSelectionNode:
     HUMAN_PROMPT = """\
 {issue_info}
 
-Bug Context:
-{bug_fix_context}
+Context:
+{context}
 
 I have generated the following patches, now please select the best patch among them:
 {patches}
@@ -124,9 +124,12 @@ class FinalPatchSelectionNode:
 - patch_index: The index of the selected patch (must be valid within the given range)
 """
 
-    def __init__(self, model: BaseChatModel, candidate_patch_key: str, final_patch_key: str):
+    def __init__(
+        self, model: BaseChatModel, candidate_patch_key: str, final_patch_key: str, context_key: str
+    ):
         self.candidate_patch_key = candidate_patch_key
         self.final_patch_key = final_patch_key
+        self.context_key = context_key
         prompt = ChatPromptTemplate.from_messages(
             [("system", self.SYS_PROMPT), ("human", "{human_prompt}")]
         )
@@ -149,7 +152,7 @@ def format_human_message(self, patches: Sequence[str], state: Dict):
             issue_info=format_issue_info(
                 state["issue_title"], state["issue_body"], state["issue_comments"]
             ),
-            bug_fix_context="\n\n".join([str(context) for context in state["bug_fix_context"]]),
+            context="\n\n".join([str(context) for context in state[self.context_key]]),
             patches=patches_str,
         )
 

prometheus/lang_graph/nodes/issue_feature_analyzer_message_node.py

@@ -0,0 +1,107 @@
+import logging
+import threading
+from typing import Dict
+
+from langchain_core.messages import HumanMessage
+
+from prometheus.utils.issue_util import format_issue_info
+
+
+class IssueFeatureAnalyzerMessageNode:
+    FIRST_HUMAN_PROMPT = """\
+I am going to share details about a feature request reported to a codebase and its related context.
+Please analyze this feature request and provide a high-level description of what needs to be implemented:
+
+1. Feature Understanding:
+- Analyze the feature request title, description, and comments provided
+- Identify the desired functionality and requirements
+- Clarify any ambiguities or edge cases
+
+2. Architecture Analysis:
+- Identify which files, modules, or components need to be created or modified
+- Determine how this feature integrates with existing code
+- Consider architectural patterns and conventions from the codebase
+
+3. Implementation Plan:
+For each needed change, describe in plain English:
+- Which file needs to be created or modified
+- Which classes, functions, or code blocks need to be added or changed
+- What needs to be implemented (e.g., "add new method to handle X", "create new service class for Y")
+- How this integrates with existing components
+
+4. Considerations:
+- Identify potential impacts on existing functionality
+- Consider backward compatibility
+- Note any dependencies or prerequisites
+
+Do NOT provide actual code snippets or diffs. Focus on describing what needs to be implemented.
+
+Here are the details for analysis:
+
+{issue_info}
+
+Feature Context:
+{feature_context}
+"""
+
+    FOLLOWUP_HUMAN_PROMPT = """\
+Given your suggestion, the edit agent generated the following patch:
+{edit_patch}
+
+The patch generated following error:
+{edit_error}
+
+Please analyze the failure and provide a revised implementation suggestion:
+
+1. Error Analysis:
+- Explain why the previous implementation failed
+- Identify what specific aspects were problematic
+
+2. Revised Implementation Suggestion:
+Describe in plain English:
+- Which file needs to be created or modified
+- Which classes, functions, or code blocks need to be added or changed
+- What needs to be implemented (e.g., "add new method to handle X", "create new service class for Y")
+- Why this change would fix the error and properly implement the feature
+
+Do NOT provide actual code snippets or diffs. Focus on describing what needs to be implemented.
+"""
+
+    def __init__(self):
+        self._logger = logging.getLogger(f"thread-{threading.get_ident()}.{__name__}")
+
+    def format_human_message(self, state: Dict):
+        edit_error = ""
+        if (
+            "tested_patch_result" in state
+            and state["tested_patch_result"]
+            and not state["tested_patch_result"][0].passed
+        ):
+            edit_error = (
+                f"The patch failed to pass the regression tests:\n"
+                f"{state['tested_patch_result'][0].regression_test_failure_log}"
+            )
+
+        if not edit_error:
+            return HumanMessage(
+                self.FIRST_HUMAN_PROMPT.format(
+                    issue_info=format_issue_info(
+                        state["issue_title"], state["issue_body"], state["issue_comments"]
+                    ),
+                    feature_context="\n\n".join(
+                        [str(context) for context in state["feature_context"]]
+                    ),
+                )
+            )
+
+        return HumanMessage(
+            self.FOLLOWUP_HUMAN_PROMPT.format(
+                edit_patch=state["edit_patch"],
+                edit_error=edit_error,
+            )
+        )
+
+    def __call__(self, state: Dict):
+        human_message = self.format_human_message(state)
+        self._logger.debug(f"Sending message to IssueFeatureAnalyzerNode:\n{human_message}")
+        return {"issue_feature_analyzer_messages": [human_message]}

prometheus/lang_graph/nodes/issue_feature_analyzer_node.py

@@ -0,0 +1,99 @@
+import functools
+import logging
+import threading
+from typing import Dict
+
+from langchain.tools import StructuredTool
+from langchain_core.language_models.chat_models import BaseChatModel
+from langchain_core.messages import SystemMessage
+
+from prometheus.tools.web_search import WebSearchTool
+
+
+class IssueFeatureAnalyzerNode:
+    SYS_PROMPT = """\
+You are an expert software engineer specializing in feature implementation and software design. Your role is to:
+
+1. Carefully analyze feature requests by:
+   - Understanding the requested functionality and user requirements
+   - Identifying how this feature fits into the existing codebase
+   - Determining integration points with current components
+
+2. Design implementation approaches through systematic analysis:
+   - Analyze similar existing features and patterns in the codebase
+   - Identify which components need to be created or modified
+   - Understand architectural constraints and conventions
+   - Consider scalability and maintainability
+
+3. Provide high-level implementation plans by describing:
+   - Which specific files need to be created or modified
+   - Which classes, functions, or modules need to be added or changed
+   - What logical changes are needed (e.g., "create new service class for X", "add method to handle Y")
+   - How this integrates with existing components
+   - Why these changes properly implement the requested feature
+
+4. For implementation failures, analyze by:
+   - Understanding error messages and test failures
+   - Identifying what went wrong with the previous attempt
+   - Suggesting revised high-level changes that avoid the previous issues
+
+MANDATORY TOOL USAGE:
+- You MUST use the web_search tool for EVERY feature analysis
+- Before providing any analysis, search for:
+  * Best practices for implementing similar features
+  * Design patterns commonly used for this type of functionality
+  * Official documentation for relevant libraries/frameworks
+  * Common pitfalls and considerations
+- Only proceed with analysis after gathering relevant web information
+
+Tools available:
+- web_search: Searches the web for technical information to aid in feature design and implementation.
+When using the web_search tool, ALWAYS include these parameters:
+    - exclude_domains: ["*swe-bench*"]
+    - include_domains: ['stackoverflow.com', 'github.com', 'developer.mozilla.org', 'learn.microsoft.com', 'fastapi.tiangolo.com'
+            'docs.python.org', 'pydantic.dev', 'pypi.org', 'readthedocs.org', 'docs.djangoproject.com','flask.palletsprojects.com']
+    - search_depth: "advanced"
+
+    Make sure to explicitly pass these parameters in your tool call.
+
+Important:
+- Do NOT provide actual code snippets or diffs
+- DO provide clear file paths and function names where changes are needed
+- Focus on describing WHAT needs to be implemented and WHY, not HOW to implement it
+- Keep descriptions precise and actionable, as they will be used by another agent to implement the changes
+- ALWAYS start your analysis with web search results
+- Consider backward compatibility and existing architectural patterns
+
+Communicate in a clear, technical manner focused on accurate analysis and practical implementation plans
+rather than implementation details.
+"""
+
+    def __init__(self, model: BaseChatModel):
+        self.web_search_tool = WebSearchTool()
+        self.model = model
+        self.system_prompt = SystemMessage(self.SYS_PROMPT)
+        self.tools = self._init_tools()
+        self.model_with_tools = model.bind_tools(self.tools)
+        self._logger = logging.getLogger(f"thread-{threading.get_ident()}.{__name__}")
+
+    def _init_tools(self):
+        """Initializes tools for the node."""
+        tools = []
+
+        web_search_fn = functools.partial(self.web_search_tool.web_search)
+        web_search_tool = StructuredTool.from_function(
+            func=web_search_fn,
+            name=self.web_search_tool.web_search.__name__,
+            description=self.web_search_tool.web_search_spec.description,
+            args_schema=self.web_search_tool.web_search_spec.input_schema,
+        )
+        tools.append(web_search_tool)
+
+        return tools
+
+    def __call__(self, state: Dict):
+        message_history = [self.system_prompt] + state["issue_feature_analyzer_messages"]
+        response = self.model_with_tools.invoke(message_history)
+
+        self._logger.debug(response)
+        return {"issue_feature_analyzer_messages": [response]}