mirror of
https://github.com/lfnovo/open-notebook.git
synced 2026-04-29 03:50:04 +00:00
69 lines
1.8 KiB
Python
69 lines
1.8 KiB
Python
from typing import Dict, Type, Union
|
|
|
|
from open_notebook.models.embedding_models import (
|
|
EmbeddingModel,
|
|
GeminiEmbeddingModel,
|
|
OllamaEmbeddingModel,
|
|
OpenAIEmbeddingModel,
|
|
VertexEmbeddingModel,
|
|
)
|
|
from open_notebook.models.llms import (
|
|
AnthropicLanguageModel,
|
|
GeminiLanguageModel,
|
|
LanguageModel,
|
|
LiteLLMLanguageModel,
|
|
OllamaLanguageModel,
|
|
OpenAILanguageModel,
|
|
OpenRouterLanguageModel,
|
|
VertexAILanguageModel,
|
|
VertexAnthropicLanguageModel,
|
|
)
|
|
from open_notebook.models.speech_to_text_models import (
|
|
OpenAISpeechToTextModel,
|
|
SpeechToTextModel,
|
|
)
|
|
from open_notebook.models.text_to_speech_models import (
|
|
ElevenLabsTextToSpeechModel,
|
|
OpenAITextToSpeechModel,
|
|
TextToSpeechModel,
|
|
)
|
|
|
|
ModelType = Union[LanguageModel, EmbeddingModel, SpeechToTextModel, TextToSpeechModel]
|
|
|
|
|
|
ProviderMap = Dict[str, Type[ModelType]]
|
|
|
|
MODEL_CLASS_MAP: Dict[str, ProviderMap] = {
|
|
"language": {
|
|
"ollama": OllamaLanguageModel,
|
|
"openrouter": OpenRouterLanguageModel,
|
|
"vertexai-anthropic": VertexAnthropicLanguageModel,
|
|
"litellm": LiteLLMLanguageModel,
|
|
"vertexai": VertexAILanguageModel,
|
|
"anthropic": AnthropicLanguageModel,
|
|
"openai": OpenAILanguageModel,
|
|
"gemini": GeminiLanguageModel,
|
|
},
|
|
"embedding": {
|
|
"openai": OpenAIEmbeddingModel,
|
|
"gemini": GeminiEmbeddingModel,
|
|
"vertexai": VertexEmbeddingModel,
|
|
"ollama": OllamaEmbeddingModel,
|
|
},
|
|
"speech_to_text": {
|
|
"openai": OpenAISpeechToTextModel,
|
|
},
|
|
"text_to_speech": {
|
|
"openai": OpenAITextToSpeechModel,
|
|
"elevenlabs": ElevenLabsTextToSpeechModel,
|
|
},
|
|
}
|
|
|
|
__all__ = [
|
|
"MODEL_CLASS_MAP",
|
|
"EmbeddingModel",
|
|
"LanguageModel",
|
|
"SpeechToTextModel",
|
|
"TextToSpeechModel",
|
|
"ModelType",
|
|
]
|