Skip to content

Commit 3327b62

Browse files
ssdeanxssdeanx
authored
Merge pull request #128 from ssdeanx/develop
feat: update Google AI model options and refactor imports across agents
2 parents 0f5f08b + fb785fa commit 3327b62

15 files changed

Lines changed: 90 additions & 165 deletions

‎src/mastra/a2a/a2aCoordinatorAgent.ts‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
import type { GoogleGenerativeAIProviderOptions } from '@ai-sdk/google'
1+
import type { GoogleLanguageModelOptions } from '@ai-sdk/google'
22
import { Agent } from '@mastra/core/agent'
33
import {
44
createAnswerRelevancyScorer,
@@ -16,7 +16,6 @@ import { editorAgent } from '../agents/editorAgent'
1616
import { knowledgeIndexingAgent } from '../agents/knowledgeIndexingAgent'
1717
import { projectManagementAgent } from '../agents/projectManagementAgent'
1818
import { researchAgent } from '../agents/researchAgent'
19-
import { googleAI, googleAIFlashLite } from '../config/google'
2019
import { pgMemory } from '../config/pg-storage'
2120
import { repoIngestionWorkflow } from '../workflows/repo-ingestion-workflow'
2221
import { researchSynthesisWorkflow } from '../workflows/research-synthesis-workflow'
@@ -96,15 +95,14 @@ Use knowledgeIndexingAgent to provide semantic context for complex queries.
9695
thinkingConfig: {
9796
thinkingLevel: 'high',
9897
includeThoughts: true,
99-
thinkingBudget: -1,
10098
},
10199
mediaResolution: 'MEDIA_RESOLUTION_MEDIUM',
102100
responseModalities: ['TEXT', 'IMAGE'],
103-
} satisfies GoogleGenerativeAIProviderOptions,
101+
} satisfies GoogleLanguageModelOptions,
104102
},
105103
}
106104
},
107-
model: googleAI,
105+
model: 'google/gemini-3.1-flash-preview',
108106
memory: pgMemory,
109107
options: {},
110108
agents: {
@@ -128,11 +126,13 @@ Use knowledgeIndexingAgent to provide semantic context for complex queries.
128126
tools: {},
129127
scorers: {
130128
relevancy: {
131-
scorer: createAnswerRelevancyScorer({ model: googleAIFlashLite }),
129+
scorer: createAnswerRelevancyScorer({ model:
130+
'google/gemini-3.1-flash-lite-preview'
131+
}),
132132
sampling: { type: 'ratio', rate: 0.4 },
133133
},
134134
safety: {
135-
scorer: createToxicityScorer({ model: googleAIFlashLite }),
135+
scorer: createToxicityScorer({ model: 'google/gemini-3.1-flash-lite-preview' }),
136136
sampling: { type: 'ratio', rate: 0.3 },
137137
},
138138
},

‎src/mastra/a2a/codingA2ACoordinator.ts‎

Lines changed: 8 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -4,32 +4,17 @@ import {
44
createToxicityScorer,
55
} from '@mastra/evals/scorers/prebuilt'
66

7-
import { googleAIFlashLite } from '../config/google'
87
import { log } from '../config/logger'
98
import { pgMemory } from '../config/pg-storage'
109

11-
import type { GoogleGenerativeAIProviderOptions } from '@ai-sdk/google'
10+
import { google, type GoogleLanguageModelOptions } from '@ai-sdk/google'
1211
import { InternalSpans } from '@mastra/core/observability'
1312
import {
1413
codeArchitectAgent,
1514
codeReviewerAgent,
1615
refactoringAgent,
1716
testEngineerAgent,
1817
} from '../agents/codingAgents'
19-
import {
20-
checkFileExists,
21-
createDirectory,
22-
createSandbox,
23-
deleteFile,
24-
getFileInfo,
25-
getFileSize,
26-
listFiles,
27-
runCode,
28-
runCommand,
29-
watchDirectory,
30-
writeFile,
31-
writeFiles,
32-
} from '../tools/e2b'
3318
import { automatedReportingWorkflow } from '../workflows/automated-reporting-workflow'
3419
import { dataAnalysisWorkflow } from '../workflows/data-analysis-workflow'
3520
import { financialReportWorkflow } from '../workflows/financial-report-workflow'
@@ -197,11 +182,11 @@ When a user's request requires prolonged, structured work across multiple subtas
197182
},
198183
mediaResolution: 'MEDIA_RESOLUTION_MEDIUM',
199184
responseModalities: ['TEXT', 'IMAGE'],
200-
} satisfies GoogleGenerativeAIProviderOptions,
185+
} satisfies GoogleLanguageModelOptions,
201186
},
202187
}
203188
},
204-
model: googleAIFlashLite,
189+
model: 'google/gemini-3.1-flash-preview',
205190
memory: pgMemory,
206191

207192
agents: {
@@ -222,18 +207,9 @@ When a user's request requires prolonged, structured work across multiple subtas
222207
automatedReportingWorkflow,
223208
},
224209
tools: {
225-
createSandbox,
226-
writeFile,
227-
writeFiles,
228-
listFiles,
229-
deleteFile,
230-
createDirectory,
231-
getFileInfo,
232-
checkFileExists,
233-
getFileSize,
234-
watchDirectory,
235-
runCommand,
236-
runCode,
210+
google_search: google.tools.googleSearch({}),
211+
url_context: google.tools.urlContext({}),
212+
code_execution: google.tools.codeExecution({}),
237213
},
238214
maxRetries: 5,
239215
options: {
@@ -244,11 +220,11 @@ When a user's request requires prolonged, structured work across multiple subtas
244220

245221
scorers: {
246222
relevancy: {
247-
scorer: createAnswerRelevancyScorer({ model: googleAIFlashLite }),
223+
scorer: createAnswerRelevancyScorer({ model: 'google/gemini-3.1-flash-lite-preview' }),
248224
sampling: { type: 'ratio', rate: 0.4 },
249225
},
250226
safety: {
251-
scorer: createToxicityScorer({ model: googleAIFlashLite }),
227+
scorer: createToxicityScorer({ model: 'google/gemini-3.1-flash-lite-preview' }),
252228
sampling: { type: 'ratio', rate: 0.3 },
253229
},
254230
},

‎src/mastra/agents/academicResearchAgent.ts‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
import type { GoogleGenerativeAIProviderOptions } from '@ai-sdk/google'
1+
import type { GoogleLanguageModelOptions } from '@ai-sdk/google'
22
import { Agent } from '@mastra/core/agent'
33
import type { RequestContext } from '@mastra/core/request-context'
44
import type { AgentRequestContext } from './request-context'
@@ -106,7 +106,7 @@ ${
106106
includeThoughts: true,
107107
thinkingLevel: 'high',
108108
},
109-
} satisfies GoogleGenerativeAIProviderOptions,
109+
} satisfies GoogleLanguageModelOptions,
110110
},
111111
}
112112
},

‎src/mastra/agents/acpAgent.ts‎

Lines changed: 3 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -2,14 +2,6 @@ import { Agent } from '@mastra/core/agent'
22
import { pgMemory, pgQueryTool } from '../config'
33
import { arxivTool } from '../tools/arxiv.tool'
44
import { csvToJsonTool } from '../tools/csv-to-json.tool'
5-
import {
6-
createDataDirTool,
7-
getDataFileInfoTool,
8-
listDataDirTool,
9-
moveDataFileTool,
10-
searchDataFilesTool,
11-
writeDataFileTool,
12-
} from '../tools/data-file-manager'
135
import {
146
csvToExcalidrawTool,
157
readCSVDataTool,
@@ -25,7 +17,7 @@ import {
2517
} from '../tools/github'
2618
import { jsonToCsvTool } from '../tools/json-to-csv.tool'
2719
import { pdfToMarkdownTool } from '../tools/pdf-data-conversion.tool'
28-
import type { GoogleGenerativeAIProviderOptions } from '@ai-sdk/google'
20+
import type { GoogleLanguageModelOptions } from '@ai-sdk/google'
2921
import type { RequestContext } from '@mastra/core/request-context'
3022
import { PGVECTOR_PROMPT } from '@mastra/pg'
3123
import { TokenLimiterProcessor } from '@mastra/core/processors'
@@ -74,10 +66,10 @@ User: ${userId} | Role: ${roleConstraint}
7466
google: {
7567
thinkingConfig: {
7668
includeThoughts: true,
77-
thinkingBudget: -1,
69+
thinkingLevel: 'medium',
7870
},
7971
responseModalities: ['TEXT'],
80-
} satisfies GoogleGenerativeAIProviderOptions,
72+
} satisfies GoogleLanguageModelOptions,
8173
},
8274
}
8375
},
@@ -96,12 +88,6 @@ User: ${userId} | Role: ${roleConstraint}
9688
csvToExcalidrawTool,
9789
readCSVDataTool,
9890
// convertDataFormatTool,
99-
writeDataFileTool,
100-
listDataDirTool,
101-
searchDataFilesTool,
102-
moveDataFileTool,
103-
getDataFileInfoTool,
104-
createDataDirTool,
10591
searchCode,
10692
getFileContent,
10793
getRepositoryInfo,

‎src/mastra/agents/codingAgents.ts‎

Lines changed: 9 additions & 49 deletions
Original file line numberDiff line numberDiff line change
@@ -1,39 +1,20 @@
11
import { Agent } from '@mastra/core/agent'
22

3-
import type { GoogleGenerativeAIProviderOptions } from '@ai-sdk/google'
3+
import type { GoogleGenerativeAIProviderOptions, GoogleLanguageModelOptions } from '@ai-sdk/google'
44
import {
55
TokenLimiterProcessor
66
} from '@mastra/core/processors'
77
import {
88
createAnswerRelevancyScorer,
99
createToxicityScorer,
1010
} from '@mastra/evals/scorers/prebuilt'
11-
import {
12-
google3,
13-
googleAI3,
14-
googleAIFlashLite,
15-
} from '../config/google'
1611
import { log } from '../config/logger'
1712

1813
import { InternalSpans } from '@mastra/core/observability'
1914
import { pgMemory } from '../config/pg-storage'
2015
import { codeAnalysisTool } from '../tools/code-analysis.tool'
2116
import { codeSearchTool } from '../tools/code-search.tool'
2217
import { diffReviewTool } from '../tools/diff-review.tool'
23-
import {
24-
checkFileExists,
25-
createDirectory,
26-
createSandbox,
27-
deleteFile,
28-
getFileInfo,
29-
getFileSize,
30-
listFiles,
31-
runCode,
32-
runCommand,
33-
watchDirectory,
34-
writeFile,
35-
writeFiles,
36-
} from '../tools/e2b'
3718
import { findReferencesTool } from '../tools/find-references.tool'
3819
import { findSymbolTool } from '../tools/find-symbol.tool'
3920
import {
@@ -87,34 +68,13 @@ const codeReviewerTools = {
8768
const testEngineerTools = {
8869
codeAnalysisTool,
8970
testGeneratorTool,
90-
createSandbox,
91-
runCommand,
92-
runCode,
93-
writeFile,
94-
writeFiles,
95-
deleteFile,
96-
listFiles,
97-
getFileInfo,
98-
getFileSize,
99-
checkFileExists,
71+
10072
}
10173

10274
const refactoringTools = {
10375
codeAnalysisTool,
10476
diffReviewTool,
10577
multiStringEditTool,
106-
createSandbox,
107-
runCode,
108-
runCommand,
109-
writeFile,
110-
writeFiles,
111-
deleteFile,
112-
listFiles,
113-
getFileInfo,
114-
getFileSize,
115-
checkFileExists,
116-
createDirectory,
117-
watchDirectory,
11878
searchCode,
11979
getFileContent,
12080
getRepositoryInfo,
@@ -178,19 +138,19 @@ Always consider maintainability, scalability, and testability in your recommenda
178138
responseModalities: ['TEXT'],
179139
cachedContent:
180140
'Repo Name, Description, Key Modules, Recent Commits',
181-
} satisfies GoogleGenerativeAIProviderOptions,
141+
} satisfies GoogleLanguageModelOptions,
182142
},
183143
}
184144
},
185145
model: ({ requestContext }) => {
186146
const userTier = getUserTierFromContext(requestContext)
187-
return userTier === 'enterprise' ? googleAI3 : google3
147+
return userTier === 'enterprise' ? 'google/gemini-3.1-flash-preview' : 'google/gemini-3.1-flash-lite-preview'
188148
},
189149
tools: codeArchitectTools,
190150
memory: pgMemory,
191151
scorers: {
192152
relevancy: {
193-
scorer: createAnswerRelevancyScorer({ model: googleAIFlashLite }),
153+
scorer: createAnswerRelevancyScorer({ model: 'google/gemini-3.1-flash-lite-preview' }),
194154
sampling: { type: 'ratio', rate: 0.5 },
195155
},
196156
},
@@ -304,11 +264,11 @@ Be constructive and educational in feedback.`,
304264
},
305265
scorers: {
306266
relevancy: {
307-
scorer: createAnswerRelevancyScorer({ model: googleAIFlashLite }),
267+
scorer: createAnswerRelevancyScorer({ model: 'google/gemini-3.1-flash-lite-preview' }),
308268
sampling: { type: 'ratio', rate: 0.5 },
309269
},
310270
safety: {
311-
scorer: createToxicityScorer({ model: googleAIFlashLite }),
271+
scorer: createToxicityScorer({ model: 'google/gemini-3.1-flash-lite-preview' }),
312272
sampling: { type: 'ratio', rate: 0.3 },
313273
},
314274
},
@@ -427,7 +387,7 @@ Always use Vitest syntax: describe, it, expect, vi.mock, vi.fn.`,
427387
},
428388
scorers: {
429389
relevancy: {
430-
scorer: createAnswerRelevancyScorer({ model: googleAIFlashLite }),
390+
scorer: createAnswerRelevancyScorer({ model: 'google/gemini-3.1-flash-lite-preview' }),
431391
sampling: { type: 'ratio', rate: 0.5 },
432392
},
433393
},
@@ -554,7 +514,7 @@ For each refactoring:
554514
},
555515
scorers: {
556516
relevancy: {
557-
scorer: createAnswerRelevancyScorer({ model: googleAIFlashLite }),
517+
scorer: createAnswerRelevancyScorer({ model: 'google/gemini-3.1-flash-lite-preview' }),
558518
sampling: { type: 'ratio', rate: 0.5 },
559519
},
560520
},

‎src/mastra/agents/contentStrategistAgent.ts‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2,10 +2,9 @@ import { Agent } from '@mastra/core/agent'
22
import { pgMemory } from '../config/pg-storage'
33
import {
44
scrapingSchedulerTool,
5-
webScraperTool,
65
} from '../tools/web-scraper-tool'
76

8-
import type { GoogleGenerativeAIProviderOptions } from '@ai-sdk/google'
7+
import { google, type GoogleGenerativeAIProviderOptions } from '@ai-sdk/google'
98
import { InternalSpans } from '@mastra/core/observability'
109
import { TokenLimiterProcessor } from '@mastra/core/processors'
1110
import {
@@ -19,6 +18,7 @@ import {
1918
createToneScorer,
2019
} from '../evals/scorers/prebuilt'
2120
import { chartSupervisorTool } from '../tools/financial-chart-tools'
21+
import { fetchTool } from '../tools/fetch.tool'
2222

2323
const STAGGERED_OUTPUT_CONTEXT_KEY = 'staggeredOutput' as const
2424
const SECTION_COUNT_CONTEXT_KEY = 'sectionCount' as const
@@ -79,9 +79,11 @@ function getBackupDataToolsFromContext(requestContext: {
7979
}
8080

8181
const contentStrategistTools = {
82-
webScraperTool,
82+
fetchTool,
8383
chartSupervisorTool,
8484
scrapingSchedulerTool,
85+
google_search: google.tools.googleSearch({}),
86+
url_context: google.tools.urlContext({}),
8587
}
8688

8789
export const contentStrategistAgent = new Agent({

0 commit comments

Comments
 (0)