Skip to content

Commit d1d86cb

Browse files
committed
Use ollama without streaming
1 parent 1986d5c commit d1d86cb

File tree

4 files changed

+55
-96
lines changed

4 files changed

+55
-96
lines changed

src/modules/lineGenerator.ts

Lines changed: 0 additions & 54 deletions
This file was deleted.

src/modules/ollamaRequest.ts

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,35 @@
1+
export async function makeOllamaRequest(url: string, data: any, bearerToken: string): Promise<string> {
2+
// Request
3+
const controller = new AbortController();
4+
let res = await fetch(url, {
5+
method: 'POST',
6+
body: JSON.stringify(data),
7+
headers: bearerToken ? {
8+
'Content-Type': 'application/json',
9+
Authorization: `Bearer ${bearerToken}`,
10+
} : {
11+
'Content-Type': 'application/json',
12+
},
13+
signal: controller.signal,
14+
});
15+
if (!res.ok || !res.body) {
16+
throw Error('Unable to connect to backend');
17+
}
18+
19+
// Reading stream
20+
let stream = res.body.getReader();
21+
const decoder = new TextDecoder();
22+
try {
23+
const { value } = await stream.read();
24+
25+
// Append chunk
26+
let chunk = decoder.decode(value);
27+
return chunk;
28+
} finally {
29+
stream.releaseLock();
30+
if (!stream.closed) { // Stop generation
31+
await stream.cancel();
32+
}
33+
controller.abort();
34+
}
35+
}

src/modules/ollamaTokenGenerator.ts

Lines changed: 0 additions & 22 deletions
This file was deleted.

src/prompts/autocomplete.ts

Lines changed: 20 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,9 @@
1-
import { ollamaTokenGenerator } from '../modules/ollamaTokenGenerator';
2-
import { countSymbol } from '../modules/text';
3-
import { info } from '../modules/log';
1+
import { makeOllamaRequest } from "../modules/ollamaRequest";
2+
3+
type OllamaToken = {
4+
model: string,
5+
response: string,
6+
};
47

58
export async function autocomplete(args: {
69
endpoint: string,
@@ -20,30 +23,27 @@ export async function autocomplete(args: {
2023
prompt: args.prefix,
2124
suffix: args.suffix,
2225
raw: true,
26+
stream: false,
2327
options: {
2428
num_predict: args.maxTokens,
2529
temperature: args.temperature
2630
}
2731
};
2832

29-
// Receiving tokens
30-
let res = '';
31-
let totalLines = 1;
32-
for await (let tokens of ollamaTokenGenerator(args.endpoint + '/api/generate', data, args.bearerToken)) {
33+
const res = await makeOllamaRequest(args.endpoint + '/api/generate', data, args.bearerToken);
34+
try {
35+
const tokens = JSON.parse(res) as OllamaToken;
3336
if (args.canceled && args.canceled()) {
34-
break;
35-
}
36-
37-
res = res + tokens.response;
38-
39-
// Update total lines
40-
totalLines += countSymbol(tokens.response, '\n');
41-
// Break if too many lines and on top level
42-
if (totalLines > args.maxLines) {
43-
info('Too many lines, breaking.');
44-
break;
37+
return "";
4538
}
39+
const response = tokens.response;
40+
41+
// take only args.maLines lines from the response
42+
let lines = response.split('\n');
43+
lines = lines.slice(0, args.maxLines);
44+
return lines.join('\n');
45+
} catch (e) {
46+
console.warn('Receive wrong line: ' + res);
47+
return "";
4648
}
47-
48-
return res;
4949
}

0 commit comments

Comments
 (0)