Refactor SDK logic and improve logging of LLM inferences

doc/drafts/GENUI.md

@@ -10,7 +10,7 @@ The core of the application is a continuous interaction loop between the user, t
 
 1. **Prompt**: The user enters a text prompt describing the desired user interface (e.g., "Create a login form with a username field, a password field, and a login button").
 2. **Generation**: The prompt is sent to the generative AI model. The model is given instructions to return a JSON object that conforms to a predefined UI schema.
-3. **Rendering**: The Flutter client receives the JSON response. A `DynamicUi` widget parses the JSON and recursively builds a native Flutter widget tree based on the definition.
+3. **Rendering**: The Flutter client receives the JSON response. A `SurfaceWidget` widget parses the JSON and recursively builds a native Flutter widget tree based on the definition.
 4. **Interaction**: The user interacts with the rendered UI (e.g., types in a text field, taps a button).
 5. **Event Feedback**: Each interaction generates a `UiEvent` object. This event is sent back to the AI model, framed as a "function call response". This makes the model believe it has invoked a tool that returned the user's action as its result.
 6. **Update**: The model processes the event and the conversation history, and can then generate a new JSON UI definition to reflect the new state of the application (e.g., showing a loading spinner after the login button is pressed). This cycle allows for creating truly interactive and stateful applications driven by the AI.

pkgs/flutter_genui/IMPLEMENTATION.md

@@ -35,7 +35,7 @@ This layer defines the data structures that represent the dynamic UI and the con
 - **`Catalog` and `CatalogItem`**: These classes define the registry of available UI components. The `Catalog` holds a list of `CatalogItem`s, and each `CatalogItem` defines a widget's name, its data schema, and a builder function to render it.
 - **`UiDefinition` and `UiEvent`**: `UiDefinition` represents a complete UI tree to be rendered, including the root widget and a map of all widget definitions. `UiEvent` is a data object representing a user interaction (e.g., a button tap), which is sent back to the `ConversationManager`.
 - **`ChatMessage`**: A sealed class representing the different types of messages in the conversation history: `UserPrompt`, `TextResponse`, `UiResponse`, and `SystemMessage`.
-- **`DynamicUi`**: The Flutter widget responsible for recursively building the UI tree from a `UiDefinition`. It uses the provided `Catalog` to find the correct widget builder for each node in the tree.
+- **`SurfaceWidget`**: The Flutter widget responsible for recursively building the UI tree from a `UiDefinition`. It uses the provided `Catalog` to find the correct widget builder for each node in the tree.
 
 ### 4. Widget Catalog Layer (`lib/src/catalog/`)
 
@@ -57,7 +57,7 @@ This layer provides a set of core, general-purpose UI widgets that can be used o
    - For an **`update`** action, it finds the existing `UiResponse` with the matching `surfaceId` and replaces it with the new definition.
    - For a **`delete`** action, it removes the corresponding `UiResponse` from the chat history.
 8. **UI Rendering**: The `_uiDataStreamController` in the `ConversationManager` emits the updated chat history. The `ConversationWidget` rebuilds, and its `ListView` now includes the new or updated `UiResponse`.
-9. **Dynamic UI Build**: The `DynamicUi` widget within the `ConversationWidget` receives the `UiDefinition` and recursively builds the Flutter widget tree using the `Catalog`.
+9. **Dynamic UI Build**: The `SurfaceWidget` widget within the `ConversationWidget` receives the `UiDefinition` and recursively builds the Flutter widget tree using the `Catalog`.
 10. **User Interaction**: The user interacts with the newly generated UI (e.g., clicks a button).
 11. **Event Dispatch**: The widget's `onPressed` handler calls the `dispatchEvent` function, creating a `UiEvent`.
 12. **Event Handling**: The `EventDebouncer` collects and sends the event to the `ConversationManager`.

pkgs/flutter_genui/lib/flutter_genui.dart

@@ -12,5 +12,5 @@ export 'src/core/event_debouncer.dart';
 export 'src/core/genui_manager.dart';
 export 'src/model/catalog.dart';
 export 'src/model/catalog_item.dart';
-export 'src/model/dynamic_ui.dart';
+export 'src/model/surface_widget.dart';
 export 'src/model/ui_models.dart';

pkgs/flutter_genui/lib/src/ai_client/ai_client.dart

@@ -275,7 +275,7 @@ class AiClient implements LlmConnection {
     List<Tool>? tools,
     ToolConfig? toolConfig,
   }) {
-    return GenerativeModelWrapper(
+    return FirebaseAiGenerativeModel(
       FirebaseAI.googleAI().generativeModel(
         model: configuration.model.value.modelName,
         systemInstruction: systemInstruction,
@@ -439,15 +439,19 @@ class AiClient implements LlmConnection {
 
     while (toolUsageCycle < maxToolUsageCycles && capturedResult == null) {
       toolUsageCycle++;
-      _log('Generating content with:');
-      for (final content in contents) {
-        _log(const JsonEncoder.withIndent('  ').convert(content.toJson()));
-      }
-      _log(
-        'With functions: '
-        '${allowedFunctionNames.join(', ')}',
-      );
+
+      final concatenatedContents = contents
+          .map((c) => const JsonEncoder.withIndent('  ').convert(c.toJson()))
+          .join('\n');
+
+      _log('''****** Performing Inference ******
+$concatenatedContents
+With functions:
+  '${allowedFunctionNames.join(', ')}',
+  ''');
+      final inferenceStartTime = DateTime.now();
       final response = await model.generateContent(contents);
+      final elapsed = DateTime.now().difference(inferenceStartTime);
 
       // If the generate call succeeds, we need to reset the delay for the next
       // retry. If the generate call throws, this won't get called, and the
@@ -458,6 +462,12 @@ class AiClient implements LlmConnection {
         inputTokenUsage += response.usageMetadata!.promptTokenCount ?? 0;
         outputTokenUsage += response.usageMetadata!.candidatesTokenCount ?? 0;
       }
+      _log(
+        '****** Completed Inference ******\n'
+        'Latency = ${elapsed.inMilliseconds}ms\n'
+        'Output tokens = ${response.usageMetadata?.candidatesTokenCount ?? 0}\n'
+        'Prompt tokens = ${response.usageMetadata?.promptTokenCount ?? 0}',
+      );
 
       if (response.candidates.isEmpty) {
         _warn('Response has no candidates: ${response.promptFeedback}');
@@ -493,13 +503,15 @@ class AiClient implements LlmConnection {
       for (final call in functionCalls) {
         if (call.name == outputToolName) {
           try {
-            capturedResult = call.args['output'] as T?;
+            capturedResult =
+                (call.args['parameters'] as Map<String, Object?>)['output']
+                    as T?;
           } catch (e, s) {
             _error('Unable to read output: $call [${call.args}]: $e', s);
           }
           _log(
-            'Invoked output tool ${call.name} with args ${call.args}. '
-            'Final result: $capturedResult',
+            '****** Gen UI Output ******.\n'
+            '${const JsonEncoder.withIndent('  ').convert(capturedResult)}}',
           );
           continue;
         }

pkgs/flutter_genui/lib/src/ai_client/generative_model_interface.dart

@@ -6,36 +6,18 @@ import 'package:firebase_ai/firebase_ai.dart';
 
 /// An interface for a generative model, allowing for mock implementations.
 abstract class GenerativeModelInterface {
-  Future<GenerateContentResponse> generateContent(
-    Iterable<Content> content, {
-    List<SafetySetting>? safetySettings,
-    GenerationConfig? generationConfig,
-    List<Tool>? tools,
-    ToolConfig? toolConfig,
-  });
+  Future<GenerateContentResponse> generateContent(Iterable<Content> content);
 }
 
 /// A wrapper for the `firebase_ai` [GenerativeModel] that implements the
 /// [GenerativeModelInterface].
-class GenerativeModelWrapper implements GenerativeModelInterface {
+class FirebaseAiGenerativeModel implements GenerativeModelInterface {
   final GenerativeModel _model;
 
-  GenerativeModelWrapper(this._model);
+  FirebaseAiGenerativeModel(this._model);
 
   @override
-  Future<GenerateContentResponse> generateContent(
-    Iterable<Content> content, {
-    List<SafetySetting>? safetySettings,
-    GenerationConfig? generationConfig,
-    List<Tool>? tools,
-    ToolConfig? toolConfig,
-  }) {
-    return _model.generateContent(
-      content,
-      safetySettings: safetySettings,
-      generationConfig: generationConfig,
-      tools: tools,
-      toolConfig: toolConfig,
-    );
+  Future<GenerateContentResponse> generateContent(Iterable<Content> content) {
+    return _model.generateContent(content);
   }
 }

pkgs/flutter_genui/lib/src/core/conversation_widget.dart

@@ -2,7 +2,7 @@ import 'package:flutter/material.dart';
 
 import '../model/catalog.dart';
 import '../model/chat_message.dart';
-import '../model/dynamic_ui.dart';
+import '../model/surface_widget.dart';
 import '../model/ui_models.dart';
 
 class ConversationWidget extends StatelessWidget {
@@ -50,7 +50,7 @@ class ConversationWidget extends StatelessWidget {
           ),
           UiResponse() => Padding(
             padding: const EdgeInsets.all(16.0),
-            child: DynamicUi(
+            child: SurfaceWidget(
               key: message.uiKey,
               catalog: catalog,
               surfaceId: message.surfaceId,

pkgs/flutter_genui/lib/src/model/dynamic_ui.dart → pkgs/flutter_genui/lib/src/model/surface_widget.dart

@@ -7,8 +7,8 @@ import 'ui_models.dart';
 ///
 /// It takes an initial [definition] and reports user interactions
 /// via the [onEvent] callback.
-class DynamicUi extends StatefulWidget {
-  const DynamicUi({
+class SurfaceWidget extends StatefulWidget {
+  const SurfaceWidget({
     super.key,
     required this.catalog,
     required this.surfaceId,
@@ -28,11 +28,12 @@ class DynamicUi extends StatefulWidget {
   final Catalog catalog;
 
   @override
-  State<DynamicUi> createState() => _DynamicUiState();
+  State<SurfaceWidget> createState() => _SurfaceWidgetState();
 }
 
-class _DynamicUiState extends State<DynamicUi> {
-  /// Dispatches an event by calling the public [DynamicUi.onEvent] callback.
+class _SurfaceWidgetState extends State<SurfaceWidget> {
+  /// Dispatches an event by calling the public [SurfaceWidget.onEvent]
+  /// callback.
   void _dispatchEvent({
     required String widgetId,
     required String eventType,

pkgs/flutter_genui/test/ai_client/ai_client_test.dart

@@ -44,10 +44,13 @@ void main() {
         description: 'd',
         invokeFunction: (_) async => <String, Object?>{},
       );
-      expect(
-        () => createClient(tools: [tool1, tool2]),
-        throwsA(isA<AiClientException>()),
-      );
+      try {
+        createClient(tools: [tool1, tool2]);
+        fail('should throw');
+      } catch (e) {
+        expect(e, isA<AiClientException>());
+        expect((e as AiClientException).message, contains('Duplicate tool(s)'));
+      }
     });
 
     test('generateContent returns structured data', () async {
@@ -56,7 +59,9 @@ void main() {
         Candidate(
           Content.model([
             FunctionCall('provideFinalOutput', {
-              'output': {'key': 'value'},
+              'parameters': {
+                'output': {'key': 'value'},
+              },
             }),
           ]),
           [],
@@ -104,7 +109,9 @@ void main() {
           Candidate(
             Content.model([
               FunctionCall('provideFinalOutput', {
-                'output': {'final': 'result'},
+                'parameters': {
+                  'output': {'final': 'result'},
+                },
               }),
             ]),
             [],
@@ -132,7 +139,9 @@ void main() {
         Candidate(
           Content.model([
             FunctionCall('provideFinalOutput', {
-              'output': {'key': 'value'},
+              'parameters': {
+                'output': {'key': 'value'},
+              },
             }),
           ]),
           [],
@@ -174,7 +183,9 @@ void main() {
           Candidate(
             Content.model([
               FunctionCall('provideFinalOutput', {
-                'output': {'final': 'result'},
+                'parameters': {
+                  'output': {'final': 'result'},
+                },
               }),
             ]),
             [],
@@ -190,6 +201,7 @@ void main() {
       ], Schema.object(properties: {'final': Schema.string()}));
 
       expect(result, isNotNull);
+      expect(result!['final'], 'result');
     });
 
     test('generateContent returns null if no candidates', () async {
@@ -273,31 +285,6 @@ void main() {
       expect(result, isNull);
     });
 
-    test('token usage is tracked', () async {
-      client = createClient();
-      fakeModel.response = GenerateContentResponse([
-        Candidate(
-          Content.model([
-            FunctionCall('provideFinalOutput', {
-              'output': {'key': 'value'},
-            }),
-          ]),
-          [],
-          null,
-          null,
-          null,
-        ),
-      ], PromptFeedback(null, null, []));
-
-      await client.generateContent<Map<String, Object?>>(
-        [],
-        Schema.object(properties: {'key': Schema.string()}),
-      );
-
-      expect(client.inputTokenUsage, 0);
-      expect(client.outputTokenUsage, 0);
-    });
-
     test('logging callback is called', () async {
       final logMessages = <String>[];
       client = createClient(

pkgs/flutter_genui/test/catalog/core_widgets_test.dart

@@ -25,7 +25,7 @@ void main() {
       await tester.pumpWidget(
         MaterialApp(
           home: Scaffold(
-            body: DynamicUi(
+            body: SurfaceWidget(
               catalog: testCatalog,
               surfaceId: 'testSurface',
               definition: UiDefinition.fromMap(definition),

pkgs/flutter_genui/test/core/conversation_widget_test.dart → pkgs/flutter_genui/test/core/surface_widget_test.dart

@@ -6,7 +6,7 @@ import 'package:flutter/material.dart';
 import 'package:flutter_genui/src/core/conversation_widget.dart';
 import 'package:flutter_genui/src/core/core_catalog.dart';
 import 'package:flutter_genui/src/model/chat_message.dart';
-import 'package:flutter_genui/src/model/dynamic_ui.dart';
+import 'package:flutter_genui/src/model/surface_widget.dart';
 import 'package:flutter_test/flutter_test.dart';
 
 void main() {
@@ -91,7 +91,7 @@ void main() {
           ),
         ),
       );
-      expect(find.byType(DynamicUi), findsOneWidget);
+      expect(find.byType(SurfaceWidget), findsOneWidget);
       expect(find.text('UI Content'), findsOneWidget);
     });
   });

pkgs/flutter_genui/test/dynamic_ui_test.dart

@@ -5,7 +5,7 @@ import 'package:flutter_test/flutter_test.dart';
 void main() {
   final testCatalog = Catalog([elevatedButtonCatalogItem, text]);
 
-  testWidgets('DynamicUi builds a widget from a definition', (
+  testWidgets('SurfaceWidget builds a widget from a definition', (
     WidgetTester tester,
   ) async {
     final definition = UiDefinition.fromMap({
@@ -29,7 +29,7 @@ void main() {
 
     await tester.pumpWidget(
       MaterialApp(
-        home: DynamicUi(
+        home: SurfaceWidget(
           catalog: testCatalog,
           surfaceId: 'testSurface',
           definition: definition,
@@ -42,7 +42,7 @@ void main() {
     expect(find.byType(ElevatedButton), findsOneWidget);
   });
 
-  testWidgets('DynamicUi handles events', (WidgetTester tester) async {
+  testWidgets('SurfaceWidget handles events', (WidgetTester tester) async {
     Map<String, Object?>? event;
 
     final definition = UiDefinition.fromMap({
@@ -66,7 +66,7 @@ void main() {
 
     await tester.pumpWidget(
       MaterialApp(
-        home: DynamicUi(
+        home: SurfaceWidget(
           catalog: testCatalog,
           surfaceId: 'testSurface',
           definition: definition,

pkgs/flutter_genui/test/test_infra/utils.dart

@@ -16,26 +16,24 @@ class FakeGenerativeModel implements GenerativeModelInterface {
   GenerateContentResponse? response;
   List<GenerateContentResponse> responses = [];
   Exception? exception;
+  PromptFeedback? promptFeedback;
 
   @override
   Future<GenerateContentResponse> generateContent(
-    Iterable<Content> content, {
-    List<SafetySetting>? safetySettings,
-    GenerationConfig? generationConfig,
-    List<Tool>? tools,
-    ToolConfig? toolConfig,
-  }) async {
+    Iterable<Content> content,
+  ) async {
     generateContentCallCount++;
     if (exception != null) {
       final e = exception;
       exception = null; // Reset for next call
       throw e!;
     }
     if (responses.isNotEmpty) {
-      return responses.removeAt(0);
+      final response = responses.removeAt(0);
+      return GenerateContentResponse(response.candidates, promptFeedback);
     }
     if (response != null) {
-      return response!;
+      return GenerateContentResponse(response!.candidates, promptFeedback);
     }
     throw StateError(
       'No response or exception configured for FakeGenerativeModel',