diff --git a/examples/contextual_memory.py b/examples/contextual_memory.py index 5797213..2e4e542 100644 --- a/examples/contextual_memory.py +++ b/examples/contextual_memory.py @@ -8,7 +8,7 @@ import os import pickle -class ContextualMemoryPlugin: +class ContextualMemoryPlugin(sm.BasePlugin): def __init__( self, api_key: str, diff --git a/examples/generate_data.py b/examples/generate_data.py index 888d694..a682644 100644 --- a/examples/generate_data.py +++ b/examples/generate_data.py @@ -25,7 +25,7 @@ class QuotesList(BaseModel): quotes: List[MovieQuote] -def gen_quotes(n=10) -> Iterator[MovieQuote]: +def gen_quotes(n: int = 10) -> Iterator[MovieQuote]: """Generate a list of quotes from famous movies.""" for q in sm.generate_data( diff --git a/examples/math_plugin.py b/examples/math_plugin.py index 8637073..436d04f 100644 --- a/examples/math_plugin.py +++ b/examples/math_plugin.py @@ -1,9 +1,11 @@ from _context import sm -class MathPlugin: +class MathPlugin(sm.BasePlugin): def send_hook(self, conversation: sm.Conversation): last_user_message = conversation.get_last_message(role="user") + if last_user_message is None: + return if "calculate" in last_user_message.text.lower(): expression = last_user_message.text.lower().replace("calculate", "").strip() try: @@ -14,7 +16,7 @@ class MathPlugin: except Exception: conversation.add_message( role="assistant", - text="I'm sorry, I couldn't compute that expression.", + text="I'm sorry, I couldn't compute that expression. Please try again.", ) diff --git a/examples/two_llms_talking.py b/examples/two_llms_talking.py index 4971867..d098599 100644 --- a/examples/two_llms_talking.py +++ b/examples/two_llms_talking.py @@ -8,7 +8,7 @@ class ConversationPlugin(sm.BasePlugin): print(f"{conversation.llm_model}:\n{response.text.strip()}\n\n------------\n") -def have_conversation(rounds=3): +def have_conversation(rounds: int = 3): # Create two conversations - one for each AI with ( sm.create_conversation( diff --git a/simplemind/providers/anthropic.py b/simplemind/providers/anthropic.py index 77e4d20..bfb31c6 100644 --- a/simplemind/providers/anthropic.py +++ b/simplemind/providers/anthropic.py @@ -57,7 +57,7 @@ class Anthropic(BaseProvider): def structured_response(self, model: str, response_model, **kwargs): response = self.structured_client.messages.create( - model=model, response_model=response_model or self.DEFAULT_MODEL, **kwargs + model=model or self.DEFAULT_MODEL, response_model=response_model, **kwargs ) return response diff --git a/simplemind/providers/gemini.py b/simplemind/providers/gemini.py index f3bdb26..5bdd25d 100644 --- a/simplemind/providers/gemini.py +++ b/simplemind/providers/gemini.py @@ -23,7 +23,7 @@ class Gemini(BaseProvider): if not self.api_key: raise ValueError("Gemini API key is required") self.model_name = model_name - return genai.GenerativeModel(model_name=model_name) + return genai.GenerativeModel(model_name=self.model_name) @property def structured_client(self): @@ -42,9 +42,13 @@ class Gemini(BaseProvider): for msg in conversation.messages ] - response = self.structured_client.chat.completions.create( - messages=messages, response_model=None - ) + try: + response = self.structured_client.chat.completions.create( + messages=messages, response_model=None + ) + except Exception as e: + # Handle the exception appropriately, e.g., log the error or raise a custom exception + raise RuntimeError(f"Failed to send conversation to Gemini API: {e}") from e # Create and return a properly formatted Message instance return Message( @@ -57,15 +61,25 @@ class Gemini(BaseProvider): def structured_response(self, prompt: str, response_model, **kwargs): """Send a structured response to the Gemini API.""" - response = self.structured_client.chat.completions.create( - messages=[{"role": "user", "content": prompt}], - response_model=response_model, - ) + try: + response = self.structured_client.chat.completions.create( + messages=[{"role": "user", "content": prompt}], + response_model=response_model, + ) + except Exception as e: + # Handle the exception appropriately, e.g., log the error or raise a custom exception + raise RuntimeError( + f"Failed to send structured response to Gemini API: {e}" + ) from e return response def generate_text(self, prompt: str, **kwargs) -> str: """Generate text using the Gemini API.""" - response = self.structured_client.chat.completions.create( - messages=[{"role": "user", "content": prompt}], response_model=None - ) + try: + response = self.structured_client.chat.completions.create( + messages=[{"role": "user", "content": prompt}], response_model=None + ) + except Exception as e: + # Handle the exception appropriately, e.g., log the error or raise a custom exception + raise RuntimeError(f"Failed to generate text with Gemini API: {e}") from e return str(response) diff --git a/simplemind/providers/ollama.py b/simplemind/providers/ollama.py index 24a1ea7..d10a036 100644 --- a/simplemind/providers/ollama.py +++ b/simplemind/providers/ollama.py @@ -26,7 +26,7 @@ class Ollama(BaseProvider): return ol.Client(timeout=self.TIMEOUT, host=self.host_url) @property - def structured_client(self): + def structured_client(self) -> instructor.Instructor: """A client patched with Instructor.""" return instructor.from_openai( OpenAI( @@ -36,7 +36,7 @@ class Ollama(BaseProvider): mode=instructor.Mode.JSON, ) - def send_conversation(self, conversation: "Conversation"): + def send_conversation(self, conversation: "Conversation") -> "Message": """Send a conversation to the Ollama API.""" from ..models import Message @@ -81,4 +81,4 @@ class Ollama(BaseProvider): messages=messages, model=llm_model or self.DEFAULT_MODEL ) - return response.get("message").get("content") + return response.get("message", {}).get("content", "") diff --git a/simplemind/providers/openai.py b/simplemind/providers/openai.py index 2a296cc..549f28a 100644 --- a/simplemind/providers/openai.py +++ b/simplemind/providers/openai.py @@ -52,7 +52,7 @@ class OpenAI(BaseProvider): ) def structured_response( - self, prompt: str, response_model, *, llm_model: str, **kwargs + self, prompt: str, response_model, *, llm_model: str | None = None, **kwargs ): # Ensure messages are provided in kwargs messages = [ @@ -67,7 +67,7 @@ class OpenAI(BaseProvider): ) return response - def generate_text(self, prompt: str, *, llm_model: str, **kwargs): + def generate_text(self, prompt: str, *, llm_model: str | None = None, **kwargs): messages = [ {"role": "user", "content": prompt}, ] diff --git a/test.py b/test.py new file mode 100644 index 0000000..0b710d6 --- /dev/null +++ b/test.py @@ -0,0 +1,9 @@ +import simplemind as sm + +res = sm.generate_text( + prompt="Wish you a happy Diwali!", + llm_model="gpt-4o", + llm_provider="openai", +) + +print(res)