| from pydantic_ai.models import Model, KnownModelName | |
| from knowlang.core.types import ModelProvider | |
| from typing import get_args | |
| def create_pydantic_model( | |
| model_provider: ModelProvider, | |
| model_name: str, | |
| ) -> Model | KnownModelName: | |
| model_str = f"{model_provider}:{model_name}" | |
| if model_str in get_args(KnownModelName): | |
| return model_str | |
| elif model_provider == ModelProvider.HUGGINGFACE: | |
| from knowlang.models.huggingface import HuggingFaceModel | |
| return HuggingFaceModel(model_name=model_name) | |
| elif model_provider == ModelProvider.TESTING: | |
| # should be used for testing purposes only | |
| pass | |
| else: | |
| raise NotImplementedError(f"Model {model_provider}:{model_name} is not supported") | |