TDD: Tool-Aware Response Decorator Agent

Issue: #268
Parent: #267 - Response Decorator
Created: 2025-10-27

---

Overview

Enhance ChatResponseDecoratorAgent to be a mini-agent with function calling capabilities, enabling context-aware link generation for implicit references like "this page" and code sections.

Problem Statement

Current decorator (text-only) limitations:

Reference Type	Example	Current Behavior	Desired Behavior
Explicit page	"File 1, Page 6"	✅ Links	✅ Links
Implicit page	"this page"	❌ No link	✅ Link to current page
Sheet ID	"Sheet E1.0"	❌ No link	✅ Link to page
ICC book	"CBC" with catalog	✅ Links	✅ Links
Code section	"Section 1013.4"	❌ No link	✅ Link with chapter slug

---

Solution: Mini-Agent with Tools

Use Gemini 2.5 Flash Lite with function calling to give decorator access to:

1. GetArchitecturalPlan - Resolve sheet IDs to file/page numbers
2. GetIccBookTableOfContents - Get chapter slugs for section URLs

Per Gemini API docs:

• ✅ Function calling supported
• ✅ Optimized for cost and latency
• ✅ 1M token context window

---

Architecture

Current (Text-Only Decorator)

Raw Response
    ↓
Gemini 2.5 Flash (text prompt)
    ↓
Decorated Response

Limitations: Can only link what's explicitly stated

New (Tool-Aware Decorator)

Raw Response + Context
    ↓
Gemini 2.5 Flash Lite (with tools)
    ↓
Calls GetArchitecturalPlan → Resolves "Sheet E1.0" = File 2, Page 3
    ↓
Calls GetIccBookTableOfContents → Gets chapter-10-means-of-egress slug
    ↓
Builds complete URLs
    ↓
Decorated Response (all references linked!)

---

Tool Definitions

Tool 1: GetArchitecturalPlan

Purpose: Resolve sheet IDs and page titles to file/page numbers

Function Declaration:

FunctionDeclaration.builder()
    .name("get_architectural_plan")
    .description("Get all pages in the project with their titles, sheet IDs, file numbers, and page numbers")
    .parameters(Schema.builder()
        .type("OBJECT")
        .properties(Map.of(
            "architecturalPlanId", Schema.builder()
                .type("STRING")
                .description("Project ID from context")
                .build()
        ))
        .required(List.of("architecturalPlanId"))
        .build())
    .build()

Implementation:

public Map<String, Object> getArchitecturalPlan(Map<String, Object> args) {
  String planId = (String) args.get("architecturalPlanId");
  
  // Call existing gRPC API
  ArchitecturalPlan plan = architecturalPlanService.getArchitecturalPlan(planId);
  
  // Return simplified structure for LLM
  return Map.of("pages", plan.getPagesList().stream()
      .map(page -> Map.of(
          "fileId", extractFileId(page),
          "pageNumber", page.getPageNumber(),
          "title", page.getTitle(),
          "sheetId", page.getPageId()
      ))
      .toList());
}

Tool 2: GetIccBookTableOfContents

Purpose: Get chapter slugs and section hierarchy for building ICC URLs

Function Declaration:

FunctionDeclaration.builder()
    .name("get_icc_book_toc")
    .description("Get table of contents for ICC book to find chapter slugs for section URLs")
    .parameters(Schema.builder()
        .type("OBJECT")
        .properties(Map.of(
            "iccBookId", Schema.builder()
                .type("STRING")
                .description("ICC book ID (e.g., 3757 for CBC 2022)")
                .build()
        ))
        .required(List.of("iccBookId"))
        .build())
    .build()

Implementation:

public Map<String, Object> getIccBookToc(Map<String, Object> args) {
  String bookId = (String) args.get("iccBookId");
  
  // Call existing gRPC API
  Chapters chapters = reviewService.getIccBookTableOfContents(bookId);
  
  // Return simplified structure for LLM
  return Map.of("chapters", chapters.getChaptersList().stream()
      .map(chapter -> Map.of(
          "number", chapter.getNumber(),
          "title", chapter.getTitle(),
          "slug", chapter.getSlug(),
          "sections", chapter.getSubSectionsList().stream()
              .map(section -> Map.of(
                  "id", section.getId(),
                  "title", section.getTitle()
              ))
              .toList()
      ))
      .toList());
}

---

Implementation

Step 1: Create Decorator Tools

File: src/main/java/.../service/ChatResponseDecoratorTools.java

public class ChatResponseDecoratorTools {
  
  @Schema(description = "Get all pages in architectural plan with titles and sheet IDs")
  public static Map<String, Object> getArchitecturalPlanPages(
      @Schema(description = "Project ID") String architecturalPlanId) {
    // Implementation
  }
  
  @Schema(description = "Get ICC book table of contents with chapter slugs")
  public static Map<String, Object> getIccBookTableOfContents(
      @Schema(description = "ICC book ID like 3757") String iccBookId) {
    // Implementation
  }
}

Step 2: Update ChatResponseDecoratorAgent

Before (text-only prompt):

GoogleGenAiClient client = GoogleGenAiClient.builder()
    .setModel(Model.getLatestGeminiFlashLiteModel())
    .singleTurn()
    .build();

String decorated = client.prompt(prompt);

After (with tools):

// Create tools
FunctionTool tool1 = FunctionTool.create(ChatResponseDecoratorTools.class, "getArchitecturalPlanPages");
FunctionTool tool2 = FunctionTool.create(ChatResponseDecoratorTools.class, "getIccBookTableOfContents");

// Create agent with tools
LlmAgent decoratorAgent = LlmAgent.builder()
    .name("decorator_agent")
    .model(Model.getLatestGeminiFlashLiteModel())  // Use Flash Lite for speed/cost
    .instruction(buildDecoratorInstruction(context))
    .tools(tool1, tool2)
    .build();

// Run agent
InMemoryRunner runner = new InMemoryRunner(decoratorAgent);
Session session = runner.createSession("decorator", "decorator-user");
Content input = Content.fromParts(Part.fromText(rawResponse));
Flowable<Event> events = runner.runAsync("decorator-user", session.id(), input);

// Extract decorated response from events
String decorated = extractFinalResponse(events);

---

URL Generation Logic

Page URLs

// Format: /projects/{projectId}/files/{fileId}/pages/{pageNumber}/preview
private String buildPageUrl(String projectId, String fileId, int pageNumber) {
  return String.format("/projects/%s/files/%s/pages/%d/preview", 
      projectId, fileId, pageNumber);
}

ICC Section URLs

// Format: https://codes.iccsafe.org/content/{documentSlug}/{chapterSlug}#{sectionId}
private String buildIccSectionUrl(String documentSlug, String chapterSlug, String sectionId) {
  return String.format("https://codes.iccsafe.org/content/%s/%s#%s",
      documentSlug, chapterSlug, sectionId);
}

---

Performance

Current (text-only):

• Latency: 200-400ms
• Cost: $0.000075 per decoration

Tool-Aware:

• Latency: 800ms-1.5s (2 tool calls + LLM)
• Cost: $0.0001 per decoration (Flash Lite cheaper)
• Still acceptable for UX

Optimization: Cache TOC and plan data (1 hour TTL)

---

Testing Strategy

Unit Tests

@Test
public void testLinkImplicitPageReference() {
  String input = "Yes, there are violations on this page (Sheet E1.0).";
  ChatContext context = ChatContext.newBuilder()
      .setProjectId("san-jose-multi-file3")
      .setPageNumber(3)
      .setFileId("2")
      .build();
  
  String output = decorator.decorateResponse(input, context, catalog);
  
  assertTrue(output.contains("[this page](/projects/san-jose-multi-file3/files/2/pages/3/preview)"));
  assertTrue(output.contains("[Sheet E1.0](/projects/san-jose-multi-file3/files/2/pages/3/preview)"));
}

@Test
public void testLinkCodeSection() {
  String input = "Section 1013.4 requires tactile exit signs.";
  
  String output = decorator.decorateResponse(input, context, catalog);
  
  assertTrue(output.contains("](https://codes.iccsafe.org/content/CABC2022P4/chapter-10-means-of-egress#CABC2022P4_Ch10_Sec1013"));
}

---

Success Metrics

• ✅ 95%+ of implicit page references linked
• ✅ 90%+ of code sections linked correctly
• ✅ P95 latency <1.5 seconds
• ✅ Cost <$0.0002 per decoration
• ✅ No regressions on explicit references

---

Priority: Medium
Complexity: Medium
Estimated Effort: 8 hours
Model: Gemini 2.5 Flash Lite