import axios from "axios"; import { EventEmitter } from "events"; import { strict as assert } from "assert"; import { CompletionCache } from "./CompletionCache"; import { sleep, cancelable } from "./utils"; import { Agent, AgentEvent } from "./types"; import { TabbyApi, CancelablePromise, ApiError, CompletionRequest, CompletionResponse, ChoiceEvent, CompletionEvent, } from "./generated"; export class TabbyAgent extends EventEmitter implements Agent { private serverUrl: string = "http://127.0.0.1:5000"; private status: "connecting" | "ready" | "disconnected" = "connecting"; private api: TabbyApi; private completionCache: CompletionCache; constructor() { super(); this.ping(); this.api = new TabbyApi({ BASE: this.serverUrl }); this.completionCache = new CompletionCache(); } private changeStatus(status: "connecting" | "ready" | "disconnected") { if (this.status != status) { this.status = status; const event: AgentEvent = { event: "statusChanged", status }; super.emit("statusChanged", event); } } private async ping(tries: number = 0): Promise { try { const response = await axios.get(this.serverUrl); this.changeStatus("ready"); return true; } catch (e) { if (tries > 5) { this.changeStatus("disconnected"); return false; } this.changeStatus("connecting"); const pingRetryDelay = 1000; await sleep(pingRetryDelay); return this.ping(tries + 1); } } private wrapApiPromise(promise: CancelablePromise): CancelablePromise { return cancelable( promise .then((resolved: T) => { this.changeStatus("ready"); return resolved; }) .catch((err: ApiError) => { this.changeStatus("disconnected"); throw err; }), () => { promise.cancel(); } ); } public setServerUrl(serverUrl: string): string { this.serverUrl = serverUrl.replace(/\/$/, ""); // Remove trailing slash this.ping(); this.api = new TabbyApi({ BASE: this.serverUrl }); return this.serverUrl; } public getServerUrl(): string { return this.serverUrl; } public getStatus(): "connecting" | "ready" | "disconnected" { return this.status; } public getCompletions(request: CompletionRequest): CancelablePromise { if (this.completionCache.has(request)) { return new CancelablePromise((resolve) => { resolve(this.completionCache.get(request)); }); } const promise = this.wrapApiPromise(this.api.default.completionsV1CompletionsPost(request)); return cancelable( promise.then((response: CompletionResponse) => { this.completionCache.set(request, response); return response; }), () => { promise.cancel(); } ); } public postEvent(request: ChoiceEvent | CompletionEvent): CancelablePromise { return this.wrapApiPromise(this.api.default.eventsV1EventsPost(request)); } }