image rendering and vision agent endpoint boilerplate (#6)

* check-in code

* image rendering and vision agent endpoint boilerplate

app/api/chat/route.ts

@@ -7,8 +7,7 @@ import {
   ChatCompletionContentPart,
   ChatCompletionContentPartImage,
 } from 'openai/resources';
-import { MessageWithSelectedDataset } from '../../../lib/types';
-// import { postAgentChat } from '@/lib/fetch';
+import { MessageBase } from '../../../lib/types';
 
 export const runtime = 'edge';
 
@@ -19,9 +18,10 @@ const openai = new OpenAI({
 export async function POST(req: Request) {
   const json = await req.json();
   const { messages } = json as {
-    messages: MessageWithSelectedDataset[];
+    messages: MessageBase[];
+    id: string;
+    url: string;
   };
-  console.log('[Ming] ~ POST ~ messages:', messages);
 
   const session = await auth();
   if (!session?.user?.email) {
@@ -32,20 +32,11 @@ export async function POST(req: Request) {
 
   const formattedMessage: ChatCompletionMessageParam[] = messages.map(
     message => {
-      const { dataset, ...rest } = message;
-
       const contentWithImage: ChatCompletionContentPart[] = [
         {
           type: 'text',
           text: message.content as string,
         },
-        ...(dataset ?? []).map(
-          entity =>
-            ({
-              type: 'image_url',
-              image_url: { url: entity.url },
-            }) satisfies ChatCompletionContentPartImage,
-        ),
       ];
       return {
         role: 'user',
@@ -53,7 +44,6 @@ export async function POST(req: Request) {
       };
     },
   );
-
   const res = await openai.chat.completions.create({
     model: 'gpt-4-vision-preview',
     messages: formattedMessage,

app/api/vision-agent/route.ts

@@ -0,0 +1,95 @@
+import { StreamingTextResponse } from 'ai';
+
+import { auth } from '@/auth';
+import {
+  AIStream,
+  type AIStreamParser,
+  type AIStreamCallbacksAndOptions,
+} from 'ai';
+import { MessageBase } from '../../../lib/types';
+import { fetcher } from '@/lib/utils';
+
+export const runtime = 'edge';
+
+export async function POST(req: Request) {
+  const json = await req.json();
+  const { messages, url } = json as {
+    messages: MessageBase[];
+    id: string;
+    url: string;
+  };
+
+  const session = await auth();
+  if (!session?.user?.email) {
+    return new Response('Unauthorized', {
+      status: 401,
+    });
+  }
+
+  function parseVisionAgentStream(): AIStreamParser {
+    let previous = '';
+
+    return data => {
+      console.log('[Ming] ~ parseVisionAgentStream ~ data:', data);
+      // const json = JSON.parse(data) as {
+      //   completion: string;
+      //   stop: string | null;
+      //   stop_reason: string | null;
+      //   truncated: boolean;
+      //   log_id: string;
+      //   model: string;
+      //   exception: string | null;
+      // };
+
+      // // Anthropic's `completion` field is cumulative unlike OpenAI's
+      // // deltas. In order to compute the delta, we must slice out the text
+      // // we previously received.
+      // const text = json.completion;
+      // console.log('[Ming] ~ parseVisionAgentStream ~ text:', text);
+      // const delta = text.slice(previous.length);
+      // previous = text;
+
+      return data;
+    };
+  }
+
+  function visionAgentStream(
+    res: Response,
+    cb?: AIStreamCallbacksAndOptions,
+  ): ReadableStream {
+    return AIStream(res, parseVisionAgentStream(), cb);
+  }
+
+  const formData = new FormData();
+  formData.append('input', JSON.stringify(messages));
+  formData.append('image', url);
+
+  const fetchResponse = await fetcher(
+    'https://api.dev.landing.ai/v1/agent/chat?agent_class=vision_agent',
+    {
+      method: 'POST',
+      headers: {
+        apikey: 'land_sk_DKeoYtaZZrYqJ9TMMiXe4BIQgJcZ0s3XAoB0JT3jv73FFqnr6k',
+      },
+      body: formData,
+    },
+  );
+  console.log('[Ming] ~ POST ~ fetchResponse:', fetchResponse);
+  const stream = visionAgentStream(fetchResponse, {
+    onStart: async () => {
+      console.log('Stream started');
+    },
+    onCompletion: async completion => {
+      console.log('Completion completed', completion);
+    },
+    onFinal: async completion => {
+      console.log('Stream completed', completion);
+    },
+    onToken: async token => {
+      console.log('Token received', token);
+    },
+  });
+  // Now you can consume the VisionAgentStream
+
+  return new StreamingTextResponse(stream);
+}

app/chat/[id]/page.tsx

@@ -1,8 +1,17 @@
 import { nanoid } from '@/lib/utils';
 import { Chat } from '@/components/chat';
+import { getKVChat } from '@/lib/kv/chat';
 
-export default async function Page() {
-  const id = nanoid();
+interface PageProps {
+  params: {
+    id: string;
+  };
+}
+
+export default async function Page({ params }: PageProps) {
+  const { id: chatId } = params;
+
+  const chat = await getKVChat(chatId);
 
-  return <Chat id={id} />;
+  return <Chat chat={chat} />;
 }

components/chat/ChatList.tsx

@@ -2,10 +2,10 @@
 
 import { Separator } from '@/components/ui/Separator';
 import { ChatMessage } from '@/components/chat/ChatMessage';
-import { MessageWithSelectedDataset } from '../../lib/types';
+import { MessageBase } from '../../lib/types';
 
 export interface ChatList {
-  messages: MessageWithSelectedDataset[];
+  messages: MessageBase[];
 }
 
 export function ChatList({ messages }: ChatList) {

components/chat/ChatMessage.tsx

@@ -9,10 +9,10 @@ import { CodeBlock } from '@/components/ui/CodeBlock';
 import { MemoizedReactMarkdown } from '@/components/chat/MemoizedReactMarkdown';
 import { IconOpenAI, IconUser } from '@/components/ui/Icons';
 import { ChatMessageActions } from '@/components/chat/ChatMessageActions';
-import { MessageWithSelectedDataset } from '../../lib/types';
+import { MessageBase } from '../../lib/types';
 
 export interface ChatMessageProps {
-  message: MessageWithSelectedDataset;
+  message: MessageBase;
 }
 
 export function ChatMessage({ message, ...props }: ChatMessageProps) {

components/chat/ChatMessageActions.tsx

@@ -6,10 +6,10 @@ import { Button } from '@/components/ui/Button';
 import { IconCheck, IconCopy } from '@/components/ui/Icons';
 import { useCopyToClipboard } from '@/lib/hooks/useCopyToClipboard';
 import { cn } from '@/lib/utils';
-import { MessageWithSelectedDataset } from '../../lib/types';
+import { MessageBase } from '../../lib/types';
 
 interface ChatMessageActionsProps extends React.ComponentProps<'div'> {
-  message: MessageWithSelectedDataset;
+  message: MessageBase;
 }
 
 export function ChatMessageActions({

components/chat/ChatPanel.tsx

@@ -5,7 +5,7 @@ import { Button } from '@/components/ui/Button';
 import { PromptForm } from '@/components/chat/PromptForm';
 import { ButtonScrollToBottom } from '@/components/chat/ButtonScrollToBottom';
 import { IconRefresh, IconShare, IconStop } from '@/components/ui/Icons';
-import { MessageWithSelectedDataset } from '../../lib/types';
+import { MessageBase } from '../../lib/types';
 
 export interface ChatPanelProps
   extends Pick<
@@ -14,7 +14,7 @@ export interface ChatPanelProps
   > {
   id?: string;
   title?: string;
-  messages: MessageWithSelectedDataset[];
+  messages: MessageBase[];
 }
 
 export function ChatPanel({

components/chat/ImageSelector.tsx

@@ -18,7 +18,11 @@ const examples: Example[] = [
   {
     url: 'https://landing-lens-support.s3.us-east-2.amazonaws.com/vision-agent-examples/cereal-example.jpg',
     initMessages: [
-      { role: 'user', content: 'how many cereals are there in the image?' },
+      {
+        role: 'user',
+        content: 'how many cereals are there in the image?',
+        id: 'fake-id-1',
+      },
     ],
   },
   // 'https://landing-lens-support.s3.us-east-2.amazonaws.com/vision-agent-examples/people-example.jpeg',

components/chat/index.tsx

@@ -1,37 +1,44 @@
 'use client';
+
 import { cn } from '@/lib/utils';
 import { ChatList } from '@/components/chat/ChatList';
 import { ChatPanel } from '@/components/chat/ChatPanel';
 import { ChatScrollAnchor } from '@/components/chat/ChatScrollAnchor';
-import ImageList from './ImageList';
-import useChatWithDataset from '../../lib/hooks/useChatWithDataset';
-import { useChat } from 'ai/react';
-import { Button } from '../ui/Button';
-import ImageSelector from './ImageSelector';
+import { ChatEntity } from '@/lib/types';
+import Image from 'next/image';
+import useVisionAgent from '@/lib/hooks/useVisionAgent';
 
 export interface ChatProps extends React.ComponentProps<'div'> {
-  id: string;
+  chat: ChatEntity;
 }
 
-export function Chat({ id, className }: ChatProps) {
+export function Chat({ chat }: ChatProps) {
+  const { url, id } = chat;
   const { messages, append, reload, stop, isLoading, input, setInput } =
-    useChat();
+    useVisionAgent(chat);
 
   return (
     <>
-      <div className={cn('pb-[150px] pt-4 md:pt-10 h-full', className)}>
+      <div className={cn('pb-[150px] pt-4 md:pt-10 h-full')}>
         <div className="flex h-full">
-          <div className="w-1/2 relative border-r border-gray-400 overflow-auto">
-            {/* <ImageList /> */}
-            <ImageSelector />
+          <div className="w-1/2 relative border-r border-gray-400 overflow-auto p-4 flex items-center justify-center">
+            <div className="max-h-[600px] size-full relative">
+              <Image
+                draggable={false}
+                src={url}
+                alt={url}
+                fill
+                objectFit="contain"
+              />
+            </div>
           </div>
           <div className="w-1/2 relative overflow-auto">
-            {/* <ChatList messages={messages} /> */}
+            <ChatList messages={messages} />
             <ChatScrollAnchor trackVisibility={isLoading} />
           </div>
         </div>
       </div>
-      {/* <ChatPanel
+      <ChatPanel
         id={id}
         isLoading={isLoading}
         stop={stop}
@@ -40,7 +47,7 @@ export function Chat({ id, className }: ChatProps) {
         messages={messages}
         input={input}
         setInput={setInput}
-      /> */}
+      />
     </>
   );
 }

lib/hooks/useImageUpload.ts

@@ -2,7 +2,6 @@ import { useAtom } from 'jotai';
 import { DropzoneOptions, useDropzone } from 'react-dropzone';
 import { datasetAtom } from '../../state';
 import { toast } from 'react-hot-toast';
-import { DatasetImageEntity } from '../types';
 
 const useImageUpload = (options?: Partial<DropzoneOptions>) => {
   const { getRootProps, getInputProps, isDragActive } = useDropzone({

lib/hooks/{useChatWithDataset.ts → useVisionAgent.tsx}

@@ -1,12 +1,10 @@
 import { useChat, type Message } from 'ai/react';
-import { useAtom } from 'jotai';
 import { toast } from 'react-hot-toast';
-import { datasetAtom } from '../../state';
 import { useEffect, useState } from 'react';
-import { MessageWithSelectedDataset } from '../types';
+import { ChatEntity, MessageBase } from '../types';
 
-const useChatWithDataset = () => {
-  const [dataset] = useAtom(datasetAtom);
+const useVisionAgent = (chat: ChatEntity) => {
+  const { messages: initialMessages, id, url } = chat;
   const {
     messages,
     append,
@@ -18,11 +16,17 @@ const useChatWithDataset = () => {
     setMessages,
   } = useChat({
     sendExtraMessageFields: true,
+    api: '/api/vision-agent',
     onResponse(response) {
       if (response.status !== 200) {
         toast.error(response.statusText);
       }
     },
+    initialMessages: initialMessages,
+    body: {
+      url,
+      id,
+    },
   });
 
   const [loadingDots, setLoadingDots] = useState('');
@@ -67,43 +71,9 @@ const useChatWithDataset = () => {
       ? [...messages, assistantLoadingMessage]
       : messages;
 
-  const selectedDataset = dataset.find(entity => entity.selected)
-    ? dataset.filter(entity => entity.selected)
-    : // If there is no selected dataset, use the entire dataset
-      dataset;
-
-  const appendWithDataset: typeof append = message => {
-    // const newSystemMessage: Message = {
-    // 	id: 'fake-id',
-    // 	content:
-    // 		'For the next prompt, here are names of images provided by user, please use these name if you need reference: ' +
-    // 		selectedDataset.map(entity => entity.name).join(', '),
-    // 	role: 'system',
-    // };
-    // const newSystemMessage: Message = {
-    // 	id: 'fake-id',
-    // 	content: `For the next prompt, please use tags provided by the user to assign to corresponding images.
-    // 		For example:
-
-    // 		Input:
-    // 		red, blue, round
-
-    // 		Answer (each in a new line):
-    // 		Image 1: red\n
-    // 		Image 2: blue,round\n`,
-    // 	role: 'system',
-    // };
-    // setMessages([...messages, newSystemMessage]);
-    return append({
-      ...message,
-      // @ts-ignore this is extra fields
-      dataset: selectedDataset,
-    } as MessageWithSelectedDataset);
-  };
-
   return {
-    messages: messageWithLoading as MessageWithSelectedDataset[],
-    append: appendWithDataset,
+    messages: messageWithLoading as MessageBase[],
+    append,
     reload,
     stop,
     isLoading,
@@ -112,4 +82,4 @@ const useChatWithDataset = () => {
   };
 };
 
-export default useChatWithDataset;
+export default useVisionAgent;

lib/kv/chat.ts

@@ -5,14 +5,19 @@ import { kv } from '@vercel/kv';
 
 import { auth } from '@/auth';
 import { ChatEntity } from '@/lib/types';
+import { redirect } from 'next/navigation';
 
-export async function getKVChats() {
+async function authCheck() {
   const session = await auth();
   const email = session?.user?.email;
-
   if (!email) {
-    return [];
+    redirect('/');
   }
+  return { email, isAdmin: email.endsWith('landing.ai') };
+}
+
+export async function getKVChats() {
+  const { email } = await authCheck();
 
   try {
     const pipeline = kv.pipeline();
@@ -33,8 +38,13 @@ export async function getKVChats() {
 }
 
 export async function getKVChat(id: string) {
+  const { email, isAdmin } = await authCheck();
   const chat = await kv.hgetall<ChatEntity>(`chat:${id}`);
 
+  if (chat?.user !== email || !isAdmin) {
+    redirect('/');
+  }
+
   return chat;
 }
 

lib/types.ts

@@ -24,8 +24,9 @@ export type MessageWithSelectedDataset = Message & {
 };
 
 export type MessageBase = {
-  role: 'user' | 'system' | 'assistant';
+  role: Message['role'];
   content: string;
+  id: string;
 };
 
 export type ChatEntity = {