Dev

example.env

@@ -17,6 +17,7 @@ PROMETHEUS_MAX_TOKEN_PER_NEO4J_RESULT=10000
 # LLM API keys and model settings
 PROMETHEUS_ADVANCED_MODEL=gpt-4o
 PROMETHEUS_BASE_MODEL=gpt-4o
+PROMETHEUS_MAX_TOKENS=64000
 PROMETHEUS_ANTHROPIC_API_KEY=anthropic_api_key
 PROMETHEUS_GEMINI_API_KEY=gemini_api_key
 PROMETHEUS_OPENAI_API_KEY=openai_api_key

prometheus/app/main.py

@@ -26,6 +26,7 @@
 logger.info(f"KNOWLEDGE_GRAPH_CHUNK_SIZE={settings.KNOWLEDGE_GRAPH_CHUNK_SIZE}")
 logger.info(f"KNOWLEDGE_GRAPH_CHUNK_OVERLAP={settings.KNOWLEDGE_GRAPH_CHUNK_OVERLAP}")
 logger.info(f"MAX_TOKEN_PER_NEO4J_RESULT={settings.MAX_TOKEN_PER_NEO4J_RESULT}")
+logger.info(f"MAX_TOKENS={settings.MAX_TOKENS}")
 
 
 @asynccontextmanager

prometheus/app/services/issue_service.py

@@ -40,6 +40,31 @@ def answer_issue(
         build_commands: Optional[Sequence[str]] = None,
         test_commands: Optional[Sequence[str]] = None,
     ):
+        """
+        Processes an issue, generates patches if needed, runs optional builds and tests, and returning the results.
+
+        Args:
+            issue_title (str): The title of the issue.
+            issue_body (str): The body of the issue.
+            issue_comments (Sequence[Mapping[str, str]]): Comments on the issue.
+            issue_type (IssueType): The type of the issue (BUG or QUESTION).
+            run_build (bool): Whether to run the build commands.
+            run_existing_test (bool): Whether to run existing tests.
+            number_of_candidate_patch (int): Number of candidate patches to generate.
+            dockerfile_content (Optional[str]): Content of the Dockerfile for user-defined environments.
+            image_name (Optional[str]): Name of the Docker image.
+            workdir (Optional[str]): Working directory for the container.
+            build_commands (Optional[Sequence[str]]): Commands to build the project.
+            test_commands (Optional[Sequence[str]]): Commands to test the project.
+        Returns:
+            Tuple containing:
+                - edit_patch (str): The generated patch for the issue.
+                - passed_reproducing_test (bool): Whether the reproducing test passed.
+                - passed_build (bool): Whether the build passed.
+                - passed_existing_test (bool): Whether the existing tests passed.
+                - issue_response (str): Response generated for the issue.
+        """
+        # Construct the working directory
         if dockerfile_content or image_name:
             container = UserDefinedContainer(
                 self.kg_service.kg.get_local_path(),
@@ -51,7 +76,7 @@ def answer_issue(
             )
         else:
             container = GeneralContainer(self.kg_service.kg.get_local_path())
-
+        # Initialize the issue graph with the necessary services and parameters
         issue_graph = IssueGraph(
             advanced_model=self.llm_service.advanced_model,
             base_model=self.llm_service.base_model,
@@ -63,7 +88,7 @@ def answer_issue(
             build_commands=build_commands,
             test_commands=test_commands,
         )
-
+        # Invoke the issue graph with the provided parameters
         output_state = issue_graph.invoke(
             issue_title,
             issue_body,
@@ -84,11 +109,13 @@ def answer_issue(
             )
         elif output_state["issue_type"] == IssueType.QUESTION:
             return (
-                "",
+                None,
                 False,
                 False,
                 False,
                 output_state["issue_response"],
             )
 
-        return "", False, False, False, ""
+        raise ValueError(
+            f"Unknown issue type: {output_state['issue_type']}. Expected BUG or QUESTION."
+        )

prometheus/app/services/service_coordinator.py

@@ -90,6 +90,34 @@ def answer_issue(
         test_commands: Optional[Sequence[str]] = None,
         push_to_remote: Optional[bool] = None,
     ):
+        """
+        Processes an issue, generates patches if needed, runs optional builds and tests,
+        and can push changes to a remote branch.
+
+        Args:
+            issue_number: The issue number to process.
+            issue_title: Title of the issue.
+            issue_body: Body of the issue.
+            issue_comments: Comments on the issue.
+            issue_type: Type of the issue (e.g., bug, feature).
+            run_build: Whether to run a build after applying the patch.
+            run_existing_test: Whether to run existing tests after applying the patch.
+            number_of_candidate_patch: Number of candidate patches to generate.
+            dockerfile_content: Optional Dockerfile content for user-defined environment.
+            image_name: Optional name for the Docker image.
+            workdir: Working directory for the container.
+            build_commands: Commands to build the project.
+            test_commands: Commands to test the project.
+            push_to_remote: Whether to push changes to a remote branch.
+        Returns:
+            A tuple containing:
+                - remote_branch_name: Name of the remote branch if changes were pushed.
+                - patch: The generated patch for the issue.
+                - passed_reproducing_test: Whether the reproducing test passed.
+                - passed_build: Whether the build passed.
+                - passed_existing_test: Whether existing tests passed.
+                - issue_response: Response from the issue service after processing.
+        """
         logger = logging.getLogger("prometheus")
         formatter = logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s")
         timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
@@ -99,6 +127,7 @@ def answer_issue(
         logger.addHandler(file_handler)
 
         try:
+            # fix issue
             patch, passed_reproducing_test, passed_build, passed_existing_test, issue_response = (
                 self.issue_service.answer_issue(
                     issue_title=issue_title,
@@ -115,7 +144,7 @@ def answer_issue(
                     test_commands=test_commands,
                 )
             )
-
+            # push to remote if requested
             remote_branch_name = None
             if patch and push_to_remote:
                 remote_branch_name = self.repository_service.push_change_to_remote(

prometheus/lang_graph/graphs/issue_graph.py

@@ -16,6 +16,13 @@
 
 
 class IssueGraph:
+    """
+    A LangGraph-based workflow to handle and triage GitHub issues with LLM assistance.
+    Attributes:
+        git_repo (GitRepository): The Git repository to work with.
+        graph (StateGraph): The state graph representing the issue handling workflow.
+    """
+
     def __init__(
         self,
         advanced_model: BaseChatModel,
@@ -30,7 +37,9 @@ def __init__(
     ):
         self.git_repo = git_repo
 
+        # Entrance point for the issue handling workflow
         issue_type_branch_node = NoopNode()
+        # Subgraph nodes for issue classification and bug handling
         issue_classification_subgraph_node = IssueClassificationSubgraphNode(
             model=base_model,
             kg=kg,
@@ -48,14 +57,16 @@ def __init__(
             build_commands=build_commands,
             test_commands=test_commands,
         )
-
+        # 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)
-
+        # Set the entry point for the workflow
         workflow.set_entry_point("issue_type_branch_node")
+        # Define the edges and conditions for the workflow
+        # Classify the issue type if not provided
         workflow.add_conditional_edges(
             "issue_type_branch_node",
             lambda state: state["issue_type"],
@@ -67,6 +78,7 @@ def __init__(
                 IssueType.QUESTION: END,
             },
         )
+        # Add edges for the issue classification subgraph
         workflow.add_conditional_edges(
             "issue_classification_subgraph_node",
             lambda state: state["issue_type"],
@@ -77,6 +89,7 @@ def __init__(
                 IssueType.QUESTION: END,
             },
         )
+        # Add edges for ending the workflow
         workflow.add_edge("issue_bug_subgraph_node", END)
 
         self.graph = workflow.compile()
@@ -91,6 +104,9 @@ def invoke(
         run_existing_test: bool,
         number_of_candidate_patch: int,
     ):
+        """
+        Invoke the issue handling workflow with the provided parameters.
+        """
         config = None
 
         input_state = {
@@ -105,6 +121,7 @@ def invoke(
 
         output_state = self.graph.invoke(input_state, config)
 
+        # Reset the git repository to its original state
         self.git_repo.reset_repository()
 
         return output_state

prometheus/lang_graph/nodes/context_provider_node.py

@@ -120,14 +120,18 @@ def __init__(
         self._logger = logging.getLogger("prometheus.lang_graph.nodes.context_provider_node")
 
     def _init_tools(self):
-        """Initializes KnowledgeGraph traversal tools.
-
+        """
+        Initializes KnowledgeGraph traversal tools.
 
         Returns:
           List of StructuredTool instances configured for KnowledgeGraph traversal.
         """
         tools = []
 
+        # === FILE SEARCH TOOLS ===
+
+        # Tool: Find file node by filename (basename)
+        # Used when only the filename (not full path) is known
         find_file_node_with_basename_fn = functools.partial(
             graph_traversal.find_file_node_with_basename,
             driver=self.neo4j_driver,
@@ -142,6 +146,8 @@ def _init_tools(self):
         )
         tools.append(find_file_node_with_basename_tool)
 
+        # Tool: Find file node by relative path
+        # Preferred method when the exact file path is known
         find_file_node_with_relative_path_fn = functools.partial(
             graph_traversal.find_file_node_with_relative_path,
             driver=self.neo4j_driver,
@@ -156,6 +162,10 @@ def _init_tools(self):
         )
         tools.append(find_file_node_with_relative_path_tool)
 
+        # === AST NODE SEARCH TOOLS ===
+
+        # Tool: Find AST node by text match in file (by basename)
+        # Useful for searching specific snippets or patterns in unknown locations
         find_ast_node_with_text_in_file_with_basename_fn = functools.partial(
             graph_traversal.find_ast_node_with_text_in_file_with_basename,
             driver=self.neo4j_driver,
@@ -170,6 +180,7 @@ def _init_tools(self):
         )
         tools.append(find_ast_node_with_text_in_file_with_basename_tool)
 
+        # Tool: Find AST node by text match in file (by relative path)
         find_ast_node_with_text_in_file_with_relative_path_fn = functools.partial(
             graph_traversal.find_ast_node_with_text_in_file_with_relative_path,
             driver=self.neo4j_driver,
@@ -184,6 +195,8 @@ def _init_tools(self):
         )
         tools.append(find_ast_node_with_text_in_file_with_relative_path_tool)
 
+        # Tool: Find AST node by type in file (by basename)
+        # Example types: FunctionDef, ClassDef, Assign, etc.
         find_ast_node_with_type_in_file_with_basename_fn = functools.partial(
             graph_traversal.find_ast_node_with_type_in_file_with_basename,
             driver=self.neo4j_driver,
@@ -198,6 +211,7 @@ def _init_tools(self):
         )
         tools.append(find_ast_node_with_type_in_file_with_basename_tool)
 
+        # Tool: Find AST node by type in file (by relative path)
         find_ast_node_with_type_in_file_with_relative_path_fn = functools.partial(
             graph_traversal.find_ast_node_with_type_in_file_with_relative_path,
             driver=self.neo4j_driver,
@@ -212,6 +226,9 @@ def _init_tools(self):
         )
         tools.append(find_ast_node_with_type_in_file_with_relative_path_tool)
 
+        # === TEXT/DOCUMENT SEARCH TOOLS ===
+
+        # Tool: Find text node globally by keyword
         find_text_node_with_text_fn = functools.partial(
             graph_traversal.find_text_node_with_text,
             driver=self.neo4j_driver,
@@ -226,6 +243,7 @@ def _init_tools(self):
         )
         tools.append(find_text_node_with_text_tool)
 
+        # Tool: Find text node by keyword in specific file
         find_text_node_with_text_in_file_fn = functools.partial(
             graph_traversal.find_text_node_with_text_in_file,
             driver=self.neo4j_driver,
@@ -240,6 +258,7 @@ def _init_tools(self):
         )
         tools.append(find_text_node_with_text_in_file_tool)
 
+        # Tool: Fetch the next text node chunk in a chain (used for long docs/comments)
         get_next_text_node_with_node_id_fn = functools.partial(
             graph_traversal.get_next_text_node_with_node_id,
             driver=self.neo4j_driver,
@@ -254,6 +273,9 @@ def _init_tools(self):
         )
         tools.append(get_next_text_node_with_node_id_tool)
 
+        # === FILE PREVIEW & READING TOOLS ===
+
+        # Tool: Preview contents of file by basename
         preview_file_content_with_basename_fn = functools.partial(
             graph_traversal.preview_file_content_with_basename,
             driver=self.neo4j_driver,
@@ -268,6 +290,7 @@ def _init_tools(self):
         )
         tools.append(preview_file_content_with_basename_tool)
 
+        # Tool: Preview contents of file by relative path
         preview_file_content_with_relative_path_fn = functools.partial(
             graph_traversal.preview_file_content_with_relative_path,
             driver=self.neo4j_driver,
@@ -282,6 +305,7 @@ def _init_tools(self):
         )
         tools.append(preview_file_content_with_relative_path_tool)
 
+        # Tool: Read entire code file by basename
         read_code_with_basename_fn = functools.partial(
             graph_traversal.read_code_with_basename,
             driver=self.neo4j_driver,
@@ -296,6 +320,7 @@ def _init_tools(self):
         )
         tools.append(read_code_with_basename_tool)
 
+        # Tool: Read entire code file by relative path
         read_code_with_relative_path_fn = functools.partial(
             graph_traversal.read_code_with_relative_path,
             driver=self.neo4j_driver,
@@ -316,13 +341,15 @@ def __call__(self, state: Dict):
         """Processes the current state and traverse the knowledge graph to retrieve context.
 
         Args:
-          state: Current state containing the human query and preivous context_messages.
+          state: Current state containing the human query and previous context_messages.
 
         Returns:
           Dictionary that will update the state with the model's response messages.
         """
+        self._logger.debug(f"Context provider messages: {state['context_provider_messages']}")
         message_history = [self.system_prompt] + state["context_provider_messages"]
         truncated_message_history = truncate_messages(message_history)
         response = self.model_with_tools.invoke(truncated_message_history)
         self._logger.debug(response)
+        # The response will be added to the bottom of the list
         return {"context_provider_messages": [response]}

prometheus/lang_graph/nodes/context_query_message_node.py

@@ -12,4 +12,5 @@ def __init__(self):
     def __call__(self, state: ContextRetrievalState):
         human_message = HumanMessage(state["query"])
         self._logger.debug(f"Sending query to ContextProviderNode:\n{human_message}")
+        # The message will be added to the end of the context provider messages
         return {"context_provider_messages": [human_message]}

prometheus/lang_graph/nodes/context_retrieval_subgraph_node.py

@@ -21,7 +21,7 @@ def __init__(
         self._logger = logging.getLogger(
             "prometheus.lang_graph.nodes.context_retrieval_subgraph_node"
         )
-        self.context_retrevial_subgraph = ContextRetrievalSubgraph(
+        self.context_retrieval_subgraph = ContextRetrievalSubgraph(
             model=model,
             kg=kg,
             neo4j_driver=neo4j_driver,
@@ -32,8 +32,9 @@ def __init__(
 
     def __call__(self, state: Dict):
         self._logger.info("Enter context retrieval subgraph")
-        output_state = self.context_retrevial_subgraph.invoke(
+        output_state = self.context_retrieval_subgraph.invoke(
             state[self.query_key_name], state["max_refined_query_loop"]
         )
 
+        self._logger.info(f"Context retrieved: {output_state['context']}")
         return {self.context_key_name: output_state["context"]}

prometheus/lang_graph/nodes/context_selection_node.py

@@ -96,6 +96,7 @@ def format_human_prompt(self, state: ContextRetrievalState, search_result: str)
         return context_info
 
     def __call__(self, state: ContextRetrievalState):
+        self._logger.info("Starting context selection process")
         context_list = state.get("context", [])
         for tool_message in extract_last_tool_messages(state["context_provider_messages"]):
             for search_result in neo4j_data_for_context_generator(tool_message.artifact):