-
Notifications
You must be signed in to change notification settings - Fork 295
/
CallbackManager.ts
90 lines (79 loc) · 2.6 KB
/
CallbackManager.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
77
78
79
80
81
82
83
84
85
86
87
88
89
90
import { NodeWithScore } from "../Node";
/*
An event is a wrapper that groups related operations.
For example, during retrieve and synthesize,
a parent event wraps both operations, and each operation has it's own
event. In this case, both sub-events will share a parentId.
*/
export type EventTag = "intermediate" | "final";
export type EventType = "retrieve" | "llmPredict" | "wrapper";
export interface Event {
id: string;
type: EventType;
tags?: EventTag[];
parentId?: string;
}
interface BaseCallbackResponse {
event: Event;
}
//Specify StreamToken per mainstream LLM
export interface DefaultStreamToken {
id: string;
object: string;
created: number;
model: string;
choices: {
index: number;
delta: {
content?: string | null;
role?: "user" | "assistant" | "system" | "function" | "tool";
};
finish_reason: string | null;
}[];
}
//OpenAI stream token schema is the default.
//Note: Anthropic and Replicate also use similar token schemas.
export type OpenAIStreamToken = DefaultStreamToken;
export type AnthropicStreamToken = {
completion: string;
model: string;
stop_reason: string | undefined;
stop?: boolean | undefined;
log_id?: string;
};
//
//Callback Responses
//
//TODO: Write Embedding Callbacks
//StreamCallbackResponse should let practitioners implement callbacks out of the box...
//When custom streaming LLMs are involved, people are expected to write their own StreamCallbackResponses
export interface StreamCallbackResponse extends BaseCallbackResponse {
index: number;
isDone?: boolean;
token?: DefaultStreamToken;
}
export interface RetrievalCallbackResponse extends BaseCallbackResponse {
query: string;
nodes: NodeWithScore[];
}
interface CallbackManagerMethods {
/*
onLLMStream is called when a token is streamed from the LLM. Defining this
callback auto sets the stream = True flag on the openAI createChatCompletion request.
*/
onLLMStream?: (params: StreamCallbackResponse) => Promise<void> | void;
/*
onRetrieve is called as soon as the retriever finishes fetching relevant nodes.
This callback allows you to handle the retrieved nodes even if the synthesizer
is still running.
*/
onRetrieve?: (params: RetrievalCallbackResponse) => Promise<void> | void;
}
export class CallbackManager implements CallbackManagerMethods {
onLLMStream?: (params: StreamCallbackResponse) => Promise<void> | void;
onRetrieve?: (params: RetrievalCallbackResponse) => Promise<void> | void;
constructor(handlers?: CallbackManagerMethods) {
this.onLLMStream = handlers?.onLLMStream;
this.onRetrieve = handlers?.onRetrieve;
}
}