Skip to content

Commit c2ebd78

Browse files
wangzuowangzuo
committed
add schema file
1 parent 4e45a11 commit c2ebd78

File tree

12 files changed

+335
-320
lines changed

12 files changed

+335
-320
lines changed

‎src/jobs/chat/index.ts‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
export * from "./job";
2+
export * from "./schema";

src/jobs/chat.ts renamed to src/jobs/chat/job.ts

Lines changed: 6 additions & 106 deletions
Original file line numberDiff line numberDiff line change
@@ -1,111 +1,11 @@
11
import { z, ZodSchema } from "zod";
22
import zodToJsonSchema from "zod-to-json-schema";
3-
import { Job } from "./job";
4-
import { JobBaseSchema } from "./schema";
5-
6-
export const MessageContentSchema = z.union([
7-
z.string(),
8-
z.array(
9-
z.object({
10-
type: z.literal("text"),
11-
text: z.string(),
12-
})
13-
),
14-
z.array(
15-
z.object({
16-
type: z.literal("image"),
17-
image_url: z.string().optional(),
18-
source: z
19-
.object({
20-
type: z.literal("base64"),
21-
data: z.string(),
22-
media_type: z.enum([
23-
"image/jpeg",
24-
"image/png",
25-
"image/webp",
26-
"image/gif",
27-
]),
28-
})
29-
.optional(),
30-
})
31-
),
32-
]);
33-
34-
export const MessageSchema = z.object({
35-
role: z.enum(["system", "user", "assistant"]),
36-
content: MessageContentSchema,
37-
});
38-
39-
export const StreamOptionsSchema = z.object({
40-
includeUsage: z.boolean().optional(),
41-
});
42-
43-
export const ResponseFormatSchema = z.object({
44-
type: z.enum(["json_object", "json_schema"]),
45-
json_schema: z.any().optional(),
46-
});
47-
48-
export const ChatToolSchema = z.object({
49-
params: z.object({
50-
name: z.string(),
51-
description: z.string().optional(),
52-
parameters: z.any().optional(),
53-
}),
54-
toJSON: z.function().returns(z.any()).optional(),
55-
});
56-
57-
export const JsonSchemaDefSchema = z.object({
58-
name: z.string(),
59-
description: z.string().optional(),
60-
schema: z.any(),
61-
});
62-
63-
export const ChatJobParamsSchema = z.object({
64-
temperature: z.number().optional(),
65-
stream: z.boolean().optional(),
66-
streamOptions: StreamOptionsSchema.optional(),
67-
maxTokens: z.number().optional(),
68-
messages: z.array(MessageSchema),
69-
tools: z.array(ChatToolSchema).optional(),
70-
toolChoice: z.string().optional(),
71-
responseFormat: ResponseFormatSchema.optional(),
72-
topP: z.number().optional(),
73-
topK: z.number().optional(),
74-
systemPrompt: z.string().optional(),
75-
jsonSchema: JsonSchemaDefSchema.optional(),
76-
});
77-
78-
export const ChatResultSchema = z.object({
79-
message: z.object({
80-
role: z.literal("assistant"),
81-
content: z.string().nullable(),
82-
}),
83-
usage: z
84-
.object({
85-
prompt_tokens: z.number(),
86-
completion_tokens: z.number(),
87-
total_tokens: z.number(),
88-
})
89-
.optional(),
90-
tool_calls: z
91-
.array(
92-
z.object({
93-
name: z.string(),
94-
arguments: z.record(z.any()),
95-
})
96-
)
97-
.optional(),
98-
});
99-
100-
export const ChatJobSchema = JobBaseSchema.extend({
101-
type: z.literal("chat"),
102-
model: z.string(),
103-
params: ChatJobParamsSchema,
104-
result: ChatResultSchema.optional(),
105-
});
106-
107-
export type ChatJobSchemaType = z.infer<typeof ChatJobSchema>;
108-
export type ChatJobParams = z.infer<typeof ChatJobParamsSchema>;
3+
import { Job } from "../job";
4+
import type {
5+
ChatJobSchemaType,
6+
ChatJobParams,
7+
ChatResultSchema,
8+
} from "./schema";
1099

11010
export function systemPrompt(content: string) {
11111
return { role: "system", content };

‎src/jobs/chat/schema.ts‎

Lines changed: 106 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,106 @@
1+
import { z } from "zod";
2+
import { JobBaseSchema } from "../schema";
3+
4+
export const MessageContentSchema = z.union([
5+
z.string(),
6+
z.array(
7+
z.object({
8+
type: z.literal("text"),
9+
text: z.string(),
10+
})
11+
),
12+
z.array(
13+
z.object({
14+
type: z.literal("image"),
15+
image_url: z.string().optional(),
16+
source: z
17+
.object({
18+
type: z.literal("base64"),
19+
data: z.string(),
20+
media_type: z.enum([
21+
"image/jpeg",
22+
"image/png",
23+
"image/webp",
24+
"image/gif",
25+
]),
26+
})
27+
.optional(),
28+
})
29+
),
30+
]);
31+
32+
export const MessageSchema = z.object({
33+
role: z.enum(["system", "user", "assistant"]),
34+
content: MessageContentSchema,
35+
});
36+
37+
export const StreamOptionsSchema = z.object({
38+
includeUsage: z.boolean().optional(),
39+
});
40+
41+
export const ResponseFormatSchema = z.object({
42+
type: z.enum(["json_object", "json_schema"]),
43+
json_schema: z.any().optional(),
44+
});
45+
46+
export const ChatToolSchema = z.object({
47+
params: z.object({
48+
name: z.string(),
49+
description: z.string().optional(),
50+
parameters: z.any().optional(),
51+
}),
52+
toJSON: z.function().returns(z.any()).optional(),
53+
});
54+
55+
export const JsonSchemaDefSchema = z.object({
56+
name: z.string(),
57+
description: z.string().optional(),
58+
schema: z.any(),
59+
});
60+
61+
export const ChatJobParamsSchema = z.object({
62+
temperature: z.number().optional(),
63+
stream: z.boolean().optional(),
64+
streamOptions: StreamOptionsSchema.optional(),
65+
maxTokens: z.number().optional(),
66+
messages: z.array(MessageSchema),
67+
tools: z.array(ChatToolSchema).optional(),
68+
toolChoice: z.string().optional(),
69+
responseFormat: ResponseFormatSchema.optional(),
70+
topP: z.number().optional(),
71+
topK: z.number().optional(),
72+
systemPrompt: z.string().optional(),
73+
jsonSchema: JsonSchemaDefSchema.optional(),
74+
});
75+
76+
export const ChatResultSchema = z.object({
77+
message: z.object({
78+
role: z.literal("assistant"),
79+
content: z.string().nullable(),
80+
}),
81+
usage: z
82+
.object({
83+
prompt_tokens: z.number(),
84+
completion_tokens: z.number(),
85+
total_tokens: z.number(),
86+
})
87+
.optional(),
88+
tool_calls: z
89+
.array(
90+
z.object({
91+
name: z.string(),
92+
arguments: z.record(z.any()),
93+
})
94+
)
95+
.optional(),
96+
});
97+
98+
export const ChatJobSchema = JobBaseSchema.extend({
99+
type: z.literal("chat"),
100+
model: z.string(),
101+
params: ChatJobParamsSchema,
102+
result: ChatResultSchema.optional(),
103+
});
104+
105+
export type ChatJobSchemaType = z.infer<typeof ChatJobSchema>;
106+
export type ChatJobParams = z.infer<typeof ChatJobParamsSchema>;

‎src/jobs/embedding.ts‎

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

‎src/jobs/embedding/index.ts‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
export * from "./job";
2+
export * from "./schema";

‎src/jobs/embedding/job.ts‎

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,39 @@
1+
import { Job } from "../job";
2+
import type { EmbeddingJobParams, EmbeddingJobSchemaType } from "./schema";
3+
4+
export class EmbeddingJob<T extends EmbeddingJobSchemaType> extends Job<T> {
5+
model: string;
6+
params: EmbeddingJobParams;
7+
8+
constructor(model: string) {
9+
super();
10+
this.model = model;
11+
this.params = {};
12+
}
13+
14+
input(_input: string) {
15+
this.params.input = _input;
16+
return this;
17+
}
18+
19+
dimensions(_dimensions: number) {
20+
this.params.dimensions = _dimensions;
21+
return this;
22+
}
23+
24+
encodingFormat(_encodingFormat: string) {
25+
this.params.encodingFormat = _encodingFormat;
26+
return this;
27+
}
28+
29+
dump() {
30+
const obj = super.dump();
31+
return {
32+
...obj,
33+
type: "embedding" as const,
34+
model: this.model,
35+
params: this.params,
36+
provider: this.provider,
37+
};
38+
}
39+
}

‎src/jobs/embedding/schema.ts‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,28 @@
1+
import { z } from "zod";
2+
import { JobBaseSchema } from "../schema";
3+
4+
export const EmbeddedJobParamsSchema = z.object({
5+
input: z.string().optional(),
6+
dimensions: z.number().optional(),
7+
encodingFormat: z.string().optional(),
8+
});
9+
10+
export type EmbeddingJobParams = z.infer<typeof EmbeddedJobParamsSchema>;
11+
12+
const EmbeddingResultSchema = z.object({
13+
embedding: z.array(z.number()),
14+
usage: z
15+
.object({
16+
prompt_tokens: z.number(),
17+
total_tokens: z.number(),
18+
})
19+
.optional(),
20+
});
21+
22+
export const EmbeddingJobSchema = JobBaseSchema.extend({
23+
type: z.literal("embedding"),
24+
model: z.string(),
25+
params: EmbeddedJobParamsSchema,
26+
result: EmbeddingResultSchema.optional(),
27+
});
28+
export type EmbeddingJobSchemaType = z.infer<typeof EmbeddingJobSchema>;

0 commit comments

Comments
 (0)