Adding a dynamic RAG system to handle build errors

agent/base_agent.py

@@ -31,6 +31,51 @@
 
 import logger
 import utils
+
+def patch_gemini_thought_signature():
+    try:
+        from google.adk.models.google_llm import Gemini
+        from google.genai import types
+
+        orig_preprocess = Gemini._preprocess_request
+        async def patched_preprocess(self, llm_request):
+            await orig_preprocess(self, llm_request)
+            if getattr(llm_request, 'contents', None) is None:
+                return
+
+            needs_flattening = False
+            for content in llm_request.contents:
+                if getattr(content, 'parts', None) is None: 
+                    continue
+                if any(getattr(p, 'function_call', None) and not getattr(p, 'thought_signature', None) for p in content.parts):
+                    needs_flattening = True
+                    break
+
+            if needs_flattening:
+                for content in llm_request.contents:
+                    if getattr(content, 'parts', None) is None: 
+                        continue
+                    new_parts = []
+                    for p in content.parts:
+                        if getattr(p, 'function_call', None):
+                            new_parts.append(types.Part(
+                                text=f"Model called tool `{p.function_call.name}` with parameters: {p.function_call.args}"
+                            ))
+                        elif getattr(p, 'function_response', None):
+                            new_parts.append(types.Part(
+                                text=f"Tool `{p.function_response.name}` returned result: {p.function_response.response}"
+                            ))
+                        else:
+                            new_parts.append(p)
+                    content.parts = new_parts
+
+        Gemini._preprocess_request = patched_preprocess
+    except Exception as e:
+        import logging
+        logging.getLogger(__name__).warning("Failed to patch Gemini: %s", e)
+
+patch_gemini_thought_signature()
+
 from data_prep import introspector
 from experiment import benchmark as benchmarklib
 from llm_toolkit.models import LLM, VertexAIModel
@@ -69,6 +114,7 @@ def get_tool(self, tool_name: str) -> Optional[BaseTool]:
   def chat_llm_with_tools(self, client: Any, prompt: Optional[Prompt], tools,
                           trial) -> Any:
     """Chat with LLM with tools."""
+    self._jitter_for_rate_limit(trial=trial)
     logger.info(
         '<CHAT WITH TOOLS PROMPT:ROUND %02d>%s</CHAT PROMPT:ROUND %02d>',
         trial,
@@ -89,6 +135,7 @@ def chat_llm_with_tools(self, client: Any, prompt: Optional[Prompt], tools,
   def chat_llm(self, cur_round: int, client: Any, prompt: Prompt,
                trial: int) -> str:
     """Chat with LLM."""
+    self._jitter_for_rate_limit(trial=trial)
     logger.info('<CHAT PROMPT:ROUND %02d>%s</CHAT PROMPT:ROUND %02d>',
                 cur_round,
                 prompt.gettext(),
@@ -104,6 +151,7 @@ def chat_llm(self, cur_round: int, client: Any, prompt: Prompt,
 
   def ask_llm(self, cur_round: int, prompt: Prompt, trial: int) -> str:
     """Ask LLM."""
+    self._jitter_for_rate_limit(trial=trial)
     logger.info('<ASK PROMPT:ROUND %02d>%s</ASK PROMPT:ROUND %02d>',
                 cur_round,
                 prompt.gettext(),
@@ -201,6 +249,17 @@ def _container_handle_bash_commands(self, response: str, tool: BaseTool,
       prompt.append(prompt_text)
     return prompt
 
+  def _jitter_for_rate_limit(
+      self,
+      trial: int,
+      min_sec: float = 1.0,
+      max_sec: float = 10.0,
+  ) -> None:
+    """Sleeps for a very short random duration to avoid 429 collisions."""
+    duration = random.uniform(min_sec, max_sec)
+    logger.debug('Jittering for %f before query', duration, trial=trial)
+    time.sleep(duration)
+
   def _sleep_random_duration(
       self,
       trial: int,
@@ -388,6 +447,7 @@ def get_xml_representation(self, response: Optional[dict]) -> str:
   def chat_llm(self, cur_round: int, client: Any, prompt: Prompt,
                trial: int) -> Any:
     """Call the agent with the given prompt, running async code in sync."""
+    self._jitter_for_rate_limit(trial=trial)
 
     self.round = cur_round
 

agent/memory_enhancer.py

@@ -0,0 +1,94 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""An LLM agent to improve a fuzz target's runtime performance.
+Use it as a usual module locally, or as script in cloud builds.
+"""
+import logger
+from agent.memory_prototyper import MemoryPrototyper
+from llm_toolkit.prompt_builder import (CoverageEnhancerTemplateBuilder,
+                                        CrashEnhancerTemplateBuilder,
+                                        EnhancerTemplateBuilder,
+                                        JvmFixingBuilder)
+from llm_toolkit.prompts import Prompt, TextPrompt
+from results import AnalysisResult, BuildResult, Result
+
+
+class MemoryEnhancer(MemoryPrototyper):
+  """The Agent to refine a compilable fuzz target for higher coverage."""
+
+  def _initial_prompt(self, results: list[Result]) -> Prompt:
+    """Constructs initial prompt of the agent."""
+    last_result = results[-1]
+    benchmark = last_result.benchmark
+
+    if not isinstance(last_result, AnalysisResult):
+      logger.error('The last result in Enhancer is not AnalysisResult: %s',
+                   results,
+                   trial=self.trial)
+      return Prompt()
+
+    last_build_result = None
+    for result in results[::-1]:
+      if isinstance(result, BuildResult):
+        last_build_result = result
+        break
+    if not last_build_result:
+      logger.error('Unable to find the last build result in Enhancer : %s',
+                   results,
+                   trial=self.trial)
+      return Prompt()
+
+    function_requirements = self.get_function_requirements()
+
+    if benchmark.language == 'jvm':
+      # TODO: Do this in a separate agent for JVM coverage.
+      builder = JvmFixingBuilder(self.llm, benchmark,
+                                 last_result.run_result.fuzz_target_source, [])
+      prompt = builder.build([], None, None)
+    else:
+      # TODO(dongge): Refine this logic.
+      if last_result.semantic_result:
+        error_desc, errors = last_result.semantic_result.get_error_info()
+        builder = EnhancerTemplateBuilder(self.llm, benchmark,
+                                          last_build_result, error_desc, errors)
+      elif last_result.crash_result:
+        crash_result = last_result.crash_result
+        context_result = last_result.crash_context_result
+        builder = CrashEnhancerTemplateBuilder(self.llm, benchmark,
+                                               last_build_result, crash_result,
+                                               context_result)
+      elif last_result.coverage_result:
+        builder = CoverageEnhancerTemplateBuilder(
+            self.llm,
+            benchmark,
+            last_build_result,
+            coverage_result=last_result.coverage_result)
+      else:
+        logger.error(
+            'Last result does not contain either semantic result or '
+            'coverage result',
+            trial=self.trial)
+        # TODO(dongge): Give some default initial prompt.
+        prompt = TextPrompt(
+            'Last result does not contain either semantic result or '
+            'coverage result')
+        return prompt
+      prompt = builder.build(example_pair=[],
+                             tool_guides=self.inspect_tool.tutorial(),
+                             project_dir=self.inspect_tool.project_dir,
+                             function_requirements=function_requirements)
+      # TODO: A different file name/dir.
+      prompt.save(self.args.work_dirs.prompt)
+
+    return prompt