feat(umem_ai): add embedding API and provider support

[vidurkhanal]

• Introduce EmbeddingModel, EmbeddingRequest, and EmbeddingResponse types.
• Add Embeds async trait for embedding functionality.
• Implement Embeds for AmazonBedrockProvider with parallel, retryable embedding requests.
• Add do_embed async method to AIProvider.
• Add embed module with retry logic and request struct.
• Add futures as a workspace dependency.
• Provide Default impls for OpenAIProvider and CohereProvider.
• Extend ResponseGeneratorError with new error variants for embedding.
• Update module exports and imports accordingly.

[greptile-apps]

Greptile Summary

This PR adds embedding API support to the umem_ai crate with Amazon Bedrock as the initial provider implementation. The implementation introduces parallel, retryable embedding requests with proper concurrency control using semaphores.

Key Changes:

• Added EmbeddingModel, EmbeddingRequest, and EmbeddingResponse types for the embedding API
• Implemented Embeds trait with Amazon Bedrock provider support using parallel requests
• Added retry logic with exponential backoff for embedding operations
• Enhanced error handling to use .meta() for better error messages from AWS SDK
• Added Default implementations for OpenAIProvider and CohereProvider for easier instantiation
• Introduced new error variants: BedrockInvokeError and InternalServerError

Issues Found:

• Minor typo in error message ("makng" instead of "making") in amazon_bedrock.rs:477
• The custom_headers field in EmbeddingRequest is defined but never used in the implementation

Confidence Score: 4/5

• This PR is safe to merge with one minor typo correction needed
• The implementation is well-structured with proper error handling, concurrency control, and retry logic. The parallel embedding requests use semaphores correctly to limit concurrent operations. Error handling has been improved to use .meta() for better diagnostics. Only one syntax error (typo) was found that needs correction. The unused custom_headers field is a minor issue that doesn't affect functionality.
• Pay attention to crates/umem_ai/src/providers/amazon_bedrock.rs for the typo fix

Important Files Changed

Filename	Overview
crates/umem_ai/src/providers/amazon_bedrock.rs	added Embeds trait implementation with parallel embedding requests and semaphore-based concurrency control, changed error handling to use .meta() for better error messages, added default() method (minor typo in error message)
crates/umem_ai/src/response_generators/embed.rs	new file introducing embedding functionality with retry logic, timeout handling, and configurable request parameters (custom_headers field unused)
crates/umem_ai/src/lib.rs	added EmbeddingModel struct, Embeds trait, and do_embed method to AIProvider enum

Sequence Diagram

sequenceDiagram
    participant Client
    participant embed_fn as embed()
    participant AIProvider
    participant AmazonBedrockProvider
    participant Semaphore
    participant BedrockRuntime as AWS Bedrock API
    
    Client->>embed_fn: EmbeddingRequest
    embed_fn->>embed_fn: Setup retry logic<br/>(ExponentialBuilder)
    
    loop Retry up to max_retries
        embed_fn->>AIProvider: do_embed(request)
        AIProvider->>AmazonBedrockProvider: embed(request)
        
        AmazonBedrockProvider->>AmazonBedrockProvider: Validate input is not empty
        AmazonBedrockProvider->>Semaphore: Create with max_parallels
        
        par Parallel embedding requests
            loop For each input text
                AmazonBedrockProvider->>Semaphore: acquire_owned()
                Semaphore-->>AmazonBedrockProvider: permit
                AmazonBedrockProvider->>BedrockRuntime: invoke_model()<br/>(inputText, dimensions, normalize)
                BedrockRuntime-->>AmazonBedrockProvider: InvokeModelOutput
                AmazonBedrockProvider->>AmazonBedrockProvider: Drop permit
            end
        end
        
        AmazonBedrockProvider->>AmazonBedrockProvider: join_all(handles)
        AmazonBedrockProvider->>AmazonBedrockProvider: Deserialize responses
        AmazonBedrockProvider->>AmazonBedrockProvider: Extract embeddings
        AmazonBedrockProvider-->>AIProvider: EmbeddingResponse
        AIProvider-->>embed_fn: EmbeddingResponse
        
        alt Success
            embed_fn-->>Client: EmbeddingResponse
        else Retryable Error
            embed_fn->>embed_fn: Wait (exponential backoff)
        end
    end

Loading

[greptile-apps]

_{8 files reviewed, 1 comment}

_{Edit Code Review Agent Settings | Greptile}