Merge pull request #29 from CommandDash/develop

Develop
CommandDash · Apr 20, 2024 · 298c343 · 298c343
2 parents bc74a79 + 00f2efa
commit 298c343
Show file tree

Hide file tree

Showing 20 changed files with 125 additions and 150 deletions.
diff --git a/commanddash/lib/agent/agent_handler.dart b/commanddash/lib/agent/agent_handler.dart
@@ -37,14 +37,14 @@ class AgentHandler {
         (json['steps'] as List).cast<Map<String, dynamic>>();
 
     final GenerationRepository generationRepository =
-        GenerationRepository.fromJson(json['authdetails']);
+        GenerationRepository.fromJson(json['auth_details']);
 
     return AgentHandler(
       inputs: inputs,
       outputs: outputs,
       steps: steps,
       generationRepository: generationRepository,
-      githubAccessToken: json['authdetails']['githubToken'],
+      githubAccessToken: json['auth_details']['github_token'],
       agentName: json['agent_name'],
       agentVersion: json['agent_version'],
     );

diff --git a/commanddash/lib/agent/output_model.dart b/commanddash/lib/agent/output_model.dart
@@ -49,6 +49,9 @@ class MultiCodeOutput extends Output {
       return code;
     } else {
       for (WorkspaceFile file in value!) {
+        if (code.length > 12000) {
+          break; //maximum char limit ?? TODO: Replace this later with a holistic counting mechanism
+        }
         code += 'File: ${file.path}\n';
         if (file.content == null) {
           code += File(file.path).readAsStringSync();

diff --git a/commanddash/lib/repositories/dash_repository.dart b/commanddash/lib/repositories/dash_repository.dart
@@ -30,26 +30,29 @@ class DashRepository {
     required List<DataSource> datasources,
   }) async {
     try {
-      final response = await dio.post('/agent/get-reference', data: {
-        "name": agentName,
-        "query": query,
-        "version": agentVersion,
-        "matching_doc_data_source_ids": datasources.map((e) => e.id).toList(),
-        "testing": true,
-      });
+      final response = await dio.post(
+        '/agent/get-reference',
+        data: {
+          "name": agentName,
+          "query": query,
+          "version": agentVersion,
+          "matching_doc_data_source_ids": datasources.map((e) => e.id).toList(),
+          "testing": false,
+        },
+      );
       return List<Map<String, dynamic>>.from(response.data['data']).map((e) {
         return DataSource.fromJson(e);
       }).toList();
     } catch (e) {
-      throw Exception('Error fetching datasource');
+      throw Exception('Error fetching datasource: ${e.toString()}');
     }
   }
 
   // TODO: to be tested
   Future<List<Map<String, dynamic>>> getAgents() async {
     try {
       final response = await dio.post('/agent/get-general-agents', data: {
-        "testing": true,
+        "testing": false,
         "min_cli_version": "1.0.0",
       });
       return response.data;

diff --git a/commanddash/lib/repositories/gemini_repository.dart b/commanddash/lib/repositories/gemini_repository.dart
@@ -34,6 +34,7 @@ class GeminiRepository implements GenerationRepository {
     try {
       final model = GenerativeModel(model: 'embedding-001', apiKey: apiKey);
       final content = Content.text(value);
+
       final result = await model.embedContent(
         content,
         taskType: TaskType
@@ -54,34 +55,16 @@ class GeminiRepository implements GenerationRepository {
   }
 
   @override
-  Future<List<List<double>>> getCodeBatchEmbeddings(List<String> code) async {
+  Future<List<List<double>>> getCodeBatchEmbeddings(
+      List<Map<String, dynamic>> code) async {
     try {
-      final response = await HttpApiClient(apiKey: apiKey).makeRequest(
-          Uri.https('generativelanguage.googleapis.com').resolveUri(Uri(
-              pathSegments: [
-                'v1',
-                'models',
-                'embedding-001:batchEmbedContents'
-              ])),
-          {
-            'requests': code
-                .map((e) => <String, Object?>{
-                      'model': 'models/embedding-001',
-                      'content': Content.text(e).toJson(),
-                      'taskType': TaskType.retrievalDocument.toJson(),
-                    })
-                .toList()
-          });
-      try {
-        return (response['embeddings'] as List)
-            .map((e) => List<double>.from(e['values']))
-            .toList();
-      } catch (e) {
-        if (response.containsKey('error')) {
-          throw parseError(response['error']!);
-        }
-        rethrow;
-      }
+      final model = GenerativeModel(model: 'embedding-001', apiKey: apiKey);
+      final embedRequest = code
+          .map((value) => EmbedContentRequest(Content.text(value['content']),
+              title: value['title'], taskType: TaskType.retrievalDocument))
+          .toList();
+      final response = await model.batchEmbedContents(embedRequest);
+      return response.embeddings.map((e) => e.values).toList();
     } on InvalidApiKey catch (_) {
       //Note: this exeception are not thrown anyway by the embedAPIs
       throw InvalidApiKeyException();
@@ -122,6 +105,7 @@ class GeminiRepository implements GenerationRepository {
   @override
   Future<List<List<double>>> getStringBatchEmbeddings(
       List<String> values) async {
+    //TODO: update to batch embed
     try {
       final response = await HttpApiClient(apiKey: apiKey).makeRequest(
           Uri.https('generativelanguage.googleapis.com').resolveUri(Uri(

diff --git a/commanddash/lib/repositories/generation_repository.dart b/commanddash/lib/repositories/generation_repository.dart
@@ -10,7 +10,8 @@ abstract class GenerationRepository {
       List<ChatMessage> messages, String lastMessage);
   // Generates embeddings for the given [code]. This should be using a tasktype of retrievalDocument.
   Future getCodeEmbeddings(String code);
-  Future<List<List<double>>> getCodeBatchEmbeddings(List<String> code);
+  Future<List<List<double>>> getCodeBatchEmbeddings(
+      List<Map<String, dynamic>> code);
   // Generates embeddings for the given [value]. This should be using a tasktype of retrievalQuery.
   Future getStringEmbeddings(String value);
   Future<List<List<double>>> getStringBatchEmbeddings(List<String> value);

diff --git a/commanddash/lib/server/temp.dart b/commanddash/lib/server/temp.dart
@@ -8,7 +8,7 @@
 //     {"id": "436621806", "type": "multi_code_object"},
 //     {"id": "90611917", "type": "default_output"}
 //   ],
-//   "authdetails": {
+//   "auth_details": {
 //     "type": "gemini" | "dash",
 //     "key": "aisjoaisj",
 //     "githubToken": "",

diff --git a/commanddash/lib/steps/find_closest_files/embedding_generator.dart b/commanddash/lib/steps/find_closest_files/embedding_generator.dart
@@ -43,7 +43,9 @@ class EmbeddingGenerator {
 
     // Use the batches API to update the embeddings
     final embeddings = await Future.wait(batches.map((batch) async {
-      final code = batch.map((file) => file.content!).toList();
+      final code = batch
+          .map((file) => {'content': file.content!, 'title': file.path})
+          .toList();
       final embeddings =
           await generationRepository.getCodeBatchEmbeddings(code);
       return embeddings;

diff --git a/commanddash/pubspec.lock b/commanddash/pubspec.lock
@@ -221,10 +221,10 @@ packages:
     dependency: "direct main"
     description:
       name: google_generative_ai
-      sha256: b2d3f7277a85e3e6be4c4392c59e73ea211b5b6c8bb21c24c71fd411a2d1822e
+      sha256: "55e1f13cd8c957be2e3644e0347471db1833c2ff3ff66ebef9fb71e9d9c9f24c"
       url: "https://pub.dev"
     source: hosted
-    version: "0.2.2"
+    version: "0.3.1"
   graphs:
     dependency: transitive
     description:

diff --git a/commanddash/pubspec.yaml b/commanddash/pubspec.yaml
@@ -18,7 +18,7 @@ dependencies:
   dio: ^5.4.1
   equatable: ^2.0.5
   crypto: ^3.0.3
-  google_generative_ai: ^0.2.2
+  google_generative_ai: ^0.3.1
   rxdart: ^0.27.7
   # path: ^1.8.0
 

diff --git a/commanddash/test/e2e/flutter_agent_test.dart b/commanddash/test/e2e/flutter_agent_test.dart
@@ -31,14 +31,14 @@ void main() {
         {
           "method": "agent-execute",
           "id": 1,
-          "version": "1.0.0",
           'params': {
-            "authdetails": {
+            "auth_details": {
               "type": "gemini",
               "key": EnvReader.get('GEMINI_KEY'),
               "githubToken": "authtoken",
             },
             "agent_name": "flutter",
+            "agent_version": "1.0.0",
             "slug": "/doc",
             "intent": "Your Flutter doc expert",
             "text_field_layout":
@@ -72,7 +72,7 @@ void main() {
               },
               {
                 "type": "prompt_query",
-                "query":
+                "prompt":
                     "You are an Flutter expert who answers user's queries related to the framework. \n\n Please find the user query <Query> and relavant references <References> picked from the Flutter docs to assist you: \n\n Query: <14340369>, \nReferences: <897806645>. Please respond to the user's query!",
                 "outputs": ["81443790"],
               },

diff --git a/commanddash/test/embedding_generator_test.dart b/commanddash/test/embedding_generator_test.dart
@@ -22,8 +22,9 @@ void main() {
             content: 'content of file $i');
       });
 
-      when(mockGenerationRepository
-              .getCodeBatchEmbeddings(files.map((e) => e.content!).toList()))
+      when(mockGenerationRepository.getCodeBatchEmbeddings(files
+              .map((file) => {'content': file.content!, 'title': file.path})
+              .toList()))
           .thenAnswer((_) async =>
               files.map((e) => getMockEmbeddingForString(e.content!)).toList());
 
@@ -41,14 +42,18 @@ void main() {
             content: 'content of file $i');
       });
 
-      when(mockGenerationRepository.getCodeBatchEmbeddings(
-              files.sublist(0, 100).map((e) => e.content!).toList()))
+      when(mockGenerationRepository.getCodeBatchEmbeddings(files
+              .sublist(0, 100)
+              .map((file) => {'content': file.content!, 'title': file.path})
+              .toList()))
           .thenAnswer((_) async => files
               .sublist(0, 100)
               .map((e) => getMockEmbeddingForString(e.content!))
               .toList());
-      when(mockGenerationRepository.getCodeBatchEmbeddings(
-              files.sublist(100).map((e) => e.content!).toList()))
+      when(mockGenerationRepository.getCodeBatchEmbeddings(files
+              .sublist(100)
+              .map((file) => {'content': file.content!, 'title': file.path})
+              .toList()))
           .thenAnswer((_) async => files
               .sublist(100)
               .map((e) => getMockEmbeddingForString(e.content!))
@@ -68,20 +73,26 @@ void main() {
             content: 'content of file $i');
       });
 
-      when(mockGenerationRepository.getCodeBatchEmbeddings(
-              files.sublist(0, 100).map((e) => e.content!).toList()))
+      when(mockGenerationRepository.getCodeBatchEmbeddings(files
+              .sublist(0, 100)
+              .map((file) => {'content': file.content!, 'title': file.path})
+              .toList()))
           .thenAnswer((_) async => files
               .sublist(0, 100)
               .map((e) => getMockEmbeddingForString(e.content!))
               .toList());
-      when(mockGenerationRepository.getCodeBatchEmbeddings(
-              files.sublist(100, 200).map((e) => e.content!).toList()))
+      when(mockGenerationRepository.getCodeBatchEmbeddings(files
+              .sublist(100, 200)
+              .map((file) => {'content': file.content!, 'title': file.path})
+              .toList()))
           .thenAnswer((_) async => files
               .sublist(100, 200)
               .map((e) => getMockEmbeddingForString(e.content!))
               .toList());
-      when(mockGenerationRepository.getCodeBatchEmbeddings(
-              files.sublist(200).map((e) => e.content!).toList()))
+      when(mockGenerationRepository.getCodeBatchEmbeddings(files
+              .sublist(200)
+              .map((file) => {'content': file.content!, 'title': file.path})
+              .toList()))
           .thenAnswer((_) async => files
               .sublist(200)
               .map((e) => getMockEmbeddingForString(e.content!))

diff --git a/commanddash/test/generation_repositories/gemini_repository_test.dart b/commanddash/test/generation_repositories/gemini_repository_test.dart
@@ -30,13 +30,16 @@ void main() {
       expect(
           () async => await geminiRepository.getCodeEmbeddings(testString),
           throwsA(isA<
-              UnknownException>())); // SDK gives format exception for incorrect key
+              InvalidApiKeyException>())); // SDK gives format exception for incorrect key
     });
 
     test('getCodeBatchEmbeddings [custom api implementation]', () async {
       final geminiRepository = GeminiRepository(apiKey!);
       final result = await geminiRepository.getCodeBatchEmbeddings(
-        ['Hello', 'World'],
+        [
+          {'content': 'Hello', 'title': 'Hello'},
+          {'content': 'World', 'title': 'World'}
+        ],
       );
       expect(result, isA<List<List<double>>>());
     });
@@ -45,14 +48,20 @@ void main() {
       final geminiRepository = GeminiRepository('AIzaSyCGXM9N6U9LkUoNou4KX-');
       expect(
           () async => await geminiRepository.getCodeBatchEmbeddings(
-                ['Hello', 'World'],
+                [
+                  {'content': 'Hello', 'title': 'Hello'},
+                  {'content': 'World', 'title': 'World'}
+                ],
               ),
           throwsA(isA<InvalidApiKeyException>()));
     });
     test('getStringBatchEmbeddings [custom api implementation]', () async {
       final geminiRepository = GeminiRepository(apiKey!);
       final result = await geminiRepository.getCodeBatchEmbeddings(
-        ['Hello', 'World'],
+        [
+          {'content': 'Hello', 'title': 'Hello'},
+          {'content': 'World', 'title': 'World'}
+        ],
       );
       expect(result, isA<List<List<double>>>());
     });
@@ -61,7 +70,10 @@ void main() {
       final geminiRepository = GeminiRepository('AIzaSyCGXM9N6U9LkUoNou4KX-');
       expect(
           () async => await geminiRepository.getCodeBatchEmbeddings(
-                ['Hello', 'World'],
+                [
+                  {'content': 'Hello', 'title': 'Hello'},
+                  {'content': 'World', 'title': 'World'}
+                ],
               ),
           throwsA(isA<InvalidApiKeyException>()));
     });

diff --git a/commanddash/test/server/task_assist_test.dart b/commanddash/test/server/task_assist_test.dart
@@ -112,6 +112,9 @@ void main() {
         "data": {"result": "success", "value": "unique_value"}
       }));
       result = await queue.next;
+      expect(result, isA<LogMessage>());
+      expect((result as LogMessage).message, 'response received');
+      result = await queue.next;
       expect(result, isA<ResultMessage>());
       expect(result.id, 1);
     });

diff --git a/commanddash/test/steps/append_to_chat_test.dart b/commanddash/test/steps/append_to_chat_test.dart
@@ -29,7 +29,9 @@ void main() {
       "method": "agent-execute",
       "id": 1,
       "params": {
-        "authdetails": {
+        "agent_name": "",
+        "agent_version": "1.0.0",
+        "auth_details": {
           "type": "gemini",
           "key": EnvReader.get('GEMINI_KEY'),
           "githubToken": ""
@@ -48,7 +50,7 @@ void main() {
         "steps": [
           {
             "type": "prompt_query",
-            "query": "736841542",
+            "prompt": "736841542",
             "post_process": {"type": "raw"},
             "outputs": ["90611917"]
           },

diff --git a/commanddash/test/steps/chat_step_test.dart b/commanddash/test/steps/chat_step_test.dart
@@ -34,7 +34,9 @@ void main() {
               "method": "agent-execute",
               "id": 1,
               "params": {
-                "authdetails": {
+                "agent_name": "",
+                "agent_version": "1.0.0",
+                "auth_details": {
                   "type": "gemini",
                   "key": EnvReader.get('GEMINI_KEY'),
                   "githubToken": ""
@@ -113,7 +115,9 @@ void main() {
               "method": "agent-execute",
               "id": 1,
               "params": {
-                "authdetails": {
+                "agent_name": "",
+                "agent_version": "1.0.0",
+                "auth_details": {
                   "type": "gemini",
                   "key": EnvReader.get('GEMINI_KEY'),
                   "githubToken": ""