Merge pull request #3318 from smakosh/feat/llmgateway-ai-provider

Docs: support LLMGateway for AI chat

apps/docs/components/ai-sdk/search.tsx

@@ -15,10 +15,25 @@ import { Loader2, MessageCircleIcon, RefreshCw, SearchIcon, Send, X } from 'luci
 import { cn } from '@/lib/cn';
 import { buttonVariants } from 'fumadocs-ui/components/ui/button';
 import { useChat, type UseChatHelpers } from '@ai-sdk/react';
-import { DefaultChatTransport, type Tool, type UIToolInvocation } from 'ai';
+import {
+  DefaultChatTransport,
+  type Tool,
+  type UIMessage,
+  type UIToolInvocation,
+} from 'ai';
 import { Markdown } from '../markdown';
 import { Presence } from '@radix-ui/react-presence';
-import type { ChatUIMessage, SearchTool } from '@/lib/openrouter/route';
+
+export type ChatUIMessage = UIMessage<
+  never,
+  {
+    client: {
+      location: string;
+    };
+  }
+>;
+
+export type SearchTool = Tool<{ query: string; limit: number }, unknown[]>;
 
 const Context = createContext<{
   open: boolean;

apps/docs/components/registry/index.ts

@@ -129,7 +129,7 @@ export const registry: Registry = {
       files: [
         {
           type: 'components',
-          path: 'components/openrouter/search.tsx',
+          path: 'components/ai-sdk/search.tsx',
           target: '<dir>/ai/search.tsx',
         },
         {
@@ -142,6 +142,26 @@ export const registry: Registry = {
         flexsearch: '^0.8.212',
       },
     },
+    {
+      name: 'ai/llmgateway',
+      title: 'AI Chat (LLMGateway)',
+      description: 'Ask AI dialog for your docs, using LLMGateway',
+      files: [
+        {
+          type: 'components',
+          path: 'components/ai-sdk/search.tsx',
+          target: '<dir>/ai/search.tsx',
+        },
+        {
+          type: 'route-handler',
+          route: 'api/chat',
+          path: 'lib/llmgateway/route.ts',
+        },
+      ],
+      dependencies: {
+        flexsearch: '^0.8.212',
+      },
+    },
     {
       name: 'markdown',
       unlisted: true,

apps/docs/content/docs/(framework)/integrations/llms.mdx

@@ -386,7 +386,7 @@ const markdownUrl = `${page.url}.mdx`
 
 You can install the AI chat dialog using Fumadocs CLI.
 
-<Tabs items={['AI SDK', 'Inkeep AI']}>
+<Tabs items={['AI SDK', 'LLMGateway', 'Inkeep AI']}>
 <Tab>
 
 ```npm
@@ -404,6 +404,23 @@ You can use other models by updating the `/api/chat` route.
 </Tab>
 <Tab>
 
+```npm
+npx @fumadocs/cli add ai/llmgateway
+```
+
+It's automatically configured for [LLMGateway](https://llmgateway.io) using Vercel AI SDK, with a `/search` tool for AI.
+
+LLMGateway is an open-source API gateway that proxies requests to 300+ language models, so the same route handler can switch between providers by changing the model ID.
+
+Add your LLMGateway API key to environment variables:
+
+```dotenv
+LLM_GATEWAY_API_KEY="..."
+```
+
+</Tab>
+<Tab>
+
 ```npm
 npx @fumadocs/cli add ai/inkeep
 ```

apps/docs/lib/llmgateway/route.ts

@@ -0,0 +1,121 @@
+import { createLLMGateway } from '@llmgateway/ai-sdk-provider';
+import { convertToModelMessages, stepCountIs, streamText, tool, type UIMessage } from 'ai';
+import { z } from 'zod';
+import { source } from '@/lib/source';
+import { Document, type DocumentData } from 'flexsearch';
+import { $routeHandler } from 'fuma-cli/macros/route-handler';
+
+interface CustomDocument extends DocumentData {
+  url: string;
+  title: string;
+  description: string;
+  content: string;
+}
+
+export type ChatUIMessage = UIMessage<
+  never,
+  {
+    client: {
+      location: string;
+    };
+  }
+>;
+
+const searchServer = createSearchServer();
+
+async function createSearchServer() {
+  const search = new Document<CustomDocument>({
+    document: {
+      id: 'url',
+      index: ['title', 'description', 'content'],
+      store: true,
+    },
+  });
+
+  const docs = await chunkedAll(
+    source.getPages().map(async (page) => {
+      if (!('getText' in page.data)) return null;
+
+      return {
+        title: page.data.title,
+        description: page.data.description,
+        url: page.url,
+        content: await page.data.getText('processed'),
+      } as CustomDocument;
+    }),
+  );
+
+  for (const doc of docs) {
+    if (doc) search.add(doc);
+  }
+
+  return search;
+}
+
+async function chunkedAll<O>(promises: Promise<O>[]): Promise<O[]> {
+  const SIZE = 50;
+  const out: O[] = [];
+  for (let i = 0; i < promises.length; i += SIZE) {
+    out.push(...(await Promise.all(promises.slice(i, i + SIZE))));
+  }
+  return out;
+}
+
+const llmgateway = createLLMGateway({
+  apiKey: process.env.LLM_GATEWAY_API_KEY,
+});
+
+/** System prompt, you can update it to provide more specific information */
+const systemPrompt = [
+  'You are an AI assistant for a documentation site.',
+  'Use the `search` tool to retrieve relevant docs context before answering when needed.',
+  'The `search` tool returns raw JSON results from documentation. Use those results to ground your answer and cite sources as markdown links using the document `url` field when available.',
+  'If you cannot find the answer in search results, say you do not know and suggest a better search query.',
+].join('\n');
+
+export const handler = $routeHandler(
+  {
+    methods: ['POST'],
+    params: [],
+  },
+  async (req) => {
+    const reqJson = await req.json();
+
+    const result = streamText({
+      model: llmgateway.chat(process.env.LLM_GATEWAY_MODEL ?? 'anthropic/claude-3.5-sonnet'),
+      stopWhen: stepCountIs(5),
+      tools: {
+        search: searchTool,
+      },
+      messages: [
+        { role: 'system', content: systemPrompt },
+        ...(await convertToModelMessages<ChatUIMessage>(reqJson.messages ?? [], {
+          convertDataPart(part) {
+            if (part.type === 'data-client')
+              return {
+                type: 'text',
+                text: `[Client Context: ${JSON.stringify(part.data)}]`,
+              };
+          },
+        })),
+      ],
+      toolChoice: 'auto',
+    });
+
+    return result.toUIMessageStreamResponse();
+  },
+);
+
+export type SearchTool = typeof searchTool;
+
+const searchTool = tool({
+  description: 'Search the docs content and return raw JSON results.',
+  inputSchema: z.object({
+    query: z.string(),
+    limit: z.number().int().min(1).max(100).default(10),
+  }),
+  async execute({ query, limit }) {
+    const search = await searchServer;
+    return await search.searchAsync(query, { limit, merge: true, enrich: true });
+  },
+});

apps/docs/package.json

@@ -18,6 +18,7 @@
     "@floating-ui/react-dom": "^2.1.8",
     "@fumadocs/story": "workspace:*",
     "@mixedbread/sdk": "^0.71.0",
+    "@llmgateway/ai-sdk-provider": "^3.8.0",
     "@openrouter/ai-sdk-provider": "^2.9.0",
     "@orama/core": "^1.2.19",
     "@orama/orama": "^3.1.18",

packages/create-app/src/bin.ts

@@ -35,7 +35,13 @@ const command = program
   .addOption(
     new Option('--og-image <name>', 'configure OG image generation').choices(['next-og', 'takumi']),
   )
-  .addOption(new Option('--ai-chat <name>', 'configure AI chat').choices(['openrouter', 'inkeep']))
+  .addOption(
+    new Option('--ai-chat <name>', 'configure AI chat').choices([
+      'openrouter',
+      'llmgateway',
+      'inkeep',
+    ]),
+  )
   .addOption(
     new Option('--template <name>', 'choose a template').choices(
       templates.map((item) => item.value),
@@ -175,7 +181,7 @@ async function main(): Promise<void> {
         )
           return false;
 
-        return select<false | 'openrouter' | 'inkeep'>({
+        return select<false | 'openrouter' | 'llmgateway' | 'inkeep'>({
           message: 'Configure AI Chat?',
           options: [
             {
@@ -187,6 +193,11 @@ async function main(): Promise<void> {
               label: 'AI SDK',
               hint: 'default to OpenRouter',
             },
+            {
+              value: 'llmgateway',
+              label: 'LLMGateway',
+              hint: 'open-source LLM gateway, API key required',
+            },
             {
               value: 'inkeep',
               label: 'Inkeep AI',

packages/create-app/src/plugins/ai.ts

@@ -7,7 +7,13 @@ import { FumadocsComponentInstaller } from '@fumadocs/cli/registry/installer';
 import { HttpRegistryConnector } from 'fuma-cli/registry/connector';
 import { getDefaultConfig } from '@fumadocs/cli/config';
 
-export function ai(provider: 'openrouter' | 'inkeep'): TemplatePlugin {
+const envKey: Record<'openrouter' | 'llmgateway' | 'inkeep', string> = {
+  openrouter: 'OPENROUTER_API_KEY',
+  llmgateway: 'LLM_GATEWAY_API_KEY',
+  inkeep: 'INKEEP_API_KEY',
+};
+
+export function ai(provider: 'openrouter' | 'llmgateway' | 'inkeep'): TemplatePlugin {
   return {
     async afterWrite() {
       const config = await getDefaultConfig(this.dest);
@@ -22,10 +28,7 @@ export function ai(provider: 'openrouter' | 'inkeep'): TemplatePlugin {
         if (deps.hasRequired()) await deps.writeRequired();
 
         await addAIChat(this);
-        await fs.writeFile(
-          path.join(this.dest, '.env.local'),
-          provider === 'openrouter' ? 'OPENROUTER_API_KEY=' : 'INKEEP_API_KEY=',
-        );
+        await fs.writeFile(path.join(this.dest, '.env.local'), `${envKey[provider]}=`);
       } catch (e) {
         console.error(e);
       }

pnpm-lock.yaml

@@ -47,6 +47,9 @@ importers:
       '@fumadocs/story':
         specifier: workspace:*
         version: link:../../packages/story
+      '@llmgateway/ai-sdk-provider':
+        specifier: ^3.8.0
+        version: 3.8.0([email protected]([email protected]))([email protected])
       '@mixedbread/sdk':
         specifier: ^0.71.0
         version: 0.71.0
@@ -5866,6 +5869,16 @@ packages:
   '@lezer/[email protected]':
     resolution: {integrity: sha512-2lrrHqxalACEbxIbsjhqGpSW8kWpUKuY6RHgnSAFZa6qK62wvnPxA8hGOwOoDbwHcOFs5M4o27mjGu+P7TvBmw==}
 
+  '@llmgateway/[email protected]':
+    resolution: {integrity: sha512-3Yr5evl5b2QtXAPSl+kHLN0Zjni+MFAv0h0bW6nQrbqPaEmC28P4gNnRNyMkrpsGptPB7ahTl66ZXg+wi0NunA==}
+    engines: {node: '>=18'}
+    peerDependencies:
+      ai: ^6.0.0
+      zod: ^3.24.1 || ^v4
+
+  '@llmgateway/[email protected]':
+    resolution: {integrity: sha512-3A81z2BYgnBOtqjSFoW1AnkgUcD8niQcev2G5lTkC8PeDFb1hLIuRn5N6OjDSFcGD0f9NMd1+sbgRObkmGhc5Q==}
+
   '@manypkg/[email protected]':
     resolution: {integrity: sha512-mki5uBvhHzO8kYYix/WRy2WX8S3B5wdVSc9D6KcU5lQNglP2yt58/VfLuAK49glRXChosY8ap2oJ1qgma3GUVA==}
 
@@ -18094,6 +18107,14 @@ snapshots:
       '@lezer/highlight': 1.2.3
       '@lezer/lr': 1.4.10
 
+  '@llmgateway/[email protected]([email protected]([email protected]))([email protected])':
+    dependencies:
+      '@llmgateway/models': 1.97.1
+      ai: 6.0.191([email protected])
+      zod: 4.4.3
+
+  '@llmgateway/[email protected]': {}
+
   '@manypkg/[email protected]':
     dependencies:
       '@babel/runtime': 7.29.2