feat(rerank): add Bedrock reranking support

.claude/settings.local.json

@@ -0,0 +1,7 @@
+{
+  "permissions": {
+    "allow": [
+      "Bash(go build *)"
+    ]
+  }
+}

README.md

@@ -7,6 +7,7 @@ A comprehensive AWS Bedrock plugin for Genkit Go that provides text generation,
 - **Text Generation**: Support for multiple foundation models via AWS Bedrock Converse API
 - **Image Generation**: Support for image generation models like Amazon Titan Image Generator
 - **Embeddings**: Support for text embedding models from Amazon Titan and Cohere
+- **Reranking**: Support for Cohere Rerank models through Bedrock `InvokeModel`
 - **Streaming**: Full streaming support for real-time responses
 - **Tool Calling**: Complete function calling capabilities with schema validation and type conversion
 - **Multimodal Support**: Support for text + image inputs (vision models)
@@ -30,6 +31,9 @@ A comprehensive AWS Bedrock plugin for Genkit Go that provides text generation,
 - **Amazon Titan Embeddings**: Text v1/v2, Multimodal v1
 - **Cohere**: Embed English/Multilingual v3
 
+### Reranking Models
+- **Cohere**: Rerank v3.5
+
 ### Multimodal Models (Text + Vision)
 - All Claude 3/3.5/4 models
 - Amazon Nova models 
@@ -137,6 +141,41 @@ func main() {
 }
 ```
 
+## Rerank Documents
+
+Genkit Go does not yet expose a first-class reranker action, so the plugin provides a standalone `Rerank` helper. It reuses the initialized Bedrock plugin configuration and returns documents ordered by relevance with scores attached to `ai.RankedDocumentMetadata`.
+
+```go
+resp, err := bedrock.Rerank(ctx, g, "cohere.rerank-v3-5:0", &ai.RerankerRequest{
+    Query: ai.DocumentFromText("Which document explains Bedrock authentication?", nil),
+    Documents: []*ai.Document{
+        ai.DocumentFromText("Configure AWS credentials with environment variables or AWS SSO.", nil),
+        ai.DocumentFromText("Titan Image Generator returns base64-encoded PNG data.", nil),
+    },
+    Options: &bedrock.RerankOptions{TopN: 1},
+})
+if err != nil {
+    log.Fatal(err)
+}
+
+for _, doc := range resp.Documents {
+    log.Printf("score=%.3f text=%s", doc.Metadata.Score, doc.Content[0].Text)
+}
+```
+
+Run the reranking example with:
+
+```bash
+cd examples/reranking
+go run main.go
+```
+
+From the repository root, run the optional live reranking test with:
+
+```bash
+go test . -run TestBedrockLive_CohereRerank -test-bedrock-rerank-model=cohere.rerank-v3-5:0
+```
+
 ## Configuration Options
 
 The plugin supports various configuration options:
@@ -222,6 +261,7 @@ The repository includes comprehensive examples:
 - **`examples/tool_calling/`** - Function calling with multiple tools
 - **`examples/image_generation/`** - Image generation and file saving
 - **`examples/embeddings/`** - Text embeddings and similarity
+- **`examples/reranking/`** - Cohere Rerank through Bedrock
 - **`examples/multimodal/`** - Vision models with image inputs
 - **`examples/advanced_schemas/`** - Complex tool schemas and validation
 - **`examples/prompt_caching`** - Several calls with prompt caching enabled to save costs
@@ -244,6 +284,10 @@ go run main.go
 # Run image generation example
 cd ../image_generation
 go run main.go
+
+# Run reranking example
+cd ../reranking
+go run main.go
 ```
 
 ## Features in Detail

bedrock.go

@@ -60,9 +60,9 @@ func (b *Bedrock) Name() string {
 // This method follows the same pattern as the Ollama plugin.
 func (b *Bedrock) Init(ctx context.Context) []api.Action {
 	b.mu.Lock()
+	defer b.mu.Unlock()
 
 	if b.initted {
-		b.mu.Unlock()
 		panic("bedrock: Init already called")
 	}
 
@@ -99,13 +99,23 @@ func (b *Bedrock) Init(ctx context.Context) []api.Action {
 
 	b.initted = true
 
-	// Release the mutex
-	b.mu.Unlock()
-
-	// Don't defer unlock since we already unlocked manually
 	return []api.Action{}
 }
 
+func (b *Bedrock) withRequestTimeout(ctx context.Context) (context.Context, context.CancelFunc) {
+	if b == nil {
+		return ctx, func() {}
+	}
+	return withRequestTimeout(ctx, b.RequestTimeout)
+}
+
+func withRequestTimeout(ctx context.Context, timeout time.Duration) (context.Context, context.CancelFunc) {
+	if timeout <= 0 {
+		return ctx, func() {}
+	}
+	return context.WithTimeout(ctx, timeout)
+}
+
 // DefineModel defines a model in the registry.
 // This follows the same pattern as the Anthropic plugin's DefineModel method.
 func (b *Bedrock) DefineModel(g *genkit.Genkit, model ModelDefinition, info *ai.ModelInfo) ai.Model {

embed.go

@@ -98,6 +98,9 @@ func (b *Bedrock) getTitanEmbedding(ctx context.Context, modelName, text string)
 		Accept:      aws.String("application/json"),
 	}
 
+	ctx, cancel := b.withRequestTimeout(ctx)
+	defer cancel()
+
 	response, err := b.client.InvokeModel(ctx, input)
 	if err != nil {
 		return nil, fmt.Errorf("failed to invoke model: %w", err)
@@ -137,6 +140,9 @@ func (b *Bedrock) getCohereEmbedding(ctx context.Context, modelName, text string
 		Accept:      aws.String("application/json"),
 	}
 
+	ctx, cancel := b.withRequestTimeout(ctx)
+	defer cancel()
+
 	response, err := b.client.InvokeModel(ctx, input)
 	if err != nil {
 		return nil, fmt.Errorf("failed to invoke model: %w", err)

examples/reranking/main.go

@@ -0,0 +1,69 @@
+// Copyright 2025 Xavier Portilla Edo
+// Copyright 2025 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+// Package main demonstrates document reranking with AWS Bedrock.
+package main
+
+import (
+	"context"
+	"log"
+
+	"github.com/firebase/genkit/go/ai"
+	"github.com/firebase/genkit/go/genkit"
+	bedrock "github.com/xavidop/genkit-aws-bedrock-go"
+)
+
+func main() {
+	ctx := context.Background()
+
+	bedrockPlugin := &bedrock.Bedrock{
+		Region: "us-east-1",
+	}
+
+	g := genkit.Init(ctx,
+		genkit.WithPlugins(bedrockPlugin),
+	)
+
+	log.Println("Starting reranking example...")
+
+	response, err := bedrock.Rerank(ctx, g, "cohere.rerank-v3-5:0", &ai.RerankerRequest{
+		Query: ai.DocumentFromText("How do I configure authentication for AWS Bedrock?", nil),
+		Documents: []*ai.Document{
+			ai.DocumentFromText("Configure AWS credentials with environment variables, shared credentials files, IAM roles, or AWS SSO.", map[string]any{"id": "auth"}),
+			ai.DocumentFromText("Amazon Titan Image Generator returns generated images as base64-encoded image data.", map[string]any{"id": "image"}),
+			ai.DocumentFromText("Bedrock model access is managed in the AWS Bedrock console for each supported region.", map[string]any{"id": "model-access"}),
+		},
+		Options: &bedrock.RerankOptions{TopN: 2},
+	})
+	if err != nil {
+		log.Fatalf("Error reranking documents: %v", err)
+	}
+
+	for i, doc := range response.Documents {
+		score := 0.0
+		if doc.Metadata != nil {
+			score = doc.Metadata.Score
+		}
+		text := ""
+		if len(doc.Content) > 0 && doc.Content[0] != nil {
+			text = doc.Content[0].Text
+		}
+		log.Printf("Rank %d score=%.4f text=%s", i+1, score, text)
+	}
+
+	log.Println("Reranking example completed")
+}

generate.go

@@ -339,6 +339,9 @@ func (b *Bedrock) buildConverseInput(modelName string, input *ai.ModelRequest) (
 
 // generateTextSync handles synchronous text generation
 func (b *Bedrock) generateTextSync(ctx context.Context, input *bedrockruntime.ConverseInput, originalInput *ai.ModelRequest) (*ai.ModelResponse, error) {
+	ctx, cancel := b.withRequestTimeout(ctx)
+	defer cancel()
+
 	// Call Bedrock Converse API
 	response, err := b.client.Converse(ctx, input)
 	if err != nil {

image.go

@@ -103,6 +103,9 @@ func (b *Bedrock) generateTitanImage(ctx context.Context, modelName, prompt stri
 		Accept:      aws.String("application/json"),
 	}
 
+	ctx, cancel := b.withRequestTimeout(ctx)
+	defer cancel()
+
 	response, err := b.client.InvokeModel(ctx, input)
 	if err != nil {
 		return nil, fmt.Errorf("failed to invoke model: %w", err)
@@ -174,6 +177,9 @@ func (b *Bedrock) generateStableDiffusionImage(ctx context.Context, modelName, p
 		Accept:      aws.String("application/json"),
 	}
 
+	ctx, cancel := b.withRequestTimeout(ctx)
+	defer cancel()
+
 	response, err := b.client.InvokeModel(ctx, input)
 	if err != nil {
 		return nil, fmt.Errorf("failed to invoke model: %w", err)
@@ -251,6 +257,9 @@ func (b *Bedrock) generateNovaCanvasImage(ctx context.Context, modelName, prompt
 		Accept:      aws.String("application/json"),
 	}
 
+	ctx, cancel := b.withRequestTimeout(ctx)
+	defer cancel()
+
 	response, err := b.client.InvokeModel(ctx, input)
 	if err != nil {
 		return nil, fmt.Errorf("failed to invoke model: %w", err)