Merge pull request #1196 from trycompai/mariano/qa-endpoint

[dev] [Marfuen] mariano/qa-endpoint

Author: Marfuen

apps/app/package.json

@@ -2,10 +2,12 @@
   "name": "@comp/app",
   "version": "0.1.0",
   "dependencies": {
-    "@ai-sdk/groq": "^1.2.8",
-    "@ai-sdk/openai": "^1.3.19",
-    "@ai-sdk/provider": "^1.1.3",
-    "@ai-sdk/react": "^1.2.9",
+    "@ai-sdk/anthropic": "^2.0.0",
+    "@ai-sdk/groq": "^2.0.0",
+    "@ai-sdk/openai": "^2.0.0",
+    "@ai-sdk/provider": "^2.0.0",
+    "@ai-sdk/react": "^2.0.0",
+    "@ai-sdk/rsc": "^1.0.0",
     "@aws-sdk/client-s3": "^3.806.0",
     "@aws-sdk/client-sts": "^3.808.0",
     "@aws-sdk/s3-request-presigner": "^3.832.0",
@@ -50,7 +52,7 @@
     "@uploadthing/react": "^7.3.0",
     "@upstash/ratelimit": "^2.0.5",
     "@vercel/sdk": "^1.7.1",
-    "ai": "^4.3.16",
+    "ai": "^5.0.0",
     "axios": "^1.9.0",
     "better-auth": "^1.2.8",
     "canvas-confetti": "^1.9.3",

apps/app/src/app/api/chat/route.ts

@@ -1,14 +1,13 @@
 import { tools } from '@/data/tools';
-import { model, type modelID } from '@/hooks/ai/providers';
 import { auth } from '@/utils/auth';
-import { type UIMessage, streamText } from 'ai';
+import { groq } from '@ai-sdk/groq';
+import { type UIMessage, convertToModelMessages, streamText } from 'ai';
 import { headers } from 'next/headers';
 
 export const maxDuration = 30;
 
 export async function POST(req: Request) {
-  const { messages, selectedModel }: { messages: UIMessage[]; selectedModel: modelID } =
-    await req.json();
+  const { messages }: { messages: UIMessage[] } = await req.json();
 
   const session = await auth.api.getSession({
     headers: await headers(),
@@ -31,11 +30,11 @@ export async function POST(req: Request) {
 `;
 
   const result = streamText({
-    model: model.languageModel(selectedModel),
+    model: groq('deepseek-r1-distill-llama-70b'),
     system: systemPrompt,
-    messages,
+    messages: convertToModelMessages(messages),
     tools,
   });
 
-  return result.toDataStreamResponse({ sendReasoning: true });
+  return result.toUIMessageStreamResponse();
 }

apps/app/src/components/ai/chat-text-area.tsx

@@ -1,4 +1,3 @@
-import type { modelID } from '@/hooks/ai/providers';
 import { Icons } from '@comp/ui/icons';
 import { Popover, PopoverContent, PopoverTrigger } from '@comp/ui/popover';
 import { Textarea as ShadcnTextarea } from '@comp/ui/textarea';
@@ -10,8 +9,6 @@ interface InputProps {
   isLoading: boolean;
   status: string;
   stop: () => void;
-  selectedModel: modelID;
-  setSelectedModel: (model: modelID) => void;
 }
 
 export const ChatTextarea = ({ input, handleInputChange, isLoading }: InputProps) => {

apps/app/src/components/ai/chat.tsx

@@ -1,9 +1,9 @@
 'use client';
 
-import type { modelID } from '@/hooks/ai/providers';
 import { useSession } from '@/utils/auth-client';
 import { useChat } from '@ai-sdk/react';
 import { ScrollArea } from '@comp/ui/scroll-area';
+import { DefaultChatTransport, lastAssistantMessageIsCompleteWithToolCalls } from 'ai';
 import { useState } from 'react';
 import { ChatEmpty } from './chat-empty';
 import { ChatTextarea } from './chat-text-area';
@@ -12,13 +12,15 @@ import { Messages } from './messages';
 export default function Chat() {
   const { data: session } = useSession();
 
-  const [selectedModel, setSelectedModel] = useState<modelID>('deepseek-r1-distill-llama-70b');
+  const [input, setInput] = useState('');
 
-  const { messages, input, handleInputChange, handleSubmit, error, status, stop } = useChat({
-    maxSteps: 5,
-    body: {
-      selectedModel,
-    },
+  const { messages, sendMessage, addToolResult, error, status, stop } = useChat({
+    transport: new DefaultChatTransport({
+      api: '/api/chat',
+    }),
+
+    // Automatically submit when all server-side tool calls are complete
+    sendAutomaticallyWhen: lastAssistantMessageIsCompleteWithToolCalls,
   });
 
   const isLoading = status === 'streaming' || status === 'submitted';
@@ -37,11 +39,17 @@ export default function Chat() {
         )}
       </ScrollArea>
 
-      <form onSubmit={handleSubmit}>
+      <form
+        onSubmit={(e) => {
+          e.preventDefault();
+          if (input.trim()) {
+            sendMessage({ text: input });
+            setInput('');
+          }
+        }}
+      >
         <ChatTextarea
-          selectedModel={selectedModel}
-          setSelectedModel={setSelectedModel}
-          handleInputChange={handleInputChange}
+          handleInputChange={(e) => setInput(e.target.value)}
           input={input}
           isLoading={isLoading}
           status={status}

apps/app/src/components/ai/message.tsx

@@ -1,9 +1,10 @@
 'use client';
 
 import { useStreamableText } from '@/hooks/use-streamable-text';
+import { type StreamableValue } from '@ai-sdk/rsc';
 import { cn } from '@comp/ui/cn';
-import type { Message as TMessage } from 'ai';
-import type { StreamableValue } from 'ai/rsc';
+import type { UIMessage } from 'ai';
+
 import equal from 'fast-deep-equal';
 import { ChevronDownIcon, ChevronUpIcon } from 'lucide-react';
 import { AnimatePresence, motion } from 'motion/react';
@@ -26,8 +27,8 @@ interface ExtendedToolInvocation extends ToolInvocation {
 
 interface ReasoningPart {
   type: 'reasoning';
-  reasoning: string;
-  details: Array<{ type: 'text'; text: string }>;
+  text: string;
+  details?: Array<{ type: 'text'; text: string }>;
 }
 
 interface ReasoningMessagePartProps {
@@ -117,13 +118,13 @@ export function ReasoningMessagePart({ part, isReasoning }: ReasoningMessagePart
             variants={variants}
             transition={{ duration: 0.2, ease: 'easeInOut' }}
           >
-            {part.details.map((detail) =>
+            {part.details?.map((detail) =>
               detail.type === 'text' ? (
                 <StreamableMarkdown key={detail.text} text={detail.text} />
               ) : (
                 '<redacted>'
               ),
-            )}
+            ) || <StreamableMarkdown text={part.text} />}
           </motion.div>
         )}
       </AnimatePresence>
@@ -159,7 +160,7 @@ const PurePreviewMessage = ({
   isLatestMessage,
   status,
 }: {
-  message: TMessage;
+  message: UIMessage;
   isLoading: boolean;
   status: 'error' | 'submitted' | 'streaming' | 'ready';
   isLatestMessage: boolean;
@@ -180,7 +181,12 @@ const PurePreviewMessage = ({
           )}
         >
           <div className="flex w-full flex-col space-y-4">
-            {message.parts?.map((part, i) => {
+            {message.parts?.map((part: any, i: number) => {
+              // Skip invalid parts
+              if (!part || !part.type) {
+                return null;
+              }
+
               switch (part.type) {
                 case 'text':
                   return message.role === 'user' ? (
@@ -196,7 +202,7 @@ const PurePreviewMessage = ({
                             message.role === 'user',
                         })}
                       >
-                        <StreamableMarkdown text={part.text} />
+                        <StreamableMarkdown text={part.text || ''} />
                       </div>
                     </motion.div>
                   ) : (
@@ -207,7 +213,7 @@ const PurePreviewMessage = ({
                       className="flex w-full flex-row items-start gap-2 pb-2"
                     >
                       <BotCard key={`message-${message.id}-part-${i}`}>
-                        <StreamableMarkdown text={part.text} />
+                        <StreamableMarkdown text={part.text || ''} />
                       </BotCard>
                     </motion.div>
                   );
@@ -216,8 +222,7 @@ const PurePreviewMessage = ({
                   return (
                     <ReasoningMessagePart
                       key={`message-${message.id}-${i}`}
-                      // @ts-expect-error part
-                      part={part}
+                      part={part as ReasoningPart}
                       isReasoning={
                         (message.parts &&
                           status === 'streaming' &&
@@ -227,7 +232,9 @@ const PurePreviewMessage = ({
                     />
                   );
                 }
+
                 default:
+                  // Skip tool parts - only show final text response
                   return null;
               }
             })}
@@ -303,8 +310,6 @@ export function UserMessage({ content }: { content: string }) {
 export const Message = memo(PurePreviewMessage, (prevProps, nextProps) => {
   if (prevProps.status !== nextProps.status) return false;
 
-  if (prevProps.message.annotations !== nextProps.message.annotations) return false;
-
   if (!equal(prevProps.message.parts, nextProps.message.parts)) return false;
 
   return true;

apps/app/src/components/ai/messages.tsx

@@ -1,13 +1,13 @@
 import { useScrollToBottom } from '@/hooks/use-scroll-to-bottom';
-import type { Message as TMessage } from 'ai';
+import type { UIMessage } from 'ai';
 import { Message } from './message';
 
 export const Messages = ({
   messages,
   isLoading,
   status,
 }: {
-  messages: TMessage[];
+  messages: UIMessage[];
   isLoading: boolean;
   status: 'error' | 'submitted' | 'streaming' | 'ready';
 }) => {

apps/app/src/data/tools/organization.ts

@@ -1,6 +1,5 @@
 import { auth } from '@/utils/auth';
 import { db } from '@db';
-import { tool } from 'ai';
 import { headers } from 'next/headers';
 import { z } from 'zod';
 
@@ -10,9 +9,9 @@ export function getOrganizationTools() {
   };
 }
 
-export const findOrganization = tool({
+export const findOrganization = {
   description: "Find the users organization and it's details",
-  parameters: z.object({}),
+  inputSchema: z.object({}),
   execute: async () => {
     const session = await auth.api.getSession({
       headers: await headers(),
@@ -40,4 +39,4 @@ export const findOrganization = tool({
       organization: org,
     };
   },
-});
+};

apps/app/src/data/tools/policies.ts

@@ -1,6 +1,5 @@
 import { auth } from '@/utils/auth';
 import { db } from '@db';
-import { tool } from 'ai';
 import { headers } from 'next/headers';
 import { z } from 'zod';
 
@@ -11,12 +10,12 @@ export function getPolicyTools() {
   };
 }
 
-export const getPolicies = tool({
+export const getPolicies = {
   description: 'Get all policies for the organization',
-  parameters: z.object({
+  inputSchema: z.object({
     status: z.enum(['draft', 'published']).optional(),
   }),
-  execute: async ({ status }) => {
+  execute: async ({ status }: { status?: 'draft' | 'published' }) => {
     const session = await auth.api.getSession({
       headers: await headers(),
     });
@@ -49,15 +48,15 @@ export const getPolicies = tool({
       policies,
     };
   },
-});
+};
 
-export const getPolicyContent = tool({
+export const getPolicyContent = {
   description:
     'Get the content of a specific policy by id. We can only acquire the policy id by running the getPolicies tool first.',
-  parameters: z.object({
+  inputSchema: z.object({
     id: z.string(),
   }),
-  execute: async ({ id }) => {
+  execute: async ({ id }: { id: string }) => {
     const session = await auth.api.getSession({
       headers: await headers(),
     });
@@ -84,4 +83,4 @@ export const getPolicyContent = tool({
       content: policy?.content,
     };
   },
-});
+};

apps/app/src/data/tools/risks-tool.ts

@@ -1,6 +1,5 @@
 import { auth } from '@/utils/auth';
 import { db, Departments, RiskCategory, RiskStatus } from '@db';
-import { tool } from 'ai';
 import { headers } from 'next/headers';
 import { z } from 'zod';
 
@@ -11,15 +10,25 @@ export function getRiskTools() {
   };
 }
 
-export const getRisks = tool({
+export const getRisks = {
   description: 'Get risks for the organization',
-  parameters: z.object({
+  inputSchema: z.object({
     status: z.enum(Object.values(RiskStatus) as [RiskStatus, ...RiskStatus[]]).optional(),
     department: z.enum(Object.values(Departments) as [Departments, ...Departments[]]).optional(),
     category: z.enum(Object.values(RiskCategory) as [RiskCategory, ...RiskCategory[]]).optional(),
     owner: z.string().optional(),
   }),
-  execute: async ({ status, department, category, owner }) => {
+  execute: async ({
+    status,
+    department,
+    category,
+    owner,
+  }: {
+    status?: RiskStatus;
+    department?: Departments;
+    category?: RiskCategory;
+    owner?: string;
+  }) => {
     const session = await auth.api.getSession({
       headers: await headers(),
     });
@@ -54,14 +63,14 @@ export const getRisks = tool({
       risks,
     };
   },
-});
+};
 
-export const getRiskById = tool({
+export const getRiskById = {
   description: 'Get a risk by id',
-  parameters: z.object({
+  inputSchema: z.object({
     id: z.string(),
   }),
-  execute: async ({ id }) => {
+  execute: async ({ id }: { id: string }) => {
     const session = await auth.api.getSession({
       headers: await headers(),
     });
@@ -85,4 +94,4 @@ export const getRiskById = tool({
       risk,
     };
   },
-});
+};