Send headers only to Yandex API host (#49)

vladcto · May 4, 2024 · 72b5b85 · 72b5b85
1 parent 6ef4084
commit 72b5b85
Show file tree

Hide file tree

Showing 5 changed files with 59 additions and 32 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,3 +1,7 @@
+## UNRELEASED
+
+- Add auth tokens only to Yandex API host requests.
+
 ## 1.0.2
 
 - Update docs.

diff --git a/lib/src/logic/client/yandex_gpt_header_interceptor.dart b/lib/src/logic/client/yandex_gpt_header_interceptor.dart
@@ -1,5 +1,6 @@
 import 'package:dio/dio.dart';
 import 'package:yandex_gpt_rest_api/src/models/auth/auth_token.dart';
+import 'package:yandex_gpt_rest_api/src/utils/constants/api_url.dart';
 import 'package:yandex_gpt_rest_api/src/utils/constants/headers.dart';
 
 /// Interceptor for adding authentication data to the request.
@@ -15,10 +16,12 @@ class YandexGptHeaderInterceptor extends Interceptor {
 
   @override
   void onRequest(RequestOptions options, RequestInterceptorHandler handler) {
-    options.headers.addAll({
-      authHeaderName: _token.toString(),
-      catalogIdHeaderName: _catalog,
-    });
+    if (options.uri.host == ApiUrl.host) {
+      options.headers.addAll({
+        authHeaderName: _token.toString(),
+        catalogIdHeaderName: _catalog,
+      });
+    }
     handler.next(options);
   }
 

diff --git a/lib/src/logic/client/yandex_gpt_http_client.dart b/lib/src/logic/client/yandex_gpt_http_client.dart
@@ -14,36 +14,34 @@ class YandexGptHttpClient {
     String url, {
     Map<String, dynamic>? body,
     CancelToken? cancelToken,
-  }) async {
-    return await _fetch(
-      _dio.post<String>(
-        url,
-        data: jsonEncode(body),
-        cancelToken: cancelToken,
-      ),
-    );
-  }
+  }) =>
+      _fetch(
+        _dio.post<String>(
+          url,
+          data: jsonEncode(body),
+          cancelToken: cancelToken,
+        ),
+      );
 
   Future<Map<String, dynamic>> get(
     String url, {
     Map<String, dynamic>? body,
     CancelToken? cancelToken,
-  }) async {
-    return await _fetch(
-      _dio.get<String>(
-        url,
-        data: jsonEncode(body),
-        cancelToken: cancelToken,
-      ),
-    );
-  }
+  }) =>
+      _fetch(
+        _dio.get<String>(
+          url,
+          data: jsonEncode(body),
+          cancelToken: cancelToken,
+        ),
+      );
 
   /// A wrapper for handling [ApiError] responses.
   Future<Map<String, dynamic>> _fetch(Future<Response<String>> request) async {
-    late final Response<String> response;
-
     try {
-      response = await request;
+      final response = await request;
+      final jsonBody = jsonDecode(response.data!) as Map<String, dynamic>;
+      return jsonBody;
     } on DioException catch (e) {
       // Check if response contains [ApiError].
       final body = jsonDecode(e.response?.data as String? ?? "{}");
@@ -52,8 +50,5 @@ class YandexGptHttpClient {
       if (apiError == null) rethrow;
       throw apiError;
     }
-
-    final jsonBody = jsonDecode(response.data!) as Map<String, dynamic>;
-    return jsonBody;
   }
 }
diff --git a/lib/src/utils/constants/api_url.dart b/lib/src/utils/constants/api_url.dart
@@ -1,7 +1,8 @@
 abstract final class ApiUrl {
-  static const _apiHost = "https://llm.api.cloud.yandex.net";
-  static const _foundationHost = "$_apiHost/foundationModels/v1";
+  static const _apiHost = "https://$host";
+  static const host = "llm.api.cloud.yandex.net";
 
+  static const _foundationHost = "$_apiHost/foundationModels/v1";
   static const textGeneration = "$_foundationHost/completion";
   static const textGenerationAsync = "$_foundationHost/completionAsync";
   static const tokenizeText = "$_foundationHost/tokenize";

diff --git a/test/logic/client/yandex_gpt_interceptor_test.dart b/test/logic/client/yandex_gpt_interceptor_test.dart
@@ -4,6 +4,7 @@ import 'package:mockito/mockito.dart';
 import 'package:test/test.dart';
 import 'package:yandex_gpt_rest_api/src/logic/client/yandex_gpt_header_interceptor.dart';
 import 'package:yandex_gpt_rest_api/src/models/auth/auth_token.dart';
+import 'package:yandex_gpt_rest_api/src/utils/constants/api_url.dart';
 import 'package:yandex_gpt_rest_api/src/utils/constants/headers.dart';
 
 import 'yandex_gpt_interceptor_test.mocks.dart';
@@ -33,7 +34,10 @@ void main() {
         token: const AuthToken.iam('iam'),
       );
 
-      interceptor.onRequest(RequestOptions(), requestInterceptorHandler);
+      interceptor.onRequest(
+        RequestOptions(path: ApiUrl.tokenizeText),
+        requestInterceptorHandler,
+      );
       final sent = verify(requestInterceptorHandler.next(captureAny))
           .captured
           .first as RequestOptions;
@@ -45,14 +49,34 @@ void main() {
       );
     });
 
+    test('Empty headers if host not Yandex API', () {
+      final interceptor = YandexGptHeaderInterceptor(
+        catalog: "catalog",
+        token: const AuthToken.iam('iam'),
+      );
+
+      interceptor.onRequest(
+        RequestOptions(path: 'yandex.ru'),
+        requestInterceptorHandler,
+      );
+      final sent = verify(requestInterceptorHandler.next(captureAny))
+          .captured
+          .first as RequestOptions;
+
+      expect(sent.headers, isEmpty);
+    });
+
     test('Change token', () {
       final interceptor = YandexGptHeaderInterceptor(
         catalog: "catalog",
         token: const AuthToken.iam('iam'),
       );
 
       interceptor.changeToken(const AuthToken.apiKey('key'));
-      interceptor.onRequest(RequestOptions(), requestInterceptorHandler);
+      interceptor.onRequest(
+        RequestOptions(path: ApiUrl.tokenizeText),
+        requestInterceptorHandler,
+      );
       final afterChange = verify(requestInterceptorHandler.next(captureAny))
           .captured
           .first as RequestOptions;