Skip to content

Commit 23f374f

Browse files
committed
Add muna.beta.chat.completions.create
1 parent edfc7d0 commit 23f374f

File tree

12 files changed

+358
-5
lines changed

12 files changed

+358
-5
lines changed

Changelog.md

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
## 0.0.47
2-
*INCOMPLETE*
2+
+ Added `muna.beta.chat.completions.create` method to create chat completions mirroring the OpenAI client.
33

44
## 0.0.46
55
+ Function is now Muna!

package.json

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,18 @@
3030
"main": "build/lib/index.js",
3131
"typings": "build/types/index.d.ts",
3232
"unpkg": "build/umd/muna.js",
33+
"exports": {
34+
".": {
35+
"types": "./build/types/index.d.ts",
36+
"import": "./build/lib/index.js",
37+
"require": "./build/lib/index.js"
38+
},
39+
"./beta": {
40+
"types": "./build/types/beta/types.d.ts",
41+
"import": "./build/lib/beta/types.js",
42+
"require": "./build/lib/beta/types.js"
43+
}
44+
},
3345
"dependencies": {
3446
"base64-arraybuffer": "^1.0.2",
3547
"chalk": "^5.3.0",

src/beta/chat/chat.ts

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,20 @@
1+
/*
2+
* Muna
3+
* Copyright © 2025 NatML Inc. All Rights Reserved.
4+
*/
5+
6+
import type { PredictionService } from "../../services"
7+
import type { RemotePredictionService } from "../remote"
8+
import { ChatCompletionsService } from "./completions"
9+
10+
export class ChatService {
11+
12+
/**
13+
* Create completions.
14+
*/
15+
public readonly completions: ChatCompletionsService;
16+
17+
public constructor(predictions: PredictionService, remotePredictions: RemotePredictionService) {
18+
this.completions = new ChatCompletionsService(predictions, remotePredictions);
19+
}
20+
}

src/beta/chat/completions.ts

Lines changed: 110 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,110 @@
1+
/*
2+
* Muna
3+
* Copyright © 2025 NatML Inc. All Rights Reserved.
4+
*/
5+
6+
import type { CreatePredictionInput, PredictionService } from "../../services"
7+
import type { Acceleration, Prediction } from "../../types"
8+
import type { CreateRemotePredictionInput, RemoteAcceleration, RemotePredictionService } from "../remote"
9+
import type { ChatCompletion, ChatCompletionChunk, ChatCompletionMessage } from "./types"
10+
11+
export interface ChatCompletionCreateParamsBase {
12+
/**
13+
* Messages comprising the conversation so far.
14+
*/
15+
messages: ChatCompletionMessage[];
16+
/**
17+
* Chat model predictor tag.
18+
*/
19+
model: string;
20+
/**
21+
* Maximum output tokens.
22+
*/
23+
max_tokens?: number | null;
24+
/**
25+
* Prediction acceleration.
26+
*/
27+
acceleration?: Acceleration | RemoteAcceleration;
28+
}
29+
30+
export interface ChatCompletionCreateParamsNonStreaming extends ChatCompletionCreateParamsBase {
31+
/**
32+
* Whether to stream responses.
33+
*/
34+
stream?: false | null;
35+
}
36+
37+
export interface ChatCompletionCreateParamsStreaming extends ChatCompletionCreateParamsBase {
38+
/**
39+
* Whether to stream responses.
40+
*/
41+
stream: true;
42+
}
43+
44+
export type ChatCompletionCreateParams = (
45+
ChatCompletionCreateParamsNonStreaming |
46+
ChatCompletionCreateParamsStreaming
47+
);
48+
49+
export class ChatCompletionsService {
50+
51+
private readonly predictions: PredictionService;
52+
private readonly remotePredictions: RemotePredictionService;
53+
54+
public constructor(predictions: PredictionService, remotePredictions: RemotePredictionService) {
55+
this.predictions = predictions;
56+
this.remotePredictions = remotePredictions;
57+
}
58+
59+
/**
60+
* Create a chat completion.
61+
*/
62+
public create(body: ChatCompletionCreateParamsNonStreaming): Promise<ChatCompletion>;
63+
public create(body: ChatCompletionCreateParamsStreaming): AsyncGenerator<ChatCompletionChunk>;
64+
public create(body: ChatCompletionCreateParamsBase): Promise<ChatCompletion> | AsyncGenerator<ChatCompletionChunk>;
65+
public create(body: ChatCompletionCreateParams): Promise<ChatCompletion> | AsyncGenerator<ChatCompletionChunk> {
66+
const { model: tag, acceleration = "auto", ...inputs } = body;
67+
inputs.stream = inputs.stream ?? false;
68+
const input = { tag, inputs, acceleration };
69+
if (inputs.stream)
70+
return this.createCompletionStreaming(input);
71+
else
72+
return this.createCompletionNonStreaming(input);
73+
}
74+
75+
private async createCompletionNonStreaming(input: CreatePredictionInput | CreateRemotePredictionInput): Promise<ChatCompletion> {
76+
const prediction = await this.createPrediction(input);
77+
const completion = this.parseResponse(prediction) as ChatCompletion;
78+
return completion;
79+
}
80+
81+
private async * createCompletionStreaming(input: CreatePredictionInput | CreateRemotePredictionInput): AsyncGenerator<ChatCompletionChunk> {
82+
const stream = await this.streamPrediction(input);
83+
for await (const prediction of stream) {
84+
const completion = this.parseResponse(prediction) as ChatCompletionChunk;
85+
yield completion;
86+
}
87+
}
88+
89+
private createPrediction(input: CreatePredictionInput | CreateRemotePredictionInput): Promise<Prediction> {
90+
// muna.beta.predictions.remote.create(...)
91+
if ((input.acceleration as string).startsWith("remote_"))
92+
return this.remotePredictions.create(input as CreateRemotePredictionInput);
93+
// muna.predictions.create(...)
94+
else
95+
return this.predictions.create(input as CreatePredictionInput);
96+
}
97+
98+
private streamPrediction(input: CreatePredictionInput | CreateRemotePredictionInput): AsyncGenerator<Prediction> {
99+
if ((input.acceleration as string).startsWith("remote_"))
100+
throw new Error("Streaming predictions are not supported with remote acceleration");
101+
// muna.predictions.stream(...)
102+
return this.predictions.stream(input as CreatePredictionInput);
103+
}
104+
105+
private parseResponse(prediction: Prediction): ChatCompletion | ChatCompletionChunk {
106+
if (prediction.error)
107+
throw new Error(prediction.error);
108+
return prediction.results[0] as ChatCompletion | ChatCompletionChunk;
109+
}
110+
}

src/beta/chat/index.ts

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
/*
2+
* Muna
3+
* Copyright © 2025 NatML Inc. All Rights Reserved.
4+
*/
5+
6+
export { ChatService } from "./chat"

src/beta/chat/types.ts

Lines changed: 156 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,156 @@
1+
/*
2+
* Muna
3+
* Copyright © 2025 NatML Inc. All Rights Reserved.
4+
*/
5+
6+
export interface ChatCompletion {
7+
/**
8+
* The object type, which is always `chat.completion`.
9+
*/
10+
object: "chat.completion";
11+
/**
12+
* A unique identifier for the chat completion.
13+
*/
14+
id: string;
15+
/**
16+
* The model used for the chat completion.
17+
*/
18+
model: string;
19+
/**
20+
* A list of chat completion choices. Can be more than one if `n` is greater
21+
* than 1.
22+
*/
23+
choices: Array<ChatCompletion.Choice>;
24+
/**
25+
* The Unix timestamp (in seconds) of when the chat completion was created.
26+
*/
27+
created: number;
28+
/**
29+
* Usage statistics for the completion request.
30+
*/
31+
usage?: CompletionUsage;
32+
}
33+
34+
export interface ChatCompletionChunk {
35+
/**
36+
* The object type, which is always `chat.completion.chunk`.
37+
*/
38+
object: "chat.completion.chunk";
39+
/**
40+
* A unique identifier for the chat completion. Each chunk has the same ID.
41+
*/
42+
id: string;
43+
/**
44+
* The model to generate the completion.
45+
*/
46+
model: string;
47+
/**
48+
* A list of chat completion choices. Can contain more than one elements if `n` is
49+
* greater than 1. Can also be empty for the last chunk if you set
50+
* `stream_options: {"include_usage": true}`.
51+
*/
52+
choices: Array<ChatCompletionChunk.Choice>;
53+
/**
54+
* The Unix timestamp (in seconds) of when the chat completion was created. Each
55+
* chunk has the same timestamp.
56+
*/
57+
created: number;
58+
/**
59+
* Usage statistics for the completion request.
60+
*/
61+
usage?: CompletionUsage;
62+
}
63+
64+
export namespace ChatCompletion {
65+
66+
export interface Choice {
67+
/**
68+
* The index of the choice in the list of choices.
69+
*/
70+
index: number;
71+
/**
72+
* A chat completion message generated by the model.
73+
*/
74+
message: ChatCompletionMessage;
75+
/**
76+
* The reason the model stopped generating tokens. This will be `stop` if the model
77+
* hit a natural stop point or a provided stop sequence, `length` if the maximum
78+
* number of tokens specified in the request was reached, `content_filter` if
79+
* content was omitted due to a flag from our content filters, `tool_calls` if the
80+
* model called a tool, or `function_call` (deprecated) if the model called a
81+
* function.
82+
*/
83+
finish_reason: "stop" | "length" | "tool_calls" | "content_filter" | "function_call";
84+
/**
85+
* Log probability information for the choice.
86+
*/
87+
logprobs?: null;
88+
}
89+
}
90+
91+
export namespace ChatCompletionChunk {
92+
93+
export interface Choice {
94+
/**
95+
* The index of the choice in the list of choices.
96+
*/
97+
index: number;
98+
/**
99+
* A chat completion delta generated by streamed model responses.
100+
*/
101+
delta: Choice.Delta;
102+
/**
103+
* The reason the model stopped generating tokens. This will be `stop` if the model
104+
* hit a natural stop point or a provided stop sequence, `length` if the maximum
105+
* number of tokens specified in the request was reached, `content_filter` if
106+
* content was omitted due to a flag from our content filters, `tool_calls` if the
107+
* model called a tool, or `function_call` (deprecated) if the model called a
108+
* function.
109+
*/
110+
finish_reason: "stop" | "length" | "tool_calls" | "content_filter" | "function_call" | null;
111+
/**
112+
* Log probability information for the choice.
113+
*/
114+
logprobs?: null;
115+
}
116+
117+
export namespace Choice {
118+
119+
export interface Delta {
120+
/**
121+
* The role of the author of this message.
122+
*/
123+
role?: "developer" | "system" | "user" | "assistant" | "tool";
124+
/**
125+
* The contents of the chunk message.
126+
*/
127+
content?: string | null;
128+
}
129+
}
130+
}
131+
132+
export interface ChatCompletionMessage {
133+
/**
134+
* The role of the author of this message.
135+
*/
136+
role: "assistant" | "user" | "system";
137+
/**
138+
* The contents of the message.
139+
*/
140+
content: string | null;
141+
}
142+
143+
export interface CompletionUsage {
144+
/**
145+
* Number of tokens in the generated completion.
146+
*/
147+
completion_tokens: number;
148+
/**
149+
* Number of tokens in the prompt.
150+
*/
151+
prompt_tokens: number;
152+
/**
153+
* Total number of tokens used in the request (prompt + completion).
154+
*/
155+
total_tokens: number;
156+
}

src/beta/client.ts

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,8 @@
44
*/
55

66
import type { MunaClient } from "../client"
7+
import type { PredictionService as EdgePredictionService } from "../services"
8+
import { ChatService } from "./chat"
79
import { PredictionService } from "./remote"
810

911
/**
@@ -16,7 +18,13 @@ export class BetaClient {
1618
*/
1719
public readonly predictions: PredictionService;
1820

19-
public constructor(client: MunaClient) {
21+
/**
22+
* Make chat conversations.
23+
*/
24+
public readonly chat: ChatService;
25+
26+
public constructor(client: MunaClient, predictions: EdgePredictionService) {
2027
this.predictions = new PredictionService(client);
28+
this.chat = new ChatService(predictions, this.predictions.remote);
2129
}
2230
}

src/beta/remote/prediction.ts

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,8 +3,8 @@
33
* Copyright © 2025 NatML Inc. All Rights Reserved.
44
*/
55

6-
import { RemotePredictionService } from "./remote"
76
import type { MunaClient } from "../../client"
7+
import { RemotePredictionService } from "./remote"
88

99
export class PredictionService {
1010

src/beta/types.ts

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
/*
2+
* Muna
3+
* Copyright © 2025 NatML Inc. All Rights Reserved.
4+
*/
5+
6+
export type { CreateRemotePredictionInput, RemoteAcceleration } from "./remote"
7+
export * from "./chat/types"

src/index.ts

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,6 @@
33
* Copyright © 2025 NatML Inc. All Rights Reserved.
44
*/
55

6-
export type { CreateRemotePredictionInput, RemoteAcceleration } from "./beta"
76
export type { MunaAPIError } from "./client"
87
export * from "./muna"
98
export * from "./types"

0 commit comments

Comments
 (0)