use annotated union to parse

llama_stack/distribution/utils/memory_bank_utils.py

@@ -1,60 +0,0 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-# All rights reserved.
-#
-# This source code is licensed under the terms described in the LICENSE file in
-# the root directory of this source tree.
-
-
-from llama_stack.apis.memory_banks.memory_banks import (
-    BankParams,
-    GraphMemoryBank,
-    KeyValueMemoryBank,
-    KeywordMemoryBank,
-    MemoryBank,
-    MemoryBankType,
-    VectorMemoryBank,
-    VectorMemoryBankParams,
-)
-
-
-def build_memory_bank(
-    memory_bank_id: str,
-    params: BankParams,
-    provider_id: str,
-    provider_memorybank_id: str,
-) -> MemoryBank:
-    if params.memory_bank_type == MemoryBankType.vector.value:
-        assert isinstance(params, VectorMemoryBankParams)
-        memory_bank = VectorMemoryBank(
-            identifier=memory_bank_id,
-            provider_id=provider_id,
-            provider_resource_id=provider_memorybank_id,
-            memory_bank_type=params.memory_bank_type,
-            embedding_model=params.embedding_model,
-            chunk_size_in_tokens=params.chunk_size_in_tokens,
-            overlap_size_in_tokens=params.overlap_size_in_tokens,
-        )
-    elif params.memory_bank_type == MemoryBankType.keyvalue.value:
-        memory_bank = KeyValueMemoryBank(
-            identifier=memory_bank_id,
-            provider_id=provider_id,
-            provider_resource_id=provider_memorybank_id,
-            memory_bank_type=params.memory_bank_type,
-        )
-    elif params.memory_bank_type == MemoryBankType.keyword.value:
-        memory_bank = KeywordMemoryBank(
-            identifier=memory_bank_id,
-            provider_id=provider_id,
-            provider_resource_id=provider_memorybank_id,
-            memory_bank_type=params.memory_bank_type,
-        )
-    elif params.memory_bank_type == MemoryBankType.graph.value:
-        memory_bank = GraphMemoryBank(
-            identifier=memory_bank_id,
-            provider_id=provider_id,
-            provider_resource_id=provider_memorybank_id,
-            memory_bank_type=params.memory_bank_type,
-        )
-    else:
-        raise ValueError(f"Unknown memory bank type: {params.memory_bank_type}")
-    return memory_bank