diff --git a/pretor/core/global_state_machine/model_provider/__init__.py b/pretor/core/global_state_machine/model_provider/__init__.py index 39eb890..77aded4 100644 --- a/pretor/core/global_state_machine/model_provider/__init__.py +++ b/pretor/core/global_state_machine/model_provider/__init__.py @@ -16,4 +16,5 @@ from pretor.core.global_state_machine.model_provider.base_provider import Provid from pretor.core.global_state_machine.model_provider.openai_provider import OpenAIProvider from pretor.core.global_state_machine.model_provider.gemini_provider import GeminiProvider from pretor.core.global_state_machine.model_provider.claude_provider import ClaudeProvider -__all__ = ["Provider", "ProviderArgs", "OpenAIProvider", "GeminiProvider", "ClaudeProvider"] +from pretor.core.global_state_machine.model_provider.deepseek_provider import DeepseekProvider +__all__ = ["Provider", "ProviderArgs", "OpenAIProvider", "GeminiProvider", "ClaudeProvider", "DeepseekProvider"] diff --git a/pretor/core/global_state_machine/provider_manager.py b/pretor/core/global_state_machine/provider_manager.py index db4b5c6..fa5ef25 100644 --- a/pretor/core/global_state_machine/provider_manager.py +++ b/pretor/core/global_state_machine/provider_manager.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from pretor.core.global_state_machine.model_provider import Provider, OpenAIProvider,GeminiProvider, ClaudeProvider +from pretor.core.global_state_machine.model_provider import Provider, OpenAIProvider,GeminiProvider, ClaudeProvider, DeepseekProvider from typing import Dict, Type class ProviderManager: @@ -29,7 +29,8 @@ class ProviderManager: def __init__(self, postgres): self.provider_mapper = {"openai": OpenAIProvider, "gemini": GeminiProvider, - "claude": ClaudeProvider} + "claude": ClaudeProvider, + "deepseek": DeepseekProvider} self.provider_register = {} async def init_provider_register(self, postgres) -> None: