feat: some optimization for customized render function

.github/PULL_REQUEST_TEMPLATE.md

@@ -6,11 +6,9 @@
 
 ### Check List
 - [ ] New functionality includes testing.
-  - [ ] All tests pass, including unit test, integration test and doctest
-- [ ] New functionality has been documented.
-  - [ ] New functionality has javadoc added
-  - [ ] New functionality has user manual doc added
-- [ ] Commits are signed per the DCO using --signoff
+  - [ ] All tests pass, including unit test, integration test.
+- [ ] New functionality has user manual doc added.
+- [ ] Commits are signed per the DCO using --signoff.
 
 By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.
 For more information on following Developer Certificate of Origin and signing off your commits, please check [here](https://github.com/opensearch-project/OpenSearch/blob/main/CONTRIBUTING.md#developer-certificate-of-origin).

common/types/chat_saved_object_attributes.ts

@@ -43,11 +43,11 @@ export interface IOutput {
   type: 'output';
   traceId?: string; // used for tracing agent calls
   toolsUsed?: string[];
-  // TODO: ppl_visualization type may need to be removed in the PR which replaces ppl query render from visualization to data grid. @suzhou
-  contentType: 'error' | 'markdown' | 'visualization' | 'ppl_visualization' | 'ppl_data_grid';
+  contentType: 'error' | 'markdown' | 'visualization' | 'ppl_data_grid' | string;
   content: string;
   suggestedActions?: ISuggestedAction[];
   messageId?: string;
+  isVisualization?: boolean;
 }
 export type IMessage = IInput | IOutput;
 

public/index.ts

@@ -12,4 +12,5 @@ export function plugin(initializerContext: PluginInitializerContext) {
   return new AssistantPlugin(initializerContext);
 }
 
-export { AssistantSetup } from './types';
+export { AssistantSetup, RenderProps } from './types';
+export { IMessage } from '../common/types/chat_saved_object_attributes';

public/tabs/chat/messages/message_bubble.test.tsx

@@ -102,6 +102,7 @@ describe('<MessageBubble />', () => {
           type: 'output',
           contentType: 'visualization',
           content: 'vis_id_mock',
+          isVisualization: true,
         }}
       />
     );
@@ -114,6 +115,7 @@ describe('<MessageBubble />', () => {
           type: 'output',
           contentType: 'ppl_visualization',
           content: 'vis_id_mock',
+          isVisualization: true,
         }}
       />
     );

public/tabs/chat/messages/message_bubble.tsx

@@ -133,13 +133,7 @@ export const MessageBubble: React.FC<MessageBubbleProps> = React.memo((props) =>
       );
     }
 
-    // if (['visualization', 'ppl_visualization'].includes(props.contentType)) {
-    //   return <>{props.children}</>;
-    // }
-
-    const isVisualization = ['visualization', 'ppl_visualization'].includes(
-      props.message.contentType
-    );
+    const isVisualization = props.message.isVisualization;
 
     return (
       <EuiFlexGroup
@@ -154,7 +148,11 @@ export const MessageBubble: React.FC<MessageBubbleProps> = React.memo((props) =>
         </EuiFlexItem>
         <EuiFlexItem className="llm-chat-bubble-wrapper">
           <EuiPanel
-            style={isVisualization ? { minWidth: '100%' } : {}}
+            /**
+             * When using minWidth the content width inside may be larger than the container itself,
+             * especially in data grid case that the content will change its size according to if fullScreen or not.
+             */
+            style={isVisualization ? { width: '100%' } : {}}
             hasShadow={false}
             hasBorder={false}
             paddingSize="l"

public/tabs/chat/messages/message_content.test.tsx

@@ -86,10 +86,45 @@ describe('<MessageContent />', () => {
           type: 'output',
           contentType: 'ppl_visualization',
           content: 'mock ppl query',
+          isVisualization: true,
         }}
       />
     );
-    expect(pplVisualizationRenderMock).toHaveBeenCalledWith({ query: 'mock ppl query' });
+    expect(pplVisualizationRenderMock.mock.calls[0]).toMatchInlineSnapshot(`
+      Array [
+        Object {
+          "content": "mock ppl query",
+          "contentType": "ppl_visualization",
+          "isVisualization": true,
+          "type": "output",
+        },
+        Object {
+          "chatContext": Object {
+            "contentRenderers": Object {
+              "ppl_visualization": [MockFunction] {
+                "calls": Array [
+                  [Circular],
+                ],
+                "results": Array [
+                  Object {
+                    "type": "return",
+                    "value": undefined,
+                  },
+                ],
+              },
+            },
+          },
+          "props": Object {
+            "message": Object {
+              "content": "mock ppl query",
+              "contentType": "ppl_visualization",
+              "isVisualization": true,
+              "type": "output",
+            },
+          },
+        },
+      ]
+    `);
   });
 
   it('should render customized render content', () => {

public/tabs/chat/messages/message_content.tsx

@@ -9,7 +9,7 @@ import { IMessage } from '../../../../common/types/chat_saved_object_attributes'
 import { CoreVisualization } from '../../../components/core_visualization';
 import { useChatContext } from '../../../contexts/chat_context';
 
-interface MessageContentProps {
+export interface MessageContentProps {
   message: IMessage;
 }
 
@@ -37,18 +37,15 @@ export const MessageContent: React.FC<MessageContentProps> = React.memo((props)
         </div>
       );
 
-    case 'ppl_visualization': {
-      const render = chatContext.contentRenderers[props.message.contentType];
-      if (!render) return null;
-      return (
-        <div className="llm-chat-visualizations">{render({ query: props.message.content })}</div>
-      );
-    }
-
     // content types registered by plugins unknown to assistant
     default: {
       const message = props.message as IMessage;
-      return chatContext.contentRenderers[message.contentType]?.(message.content) ?? null;
+      return (
+        chatContext.contentRenderers[message.contentType]?.(message, {
+          props,
+          chatContext,
+        }) ?? null
+      );
     }
   }
 });

public/types.ts

@@ -6,9 +6,16 @@
 import { DashboardStart } from '../../../src/plugins/dashboard/public';
 import { EmbeddableSetup, EmbeddableStart } from '../../../src/plugins/embeddable/public';
 import { IMessage, ISuggestedAction } from '../common/types/chat_saved_object_attributes';
+import { IChatContext } from './contexts/chat_context';
+import { MessageContentProps } from './tabs/chat/messages/message_content';
+
+export interface RenderProps {
+  props: MessageContentProps;
+  chatContext: IChatContext;
+}
 
 // TODO should pair with server side registered output parser
-export type ContentRenderer = (content: unknown) => React.ReactElement;
+export type ContentRenderer = (content: unknown, renderProps: RenderProps) => React.ReactElement;
 export type ActionExecutor = (params: Record<string, unknown>) => void;
 export interface AssistantActions {
   send: (input: IMessage) => Promise<void>;

server/parsers/visualization_card_parser.test.ts

@@ -29,18 +29,21 @@ describe('VisualizationCardParser', () => {
         contentType: 'visualization',
         suggestedActions: [{ actionType: 'view_in_dashboards', message: 'View in Visualize' }],
         type: 'output',
+        isVisualization: true,
       },
       {
         content: 'id2',
         contentType: 'visualization',
         suggestedActions: [{ actionType: 'view_in_dashboards', message: 'View in Visualize' }],
         type: 'output',
+        isVisualization: true,
       },
       {
         content: 'id3',
         contentType: 'visualization',
         suggestedActions: [{ actionType: 'view_in_dashboards', message: 'View in Visualize' }],
         type: 'output',
+        isVisualization: true,
       },
     ]);
   });
@@ -66,12 +69,14 @@ describe('VisualizationCardParser', () => {
         contentType: 'visualization',
         suggestedActions: [{ actionType: 'view_in_dashboards', message: 'View in Visualize' }],
         type: 'output',
+        isVisualization: true,
       },
       {
         content: 'id2',
         contentType: 'visualization',
         suggestedActions: [{ actionType: 'view_in_dashboards', message: 'View in Visualize' }],
         type: 'output',
+        isVisualization: true,
       },
     ]);
   });
@@ -110,6 +115,7 @@ describe('VisualizationCardParser', () => {
         contentType: 'visualization',
         suggestedActions: [{ actionType: 'view_in_dashboards', message: 'View in Visualize' }],
         type: 'output',
+        isVisualization: true,
       },
     ]);
   });
@@ -137,12 +143,14 @@ describe('VisualizationCardParser', () => {
         contentType: 'visualization',
         suggestedActions: [{ actionType: 'view_in_dashboards', message: 'View in Visualize' }],
         type: 'output',
+        isVisualization: true,
       },
       {
         content: 'id2',
         contentType: 'visualization',
         suggestedActions: [{ actionType: 'view_in_dashboards', message: 'View in Visualize' }],
         type: 'output',
+        isVisualization: true,
       },
     ]);
   });

server/parsers/visualization_card_parser.ts

@@ -35,6 +35,7 @@ export const VisualizationCardParser = {
         type: 'output',
         content: id,
         contentType: 'visualization',
+        isVisualization: true,
         suggestedActions: [
           {
             message: 'View in Visualize',