Skip to content

Commit

Permalink
Allow the developer to use their own http.Client (#17)
Browse files Browse the repository at this point in the history
This lets them:
1. pick the implementation
2. pool connections
  • Loading branch information
brianquinlan authored Jan 29, 2024
1 parent 413000f commit 616a97c
Show file tree
Hide file tree
Showing 2 changed files with 14 additions and 5 deletions.
11 changes: 8 additions & 3 deletions pkgs/google_generative_ai/lib/src/client.dart
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand All @@ -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())
Expand Down
8 changes: 6 additions & 2 deletions pkgs/google_generative_ai/lib/src/model.dart
Original file line number Diff line number Diff line change
Expand Up @@ -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';
Expand Down Expand Up @@ -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 {
Expand Down

0 comments on commit 616a97c

Please sign in to comment.