RPP: store pure function calls as facts

This CL implements some basic caching for the pure functions that we
expose to the AI. This avoids them continually recalling the same
functions which is inefficient and also a poor user experience to see
the AI continually repeat the exact same calls.

The data is cached based on the insight, ensuring that if the insight
changes, we invalidate the cache.

We could do this better by caching based on the bounds of the insight
(because if the bounds are the same, we don't need to recall the
function, regardless of which insight it is) but that will be harder for
the AI to understand, whereas this works reliably based on my testing.

Bug: 408172181
Change-Id: I243becee05140a2c4ca148cb3aed25936285f29d
Reviewed-on: https://chromium-review.googlesource.com/c/devtools/devtools-frontend/+/6439198
Reviewed-by: Alex Rudenko <[email protected]>
Commit-Queue: Jack Franklin <[email protected]>
Auto-Submit: Jack Franklin <[email protected]>
Commit-Queue: Alex Rudenko <[email protected]>

Author: jackfranklin

front_end/models/ai_assistance/agents/PerformanceInsightsAgent.test.ts

@@ -312,5 +312,131 @@ Help me understand?`;
         canceled: false
       });
     });
+
+    it('caches getNetworkActivitySummary calls and passes them to future requests as facts', async function() {
+      const {parsedTrace, insights} = await TraceLoader.traceEngine(this, 'lcp-images.json.gz');
+      assert.isOk(insights);
+      const [firstNav] = parsedTrace.Meta.mainFrameNavigations;
+      const lcpPhases = getInsightOrError('LCPPhases', insights, firstNav);
+      const agent = new PerformanceInsightsAgent({
+        aidaClient: mockAidaClient([
+          [{explanation: '', functionCalls: [{name: 'getNetworkActivitySummary', args: {}}]}], [{explanation: 'done'}]
+        ])
+      });
+      const activeInsight = new TimelineUtils.InsightAIContext.ActiveInsight(lcpPhases, parsedTrace);
+      const context = new InsightContext(activeInsight);
+
+      // Make the first query to trigger the getNetworkActivitySummary function
+      const responses = await Array.fromAsync(agent.run('test', {selected: context}));
+      const action = responses.find(response => response.type === ResponseType.ACTION);
+      assert.exists(action);
+      assert.strictEqual(action.code, 'getNetworkActivitySummary()');
+
+      // Trigger another request so that the agent populates the facts.
+      await Array.fromAsync(agent.run('test 2', {selected: context}));
+
+      assert.strictEqual(agent.currentFacts().size, 1);
+      const networkSummaryFact = Array.from(agent.currentFacts()).at(0);
+      assert.exists(networkSummaryFact);
+
+      const expectedRequestUrls = [
+        'https://chromedevtools.github.io/performance-stories/lcp-large-image/index.html',
+        'https://fonts.googleapis.com/css2?family=Poppins:ital,wght@1,800',
+        'https://chromedevtools.github.io/performance-stories/lcp-large-image/app.css',
+        'https://via.placeholder.com/50.jpg', 'https://via.placeholder.com/2000.jpg'
+      ];
+      // Ensure that each URL was in the fact as a way to validate the fact is accurate.
+      assert.isTrue(expectedRequestUrls.every(url => {
+        return networkSummaryFact.text.includes(url);
+      }));
+
+      // Now we make one more request; we do this to ensure that we don't add the same fact again.
+      await Array.fromAsync(agent.run('test 3', {selected: context}));
+
+      assert.strictEqual(agent.currentFacts().size, 1);
+    });
+
+    it('caches getMainThreadActivity calls and passes them to future requests as facts', async function() {
+      const {parsedTrace, insights} = await TraceLoader.traceEngine(this, 'lcp-discovery-delay.json.gz');
+      assert.isOk(insights);
+      const [firstNav] = parsedTrace.Meta.mainFrameNavigations;
+      const lcpPhases = getInsightOrError('LCPPhases', insights, firstNav);
+      const agent = new PerformanceInsightsAgent({
+        aidaClient: mockAidaClient(
+            [[{explanation: '', functionCalls: [{name: 'getMainThreadActivity', args: {}}]}], [{explanation: 'done'}]])
+      });
+      const activeInsight = new TimelineUtils.InsightAIContext.ActiveInsight(lcpPhases, parsedTrace);
+      const context = new InsightContext(activeInsight);
+
+      // Make the first query to trigger the getMainThreadActivity function
+      const responses = await Array.fromAsync(agent.run('test', {selected: context}));
+      const action = responses.find(response => response.type === ResponseType.ACTION);
+      assert.exists(action);
+      assert.strictEqual(action.code, 'getMainThreadActivity()');
+
+      // Trigger another request so that the agent populates the facts.
+      await Array.fromAsync(agent.run('test 2', {selected: context}));
+
+      assert.strictEqual(agent.currentFacts().size, 1);
+      const mainThreadActivityFact = Array.from(agent.currentFacts()).at(0);
+      assert.exists(mainThreadActivityFact);
+
+      const expectedTree = TimelineUtils.InsightAIContext.AIQueries.mainThreadActivity(lcpPhases, parsedTrace);
+      assert.isOk(expectedTree);
+      assert.include(mainThreadActivityFact.text, expectedTree.serialize());
+
+      // Now we make one more request; we do this to ensure that we don't add the same fact again.
+      await Array.fromAsync(agent.run('test 3', {selected: context}));
+
+      assert.strictEqual(agent.currentFacts().size, 1);
+    });
+
+    it('will not send facts from a previous insight if the context changes', async function() {
+      const {parsedTrace, insights} = await TraceLoader.traceEngine(this, 'lcp-discovery-delay.json.gz');
+      assert.isOk(insights);
+      const [firstNav] = parsedTrace.Meta.mainFrameNavigations;
+      const lcpPhases = getInsightOrError('LCPPhases', insights, firstNav);
+      const renderBlocking = getInsightOrError('RenderBlocking', insights, firstNav);
+      const agent = new PerformanceInsightsAgent({
+        aidaClient: mockAidaClient([
+          [{explanation: '', functionCalls: [{name: 'getMainThreadActivity', args: {}}]}],
+        ])
+      });
+      const lcpPhasesActiveInsight = new TimelineUtils.InsightAIContext.ActiveInsight(lcpPhases, parsedTrace);
+      const lcpContext = new InsightContext(lcpPhasesActiveInsight);
+      const renderBlockingActiveInsight = new TimelineUtils.InsightAIContext.ActiveInsight(renderBlocking, parsedTrace);
+      const renderBlockingContext = new InsightContext(renderBlockingActiveInsight);
+
+      // Populate the function calls for the LCP Context
+      await Array.fromAsync(agent.run('test 1 LCP', {selected: lcpContext}));
+      await Array.fromAsync(agent.run('test 2 LCP', {selected: lcpContext}));
+      assert.strictEqual(agent.currentFacts().size, 1);
+      // Now change the context and send a request.
+      await Array.fromAsync(agent.run('test 1 RenderBlocking', {selected: renderBlockingContext}));
+      // Because the context changed, we should now not have any facts.
+      assert.strictEqual(agent.currentFacts().size, 0);
+    });
+
+    it('will send multiple facts', async function() {
+      const {parsedTrace, insights} = await TraceLoader.traceEngine(this, 'lcp-discovery-delay.json.gz');
+      assert.isOk(insights);
+      const [firstNav] = parsedTrace.Meta.mainFrameNavigations;
+      const lcpPhases = getInsightOrError('LCPPhases', insights, firstNav);
+      const agent = new PerformanceInsightsAgent({
+        aidaClient: mockAidaClient([
+          [{explanation: '', functionCalls: [{name: 'getMainThreadActivity', args: {}}]}],
+          [{explanation: '', functionCalls: [{name: 'getNetworkActivitySummary', args: {}}]}], [{explanation: 'done'}]
+        ])
+      });
+      const activeInsight = new TimelineUtils.InsightAIContext.ActiveInsight(lcpPhases, parsedTrace);
+      const context = new InsightContext(activeInsight);
+      // First query to populate the function calls
+      await Array.fromAsync(agent.run('test 1', {selected: context}));
+      // Second query should have two facts
+      await Array.fromAsync(agent.run('test 2', {selected: context}));
+      assert.deepEqual(Array.from(agent.currentFacts(), fact => {
+        return fact.metadata.source;
+      }), ['getMainThreadActivity()', 'getNetworkActivitySummary()']);
+    });
   });
 });

front_end/models/ai_assistance/agents/PerformanceInsightsAgent.ts

@@ -196,6 +196,23 @@ export class PerformanceInsightsAgent extends AiAgent<TimelineUtils.InsightAICon
 
   #lastContextForEnhancedQuery: ConversationContext<TimelineUtils.InsightAIContext.ActiveInsight>|undefined;
 
+  /**
+   * Store results (as facts) for the functions that are pure and return the
+   * same data for the same insight.
+   * This fact is then passed into the request on all future
+   * queries for the conversation. This means that the LLM is far less likely to
+   * call the function again, because we have provided the same data as a
+   * fact. We cache based on the active insight to ensure that if the user
+   * changes which insight they are focusing we will call the function again.
+   * It's important that we store it as a Fact in the cache, because the AI
+   * Agent stores facts in a set, and we need to pass the same object through to
+   * make sure it isn't mistakenly duplicated in the request.
+   */
+  #functionCallCache = new Map<TimelineUtils.InsightAIContext.ActiveInsight, {
+    getNetworkActivitySummary?: Host.AidaClient.RequestFact,
+    getMainThreadActivity?: Host.AidaClient.RequestFact,
+  }>();
+
   override async *
       handleContextDetails(activeContext: ConversationContext<TimelineUtils.InsightAIContext.ActiveInsight>|null):
           AsyncGenerator<ContextResponse, void, void> {
@@ -263,6 +280,15 @@ export class PerformanceInsightsAgent extends AiAgent<TimelineUtils.InsightAICon
         );
         const formatted =
             requests.map(r => TraceEventFormatter.networkRequest(r, activeInsight.parsedTrace, {verbose: false}));
+        const summaryFact: Host.AidaClient.RequestFact = {
+          text:
+              `This is the network summary for this insight. You can use this and not call getNetworkActivitySummary again:\n${
+                  formatted.join('\n')}`,
+          metadata: {source: 'getNetworkActivitySummary()'}
+        };
+        const cacheForInsight = this.#functionCallCache.get(activeInsight) ?? {};
+        cacheForInsight.getNetworkActivitySummary = summaryFact;
+        this.#functionCallCache.set(activeInsight, cacheForInsight);
         return {result: {requests: formatted}};
       },
     });
@@ -349,7 +375,18 @@ The fields are:
         if (!tree) {
           return {error: 'No main thread activity found'};
         }
-        return {result: {activity: tree.serialize()}};
+        const activity = tree.serialize();
+        const activityFact: Host.AidaClient.RequestFact = {
+          text:
+              `This is the main thread activity for this insight. You can use this and not call getMainThreadActivity again:\n${
+                  activity}`,
+          metadata: {source: 'getMainThreadActivity()'},
+        };
+        const cacheForInsight = this.#functionCallCache.get(activeInsight) ?? {};
+        cacheForInsight.getMainThreadActivity = activityFact;
+        this.#functionCallCache.set(activeInsight, cacheForInsight);
+
+        return {result: {activity}};
       },
 
     });
@@ -398,6 +435,15 @@ The fields are:
   }): AsyncGenerator<ResponseData, void, void> {
     this.#insight = options.selected ?? undefined;
 
+    // Clear any previous facts in case the user changed the active context.
+    this.clearFacts();
+    const cachedFunctionCalls = this.#insight ? this.#functionCallCache.get(this.#insight.getItem()) : null;
+    if (cachedFunctionCalls) {
+      for (const fact of Object.values(cachedFunctionCalls)) {
+        this.addFact(fact);
+      }
+    }
+
     return yield* super.run(initialQuery, options);
   }
 }