feat: add new model in gen image

redevrx · Jan 18, 2024 · f359c13 · f359c13
1 parent 1f711f1
commit f359c13
Show file tree

Hide file tree

Showing 9 changed files with 68 additions and 31 deletions.
diff --git a/example/lib/generate_img_screen.dart b/example/lib/generate_img_screen.dart
@@ -28,7 +28,7 @@ class _GenImgScreenState extends State<GenImgScreen> {
     const prompt = "Snake eat cat.";
 
     final request = GenerateImage(prompt, 1,
-        size: ImageSize.size256, responseFormat: Format.url);
+        model: DallE3(), size: ImageSize.size256, responseFormat: Format.url);
     final response = await openAI.generateImage(request);
     setState(() {
       img = "${response?.data?.last?.url}";

diff --git a/example_app/openai_app/lib/bloc/openai/openai_bloc.dart b/example_app/openai_app/lib/bloc/openai/openai_bloc.dart
@@ -153,7 +153,7 @@ class OpenAIBloc extends Cubit<OpenAIState> {
   ///[generateImage]
   void generateImage() async {
     final request = GenerateImage(_txtInput.value.text, 1,
-        size: ImageSize.size1024, responseFormat: Format.url);
+        model: DallE3(), size: ImageSize.size1024, responseFormat: Format.url);
 
     ///update user chat message
     list.add(Message(isBot: false, message: getTextInput().value.text));

diff --git a/lib/chat_gpt_sdk.dart b/lib/chat_gpt_sdk.dart
@@ -45,3 +45,4 @@ export 'src/model/chat_complete/enum/function_call.dart';
 export 'src/model/chat_complete/request/messages.dart';
 export 'src/model/chat_complete/request/function_data.dart';
 export 'src/model/chat_complete/request/response_format.dart';
+export 'src/model/gen_image/enum/generate_image_model.dart';
diff --git a/lib/src/model/complete_text/request/complete_text.dart b/lib/src/model/complete_text/request/complete_text.dart
@@ -4,16 +4,6 @@ class CompleteText {
   final String prompt;
 
   /// ## completion models
-  /// - TextDavinci3Model();
-  /// - TextDavinci2Model();
-  /// - CodeDavinci2Model();
-  /// - TextCurie001Model();
-  /// - TextBabbage001Model();
-  /// - TextAda001Model();
-  /// - DavinciModel();
-  /// - CurieModel();
-  /// - BabbageModel();
-  /// - AdaModel();
   /// - ModelFromValue(model: 'your-model-name');
   final Model model;
   final double temperature;

diff --git a/lib/src/model/gen_image/enum/generate_image_model.dart b/lib/src/model/gen_image/enum/generate_image_model.dart
@@ -0,0 +1,15 @@
+import 'package:chat_gpt_sdk/chat_gpt_sdk.dart';
+
+sealed class GenerateImageModel {
+  final String model;
+
+  GenerateImageModel({required this.model});
+}
+
+class DallE2 extends GenerateImageModel {
+  DallE2() : super(model: kDallE2);
+}
+
+class DallE3 extends GenerateImageModel {
+  DallE3() : super(model: kDallE3);
+}
diff --git a/lib/src/model/gen_image/request/generate_image.dart b/lib/src/model/gen_image/request/generate_image.dart
@@ -1,4 +1,5 @@
 import 'package:chat_gpt_sdk/src/model/gen_image/enum/format.dart';
+import 'package:chat_gpt_sdk/src/model/gen_image/enum/generate_image_model.dart';
 import 'package:chat_gpt_sdk/src/model/gen_image/enum/image_size.dart';
 
 class GenerateImage {
@@ -17,16 +18,20 @@ class GenerateImage {
   ///A unique identifier representing your end-user, which can help OpenAI to monitor and detect abuse.
   final String user;
 
+  final GenerateImageModel model;
+
   GenerateImage(
     this.prompt,
     this.n, {
+    required this.model,
     this.size = ImageSize.size1024,
     this.responseFormat = Format.url,
     this.user = "",
   }) : assert(1 <= n && n <= 10, 'n must be between 1 and 10.');
 
   Map<String, dynamic> toJson() => Map.of({
         "prompt": prompt,
+        "model": model.model,
         "n": n,
         "size": size?.size,
         "response_format": responseFormat?.getName(),

diff --git a/lib/src/utils/constants.dart b/lib/src/utils/constants.dart
@@ -73,6 +73,10 @@ const kDavinci002Model = 'davinci-002';
 const kTextMStable = 'text-moderation-stable';
 const kTextMLast = 'text-moderation-latest';
 
+///generate image model
+const kDallE2 = 'dall-e-2';
+const kDallE3 = 'dall-e-3';
+
 ///default header
 Map<String, String> kHeader(
   String? token,

diff --git a/test/model/gen_image/request/generate_image_test.dart b/test/model/gen_image/request/generate_image_test.dart
@@ -6,6 +6,7 @@ main() {
     final target = GenerateImage(
       'test',
       2,
+      model: DallE3(),
       size: ImageSize.size1024,
       responseFormat: Format.url,
     );
@@ -19,6 +20,7 @@ main() {
 
   test('set value with enum', () {
     final target1 = GenerateImage(
+      model: DallE3(),
       'test',
       1,
       size: ImageSize.size256,
@@ -33,6 +35,7 @@ main() {
     expect(target1.user, 'user');
 
     final target2 = GenerateImage(
+      model: DallE3(),
       'test',
       2,
       size: ImageSize.size512,
@@ -42,41 +45,56 @@ main() {
     expect(target2.size?.size, '512x512');
     expect(target2.responseFormat?.name, 'url');
 
-    final target3 =
-        GenerateImage('test', 1, size: ImageSize.size1024, user: 'user');
+    final target3 = GenerateImage('test', 1,
+        model: DallE3(), size: ImageSize.size1024, user: 'user');
     expect(target3.size?.size, '1024x1024');
   });
 
   group('GeneratedImageSize', () {
     test('normal', () {
-      expect(GenerateImage('test', 2).size?.size, '1024x1024');
       expect(
-        GenerateImage('test', 2, size: ImageSize.size256).size?.size,
+          GenerateImage(
+            'test',
+            2,
+            model: DallE3(),
+          ).size?.size,
+          '1024x1024');
+      expect(
+        GenerateImage('test', 2, model: DallE3(), size: ImageSize.size256)
+            .size
+            ?.size,
         '256x256',
       );
       expect(
-        GenerateImage('test', 2, size: ImageSize.size512).size?.size,
+        GenerateImage('test', 2, model: DallE3(), size: ImageSize.size512)
+            .size
+            ?.size,
         '512x512',
       );
       expect(
-        GenerateImage('test', 2, size: ImageSize.size1024).size?.size,
+        GenerateImage('test', 2, model: DallE3(), size: ImageSize.size1024)
+            .size
+            ?.size,
         '1024x1024',
       );
     });
   });
 
   test('n must be between 1 and 10', () {
-    expect(() => GenerateImage('test', 0), throwsA(isA<AssertionError>()));
-    expect(() => GenerateImage('test', 11), throwsA(isA<AssertionError>()));
+    expect(() => GenerateImage('test', model: DallE3(), 0),
+        throwsA(isA<AssertionError>()));
+    expect(() => GenerateImage('test', model: DallE3(), 11),
+        throwsA(isA<AssertionError>()));
 
-    expect(GenerateImage('test', 1).n, 1);
-    expect(GenerateImage('test', 10).n, 10);
+    expect(GenerateImage('test', model: DallE3(), 1).n, 1);
+    expect(GenerateImage('test', model: DallE3(), 10).n, 10);
   });
 
   group('toJson', () {
     test('example', () {
       final json = GenerateImage(
         'test',
+        model: DallE3(),
         1,
         size: ImageSize.size256,
         responseFormat: Format.b64Json,

diff --git a/test/openai_test.dart b/test/openai_test.dart
@@ -73,6 +73,7 @@ void main() async {
     expect(
       () => ai.generateImage(GenerateImage(
         'prompt',
+        model: DallE3(),
         1,
         size: ImageSize.size256,
         responseFormat: Format.url,
@@ -631,7 +632,8 @@ void main() async {
 
   group('chatGPT Image Generate With Prompt test case', () {
     test('chatGPT Image Generate With Prompt success case  test', () async {
-      final request = GenerateImage('snake red eating cat.', 2);
+      final request =
+          GenerateImage('snake red eating cat.', model: DallE3(), 2);
 
       when(openAI.generateImage(request))
           .thenAnswer((_) async => GenImgResponse());
@@ -645,7 +647,8 @@ void main() async {
     test(
       'chatGPT Image Generate With Prompt success case return value test',
       () async {
-        final request = GenerateImage('snake red eating cat.', 2);
+        final request =
+            GenerateImage('snake red eating cat.', model: DallE3(), 2);
 
         when(openAI.generateImage(request))
             .thenAnswer((_) async => GenImgResponse(created: 1221120));
@@ -661,7 +664,8 @@ void main() async {
     test(
       'chatGPT Image Generate With Prompt error case with n is 0 test',
       () async {
-        final request = GenerateImage('snake red eating cat.', 1);
+        final request =
+            GenerateImage('snake red eating cat.', model: DallE3(), 1);
 
         when(openAI.generateImage(request))
             .thenAnswer((_) async => GenImgResponse(created: 1221120));
@@ -673,7 +677,7 @@ void main() async {
         expect(response?.created, 1221120);
         expect(response?.data, null);
         expect(
-          () => GenerateImage('snake red eating cat.', 0),
+          () => GenerateImage('snake red eating cat.', model: DallE3(), 0),
           throwsA(isA<AssertionError>()),
         );
       },
@@ -760,7 +764,7 @@ void main() async {
     test(
       'chatGPT Image Generate with success case',
       () async {
-        final request = GenerateImage('cat eating snake', 1);
+        final request = GenerateImage('cat eating snake', model: DallE3(), 1);
 
         when(openAI.generateImage(request))
             .thenAnswer((realInvocation) async => GenImgResponse(
@@ -779,7 +783,7 @@ void main() async {
     test(
       'chatGPT Image Generate with cancel gen success case',
       () {
-        final request = GenerateImage('cat eating snake', 1);
+        final request = GenerateImage('cat eating snake', model: DallE3(), 1);
 
         when(openAI.generateImage(request))
             .thenAnswer((realInvocation) async => GenImgResponse(
@@ -796,7 +800,7 @@ void main() async {
     test(
       'chatGPT Image Generate with return two image success case',
       () async {
-        final request = GenerateImage('cat eating snake', 2);
+        final request = GenerateImage('cat eating snake', model: DallE3(), 2);
 
         when(openAI.generateImage(request)).thenAnswer(
           (realInvocation) async => GenImgResponse(created: 912312, data: [
@@ -816,7 +820,7 @@ void main() async {
     test(
       'chatGPT Image Generate with error case',
       () async {
-        final request = GenerateImage('', 1);
+        final request = GenerateImage('', model: DallE3(), 1);
 
         when(openAI.generateImage(request))
             .thenAnswer((realInvocation) async => GenImgResponse(
@@ -833,7 +837,7 @@ void main() async {
     test(
       'chatGPT Image Generate with openai auth error case',
       () async {
-        final request = GenerateImage('snake', 1);
+        final request = GenerateImage('snake', model: DallE3(), 1);
 
         when(openAI.generateImage(request))
             .thenThrow(OpenAIAuthError(code: 404));