diff --git a/CHANGELOG.md b/CHANGELOG.md index f0e88da..48cc0d8 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,12 @@ Release History =============== +## 0.1.7 (2024-11-01) + +- Add `logger` decorator. +- Add `sm.enable_logfire()` function. +- General improvements. + ## 0.1.6 (2024-10-31) - Add `sm.Plugin` syntax sugar. diff --git a/README.md b/README.md index cedcd16..6210490 100644 --- a/README.md +++ b/README.md @@ -91,6 +91,33 @@ class Poem(BaseModel): title='Eternal Embrace' content='In the quiet hours of the night,\nWhen stars whisper secrets bright,\nTwo hearts beat in a gentle rhyme,\nDancing through the sands of time.\n\nWith every glance, a spark ignites,\nA flame that warms the coldest nights,\nIn laughter shared and whispers sweet,\nLove paints the world, a masterpiece.\n\nThrough stormy skies and sunlit days,\nIn myriad forms, it finds its ways,\nA tender touch, a knowing sigh,\nIn love’s embrace, we learn to fly.\n\nAs seasons change and moments fade,\nIn the tapestry of dreams we’ve laid,\nLove’s threads endure, forever bind,\nA timeless bond, two souls aligned.\n\nSo here’s to love, both bright and true,\nA gift we give, anew, anew,\nIn every heartbeat, every prayer,\nA story written in the air.' ``` +#### A more complex example + +```python +class InstructionStep(BaseModel): + step_number: int + instruction: str + +class RecipeIngredient(BaseModel): + name: str + quantity: float + unit: str + +class Recipe(BaseModel): + name: str + ingredients: list[RecipeIngredient] + instructions: list[InstructionStep] + +recipe = sm.generate_data( + "Write a recipe for chocolate chip cookies", + llm_model="gpt-4o-mini", + llm_provider="openai", + response_model=Recipe, +) +``` + +Special thanks to [@jxnl](https://github.com/jxnl) for building [Instructor](https://github.com/jxnl/instructor), which makes this possible! + ### Conversational AI SimpleMind also allows for easy conversational flows: @@ -200,6 +227,12 @@ The universe is never done. Simple, yet effective. +### Logging + +Simplemind uses [logfire](https://logfire.ai) for logging. To enable logging, call `sm.enable_logfire()`. + +### More Examples + Please see the [examples](examples) directory for executable examples. --- diff --git a/docs/conf.py b/docs/conf.py index 30c9fd6..a123700 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -16,7 +16,7 @@ import simplemind project = "simplemind" copyright = "2024 Kenneth Reitz" author = "Kenneth Reitz" -release = "v0.1.6" +release = "v0.1.7" # -- General configuration --------------------------------------------------- # https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration diff --git a/pyproject.toml b/pyproject.toml index 6da1744..b288937 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,10 +1,10 @@ [project] name = "simplemind" -version = "0.1.6" +version = "0.1.7" description = "An experimental client for AI providers that intends to replace LangChain and LangGraph for most common use cases." readme = "README.md" requires-python = ">=3.10" -dependencies = ["pydantic", "pydantic-settings", "instructor", "openai", "anthropic", "ollama", "groq", "google-generativeai"] +dependencies = ["pydantic", "pydantic-settings", "instructor", "openai", "anthropic", "ollama", "groq", "google-generativeai", "logfire"] [build-system] requires = ["hatchling"] diff --git a/simplemind/__init__.py b/simplemind/__init__.py index cf725c8..b9e4d1a 100644 --- a/simplemind/__init__.py +++ b/simplemind/__init__.py @@ -113,6 +113,11 @@ def generate_text( return provider.generate_text(prompt=prompt, llm_model=llm_model, **kwargs) +def enable_logfire() -> None: + """Enable logfire logging.""" + settings.logging.enable_logfire() + + # Syntax sugar. Plugin = BasePlugin @@ -125,4 +130,5 @@ __all__ = [ "BasePlugin", "Session", "Plugin", + "enable_logfire", ] diff --git a/simplemind/logging.py b/simplemind/logging.py new file mode 100644 index 0000000..06defc3 --- /dev/null +++ b/simplemind/logging.py @@ -0,0 +1,33 @@ +import time +from typing import Any, Callable + +import logfire + +from .settings import settings + + +def logger(func: Callable[..., Any]) -> Callable[..., Any]: + """A decorator that logs the function parameters, function returns, + and exceptions raised if logging is enabled, using logfire. + """ + + def wrapper(*args, **kwargs) -> Any: + if not settings.logging.is_enabled: + return func(*args, **kwargs) + + logfire.info(f"Calling {func.__name__} with args: {args}, kwargs: {kwargs}") + t1 = time.perf_counter() + + try: + result = func(*args, **kwargs) + t2 = time.perf_counter() + logfire.info(f"{func.__name__} returned: {result} in {t2-t1} seconds") + + return result + + except Exception as e: + t2 = time.perf_counter() + logfire.error(f"Error in {func.__name__}: {e} in {t2-t1} seconds") + raise e + + return wrapper diff --git a/simplemind/models.py b/simplemind/models.py index 95df38c..54eade1 100644 --- a/simplemind/models.py +++ b/simplemind/models.py @@ -1,6 +1,6 @@ -from types import TracebackType import uuid from datetime import datetime +from types import TracebackType from typing import Any, Dict, List, Literal, Optional from pydantic import BaseModel, Field diff --git a/simplemind/providers/_base.py b/simplemind/providers/_base.py index 42ffb9b..4485246 100644 --- a/simplemind/providers/_base.py +++ b/simplemind/providers/_base.py @@ -1,10 +1,13 @@ from abc import ABC, abstractmethod from functools import cached_property -from typing import Any, Type, TypeVar +from typing import TYPE_CHECKING, Any, Type, TypeVar from instructor import Instructor from pydantic import BaseModel +if TYPE_CHECKING: + from ..models import Conversation, Message + T = TypeVar("T", bound=BaseModel) diff --git a/simplemind/providers/anthropic.py b/simplemind/providers/anthropic.py index 5c76160..4798933 100644 --- a/simplemind/providers/anthropic.py +++ b/simplemind/providers/anthropic.py @@ -1,13 +1,16 @@ from functools import cached_property -from typing import Type, TypeVar +from typing import TYPE_CHECKING, Type, TypeVar -import anthropic import instructor from pydantic import BaseModel +from ..logging import logger from ..settings import settings from ._base import BaseProvider +if TYPE_CHECKING: + from ..models import Conversation, Message + T = TypeVar("T", bound=BaseModel) @@ -30,6 +33,13 @@ class Anthropic(BaseProvider): """The raw Anthropic client.""" if not self.api_key: raise ValueError("Anthropic API key is required") + try: + import anthropic + except ImportError as exc: + raise ImportError( + "Please install the `anthropic` package: `pip install anthropic`" + ) from exc + return anthropic.Anthropic(api_key=self.api_key) @cached_property @@ -37,7 +47,8 @@ class Anthropic(BaseProvider): """A client patched with Instructor.""" return instructor.from_anthropic(self.client) - def send_conversation(self, conversation: "Conversation", **kwargs): + @logger + def send_conversation(self, conversation: "Conversation", **kwargs) -> "Message": """Send a conversation to the Anthropic API.""" from ..models import Message @@ -63,6 +74,7 @@ class Anthropic(BaseProvider): llm_provider=PROVIDER_NAME, ) + @logger def structured_response( self, response_model: Type[T], *, llm_model: str | None = None, **kwargs ) -> T: @@ -80,8 +92,9 @@ class Anthropic(BaseProvider): response_model=response_model, **{**self.DEFAULT_KWARGS, **kwargs}, ) - return response + return response_model.model_validate(response) + @logger def generate_text(self, prompt: str, *, llm_model: str, **kwargs): messages = [ {"role": "user", "content": prompt}, diff --git a/simplemind/providers/gemini.py b/simplemind/providers/gemini.py index c902141..aa08db8 100644 --- a/simplemind/providers/gemini.py +++ b/simplemind/providers/gemini.py @@ -2,21 +2,25 @@ # IT is not currently working as desired. from functools import cached_property -from typing import Type, TypeVar +from typing import TYPE_CHECKING, Type, TypeVar -import google.generativeai as genai import instructor from pydantic import BaseModel +from ..logging import logger from ..settings import settings from ._base import BaseProvider -PROVIDER_NAME = "gemini" -DEFAULT_MODEL = "models/gemini-1.5-flash-latest" +if TYPE_CHECKING: + from ..models import Conversation, Message T = TypeVar("T", bound=BaseModel) +PROVIDER_NAME = "gemini" +DEFAULT_MODEL = "models/gemini-1.5-flash-latest" + + class Gemini(BaseProvider): NAME = PROVIDER_NAME DEFAULT_MODEL = DEFAULT_MODEL @@ -25,12 +29,21 @@ class Gemini(BaseProvider): self.api_key = api_key or settings.get_api_key(PROVIDER_NAME) self.model_name = DEFAULT_MODEL + def set_model(self, model_name: str): + self.model_name = model_name + @cached_property - def client(self, model_name: str = DEFAULT_MODEL): + def client(self): """The raw Gemini client.""" if not self.api_key: raise ValueError("Gemini API key is required") - self.model_name = model_name + try: + import google.generativeai as genai + except ImportError as exc: + raise ImportError( + "Please install the `google-generativeai` package: `pip install google-generativeai`" + ) from exc + genai.configure(api_key=self.api_key) return genai.GenerativeModel(model_name=self.model_name) @cached_property @@ -38,6 +51,7 @@ class Gemini(BaseProvider): """A Gemini client patched with Instructor.""" return instructor.from_gemini(self.client) + @logger def send_conversation(self, conversation: "Conversation") -> "Message": """Send a conversation to the Gemini API.""" from ..models import Message @@ -64,9 +78,11 @@ class Gemini(BaseProvider): llm_provider=PROVIDER_NAME, ) + @logger def structured_response(self, prompt: str, response_model: Type[T], **kwargs) -> T: """Send a structured response to the Gemini API.""" - llm_model = kwargs.pop("llm_model", self.model_name) + # Only try to pop if the key exists + kwargs.pop("llm_model", None) # Add default value of None try: response = self.structured_client.chat.completions.create( @@ -79,12 +95,12 @@ class Gemini(BaseProvider): raise RuntimeError( f"Failed to send structured response to Gemini API: {e}" ) from e - return response + return response_model.model_validate(response) + @logger def generate_text(self, prompt: str, **kwargs) -> str: """Generate text using the Gemini API.""" - llm_model = kwargs.pop("llm_model", self.model_name) - + kwargs.pop("llm_model") try: response = self.client.generate_content(prompt, **kwargs) except Exception as e: diff --git a/simplemind/providers/groq.py b/simplemind/providers/groq.py index 42d8090..5b2801f 100644 --- a/simplemind/providers/groq.py +++ b/simplemind/providers/groq.py @@ -1,22 +1,29 @@ from functools import cached_property -from typing import Type, TypeVar +from typing import TYPE_CHECKING, Type, TypeVar -import groq import instructor from pydantic import BaseModel +from ..logging import logger from ..settings import settings from ._base import BaseProvider -PROVIDER_NAME = "groq" -DEFAULT_MODEL = "llama3-8b-8192" +if TYPE_CHECKING: + from ..models import Conversation, Message T = TypeVar("T", bound=BaseModel) +PROVIDER_NAME = "groq" +DEFAULT_MODEL = "llama3-8b-8192" +DEFAULT_MAX_TOKENS = 1_000 +DEFAULT_KWARGS = {"max_tokens": DEFAULT_MAX_TOKENS} + + class Groq(BaseProvider): NAME = PROVIDER_NAME DEFAULT_MODEL = DEFAULT_MODEL + DEFAULT_KWARGS = DEFAULT_KWARGS def __init__(self, api_key: str | None = None): self.api_key = api_key or settings.get_api_key(PROVIDER_NAME) @@ -26,6 +33,12 @@ class Groq(BaseProvider): """The raw Groq client.""" if not self.api_key: raise ValueError("Groq API key is required") + try: + import groq + except ImportError as exc: + raise ImportError( + "Please install the `groq` package: `pip install groq`" + ) from exc return groq.Groq(api_key=self.api_key) @cached_property @@ -33,6 +46,7 @@ class Groq(BaseProvider): """A client patched with Instructor.""" return instructor.from_groq(self.client) + @logger def send_conversation( self, conversation: "Conversation", @@ -48,7 +62,7 @@ class Groq(BaseProvider): response = self.client.chat.completions.create( model=conversation.llm_model or self.DEFAULT_MODEL, messages=messages, - **kwargs, + **{**self.DEFAULT_KWARGS, **kwargs}, ) # Get the response content from the Groq response @@ -63,6 +77,7 @@ class Groq(BaseProvider): llm_provider=PROVIDER_NAME, ) + @logger def structured_response(self, prompt: str, response_model: Type[T], **kwargs) -> T: # Ensure messages are provided in kwargs messages = [ @@ -73,17 +88,18 @@ class Groq(BaseProvider): messages=messages, response_model=response_model, model=kwargs.pop("llm_model", self.DEFAULT_MODEL), - **kwargs, + **{**self.DEFAULT_KWARGS, **kwargs}, ) - return response + return response_model.model_validate(response) + @logger def generate_text( self, prompt: str, *, llm_model: str, **kwargs, - ): + ) -> str: messages = [ {"role": "user", "content": prompt}, ] @@ -91,7 +107,7 @@ class Groq(BaseProvider): response = self.client.chat.completions.create( messages=messages, model=llm_model or self.DEFAULT_MODEL, - **kwargs, + **{**self.DEFAULT_KWARGS, **kwargs}, ) - return response.choices[0].message.content + return str(response.choices[0].message.content) diff --git a/simplemind/providers/ollama.py b/simplemind/providers/ollama.py index 12f241c..3e00c25 100644 --- a/simplemind/providers/ollama.py +++ b/simplemind/providers/ollama.py @@ -1,25 +1,30 @@ from functools import cached_property -from typing import Type, TypeVar +from typing import TYPE_CHECKING, Type, TypeVar import instructor -import ollama as ol from openai import OpenAI from pydantic import BaseModel +from ..logging import logger from ..settings import settings from ._base import BaseProvider +if TYPE_CHECKING: + from ..models import Conversation, Message + T = TypeVar("T", bound=BaseModel) PROVIDER_NAME = "ollama" DEFAULT_MODEL = "llama3.2" DEFAULT_TIMEOUT = 60 +DEFAULT_KWARGS = {} class Ollama(BaseProvider): NAME = PROVIDER_NAME DEFAULT_MODEL = DEFAULT_MODEL + DEFAULT_KWARGS = DEFAULT_KWARGS TIMEOUT = DEFAULT_TIMEOUT def __init__(self, host_url: str | None = None): @@ -30,6 +35,12 @@ class Ollama(BaseProvider): """The raw Ollama client.""" if not self.host_url: raise ValueError("No ollama host url provided") + try: + import ollama as ol + except ImportError as exc: + raise ImportError( + "Please install the `ollama` package: `pip install ollama`" + ) from exc return ol.Client(timeout=self.TIMEOUT, host=self.host_url) @cached_property @@ -43,7 +54,8 @@ class Ollama(BaseProvider): mode=instructor.Mode.JSON, ) - def send_conversation(self, conversation: "Conversation") -> "Message": + @logger + def send_conversation(self, conversation: "Conversation", **kwargs) -> "Message": """Send a conversation to the Ollama API.""" from ..models import Message @@ -51,7 +63,9 @@ class Ollama(BaseProvider): {"role": msg.role, "content": msg.text} for msg in conversation.messages ] response = self.client.chat( - model=conversation.llm_model or DEFAULT_MODEL, messages=messages + model=conversation.llm_model or DEFAULT_MODEL, + messages=messages, + **{**self.DEFAULT_KWARGS, **kwargs}, ) assistant_message = response.get("message") @@ -64,6 +78,7 @@ class Ollama(BaseProvider): llm_provider=PROVIDER_NAME, ) + @logger def structured_response( self, prompt: str, @@ -81,18 +96,23 @@ class Ollama(BaseProvider): messages=messages, model=llm_model or self.DEFAULT_MODEL, response_model=response_model, - **kwargs, + **{**self.DEFAULT_KWARGS, **kwargs}, ) - return response + return response_model.model_validate(response) - def generate_text(self, prompt: str, *, llm_model: str | None = None) -> str: + @logger + def generate_text( + self, prompt: str, *, llm_model: str | None = None, **kwargs + ) -> str: """Generate text using the Ollama API.""" messages = [ {"role": "user", "content": prompt}, ] response = self.client.chat( - messages=messages, model=llm_model or self.DEFAULT_MODEL + messages=messages, + model=llm_model or self.DEFAULT_MODEL, + **{**self.DEFAULT_KWARGS, **kwargs}, ) return response.get("message", {}).get("content", "") diff --git a/simplemind/providers/openai.py b/simplemind/providers/openai.py index b6fa568..fb197e5 100644 --- a/simplemind/providers/openai.py +++ b/simplemind/providers/openai.py @@ -1,22 +1,28 @@ from functools import cached_property -from typing import Type, TypeVar +from typing import TYPE_CHECKING, Type, TypeVar import instructor -import openai as oa from pydantic import BaseModel +from ..logging import logger from ..settings import settings from ._base import BaseProvider +if TYPE_CHECKING: + from ..models import Conversation, Message + T = TypeVar("T", bound=BaseModel) PROVIDER_NAME = "openai" DEFAULT_MODEL = "gpt-4o-mini" +DEFAULT_MAX_TOKENS = 1_000 +DEFAULT_KWARGS = {"max_tokens": DEFAULT_MAX_TOKENS} class OpenAI(BaseProvider): NAME = PROVIDER_NAME DEFAULT_MODEL = DEFAULT_MODEL + DEFAULT_KWARGS = DEFAULT_KWARGS def __init__(self, api_key: str | None = None): self.api_key = api_key or settings.get_api_key(PROVIDER_NAME) @@ -26,6 +32,12 @@ class OpenAI(BaseProvider): """The raw OpenAI client.""" if not self.api_key: raise ValueError("OpenAI API key is required") + try: + import openai as oa + except ImportError as exc: + raise ImportError( + "Please install the `openai` package: `pip install openai`" + ) from exc return oa.OpenAI(api_key=self.api_key) @cached_property @@ -33,7 +45,8 @@ class OpenAI(BaseProvider): """A OpenAI client with Instructor.""" return instructor.from_openai(self.client) - def send_conversation(self, conversation: "Conversation", **kwargs): + @logger + def send_conversation(self, conversation: "Conversation", **kwargs) -> "Message": """Send a conversation to the OpenAI API.""" from ..models import Message @@ -42,7 +55,9 @@ class OpenAI(BaseProvider): ] response = self.client.chat.completions.create( - model=conversation.llm_model or DEFAULT_MODEL, messages=messages, **kwargs + model=conversation.llm_model or DEFAULT_MODEL, + messages=messages, + **{**self.DEFAULT_KWARGS, **kwargs}, ) # Get the response content from the OpenAI response @@ -57,6 +72,7 @@ class OpenAI(BaseProvider): llm_provider=PROVIDER_NAME, ) + @logger def structured_response( self, prompt: str, @@ -74,16 +90,19 @@ class OpenAI(BaseProvider): messages=messages, model=llm_model or self.DEFAULT_MODEL, response_model=response_model, - **kwargs, + **{**self.DEFAULT_KWARGS, **kwargs}, ) - return response + return response_model.model_validate(response) + @logger def generate_text(self, prompt: str, *, llm_model: str | None = None, **kwargs): """Generate text using the OpenAI API.""" messages = [ {"role": "user", "content": prompt}, ] response = self.client.chat.completions.create( - messages=messages, model=llm_model or self.DEFAULT_MODEL, **kwargs + messages=messages, + model=llm_model or self.DEFAULT_MODEL, + **{**self.DEFAULT_KWARGS, **kwargs}, ) return response.choices[0].message.content diff --git a/simplemind/providers/xai.py b/simplemind/providers/xai.py index dd32e0b..c19a0ae 100644 --- a/simplemind/providers/xai.py +++ b/simplemind/providers/xai.py @@ -1,20 +1,30 @@ from functools import cached_property +from typing import TYPE_CHECKING, Type, TypeVar import instructor -import openai as oa +from pydantic import BaseModel +from ..logging import logger from ..settings import settings from ._base import BaseProvider +if TYPE_CHECKING: + from ..models import Conversation, Message + +T = TypeVar("T", bound=BaseModel) + + PROVIDER_NAME = "xai" DEFAULT_MODEL = "grok-beta" BASE_URL = "https://api.x.ai/v1" DEFAULT_MAX_TOKENS = 1000 +DEFAULT_KWARGS = {"max_tokens": DEFAULT_MAX_TOKENS} class XAI(BaseProvider): NAME = PROVIDER_NAME DEFAULT_MODEL = DEFAULT_MODEL + DEFAULT_KWARGS = DEFAULT_KWARGS def __init__(self, api_key: str | None = None): self.api_key = api_key or settings.get_api_key(PROVIDER_NAME) @@ -24,6 +34,12 @@ class XAI(BaseProvider): """The raw OpenAI client.""" if not self.api_key: raise ValueError("XAI API key is required") + try: + import openai as oa + except ImportError as exc: + raise ImportError( + "Please install the `openai` package: `pip install openai`" + ) from exc return oa.OpenAI( api_key=self.api_key, base_url=BASE_URL, @@ -34,7 +50,8 @@ class XAI(BaseProvider): """A client patched with Instructor.""" return instructor.from_openai(self.client) - def send_conversation(self, conversation: "Conversation", **kwargs): + @logger + def send_conversation(self, conversation: "Conversation", **kwargs) -> "Message": """Send a conversation to the OpenAI API.""" from ..models import Message @@ -45,7 +62,7 @@ class XAI(BaseProvider): response = self.client.chat.completions.create( model=conversation.llm_model or self.DEFAULT_MODEL, messages=messages, - **kwargs, + **{**self.DEFAULT_KWARGS, **kwargs}, ) # Get the response content from the OpenAI response @@ -60,10 +77,14 @@ class XAI(BaseProvider): llm_provider=PROVIDER_NAME, ) - def structured_response(self, prompt: str, response_model, *, llm_model: str): + @logger + def structured_response( + self, prompt: str, response_model: Type[T], *, llm_model: str + ) -> T: raise NotImplementedError("XAI does not support structured responses") - def generate_text(self, prompt: str, *, llm_model: str, **kwargs): + @logger + def generate_text(self, prompt: str, *, llm_model: str, **kwargs) -> str: messages = [ {"role": "user", "content": prompt}, ] @@ -71,7 +92,7 @@ class XAI(BaseProvider): response = self.client.chat.completions.create( messages=messages, model=llm_model or self.DEFAULT_MODEL, - **kwargs, + **{**self.DEFAULT_KWARGS, **kwargs}, ) - return response.choices[0].message.content + return str(response.choices[0].message.content) diff --git a/simplemind/settings.py b/simplemind/settings.py index 5e28629..30efc26 100644 --- a/simplemind/settings.py +++ b/simplemind/settings.py @@ -1,19 +1,42 @@ -from typing import Literal, Optional, Union +from typing import Optional, Union from pydantic import Field, SecretStr, field_validator from pydantic_settings import BaseSettings, SettingsConfigDict -logging_level = Literal["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"] - class LoggingConfig(BaseSettings): """The class that holds all the logging settings for the application.""" - enabled: bool = Field(False, description="Enable logging") - level: logging_level = Field("INFO", description="The logging level") + is_enabled: bool = Field(False, description="Enable logging") model_config = SettingsConfigDict(extra="forbid") + def enable_logfire(self, **kwargs) -> None: + """Enable logging for the application.""" + # adding imports here to avoid forced dependencies + try: + import logfire + from logging import basicConfig + except ImportError as e: + raise ImportError( + "To enable logging, please install logfire: `pip install logfire`" + ) from e + + self.is_enabled = True + logfire.configure(**kwargs) + basicConfig(handlers=[logfire.LogfireLoggingHandler()]) + + try: + logfire.configure(**kwargs) + basicConfig(handlers=[logfire.LogfireLoggingHandler()]) + except Exception as e: + self.is_enabled = False # Reset flag on failure + raise RuntimeError("Failed to configure logging") from e + + def disable_logfire(self) -> None: + """Disable logging for the application.""" + self.is_enabled = False + class Settings(BaseSettings): """The class that holds all the API keys for the application.""" diff --git a/tests/conftest.py b/tests/conftest.py index 663d5b5..849d89d 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,8 +1,8 @@ -import pytest - import os import sys +import pytest + # Add the project root to the Python path. sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) diff --git a/tests/test_generate_data.py b/tests/test_generate_data.py index b622cad..dc07300 100644 --- a/tests/test_generate_data.py +++ b/tests/test_generate_data.py @@ -1,6 +1,8 @@ + import pytest from simplemind.providers import Anthropic, Gemini, OpenAI, Groq, Ollama, Amazon + from pydantic import BaseModel @@ -26,4 +28,4 @@ def test_generate_data(provider_cls): data = provider.structured_response(prompt=prompt, response_model=ResponseModel) assert isinstance(data, ResponseModel) - assert type(data.result) == int + assert isinstance(data.result, int)