Feature: Allow passing more parameters to cohere embed models (#36)

Author: rhysemmerson

src/Schemas/Cohere/CohereEmbeddingsHandler.php

@@ -3,6 +3,7 @@
 namespace Prism\Bedrock\Schemas\Cohere;
 
 use Illuminate\Http\Client\Response;
+use Illuminate\Support\Arr;
 use Prism\Bedrock\Contracts\BedrockEmbeddingsHandler;
 use Prism\Prism\Embeddings\Request;
 use Prism\Prism\Embeddings\Response as EmbeddingsResponse;
@@ -39,9 +40,15 @@ public static function buildPayload(Request $request): array
     {
         return array_filter([
             'texts' => $request->inputs(),
-            'input_type' => 'search_document', // TODO: Need to PR providerOptions onto embeddings request to allow override.
-            'truncate' => null, // TODO: Need to PR providerOptions onto embeddings request to allow override. Default for now.
-            'embedding_types' => null, // TODO: Need to PR providerOptions onto embeddings request to allow override. Default for now.
+            'input_type' => 'search_document',
+            'truncate' => null,
+            'embedding_types' => null,
+            ...Arr::only($request->providerOptions(), [
+                'input_type',
+                'embedding_types',
+                'truncate',
+                'output_dimension',
+            ]),
         ]);
     }
 

tests/Schemas/Cohere/CohereEmbeddingsTest.php

@@ -4,6 +4,7 @@
 
 namespace Tests\Schemas\Cohere;
 
+use Illuminate\Support\Facades\Http;
 use Prism\Prism\Prism;
 use Prism\Prism\ValueObjects\Embedding;
 use Tests\Fixtures\FixtureResponse;
@@ -63,3 +64,38 @@
     expect($response->embeddings[1]->embedding)->toEqual($embeddings[1]->embedding);
     expect($response->usage->tokens)->toBe(1);
 });
+
+it('can set request params', function (): void {
+    FixtureResponse::fakeResponseSequence('invoke', 'cohere/generate-embeddings-from-input', [
+        'X-Amzn-Bedrock-Input-Token-Count' => 4,
+    ]);
+
+    $response = Prism::embeddings()
+        ->using('bedrock', 'cohere.embed-english-v3')
+        ->withProviderOptions([
+            'input_type' => 'search_query',
+            'truncate' => 'RIGHT',
+            'embedding_types' => ['sparse', 'dense'],
+            'output_dimension' => 1536,
+            'some_other_option' => 'should be filtered out',
+        ])
+        ->fromInput('Hello, world!')
+        ->asEmbeddings();
+
+    $embeddings = json_decode(file_get_contents('tests/Fixtures/cohere/generate-embeddings-from-input-1.json'), true);
+    $embeddings = array_map(fn (array $item): Embedding => Embedding::fromArray($item), data_get($embeddings, 'embeddings'));
+
+    Http::assertSent(function ($request): bool {
+        $body = $request->data();
+
+        return $body['input_type'] === 'search_query'
+            && $body['truncate'] === 'RIGHT'
+            && $body['embedding_types'] === ['sparse', 'dense']
+            && $body['output_dimension'] === 1536
+            && ! array_key_exists('some_other_option', $body);
+    });
+
+    expect($response->embeddings)->toBeArray();
+    expect($response->embeddings[0]->embedding)->toEqual($embeddings[0]->embedding);
+    expect($response->usage->tokens)->toBe(4);
+});