cleanup

src/app/search/web/page.jsx

@@ -42,5 +42,57 @@ export default function WebSearchPage({ searchParams }) {
 
   console.log(aiResponse);
 
-  return <>{aiResponse}</>;
+  return (
+      <div className="flex flex-row">
+        <MemoizedReactMarkdown
+          className="prose dark:prose-invert flex-1"
+          components={{
+            code({ node, inline, className, children, ...props }) {
+              if (children.length) {
+                if (children[0] == '▍') {
+                  return <span className="animate-pulse cursor-default mt-1">▍</span>
+                }
+                children[0] = (children[0]).replace("`▍`", "▍")
+              }
+              const match = /language-(\w+)/.exec(className || '');
+              return !inline ? (
+                <CodeBlock
+                  key={Math.random()}
+                  language={(match && match[1]) || ''}
+                  value={String(children).replace(/\n$/, '')}
+                  {...props}
+                />
+              ) : (
+                <code className={className} {...props}>
+                  {children}
+                </code>
+              );
+            },
+            table({ children }) {
+              return (
+                <table className="border-collapse border border-black px-3 py-1 dark:border-white">
+                  {children}
+                </table>
+              );
+            },
+            th({ children }) {
+              return (
+                <th className="break-words border border-black bg-gray-500 px-3 py-1 text-white dark:border-white">
+                  {children}
+                </th>
+              );
+            },
+            td({ children }) {
+              return (
+                <td className="break-words border border-black px-3 py-1 dark:border-white">
+                  {children}
+                </td>
+              );
+            },
+          }}
+        >
+          {aiResponse}
+        </MemoizedReactMarkdown>
+      </div>
+    );
 }

src/pages/api/llm.js

@@ -1,92 +1,83 @@
-import { Configuration, OpenAIApi } from "openai";
 import { GoogleCustomSearch } from "openai-function-calling-tools";
+import { DEFAULT_SYSTEM_PROMPT, DEFAULT_TEMPERATURE } from '@/utils/app/const';
+import { LLMError, LLMStream } from './stream';
 
-export default function handler(req, res) {
-  if (req.method !== 'POST') {
-    res.status(405).send({ error: 'Method Not Allowed', method: req.method });
-    return;
-  }
+// @ts-expect-error
+import wasm from '../../node_modules/@dqbd/tiktoken/lite/tiktoken_bg.wasm?module';
 
-  const QUESTION = req.body.question;
+import tiktokenModel from '@dqbd/tiktoken/encoders/cl100k_base.json';
+import { Tiktoken, init } from '@dqbd/tiktoken/lite/init';
 
-  if (!QUESTION) {
-    res.status(400).send({ error: 'Question is missing in request body' });
-    return;
-  }
+export const config = {
+  runtime: 'edge',
+};
+
+const handler = async (req) => {
+  try {
+    const { question } = (await req.json());
 
-  const configuration = new Configuration({
-    apiKey: process.env.OPENAI_API_KEY,
-  });
-  const openai = new OpenAIApi(configuration);
-
-  const messages = [
-    {
-      role: "user",
-      content: QUESTION,
-    },
-  ];
-
-  const googleCustomSearch = new GoogleCustomSearch({
-    apiKey: process.env.API_KEY,
-    googleCSEId: process.env.CONTEXT_KEY,
-  });
-
-  const functions = {
-    googleCustomSearch,
-  };
-
-  const getCompletion = async (messages) => {
-    const response = await openai.createChatCompletion({
-      model: "gpt-3.5-turbo-0613",
-      messages,
-      functions: [googleCustomSearchSchema],
-      temperature: 0,
-      stream: true
+    await init((imports) => WebAssembly.instantiate(wasm, imports));
+
+    const googleCustomSearch = new GoogleCustomSearch({
+        apiKey: process.env.API_KEY,
+        googleCSEId: process.env.CONTEXT_KEY,
     });
 
-    return response;
-  };
+    const messages = [
+        {
+          role: "user",
+          content: question,
+        },
+      ];
+    
+    const functions = {
+        googleCustomSearch,
+    };
+    const encoding = new Tiktoken(
+      tiktokenModel.bpe_ranks,
+      tiktokenModel.special_tokens,
+      tiktokenModel.pat_str,
+    );
+
+    let promptToSend = question;
+    if (!promptToSend) {
+      promptToSend = DEFAULT_SYSTEM_PROMPT;
+    }
+
+    let temperatureToUse = temperature;
+    if (temperatureToUse == null) {
+      temperatureToUse = DEFAULT_TEMPERATURE;
+    }
 
-  res.setHeader('Content-Type', 'text/event-stream');
-  res.setHeader('Cache-Control', 'no-cache');
-  res.setHeader('Connection', 'keep-alive');
+    const prompt_tokens = encoding.encode(promptToSend);
 
-  const processCompletions = async () => {
-    let response;
+    let tokenCount = prompt_tokens.length;
+    let messagesToSend = [];
 
-    while (true) {
-      response = await getCompletion(messages);
+    for (let i = messages.length - 1; i >= 0; i--) {
+      const message = messages[i];
+      const tokens = encoding.encode(message.content);
 
-      if (response.data.choices[0].finish_reason === "stop") {
-        res.write(`data: ${JSON.stringify({ result: response.data.choices[0].message.content })}\n\n`);
+      if (tokenCount + tokens.length + 1000 > model.tokenLimit) {
         break;
-      } else if (response.data.choices[0].finish_reason === "function_call") {
-        const fnName = response.data.choices[0].message.function_call.name;
-        const args = response.data.choices[0].message.function_call.arguments;
-
-        const fn = functions[fnName];
-        const result = await fn(...Object.values(JSON.parse(args)));
-
-        messages.push({
-          role: "assistant",
-          content: "",
-          function_call: {
-            name: fnName,
-            arguments: args,
-          },
-        });
-
-        messages.push({
-          role: "function",
-          name: fnName,
-          content: JSON.stringify({ result: result }),
-        });
       }
+      tokenCount += tokens.length;
+      messagesToSend = [message, ...messagesToSend];
+    }
+
+    encoding.free();
+
+    const stream = await LLMStream(model, promptToSend, temperatureToUse, key, messagesToSend, functions);
+
+    return new Response(stream);
+  } catch (error) {
+    console.error(error);
+    if (error instanceof LLMError) {
+      return new Response('Error', { status: 500, statusText: error.message });
+    } else {
+      return new Response('Error', { status: 500 });
     }
-  };
+  }
+};
 
-  processCompletions().catch(err => {
-    console.error(err);
-    res.status(500).send({ error: 'Internal Server Error' });
-  });
-}
+export default handler;

src/{stream.js → pages/api/stream.js}

@@ -1,10 +1,10 @@
-import { AZURE_DEPLOYMENT_ID, OPENAI_API_HOST, OPENAI_API_TYPE, OPENAI_API_VERSION, OPENAI_ORGANIZATION } from './pages/api/llm';
+import { AZURE_DEPLOYMENT_ID, OPENAI_API_HOST, OPENAI_API_TYPE, OPENAI_API_VERSION, OPENAI_ORGANIZATION } from './llm';
 import { createParser } from 'eventsource-parser';
 
-export class OpenAIError extends Error {
+export class LLMError extends Error {
   constructor(message, type, param, code) {
     super(message);
-    this.name = 'OpenAIError';
+    this.name = 'LLMError';
     this.type = type;
     this.param = param;
     this.code = code;