From 616a97cfa621b91e73eb5726408a96be9d841f9d Mon Sep 17 00:00:00 2001
From: Brian Quinlan <bquinlan@google.com>
Date: Mon, 29 Jan 2024 11:35:55 -0800
Subject: [PATCH] Allow the developer to use their own http.Client (#17)

This lets them:
1. pick the implementation
2. pool connections
---
 pkgs/google_generative_ai/lib/src/client.dart | 11 ++++++++---
 pkgs/google_generative_ai/lib/src/model.dart  |  8 ++++++--
 2 files changed, 14 insertions(+), 5 deletions(-)

diff --git a/pkgs/google_generative_ai/lib/src/client.dart b/pkgs/google_generative_ai/lib/src/client.dart
index 8f80afb..8ecf94f 100644
--- a/pkgs/google_generative_ai/lib/src/client.dart
+++ b/pkgs/google_generative_ai/lib/src/client.dart
@@ -28,13 +28,16 @@ const _clientName = 'genai-dart/$_packageVersion';
 
 final class HttpApiClient implements ApiClient {
   final String _apiKey;
+  final http.Client? _httpClient;
   late final _headers = {
     'x-goog-api-key': _apiKey,
     'x-goog-api-client': _clientName
   };
 
-  HttpApiClient({required String model, required String apiKey})
-      : _apiKey = apiKey;
+  HttpApiClient(
+      {required String model, required String apiKey, http.Client? httpClient})
+      : _apiKey = apiKey,
+        _httpClient = httpClient ?? http.Client();
 
   @override
   Future<String> makeRequest(Uri uri, Uint8List body) async {
@@ -52,7 +55,9 @@ final class HttpApiClient implements ApiClient {
         ..bodyBytes = body
         ..headers.addAll(_headers)
         ..headers['Content-Type'] = 'application/json';
-      final response = await request.send();
+      final response = _httpClient == null
+          ? await request.send()
+          : await _httpClient.send(request);
       await response.stream
           .toStringStream()
           .transform(const LineSplitter())
diff --git a/pkgs/google_generative_ai/lib/src/model.dart b/pkgs/google_generative_ai/lib/src/model.dart
index cbdcc5c..4bca087 100644
--- a/pkgs/google_generative_ai/lib/src/model.dart
+++ b/pkgs/google_generative_ai/lib/src/model.dart
@@ -15,6 +15,8 @@
 import 'dart:async';
 import 'dart:convert';
 
+import 'package:http/http.dart' as http;
+
 import 'api.dart';
 import 'client.dart';
 import 'content.dart';
@@ -46,14 +48,16 @@ final class GenerativeModel {
       {required String model,
       required String apiKey,
       List<SafetySetting> safetySettings = const [],
-      GenerationConfig? generationConfig})
+      GenerationConfig? generationConfig,
+      http.Client? httpClient})
       :
         // TODO: Allow `models/` prefix and strip it.
         // https://github.com/google/generative-ai-js/blob/2be48f8e5427f2f6191f24bcb8000b450715a0de/packages/main/src/models/generative-model.ts#L59
         _model = model,
         _safetySettings = safetySettings,
         _generationConfig = generationConfig,
-        _client = HttpApiClient(model: model, apiKey: apiKey);
+        _client =
+            HttpApiClient(model: model, apiKey: apiKey, httpClient: httpClient);
 
   Future<Object> _makeRequest(
       Task task, Map<String, Object?> parameters) async {