From e9354d915c5131624a3a8c9522cdc87b0daeb37c Mon Sep 17 00:00:00 2001 From: Kenneth Reitz Date: Mon, 28 Oct 2024 07:57:03 -0400 Subject: [PATCH] test --- simplemind/integrations/__init__.py | 2 +- simplemind/integrations/anthropic.py | 41 ++++++++++++++++++++++++++++ simplemind/integrations/base.py | 35 ++++++++++++++++++------ simplemind/integrations/claude.py | 4 --- simplemind/integrations/openai.py | 28 +++++++++++++++---- simplemind/models.py | 12 +++++++- t.py | 16 ++++++++++- 7 files changed, 117 insertions(+), 21 deletions(-) create mode 100644 simplemind/integrations/anthropic.py delete mode 100644 simplemind/integrations/claude.py diff --git a/simplemind/integrations/__init__.py b/simplemind/integrations/__init__.py index d17c4bc..6942e5c 100644 --- a/simplemind/integrations/__init__.py +++ b/simplemind/integrations/__init__.py @@ -1,2 +1,2 @@ -from .claude import Claude +from .anthropic import Anthropic from .openai import OpenAI diff --git a/simplemind/integrations/anthropic.py b/simplemind/integrations/anthropic.py new file mode 100644 index 0000000..b96a033 --- /dev/null +++ b/simplemind/integrations/anthropic.py @@ -0,0 +1,41 @@ +import os + +import instructor +from anthropic import Anthropic as BaseAnthropic + +from .base import BaseClientProvider + + +class Anthropic(BaseClientProvider): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.login() + + def login(self): + """Initialize Anthropic client, with Instructor enabled.""" + + # Default to environment variable if not provided. + if self._api_key is None: + self._api_key = os.getenv("ANTHROPIC_API_KEY") + + base_client = BaseAnthropic(api_key=self._api_key) + self.client = instructor.from_anthropic(base_client) + # assert self.test_connection() + + @property + def available_models(self): + """Returns the available models from the Anthropic client.""" + + # TODO: scrape from website or embed + return [ + "claude-3-opus-20240229", + "claude-3-5-sonnet-20240620", + "claude-3-haiku-20240307", + "claude-3-5-sonnet-20240620", + "claude-3-5-sonnet-20240620", + ] + + # def test_connection(self): + # """Test the connection to Anthropic. Returns True if successful.""" + + # raise NotImplementedError("Anthropic test_connection not implemented.") diff --git a/simplemind/integrations/base.py b/simplemind/integrations/base.py index b8d5848..a0463fc 100644 --- a/simplemind/integrations/base.py +++ b/simplemind/integrations/base.py @@ -2,11 +2,16 @@ from pydantic import BaseModel + +DEFAULT_MODEL = "gpt-4o" + + class BaseClientProvider: - def __init__(self, *, api_key=None): + def __init__(self, *, model=DEFAULT_MODEL, api_key=None): # self.logger = logging.getLogger(self.__class__.__name__) self.client = None + self.model = model # Load API key from environment if not provided self._api_key = api_key @@ -23,11 +28,11 @@ class BaseClientProvider: msg = "This method must be implemented by the AI provider client." raise NotImplementedError(msg) - def generate_response(self, request): - """Generates a response from the AI provider client.""" + # def generate_response(self, request): + # """Generates a response from the AI provider client.""" - msg = "This method must be implemented by the AI provider client." - raise NotImplementedError(msg) + # msg = "This method must be implemented by the AI provider client." + # raise NotImplementedError(msg) def health_check(self): """Checks the health of the AI provider client.""" @@ -43,8 +48,20 @@ class BaseClientProvider: raise NotImplementedError(msg) - def features(self): - """Returns the features of the AI provider client.""" + # def features(self): + # """Returns the features of the AI provider client.""" - msg = "This method must be implemented by the AI provider client." - raise NotImplementedError(msg) + # msg = "This method must be implemented by the AI provider client." + # raise NotImplementedError(msg) + + # def structured_response(self, model, message, **kwargs): + # pass + + # def structured_conversation(self, model, message, **kwargs): + # pass + + # def single_message(self, model, message, **kwargs): + # return self.generate_response(message) + + # def start_conversation(self, model, message, **kwargs): + # pass diff --git a/simplemind/integrations/claude.py b/simplemind/integrations/claude.py deleted file mode 100644 index 55286da..0000000 --- a/simplemind/integrations/claude.py +++ /dev/null @@ -1,4 +0,0 @@ -from .base import BaseClientProvider - -class Claude(BaseClientProvider): - pass diff --git a/simplemind/integrations/openai.py b/simplemind/integrations/openai.py index 055cd3c..87d0af9 100644 --- a/simplemind/integrations/openai.py +++ b/simplemind/integrations/openai.py @@ -4,6 +4,7 @@ import instructor from openai import OpenAI as BaseOpenAI from .base import BaseClientProvider +from ..models import AIResponse class OpenAI(BaseClientProvider): @@ -19,10 +20,9 @@ class OpenAI(BaseClientProvider): if self._api_key is None: self._api_key = os.getenv("OPENAI_API_KEY") - base_client = BaseOpenAI(api_key=self._api_key) - self.client = instructor.from_openai(base_client) - self.test_connection() - + self.client = BaseOpenAI(api_key=self._api_key) + self.instructor_client = instructor.from_openai(self.client) + assert self.test_connection() @property def available_models(self): @@ -35,4 +35,22 @@ class OpenAI(BaseClientProvider): return [g for g in gen()] def test_connection(self): - pass + """Test the connection to OpenAI. Returns True if successful.""" + + return bool(len(self.available_models)) + + def single_message(self, message, *, response_model=False, **kwargs): + """Generates a response from the OpenAI client.""" + + client = self.client if not response_model else self.instructor_client + + completion = self.client.chat.completions.create( + model=self.model, + messages=[{"role": "user", "content": message}], + **kwargs, + ) + + return AIResponse( + response=completion, + text=completion.choices[0].message.content, + ) diff --git a/simplemind/models.py b/simplemind/models.py index 40ec8b1..6caa2f2 100644 --- a/simplemind/models.py +++ b/simplemind/models.py @@ -1,9 +1,19 @@ from pydantic import BaseModel +from typing import Any, ClassVar + class AIRequest(BaseModel): - prompt: str + text: str parameters: dict = {} + def __str__(self): + return self.text + + class AIResponse(BaseModel): + text: str response: Any metadata: dict = {} + + def __str__(self): + return self.text diff --git a/t.py b/t.py index adcb502..da96e2a 100644 --- a/t.py +++ b/t.py @@ -3,6 +3,20 @@ import simplemind context = None openai = simplemind.integrations.OpenAI() -openai.login() +print(openai.test_connection()) print(openai.available_models) + +print() +print() +message = "who is kennethreitz?" + +print(f"> {message}") +print(openai.single_message(message)) + +# claude = simplemind.integrations.Anthropic() + +# # print(claude.test_connection()) +# # print(claude.available_models) + +# claude.login()