Feat/web search integration

app/(chat)/api/chat/route.ts

@@ -21,7 +21,7 @@ import type { VisibilityType } from "@/components/visibility-selector";
 import { entitlementsByUserType } from "@/lib/ai/entitlements";
 import type { ChatModel } from "@/lib/ai/models";
 import { type RequestHints, systemPrompt } from "@/lib/ai/prompts";
-import { myProvider } from "@/lib/ai/providers";
+import { buildProviderOptions, myProvider } from "@/lib/ai/providers";
 import { createDocument } from "@/lib/ai/tools/create-document";
 import { getWeather } from "@/lib/ai/tools/get-weather";
 import { requestSuggestions } from "@/lib/ai/tools/request-suggestions";
@@ -39,7 +39,7 @@ import {
 } from "@/lib/db/queries";
 import type { DBMessage } from "@/lib/db/schema";
 import { ChatSDKError } from "@/lib/errors";
-import type { ChatMessage } from "@/lib/types";
+import type { ChatMessage, SearchSource } from "@/lib/types";
 import type { AppUsage } from "@/lib/usage";
 import { convertToUIMessages, generateUUID } from "@/lib/utils";
 import { generateTitleFromUserMessage } from "../../actions";
@@ -168,6 +168,7 @@ export async function POST(request: Request) {
           parts: message.parts,
           attachments: [],
           createdAt: new Date(),
+          searchResults: null,
         },
       ],
     });
@@ -176,6 +177,7 @@ export async function POST(request: Request) {
     await createStreamId({ streamId, chatId: id });
 
     let finalMergedUsage: AppUsage | undefined;
+    let searchResultsData: SearchSource[] | null = null;
 
     const stream = createUIMessageStream({
       execute: ({ writer: dataStream }) => {
@@ -203,6 +205,7 @@ export async function POST(request: Request) {
               dataStream,
             }),
           },
+          providerOptions: buildProviderOptions(selectedChatModel),
           experimental_telemetry: {
             isEnabled: isProductionEnvironment,
             functionId: "stream-text",
@@ -241,6 +244,46 @@ export async function POST(request: Request) {
           },
         });
 
+        (async () => {
+          try {
+            const sources = await result.sources;
+
+            if (sources && sources.length > 0) {
+
+              const validSources = sources.filter(
+                (source): source is typeof source & { url: string } =>
+                  "url" in source && typeof source.url === "string"
+              );
+
+              // Deduplicate by URL - workaround for AI SDK duplicate sources issue
+              // TODO: Remove once https://github.com/vercel/ai/issues/9771 is fixed
+              const seenUrls = new Set<string>();
+              const uniqueSources = validSources.filter((source) => {
+                if (seenUrls.has(source.url)) {
+                  return false;
+                }
+                seenUrls.add(source.url);
+                return true;
+              });
+
+              if (uniqueSources.length > 0) {
+                searchResultsData = uniqueSources.map((source) => ({
+                  title: source.title || source.url,
+                  url: source.url,
+                  favicon: `https://www.google.com/s2/favicons?domain=${new URL(source.url).hostname}&sz=32`,
+                }));
+
+                dataStream.write({
+                  type: "data-searchResults",
+                  data: searchResultsData,
+                });
+              }
+            }
+          } catch (error) {
+            console.error("Error processing search results:", error);
+          }
+        })();
+
         result.consumeStream();
 
         dataStream.merge(
@@ -259,6 +302,9 @@ export async function POST(request: Request) {
             createdAt: new Date(),
             attachments: [],
             chatId: id,
+            // Add search results to assistant messages
+            searchResults:
+              currentMessage.role === "assistant" ? searchResultsData : null,
           })),
         });
 

components/message.tsx

@@ -24,6 +24,7 @@ import { MessageEditor } from "./message-editor";
 import { MessageReasoning } from "./message-reasoning";
 import { PreviewAttachment } from "./preview-attachment";
 import { Weather } from "./weather";
+import { SearchResults, SearchingIndicator } from "./search-results";
 
 const PurePreviewMessage = ({
   chatId,
@@ -52,6 +53,18 @@ const PurePreviewMessage = ({
 
   useDataStream();
 
+  // Check if search results are in the message parts
+  const searchResultsPart = message.parts?.find(
+    (part) => part.type === "data-searchResults"
+  );
+
+  const searchResults =
+    searchResultsPart && "data" in searchResultsPart && searchResultsPart.data
+      ? searchResultsPart.data
+      : undefined;
+
+  // Check if we should show searching indicator (message is loading and assistant role)
+  const isSearching = isLoading && message.role === "assistant" && (!searchResults || searchResults.length === 0);
   return (
     <motion.div
       animate={{ opacity: 1 }}
@@ -106,6 +119,16 @@ const PurePreviewMessage = ({
             </div>
           )}
 
+          {message.role === "assistant" && (isSearching || searchResults) && (
+            <div className="w-full">
+              {searchResults ? (
+                <SearchResults searchResult={searchResults} />
+              ) : (
+                <SearchingIndicator />
+              )}
+            </div>
+          )}
+
           {message.parts?.map((part, index) => {
             const { type } = part;
             const key = `message-${message.id}-part-${index}`;

components/search-results.tsx

@@ -0,0 +1,92 @@
+"use client";
+
+import { GlobeIcon } from "lucide-react";
+import type { SearchSource } from "@/lib/types";
+import {
+  Collapsible,
+  CollapsibleContent,
+  CollapsibleTrigger,
+} from "@/components/ui/collapsible";
+import { cn } from "@/lib/utils";
+import { ChevronDownIcon } from "lucide-react";
+import { useState } from "react";
+import Link from "next/link";
+
+
+export function SearchResults({ searchResult }: { searchResult: SearchSource[] }) {
+  const [isOpen, setIsOpen] = useState(false);
+
+  if (!searchResult || !Array.isArray(searchResult) || searchResult.length === 0) {
+    return null;
+  }
+
+  return (
+    <Collapsible
+      className="not-prose w-full"
+      onOpenChange={setIsOpen}
+      open={isOpen}
+    >
+      <CollapsibleTrigger className="flex w-full min-w-0 items-center gap-2 text-muted-foreground text-xs transition-colors hover:text-foreground">
+        <GlobeIcon className="size-4 shrink-0" />
+        <span className="font-medium">
+          Searched web
+        </span>
+        <div className="flex shrink-0 items-center gap-1.5">
+          <span className="text-muted-foreground">
+            {searchResult.length} {searchResult.length === 1 ? "result" : "results"}
+          </span>
+          <ChevronDownIcon
+            className={cn(
+              "size-3 text-muted-foreground transition-transform",
+              isOpen && "rotate-180"
+            )}
+          />
+        </div>
+      </CollapsibleTrigger>
+
+      <CollapsibleContent className="mt-3 data-[state=closed]:fade-out-0 data-[state=closed]:slide-out-to-top-2 data-[state=open]:slide-in-from-top-2 outline-hidden data-[state=closed]:animate-out data-[state=open]:animate-in">
+        <div className="grid gap-1.5">
+          {searchResult.map((source, index) => (
+            <Link
+              className="flex items-center gap-2 text-xs transition-colors hover:text-foreground min-w-0 overflow-hidden animate-in fade-in-0 slide-in-from-top-1"
+              href={source.url}
+              key={index}
+              rel="noopener noreferrer"
+              target="_blank"
+            >
+              <div className="flex size-4 shrink-0 items-center justify-center overflow-hidden rounded">
+                {source.favicon ? (
+                  <img
+                    alt=""
+                    className="size-4"
+                    src={source.favicon}
+                    onError={(e) => {
+                      e.currentTarget.style.display = "none";
+                    }}
+                  />
+                ) : (
+                  <GlobeIcon className="size-3 text-muted-foreground" />
+                )}
+              </div>
+              <span className="truncate font-medium text-foreground">
+                {source.title}
+              </span>
+              <span className="shrink-0 truncate text-muted-foreground max-w-[120px]">
+                {new URL(source.url).hostname.replace('www.', '')}
+              </span>
+            </Link>
+          ))}
+        </div>
+      </CollapsibleContent>
+    </Collapsible>
+  );
+}
+
+export function SearchingIndicator() {
+  return (
+    <div className="flex w-full items-center gap-2 text-muted-foreground text-xs animate-in fade-in-0">
+      <GlobeIcon className="size-4 shrink-0 animate-pulse" />
+      <span>Searching the web</span>
+    </div>
+  );
+}

lib/ai/models.ts

@@ -4,18 +4,36 @@ export type ChatModel = {
   id: string;
   name: string;
   description: string;
+  capabilities?: {
+    webSearch?: boolean;
+  };
 };
 
 export const chatModels: ChatModel[] = [
   {
     id: "chat-model",
     name: "Grok Vision",
     description: "Advanced multimodal model with vision and text capabilities",
+    capabilities: {
+      webSearch: true,
+    },
   },
   {
     id: "chat-model-reasoning",
     name: "Grok Reasoning",
     description:
       "Uses advanced chain-of-thought reasoning for complex problems",
+    capabilities: {
+      webSearch: true,
+    },
   },
 ];
+
+export function getModelCapabilities(modelId: string) {
+  const model = chatModels.find((m) => m.id === modelId);
+  return model?.capabilities || {};
+}
+
+export function supportsWebSearch(modelId: string): boolean {
+  return getModelCapabilities(modelId).webSearch === true;
+}

lib/ai/providers.ts

@@ -5,6 +5,7 @@ import {
   wrapLanguageModel,
 } from "ai";
 import { isTestEnvironment } from "../constants";
+import { supportsWebSearch } from "./models";
 
 export const myProvider = isTestEnvironment
   ? (() => {
@@ -34,3 +35,24 @@ export const myProvider = isTestEnvironment
         "artifact-model": gateway.languageModel("xai/grok-2-1212"),
       },
     });
+
+/**
+ * Builds provider-specific options based on model capabilities
+ */
+export function buildProviderOptions(modelId: string) {
+  if (!supportsWebSearch(modelId)) {
+    return;
+  }
+
+  const options = {
+    xai: {
+      searchParameters: {
+        mode: "auto" as const,
+        returnCitations: true,
+        maxSearchResults: 10,
+      },
+    },
+  };
+
+  return options;
+}

lib/db/migrations/0008_left_giant_girl.sql

@@ -0,0 +1 @@
+ALTER TABLE "Message_v2" ADD COLUMN "searchResults" jsonb;