matt HOFFNER commited on
Commit
3816441
·
1 Parent(s): b4f7005

support passing a llm url directly

Browse files
Files changed (2) hide show
  1. pages/api/chat/index.ts +5 -1
  2. utils/llm.ts +88 -0
pages/api/chat/index.ts CHANGED
@@ -3,6 +3,7 @@ import { Configuration, OpenAIApi } from "openai-edge";
3
  import { HfInference } from '@huggingface/inference';
4
  import { HuggingFaceStream } from 'ai';
5
  import { experimental_buildLlama2Prompt } from 'ai/prompts';
 
6
 
7
  export const runtime = 'edge';
8
 
@@ -23,7 +24,10 @@ export default async function(req: Request) {
23
  content: systemMessage
24
  }, ...messages];
25
  }
26
- if (aiProvider === 'openai') {
 
 
 
27
  const response = await openai.createChatCompletion({
28
  model: 'gpt-4',
29
  stream: true,
 
3
  import { HfInference } from '@huggingface/inference';
4
  import { HuggingFaceStream } from 'ai';
5
  import { experimental_buildLlama2Prompt } from 'ai/prompts';
6
+ import { LLMStream } from "../../../utils/llm";
7
 
8
  export const runtime = 'edge';
9
 
 
24
  content: systemMessage
25
  }, ...messages];
26
  }
27
+ if (url) {
28
+ const stream = await LLMStream(url, messages);
29
+ return new StreamingTextResponse(stream);
30
+ } else if (aiProvider === 'openai') {
31
  const response = await openai.createChatCompletion({
32
  model: 'gpt-4',
33
  stream: true,
utils/llm.ts ADDED
@@ -0,0 +1,88 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import {
2
+ ParsedEvent,
3
+ ReconnectInterval,
4
+ createParser,
5
+ } from 'eventsource-parser';
6
+
7
+ export class LLMError extends Error {
8
+ type: string;
9
+ param: string;
10
+ code: string;
11
+
12
+ constructor(message: string, type: string, param: string, code: string) {
13
+ super(message);
14
+ this.name = 'LLMError';
15
+ this.type = type;
16
+ this.param = param;
17
+ this.code = code;
18
+ }
19
+ }
20
+
21
+ export const LLMStream = async (baseUrl: string, messages: any[]) => {
22
+ let url = `${baseUrl}/v1/chat/completions`;
23
+ const res = await fetch(url, {
24
+ headers: {
25
+ 'Content-Type': 'application/json'
26
+ },
27
+ method: 'POST',
28
+ body: JSON.stringify({
29
+ messages,
30
+ stream: true,
31
+ }),
32
+ });
33
+
34
+ const encoder = new TextEncoder();
35
+ const decoder = new TextDecoder();
36
+
37
+ if (res.status !== 200) {
38
+ const result = await res.json();
39
+ if (result.error) {
40
+ throw new LLMError(
41
+ result.error.message,
42
+ result.error.type,
43
+ result.error.param,
44
+ result.error.code,
45
+ );
46
+ } else {
47
+ throw new Error(
48
+ `API returned an error: ${
49
+ decoder.decode(result?.value) || result.statusText
50
+ }`,
51
+ );
52
+ }
53
+ }
54
+
55
+ const stream = new ReadableStream({
56
+ async start(controller) {
57
+ let accumulatedContent = ""; // To accumulate message content
58
+
59
+ const onParse = (event: ParsedEvent | ReconnectInterval) => {
60
+ if (event.type === 'event') {
61
+ const data = event.data;
62
+
63
+ if (data === '[DONE]') {
64
+ const queue = encoder.encode(accumulatedContent);
65
+ controller.enqueue(queue);
66
+ controller.close();
67
+ return;
68
+ }
69
+ try {
70
+ const parsedData = JSON.parse(data);
71
+ const content = parsedData?.choices?.[0]?.message?.content;
72
+ accumulatedContent += content ? content + " " : "";
73
+ } catch (e) {
74
+ controller.error(`Error parsing message: ${e}`);
75
+ }
76
+ }
77
+ };
78
+
79
+ const parser = createParser(onParse);
80
+
81
+ for await (const chunk of res.body as any) {
82
+ parser.feed(decoder.decode(chunk));
83
+ }
84
+ },
85
+ });
86
+
87
+ return stream;
88
+ };