@trymirai/uzu 0.1.19 → 0.1.20

package/README.md

@@ -26,7 +26,7 @@ Add the `uzu` dependency to your project's `package.json`:
 
 ```json
 "dependencies": {
-    "@trymirai/uzu": "0.1.19"
+    "@trymirai/uzu": "0.1.20"
 }
 ```
 
@@ -41,7 +41,7 @@ const licenseStatus: LicenseStatus = await engine.activate(apiKey)
 
 ```ts
 const localModels: LocalModel[] = await engine.updateRegistry()
-const localModelId = 'Meta-Llama-3.2-1B-Instruct-bfloat16'
+const localModelId = 'Meta-Llama-3.2-1B-Instruct'
 ```
 
 ### Download with progress handle
@@ -51,6 +51,7 @@ const donwloadHandle = engine.downloadHandle(localModelId)
 donwloadHandle.start()
 
 for await (const donwloadProgress of donwloadHandle.progress()) {
+    // Implement a custom download progress handler
     handleProgress(donwloadProgress)
 }
 ```
@@ -58,13 +59,13 @@ for await (const donwloadProgress of donwloadHandle.progress()) {
 Alternatively, you may use engine to control and observe model download:
 
 ```ts
+const modelState: ModelDownloadState = engine.getState(localModelId)
+
 // engine.download(localModelId)
 // engine.pause(localModelId)
 // engine.resume(localModelId)
 // engine.stop(localModelId)
 // engine.delete(localModelId)
-
-const modelState: ModelDownloadState = engine.getState(localModelId)
 ```
 
 ### Session
@@ -72,35 +73,30 @@ const modelState: ModelDownloadState = engine.getState(localModelId)
 `Session` is the core entity used to communicate with the model:
 
 ```ts
-// Choose one of the two options below by commenting/uncommenting:
-// const modelId: ModelID = { type: 'Cloud', id: cloudRepoId }
 const modelId: ModelID = { type: 'Local', id: localModelId }
-const session: Session = engine.createSession(modelId)
-```
-
-### Chat
-
-Load `Session` with a chat-configured config and run it with a specific prompt or a list of messages:
-
-```ts
-const config: SessionConfig = {
+const config: Config = {
     preset: { type: 'General' },
-    samplingSeed: { type: 'Custom', seed: 12345 },
+    prefillStepSize: { type: 'Default' },
     contextLength: { type: 'Default' },
+    samplingSeed: { type: 'Default' },
 }
-session.load(config)
+const session: Session = engine.createSession(modelId, config)
 ```
 
+### Chat
+
+After creating it, you can run the `Session` with a specific prompt or a list of messages:
+
 ```ts
-const input: SessionInput = {
+const input: Input = {
     type: 'Messages',
     messages: [
         {
-            role: SessionMessageRole.System,
+            role: Role.System,
             content: 'You are a helpful assistant.'
         },
         {
-            role: SessionMessageRole.User,
+            role: Role.User,
             content: 'Tell me a short, funny story about a robot.'
         },
     ],
@@ -108,14 +104,14 @@ const input: SessionInput = {
 ```
 
 ```ts
-const runConfig: SessionRunConfig = {
+const runConfig: RunConfig = {
     tokensLimit: 128,
-    samplingConfig: { type: 'Argmax' },
+    enableThinking: true,
+    samplingPolicy: { type: 'Default' },
 }
-```
 
-```ts
 const output = session.run(input, runConfig, (partialOutput) => {
+    // Implement a custom partial output handler
     return handlePartialOutput(partialOutput)
 })
 ```
@@ -125,30 +121,34 @@ const output = session.run(input, runConfig, (partialOutput) => {
 In this example, we will extract a summary of the input text:
 
 ```ts
-const config: SessionConfig = {
+const modelId: ModelID = { type: 'Local', id: localModelId }
+const config: Config = {
     preset: { type: 'Summarization' },
-    samplingSeed: { type: 'Default' },
+    prefillStepSize: { type: 'Default' },
     contextLength: { type: 'Default' },
+    samplingSeed: { type: 'Default' },
 }
-session.load(config)
+const session = engine.createSession(modelId, config)
 ```
 
 ```ts
-const input: SessionInput = {
+const textToSummarize =
+    'A Large Language Model (LLM) is a type of AI that processes and generates text using transformer-based architectures trained on vast datasets. They power chatbots, translation, code assistants, and more.'
+const input: Input = {
     type: 'Text',
     text: `Text is: "${textToSummarize}". Write only summary itself.`,
 }
 ```
 
 ```ts
-const runConfig: SessionRunConfig = {
+const runConfig: RunConfig = {
     tokensLimit: 256,
-    samplingConfig: { type: 'Argmax' },
+    enableThinking: true,
+    samplingPolicy: { type: 'Custom', value: { type: 'Greedy' } },
 }
-```
 
-```ts
 const output = session.run(input, runConfig, (partialOutput) => {
+    // Implement a custom partial output handler
     return handlePartialOutput(partialOutput)
 })
 ```
@@ -158,19 +158,19 @@ const output = session.run(input, runConfig, (partialOutput) => {
 Let’s look at a case where you need to classify input text based on a specific feature, such as `sentiment`:
 
 ```ts
-const feature: SessionClassificationFeature = {
+const feature: ClassificationFeature = {
     name: 'sentiment',
     values: ['Happy', 'Sad', 'Angry', 'Fearful', 'Surprised', 'Disgusted'],
 }
-```
-
-```ts
-const config: SessionConfig = {
+const config: Config = {
     preset: { type: 'Classification', feature },
-    samplingSeed: { type: 'Default' },
+    prefillStepSize: { type: 'Default' },
     contextLength: { type: 'Default' },
+    samplingSeed: { type: 'Default' },
 }
-session.load(config)
+
+const modelId: ModelID = { type: 'Local', id: localModelId }
+const session = engine.createSession(modelId, config)
 ```
 
 ```ts
@@ -179,18 +179,18 @@ const textToDetectFeature =
 const classificationPrompt =
     `Text is: "${textToDetectFeature}". Choose ${feature.name} from the list: ${feature.values.join(', ')}. ` +
     "Answer with one word. Don't add a dot at the end."
-const input: SessionInput = { type: 'Text', text: classificationPrompt }
+const input: Input = { type: 'Text', text: classificationPrompt }
 ```
 
 ```ts
-const runConfig: SessionRunConfig = {
+const runConfig: RunConfig = {
     tokensLimit: 32,
-    samplingConfig: { type: 'Argmax' },
+    enableThinking: true,
+    samplingPolicy: { type: 'Custom', value: { type: 'Greedy' } },
 }
-```
 
-```ts
 const output = session.run(input, runConfig, (partialOutput) => {
+    // Implement a custom partial output handler
     return handlePartialOutput(partialOutput)
 })
 ```
@@ -199,6 +199,4 @@ In this example, you will get the answer `Happy` immediately after the prefill s
 
 ## License
 
-This project is licensed under the MIT License. See the [LICENSE](LICENSE) file for details.
-
-
+This project is licensed under the MIT License. See the [LICENSE](LICENSE) file for details.

package/README.orig.md

@@ -25,7 +25,7 @@ Node package for [uzu](https://github.com/trymirai/uzu), a **high-performance**
 Set up your project via [Platform](https://platform.trymirai.com), obtain an `MIRAI_API_KEY`.
 
 ```bash
-# Set your API key in `examples/api_key.ts`, then run the examples
+# Set your API key in `examples/common.ts`, then run the examples
 pnpm run chat
 pnpm run summarisation
 pnpm run classification
@@ -37,7 +37,7 @@ Add the `uzu` dependency to your project's `package.json`:
 
 ```json
 "dependencies": {
-    "@trymirai/uzu": "0.1.19"
+    "@trymirai/uzu": "0.1.20"
 }
 ```
 
@@ -52,7 +52,7 @@ const licenseStatus: LicenseStatus = await engine.activate(apiKey)
 
 ```ts
 const localModels: LocalModel[] = await engine.updateRegistry()
-const localModelId = 'Meta-Llama-3.2-1B-Instruct-bfloat16'
+const localModelId = 'Meta-Llama-3.2-1B-Instruct'
 ```
 
 ### Download with progress handle
@@ -62,6 +62,7 @@ const donwloadHandle = engine.downloadHandle(localModelId)
 donwloadHandle.start()
 
 for await (const donwloadProgress of donwloadHandle.progress()) {
+    // Implement a custom download progress handler
     handleProgress(donwloadProgress)
 }
 ```
@@ -69,13 +70,13 @@ for await (const donwloadProgress of donwloadHandle.progress()) {
 Alternatively, you may use engine to control and observe model download:
 
 ```ts
+const modelState: ModelDownloadState = engine.getState(localModelId)
+
 // engine.download(localModelId)
 // engine.pause(localModelId)
 // engine.resume(localModelId)
 // engine.stop(localModelId)
 // engine.delete(localModelId)
-
-const modelState: ModelDownloadState = engine.getState(localModelId)
 ```
 
 ### Session
@@ -83,35 +84,30 @@ const modelState: ModelDownloadState = engine.getState(localModelId)
 `Session` is the core entity used to communicate with the model:
 
 ```ts
-// Choose one of the two options below by commenting/uncommenting:
-// const modelId: ModelID = { type: 'Cloud', id: cloudRepoId }
 const modelId: ModelID = { type: 'Local', id: localModelId }
-const session: Session = engine.createSession(modelId)
-```
-
-### Chat
-
-Load `Session` with a chat-configured config and run it with a specific prompt or a list of messages:
-
-```ts
-const config: SessionConfig = {
+const config: Config = {
     preset: { type: 'General' },
-    samplingSeed: { type: 'Custom', seed: 12345 },
+    prefillStepSize: { type: 'Default' },
     contextLength: { type: 'Default' },
+    samplingSeed: { type: 'Default' },
 }
-session.load(config)
+const session: Session = engine.createSession(modelId, config)
 ```
 
+### Chat
+
+After creating it, you can run the `Session` with a specific prompt or a list of messages:
+
 ```ts
-const input: SessionInput = {
+const input: Input = {
     type: 'Messages',
     messages: [
         {
-            role: SessionMessageRole.System,
+            role: Role.System,
             content: 'You are a helpful assistant.'
         },
         {
-            role: SessionMessageRole.User,
+            role: Role.User,
             content: 'Tell me a short, funny story about a robot.'
         },
     ],
@@ -119,14 +115,14 @@ const input: SessionInput = {
 ```
 
 ```ts
-const runConfig: SessionRunConfig = {
+const runConfig: RunConfig = {
     tokensLimit: 128,
-    samplingConfig: { type: 'Argmax' },
+    enableThinking: true,
+    samplingPolicy: { type: 'Default' },
 }
-```
 
-```ts
 const output = session.run(input, runConfig, (partialOutput) => {
+    // Implement a custom partial output handler
     return handlePartialOutput(partialOutput)
 })
 ```
@@ -136,30 +132,34 @@ const output = session.run(input, runConfig, (partialOutput) => {
 In this example, we will extract a summary of the input text:
 
 ```ts
-const config: SessionConfig = {
+const modelId: ModelID = { type: 'Local', id: localModelId }
+const config: Config = {
     preset: { type: 'Summarization' },
-    samplingSeed: { type: 'Default' },
+    prefillStepSize: { type: 'Default' },
     contextLength: { type: 'Default' },
+    samplingSeed: { type: 'Default' },
 }
-session.load(config)
+const session = engine.createSession(modelId, config)
 ```
 
 ```ts
-const input: SessionInput = {
+const textToSummarize =
+    'A Large Language Model (LLM) is a type of AI that processes and generates text using transformer-based architectures trained on vast datasets. They power chatbots, translation, code assistants, and more.'
+const input: Input = {
     type: 'Text',
     text: `Text is: "${textToSummarize}". Write only summary itself.`,
 }
 ```
 
 ```ts
-const runConfig: SessionRunConfig = {
+const runConfig: RunConfig = {
     tokensLimit: 256,
-    samplingConfig: { type: 'Argmax' },
+    enableThinking: true,
+    samplingPolicy: { type: 'Custom', value: { type: 'Greedy' } },
 }
-```
 
-```ts
 const output = session.run(input, runConfig, (partialOutput) => {
+    // Implement a custom partial output handler
     return handlePartialOutput(partialOutput)
 })
 ```
@@ -169,19 +169,19 @@ const output = session.run(input, runConfig, (partialOutput) => {
 Let’s look at a case where you need to classify input text based on a specific feature, such as `sentiment`:
 
 ```ts
-const feature: SessionClassificationFeature = {
+const feature: ClassificationFeature = {
     name: 'sentiment',
     values: ['Happy', 'Sad', 'Angry', 'Fearful', 'Surprised', 'Disgusted'],
 }
-```
-
-```ts
-const config: SessionConfig = {
+const config: Config = {
     preset: { type: 'Classification', feature },
-    samplingSeed: { type: 'Default' },
+    prefillStepSize: { type: 'Default' },
     contextLength: { type: 'Default' },
+    samplingSeed: { type: 'Default' },
 }
-session.load(config)
+
+const modelId: ModelID = { type: 'Local', id: localModelId }
+const session = engine.createSession(modelId, config)
 ```
 
 ```ts
@@ -190,18 +190,18 @@ const textToDetectFeature =
 const classificationPrompt =
     `Text is: "${textToDetectFeature}". Choose ${feature.name} from the list: ${feature.values.join(', ')}. ` +
     "Answer with one word. Don't add a dot at the end."
-const input: SessionInput = { type: 'Text', text: classificationPrompt }
+const input: Input = { type: 'Text', text: classificationPrompt }
 ```
 
 ```ts
-const runConfig: SessionRunConfig = {
+const runConfig: RunConfig = {
     tokensLimit: 32,
-    samplingConfig: { type: 'Argmax' },
+    enableThinking: true,
+    samplingPolicy: { type: 'Custom', value: { type: 'Greedy' } },
 }
-```
 
-```ts
 const output = session.run(input, runConfig, (partialOutput) => {
+    // Implement a custom partial output handler
     return handlePartialOutput(partialOutput)
 })
 ```
@@ -210,6 +210,4 @@ In this example, you will get the answer `Happy` immediately after the prefill s
 
 ## License
 
-This project is licensed under the MIT License. See the [LICENSE](LICENSE) file for details.
-
-
+This project is licensed under the MIT License. See the [LICENSE](LICENSE) file for details.

package/README.src.md

@@ -25,7 +25,7 @@ Node package for [uzu](https://github.com/trymirai/uzu), a **high-performance**
 Set up your project via [Platform](https://platform.trymirai.com), obtain an `MIRAI_API_KEY`.
 
 ```bash
-# Set your API key in `examples/api_key.ts`, then run the examples
+# Set your API key in `examples/common.ts`, then run the examples
 pnpm run chat
 pnpm run summarisation
 pnpm run classification
@@ -56,13 +56,13 @@ Create and activate engine:
 ### Download with progress handle
 
 ```ts
-// include:examples/chat.ts#download lang=ts
+// include:examples/chat.ts#download-handle lang=ts
 ```
 
 Alternatively, you may use engine to control and observe model download:
 
 ```ts
-// include:examples/chat.ts#model-state lang=ts
+// include:examples/chat.ts#storage-methods lang=ts
 ```
 
 ### Session
@@ -70,27 +70,19 @@ Alternatively, you may use engine to control and observe model download:
 `Session` is the core entity used to communicate with the model:
 
 ```ts
-// include:examples/chat.ts#session-create lang=ts
+// include:examples/chat.ts#session-create-general lang=ts
 ```
 
 ### Chat
 
-Load `Session` with a chat-configured config and run it with a specific prompt or a list of messages:
+After creating it, you can run the `Session` with a specific prompt or a list of messages:
 
 ```ts
-// include:examples/chat.ts#session-load lang=ts
+// include:examples/chat.ts#session-input-general lang=ts
 ```
 
 ```ts
-// include:examples/chat.ts#session-input lang=ts
-```
-
-```ts
-// include:examples/chat.ts#session-run-config lang=ts
-```
-
-```ts
-// include:examples/chat.ts#session-run lang=ts
+// include:examples/chat.ts#session-run-general lang=ts
 ```
 
 ### Summarization
@@ -98,19 +90,15 @@ Load `Session` with a chat-configured config and run it with a specific prompt o
 In this example, we will extract a summary of the input text:
 
 ```ts
-// include:examples/summarisation.ts#session-load lang=ts
-```
-
-```ts
-// include:examples/summarisation.ts#session-input lang=ts
+// include:examples/summarisation.ts#session-create-summarization lang=ts
 ```
 
 ```ts
-// include:examples/summarisation.ts#session-run-config lang=ts
+// include:examples/summarisation.ts#session-input-summarization lang=ts
 ```
 
 ```ts
-// include:examples/summarisation.ts#session-run lang=ts
+// include:examples/summarisation.ts#session-run-summarization lang=ts
 ```
 
 ### Classification
@@ -118,29 +106,19 @@ In this example, we will extract a summary of the input text:
 Let’s look at a case where you need to classify input text based on a specific feature, such as `sentiment`:
 
 ```ts
-// include:examples/classification.ts#classification-feature lang=ts
+// include:examples/classification.ts#session-create-classification lang=ts
 ```
 
 ```ts
-// include:examples/classification.ts#session-load lang=ts
+// include:examples/classification.ts#session-input-classification lang=ts
 ```
 
 ```ts
-// include:examples/classification.ts#session-input lang=ts
-```
-
-```ts
-// include:examples/classification.ts#session-run-config lang=ts
-```
-
-```ts
-// include:examples/classification.ts#session-run lang=ts
+// include:examples/classification.ts#session-run-classification lang=ts
 ```
 
 In this example, you will get the answer `Happy` immediately after the prefill step, and the actual generation won't even start.
 
 ## License
 
-This project is licensed under the MIT License. See the [LICENSE](LICENSE) file for details.
-
-
+This project is licensed under the MIT License. See the [LICENSE](LICENSE) file for details.

package/package.json

@@ -1,11 +1,12 @@
 {
   "name": "@trymirai/uzu",
-  "version": "0.1.19",
+  "version": "0.1.20",
   "private": false,
   "main": "uzu.node",
   "types": "uzu.d.ts",
   "packageManager": "pnpm@10.14.0",
   "scripts": {
+    "quick-start": "ts-node examples/quick-start.ts",
     "chat": "ts-node examples/chat.ts",
     "summarisation": "ts-node examples/summarisation.ts",
     "classification": "ts-node examples/classification.ts"

package/uzu.d.ts

@@ -27,7 +27,7 @@ export declare class Engine {
   updateRegistry(): Promise<Array<LocalModel>>
   activate(apiKey: string): Promise<LicenseStatus>
   constructor()
-  createSession(modelId: ModelID): Session
+  createSession(modelId: ModelID, config: Config): Session
 }
 
 export declare class ProgressStream {
@@ -44,8 +44,12 @@ export declare class ProgressUpdate {
 }
 
 export declare class Session {
-  load(config: SessionConfig): void
-  run(input: SessionInput, runConfig: SessionRunConfig, progressCallback?: (arg: SessionOutput) => boolean | undefined | null): SessionOutput
+  run(input: Input, config: RunConfig, progress?: (arg: Output) => boolean | undefined | null): Output
+}
+
+export interface ClassificationFeature {
+  name: string
+  values: Array<string>
 }
 
 export interface CloudModel {
@@ -54,25 +58,35 @@ export interface CloudModel {
   readonly vendor: string
 }
 
+export interface Config {
+  preset: Preset
+  prefillStepSize: PrefillStepSize
+  contextLength: ContextLength
+  samplingSeed: SamplingSeed
+}
+
 export type ContextLength =
   | { type: 'Default' }
+  | { type: 'Maximal' }
   | { type: 'Custom', length: number }
 
-export interface DecoderTestResult {
-  placementLog: string
-  iterations: number
-  timePerToken: number
-  tokensPerSecond: number
-  success: boolean
-  error?: string
-}
-
 export type DownloaderError =
   | { type: 'Http', message: string }
   | { type: 'Io', message: string }
   | { type: 'AlreadyExists', path: string }
   | { type: 'Generic', message: string }
 
+export declare const enum FinishReason {
+  Stop = 0,
+  Length = 1,
+  Cancelled = 2,
+  ContextLimitReached = 3
+}
+
+export type Input =
+  | { type: 'Text', text: string }
+  | { type: 'Messages', messages: Array<Message> }
+
 export type LicenseStatus =
   | { type: 'NotActivated' }
   | { type: 'PaymentRequired' }
@@ -85,14 +99,12 @@ export type LicenseStatus =
   | { type: 'HttpError', code: number }
 
 export interface LocalModel {
-  /** Unique identifier of the model in the form `<vendor>-<name>-<precision>`. */
+  /** Unique identifier of the model in the form `<vendor>-<name>`. */
   readonly identifier: string
   /** Vendor/author of the model (e.g. "Llama"). */
   readonly vendor: string
-  /** Human-readable model name without vendor/precision (e.g. "3B-Instruct"). */
+  /** Human-readable model name without vendor (e.g. "3B-Instruct"). */
   readonly name: string
-  /** Numerical precision of the weights (e.g. "float16"). */
-  readonly precision: string
   /** Quantization type if the model is quantized (e.g. "uint4"). */
   readonly quantization?: string
   /** Optional regex to parse model output provided by the backend. */
@@ -101,6 +113,12 @@ export interface LocalModel {
   readonly state: ModelDownloadState
 }
 
+export interface Message {
+  role: Role
+  content: string
+  reasoningContent?: string
+}
+
 export interface ModelDownloadState {
   /** Total size of all model files in kilobytes. */
   readonly totalKbytes: number
@@ -116,6 +134,17 @@ export type ModelID =
   | { type: 'Local', id: string }
   | { type: 'Cloud', id: string }
 
+export interface Output {
+  text: Text
+  stats: Stats
+  finishReason?: FinishReason
+}
+
+export interface ParsedText {
+  chainOfThought?: string
+  response?: string
+}
+
 export declare const enum Phase {
   NotDownloaded = 0,
   Downloading = 1,
@@ -124,87 +153,59 @@ export declare const enum Phase {
   Error = 4
 }
 
-export type SamplingConfig =
-  | { type: 'Argmax' }
-  | { type: 'TopP', topP: number }
-  | { type: 'Categorical', temperature: number }
-
-export type SamplingSeed =
+export type PrefillStepSize =
   | { type: 'Default' }
-  | { type: 'Custom', seed: number }
-
-export interface SessionClassificationFeature {
-  name: string
-  values: Array<string>
-}
-
-export interface SessionConfig {
-  preset: SessionPreset
-  samplingSeed: SamplingSeed
-  contextLength: ContextLength
-}
-
-export type SessionInput =
-  | { type: 'Text', text: string }
-  | { type: 'Messages', messages: Array<SessionMessage> }
+  | { type: 'Maximal' }
+  | { type: 'Custom', length: number }
 
-export interface SessionMessage {
-  role: SessionMessageRole
-  content: string
-}
+export type Preset =
+  | { type: 'General' }
+  | { type: 'Classification', feature: ClassificationFeature }
+  | { type: 'Summarization' }
 
-export declare const enum SessionMessageRole {
+export declare const enum Role {
   System = 0,
   User = 1,
   Assistant = 2
 }
 
-export interface SessionOutput {
-  text: string
-  stats: SessionOutputStats
-  finishReason?: SessionOutputFinishReason
-}
-
-export declare const enum SessionOutputFinishReason {
-  Stop = 0,
-  Length = 1,
-  Cancelled = 2
+export interface RunConfig {
+  tokensLimit: number
+  enableThinking: boolean
+  samplingPolicy: SamplingPolicy
 }
 
-export interface SessionOutputRunStats {
+export interface RunStats {
   count: bigint
   averageDuration: number
 }
 
-export interface SessionOutputStats {
-  prefillStats: SessionOutputStepStats
-  generateStats?: SessionOutputStepStats
-  totalStats: SessionOutputTotalStats
+export type SamplingMethod =
+  | { type: 'Greedy' }
+  | { type: 'Temperature', temperature: number }
+  | { type: 'TopP', topP: number }
+
+export type SamplingPolicy =
+  | { type: 'Default' }
+  | { type: 'Custom', value: SamplingMethod }
+
+export type SamplingSeed =
+  | { type: 'Default' }
+  | { type: 'Custom', seed: number }
+
+export interface Stats {
+  prefillStats: StepStats
+  generateStats?: StepStats
+  totalStats: TotalStats
 }
 
-export interface SessionOutputStepStats {
+export interface StepStats {
   duration: number
   suffixLength: bigint
   tokensCount: bigint
   tokensPerSecond: number
-  modelRun: SessionOutputRunStats
-  run?: SessionOutputRunStats
-}
-
-export interface SessionOutputTotalStats {
-  duration: number
-  tokensCountInput: bigint
-  tokensCountOutput: bigint
-}
-
-export type SessionPreset =
-  | { type: 'General' }
-  | { type: 'Classification', feature: SessionClassificationFeature }
-  | { type: 'Summarization' }
-
-export interface SessionRunConfig {
-  tokensLimit: number
-  samplingConfig?: SamplingConfig
+  modelRun: RunStats
+  run?: RunStats
 }
 
 export type StorageError =
@@ -214,4 +215,15 @@ export type StorageError =
   | { type: 'MutexPoisoned', message: string }
   | { type: 'LicenseNotActivated' }
 
+export interface Text {
+  original: string
+  parsed: ParsedText
+}
+
+export interface TotalStats {
+  duration: number
+  tokensCountInput: bigint
+  tokensCountOutput: bigint
+}
+
 export declare function version(): string