Skip to content

Commit ea8e1a3

Browse files
WauplinWauplin
committed
Update tiny-agent format to follow VSCode format
1 parent 56e9f9d commit ea8e1a3

File tree

2 files changed

+58
-62
lines changed

2 files changed

+58
-62
lines changed

‎packages/tiny-agents/src/cli.ts‎

Lines changed: 50 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ import * as readline from "node:readline/promises";
44
import { stdin, stdout } from "node:process";
55
import { z } from "zod";
66
import { PROVIDERS_OR_POLICIES } from "@huggingface/inference";
7+
import type { ServerConfig } from "@huggingface/mcp-client";
78
import { Agent } from "@huggingface/mcp-client";
89
import { version as packageVersion } from "../package.json";
910
import { InputConfigSchema, ServerConfigSchema } from "./lib/types";
@@ -107,15 +108,15 @@ async function main() {
107108
// Check env variables that will use this input
108109
const inputVars = new Set<string>();
109110
for (const server of config.servers) {
110-
if (server.type === "stdio" && server.config.env) {
111-
for (const [key, value] of Object.entries(server.config.env)) {
111+
if (server.type === "stdio" && server.env) {
112+
for (const [key, value] of Object.entries(server.env)) {
112113
if (value === envSpecialValue) {
113114
inputVars.add(key);
114115
}
115116
}
116117
}
117-
if ((server.type === "http" || server.type === "sse") && server.config.options?.requestInit?.headers) {
118-
for (const [key, value] of Object.entries(server.config.options.requestInit.headers)) {
118+
if ((server.type === "http" || server.type === "sse") && server.headers) {
119+
for (const [key, value] of Object.entries(server.headers)) {
119120
if (value.includes(envSpecialValue)) {
120121
inputVars.add(key);
121122
}
@@ -132,54 +133,56 @@ async function main() {
132133
}
133134

134135
// Prompt user for input
136+
const envVariableKey = inputId.replaceAll("-", "_").toUpperCase();
135137
stdout.write(ANSI.BLUE);
136138
stdout.write(` • ${inputId}`);
137139
stdout.write(ANSI.RESET);
138-
stdout.write(`: ${description}. (default: load from ${Array.from(inputVars).join(", ")}) `);
140+
stdout.write(`: ${description}. (default: load from ${envVariableKey}) `);
141+
stdout.write("\n");
139142

140143
const userInput = (await rl.question("")).trim();
141144

142145
// Inject user input (or env variable) into servers' env
143146
for (const server of config.servers) {
144-
if (server.type === "stdio" && server.config.env) {
145-
for (const [key, value] of Object.entries(server.config.env)) {
147+
if (server.type === "stdio" && server.env) {
148+
for (const [key, value] of Object.entries(server.env)) {
146149
if (value === envSpecialValue) {
147150
if (userInput) {
148-
server.config.env[key] = userInput;
151+
server.env[key] = userInput;
149152
} else {
150-
const valueFromEnv = process.env[key] || "";
151-
server.config.env[key] = valueFromEnv;
153+
const valueFromEnv = process.env[envVariableKey] || "";
154+
server.env[key] = valueFromEnv;
152155
if (valueFromEnv) {
153156
stdout.write(ANSI.GREEN);
154-
stdout.write(`Value successfully loaded from '${key}'`);
157+
stdout.write(`Value successfully loaded from '${envVariableKey}'`);
155158
stdout.write(ANSI.RESET);
156159
stdout.write("\n");
157160
} else {
158161
stdout.write(ANSI.YELLOW);
159-
stdout.write(`No value found for '${key}' in environment variables. Continuing.`);
162+
stdout.write(`No value found for '${envVariableKey}' in environment variables. Continuing.`);
160163
stdout.write(ANSI.RESET);
161164
stdout.write("\n");
162165
}
163166
}
164167
}
165168
}
166169
}
167-
if ((server.type === "http" || server.type === "sse") && server.config.options?.requestInit?.headers) {
168-
for (const [key, value] of Object.entries(server.config.options.requestInit.headers)) {
170+
if ((server.type === "http" || server.type === "sse") && server.headers) {
171+
for (const [key, value] of Object.entries(server.headers)) {
169172
if (value.includes(envSpecialValue)) {
170173
if (userInput) {
171-
server.config.options.requestInit.headers[key] = value.replace(envSpecialValue, userInput);
174+
server.headers[key] = value.replace(envSpecialValue, userInput);
172175
} else {
173-
const valueFromEnv = process.env[key] || "";
174-
server.config.options.requestInit.headers[key] = value.replace(envSpecialValue, valueFromEnv);
176+
const valueFromEnv = process.env[envVariableKey] || "";
177+
server.headers[key] = value.replace(envSpecialValue, valueFromEnv);
175178
if (valueFromEnv) {
176179
stdout.write(ANSI.GREEN);
177-
stdout.write(`Value successfully loaded from '${key}'`);
180+
stdout.write(`Value successfully loaded from '${envVariableKey}'`);
178181
stdout.write(ANSI.RESET);
179182
stdout.write("\n");
180183
} else {
181184
stdout.write(ANSI.YELLOW);
182-
stdout.write(`No value found for '${key}' in environment variables. Continuing.`);
185+
stdout.write(`No value found for '${envVariableKey}' in environment variables. Continuing.`);
183186
stdout.write(ANSI.RESET);
184187
stdout.write("\n");
185188
}
@@ -194,21 +197,47 @@ async function main() {
194197
rl.close();
195198
}
196199

200+
const formattedServers: ServerConfig[] = config.servers.map((server) => {
201+
switch (server.type) {
202+
case "stdio":
203+
return {
204+
type: "stdio",
205+
config: {
206+
command: server.command,
207+
args: server.args ?? [],
208+
env: server.env ?? {},
209+
cwd: server.cwd ?? process.cwd(),
210+
},
211+
};
212+
case "http":
213+
case "sse":
214+
return {
215+
type: server.type,
216+
config: {
217+
url: server.url,
218+
requestInit: {
219+
headers: server.headers,
220+
},
221+
},
222+
};
223+
}
224+
});
225+
197226
const agent = new Agent(
198227
config.endpointUrl
199228
? {
200229
endpointUrl: config.endpointUrl,
201230
model: config.model,
202231
apiKey: config.apiKey ?? process.env.API_KEY ?? process.env.HF_TOKEN,
203-
servers: config.servers,
232+
servers: formattedServers,
204233
prompt,
205234
}
206235
: {
207236
// eslint-disable-next-line @typescript-eslint/no-non-null-assertion
208237
provider: config.provider!,
209238
model: config.model,
210239
apiKey: config.apiKey ?? process.env.API_KEY ?? process.env.HF_TOKEN,
211-
servers: config.servers,
240+
servers: formattedServers,
212241
prompt,
213242
}
214243
);

‎packages/tiny-agents/src/lib/types.ts‎

Lines changed: 8 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -8,53 +8,20 @@ export interface TinyAgentConfig {
88
export const ServerConfigSchema = z.discriminatedUnion("type", [
99
z.object({
1010
type: z.literal("stdio"),
11-
config: z.object({
12-
command: z.string(),
13-
args: z.array(z.string()).optional(),
14-
env: z.record(z.string()).optional(),
15-
cwd: z.string().optional(),
16-
}),
11+
command: z.string(),
12+
args: z.array(z.string()).optional(),
13+
env: z.record(z.string()).optional(),
14+
cwd: z.string().optional(),
1715
}),
1816
z.object({
1917
type: z.literal("http"),
20-
config: z.object({
21-
url: z.union([z.string(), z.string().url()]),
22-
options: z
23-
.object({
24-
/**
25-
* Customizes HTTP requests to the server.
26-
*/
27-
requestInit: z
28-
.object({
29-
headers: z.record(z.string()).optional(),
30-
})
31-
.optional(),
32-
/**
33-
* Session ID for the connection. This is used to identify the session on the server.
34-
* When not provided and connecting to a server that supports session IDs, the server will generate a new session ID.
35-
*/
36-
sessionId: z.string().optional(),
37-
})
38-
.optional(),
39-
}),
18+
url: z.union([z.string(), z.string().url()]),
19+
headers: z.record(z.string()).optional(),
4020
}),
4121
z.object({
4222
type: z.literal("sse"),
43-
config: z.object({
44-
url: z.union([z.string(), z.string().url()]),
45-
options: z
46-
.object({
47-
/**
48-
* Customizes HTTP requests to the server.
49-
*/
50-
requestInit: z
51-
.object({
52-
headers: z.record(z.string()).optional(),
53-
})
54-
.optional(),
55-
})
56-
.optional(),
57-
}),
23+
url: z.union([z.string(), z.string().url()]),
24+
headers: z.record(z.string()).optional(),
5825
}),
5926
]);
6027

0 commit comments

Comments
 (0)