diff --git a/simplemind/integrations/base.py b/simplemind/integrations/base.py index 6b227c6..b8d5848 100644 --- a/simplemind/integrations/base.py +++ b/simplemind/integrations/base.py @@ -1,5 +1,6 @@ # import logging +from pydantic import BaseModel class BaseClientProvider: @@ -34,6 +35,7 @@ class BaseClientProvider: msg = "This method must be implemented by the AI provider client." raise NotImplementedError(msg) + @property def available_models(self): """Returns the available models from the AI provider client.""" diff --git a/simplemind/integrations/openai.py b/simplemind/integrations/openai.py index 49c3dd9..055cd3c 100644 --- a/simplemind/integrations/openai.py +++ b/simplemind/integrations/openai.py @@ -1,7 +1,7 @@ import os import instructor -from openai import OpenAI +from openai import OpenAI as BaseOpenAI from .base import BaseClientProvider @@ -9,11 +9,9 @@ from .base import BaseClientProvider class OpenAI(BaseClientProvider): def __init__(self, *args, **kwargs): - return super().__init__(*args, **kwargs) + super().__init__(*args, **kwargs) self.login() - - def login(self): """Initialize OpenAI client, with Instructor enabled.""" @@ -21,19 +19,20 @@ class OpenAI(BaseClientProvider): if self._api_key is None: self._api_key = os.getenv("OPENAI_API_KEY") - base_client = OpenAI(api_key=self._api_key) + base_client = BaseOpenAI(api_key=self._api_key) self.client = instructor.from_openai(base_client) self.test_connection() + @property def available_models(self): - pass + """Returns the available models from the OpenAI client.""" + + def gen(): + for model in self.client.models.list(): + yield model.id + + return [g for g in gen()] def test_connection(self): - try: - # openai.api_key = self._api_key - self.client.models.list() - # self.logger.info("OpenAI connection test successful") - except Exception as e: - # self.logger.error(f"OpenAI connection test failed: {str(e)}") - raise + pass diff --git a/t.py b/t.py index 133af6d..adcb502 100644 --- a/t.py +++ b/t.py @@ -3,4 +3,6 @@ import simplemind context = None openai = simplemind.integrations.OpenAI() -print(openai.login()) +openai.login() + +print(openai.available_models)