feat(ai): Enhance Google provider telemetry and AI object response handling

This commit introduces two key improvements:

1.  **Google Provider Telemetry:**
    - Updated  to include token usage data (, ) in the responses from  and .
    - This aligns the Google provider with others for consistent AI usage telemetry.

2.  **Robust AI Object Response Handling:**
    - Modified  to more flexibly handle responses from .
    - The add-task module now check for the AI-generated object in both  and , improving compatibility with different AI provider response structures (e.g., Gemini).

These changes enhance the reliability of AI interactions, particularly with the Google provider, and ensure accurate telemetry collection.
This commit is contained in:
Eyal Toledano
2025-05-13 12:13:35 -04:00
parent 59230c4d91
commit e53d5e1577
5 changed files with 72 additions and 59 deletions

View File

@@ -9,7 +9,7 @@ import { generateText, streamText, generateObject } from 'ai'; // Import from ma
import { log } from '../../scripts/modules/utils.js'; // Import logging utility
// Consider making model configurable via config-manager.js later
const DEFAULT_MODEL = 'gemini-2.0-pro'; // Or a suitable default
const DEFAULT_MODEL = 'gemini-2.5-pro-exp-03-25'; // Or a suitable default
const DEFAULT_TEMPERATURE = 0.2; // Or a suitable default
/**
@@ -52,7 +52,15 @@ async function generateGoogleText({
});
// Assuming result structure provides text directly or within a property
return result.text; // Adjust based on actual SDK response
// return result.text; // Adjust based on actual SDK response
// Return both text and usage
return {
text: result.text,
usage: {
inputTokens: result.usage.promptTokens,
outputTokens: result.usage.completionTokens
}
};
} catch (error) {
log(
'error',
@@ -143,7 +151,7 @@ async function generateGoogleObject({
// const model = google.getGenerativeModel({ model: modelId }); // Incorrect model retrieval
const model = googleProvider(modelId); // Correct model retrieval
const { object } = await generateObject({
const result = await generateObject({
model, // Pass the model instance
schema,
messages,
@@ -154,7 +162,15 @@ async function generateGoogleObject({
// Check SDK docs if specific tool calling/JSON mode needs explicit setup.
});
return object; // Return the parsed object
// return object; // Return the parsed object
// Return both object and usage
return {
object: result.object,
usage: {
inputTokens: result.usage.promptTokens,
outputTokens: result.usage.completionTokens
}
};
} catch (error) {
log(
'error',