Skip to content

Commit

Permalink
refactors changes according to suggestions
Browse files Browse the repository at this point in the history
  • Loading branch information
gtang31-te authored Jan 31, 2025
1 parent 4d7e1f2 commit 3903914
Showing 1 changed file with 6 additions and 4 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -37,10 +37,12 @@ class BedrockModelProvider(Enum):
@classmethod
def to_model_provider(cls, model_id: str) -> "BedrockModelProvider":
"""Convert a model ID to a model provider."""
if model_id.startswith("us."):
model_id = model_id.removeprefix("us.")
provider = model_id.split(".")[0]
return cls(provider)
try:
return next(provider for provider in cls if provider.value in model_id)
except StopAsyncIteration:
raise ValueError(
f"Model ID {model_id} does not contain a valid model provider name."
)


# region Text Completion
Expand Down

0 comments on commit 3903914

Please sign in to comment.