/
index.ts
76 lines (68 loc) · 2.59 KB
/
index.ts
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
import { serve } from "http/server.ts";
import { ChatOpenAI } from "langchain/chat_models/openai";
import { LLMChain } from "langchain/chains";
import { CallbackManager } from "langchain/callbacks";
import {
ChatPromptTemplate,
HumanMessagePromptTemplate,
} from "langchain/prompts";
import { corsHeaders } from "../_shared/cors.ts";
const prompt = ChatPromptTemplate.fromPromptMessages([
HumanMessagePromptTemplate.fromTemplate("{input}"),
]);
serve(async (req) => {
// This is needed if you're planning to invoke your function from a browser.
if (req.method === "OPTIONS") {
return new Response("ok", { headers: corsHeaders });
}
try {
const { input } = await req.json();
// Check if the request is for a streaming response.
const streaming = req.headers.get("accept") === "text/event-stream";
if (streaming) {
// For a streaming response we need to use a TransformStream to
// convert the LLM's callback-based API into a stream-based API.
const encoder = new TextEncoder();
const stream = new TransformStream();
const writer = stream.writable.getWriter();
const llm = new ChatOpenAI({
streaming,
callbackManager: CallbackManager.fromHandlers({
handleLLMNewToken: async (token) => {
await writer.ready;
await writer.write(encoder.encode(`data: ${token}\n\n`));
},
handleLLMEnd: async () => {
await writer.ready;
await writer.close();
},
handleLLMError: async (e) => {
await writer.ready;
await writer.abort(e);
},
}),
});
const chain = new LLMChain({ prompt, llm });
// We don't need to await the result of the chain.run() call because
// the LLM will invoke the callbackManager's handleLLMEnd() method
chain.call({ input }).catch((e) => console.error(e));
return new Response(stream.readable, {
headers: { ...corsHeaders, "Content-Type": "text/event-stream" },
});
} else {
// For a non-streaming response we can just await the result of the
// chain.run() call and return it.
const llm = new ChatOpenAI();
const chain = new LLMChain({ prompt, llm });
const response = await chain.call({ input });
return new Response(JSON.stringify(response), {
headers: { ...corsHeaders, "Content-Type": "application/json" },
});
}
} catch (e) {
return new Response(JSON.stringify({ error: e.message }), {
status: 500,
headers: { ...corsHeaders, "Content-Type": "application/json" },
});
}
});