@rlabs-inc/gemini-mcp 0.6.1 → 0.6.3
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- package/dist/gemini-client.d.ts +35 -0
- package/dist/gemini-client.js +117 -0
- package/dist/index.js +1 -1
- package/dist/tools/deep-research.d.ts +1 -1
- package/dist/tools/deep-research.js +79 -87
- package/dist/tools/token-count.d.ts +1 -1
- package/dist/tools/token-count.js +17 -23
- package/package.json +1 -1
package/dist/gemini-client.d.ts
CHANGED
|
@@ -120,3 +120,38 @@ export declare function checkVideoStatus(operationName: string): Promise<VideoGe
|
|
|
120
120
|
* Get the output directory path
|
|
121
121
|
*/
|
|
122
122
|
export declare function getOutputDir(): string;
|
|
123
|
+
/**
|
|
124
|
+
* Token count result
|
|
125
|
+
*/
|
|
126
|
+
export interface TokenCountResult {
|
|
127
|
+
totalTokens: number;
|
|
128
|
+
modelName: string;
|
|
129
|
+
}
|
|
130
|
+
/**
|
|
131
|
+
* Count tokens for content using specified model
|
|
132
|
+
*/
|
|
133
|
+
export declare function countTokens(content: string, model?: 'pro' | 'flash'): Promise<TokenCountResult>;
|
|
134
|
+
/**
|
|
135
|
+
* Deep Research interaction result
|
|
136
|
+
*/
|
|
137
|
+
export interface DeepResearchResult {
|
|
138
|
+
id: string;
|
|
139
|
+
status: 'pending' | 'processing' | 'completed' | 'failed';
|
|
140
|
+
outputs?: {
|
|
141
|
+
text?: string;
|
|
142
|
+
}[];
|
|
143
|
+
error?: string;
|
|
144
|
+
savedPath?: string;
|
|
145
|
+
}
|
|
146
|
+
/**
|
|
147
|
+
* Start a deep research task
|
|
148
|
+
*/
|
|
149
|
+
export declare function startDeepResearch(prompt: string): Promise<DeepResearchResult>;
|
|
150
|
+
/**
|
|
151
|
+
* Check deep research status
|
|
152
|
+
*/
|
|
153
|
+
export declare function checkDeepResearch(researchId: string): Promise<DeepResearchResult>;
|
|
154
|
+
/**
|
|
155
|
+
* Follow up on completed research
|
|
156
|
+
*/
|
|
157
|
+
export declare function followUpResearch(researchId: string, question: string): Promise<string>;
|
package/dist/gemini-client.js
CHANGED
|
@@ -397,3 +397,120 @@ export async function checkVideoStatus(operationName) {
|
|
|
397
397
|
export function getOutputDir() {
|
|
398
398
|
return outputDir;
|
|
399
399
|
}
|
|
400
|
+
/**
|
|
401
|
+
* Count tokens for content using specified model
|
|
402
|
+
*/
|
|
403
|
+
export async function countTokens(content, model = 'flash') {
|
|
404
|
+
const modelName = model === 'pro' ? proModelName : flashModelName;
|
|
405
|
+
const result = await genAI.models.countTokens({
|
|
406
|
+
model: modelName,
|
|
407
|
+
contents: content,
|
|
408
|
+
});
|
|
409
|
+
return {
|
|
410
|
+
totalTokens: result.totalTokens || 0,
|
|
411
|
+
modelName,
|
|
412
|
+
};
|
|
413
|
+
}
|
|
414
|
+
// Deep Research agent model
|
|
415
|
+
const DEEP_RESEARCH_AGENT = 'deep-research-pro-preview-12-2025';
|
|
416
|
+
/**
|
|
417
|
+
* Start a deep research task
|
|
418
|
+
*/
|
|
419
|
+
export async function startDeepResearch(prompt) {
|
|
420
|
+
try {
|
|
421
|
+
// The Interactions API is properly typed in @google/genai v1.34.0+
|
|
422
|
+
const interaction = await genAI.interactions.create({
|
|
423
|
+
input: prompt,
|
|
424
|
+
agent: DEEP_RESEARCH_AGENT,
|
|
425
|
+
background: true,
|
|
426
|
+
agent_config: {
|
|
427
|
+
type: 'deep-research',
|
|
428
|
+
thinking_summaries: 'auto',
|
|
429
|
+
},
|
|
430
|
+
});
|
|
431
|
+
return {
|
|
432
|
+
id: interaction.id || `research-${Date.now()}`,
|
|
433
|
+
status: 'pending',
|
|
434
|
+
};
|
|
435
|
+
}
|
|
436
|
+
catch (error) {
|
|
437
|
+
const message = error instanceof Error ? error.message : String(error);
|
|
438
|
+
throw new Error(`Deep research not available: ${message}`);
|
|
439
|
+
}
|
|
440
|
+
}
|
|
441
|
+
/**
|
|
442
|
+
* Check deep research status
|
|
443
|
+
*/
|
|
444
|
+
export async function checkDeepResearch(researchId) {
|
|
445
|
+
try {
|
|
446
|
+
const interaction = await genAI.interactions.get(researchId);
|
|
447
|
+
const status = interaction.status || 'unknown';
|
|
448
|
+
if (status === 'completed') {
|
|
449
|
+
// Save the FULL raw response to the output directory
|
|
450
|
+
const timestamp = new Date().toISOString().replace(/[:.]/g, '-');
|
|
451
|
+
const outputPath = path.join(getOutputDir(), `deep-research-${timestamp}.json`);
|
|
452
|
+
const fullResponse = {
|
|
453
|
+
id: researchId,
|
|
454
|
+
status: interaction.status,
|
|
455
|
+
created: interaction.created,
|
|
456
|
+
agent: interaction.agent,
|
|
457
|
+
model: interaction.model,
|
|
458
|
+
outputs: interaction.outputs,
|
|
459
|
+
rawInteraction: interaction,
|
|
460
|
+
};
|
|
461
|
+
fs.writeFileSync(outputPath, JSON.stringify(fullResponse, null, 2));
|
|
462
|
+
logger.info(`Full deep research response saved to: ${outputPath}`);
|
|
463
|
+
// Extract text for the summary (but full data is saved)
|
|
464
|
+
const textOutputs = (interaction.outputs || [])
|
|
465
|
+
.filter((output) => 'type' in output && output.type === 'text')
|
|
466
|
+
.map(output => ({ text: output.text }));
|
|
467
|
+
return {
|
|
468
|
+
id: researchId,
|
|
469
|
+
status: 'completed',
|
|
470
|
+
outputs: textOutputs,
|
|
471
|
+
savedPath: outputPath,
|
|
472
|
+
};
|
|
473
|
+
}
|
|
474
|
+
else if (status === 'failed' || status === 'cancelled') {
|
|
475
|
+
return {
|
|
476
|
+
id: researchId,
|
|
477
|
+
status: 'failed',
|
|
478
|
+
error: 'Research task failed or was cancelled',
|
|
479
|
+
};
|
|
480
|
+
}
|
|
481
|
+
return {
|
|
482
|
+
id: researchId,
|
|
483
|
+
status: 'processing',
|
|
484
|
+
};
|
|
485
|
+
}
|
|
486
|
+
catch (error) {
|
|
487
|
+
const message = error instanceof Error ? error.message : String(error);
|
|
488
|
+
throw new Error(`Failed to check research status: ${message}`);
|
|
489
|
+
}
|
|
490
|
+
}
|
|
491
|
+
/**
|
|
492
|
+
* Follow up on completed research
|
|
493
|
+
*/
|
|
494
|
+
export async function followUpResearch(researchId, question) {
|
|
495
|
+
try {
|
|
496
|
+
const interaction = await genAI.interactions.create({
|
|
497
|
+
input: question,
|
|
498
|
+
model: proModelName,
|
|
499
|
+
previous_interaction_id: researchId,
|
|
500
|
+
});
|
|
501
|
+
// Extract text from TextContent outputs
|
|
502
|
+
const outputs = interaction.outputs || [];
|
|
503
|
+
const textOutputs = outputs
|
|
504
|
+
.filter((output) => 'type' in output && output.type === 'text')
|
|
505
|
+
.map(output => output.text)
|
|
506
|
+
.filter((text) => !!text);
|
|
507
|
+
if (textOutputs.length > 0) {
|
|
508
|
+
return textOutputs[textOutputs.length - 1];
|
|
509
|
+
}
|
|
510
|
+
return 'No text response received';
|
|
511
|
+
}
|
|
512
|
+
catch (error) {
|
|
513
|
+
const message = error instanceof Error ? error.message : String(error);
|
|
514
|
+
throw new Error(`Research follow-up failed: ${message}`);
|
|
515
|
+
}
|
|
516
|
+
}
|
package/dist/index.js
CHANGED
|
@@ -4,7 +4,7 @@
|
|
|
4
4
|
* Uses the Gemini Deep Research Agent for complex research tasks.
|
|
5
5
|
* The agent autonomously plans, searches, reads, and synthesizes research.
|
|
6
6
|
*/
|
|
7
|
-
import { McpServer } from
|
|
7
|
+
import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js';
|
|
8
8
|
/**
|
|
9
9
|
* Register deep research tools with the MCP server
|
|
10
10
|
*/
|
|
@@ -4,20 +4,18 @@
|
|
|
4
4
|
* Uses the Gemini Deep Research Agent for complex research tasks.
|
|
5
5
|
* The agent autonomously plans, searches, reads, and synthesizes research.
|
|
6
6
|
*/
|
|
7
|
-
import { z } from
|
|
8
|
-
import { logger } from
|
|
9
|
-
import {
|
|
7
|
+
import { z } from 'zod';
|
|
8
|
+
import { logger } from '../utils/logger.js';
|
|
9
|
+
import { startDeepResearch, checkDeepResearch, followUpResearch, } from '../gemini-client.js';
|
|
10
10
|
// Store active research operations for polling
|
|
11
11
|
const activeResearchOperations = new Map();
|
|
12
|
-
// Deep Research agent model
|
|
13
|
-
const DEEP_RESEARCH_AGENT = "deep-research-pro-preview-12-2025";
|
|
14
12
|
/**
|
|
15
13
|
* Register deep research tools with the MCP server
|
|
16
14
|
*/
|
|
17
15
|
export function registerDeepResearchTool(server) {
|
|
18
16
|
// Start a deep research task
|
|
19
|
-
server.tool(
|
|
20
|
-
query: z.string().describe(
|
|
17
|
+
server.tool('gemini-deep-research', {
|
|
18
|
+
query: z.string().describe('The research question or topic to investigate'),
|
|
21
19
|
format: z
|
|
22
20
|
.string()
|
|
23
21
|
.optional()
|
|
@@ -30,33 +28,22 @@ export function registerDeepResearchTool(server) {
|
|
|
30
28
|
if (format) {
|
|
31
29
|
researchPrompt = `${query}\n\nFormat the output as: ${format}`;
|
|
32
30
|
}
|
|
33
|
-
|
|
34
|
-
// The Interactions API is accessed via genAI.interactions
|
|
35
|
-
const interaction = await genAI.interactions.create({
|
|
36
|
-
input: researchPrompt,
|
|
37
|
-
agent: DEEP_RESEARCH_AGENT,
|
|
38
|
-
background: true,
|
|
39
|
-
agentConfig: {
|
|
40
|
-
type: "deep-research",
|
|
41
|
-
thinkingSummaries: "auto"
|
|
42
|
-
}
|
|
43
|
-
});
|
|
44
|
-
const interactionId = interaction.id || `research-${Date.now()}`;
|
|
31
|
+
const result = await startDeepResearch(researchPrompt);
|
|
45
32
|
// Store for later polling
|
|
46
|
-
activeResearchOperations.set(
|
|
47
|
-
interactionId,
|
|
33
|
+
activeResearchOperations.set(result.id, {
|
|
48
34
|
startedAt: new Date(),
|
|
49
|
-
prompt: query
|
|
35
|
+
prompt: query,
|
|
50
36
|
});
|
|
51
|
-
logger.info(`Deep research started: ${
|
|
37
|
+
logger.info(`Deep research started: ${result.id}`);
|
|
52
38
|
return {
|
|
53
|
-
content: [
|
|
54
|
-
|
|
39
|
+
content: [
|
|
40
|
+
{
|
|
41
|
+
type: 'text',
|
|
55
42
|
text: `**Deep Research Started**
|
|
56
43
|
|
|
57
44
|
| Field | Value |
|
|
58
45
|
|-------|-------|
|
|
59
|
-
| **Research ID** | \`${
|
|
46
|
+
| **Research ID** | \`${result.id}\` |
|
|
60
47
|
| **Query** | ${query.substring(0, 100)}${query.length > 100 ? '...' : ''} |
|
|
61
48
|
| **Status** | In Progress |
|
|
62
49
|
| **Started** | ${new Date().toISOString()} |
|
|
@@ -69,18 +56,20 @@ export function registerDeepResearchTool(server) {
|
|
|
69
56
|
**To check progress:**
|
|
70
57
|
Use \`gemini-check-research\` with the Research ID above.
|
|
71
58
|
|
|
72
|
-
**Note:** Deep research tasks run in the background. You can continue working while waiting
|
|
73
|
-
}
|
|
59
|
+
**Note:** Deep research tasks run in the background. You can continue working while waiting.`,
|
|
60
|
+
},
|
|
61
|
+
],
|
|
74
62
|
};
|
|
75
63
|
}
|
|
76
64
|
catch (error) {
|
|
77
65
|
const errorMessage = error instanceof Error ? error.message : String(error);
|
|
78
66
|
logger.error(`Error starting deep research: ${errorMessage}`);
|
|
79
67
|
// Check if it's an API availability issue
|
|
80
|
-
if (errorMessage.includes(
|
|
68
|
+
if (errorMessage.includes('not available') || errorMessage.includes('interactions')) {
|
|
81
69
|
return {
|
|
82
|
-
content: [
|
|
83
|
-
|
|
70
|
+
content: [
|
|
71
|
+
{
|
|
72
|
+
type: 'text',
|
|
84
73
|
text: `**Deep Research Not Available**
|
|
85
74
|
|
|
86
75
|
The Interactions API required for Deep Research may not be available yet in your SDK version or API access.
|
|
@@ -90,44 +79,46 @@ The Interactions API required for Deep Research may not be available yet in your
|
|
|
90
79
|
**Alternatives:**
|
|
91
80
|
- Use \`gemini-search\` for real-time web search
|
|
92
81
|
- Use \`gemini-query\` with a detailed research prompt
|
|
93
|
-
- Wait for Interactions API to become available in your region
|
|
94
|
-
}
|
|
95
|
-
|
|
82
|
+
- Wait for Interactions API to become available in your region`,
|
|
83
|
+
},
|
|
84
|
+
],
|
|
85
|
+
isError: true,
|
|
96
86
|
};
|
|
97
87
|
}
|
|
98
88
|
return {
|
|
99
|
-
content: [{ type:
|
|
100
|
-
isError: true
|
|
89
|
+
content: [{ type: 'text', text: `Error starting deep research: ${errorMessage}` }],
|
|
90
|
+
isError: true,
|
|
101
91
|
};
|
|
102
92
|
}
|
|
103
93
|
});
|
|
104
94
|
// Check research status
|
|
105
|
-
server.tool(
|
|
106
|
-
researchId: z.string().describe(
|
|
95
|
+
server.tool('gemini-check-research', {
|
|
96
|
+
researchId: z.string().describe('The research ID returned from gemini-deep-research'),
|
|
107
97
|
}, async ({ researchId }) => {
|
|
108
98
|
logger.info(`Checking research status: ${researchId}`);
|
|
109
99
|
try {
|
|
110
100
|
// Get stored operation info
|
|
111
101
|
const operationInfo = activeResearchOperations.get(researchId);
|
|
112
|
-
|
|
113
|
-
const interaction = await genAI.interactions.get(researchId);
|
|
114
|
-
const status = interaction.status || "unknown";
|
|
115
|
-
const elapsedMs = operationInfo
|
|
116
|
-
? Date.now() - operationInfo.startedAt.getTime()
|
|
117
|
-
: 0;
|
|
102
|
+
const elapsedMs = operationInfo ? Date.now() - operationInfo.startedAt.getTime() : 0;
|
|
118
103
|
const elapsedMinutes = Math.floor(elapsedMs / 60000);
|
|
119
104
|
const elapsedSeconds = Math.floor((elapsedMs % 60000) / 1000);
|
|
120
|
-
|
|
105
|
+
const result = await checkDeepResearch(researchId);
|
|
106
|
+
if (result.status === 'completed') {
|
|
121
107
|
// Research is done - extract the result
|
|
122
108
|
activeResearchOperations.delete(researchId);
|
|
123
|
-
const outputs =
|
|
124
|
-
const
|
|
125
|
-
? outputs[outputs.length - 1].text ||
|
|
126
|
-
:
|
|
109
|
+
const outputs = result.outputs || [];
|
|
110
|
+
const resultText = outputs.length > 0
|
|
111
|
+
? outputs[outputs.length - 1].text || 'No text output'
|
|
112
|
+
: 'Research completed but no output found';
|
|
127
113
|
logger.info(`Research completed: ${researchId}`);
|
|
114
|
+
// Build the saved path info if available
|
|
115
|
+
const savedPathInfo = result.savedPath
|
|
116
|
+
? `| **Full Response** | \`${result.savedPath}\` |`
|
|
117
|
+
: '';
|
|
128
118
|
return {
|
|
129
|
-
content: [
|
|
130
|
-
|
|
119
|
+
content: [
|
|
120
|
+
{
|
|
121
|
+
type: 'text',
|
|
131
122
|
text: `**Deep Research Complete**
|
|
132
123
|
|
|
133
124
|
| Field | Value |
|
|
@@ -135,55 +126,62 @@ The Interactions API required for Deep Research may not be available yet in your
|
|
|
135
126
|
| **Research ID** | \`${researchId}\` |
|
|
136
127
|
| **Status** | ✅ Completed |
|
|
137
128
|
| **Duration** | ${elapsedMinutes}m ${elapsedSeconds}s |
|
|
129
|
+
${savedPathInfo}
|
|
130
|
+
|
|
131
|
+
> **Note:** The full response (including citations, images, and all metadata) has been saved to the file above.
|
|
138
132
|
|
|
139
133
|
---
|
|
140
134
|
|
|
141
135
|
## Research Results
|
|
142
136
|
|
|
143
|
-
${
|
|
144
|
-
}
|
|
137
|
+
${resultText}`,
|
|
138
|
+
},
|
|
139
|
+
],
|
|
145
140
|
};
|
|
146
141
|
}
|
|
147
|
-
else if (status ===
|
|
142
|
+
else if (result.status === 'failed') {
|
|
148
143
|
activeResearchOperations.delete(researchId);
|
|
149
|
-
|
|
150
|
-
logger.error(`Research failed: ${researchId} - ${errorInfo}`);
|
|
144
|
+
logger.error(`Research failed: ${researchId} - ${result.error}`);
|
|
151
145
|
return {
|
|
152
|
-
content: [
|
|
153
|
-
|
|
146
|
+
content: [
|
|
147
|
+
{
|
|
148
|
+
type: 'text',
|
|
154
149
|
text: `**Deep Research Failed**
|
|
155
150
|
|
|
156
151
|
| Field | Value |
|
|
157
152
|
|-------|-------|
|
|
158
153
|
| **Research ID** | \`${researchId}\` |
|
|
159
154
|
| **Status** | ❌ Failed |
|
|
160
|
-
| **Error** | ${
|
|
155
|
+
| **Error** | ${result.error} |
|
|
161
156
|
|
|
162
157
|
The research task encountered an error. You can try:
|
|
163
158
|
- Starting a new research task with a different query
|
|
164
|
-
- Using \`gemini-search\` for simpler web searches
|
|
165
|
-
}
|
|
166
|
-
|
|
159
|
+
- Using \`gemini-search\` for simpler web searches`,
|
|
160
|
+
},
|
|
161
|
+
],
|
|
162
|
+
isError: true,
|
|
167
163
|
};
|
|
168
164
|
}
|
|
169
165
|
else {
|
|
170
166
|
// Still in progress
|
|
171
167
|
return {
|
|
172
|
-
content: [
|
|
173
|
-
|
|
168
|
+
content: [
|
|
169
|
+
{
|
|
170
|
+
type: 'text',
|
|
174
171
|
text: `**Deep Research In Progress**
|
|
175
172
|
|
|
176
173
|
| Field | Value |
|
|
177
174
|
|-------|-------|
|
|
178
175
|
| **Research ID** | \`${researchId}\` |
|
|
179
|
-
| **Status** | ⏳ ${status} |
|
|
176
|
+
| **Status** | ⏳ ${result.status} |
|
|
180
177
|
| **Elapsed** | ${elapsedMinutes}m ${elapsedSeconds}s |
|
|
181
178
|
| **Query** | ${operationInfo?.prompt.substring(0, 50) || 'Unknown'}... |
|
|
182
179
|
|
|
183
180
|
The agent is still working. Deep research typically takes 2-10 minutes.
|
|
184
181
|
|
|
185
|
-
Check again in 30-60 seconds using \`gemini-check-research
|
|
186
|
-
}
|
|
182
|
+
Check again in 30-60 seconds using \`gemini-check-research\`.`,
|
|
183
|
+
},
|
|
184
|
+
],
|
|
187
185
|
};
|
|
188
186
|
}
|
|
189
187
|
}
|
|
@@ -191,45 +189,39 @@ Check again in 30-60 seconds using \`gemini-check-research\`.`
|
|
|
191
189
|
const errorMessage = error instanceof Error ? error.message : String(error);
|
|
192
190
|
logger.error(`Error checking research status: ${errorMessage}`);
|
|
193
191
|
return {
|
|
194
|
-
content: [{ type:
|
|
195
|
-
isError: true
|
|
192
|
+
content: [{ type: 'text', text: `Error checking research status: ${errorMessage}` }],
|
|
193
|
+
isError: true,
|
|
196
194
|
};
|
|
197
195
|
}
|
|
198
196
|
});
|
|
199
197
|
// Follow-up on completed research
|
|
200
|
-
server.tool(
|
|
201
|
-
researchId: z.string().describe(
|
|
202
|
-
question: z.string().describe(
|
|
198
|
+
server.tool('gemini-research-followup', {
|
|
199
|
+
researchId: z.string().describe('The research ID from a completed research task'),
|
|
200
|
+
question: z.string().describe('Follow-up question about the research results'),
|
|
203
201
|
}, async ({ researchId, question }) => {
|
|
204
202
|
logger.info(`Research follow-up on ${researchId}: ${question.substring(0, 50)}...`);
|
|
205
203
|
try {
|
|
206
|
-
const
|
|
207
|
-
input: question,
|
|
208
|
-
model: "gemini-3-pro-preview",
|
|
209
|
-
previousInteractionId: researchId
|
|
210
|
-
});
|
|
211
|
-
const outputs = interaction.outputs || [];
|
|
212
|
-
const result = outputs.length > 0
|
|
213
|
-
? outputs[outputs.length - 1].text || "No response"
|
|
214
|
-
: "No response received";
|
|
204
|
+
const result = await followUpResearch(researchId, question);
|
|
215
205
|
return {
|
|
216
|
-
content: [
|
|
217
|
-
|
|
206
|
+
content: [
|
|
207
|
+
{
|
|
208
|
+
type: 'text',
|
|
218
209
|
text: `**Research Follow-up**
|
|
219
210
|
|
|
220
211
|
**Question:** ${question}
|
|
221
212
|
|
|
222
213
|
**Answer:**
|
|
223
|
-
${result}
|
|
224
|
-
}
|
|
214
|
+
${result}`,
|
|
215
|
+
},
|
|
216
|
+
],
|
|
225
217
|
};
|
|
226
218
|
}
|
|
227
219
|
catch (error) {
|
|
228
220
|
const errorMessage = error instanceof Error ? error.message : String(error);
|
|
229
221
|
logger.error(`Error with research follow-up: ${errorMessage}`);
|
|
230
222
|
return {
|
|
231
|
-
content: [{ type:
|
|
232
|
-
isError: true
|
|
223
|
+
content: [{ type: 'text', text: `Error with follow-up: ${errorMessage}` }],
|
|
224
|
+
isError: true,
|
|
233
225
|
};
|
|
234
226
|
}
|
|
235
227
|
});
|
|
@@ -3,7 +3,7 @@
|
|
|
3
3
|
*
|
|
4
4
|
* Helps users estimate costs and manage context windows.
|
|
5
5
|
*/
|
|
6
|
-
import { McpServer } from
|
|
6
|
+
import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js';
|
|
7
7
|
/**
|
|
8
8
|
* Register token counting tool with the MCP server
|
|
9
9
|
*/
|
|
@@ -3,37 +3,31 @@
|
|
|
3
3
|
*
|
|
4
4
|
* Helps users estimate costs and manage context windows.
|
|
5
5
|
*/
|
|
6
|
-
import { z } from
|
|
7
|
-
import { logger } from
|
|
8
|
-
import {
|
|
6
|
+
import { z } from 'zod';
|
|
7
|
+
import { logger } from '../utils/logger.js';
|
|
8
|
+
import { countTokens } from '../gemini-client.js';
|
|
9
9
|
/**
|
|
10
10
|
* Register token counting tool with the MCP server
|
|
11
11
|
*/
|
|
12
12
|
export function registerTokenCountTool(server) {
|
|
13
|
-
server.tool(
|
|
14
|
-
content: z.string().describe(
|
|
13
|
+
server.tool('gemini-count-tokens', {
|
|
14
|
+
content: z.string().describe('The text content to count tokens for'),
|
|
15
15
|
model: z
|
|
16
|
-
.enum([
|
|
17
|
-
.default(
|
|
18
|
-
.describe(
|
|
19
|
-
}, async ({ content, model =
|
|
16
|
+
.enum(['pro', 'flash'])
|
|
17
|
+
.default('flash')
|
|
18
|
+
.describe('Which model to use for counting (affects tokenization)'),
|
|
19
|
+
}, async ({ content, model = 'flash' }) => {
|
|
20
20
|
logger.info(`Counting tokens for ${content.length} characters using ${model} model`);
|
|
21
21
|
try {
|
|
22
|
-
const
|
|
23
|
-
|
|
24
|
-
: (process.env.GEMINI_FLASH_MODEL || "gemini-3-flash-preview");
|
|
25
|
-
const result = await genAI.models.countTokens({
|
|
26
|
-
model: modelName,
|
|
27
|
-
contents: content
|
|
28
|
-
});
|
|
29
|
-
const totalTokens = result.totalTokens || 0;
|
|
22
|
+
const result = await countTokens(content, model);
|
|
23
|
+
const totalTokens = result.totalTokens;
|
|
30
24
|
// Estimate costs (approximate, based on typical pricing)
|
|
31
25
|
// Gemini 3 Pro: ~$1.25 per 1M input tokens
|
|
32
26
|
// Gemini 3 Flash: ~$0.075 per 1M input tokens
|
|
33
|
-
const costPer1M = model ===
|
|
27
|
+
const costPer1M = model === 'pro' ? 1.25 : 0.075;
|
|
34
28
|
const estimatedCost = (totalTokens / 1_000_000) * costPer1M;
|
|
35
29
|
// Context window info
|
|
36
|
-
const contextWindow = model ===
|
|
30
|
+
const contextWindow = model === 'pro' ? 1_000_000 : 1_000_000;
|
|
37
31
|
const percentUsed = (totalTokens / contextWindow) * 100;
|
|
38
32
|
const response = `**Token Count Results**
|
|
39
33
|
|
|
@@ -41,7 +35,7 @@ export function registerTokenCountTool(server) {
|
|
|
41
35
|
|--------|-------|
|
|
42
36
|
| **Total Tokens** | ${totalTokens.toLocaleString()} |
|
|
43
37
|
| **Characters** | ${content.length.toLocaleString()} |
|
|
44
|
-
| **Model** | ${modelName} |
|
|
38
|
+
| **Model** | ${result.modelName} |
|
|
45
39
|
|
|
46
40
|
**Context Window Usage:**
|
|
47
41
|
- Context window: ${contextWindow.toLocaleString()} tokens
|
|
@@ -55,15 +49,15 @@ export function registerTokenCountTool(server) {
|
|
|
55
49
|
*Note: Actual costs may vary. Check [Google AI pricing](https://ai.google.dev/pricing) for current rates.*`;
|
|
56
50
|
logger.info(`Token count: ${totalTokens}`);
|
|
57
51
|
return {
|
|
58
|
-
content: [{ type:
|
|
52
|
+
content: [{ type: 'text', text: response }],
|
|
59
53
|
};
|
|
60
54
|
}
|
|
61
55
|
catch (error) {
|
|
62
56
|
const errorMessage = error instanceof Error ? error.message : String(error);
|
|
63
57
|
logger.error(`Error counting tokens: ${errorMessage}`);
|
|
64
58
|
return {
|
|
65
|
-
content: [{ type:
|
|
66
|
-
isError: true
|
|
59
|
+
content: [{ type: 'text', text: `Error counting tokens: ${errorMessage}` }],
|
|
60
|
+
isError: true,
|
|
67
61
|
};
|
|
68
62
|
}
|
|
69
63
|
});
|