Colocate data fetching; use proper server components

Performing all image search requests upfront, before rendering any image
components, wasn't really in line with how server components should
work. This approach was initially chosen to obtain a single `function`
AI state that contains data from all searches.

With the introduction of the `onDataFetched` prop, we can achieve the
same outcome while moving data fetching directly into the server
components, where it belongs.

This change also makes it easier to add loading skeletons later on.

src/app/google-image-search.ts

@@ -2,6 +2,7 @@ import {z} from 'zod';
 
 export type Image = z.TypeOf<typeof image>;
 export type ImageSearchResponse = z.TypeOf<typeof imageSearchResponse>;
+export type ImageSearchParams = z.TypeOf<typeof imageSearchParams>;
 
 const apiKey = process.env.GOOGLE_SEARCH_API_KEY!;
 const searchEngineId = process.env.GOOGLE_SEARCH_SEARCH_ENGINE_ID!;

src/app/image-search-utils.ts → src/app/image-search-result.ts

@@ -37,13 +37,9 @@ export function createImageSearchResult(
     : {status: `found`, images: response, title};
 }
 
-export function serializeImageSearchResults(
-  imageSearchResults: ImageSearchResult[],
-): string {
-  return JSON.stringify(imageSearchResults.map(reduceImageSarchResult));
-}
-
-function reduceImageSarchResult(imageSearchResult: ImageSearchResult) {
+export function prepareImageSearchResultForAiState(
+  imageSearchResult: ImageSearchResult,
+): unknown {
   const {status, title} = imageSearchResult;
 
   if (imageSearchResult.status === `found`) {

src/app/images.tsx

@@ -0,0 +1,65 @@
+import * as React from 'react';
+import {type ImageSearchParams, searchImages} from './google-image-search.js';
+import {
+  createImageSearchResult,
+  prepareImageSearchResultForAiState,
+} from './image-search-result.js';
+import {ImageSelector} from './image-selector.js';
+import {ProgressiveImage} from './progressive-image.js';
+
+export interface ImagesProps {
+  readonly title: string;
+  readonly notFoundMessage: string;
+  readonly errorMessage: string;
+  readonly searchParams: ImageSearchParams;
+  readonly onDataFetched: (dataForAiState: unknown) => void;
+}
+
+export async function Images({
+  title,
+  notFoundMessage,
+  errorMessage,
+  searchParams,
+  onDataFetched,
+}: ImagesProps): Promise<React.ReactElement> {
+  const response = await searchImages(searchParams);
+
+  const result = createImageSearchResult({
+    response,
+    title,
+    notFoundMessage,
+    errorMessage,
+  });
+
+  onDataFetched(prepareImageSearchResultForAiState(result));
+
+  if (result.status !== `found`) {
+    const message =
+      result.status === `not-found`
+        ? result.notFoundMessage
+        : result.errorMessage;
+
+    return (
+      <p className="text-sm">
+        <em>{message}</em>
+      </p>
+    );
+  }
+
+  return (
+    <div className="space-y-3">
+      <h4 className="text-l font-bold">{result.title}</h4>
+      {result.images.map(({thumbnailUrl, url, width, height}) => (
+        <ImageSelector key={thumbnailUrl} url={url}>
+          <ProgressiveImage
+            thumbnailUrl={thumbnailUrl}
+            url={url}
+            width={width}
+            height={height}
+            alt={result.title}
+          />
+        </ImageSelector>
+      ))}
+    </div>
+  );
+}

src/app/submit-user-message.tsx

@@ -6,15 +6,10 @@ import * as React from 'react';
 import {z} from 'zod';
 import {type UserInput, fromUserInput} from './ai-state.js';
 import type {AI, UIStateItem} from './ai.js';
-import {imageSearchParams, searchImages} from './google-image-search.js';
-import {
-  createImageSearchResult,
-  serializeImageSearchResults,
-} from './image-search-utils.js';
-import {ImageSelector} from './image-selector.js';
+import {imageSearchParams} from './google-image-search.js';
+import {Images} from './images.js';
 import {LoadingIndicator} from './loading-indicator.js';
 import {Markdown} from './markdown.js';
-import {ProgressiveImage} from './progressive-image.js';
 import {UserChoiceButton} from './user-choice-button.js';
 
 const openai = new OpenAI({apiKey: process.env.OPENAI_API_KEY});
@@ -173,8 +168,8 @@ export async function submitUserMessage(
             )
             .describe(`Use multiple sets of search parameters if needed.`),
         }),
-        async *render({loadingText, searches: searchParamsList}) {
-          console.log(`search_and_show_images`, searchParamsList);
+        async *render({loadingText, searches}) {
+          console.log(`search_and_show_images`, searches);
 
           const text = lastTextContent ? (
             <div>
@@ -189,69 +184,60 @@ export async function submitUserMessage(
             </div>
           );
 
-          const imageSearchResults = await Promise.all(
-            searchParamsList.map(
-              async ({searchParams, title, notFoundMessage, errorMessage}) =>
-                createImageSearchResult({
-                  response: await searchImages(searchParams),
-                  title,
-                  notFoundMessage,
-                  errorMessage,
-                }),
-            ),
-          );
-
           if (lastTextContent) {
             aiState.update((prevAiState) => [
               ...prevAiState,
               {role: `assistant`, content: lastTextContent!},
             ]);
           }
 
-          aiState.done([
-            ...aiState.get(),
-            {
-              role: `function`,
-              name: `search_and_show_images`,
-              content: serializeImageSearchResults(imageSearchResults),
+          const elementsWithData = searches.map(
+            ({title, notFoundMessage, errorMessage, searchParams}) => {
+              let resolveDataPromise: (data: unknown) => void;
+
+              return {
+                element: (
+                  <Images
+                    key={title}
+                    title={title}
+                    notFoundMessage={notFoundMessage}
+                    errorMessage={errorMessage}
+                    searchParams={searchParams}
+                    onDataFetched={(data) => resolveDataPromise(data)}
+                  />
+                ),
+                dataPromise: new Promise(
+                  (resolve) => (resolveDataPromise = resolve),
+                ),
+              };
             },
-          ]);
+          );
 
-          return (
+          const finalUi = (
             <div className="space-y-4">
               {text}
               <div className="space-y-3">
-                {imageSearchResults.map((result) => (
-                  <React.Fragment key={result.title}>
-                    <h4 className="text-l font-bold">{result.title}</h4>
-                    {result.status === `found` ? (
-                      result.images.map(
-                        ({thumbnailUrl, url, width, height}) => (
-                          <ImageSelector key={thumbnailUrl} url={url}>
-                            <ProgressiveImage
-                              thumbnailUrl={thumbnailUrl}
-                              url={url}
-                              width={width}
-                              height={height}
-                              alt={result.title}
-                            />
-                          </ImageSelector>
-                        ),
-                      )
-                    ) : (
-                      <p className="text-sm">
-                        <em>
-                          {result.status === `not-found`
-                            ? result.notFoundMessage
-                            : result.errorMessage}
-                        </em>
-                      </p>
-                    )}
-                  </React.Fragment>
-                ))}
+                {elementsWithData.map(({element}) => element)}
               </div>
             </div>
           );
+
+          yield finalUi;
+
+          const dataItems = await Promise.all(
+            elementsWithData.map(async ({dataPromise}) => dataPromise),
+          );
+
+          aiState.done([
+            ...aiState.get(),
+            {
+              role: `function`,
+              name: `search_and_show_images`,
+              content: JSON.stringify(dataItems),
+            },
+          ]);
+
+          return finalUi;
         },
       },
     },