chore: switch to rerank in umem_ai

[Sang-it]

No description provided.

[greptile-apps]

Greptile Summary

This PR consolidates reranking functionality from the standalone umem_rerank crate into the umem_ai module, unifying AI operations under a single interface. The refactoring removes provider-specific abstractions (Pinecone, Cohere) in favor of a generic RerankingModel that uses the existing Provider enum.

Major changes:

• Deleted entire umem_rerank crate including Pinecone and Cohere implementations
• Added RerankingModel struct to umem_ai with singleton pattern via OnceCell
• Replaced trait-based RerankerBase with direct rerank() function using builder pattern
• Updated MemoryController to use Arc<RerankingModel> instead of trait object
• Renamed Embedder config enum to EmbeddingModel for naming consistency
• Commented out all code in seed_db.rs

Critical issue found:

• RerankingModel::get_model() accepts both OpenAI and AmazonBedrock providers from config, but AIProvider::do_reranking() only implements Cohere and AmazonBedrock, causing runtime panic if OpenAI is configured for reranking

Cleanup needed:

• Unused Pinecone and Cohere config structs remain in umem_config

Confidence Score: 1/5

• This PR contains a critical runtime bug that will crash when OpenAI is used as reranking provider
• The code has a severe logic error in crates/umem_ai/src/lib.rs where RerankingModel::get_model() allows initialization with OpenAI provider, but AIProvider::do_reranking() will panic with unimplemented!() for OpenAI. This will cause runtime crashes for any configuration using OpenAI for reranking. Only AmazonBedrock is safe to use (Cohere requires CohereProvider initialization which isn't implemented in get_model).
• Pay close attention to crates/umem_ai/src/lib.rs - the RerankingModel::get_model() function must validate that only supported providers (AmazonBedrock) are used for reranking

Important Files Changed

Filename	Overview
crates/umem_ai/src/lib.rs	Added RerankingModel with singleton pattern, but allows OpenAI provider which doesn't support reranking (will panic at runtime)
crates/umem_config/src/lib.rs	Replaced Reranker enum with RerankingModel struct using generic Provider, leaving unused Pinecone and Cohere structs
crates/umem_controller/src/search_memory.rs	Replaced old reranker trait calls with new unified rerank() function using RerankRequest builder pattern
crates/umem_controller/src/lib.rs	Updated MemoryController to use RerankingModel instead of RerankerBase trait
crates/umem_memory_machine/src/lib.rs	Updated service initialization to use RerankingModel::get_model() instead of Reranker::get_reranker()

Sequence Diagram

sequenceDiagram
    participant Client
    participant MemoryController
    participant Embedder
    participant VectorStore
    participant RerankingModel
    participant AIProvider
    participant RerankAPI

    Client->>MemoryController: search_with_context(query)
    MemoryController->>Embedder: generate_embedding(query)
    Embedder-->>MemoryController: vector
    MemoryController->>VectorStore: search(vector_query, limit=20)
    VectorStore-->>MemoryController: memories
    MemoryController->>MemoryController: extract summaries from memories
    MemoryController->>RerankingModel: get_model()
    RerankingModel-->>MemoryController: Arc<RerankingModel>
    MemoryController->>MemoryController: build RerankRequest
    MemoryController->>AIProvider: do_reranking(request)
    alt Cohere Provider
        AIProvider->>RerankAPI: POST /rerank
        RerankAPI-->>AIProvider: rankings with scores
    else AmazonBedrock Provider
        AIProvider->>RerankAPI: invoke rerank model
        RerankAPI-->>AIProvider: rankings with scores
    else OpenAI Provider (BUG)
        AIProvider-->>AIProvider: panic! unimplemented!()
    end
    AIProvider-->>MemoryController: RerankResponse
    MemoryController->>MemoryController: reorder memories by rankings
    MemoryController-->>Client: top-k reranked memories

Loading

[greptile-apps]

Additional Comments (1)

1. crates/umem_config/src/lib.rs, line 94-104 (link)

   style: Pinecone and Cohere structs are unused after removing the Reranker enum

_{16 files reviewed, 2 comments}

_{Edit Code Review Agent Settings | Greptile}