feat(server): migrate copilot provider sdk (#11584)

fix AI-15
fix AI-16

packages/backend/server/src/plugins/copilot/providers/perplexity.ts

@@ -1,4 +1,8 @@
-import { EventSourceParserStream } from 'eventsource-parser/stream';
+import {
+  createPerplexity,
+  type PerplexityProvider as VercelPerplexityProvider,
+} from '@ai-sdk/perplexity';
+import { generateText, streamText } from 'ai';
 import { z } from 'zod';
 
 import {
@@ -14,6 +18,7 @@ import {
   CopilotTextToTextProvider,
   PromptMessage,
 } from './types';
+import { chatToGPTMessage, CitationParser } from './utils';
 
 export type PerplexityConfig = {
   apiKey: string;
@@ -39,130 +44,8 @@ const PerplexityErrorSchema = z.union([
   }),
 ]);
 
-const PerplexityDataSchema = z.object({
-  citations: z.array(z.string()),
-  choices: z.array(
-    z.object({
-      message: z.object({
-        content: z.string(),
-        role: z.literal('assistant'),
-      }),
-      delta: z.object({
-        content: z.string(),
-        role: z.literal('assistant'),
-      }),
-      finish_reason: z.union([z.literal('stop'), z.literal(null)]),
-    })
-  ),
-});
-
-const PerplexitySchema = z.union([PerplexityDataSchema, PerplexityErrorSchema]);
-
 type PerplexityError = z.infer<typeof PerplexityErrorSchema>;
 
-export class CitationParser {
-  private readonly SQUARE_BRACKET_OPEN = '[';
-
-  private readonly SQUARE_BRACKET_CLOSE = ']';
-
-  private readonly PARENTHESES_OPEN = '(';
-
-  private startToken: string[] = [];
-
-  private endToken: string[] = [];
-
-  private numberToken: string[] = [];
-
-  private citations: string[] = [];
-
-  public parse(content: string, citations: string[]) {
-    this.citations = citations;
-    let result = '';
-    const contentArray = content.split('');
-    for (const [index, char] of contentArray.entries()) {
-      if (char === this.SQUARE_BRACKET_OPEN) {
-        if (this.numberToken.length === 0) {
-          this.startToken.push(char);
-        } else {
-          result += this.flush() + char;
-        }
-        continue;
-      }
-
-      if (char === this.SQUARE_BRACKET_CLOSE) {
-        this.endToken.push(char);
-        if (this.startToken.length === this.endToken.length) {
-          const cIndex = Number(this.numberToken.join('').trim());
-          if (
-            cIndex > 0 &&
-            cIndex <= citations.length &&
-            contentArray[index + 1] !== this.PARENTHESES_OPEN
-          ) {
-            const content = `[^${cIndex}]`;
-            result += content;
-            this.resetToken();
-          } else {
-            result += this.flush();
-          }
-        } else if (this.startToken.length < this.endToken.length) {
-          result += this.flush();
-        }
-        continue;
-      }
-
-      if (this.isNumeric(char)) {
-        if (this.startToken.length > 0) {
-          this.numberToken.push(char);
-        } else {
-          result += this.flush() + char;
-        }
-        continue;
-      }
-
-      if (this.startToken.length > 0) {
-        result += this.flush() + char;
-      } else {
-        result += char;
-      }
-    }
-
-    return result;
-  }
-
-  public end() {
-    return this.flush() + '\n' + this.getFootnotes();
-  }
-
-  private flush() {
-    const content = this.getTokenContent();
-    this.resetToken();
-    return content;
-  }
-
-  private getFootnotes() {
-    const footnotes = this.citations.map((citation, index) => {
-      return `[^${index + 1}]: {"type":"url","url":"${encodeURIComponent(
-        citation
-      )}"}`;
-    });
-    return footnotes.join('\n');
-  }
-
-  private getTokenContent() {
-    return this.startToken.concat(this.numberToken, this.endToken).join('');
-  }
-
-  private resetToken() {
-    this.startToken = [];
-    this.endToken = [];
-    this.numberToken = [];
-  }
-
-  private isNumeric(str: string) {
-    return !isNaN(Number(str)) && str.trim() !== '';
-  }
-}
-
 export class PerplexityProvider
   extends CopilotProvider<PerplexityConfig>
   implements CopilotTextToTextProvider
@@ -176,10 +59,20 @@ export class PerplexityProvider
     'sonar-reasoning-pro',
   ];
 
+  #instance!: VercelPerplexityProvider;
+
   override configured(): boolean {
     return !!this.config.apiKey;
   }
 
+  protected override setup() {
+    super.setup();
+    this.#instance = createPerplexity({
+      apiKey: this.config.apiKey,
+      baseURL: this.config.endpoint,
+    });
+  }
+
   async generateText(
     messages: PromptMessage[],
     model: string = 'sonar',
@@ -188,38 +81,26 @@ export class PerplexityProvider
     await this.checkParams({ messages, model, options });
     try {
       metrics.ai.counter('chat_text_calls').add(1, { model });
-      const sMessages = messages
-        .map(({ content, role }) => ({ content, role }))
-        .filter(({ content }) => typeof content === 'string');
 
-      const params = {
-        method: 'POST',
-        headers: {
-          Authorization: `Bearer ${this.config.apiKey}`,
-          'Content-Type': 'application/json',
-        },
-        body: JSON.stringify({
-          model,
-          messages: sMessages,
-          max_tokens: options.maxTokens || 4096,
-        }),
-      };
-      const response = await fetch(
-        this.config.endpoint || 'https://api.perplexity.ai/chat/completions',
-        params
-      );
-      const data = PerplexitySchema.parse(await response.json());
-      if ('detail' in data || 'error' in data) {
-        throw this.convertError(data);
-      } else {
-        const citationParser = new CitationParser();
-        const { content } = data.choices[0].message;
-        const { citations } = data;
-        let result = content.replaceAll(/<\/?think>\n/g, '\n---\n');
-        result = citationParser.parse(result, citations);
-        result += citationParser.end();
-        return result;
-      }
+      const [system, msgs] = await chatToGPTMessage(messages);
+
+      const modelInstance = this.#instance(model);
+
+      const { text, sources } = await generateText({
+        model: modelInstance,
+        system,
+        messages: msgs,
+        temperature: options.temperature || 0,
+        maxTokens: options.maxTokens || 4096,
+        abortSignal: options.signal,
+      });
+
+      const citationParser = new CitationParser();
+      const citations = sources.map(s => s.url);
+      let result = text.replaceAll(/<\/?think>\n/g, '\n---\n');
+      result = citationParser.parse(result, citations);
+      result += citationParser.end();
+      return result;
     } catch (e: any) {
       metrics.ai.counter('chat_text_errors').add(1, { model });
       throw this.handleError(e);
@@ -234,69 +115,54 @@ export class PerplexityProvider
     await this.checkParams({ messages, model, options });
     try {
       metrics.ai.counter('chat_text_stream_calls').add(1, { model });
-      const sMessages = messages
-        .map(({ content, role }) => ({ content, role }))
-        .filter(({ content }) => typeof content === 'string');
 
-      const params = {
-        method: 'POST',
-        headers: {
-          Authorization: `Bearer ${this.config.apiKey}`,
-          'Content-Type': 'application/json',
-        },
-        body: JSON.stringify({
-          model,
-          messages: sMessages,
-          max_tokens: options.maxTokens || 4096,
-          stream: true,
-        }),
-      };
-      const response = await fetch(
-        this.config.endpoint || 'https://api.perplexity.ai/chat/completions',
-        params
-      );
-      const errorHandler = this.convertError;
-      if (response.ok && response.body) {
-        const citationParser = new CitationParser();
-        const eventStream = response.body
-          .pipeThrough(new TextDecoderStream())
-          .pipeThrough(new EventSourceParserStream())
-          .pipeThrough(
-            new TransformStream({
-              transform(chunk, controller) {
-                if (options.signal?.aborted) {
-                  controller.enqueue(null);
-                  return;
-                }
-                const json = JSON.parse(chunk.data);
-                if (json) {
-                  const data = PerplexitySchema.parse(json);
-                  if ('detail' in data || 'error' in data) {
-                    throw errorHandler(data);
-                  }
-                  const { content } = data.choices[0].delta;
-                  const { citations } = data;
-                  let result = content.replaceAll(/<\/?think>\n?/g, '\n---\n');
-                  result = citationParser.parse(result, citations);
-                  controller.enqueue(result);
-                }
-              },
-              flush(controller) {
-                controller.enqueue(citationParser.end());
-                controller.enqueue(null);
-              },
-            })
-          );
+      const [system, msgs] = await chatToGPTMessage(messages);
 
-        const reader = eventStream.getReader();
-        while (true) {
-          const { done, value } = await reader.read();
-          if (done) break;
-          yield value;
+      const modelInstance = this.#instance(model);
+
+      const stream = streamText({
+        model: modelInstance,
+        system,
+        messages: msgs,
+        temperature: options.temperature || 0,
+        maxTokens: options.maxTokens || 4096,
+        abortSignal: options.signal,
+      });
+
+      const citationParser = new CitationParser();
+      const citations = [];
+      for await (const chunk of stream.fullStream) {
+        switch (chunk.type) {
+          case 'source': {
+            citations.push(chunk.source.url);
+            break;
+          }
+          case 'text-delta': {
+            const result = citationParser.parse(
+              chunk.textDelta.replaceAll(/<\/?think>\n?/g, '\n---\n'),
+              citations
+            );
+            yield result;
+            break;
+          }
+          case 'step-finish': {
+            const result = citationParser.end();
+            yield result;
+            break;
+          }
+          case 'error': {
+            const json =
+              typeof chunk.error === 'string'
+                ? JSON.parse(chunk.error)
+                : chunk.error;
+            if (json && typeof json === 'object') {
+              const data = PerplexityErrorSchema.parse(json);
+              if ('detail' in data || 'error' in data) {
+                throw this.convertError(data);
+              }
+            }
+          }
         }
-      } else {
-        const result = await this.generateText(messages, model, options);
-        yield result;
       }
     } catch (e) {
       metrics.ai.counter('chat_text_stream_errors').add(1, { model });