diff --git a/simplemind/__init__.py b/simplemind/__init__.py index 29da250..b3bc4e2 100644 --- a/simplemind/__init__.py +++ b/simplemind/__init__.py @@ -1,10 +1,59 @@ -from typing import List, Optional +from typing import List, Optional, Type -from .models import Conversation, BasePlugin +from .models import Conversation, BasePlugin, BaseModel from .utils import find_provider from .settings import settings +class Session: + """A session object that maintains configuration across multiple API calls. + + Similar to `requests.Session`, this allows you to specify default settings + that will be used for all operations within the session. + """ + + def __init__( + self, + *, + llm_provider: str = settings.DEFAULT_LLM_PROVIDER, + llm_model: str = settings.DEFAULT_LLM_MODEL, + **kwargs, + ): + self.llm_provider = llm_provider + self.llm_model = llm_model + self.default_kwargs = kwargs + + def generate_text(self, prompt: str, **kwargs) -> str: + """Generate text using the session's default provider and model.""" + merged_kwargs = {**self.default_kwargs, **kwargs} + return generate_text( + prompt=prompt, + llm_provider=self.llm_provider, + llm_model=self.llm_model, + **merged_kwargs, + ) + + def generate_data( + self, prompt: str, response_model: Type[BaseModel], **kwargs + ) -> BaseModel: + """Generate structured data using the session's default provider and model.""" + merged_kwargs = {**self.default_kwargs, **kwargs} + return generate_data( + prompt=prompt, + response_model=response_model, + llm_provider=self.llm_provider, + llm_model=self.llm_model, + **merged_kwargs, + ) + + def create_conversation(self, **kwargs) -> "Conversation": + """Create a conversation using the session's default provider and model.""" + merged_kwargs = {**self.default_kwargs, **kwargs} + return create_conversation( + llm_provider=self.llm_provider, llm_model=self.llm_model, **merged_kwargs + ) + + def create_conversation( llm_model=None, llm_provider=None, *, plugins: Optional[List[BasePlugin]] = None ): @@ -53,4 +102,5 @@ __all__ = [ "generate_text", "settings", "BasePlugin", + "Session", ] diff --git a/simplemind/settings.py b/simplemind/settings.py index b5e2e28..4068ebd 100644 --- a/simplemind/settings.py +++ b/simplemind/settings.py @@ -17,6 +17,7 @@ class Settings(BaseSettings): ) XAI_API_KEY: Optional[SecretStr] = Field(None, description="API key for xAI") DEFAULT_LLM_PROVIDER: str = Field("openai", description="The default LLM provider") + DEFAULT_LLM_MODEL: str = Field("gpt-4o-mini", description="The default LLM model") model_config = SettingsConfigDict( env_file=".env", env_file_encoding="utf-8", case_sensitive=True, extra="ignore"