-
Notifications
You must be signed in to change notification settings - Fork 6
/
chat.ts
125 lines (122 loc) · 5.03 KB
/
chat.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
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
import { catchError, map, Observable, of, tap } from 'rxjs';
import type { IErrorResponse, INuclia } from '../../models';
import { Ask, Citations } from './ask.models';
import { ChatOptions, Search } from './search.models';
import { ResourceProperties } from '../db.models';
const END_OF_SEARCH_RESULTS = '_END_';
const START_OF_CITATIONS = '_CIT_';
export function chat(
nuclia: INuclia,
kbid: string,
path: string,
query: string,
context: Ask.ContextEntry[] = [],
features: Ask.Features[] = [Ask.Features.VECTORS, Ask.Features.PARAGRAPHS],
options: ChatOptions = {},
): Observable<Ask.Answer | IErrorResponse> {
let sourcesLength = 0;
let sources: Search.FindResults | undefined;
let relations: Search.Relations | undefined;
let citations: Citations | undefined;
let text = '';
const { synchronous, ...searchOptions } = options;
const noEmptyValues = Object.entries(searchOptions).reduce((acc, [key, value]) => {
if (value !== undefined && value !== null && value !== '') {
// eslint-disable-next-line @typescript-eslint/no-explicit-any
(acc as any)[key] = value;
}
return acc;
}, {} as ChatOptions);
const endpoint = `${path}/chat`;
const body = {
query,
context,
show: [ResourceProperties.BASIC, ResourceProperties.VALUES],
features: features.length > 0 ? features : undefined,
...noEmptyValues,
};
body['shards'] = nuclia.currentShards?.[kbid] || [];
nuclia.events?.log('lastQuery', { endpoint, params: body, nucliaOptions: nuclia.options });
return synchronous
? nuclia.rest
.post<{ answer: string; relations: Search.Relations; results: Search.FindResults; citations: Citations }>(
endpoint,
body,
undefined,
undefined,
true,
)
.pipe(
map(({ answer, relations, results, citations }) => {
return {
type: 'answer',
text: answer,
sources: { ...results, relations },
citations,
incomplete: false,
id: '',
} as Ask.Answer;
}),
tap((res) => nuclia.events?.log('lastResults', res)),
)
: nuclia.rest.getStreamedResponse(endpoint, body).pipe(
map(({ data, incomplete, headers }) => {
const searchId = headers.get('X-Nuclia-Trace-Id') || '';
const id = headers.get('NUCLIA-LEARNING-ID') || '';
// /chat returns a readable stream structured as follows:
// - 1st block: 4 first bytes indicates the size of the 2nd block
// - 2nd block: a base64 encoded JSON containing the sources used to build the answer
// - 3rd block: the answer text, ended by "_CIT_" (or "_END_" if no citations)
// - 4th block (optional):
// - 4 first bytes indicates the size of the block,
// - followed by a base64 encoded JSON containing the citations
// - ended by "_END_"
// - 5th block (optional): base64 encoded JSON containing the relations
if (sourcesLength === 0 && data.length >= 4) {
sourcesLength = new DataView(data.buffer.slice(0, 4)).getUint32(0);
}
if (!sources && sourcesLength > 0 && data.length > sourcesLength + 4) {
const sourcesData = data.slice(4, sourcesLength + 4);
sources = JSON.parse(atob(new TextDecoder().decode(sourcesData.buffer)));
}
if (sources) {
sources.searchId = searchId;
}
if (sources && data.length > sourcesLength + 4) {
text = new TextDecoder().decode(data.slice(sourcesLength + 4).buffer);
if (text.includes(END_OF_SEARCH_RESULTS)) {
let relationsBase64;
[text, relationsBase64] = text.split(END_OF_SEARCH_RESULTS);
if (text.includes(START_OF_CITATIONS)) {
let citationsBlock;
[text, citationsBlock] = text.split(START_OF_CITATIONS);
const citationsBase64 = citationsBlock.slice(4); // remove the 4 first bytes indicating the size of the block
try {
citations = JSON.parse(atob(citationsBase64));
} catch (e) {
// block is not complete yet
}
}
if (relationsBase64) {
try {
relations = JSON.parse(atob(relationsBase64));
sources.relations = relations;
} catch (e) {
// block is not complete yet
}
}
}
}
return { type: 'answer', text, sources, incomplete, id, citations } as Ask.Answer;
}),
catchError((error) =>
of({ type: 'error', status: error.status, detail: error.detail || '' } as IErrorResponse),
),
tap((res) => {
nuclia.events?.log('lastResults', res);
if (res.type === 'answer' && res.sources) {
nuclia.currentShards = { ...nuclia.currentShards, [kbid]: res.sources.shards || [] };
}
}),
);
}