From 1367cce382f791938dabcda8ff9ac84644d8e579 Mon Sep 17 00:00:00 2001 From: Yue Pan <79363355+dcloud347@users.noreply.github.com> Date: Fri, 24 Oct 2025 21:53:03 +0100 Subject: [PATCH 1/2] feat: implement issue documentation subgraph for handling documentation update requests --- prometheus/lang_graph/graphs/issue_graph.py | 18 +- ...sue_documentation_analyzer_message_node.py | 57 ++++++ .../issue_documentation_analyzer_node.py | 71 +++++++ ...ssue_documentation_context_message_node.py | 32 ++++ .../issue_documentation_edit_message_node.py | 51 +++++ .../issue_documentation_responder_node.py | 71 +++++++ .../issue_documentation_subgraph_node.py | 69 +++++++ .../subgraphs/issue_documentation_state.py | 23 +++ .../subgraphs/issue_documentation_subgraph.py | 177 ++++++++++++++++++ 9 files changed, 567 insertions(+), 2 deletions(-) create mode 100644 prometheus/lang_graph/nodes/issue_documentation_analyzer_message_node.py create mode 100644 prometheus/lang_graph/nodes/issue_documentation_analyzer_node.py create mode 100644 prometheus/lang_graph/nodes/issue_documentation_context_message_node.py create mode 100644 prometheus/lang_graph/nodes/issue_documentation_edit_message_node.py create mode 100644 prometheus/lang_graph/nodes/issue_documentation_responder_node.py create mode 100644 prometheus/lang_graph/nodes/issue_documentation_subgraph_node.py create mode 100644 prometheus/lang_graph/subgraphs/issue_documentation_state.py create mode 100644 prometheus/lang_graph/subgraphs/issue_documentation_subgraph.py diff --git a/prometheus/lang_graph/graphs/issue_graph.py b/prometheus/lang_graph/graphs/issue_graph.py index 1ad007b..4c51ea0 100644 --- a/prometheus/lang_graph/graphs/issue_graph.py +++ b/prometheus/lang_graph/graphs/issue_graph.py @@ -11,6 +11,9 @@ from prometheus.lang_graph.nodes.issue_classification_subgraph_node import ( IssueClassificationSubgraphNode, ) +from prometheus.lang_graph.nodes.issue_documentation_subgraph_node import ( + IssueDocumentationSubgraphNode, +) 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 @@ -78,6 +81,15 @@ def __init__( repository_id=repository_id, ) + # Subgraph node for handling documentation issues + issue_documentation_subgraph_node = IssueDocumentationSubgraphNode( + advanced_model=advanced_model, + base_model=base_model, + 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 @@ -86,6 +98,7 @@ def __init__( 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) + workflow.add_node("issue_documentation_subgraph_node", issue_documentation_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 @@ -97,7 +110,7 @@ def __init__( IssueType.AUTO: "issue_classification_subgraph_node", IssueType.BUG: "issue_bug_subgraph_node", IssueType.FEATURE: "issue_feature_subgraph_node", - IssueType.DOCUMENTATION: END, + IssueType.DOCUMENTATION: "issue_documentation_subgraph_node", IssueType.QUESTION: "issue_question_subgraph_node", }, ) @@ -108,7 +121,7 @@ def __init__( { IssueType.BUG: "issue_bug_subgraph_node", IssueType.FEATURE: "issue_feature_subgraph_node", - IssueType.DOCUMENTATION: END, + IssueType.DOCUMENTATION: "issue_documentation_subgraph_node", IssueType.QUESTION: "issue_question_subgraph_node", }, ) @@ -116,6 +129,7 @@ def __init__( workflow.add_edge("issue_bug_subgraph_node", END) workflow.add_edge("issue_question_subgraph_node", END) workflow.add_edge("issue_feature_subgraph_node", END) + workflow.add_edge("issue_documentation_subgraph_node", END) self.graph = workflow.compile() diff --git a/prometheus/lang_graph/nodes/issue_documentation_analyzer_message_node.py b/prometheus/lang_graph/nodes/issue_documentation_analyzer_message_node.py new file mode 100644 index 0000000..6a4a3fe --- /dev/null +++ b/prometheus/lang_graph/nodes/issue_documentation_analyzer_message_node.py @@ -0,0 +1,57 @@ +import logging +import threading + +from langchain_core.messages import HumanMessage + +from prometheus.lang_graph.subgraphs.issue_documentation_state import IssueDocumentationState +from prometheus.utils.issue_util import format_issue_info + + +class IssueDocumentationAnalyzerMessageNode: + FIRST_HUMAN_PROMPT = """\ +I am going to share details about a documentation update request and its related context. +Please analyze this request and provide a detailed plan for updating the documentation: + +1. Issue Understanding: +- Analyze the issue title, description, and comments to understand what documentation needs to be updated +- Identify the type of documentation change (new documentation, updates, fixes, improvements) + +2. Context Analysis: +- Review the retrieved documentation files and source code context +- Identify which files need to be modified, created, or updated +- Understand the current documentation structure and style + +3. Documentation Plan: +- Provide a step-by-step plan for updating the documentation +- Specify which files to create, edit, or delete +- Describe what content needs to be added or changed +- Ensure consistency with existing documentation style +- Include any code examples, API references, or diagrams if needed + +Here is the documentation update request: +-- BEGIN ISSUE -- +{issue_info} +-- END ISSUE -- + +Here is the relevant context (existing documentation and source code): +--- BEGIN CONTEXT -- +{documentation_context} +--- END CONTEXT -- + +Based on the above information, please provide a detailed analysis and plan for updating the documentation. +""" + + def __init__(self): + self._logger = logging.getLogger(f"thread-{threading.get_ident()}.{__name__}") + + def __call__(self, state: IssueDocumentationState): + human_message = self.FIRST_HUMAN_PROMPT.format( + issue_info=format_issue_info( + state["issue_title"], state["issue_body"], state["issue_comments"] + ), + documentation_context="\n\n".join( + [str(context) for context in state["documentation_context"]] + ), + ) + self._logger.debug(f"Sending message to IssueDocumentationAnalyzerNode:\n{human_message}") + return {"issue_documentation_analyzer_messages": [HumanMessage(human_message)]} diff --git a/prometheus/lang_graph/nodes/issue_documentation_analyzer_node.py b/prometheus/lang_graph/nodes/issue_documentation_analyzer_node.py new file mode 100644 index 0000000..ee1e9bb --- /dev/null +++ b/prometheus/lang_graph/nodes/issue_documentation_analyzer_node.py @@ -0,0 +1,71 @@ +import functools +import logging +import threading + +from langchain_core.language_models.chat_models import BaseChatModel +from langchain_core.messages import SystemMessage +from langchain_core.tools import StructuredTool + +from prometheus.lang_graph.subgraphs.issue_documentation_state import IssueDocumentationState +from prometheus.tools.web_search import WebSearchTool + + +class IssueDocumentationAnalyzerNode: + SYS_PROMPT = """ +You are an expert technical writer and software documentation specialist. Your role is to: + +1. Carefully analyze documentation update requests by: + - Understanding the issue description and what documentation changes are needed + - Identifying which documentation files need to be created, updated, or modified + - Understanding the context of existing documentation and source code + +2. Create a comprehensive documentation plan through systematic analysis: + - Identify specific documentation files that need changes + - Determine what content needs to be added, updated, or removed + - Ensure consistency with existing documentation style and structure + - Consider the target audience and documentation purpose + +3. Provide a clear, actionable plan that includes: + - List of files to create, edit, or delete + - Specific content changes for each file + - Code examples, API references, or diagrams if needed + - Recommendations for improving documentation clarity and completeness + +Important: +- Keep descriptions precise and actionable +- Follow existing documentation conventions and style +- Ensure technical accuracy +- Only leave your direct final analysis and plan in the last response! + +Your analysis should be thorough enough that an editor can implement the changes directly. +""" + + def __init__(self, model: BaseChatModel): + self.system_prompt = SystemMessage(self.SYS_PROMPT) + self.web_search_tool = WebSearchTool() + 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: IssueDocumentationState): + message_history = [self.system_prompt] + state["issue_documentation_analyzer_messages"] + response = self.model_with_tools.invoke(message_history) + + self._logger.debug(response) + return {"issue_documentation_analyzer_messages": [response]} diff --git a/prometheus/lang_graph/nodes/issue_documentation_context_message_node.py b/prometheus/lang_graph/nodes/issue_documentation_context_message_node.py new file mode 100644 index 0000000..4a5a0ad --- /dev/null +++ b/prometheus/lang_graph/nodes/issue_documentation_context_message_node.py @@ -0,0 +1,32 @@ +import logging +import threading + +from prometheus.lang_graph.subgraphs.issue_documentation_state import IssueDocumentationState +from prometheus.utils.issue_util import format_issue_info + + +class IssueDocumentationContextMessageNode: + DOCUMENTATION_QUERY = """\ +{issue_info} + +Find all relevant documentation files and source code context needed to update the documentation according to this issue. +Focus on: +1. Existing documentation files (README.md, docs/, etc.) +2. Source code that needs to be documented or referenced +3. Related configuration files or examples +4. Any existing documentation that needs to be updated or extended + +Include both documentation files and relevant source code context. +""" + + def __init__(self): + self._logger = logging.getLogger(f"thread-{threading.get_ident()}.{__name__}") + + def __call__(self, state: IssueDocumentationState): + documentation_query = self.DOCUMENTATION_QUERY.format( + issue_info=format_issue_info( + state["issue_title"], state["issue_body"], state["issue_comments"] + ), + ) + self._logger.debug(f"Sending query to context provider:\n{documentation_query}") + return {"documentation_query": documentation_query} diff --git a/prometheus/lang_graph/nodes/issue_documentation_edit_message_node.py b/prometheus/lang_graph/nodes/issue_documentation_edit_message_node.py new file mode 100644 index 0000000..83d5953 --- /dev/null +++ b/prometheus/lang_graph/nodes/issue_documentation_edit_message_node.py @@ -0,0 +1,51 @@ +import logging +import threading + +from langchain_core.messages import HumanMessage + +from prometheus.lang_graph.subgraphs.issue_documentation_state import IssueDocumentationState +from prometheus.utils.lang_graph_util import get_last_message_content + + +class IssueDocumentationEditMessageNode: + EDIT_PROMPT = """\ +Based on the following documentation analysis and plan, please implement the documentation changes using the available file operation tools. + +Documentation Analysis and Plan: +--- BEGIN PLAN --- +{documentation_plan} +--- END PLAN --- + +Context Collected: +--- BEGIN CONTEXT -- +{documentation_context} +--- END CONTEXT -- + +Please proceed to implement the documentation changes: +1. Read existing files first to understand the current state +2. Make precise edits that preserve existing formatting and style +3. Create new files if specified in the plan +4. Verify your changes by reading the files again after editing + +Remember: +- Follow the plan provided in the analysis +- Maintain consistency with existing documentation style +- Ensure all edits are precise and accurate +- Do NOT write or run any tests - focus only on documentation updates +""" + + def __init__(self): + self._logger = logging.getLogger(f"thread-{threading.get_ident()}.{__name__}") + + def __call__(self, state: IssueDocumentationState): + documentation_plan = get_last_message_content( + state["issue_documentation_analyzer_messages"] + ) + documentation_context = "\n\n".join( + [str(context) for context in state["documentation_context"]] + ) + human_message = self.EDIT_PROMPT.format( + documentation_plan=documentation_plan, documentation_context=documentation_context + ) + self._logger.debug(f"Sending message to EditNode:\n{human_message}") + return {"edit_messages": [HumanMessage(human_message)]} diff --git a/prometheus/lang_graph/nodes/issue_documentation_responder_node.py b/prometheus/lang_graph/nodes/issue_documentation_responder_node.py new file mode 100644 index 0000000..407c757 --- /dev/null +++ b/prometheus/lang_graph/nodes/issue_documentation_responder_node.py @@ -0,0 +1,71 @@ +import logging +import threading + +from langchain_core.language_models.chat_models import BaseChatModel +from langchain_core.messages import HumanMessage, SystemMessage + +from prometheus.lang_graph.subgraphs.issue_documentation_state import IssueDocumentationState +from prometheus.utils.issue_util import format_issue_info + + +class IssueDocumentationResponderNode: + SYS_PROMPT = """\ +You are a technical writer summarizing documentation updates for a GitHub issue. + +Your task is to: +1. Review the documentation update request and the analysis/plan that was created +2. Review the patch that was generated +3. Create a clear, professional response explaining what documentation was updated + +The response should: +- Summarize what documentation changes were made +- Explain how the changes address the issue request +- Be concise but informative +- Use a professional, helpful tone + +Keep the response focused on what was accomplished, not implementation details. +""" + + USER_PROMPT = """\ +Here is the documentation update request: +-- BEGIN ISSUE -- +{issue_info} +-- END ISSUE -- + +Here is the analysis and plan that was created: +-- BEGIN PLAN -- +{documentation_plan} +-- END PLAN -- + +Here is the patch with the documentation changes: +-- BEGIN PATCH -- +{edit_patch} +-- END PATCH -- + +Please provide a clear, professional response summarizing the documentation updates for this issue. +""" + + def __init__(self, model: BaseChatModel): + self.model = model + self._logger = logging.getLogger(f"thread-{threading.get_ident()}.{__name__}") + + def __call__(self, state: IssueDocumentationState): + from prometheus.utils.lang_graph_util import get_last_message_content + + documentation_plan = get_last_message_content( + state["issue_documentation_analyzer_messages"] + ) + + user_message = self.USER_PROMPT.format( + issue_info=format_issue_info( + state["issue_title"], state["issue_body"], state["issue_comments"] + ), + documentation_plan=documentation_plan, + edit_patch=state.get("edit_patch", "No changes were made."), + ) + + messages = [SystemMessage(self.SYS_PROMPT), HumanMessage(user_message)] + response = self.model.invoke(messages) + + self._logger.info(f"Documentation update response:\n{response.content}") + return {"issue_response": response.content} diff --git a/prometheus/lang_graph/nodes/issue_documentation_subgraph_node.py b/prometheus/lang_graph/nodes/issue_documentation_subgraph_node.py new file mode 100644 index 0000000..d7740ce --- /dev/null +++ b/prometheus/lang_graph/nodes/issue_documentation_subgraph_node.py @@ -0,0 +1,69 @@ +import logging +import threading + +from langchain_core.language_models.chat_models import BaseChatModel +from langgraph.errors import GraphRecursionError + +from prometheus.git.git_repository import GitRepository +from prometheus.graph.knowledge_graph import KnowledgeGraph +from prometheus.lang_graph.graphs.issue_state import IssueState +from prometheus.lang_graph.subgraphs.issue_documentation_subgraph import ( + IssueDocumentationSubgraph, +) + + +class IssueDocumentationSubgraphNode: + """ + A LangGraph node that handles the issue documentation subgraph, + which is responsible for updating documentation based on a GitHub issue. + """ + + def __init__( + self, + advanced_model: BaseChatModel, + base_model: BaseChatModel, + kg: KnowledgeGraph, + git_repo: GitRepository, + repository_id: int, + ): + self._logger = logging.getLogger(f"thread-{threading.get_ident()}.{__name__}") + self.issue_documentation_subgraph = IssueDocumentationSubgraph( + advanced_model=advanced_model, + base_model=base_model, + kg=kg, + git_repo=git_repo, + repository_id=repository_id, + ) + + def __call__(self, state: IssueState): + # Logging entry into the node + self._logger.info("Enter IssueDocumentationSubgraphNode") + + try: + output_state = self.issue_documentation_subgraph.invoke( + issue_title=state["issue_title"], + issue_body=state["issue_body"], + issue_comments=state["issue_comments"], + ) + except GraphRecursionError: + # Handle recursion error gracefully + self._logger.critical( + "Please increase the recursion limit of IssueDocumentationSubgraph" + ) + return { + "edit_patch": None, + "passed_reproducing_test": False, + "passed_regression_test": False, + "passed_existing_test": False, + "issue_response": None, + } + + # Logging the issue response for debugging + self._logger.info(f"issue_response:\n{output_state['issue_response']}") + return { + "edit_patch": output_state["edit_patch"], + "passed_reproducing_test": output_state["passed_reproducing_test"], + "passed_regression_test": output_state["passed_regression_test"], + "passed_existing_test": output_state["passed_existing_test"], + "issue_response": output_state["issue_response"], + } diff --git a/prometheus/lang_graph/subgraphs/issue_documentation_state.py b/prometheus/lang_graph/subgraphs/issue_documentation_state.py new file mode 100644 index 0000000..d915b33 --- /dev/null +++ b/prometheus/lang_graph/subgraphs/issue_documentation_state.py @@ -0,0 +1,23 @@ +from typing import Annotated, Mapping, Sequence, TypedDict + +from langchain_core.messages import BaseMessage +from langgraph.graph import add_messages + +from prometheus.models.context import Context + + +class IssueDocumentationState(TypedDict): + issue_title: str + issue_body: str + issue_comments: Sequence[Mapping[str, str]] + + max_refined_query_loop: int + + documentation_query: str + documentation_context: Sequence[Context] + + issue_documentation_analyzer_messages: Annotated[Sequence[BaseMessage], add_messages] + edit_messages: Annotated[Sequence[BaseMessage], add_messages] + + edit_patch: str + issue_response: str diff --git a/prometheus/lang_graph/subgraphs/issue_documentation_subgraph.py b/prometheus/lang_graph/subgraphs/issue_documentation_subgraph.py new file mode 100644 index 0000000..874a6d9 --- /dev/null +++ b/prometheus/lang_graph/subgraphs/issue_documentation_subgraph.py @@ -0,0 +1,177 @@ +import functools +from typing import Mapping, Sequence + +from langchain_core.language_models.chat_models import BaseChatModel +from langgraph.constants import END +from langgraph.graph import StateGraph +from langgraph.prebuilt import ToolNode, tools_condition + +from prometheus.git.git_repository import GitRepository +from prometheus.graph.knowledge_graph import KnowledgeGraph +from prometheus.lang_graph.nodes.context_retrieval_subgraph_node import ContextRetrievalSubgraphNode +from prometheus.lang_graph.nodes.edit_node import EditNode +from prometheus.lang_graph.nodes.git_diff_node import GitDiffNode +from prometheus.lang_graph.nodes.issue_documentation_analyzer_message_node import ( + IssueDocumentationAnalyzerMessageNode, +) +from prometheus.lang_graph.nodes.issue_documentation_analyzer_node import ( + IssueDocumentationAnalyzerNode, +) +from prometheus.lang_graph.nodes.issue_documentation_context_message_node import ( + IssueDocumentationContextMessageNode, +) +from prometheus.lang_graph.nodes.issue_documentation_edit_message_node import ( + IssueDocumentationEditMessageNode, +) +from prometheus.lang_graph.nodes.issue_documentation_responder_node import ( + IssueDocumentationResponderNode, +) +from prometheus.lang_graph.subgraphs.issue_documentation_state import IssueDocumentationState + + +class IssueDocumentationSubgraph: + """ + A LangGraph-based subgraph to handle documentation update requests for GitHub issues. + + This subgraph processes documentation requests by: + 1. Retrieving relevant documentation and source code context + 2. Analyzing the request and creating a documentation update plan + 3. Implementing the documentation changes using file operations + 4. Generating a patch with the documentation updates + 5. Creating a response summarizing the documentation changes + """ + + def __init__( + self, + advanced_model: BaseChatModel, + base_model: BaseChatModel, + kg: KnowledgeGraph, + git_repo: GitRepository, + repository_id: int, + ): + # Step 1: Retrieve relevant context (documentation files and source code) + issue_documentation_context_message_node = IssueDocumentationContextMessageNode() + context_retrieval_subgraph_node = ContextRetrievalSubgraphNode( + base_model=base_model, + advanced_model=advanced_model, + kg=kg, + local_path=git_repo.playground_path, + query_key_name="documentation_query", + context_key_name="documentation_context", + repository_id=repository_id, + ) + + # Step 2: Analyze the documentation request and create a plan + issue_documentation_analyzer_message_node = IssueDocumentationAnalyzerMessageNode() + issue_documentation_analyzer_node = IssueDocumentationAnalyzerNode(model=advanced_model) + issue_documentation_analyzer_tools = ToolNode( + tools=issue_documentation_analyzer_node.tools, + name="issue_documentation_analyzer_tools", + messages_key="issue_documentation_analyzer_messages", + ) + + # Step 3: Implement the documentation changes + issue_documentation_edit_message_node = IssueDocumentationEditMessageNode() + edit_node = EditNode(advanced_model, git_repo.playground_path, kg) + edit_tools = ToolNode( + tools=edit_node.tools, + name="edit_tools", + messages_key="edit_messages", + ) + + # Step 4: Generate patch from changes + git_diff_node = GitDiffNode(git_repo, "edit_patch", return_list=False) + + # Step 5: Generate response summarizing the changes + issue_documentation_responder_node = IssueDocumentationResponderNode(model=base_model) + + # Define the subgraph structure + workflow = StateGraph(IssueDocumentationState) + + # Add all nodes + workflow.add_node( + "issue_documentation_context_message_node", + issue_documentation_context_message_node, + ) + workflow.add_node("context_retrieval_subgraph_node", context_retrieval_subgraph_node) + + workflow.add_node( + "issue_documentation_analyzer_message_node", + issue_documentation_analyzer_message_node, + ) + workflow.add_node("issue_documentation_analyzer_node", issue_documentation_analyzer_node) + workflow.add_node("issue_documentation_analyzer_tools", issue_documentation_analyzer_tools) + + workflow.add_node( + "issue_documentation_edit_message_node", issue_documentation_edit_message_node + ) + workflow.add_node("edit_node", edit_node) + workflow.add_node("edit_tools", edit_tools) + + workflow.add_node("git_diff_node", git_diff_node) + workflow.add_node("issue_documentation_responder_node", issue_documentation_responder_node) + + # Define the entry point + workflow.set_entry_point("issue_documentation_context_message_node") + + # Define the workflow transitions + workflow.add_edge( + "issue_documentation_context_message_node", "context_retrieval_subgraph_node" + ) + workflow.add_edge( + "context_retrieval_subgraph_node", "issue_documentation_analyzer_message_node" + ) + + workflow.add_edge( + "issue_documentation_analyzer_message_node", "issue_documentation_analyzer_node" + ) + workflow.add_conditional_edges( + "issue_documentation_analyzer_node", + functools.partial( + tools_condition, messages_key="issue_documentation_analyzer_messages" + ), + { + "tools": "issue_documentation_analyzer_tools", + END: "issue_documentation_edit_message_node", + }, + ) + workflow.add_edge("issue_documentation_analyzer_tools", "issue_documentation_analyzer_node") + + workflow.add_edge("issue_documentation_edit_message_node", "edit_node") + workflow.add_conditional_edges( + "edit_node", + functools.partial(tools_condition, messages_key="edit_messages"), + {"tools": "edit_tools", END: "git_diff_node"}, + ) + workflow.add_edge("edit_tools", "edit_node") + + workflow.add_edge("git_diff_node", "issue_documentation_responder_node") + workflow.add_edge("issue_documentation_responder_node", END) + + # Compile the workflow into an executable subgraph + self.subgraph = workflow.compile() + + def invoke( + self, + issue_title: str, + issue_body: str, + issue_comments: Sequence[Mapping[str, str]], + recursion_limit: int = 150, + ): + config = {"recursion_limit": recursion_limit} + + input_state = { + "issue_title": issue_title, + "issue_body": issue_body, + "issue_comments": issue_comments, + "max_refined_query_loop": 3, + } + + output_state = self.subgraph.invoke(input_state, config) + return { + "edit_patch": output_state["edit_patch"], + "passed_reproducing_test": False, + "passed_existing_test": False, + "passed_regression_test": False, + "issue_response": output_state["issue_response"], + } From febdf7515b61d75dc55390b7abd89d743a5e897d Mon Sep 17 00:00:00 2001 From: Yue Pan <79363355+dcloud347@users.noreply.github.com> Date: Fri, 24 Oct 2025 22:36:38 +0100 Subject: [PATCH 2/2] test: add unit tests for issue documentation analyzer and responder nodes --- ...sue_documentation_analyzer_message_node.py | 98 +++++++++++++++ .../test_issue_documentation_analyzer_node.py | 81 ++++++++++++ ...ssue_documentation_context_message_node.py | 80 ++++++++++++ ...t_issue_documentation_edit_message_node.py | 117 ++++++++++++++++++ ...test_issue_documentation_responder_node.py | 96 ++++++++++++++ .../test_issue_documentation_subgraph.py | 51 ++++++++ 6 files changed, 523 insertions(+) create mode 100644 tests/lang_graph/nodes/test_issue_documentation_analyzer_message_node.py create mode 100644 tests/lang_graph/nodes/test_issue_documentation_analyzer_node.py create mode 100644 tests/lang_graph/nodes/test_issue_documentation_context_message_node.py create mode 100644 tests/lang_graph/nodes/test_issue_documentation_edit_message_node.py create mode 100644 tests/lang_graph/nodes/test_issue_documentation_responder_node.py create mode 100644 tests/lang_graph/subgraphs/test_issue_documentation_subgraph.py diff --git a/tests/lang_graph/nodes/test_issue_documentation_analyzer_message_node.py b/tests/lang_graph/nodes/test_issue_documentation_analyzer_message_node.py new file mode 100644 index 0000000..56ecab8 --- /dev/null +++ b/tests/lang_graph/nodes/test_issue_documentation_analyzer_message_node.py @@ -0,0 +1,98 @@ +import pytest +from langchain_core.messages import HumanMessage + +from prometheus.lang_graph.nodes.issue_documentation_analyzer_message_node import ( + IssueDocumentationAnalyzerMessageNode, +) +from prometheus.lang_graph.subgraphs.issue_documentation_state import IssueDocumentationState +from prometheus.models.context import Context + + +@pytest.fixture +def basic_state(): + return IssueDocumentationState( + issue_title="Update API documentation", + issue_body="The API documentation needs to be updated", + issue_comments=[ + {"username": "user1", "comment": "Please add examples"}, + ], + max_refined_query_loop=3, + documentation_query="Find API docs", + documentation_context=[ + Context( + relative_path="/docs/api.md", + content="# API Documentation\n\nAPI Documentation content", + ) + ], + issue_documentation_analyzer_messages=[], + edit_messages=[], + edit_patch="", + issue_response="", + ) + + +def test_init_issue_documentation_analyzer_message_node(): + """Test IssueDocumentationAnalyzerMessageNode initialization.""" + node = IssueDocumentationAnalyzerMessageNode() + assert node is not None + + +def test_call_method_creates_message(basic_state): + """Test that the node creates a human message.""" + node = IssueDocumentationAnalyzerMessageNode() + result = node(basic_state) + + assert "issue_documentation_analyzer_messages" in result + assert len(result["issue_documentation_analyzer_messages"]) == 1 + assert isinstance(result["issue_documentation_analyzer_messages"][0], HumanMessage) + + +def test_message_contains_issue_info(basic_state): + """Test that the message contains issue information.""" + node = IssueDocumentationAnalyzerMessageNode() + result = node(basic_state) + + message_content = result["issue_documentation_analyzer_messages"][0].content + assert "Update API documentation" in message_content + + +def test_message_contains_context(basic_state): + """Test that the message includes documentation context.""" + node = IssueDocumentationAnalyzerMessageNode() + result = node(basic_state) + + message_content = result["issue_documentation_analyzer_messages"][0].content + # Should include context or reference to it + assert "context" in message_content.lower() or "API Documentation" in message_content + + +def test_message_includes_analysis_instructions(basic_state): + """Test that the message includes analysis instructions.""" + node = IssueDocumentationAnalyzerMessageNode() + result = node(basic_state) + + message_content = result["issue_documentation_analyzer_messages"][0].content + # Should include instructions for analysis + assert "plan" in message_content.lower() or "analyze" in message_content.lower() + + +def test_call_with_empty_context(): + """Test the node with empty documentation context.""" + state = IssueDocumentationState( + issue_title="Create new docs", + issue_body="Create documentation for new feature", + issue_comments=[], + max_refined_query_loop=3, + documentation_query="", + documentation_context=[], + issue_documentation_analyzer_messages=[], + edit_messages=[], + edit_patch="", + issue_response="", + ) + + node = IssueDocumentationAnalyzerMessageNode() + result = node(state) + + assert "issue_documentation_analyzer_messages" in result + assert len(result["issue_documentation_analyzer_messages"]) == 1 diff --git a/tests/lang_graph/nodes/test_issue_documentation_analyzer_node.py b/tests/lang_graph/nodes/test_issue_documentation_analyzer_node.py new file mode 100644 index 0000000..e2479c4 --- /dev/null +++ b/tests/lang_graph/nodes/test_issue_documentation_analyzer_node.py @@ -0,0 +1,81 @@ +import pytest +from langchain_core.messages import HumanMessage + +from prometheus.lang_graph.nodes.issue_documentation_analyzer_node import ( + IssueDocumentationAnalyzerNode, +) +from prometheus.lang_graph.subgraphs.issue_documentation_state import IssueDocumentationState +from tests.test_utils.util import FakeListChatWithToolsModel + + +@pytest.fixture +def fake_llm(): + return FakeListChatWithToolsModel( + responses=[ + "Documentation Plan:\n1. Update README.md with new API documentation\n" + "2. Add code examples\n3. Update table of contents" + ] + ) + + +@pytest.fixture +def basic_state(): + return IssueDocumentationState( + issue_title="Update API documentation", + issue_body="The API documentation needs to be updated with the new endpoints", + issue_comments=[ + {"username": "user1", "comment": "Please include examples"}, + {"username": "user2", "comment": "Add authentication details"}, + ], + max_refined_query_loop=3, + documentation_query="Find API documentation files", + documentation_context=[], + issue_documentation_analyzer_messages=[ + HumanMessage(content="Please analyze the documentation request and provide a plan") + ], + edit_messages=[], + edit_patch="", + issue_response="", + ) + + +def test_init_issue_documentation_analyzer_node(fake_llm): + """Test IssueDocumentationAnalyzerNode initialization.""" + node = IssueDocumentationAnalyzerNode(fake_llm) + + assert node.system_prompt is not None + assert node.web_search_tool is not None + assert len(node.tools) == 1 # Should have web search tool + assert node.model_with_tools is not None + + +def test_call_method_basic(fake_llm, basic_state): + """Test basic call functionality.""" + node = IssueDocumentationAnalyzerNode(fake_llm) + result = node(basic_state) + + assert "issue_documentation_analyzer_messages" in result + assert len(result["issue_documentation_analyzer_messages"]) == 1 + assert "Documentation Plan" in result["issue_documentation_analyzer_messages"][0].content + + +def test_call_method_with_empty_messages(fake_llm): + """Test call method with empty message history.""" + state = IssueDocumentationState( + issue_title="Test", + issue_body="Test body", + issue_comments=[], + max_refined_query_loop=3, + documentation_query="", + documentation_context=[], + issue_documentation_analyzer_messages=[], + edit_messages=[], + edit_patch="", + issue_response="", + ) + + node = IssueDocumentationAnalyzerNode(fake_llm) + result = node(state) + + assert "issue_documentation_analyzer_messages" in result + assert len(result["issue_documentation_analyzer_messages"]) == 1 diff --git a/tests/lang_graph/nodes/test_issue_documentation_context_message_node.py b/tests/lang_graph/nodes/test_issue_documentation_context_message_node.py new file mode 100644 index 0000000..146b905 --- /dev/null +++ b/tests/lang_graph/nodes/test_issue_documentation_context_message_node.py @@ -0,0 +1,80 @@ +import pytest + +from prometheus.lang_graph.nodes.issue_documentation_context_message_node import ( + IssueDocumentationContextMessageNode, +) +from prometheus.lang_graph.subgraphs.issue_documentation_state import IssueDocumentationState + + +@pytest.fixture +def basic_state(): + return IssueDocumentationState( + issue_title="Update API documentation", + issue_body="The API documentation needs to be updated with new endpoints", + issue_comments=[ + {"username": "user1", "comment": "Please include examples"}, + ], + max_refined_query_loop=3, + documentation_query="", + documentation_context=[], + issue_documentation_analyzer_messages=[], + edit_messages=[], + edit_patch="", + issue_response="", + ) + + +def test_init_issue_documentation_context_message_node(): + """Test IssueDocumentationContextMessageNode initialization.""" + node = IssueDocumentationContextMessageNode() + assert node is not None + + +def test_call_method_generates_query(basic_state): + """Test that the node generates a documentation query.""" + node = IssueDocumentationContextMessageNode() + result = node(basic_state) + + assert "documentation_query" in result + assert len(result["documentation_query"]) > 0 + + +def test_query_contains_issue_info(basic_state): + """Test that the query contains issue information.""" + node = IssueDocumentationContextMessageNode() + result = node(basic_state) + + query = result["documentation_query"] + assert "Update API documentation" in query or "API documentation" in query + + +def test_query_includes_instructions(basic_state): + """Test that the query includes documentation finding instructions.""" + node = IssueDocumentationContextMessageNode() + result = node(basic_state) + + query = result["documentation_query"] + # Should include instructions about finding documentation + assert "documentation" in query.lower() or "find" in query.lower() + + +def test_call_with_empty_comments(): + """Test the node with empty comments.""" + state = IssueDocumentationState( + issue_title="Test title", + issue_body="Test body", + issue_comments=[], + max_refined_query_loop=3, + documentation_query="", + documentation_context=[], + issue_documentation_analyzer_messages=[], + edit_messages=[], + edit_patch="", + issue_response="", + ) + + node = IssueDocumentationContextMessageNode() + result = node(state) + + assert "documentation_query" in result + assert len(result["documentation_query"]) > 0 diff --git a/tests/lang_graph/nodes/test_issue_documentation_edit_message_node.py b/tests/lang_graph/nodes/test_issue_documentation_edit_message_node.py new file mode 100644 index 0000000..dd41792 --- /dev/null +++ b/tests/lang_graph/nodes/test_issue_documentation_edit_message_node.py @@ -0,0 +1,117 @@ +import pytest +from langchain_core.messages import AIMessage, HumanMessage + +from prometheus.lang_graph.nodes.issue_documentation_edit_message_node import ( + IssueDocumentationEditMessageNode, +) +from prometheus.lang_graph.subgraphs.issue_documentation_state import IssueDocumentationState +from prometheus.models.context import Context + + +@pytest.fixture +def basic_state(): + return IssueDocumentationState( + issue_title="Update API documentation", + issue_body="The API documentation needs to be updated", + issue_comments=[], + max_refined_query_loop=3, + documentation_query="Find API docs", + documentation_context=[ + Context( + relative_path="/docs/api.md", + content="# API Documentation\n\nAPI Documentation content", + ) + ], + issue_documentation_analyzer_messages=[ + AIMessage(content="Plan:\n1. Update README.md\n2. Add new examples\n3. Fix typos") + ], + edit_messages=[], + edit_patch="", + issue_response="", + ) + + +def test_init_issue_documentation_edit_message_node(): + """Test IssueDocumentationEditMessageNode initialization.""" + node = IssueDocumentationEditMessageNode() + assert node is not None + + +def test_call_method_creates_message(basic_state): + """Test that the node creates a human message.""" + node = IssueDocumentationEditMessageNode() + result = node(basic_state) + + assert "edit_messages" in result + assert len(result["edit_messages"]) == 1 + assert isinstance(result["edit_messages"][0], HumanMessage) + + +def test_message_contains_plan(basic_state): + """Test that the message contains the documentation plan.""" + node = IssueDocumentationEditMessageNode() + result = node(basic_state) + + message_content = result["edit_messages"][0].content + assert "Plan:" in message_content or "Update README.md" in message_content + + +def test_message_contains_context(basic_state): + """Test that the message includes documentation context.""" + node = IssueDocumentationEditMessageNode() + result = node(basic_state) + + message_content = result["edit_messages"][0].content + # Should include context + assert "context" in message_content.lower() or "API Documentation" in message_content + + +def test_message_includes_edit_instructions(basic_state): + """Test that the message includes editing instructions.""" + node = IssueDocumentationEditMessageNode() + result = node(basic_state) + + message_content = result["edit_messages"][0].content + # Should include instructions for implementing changes + assert any( + keyword in message_content.lower() for keyword in ["implement", "edit", "changes", "file"] + ) + + +def test_call_with_empty_context(): + """Test the node with empty documentation context.""" + state = IssueDocumentationState( + issue_title="Create docs", + issue_body="Create new documentation", + issue_comments=[], + max_refined_query_loop=3, + documentation_query="", + documentation_context=[], + issue_documentation_analyzer_messages=[AIMessage(content="Create new documentation files")], + edit_messages=[], + edit_patch="", + issue_response="", + ) + + node = IssueDocumentationEditMessageNode() + result = node(state) + + assert "edit_messages" in result + assert len(result["edit_messages"]) == 1 + + +def test_extracts_last_analyzer_message(basic_state): + """Test that the node extracts the last message from analyzer history.""" + # Add multiple messages to analyzer history + basic_state["issue_documentation_analyzer_messages"] = [ + AIMessage(content="First message"), + AIMessage(content="Second message"), + AIMessage(content="Final plan: Update docs"), + ] + + node = IssueDocumentationEditMessageNode() + result = node(basic_state) + + message_content = result["edit_messages"][0].content + # Should contain the final plan + assert "Final plan" in message_content diff --git a/tests/lang_graph/nodes/test_issue_documentation_responder_node.py b/tests/lang_graph/nodes/test_issue_documentation_responder_node.py new file mode 100644 index 0000000..29b8457 --- /dev/null +++ b/tests/lang_graph/nodes/test_issue_documentation_responder_node.py @@ -0,0 +1,96 @@ +import pytest +from langchain_core.messages import AIMessage + +from prometheus.lang_graph.nodes.issue_documentation_responder_node import ( + IssueDocumentationResponderNode, +) +from prometheus.lang_graph.subgraphs.issue_documentation_state import IssueDocumentationState +from tests.test_utils.util import FakeListChatWithToolsModel + + +@pytest.fixture +def fake_llm(): + return FakeListChatWithToolsModel( + responses=[ + "The documentation has been successfully updated. " + "I've added new API endpoint documentation and included examples as requested." + ] + ) + + +@pytest.fixture +def basic_state(): + return IssueDocumentationState( + issue_title="Update API documentation", + issue_body="The API documentation needs to be updated with the new endpoints", + issue_comments=[ + {"username": "user1", "comment": "Please include examples"}, + ], + max_refined_query_loop=3, + documentation_query="Find API documentation", + documentation_context=[], + issue_documentation_analyzer_messages=[ + AIMessage(content="Plan: Update README.md with new API endpoints and add examples") + ], + edit_messages=[], + edit_patch="diff --git a/README.md b/README.md\n+New API documentation", + issue_response="", + ) + + +def test_init_issue_documentation_responder_node(fake_llm): + """Test IssueDocumentationResponderNode initialization.""" + node = IssueDocumentationResponderNode(fake_llm) + + assert node.model is not None + + +def test_call_method_basic(fake_llm, basic_state): + """Test basic call functionality.""" + node = IssueDocumentationResponderNode(fake_llm) + result = node(basic_state) + + assert "issue_response" in result + assert "successfully updated" in result["issue_response"] + assert len(result["issue_response"]) > 0 + + +def test_call_method_with_patch(fake_llm, basic_state): + """Test response generation with patch.""" + node = IssueDocumentationResponderNode(fake_llm) + result = node(basic_state) + + assert "issue_response" in result + assert isinstance(result["issue_response"], str) + + +def test_call_method_without_patch(fake_llm): + """Test response generation without patch.""" + state = IssueDocumentationState( + issue_title="Update docs", + issue_body="Please update the documentation", + issue_comments=[], + max_refined_query_loop=3, + documentation_query="", + documentation_context=[], + issue_documentation_analyzer_messages=[AIMessage(content="Documentation plan created")], + edit_messages=[], + edit_patch="", + issue_response="", + ) + + node = IssueDocumentationResponderNode(fake_llm) + result = node(state) + + assert "issue_response" in result + assert len(result["issue_response"]) > 0 + + +def test_response_includes_issue_details(fake_llm, basic_state): + """Test that the generated response is relevant to the issue.""" + node = IssueDocumentationResponderNode(fake_llm) + result = node(basic_state) + + assert "issue_response" in result + # The response should be a string with meaningful content + assert len(result["issue_response"]) > 10 diff --git a/tests/lang_graph/subgraphs/test_issue_documentation_subgraph.py b/tests/lang_graph/subgraphs/test_issue_documentation_subgraph.py new file mode 100644 index 0000000..7143364 --- /dev/null +++ b/tests/lang_graph/subgraphs/test_issue_documentation_subgraph.py @@ -0,0 +1,51 @@ +from unittest.mock import Mock + +import pytest + +from prometheus.git.git_repository import GitRepository +from prometheus.graph.knowledge_graph import KnowledgeGraph +from prometheus.lang_graph.subgraphs.issue_documentation_subgraph import ( + IssueDocumentationSubgraph, +) +from tests.test_utils.util import FakeListChatWithToolsModel + + +@pytest.fixture +def mock_kg(): + kg = Mock(spec=KnowledgeGraph) + # Configure the mock to return a list of AST node types + kg.get_all_ast_node_types.return_value = [ + "FunctionDef", + "ClassDef", + "Module", + "Import", + "Call", + ] + kg.root_node_id = 0 + return kg + + +@pytest.fixture +def mock_git_repo(): + git_repo = Mock(spec=GitRepository) + git_repo.playground_path = "mock/playground/path" + return git_repo + + +def test_issue_documentation_subgraph_basic_initialization(mock_kg, mock_git_repo): + """Test that IssueDocumentationSubgraph initializes correctly with basic components.""" + # Initialize fake model with empty responses + fake_advanced_model = FakeListChatWithToolsModel(responses=[]) + fake_base_model = FakeListChatWithToolsModel(responses=[]) + + # Initialize the subgraph with required parameters + subgraph = IssueDocumentationSubgraph( + advanced_model=fake_advanced_model, + base_model=fake_base_model, + kg=mock_kg, + git_repo=mock_git_repo, + repository_id=1, + ) + + # Verify the subgraph was created + assert subgraph.subgraph is not None