Merge pull request #4 from unstubbable/colocate-data-fetching

Colocate data fetching; use proper server components

package.json

@@ -32,10 +32,10 @@
     "ai": "^3.0.13",
     "clsx": "^1.2.1",
     "openai": "^4.29.0",
-    "react": "^18.3.0-canary-a870b2d54-20240314",
-    "react-dom": "^18.3.0-canary-a870b2d54-20240314",
+    "react": "^19.0.0-canary-7a2609eed-20240403",
+    "react-dom": "^19.0.0-canary-7a2609eed-20240403",
     "react-markdown": "^9.0.1",
-    "react-server-dom-webpack": "^18.3.0-canary-a870b2d54-20240314",
+    "react-server-dom-webpack": "^19.0.0-canary-7a2609eed-20240403",
     "react-textarea-autosize": "^8.5.3",
     "server-only": "^0.0.1",
     "zod": "^3.22.4"

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,79 @@
+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;
+}
+
+type ImagesContainerProps = React.PropsWithChildren<{
+  readonly title: string;
+}>;
+
+function ImagesContainer({title, children}: ImagesContainerProps) {
+  return (
+    <div className="space-y-3">
+      <h4 className="text-l font-bold">{title}</h4>
+      {children}
+    </div>
+  );
+}
+
+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 (
+      <ImagesContainer title={result.title}>
+        <p className="text-sm">
+          <em>{message}</em>
+        </p>
+      </ImagesContainer>
+    );
+  }
+
+  return (
+    <ImagesContainer title={result.title}>
+      {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>
+      ))}
+    </ImagesContainer>
+  );
+}

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;
         },
       },
     },