✨ [WIP] Agents now take parent task history into account

- The AnalyzerAgent no longer uses hardcoded language, source, and target, but instead gathers it from the task context

Signed-off-by: Fabian von Feilitzsch <[email protected]>

Author: fabianvf

kai/reactive_codeplanner/agent/analyzer_fix/agent.py

@@ -23,30 +23,27 @@
 @dataclass
 class _llm_response:
     reasoning: str | None
-    java_file: str | None
+    source_file: str | None
     addional_information: str | None
 
 
 class AnalyzerAgent(Agent):
-    system_message = SystemMessage(
-        content="""
-    You are an experienced java developer, who specializes in migrating code to the Quarkus Framework
+    system_message_template = Template(
+        """
+    You are an experienced {{ language }} developer, who specializes in migrating code from {{ source }} to {{ target }}
+    {{ background }}
     """
     )
 
     chat_message_template = Template(
         """
-    I will give you a JavaEE file for which I want to take one step towards migrating to Quarkus.
+    I will give you a {{ source }} file for which I want to take one step towards migrating to {{ target }}.
 
 I will provide you with static source code analysis information highlighting an issue which needs to be addressed.
 
-I will also provide you with an example of how a similar issue was solved in the past via a solved example.
-
-You can refer to the solved example for a pattern of how to update the input Java EE file to Quarkus.
-
 Fix only the problem described. Other problems will be solved in subsequent steps so it is unnecessary to handle them now.
 
-Before attempting to migrate the code to Quarkus reason through what changes are required and why.
+Before attempting to migrate the code to {{ target }} reason through what changes are required and why.
 
 Pay attention to changes you make and impacts to external dependencies in the pom.xml as well as changes to imports we need to consider.
 
@@ -83,8 +80,8 @@ class AnalyzerAgent(Agent):
 ## Reasoning
 Write the step by step reasoning in this markdown section. If you are unsure of a step or reasoning, clearly state you are unsure and why.
 
-## Updated Java File
-```java
+## Updated {{ language }} File
+```{{ language }}
 // Write the updated file for Quarkus in this section. If the file should be removed, make the content of the updated file a comment explaining it should be removed.
 ```
 
@@ -118,64 +115,79 @@ def execute(self, ask: AgentRequest) -> AnalyzerFixResponse:
 
         language = guess_language(ask.file_content, file_name)
 
+        source = " and ".join(ask.sources)
+        target = " and ".join(ask.targets)
+
+        system_message = SystemMessage(
+            content=self.system_message_template.render(
+                language=language,
+                source=source,
+                target=target,
+                background=ask.background,
+            )
+        )
+
         content = self.chat_message_template.render(
+            source=source,
+            target=target,
+            language=language,
             src_file_contents=ask.file_content,
             src_file_name=file_name,
             src_file_language=language,
             incidents=ask.incidents,
         )
 
         aimessage = self._model_provider.invoke(
-            [self.system_message, HumanMessage(content=content)]
+            [system_message, HumanMessage(content=content)]
         )
 
-        resp = self.parse_llm_response(aimessage)
+        resp = self.parse_llm_response(aimessage, language)
         return AnalyzerFixResponse(
             encountered_errors=[],
             file_to_modify=Path(os.path.abspath(ask.file_path)),
             reasoning=resp.reasoning,
             additional_information=resp.addional_information,
-            updated_file_content=resp.java_file,
+            updated_file_content=resp.source_file,
         )
 
-    def parse_llm_response(self, message: BaseMessage) -> _llm_response:
+    def parse_llm_response(self, message: BaseMessage, language: str) -> _llm_response:
         """Private method that will be used to parse the contents and get the results"""
 
         lines_of_output = cast(str, message.content).splitlines()
 
-        in_java_file = False
+        in_source_file = False
         in_reasoning = False
         in_additional_details = False
-        java_file = ""
+        source_file = ""
         reasoning = ""
         additional_details = ""
         for line in lines_of_output:
             if line.strip() == "## Reasoning":
                 in_reasoning = True
-                in_java_file = False
+                in_source_file = False
                 in_additional_details = False
                 continue
-            if line.strip() == "## Updated Java File":
-                in_java_file = True
+            if line.strip() == f"## Updated {language} File":
+                in_source_file = True
                 in_reasoning = False
                 in_additional_details = False
                 continue
             if "## Additional Information" in line.strip():
                 in_reasoning = False
-                in_java_file = False
+                in_source_file = False
                 in_additional_details = True
                 continue
-            if in_java_file:
-                if "```java" in line or "```" in line:
+            if in_source_file:
+                if f"```{language}" in line or "```" in line:
                     continue
-                java_file = "\n".join([java_file, line])
+                source_file = "\n".join([source_file, line])
             if in_reasoning:
                 reasoning = "\n".join([reasoning, line])
             if in_additional_details:
                 additional_details = "\n".join([additional_details, line])
         return _llm_response(
             reasoning=reasoning,
-            java_file=java_file,
+            source_file=source_file,
             addional_information=additional_details,
         )
 

kai/reactive_codeplanner/agent/analyzer_fix/api.py

@@ -8,6 +8,8 @@
 class AnalyzerFixRequest(AgentRequest):
     file_content: str
     incidents: list[Incident]
+    sources: list[str]
+    targets: list[str]
 
 
 @dataclass

kai/reactive_codeplanner/agent/api.py

@@ -1,11 +1,13 @@
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
 from pathlib import Path
+from typing import Optional
 
 
 @dataclass
 class AgentRequest:
     file_path: Path
+    background: Optional[str]
 
 
 @dataclass

kai/reactive_codeplanner/agent/dependency_agent/dependency_agent.py

@@ -2,6 +2,7 @@
 from pathlib import Path
 from typing import Any, Callable, Optional, TypedDict, Union
 
+from jinja2 import Template
 from langchain.prompts.chat import HumanMessagePromptTemplate
 from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
 
@@ -53,10 +54,10 @@ class _llm_response:
 
 class MavenDependencyAgent(Agent):
 
-    sys_msg = SystemMessage(
+    system_message_template = Template(
         """
 You are an excellent java developer focused on updating dependencies in a maven `pom.xml` file. 
-
+{{ background }}
 ### Guidelines:
 1  Only use the provided and predefined functions as the functions. Do not use any other functions.
 2 always search for the fqdn for the dependency to be added or updated
@@ -188,7 +189,13 @@ def execute(self, ask: AgentRequest) -> AgentResult:
         if not request.message:
             return AgentResult()
 
-        msg = [self.sys_msg, self.inst_msg_template.format(message=request.message)]
+        system_message = SystemMessage(
+            content=self.system_message_template.render(background=ask.background)
+        )
+
+        content = self.inst_msg_template.format(message=request.message)
+
+        msg = [system_message, content]
         fix_gen_attempts = 0
         llm_response: Optional[_llm_response] = None
         maven_search: Optional[FQDNResponse] = None
@@ -252,11 +259,12 @@ def execute(self, ask: AgentRequest) -> AgentResult:
                         logger.info("Need to call sub-agent for selecting FQDN")
                         r = self.child_agent.execute(
                             FQDNDependencySelectorRequest(
-                                request.file_path,
-                                request.message,
-                                a.code,
+                                file_path=request.file_path,
+                                msg=request.message,
+                                code=a.code,
                                 query=[],
                                 times=0,
+                                background=ask.background,
                             )
                         )
                         if r.response is not None and isinstance(r.response, list):

kai/reactive_codeplanner/agent/dependency_agent/dependency_fqdn_selection.py

@@ -101,6 +101,7 @@ def execute(self, ask: AgentRequest) -> FQDNDependencySelectorResult:
                     code="",
                     query=query,
                     times=ask.times + 1,
+                    background=ask.background,
                 )
             )
         if isinstance(response, list):

kai/reactive_codeplanner/agent/maven_compiler_fix/agent.py

@@ -2,16 +2,19 @@
 from langchain_core.messages import BaseMessage, HumanMessage, SystemMessage
 
 from kai.llm_interfacing.model_provider import ModelProvider
+from kai.logging.logging import get_logger
 from kai.reactive_codeplanner.agent.api import Agent, AgentRequest, AgentResult
 from kai.reactive_codeplanner.agent.maven_compiler_fix.api import (
     MavenCompilerAgentRequest,
     MavenCompilerAgentResult,
 )
 
+logger = get_logger(__name__)
+
 
 class MavenCompilerAgent(Agent):
-    system_message = SystemMessage(
-        content="""
+    system_message_template = Template(
+        """{{ background }}
     I will give you compiler errors and the offending line of code, and you will need to use the file to determine how to fix them. You should only use compiler errors to determine what to fix.
 
     Make sure that the references to any changed types are kept.
@@ -54,6 +57,10 @@ def execute(self, ask: AgentRequest) -> AgentResult:
         if not isinstance(ask, MavenCompilerAgentRequest):
             return AgentResult()
 
+        system_message = SystemMessage(
+            content=self.system_message_template.render(background=ask.background)
+        )
+
         line_of_code = ask.file_contents.split("\n")[ask.line_number]
 
         compile_errors = f"Line of code: {line_of_code};\n{ask.message}"
@@ -62,7 +69,7 @@ def execute(self, ask: AgentRequest) -> AgentResult:
         )
 
         ai_message = self.model_provider.invoke(
-            [self.system_message, HumanMessage(content=content)]
+            [system_message, HumanMessage(content=content)]
         )
 
         resp = self.parse_llm_response(ai_message)

kai/reactive_codeplanner/agent/maven_compiler_fix/api.py

@@ -38,4 +38,5 @@ def to_reflection_task(self) -> Optional[ReflectionTask]:
             reasoning=self.reasoning,
             updated_file_contents=self.updated_file_contents,
             original_file_contents=self.original_file,
+            background="",
         )

kai/reactive_codeplanner/main.py

@@ -181,6 +181,7 @@ def main() -> None:
             task_manager.supply_result(result)
         except Exception as e:
             logger.error("Failed to supply result %s: %s", result, e)
+
         supply_time = time.time() - start_supply_time
         logger.info("PERFORMANCE: %.6f seconds to supply result", supply_time)
 

kai/reactive_codeplanner/task_manager/api.py

@@ -76,6 +76,10 @@ def truncate(value: str, max_length: int = 30) -> str:
 
         return f"{class_name}<" + ", ".join(field_strings) + ">"
 
+    def background(self) -> str:
+        """Used by Agents to provide context when solving child issues"""
+        raise NotImplementedError
+
     __repr__ = __str__
 
 
@@ -130,6 +134,12 @@ def __str__(self) -> str:
 
         return f"{self.__class__.__name__}<loc={self.file}:{self.line}:{self.column}, message={self.message}>(priority={self.priority}({shadowed_priority}), depth={self.depth}, retries={self.retry_count})"
 
+    def background(self) -> str:
+        """Used by Agents to provide context when solving child issues"""
+        if self.parent is None:
+            return ""
+        return self.oldest_ancestor().background()
+
     __repr__ = __str__
 
 

kai/reactive_codeplanner/task_runner/analyzer_lsp/api.py

@@ -1,4 +1,6 @@
+import itertools
 from dataclasses import dataclass
+from functools import cached_property
 
 from kai.analyzer_types import Incident, RuleSet, Violation
 from kai.logging.logging import get_logger
@@ -29,6 +31,40 @@ def __str__(self) -> str:
 
         return f"{self.__class__.__name__}<loc={self.file}:{self.line}:{self.column}, message={self.violation.description}>(priority={self.priority}({shadowed_priority}), depth={self.depth}, retries={self.retry_count})"
 
+    def background(self) -> str:
+        if self.parent is not None:
+            return self.oldest_ancestor().background()
+        if self.children:
+            return f"""You are a software developer who specializes in migrating from {" and ".join(self.sources)} to {" and ".join(self.targets)}
+You attempted to solve an issue in a repository you are migrating:
+
+Location: {self.incident.uri}
+Message:
+{self.incident.message}
+
+However your solution caused additional problems elsewhere in the repository, which you are now going to solve."""
+        return ""
+
+    @cached_property
+    def sources(self) -> list[str]:
+        labels = set(
+            itertools.chain(*[v.labels for v in self.ruleset.violations.values()])
+        )
+        source_key = "konveyor.io/source="
+        return [
+            label.replace(source_key, "") for label in labels if source_key in label
+        ]
+
+    @cached_property
+    def targets(self) -> list[str]:
+        labels = set(
+            itertools.chain(*[v.labels for v in self.ruleset.violations.values()])
+        )
+        target_key = "konveyor.io/target="
+        return [
+            label.replace(target_key, "") for label in labels if target_key in label
+        ]
+
     __repr__ = __str__
 
     def fuzzy_equals(self, error2: Task, offset: int = 1) -> bool:

kai/reactive_codeplanner/task_runner/analyzer_lsp/task_runner.py

@@ -42,7 +42,7 @@ def can_handle_task(self, task: Task) -> bool:
         """Will determine if the task if a MavenCompilerError, and if we can handle these issues."""
         return isinstance(task, AnalyzerRuleViolation)
 
-    @tracer.start_as_current_span("analyzer_execute_task")
+    @tracer.start_as_current_span("analyzer_execute_task")  # type:ignore
     def execute_task(self, rcm: RepoContextManager, task: Task) -> TaskResult:
         """This will be responsible for getting the full file from LLM and updating the file on disk"""
 
@@ -58,6 +58,9 @@ def execute_task(self, rcm: RepoContextManager, task: Task) -> TaskResult:
             file_path=Path(os.path.abspath(task.file)),
             file_content=src_file_contents,
             incidents=[task.incident],
+            background=task.background(),
+            sources=task.sources,
+            targets=task.targets,
         )
         result = self.agent.execute(agent_request)
 

kai/reactive_codeplanner/task_runner/analyzer_lsp/validator.py

@@ -42,7 +42,7 @@ def __init__(self, config: RpcClientConfig, analyzer: AnalyzerLSP) -> None:
         self.incident_selector = config.incident_selector or ""
         super().__init__(config)
 
-    @tracer.start_as_current_span("analyzer_run_validation")
+    @tracer.start_as_current_span("analyzer_run_validation")  # type:ignore
     def run(self, scoped_paths: Optional[list[Path]] = None) -> ValidationResult:
         logger.debug("Running analyzer-lsp")