@trymirai/uzu 0.1.38 → 0.1.40

package/napi/uzu.d.ts

@@ -1,19 +1,19 @@
 /* auto-generated by NAPI-RS */
 /* eslint-disable */
-/** Thin FFI wrapper around `ModelStorage`. */
+export declare class ChatSession {
+  run(input: Input, config: RunConfig, progress?: (arg: Output) => boolean | undefined | null): Output
+}
+
 export declare class Engine {
-  getCloudModels(): Promise<Array<CloudModel>>
-  fetchCloudModels(): Promise<Array<CloudModel>>
-  /** Returns a `DownloadHandle` for the given model identifier. */
-  downloadHandle(repoId: string): ModelDownloadHandle
-  getLocalModels(): Array<LocalModel>
-  getState(repoId: string): ModelDownloadState
   activate(apiKey: string): Promise<LicenseStatus>
-  downloadModel(repoId: string): Promise<void>
-  pauseModel(repoId: string): Promise<void>
-  deleteModel(repoId: string): Promise<void>
+  getChatModels(types: Array<ModelType>): Promise<Array<ChatModel>>
+  getModelDownloadState(repoId: string): ModelDownloadState
+  createModelDownloadHandle(repoId: string): ModelDownloadHandle
   constructor()
-  createSession(repoId: string, modelType: ModelType, config: Config): Session
+  createChatSession(model: ChatModel, config: Config): ChatSession
+  registerLicenseStatusHandler(callback?: (arg: LicenseStatus) => void | undefined | null): void
+  registerModelDownloadStateHandler(callback?: (arg0: string, arg1: ModelDownloadState) => void | undefined | null): void
+  registerChatModelsHandler(callback?: (arg: Array<ChatModel>) => void | undefined | null): void
 }
 
 export declare class ModelDownloadHandle {
@@ -38,8 +38,13 @@ export declare class ProgressUpdate {
   progress: number
 }
 
-export declare class Session {
-  run(input: Input, config: RunConfig, progress?: (arg: Output) => boolean | undefined | null): Output
+export interface ChatModel {
+  readonly repoId: string
+  readonly type: ModelType
+  readonly name: string
+  readonly vendor: string
+  readonly quantization?: string
+  readonly outputParserRegex?: string
 }
 
 export interface ClassificationFeature {
@@ -47,13 +52,6 @@ export interface ClassificationFeature {
   values: Array<string>
 }
 
-export interface CloudModel {
-  readonly repoId: string
-  readonly name: string
-  readonly vendor: string
-  readonly outputParserRegex?: string
-}
-
 export interface Config {
   preset: Preset
   prefillStepSize: PrefillStepSize
@@ -88,18 +86,6 @@ export type LicenseStatus =
   | { type: 'Timeout' }
   | { type: 'HttpError', code: number }
 
-export interface LocalModel {
-  readonly identifier: string
-  readonly repoId: string
-  readonly family: string
-  readonly vendor: string
-  readonly name: string
-  readonly size: string
-  readonly quantization?: string
-  readonly outputParserRegex?: string
-  readonly state: ModelDownloadState
-}
-
 export interface Message {
   role: Role
   content: string

package/napi/uzu.js

@@ -21,20 +21,19 @@ try {
 
 // Re-export all runtime exports
 const {
+  ChatSession,
   Engine,
   ModelDownloadHandle,
   ProgressStream,
   ProgressUpdate,
-  Session,
   FinishReason,
   ModelDownloadPhase,
   ModelType,
   Role,
   version,
+  ChatModel,
   ClassificationFeature,
-  CloudModel,
   Config,
-  LocalModel,
   Message,
   ModelDownloadState,
   Output,
@@ -58,20 +57,19 @@ const {
 } = nativeBinding;
 
 module.exports = {
+  ChatSession,
   Engine,
   ModelDownloadHandle,
   ProgressStream,
   ProgressUpdate,
-  Session,
   FinishReason,
   ModelDownloadPhase,
   ModelType,
   Role,
   version,
+  ChatModel,
   ClassificationFeature,
-  CloudModel,
   Config,
-  LocalModel,
   Message,
   ModelDownloadState,
   Output,

package/napi/uzu.mjs

@@ -21,20 +21,19 @@ try {
 
 // Re-export all runtime exports as named exports  
 export const {
+  ChatSession,
   Engine,
   ModelDownloadHandle,
   ProgressStream,
   ProgressUpdate,
-  Session,
   FinishReason,
   ModelDownloadPhase,
   ModelType,
   Role,
   version,
+  ChatModel,
   ClassificationFeature,
-  CloudModel,
   Config,
-  LocalModel,
   Message,
   ModelDownloadState,
   Output,

package/package.json

@@ -98,5 +98,5 @@
   },
   "type": "commonjs",
   "types": "./index.d.ts",
-  "version": "0.1.38"
+  "version": "0.1.40"
 }

package/src/bridging/chatModel.ts

@@ -0,0 +1,54 @@
+import { ChatModel as NapiChatModel } from '../napi/uzu';
+import { ModelType, modelTypeFromNapiModelType, modelTypeToNapiModelType } from './modelType';
+
+export class ChatModel {
+    readonly repoId: string;
+    readonly type: ModelType;
+    readonly name: string;
+    readonly vendor: string;
+    readonly quantization: string | null;
+    readonly outputParserRegex: string | null;
+
+    get isThinking(): boolean {
+        return this.outputParserRegex !== null;
+    }
+
+    private constructor(
+        repoId: string,
+        type: ModelType,
+        name: string,
+        vendor: string,
+        quantization: string | null,
+        outputParserRegex: string | null,
+    ) {
+        this.repoId = repoId;
+        this.type = type;
+        this.name = name;
+        this.vendor = vendor;
+        this.quantization = quantization;
+        this.outputParserRegex = outputParserRegex;
+    }
+
+    static fromNapiChatModel(napiChatModel: NapiChatModel): ChatModel {
+        return new ChatModel(
+            napiChatModel.repoId,
+            modelTypeFromNapiModelType(napiChatModel.type),
+            napiChatModel.name,
+            napiChatModel.vendor,
+            napiChatModel.quantization ?? null,
+            napiChatModel.outputParserRegex ?? null,
+        );
+    }
+
+    toNapi(): NapiChatModel {
+        const napiChatModel: NapiChatModel = {
+            repoId: this.repoId,
+            type: modelTypeToNapiModelType(this.type),
+            name: this.name,
+            vendor: this.vendor,
+            ...(this.quantization !== null && { quantization: this.quantization }),
+            ...(this.outputParserRegex !== null && { outputParserRegex: this.outputParserRegex }),
+        };
+        return napiChatModel;
+    }
+}

package/src/bridging/{session.ts → chatSession.ts}

@@ -1,13 +1,13 @@
-import { Output as NapiOutput, Session as NapiSession } from '../napi/uzu';
+import { ChatSession as NapiChatSession, Output as NapiOutput } from '../napi/uzu';
 import { Input } from './input';
 import { Output } from './output';
 import { RunConfig } from './runConfig';
 
-export class Session {
-    private readonly napiSession: NapiSession;
+export class ChatSession {
+    private readonly napiChatSession: NapiChatSession;
 
-    constructor(napiSession: NapiSession) {
-        this.napiSession = napiSession;
+    constructor(napiChatSession: NapiChatSession) {
+        this.napiChatSession = napiChatSession;
     }
 
     run(
@@ -18,7 +18,7 @@ export class Session {
         const napiInput = input.toNapi();
         const napiConfig = config.toNapi();
         const napiProgress = (partialOutput: NapiOutput) => progress(Output.fromNapi(partialOutput));
-        const napiOutput = this.napiSession.run(napiInput, napiConfig, napiProgress);
+        const napiOutput = this.napiChatSession.run(napiInput, napiConfig, napiProgress);
         return Output.fromNapi(napiOutput);
     }
 }

package/src/bridging/engine.ts

@@ -2,12 +2,14 @@ import { EngineError, EngineErrorCode, unwrapOrThrow } from '../error';
 import { EngineInteractor } from '../interactors/engineInteractor';
 import { Engine as NapiEngine } from '../napi/uzu';
 import { readEnv } from '../utilities/env';
+import { ChatModel } from './chatModel';
+import { ChatSession } from './chatSession';
 import { Config } from './config';
 import { DownloadHandle } from './downloadHandle';
+import { DownloadProgressUpdate } from './downloadProgressUpdate';
+import { DownloadPhase, DownloadState } from './downloadState';
 import { LicenseStatus, licenseStatusFromNapiLicenseStatus } from './licenseStatus';
-import { Model, ModelKind } from './model';
-import { ModelType, modelTypeToNapi } from './modelType';
-import { Session } from './session';
+import { ModelType, modelTypeToNapiModelType } from './modelType';
 
 export class Engine {
     private readonly napiEngine: NapiEngine;
@@ -44,62 +46,60 @@ export class Engine {
         return new EngineInteractor(enginePromise);
     }
 
-    async models(types: ModelType[] | null = null, kinds: ModelKind[] | null = null): Promise<Model[]> {
-        function isTypeAllowed(type: ModelType): boolean {
-            return types === null || types.includes(type);
-        }
+    getDownloadState(repoId: string): DownloadState {
+        let downloadState = this.napiEngine.getModelDownloadState(repoId);
+        return DownloadState.fromNapi(downloadState);
+    }
 
-        function isKindAllowed(kind: ModelKind): boolean {
-            return kinds === null || kinds.includes(kind);
-        }
+    downloadHandle(repoId: string): DownloadHandle {
+        return new DownloadHandle(this.napiEngine.createModelDownloadHandle(repoId));
+    }
 
-        let results: Model[] = [];
-        if (isTypeAllowed(ModelType.Local) && isKindAllowed(ModelKind.Text)) {
-            const napiLocalModels = this.napiEngine.getLocalModels();
-            results.push(...napiLocalModels.map(Model.fromNapiLocalModel));
-        }
-        if (isTypeAllowed(ModelType.Cloud) && isKindAllowed(ModelKind.Text)) {
-            const napiCloudModels = await this.napiEngine.getCloudModels();
-            results.push(...napiCloudModels.map(Model.fromNapiCloudModel));
-        }
+    async chatModels(types: ModelType[] | null = null): Promise<ChatModel[]> {
+        let typesToFilter: ModelType[] = types ?? [ModelType.Local, ModelType.Cloud];
+        let napiModelTypes = typesToFilter.map(modelTypeToNapiModelType);
+        let napiChatModels = await this.napiEngine.getChatModels(napiModelTypes);
+        let results: ChatModel[] = napiChatModels.map(ChatModel.fromNapiChatModel);
         return results;
     }
 
-    async model(repoId: string): Promise<Model> {
-        const models = await this.models();
-        const model = models.find((model) => model.repoId === repoId);
-        return unwrapOrThrow(model, EngineErrorCode.ModelNotFound);
+    async chatModel(repoId: string, types: ModelType[] | null = null): Promise<ChatModel> {
+        const chatModels = await this.chatModels(types);
+        const chatModel = chatModels.find((model) => model.repoId === repoId);
+        return unwrapOrThrow(chatModel, EngineErrorCode.ModelNotFound);
     }
 
-    downloadHandle(model: Model): DownloadHandle {
-        if (model.type !== ModelType.Local) {
-            throw new EngineError(EngineErrorCode.ExpectedLocalModel);
+    async downloadChatModel(
+        chatModel: ChatModel,
+        progress: (progressUpdate: DownloadProgressUpdate) => void = () => {},
+    ): Promise<DownloadState> {
+        switch (chatModel.type) {
+            case ModelType.Local:
+                break;
+            case ModelType.Cloud:
+                throw new EngineError(EngineErrorCode.UnexpectedModelType);
         }
-        return new DownloadHandle(this.napiEngine.downloadHandle(model.repoId));
-    }
 
-    session(model: Model, config: Config = Config.default()): Session {
-        if (model.kind !== ModelKind.Text) {
-            throw new EngineError(EngineErrorCode.ExpectedTextModel);
+        const downloadHandle = this.downloadHandle(chatModel.repoId);
+        const state = await downloadHandle.state();
+        switch (state.phase) {
+            case DownloadPhase.Downloaded:
+                break;
+            default:
+                await downloadHandle.download();
+                for await (const progressUpdate of downloadHandle.progressUpdate()) {
+                    progress(progressUpdate);
+                }
+                break;
         }
-        const napiModelType = modelTypeToNapi(model.type);
-        const napiConfig = config.toNapi();
-        return new Session(this.napiEngine.createSession(model.repoId, napiModelType, napiConfig));
-    }
 
-    async downloadModel(repoId: string): Promise<void> {
-        await this.napiEngine.downloadModel(repoId);
+        const finalState = await downloadHandle.state();
+        return finalState;
     }
 
-    async pauseModel(repoId: string): Promise<void> {
-        await this.napiEngine.pauseModel(repoId);
-    }
-
-    async deleteModel(repoId: string): Promise<void> {
-        await this.napiEngine.deleteModel(repoId);
-    }
-
-    getState(repoId: string) {
-        return this.napiEngine.getState(repoId);
+    chatSession(chatModel: ChatModel, config: Config = Config.default()): ChatSession {
+        const napiChatModel = chatModel.toNapi();
+        const napiConfig = config.toNapi();
+        return new ChatSession(this.napiEngine.createChatSession(napiChatModel, napiConfig));
     }
 }

package/src/bridging/modelType.ts

@@ -1,3 +1,4 @@
+import { EngineError, EngineErrorCode } from '../error';
 import { ModelType as NapiModelType } from '../napi/uzu';
 
 export enum ModelType {
@@ -5,7 +6,18 @@ export enum ModelType {
     Cloud = 'Cloud',
 }
 
-export function modelTypeToNapi(modelType: ModelType): NapiModelType {
+export function modelTypeFromNapiModelType(napiModelType: NapiModelType): ModelType {
+    switch (napiModelType) {
+        case 0:
+            return ModelType.Local;
+        case 1:
+            return ModelType.Cloud;
+        default:
+            throw new EngineError(EngineErrorCode.UnexpectedModelType);
+    }
+}
+
+export function modelTypeToNapiModelType(modelType: ModelType): NapiModelType {
     switch (modelType) {
         case ModelType.Local:
             return 0;

package/src/error.ts

@@ -2,10 +2,10 @@ export enum EngineErrorCode {
     APIKeyNotFound = 'apiKeyNotFound',
     ModelNotFound = 'modelNotFound',
     ExpectedLocalModel = 'expectedLocalModel',
-    ExpectedTextModel = 'expectedTextModel',
     UnexpectedDownloadPhase = 'unexpectedDownloadPhase',
     UnexpectedFinishReason = 'unexpectedFinishReason',
     UnexpectedLicenseStatus = 'unexpectedLicenseStatus',
+    UnexpectedModelType = 'unexpectedModelType',
     LicenseNotActivated = 'licenseNotActivated',
     Unknown = 'unknown',
 }
@@ -14,10 +14,10 @@ const errorCodeToMessage = {
     [EngineErrorCode.APIKeyNotFound]: 'API key not found',
     [EngineErrorCode.ModelNotFound]: 'Model not found',
     [EngineErrorCode.ExpectedLocalModel]: 'Expected local model',
-    [EngineErrorCode.ExpectedTextModel]: 'Expected text model',
     [EngineErrorCode.UnexpectedDownloadPhase]: 'Unexpected download phase',
     [EngineErrorCode.UnexpectedFinishReason]: 'Unexpected finish reason',
     [EngineErrorCode.UnexpectedLicenseStatus]: 'Unexpected license status',
+    [EngineErrorCode.UnexpectedModelType]: 'Unexpected model type',
     [EngineErrorCode.LicenseNotActivated]: 'License not activated',
     [EngineErrorCode.Unknown]: 'Unknown error',
 };

package/src/index.ts

@@ -1,3 +1,5 @@
+export { ChatModel } from './bridging/chatModel';
+export { ChatSession } from './bridging/chatSession';
 export { ClassificationFeature } from './bridging/classificationFeature';
 export { Config } from './bridging/config';
 export { ContextLength } from './bridging/contextLength';
@@ -8,7 +10,6 @@ export { Engine as default } from './bridging/engine';
 export { FinishReason } from './bridging/finishReason';
 export { Input } from './bridging/input';
 export { Message } from './bridging/message';
-export { Model, ModelKind } from './bridging/model';
 export { ModelType } from './bridging/modelType';
 export { Output } from './bridging/output';
 export { ParsedText } from './bridging/parsedText';
@@ -20,14 +21,13 @@ export { RunStats } from './bridging/runStats';
 export { SamplingMethod } from './bridging/samplingMethod';
 export { SamplingPolicy } from './bridging/samplingPolicy';
 export { SamplingSeed } from './bridging/samplingSeed';
-export { Session } from './bridging/session';
 export { Stats } from './bridging/stats';
 export { StepStats } from './bridging/stepStats';
 export { Text } from './bridging/text';
 export { TotalStats } from './bridging/totalStats';
 export { EngineError } from './error';
+export { ChatModelInteractor } from './interactors/chatModelInteractor';
+export { ChatModelsInteractor } from './interactors/chatModelsInteractor';
+export { ChatSessionInteractor } from './interactors/chatSessionInteractor';
 export { DownloadInteractor } from './interactors/downloadInteractor';
 export { EngineInteractor } from './interactors/engineInteractor';
-export { ModelInteractor } from './interactors/modelInteractor';
-export { ModelsInteractor } from './interactors/modelsInteractor';
-export { SessionInteractor } from './interactors/sessionInteractor';

package/src/interactors/{modelInteractor.ts → chatModelInteractor.ts}

@@ -1,32 +1,32 @@
+import { ChatModel } from '../bridging/chatModel';
 import { Config } from '../bridging/config';
 import { ContextLength } from '../bridging/contextLength';
 import { DownloadProgressUpdate } from '../bridging/downloadProgressUpdate';
 import { DownloadPhase } from '../bridging/downloadState';
 import { Message } from '../bridging/message';
-import { Model } from '../bridging/model';
 import { ModelType } from '../bridging/modelType';
 import { Output } from '../bridging/output';
 import { PrefillStepSize } from '../bridging/prefillStepSize';
 import { Preset } from '../bridging/preset';
 import { RunConfig } from '../bridging/runConfig';
 import { SamplingSeed } from '../bridging/samplingSeed';
+import { ChatModelsInteractor } from './chatModelsInteractor';
+import { ChatSessionInteractor } from './chatSessionInteractor';
 import { DownloadInteractor } from './downloadInteractor';
 import { Interactor, InteractorEntity } from './interactor';
-import { ModelsInteractor } from './modelsInteractor';
-import { SessionInteractor } from './sessionInteractor';
 
-export class ModelInteractor implements Interactor<Model> {
-    readonly modelsInteractor: ModelsInteractor;
-    readonly entity: InteractorEntity<Model>;
+export class ChatModelInteractor implements Interactor<ChatModel> {
+    readonly modelsInteractor: ChatModelsInteractor;
+    readonly entity: InteractorEntity<ChatModel>;
     readonly config: Config;
 
-    constructor(modelsInteractor: ModelsInteractor, model: InteractorEntity<Model>, config: Config) {
+    constructor(modelsInteractor: ChatModelsInteractor, model: InteractorEntity<ChatModel>, config: Config) {
         this.modelsInteractor = modelsInteractor;
         this.entity = model;
         this.config = config;
     }
 
-    async finalize(): Promise<Model> {
+    async finalize(): Promise<ChatModel> {
         return await this.entity;
     }
 
@@ -36,13 +36,13 @@ export class ModelInteractor implements Interactor<Model> {
         const handlePromise = (async () => {
             const model = await this.finalize();
             const engine = await this.modelsInteractor.engineInteractor.finalize();
-            const handle = engine.downloadHandle(model);
+            const handle = engine.downloadHandle(model.repoId);
             return handle;
         })();
         return new DownloadInteractor(this, handlePromise);
     }
 
-    download(callback?: (progressUpdate: DownloadProgressUpdate) => void): ModelInteractor {
+    download(callback?: (progressUpdate: DownloadProgressUpdate) => void): ChatModelInteractor {
         const modelPromise = (async () => {
             const downloadInteractor = this.downloading();
             const handle = await downloadInteractor.finalize();
@@ -59,34 +59,34 @@ export class ModelInteractor implements Interactor<Model> {
             }
             return this.finalize();
         })();
-        return new ModelInteractor(this.modelsInteractor, modelPromise, this.config);
+        return new ChatModelInteractor(this.modelsInteractor, modelPromise, this.config);
     }
 
     /* Config */
 
-    preset(preset: Preset): ModelInteractor {
+    preset(preset: Preset): ChatModelInteractor {
         const config = this.config.withPreset(preset);
-        return new ModelInteractor(this.modelsInteractor, this.entity, config);
+        return new ChatModelInteractor(this.modelsInteractor, this.entity, config);
     }
 
-    prefillStepSize(prefillStepSize: PrefillStepSize): ModelInteractor {
+    prefillStepSize(prefillStepSize: PrefillStepSize): ChatModelInteractor {
         const config = this.config.withPrefillStepSize(prefillStepSize);
-        return new ModelInteractor(this.modelsInteractor, this.entity, config);
+        return new ChatModelInteractor(this.modelsInteractor, this.entity, config);
     }
 
-    contextLength(contextLength: ContextLength): ModelInteractor {
+    contextLength(contextLength: ContextLength): ChatModelInteractor {
         const config = this.config.withContextLength(contextLength);
-        return new ModelInteractor(this.modelsInteractor, this.entity, config);
+        return new ChatModelInteractor(this.modelsInteractor, this.entity, config);
     }
 
-    samplingSeed(samplingSeed: SamplingSeed): ModelInteractor {
+    samplingSeed(samplingSeed: SamplingSeed): ChatModelInteractor {
         const config = this.config.withSamplingSeed(samplingSeed);
-        return new ModelInteractor(this.modelsInteractor, this.entity, config);
+        return new ChatModelInteractor(this.modelsInteractor, this.entity, config);
     }
 
     /* Session */
 
-    session(): SessionInteractor {
+    session(): ChatSessionInteractor {
         const sessionPromise = (async () => {
             const engine = await this.modelsInteractor.engineInteractor.finalize();
 
@@ -95,10 +95,10 @@ export class ModelInteractor implements Interactor<Model> {
                 model = await this.download().finalize();
             }
 
-            const session = engine.session(model, this.config);
+            const session = engine.chatSession(model, this.config);
             return session;
         })();
-        return new SessionInteractor(this, sessionPromise, RunConfig.default());
+        return new ChatSessionInteractor(this, sessionPromise, RunConfig.default());
     }
 
     async reply(text: string, progress: (partialOutput: Output) => boolean = () => true): Promise<Output> {

package/src/interactors/{modelsInteractor.ts → chatModelsInteractor.ts}

@@ -1,29 +1,29 @@
+import { ChatModel } from '../bridging/chatModel';
 import { Config } from '../bridging/config';
-import { Model } from '../bridging/model';
 import { EngineErrorCode, unwrapOrThrow } from '../error';
+import { ChatModelInteractor } from './chatModelInteractor';
 import { EngineInteractor } from './engineInteractor';
 import { Interactor, InteractorEntity } from './interactor';
-import { ModelInteractor } from './modelInteractor';
 
-export class ModelsInteractor implements Interactor<Model[]> {
+export class ChatModelsInteractor implements Interactor<ChatModel[]> {
     readonly engineInteractor: EngineInteractor;
-    readonly entity: InteractorEntity<Model[]>;
+    readonly entity: InteractorEntity<ChatModel[]>;
 
-    constructor(engineInteractor: EngineInteractor, models: InteractorEntity<Model[]>) {
+    constructor(engineInteractor: EngineInteractor, models: InteractorEntity<ChatModel[]>) {
         this.engineInteractor = engineInteractor;
         this.entity = models;
     }
 
-    async finalize(): Promise<Model[]> {
+    async finalize(): Promise<ChatModel[]> {
         return await this.entity;
     }
 
-    get(repoId: string): ModelInteractor {
+    get(repoId: string): ChatModelInteractor {
         const modelPromise = (async () => {
             const models = await this.finalize();
             const model = models.find((model) => model.repoId === repoId);
             return unwrapOrThrow(model, EngineErrorCode.ModelNotFound);
         })();
-        return new ModelInteractor(this, modelPromise, Config.default());
+        return new ChatModelInteractor(this, modelPromise, Config.default());
     }
 }