feat(ai): Improve agent history recording (#14378)

* Simplify recording for chat requests/responses by providing a util
* Introduce explicit optional system message field
* Show messages and system message in history view as a context
* A few minor typos and nitpicks
* Guard undefined agent description during load

packages/ai-chat/src/common/chat-agents.ts

@@ -52,6 +52,7 @@ import {
 } from './chat-model';
 import { findFirstMatch, parseContents } from './parse-contents';
 import { DefaultResponseContentFactory, ResponseContentMatcher, ResponseContentMatcherProvider } from './response-content-matcher';
+import { ChatHistoryEntry } from './chat-history-entry';
 
 /**
  * A conversation consists of a sequence of ChatMessages.
@@ -152,19 +153,19 @@ export abstract class AbstractChatAgent {
             if (!languageModel) {
                 throw new Error('Couldn\'t find a matching language model. Please check your setup!');
             }
+
+            const systemMessageDescription = await this.getSystemMessageDescription();
             const messages = await this.getMessages(request.session);
             if (this.defaultLogging) {
-                this.recordingService.recordRequest({
-                    agentId: this.id,
-                    sessionId: request.session.id,
-                    timestamp: Date.now(),
-                    requestId: request.id,
-                    request: request.request.text,
-                    messages
-                });
+                this.recordingService.recordRequest(
+                    ChatHistoryEntry.fromRequest(
+                        this.id, request, {
+                        messages,
+                        systemMessage: systemMessageDescription?.text
+                    })
+                );
             }
 
-            const systemMessageDescription = await this.getSystemMessageDescription();
             const tools: Map<string, ToolRequest> = new Map();
             if (systemMessageDescription) {
                 const systemMsg: ChatMessage = {
@@ -196,13 +197,7 @@ export abstract class AbstractChatAgent {
             await this.addContentsToResponse(languageModelResponse, request);
             request.response.complete();
             if (this.defaultLogging) {
-                this.recordingService.recordResponse({
-                    agentId: this.id,
-                    sessionId: request.session.id,
-                    timestamp: Date.now(),
-                    requestId: request.response.requestId,
-                    response: request.response.response.asString()
-                });
+                this.recordingService.recordResponse(ChatHistoryEntry.fromResponse(this.id, request));
             }
         } catch (e) {
             this.handleError(request, e);
@@ -322,28 +317,15 @@ export abstract class AbstractStreamParsingChatAgent extends AbstractChatAgent {
             request.response.response.addContents(contents);
             request.response.complete();
             if (this.defaultLogging) {
-                this.recordingService.recordResponse({
-                    agentId: this.id,
-                    sessionId: request.session.id,
-                    timestamp: Date.now(),
-                    requestId: request.response.requestId,
-                    response: request.response.response.asString()
-
-                });
+                this.recordingService.recordResponse(ChatHistoryEntry.fromResponse(this.id, request));
             }
             return;
         }
         if (isLanguageModelStreamResponse(languageModelResponse)) {
             await this.addStreamResponse(languageModelResponse, request);
             request.response.complete();
             if (this.defaultLogging) {
-                this.recordingService.recordResponse({
-                    agentId: this.id,
-                    sessionId: request.session.id,
-                    timestamp: Date.now(),
-                    requestId: request.response.requestId,
-                    response: request.response.response.asString()
-                });
+                this.recordingService.recordResponse(ChatHistoryEntry.fromResponse(this.id, request));
             }
             return;
         }

packages/ai-chat/src/common/chat-history-entry.ts

@@ -0,0 +1,47 @@
+// *****************************************************************************
+// Copyright (C) 2024 EclipseSource GmbH.
+//
+// This program and the accompanying materials are made available under the
+// terms of the Eclipse Public License v. 2.0 which is available at
+// http://www.eclipse.org/legal/epl-2.0.
+//
+// This Source Code may also be made available under the following Secondary
+// Licenses when the conditions for such availability set forth in the Eclipse
+// Public License v. 2.0 are satisfied: GNU General Public License, version 2
+// with the GNU Classpath Exception which is available at
+// https://www.gnu.org/software/classpath/license.html.
+//
+// SPDX-License-Identifier: EPL-2.0 OR GPL-2.0-only WITH Classpath-exception-2.0
+// *****************************************************************************
+
+import { CommunicationRequestEntryParam, CommunicationResponseEntryParam } from '@theia/ai-core/lib/common/communication-recording-service';
+import { ChatRequestModel } from './chat-model';
+
+export namespace ChatHistoryEntry {
+    export function fromRequest(
+        agentId: string,
+        request: ChatRequestModel,
+        args: Partial<CommunicationRequestEntryParam> = {}
+    ): CommunicationRequestEntryParam {
+        return {
+            agentId: agentId,
+            sessionId: request.session.id,
+            requestId: request.id,
+            request: request.request.text,
+            ...args,
+        };
+    }
+    export function fromResponse(
+        agentId: string,
+        request: ChatRequestModel,
+        args: Partial<CommunicationResponseEntryParam> = {}
+    ): CommunicationResponseEntryParam {
+        return {
+            agentId: agentId,
+            sessionId: request.session.id,
+            requestId: request.id,
+            response: request.response.response.asString(),
+            ...args,
+        };
+    }
+}

packages/ai-chat/src/common/orchestrator-chat-agent.ts

@@ -23,6 +23,7 @@ import { ChatAgentService } from './chat-agent-service';
 import { AbstractStreamParsingChatAgent, ChatAgent, SystemMessageDescription } from './chat-agents';
 import { ChatRequestModelImpl, InformationalChatResponseContentImpl } from './chat-model';
 import { generateUuid } from '@theia/core';
+import { ChatHistoryEntry } from './chat-history-entry';
 
 export const orchestratorTemplate: PromptTemplate = {
     id: 'orchestrator-system',
@@ -60,7 +61,7 @@ You must only use the \`id\` attribute of the agent, never the name.
 `};
 
 export const OrchestratorChatAgentId = 'Orchestrator';
-const OrchestatorRequestIdKey = 'orchestatorRequestIdKey';
+const OrchestratorRequestIdKey = 'orchestatorRequestIdKey';
 
 @injectable()
 export class OrchestratorChatAgent extends AbstractStreamParsingChatAgent implements ChatAgent {
@@ -93,16 +94,17 @@ export class OrchestratorChatAgent extends AbstractStreamParsingChatAgent implem
     override async invoke(request: ChatRequestModelImpl): Promise<void> {
         request.response.addProgressMessage({ content: 'Determining the most appropriate agent', status: 'inProgress' });
         // We generate a dedicated ID for recording the orchestrator request/response, as we will forward the original request to another agent
-        const orchestartorRequestId = generateUuid();
-        request.addData(OrchestatorRequestIdKey, orchestartorRequestId);
-        const userPrompt = request.request.text;
-        this.recordingService.recordRequest({
-            agentId: this.id,
-            sessionId: request.session.id,
-            timestamp: Date.now(),
-            requestId: orchestartorRequestId,
-            request: userPrompt,
-        });
+        const orchestratorRequestId = generateUuid();
+        request.addData(OrchestratorRequestIdKey, orchestratorRequestId);
+        const messages = await this.getMessages(request.session);
+        const systemMessage = (await this.getSystemMessageDescription())?.text;
+        this.recordingService.recordRequest(
+            ChatHistoryEntry.fromRequest(this.id, request, {
+                requestId: orchestratorRequestId,
+                messages,
+                systemMessage
+            })
+        );
         return super.invoke(request);
     }
 
@@ -115,12 +117,11 @@ export class OrchestratorChatAgent extends AbstractStreamParsingChatAgent implem
         let agentIds: string[] = [];
         const responseText = await getTextOfResponse(response);
         // We use the previously generated, dedicated ID to log the orchestrator response before we forward the original request
-        const orchestratorRequestId = request.getDataByKey(OrchestatorRequestIdKey);
+        const orchestratorRequestId = request.getDataByKey(OrchestratorRequestIdKey);
         if (typeof orchestratorRequestId === 'string') {
             this.recordingService.recordResponse({
                 agentId: this.id,
                 sessionId: request.session.id,
-                timestamp: Date.now(),
                 requestId: orchestratorRequestId,
                 response: responseText,
             });

packages/ai-code-completion/src/common/code-completion-agent.ts

@@ -15,7 +15,7 @@
 // *****************************************************************************
 
 import {
-    Agent, AgentSpecificVariables, CommunicationHistoryEntry, CommunicationRecordingService, getTextOfResponse,
+    Agent, AgentSpecificVariables, CommunicationRecordingService, getTextOfResponse,
     LanguageModelRegistry, LanguageModelRequest, LanguageModelRequirement, PromptService, PromptTemplate
 } from '@theia/ai-core/lib/common';
 import { generateUuid, ILogger } from '@theia/core';
@@ -84,17 +84,15 @@ export class CodeCompletionAgentImpl implements CodeCompletionAgent {
         const request: LanguageModelRequest = {
             messages: [{ type: 'text', actor: 'user', query: prompt }],
         };
-        const requestEntry: CommunicationHistoryEntry = {
+        if (token.isCancellationRequested) {
+            return undefined;
+        }
+        this.recordingService.recordRequest({
             agentId: this.id,
             sessionId,
-            timestamp: Date.now(),
             requestId,
             request: prompt,
-        };
-        if (token.isCancellationRequested) {
-            return undefined;
-        }
-        this.recordingService.recordRequest(requestEntry);
+        });
         const response = await languageModel.request(request, token);
         if (token.isCancellationRequested) {
             return undefined;
@@ -106,7 +104,6 @@ export class CodeCompletionAgentImpl implements CodeCompletionAgent {
         this.recordingService.recordResponse({
             agentId: this.id,
             sessionId,
-            timestamp: Date.now(),
             requestId,
             response: completionText,
         });

packages/ai-core/src/common/communication-recording-service.ts

@@ -26,18 +26,22 @@ export interface CommunicationHistoryEntry {
     request?: string;
     response?: string;
     responseTime?: number;
+    systemMessage?: string;
     messages?: unknown[];
 }
 
 export type CommunicationRequestEntry = Omit<CommunicationHistoryEntry, 'response' | 'responseTime'>;
 export type CommunicationResponseEntry = Omit<CommunicationHistoryEntry, 'request'>;
 
+export type CommunicationRequestEntryParam = Omit<CommunicationRequestEntry, 'timestamp'> & Partial<Pick<CommunicationHistoryEntry, 'timestamp'>>;
+export type CommunicationResponseEntryParam = Omit<CommunicationResponseEntry, 'timestamp'> & Partial<Pick<CommunicationHistoryEntry, 'timestamp'>>;
+
 export const CommunicationRecordingService = Symbol('CommunicationRecordingService');
 export interface CommunicationRecordingService {
-    recordRequest(requestEntry: CommunicationRequestEntry): void;
+    recordRequest(requestEntry: CommunicationRequestEntryParam): void;
     readonly onDidRecordRequest: Event<CommunicationRequestEntry>;
 
-    recordResponse(responseEntry: CommunicationResponseEntry): void;
+    recordResponse(responseEntry: CommunicationResponseEntryParam): void;
     readonly onDidRecordResponse: Event<CommunicationResponseEntry>;
 
     getHistory(agentId: string): CommunicationHistory;

packages/ai-history/README.md

@@ -15,7 +15,6 @@
 The `@theia/ai-history` extension offers a framework for agents to record their requests and responses.
 It also offers a view to inspect the history.
 
-
 ## Additional Information
 
 - [Theia - GitHub](https://github.com/eclipse-theia/theia)

packages/ai-history/src/browser/ai-history-communication-card.tsx

@@ -29,16 +29,39 @@ export const CommunicationCard: React.FC<CommunicationCardProps> = ({ entry }) =
         <div className='theia-card-content'>
             {entry.request && (
                 <div className='theia-card-request'>
-                    <p><strong>Request</strong></p>
+                    <h2>Request</h2>
                     <pre>{entry.request}</pre>
                 </div>
             )}
             {entry.response && (
                 <div className='theia-card-response'>
-                    <p><strong>Response</strong></p>
+                    <h2>Response</h2>
                     <pre>{entry.response}</pre>
                 </div>
             )}
+            {(entry.systemMessage || (entry.messages && entry.messages.length > 0)) && (
+                <div className='theia-card-context'>
+                    <details>
+                        <summary><h2>Context</h2></summary>
+                        {(entry.systemMessage && (
+                            <div className='theia-context-system-message'>
+                                <h3>System Message</h3>
+                                <pre>{entry.systemMessage}</pre>
+                            </div>
+                        ))}
+                        {(entry.messages && entry.messages.length > 0) && (
+                            <div className='theia-context-messages'>
+                                <h3>Messages</h3>
+                                <ul>
+                                    {entry.messages.map((message, index) => (
+                                        <li key={index}><pre>{JSON.stringify(message, undefined, 2)}</pre></li>
+                                    ))}
+                                </ul>
+                            </div>
+                        )}
+                    </details>
+                </div>
+            )}
         </div>
         <div className='theia-card-meta'>
             <span className='theia-card-timestamp'>Timestamp: {new Date(entry.timestamp).toLocaleString()}</span>

packages/ai-history/src/browser/ai-history-widget.tsx

@@ -102,13 +102,17 @@ export class AIHistoryView extends ReactWidget implements StatefulWidget {
         return (
             <div className='agent-history-widget'>
                 <SelectComponent
-                    options={this.agentService.getAllAgents().map(agent => ({ value: agent.id, label: agent.name, description: agent.description }))}
+                    options={this.agentService.getAllAgents().map(agent => ({
+                        value: agent.id,
+                        label: agent.name,
+                        description: agent.description || ''
+                    }))}
                     onChange={selectionChange}
                     defaultValue={this.selectedAgent?.id} />
                 <div className='agent-history'>
                     {this.renderHistory()}
                 </div>
-            </div >
+            </div>
         );
     }
 

packages/ai-history/src/browser/style/ai-history.css

@@ -31,7 +31,6 @@
     display: flex;
     justify-content: space-between;
     font-size: 0.9em;
-    margin-bottom: var(--theia-ui-padding);
     padding: var(--theia-ui-padding) 0;
 }
 
@@ -40,6 +39,24 @@
     margin-bottom: 10px;
 }
 
+.theia-card-content h2 {
+    font-size: var(--theia-ui-font-size2);
+}
+.theia-card-content h3 {
+    font-size: var(--theia-ui-font-size1);
+    font-weight: bold;
+}
+
+.theia-card-content .theia-card-context h2 {
+    display: inline;
+}
+.theia-card-content .theia-card-context ul {
+    padding-inline-start: 20px;
+}
+.theia-card-content .theia-card-context pre {
+    overflow: scroll;
+}
+
 .theia-card-content p {
     margin: var(--theia-ui-padding) 0;
 }