Add streaming support to Groq provider and implement generate_stream_text method

This commit refactors the `create_conversation` function in `simplemind/__init__.py` to use a more descriptive variable name (`conv`) for the conversation object. It also updates the `add_plugin` method to use the new variable name (`conv`) instead of `conversation`.

In `simplemind/models.py`, the `prepend_system_message` method now accepts an optional `meta` parameter. The method also adds a system message to the conversation by prepending it to the list of messages.

Additionally, the `add_message` method in `simplemind/models.py` has been modified to include type annotations and a default value for the `role` parameter. The method now requires the `text` parameter to be provided explicitly.

A new test file, `tests/test_conversations.py`, has been added to the repository. This file contains a test case for the `generate_data` function, which tests the functionality of different LLM providers.

Lastly, the test files `tests/test_generate_data.py` and `tests/test_generate_text.py` have been modified to remove the unused `Amazon` provider from the list of test cases.

.envrc.template

@@ -1,4 +1,7 @@
-export OPENAI_API_KEY=""
 export ANTHROPIC_API_KEY=""
+export GEMINI_API_KEY=""
+export GROQ_API_KEY=""
+export OLLAMA_HOST_URL=""
+export OPENAI_API_KEY=""
 export XAI_API_KEY=""
-export GROQ_API_KEY=""
+export AMAZON_PROFILE_NAME=""

.gitignore

@@ -166,3 +166,5 @@ cython_debug/
 .env
 
 src/**
+requirements.txt
+Pipfile

CHANGELOG.md

@@ -1,6 +1,64 @@
 Release History
 ===============
 
+## 0.2.2 (2024-11-02)
+
+- Add openai streaming support (set `stream=True` to `generate_text`).
+- `conv.prepend_system_message` now uses system role by default.
+- Add `provider.supports_streaming` property.
+- Add `provider.supports_structured_response` property.
+- General improvements.
+
+## 0.2.1 (2024-11-01)
+
+- Add `cached_property` to Amazon provider.
+
+## 0.2.0 (2024-11-01)
+
+- Add Amazon Bedrock provider.
+- Make all provider optional dependencies. Use `$ pip install 'simplemind[full]'` to install all providers.
+- General improvements.
+
+## 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.
+- Improvements to Anthropic provider, related to max tokens.
+- General improvements.
+- Add tests for structured response.
+- Add `llm_model` to `structured_response`.
+
+## 0.1.5 (2024-10-31)
+
+- Add Gemini provider.
+- Add structured response to Gemini provider.
+- Support for Python 3.10.
+
+## 0.1.4 (2024-10-30)
+
+- Introduce `Session` class to manage repeatability.
+- General improvements.
+
+## 0.1.3 (2024-10-30)
+
+- Make Conversation a context manager.
+- Add more robust conversation plugin hooks — replace `send_hook` with `pre_send_hook` and `post_send_hook`.
+- Change plugin hooks to try/except NotImplementedError.
+- Implement 'did you mean' with provider names. Can do this eventually with model names, as well.
+
+## 0.1.2 (2024-10-29)
+
+- Add ollama provider.
+
+## 0.1.1 (2024-10-29)
+
+- Fix Groq provider.
+
 ## 0.1.0 (2024-10-29)
 
 - Initial release.

Dockerfile

@@ -1,12 +1,21 @@
-FROM python:3.12.0
+FROM python:3.12-slim
 
-RUN apt-get update -y && apt-get upgrade -y
-RUN pip install --upgrade pip
+# Install system dependencies
+RUN apt-get update && apt-get install -y \
+    git \
+    && rm -rf /var/lib/apt/lists/*
 
-COPY requirements.txt /src/requirements.txt
+# Install uv
+RUN pip install uv
 
-WORKDIR /src
+# Create and set working directory
+WORKDIR /app
 
-RUN pip install -r requirements.txt
+# Copy requirements/project files
+ONBUILD COPY . .
 
-ENTRYPOINT ["python", "build.py"]
+# Install dependencies using uv
+RUN uv pip install "simplemind[full]" --system
+
+# Set default command
+CMD ["python"]

README.md

@@ -1,38 +1,81 @@
-# SimpleMind: AI for Humans™
+# Simplemind: AI for Humans™
 
-[![Auto Wiki](https://img.shields.io/badge/Auto_Wiki-Mutable.ai-blue)](https://mutable.ai/kennethreitz/simplemind)
+**Keep it simple, keep it human.**
 
-SimpleMind is an AI library designed to simplify your experience with AI APIs in Python. Inspired by a "for humans" philosophy, it abstracts away complexity, giving developers an intuitive and human-friendly way to interact with powerful AI capabilities. With SimpleMind, tapping into AI is as easy as a friendly conversation.
+Simplemind is AI library designed to simplify your experience with AI APIs in Python. Inspired by a "for humans" philosophy, it abstracts away complexity, giving developers an intuitive and human-friendly way to interact with powerful AI capabilities.
 
-```bash
-$ pip install simplemind
-```
+![simplemind](https://github.com/user-attachments/assets/36df2103-2583-4958-ad5e-19cda7740256)
 
 ## Features
-- **Easy-to-use AI tools**: SimpleMind provides simple interfaces to popular AI services.
+
+With Simplemind, tapping into AI is as easy as a friendly conversation.
+
+- **Easy-to-use AI tools**: Simplemind provides simple interfaces to most popular AI services.
 - **Human-centered design**: The library prioritizes readability and usability—no need to be an expert to start experimenting.
 - **Minimal configuration**: Get started quickly, without worrying about configuration headaches.
 
 ## Supported APIs
 
+The APIs remain identical between all supported providers / models:
+
+<table>
+  <thead>
+    <tr>
+      <th></th>
+      <th><code>llm_provider</code></th>
+      <th>Default <code>llm_model</code></th>
+    </tr>
+  </thead>
+  <tbody>
+    <tr>
+      <td><a href="https://www.anthropic.com/claude">Anthropic's Claude</a></td>
+      <td><code>"anthropic"</code></td>
+      <td><code>"claude-3-5-sonnet-20241022"</code></td>
+    </tr>
+    <tr>
+      <td><a href="https://aws.amazon.com/bedrock/">Amazon's Bedrock</a></td>
+      <td><code>"amazon"</code></td>
+      <td><code>"anthropic.claude-3-sonnet-20240229-v1:0"</code></td>
+    </tr>
+    <tr>
+      <td><a href="https://gemini.google/">Google's Gemini</a></td>
+      <td><code>"gemini"</code></td>
+      <td><code>"models/gemini-1.5-pro"</code></td>
+    </tr>
+    <tr>
+      <td><a href="https://groq.com/">Groq's Groq</a></td>
+      <td><code>"groq"</code></td>
+      <td><code>"llama3-8b-8192"</code></td>
+    </tr>
+    <tr>
+      <td><a href="https://ollama.com">Ollama</a></td>
+      <td><code>"ollama"</code></td>
+      <td><code>"llama3.2"</code></td>
+    </tr>
+    <tr>
+      <td><a href="https://openai.com/gpt">OpenAI's GPT</a></td>
+      <td><code>"openai"</code></td>
+      <td><code>"gpt-4o-mini"</code></td>
+    </tr>
+    <tr>
+      <td><a href="https://x.ai/">xAI's Grok</a></td>
+      <td><code>"xai"</code></td>
+      <td><code>"grok-beta"</code></td>
+    </tr>
+  </tbody>
+</table>
+
 To specify a specific provider or model, you can use the `llm_provider` and `llm_model` parameters when calling: `generate_text`, `generate_data`, or `create_conversation`.
 
-- **[OpenAI's GPT](https://openai.com/gpt)**
-- **[Anthropic's Claude](https://www.anthropic.com/claude)**
-- **[xAI's Grok](https://x.ai/)**
-- **[Groq's Groq](https://groq.com/)**
-
-If you'd like to see SimpleMind support additional providers or models, please send a pull request!
-
-## Why SimpleMind?
-- **Intuitive**: Built with Pythonic simplicity and readability in mind.
-- **For Humans**: Emphasizes a human-friendly interface, just like `requests` for HTTP.
-- **Open Source**: SimpleMind is open source, and contributions are always welcome!
-
+If you want to see Simplemind support additional providers or models, please send a pull request!
 
 ## Quickstart
 
-SimpleMind takes care of the complex API calls so you can focus on what matters—building, experimenting, and creating.
+Simplemind takes care of the complex API calls so you can focus on what matters—building, experimenting, and creating.
+
+```bash
+$ pip install 'simplemind[full]'
+```
 
 First, authenticate your API keys by setting them in the environment variables:
 
@@ -40,18 +83,17 @@ First, authenticate your API keys by setting them in the environment variables:
 $ export OPENAI_API_KEY="sk-..."
 ```
 
-This pattern allows you to keep your API keys private and out of your codebase. Other supported environment variables: `ANTHROPIC_API_KEY`, `GROK_API_KEY`, `XAI_API_KEY`, and `GROQ_API_KEY`.
+This pattern allows you to keep your API keys private and out of your codebase. Other supported environment variables: `ANTHROPIC_API_KEY`, `XAI_API_KEY`, `GROQ_API_KEY`, and `GEMINI_API_KEY`.
 
-Next, import SimpleMind and start using it:
+Next, import Simplemind and start using it:
 
 ```python
 import simplemind as sm
 ```
 
-
 ## Examples
 
-Here are some examples of how to use SimpleMind:
+Here are some examples of how to use Simplemind:
 
 ### Text Completion
 
@@ -82,37 +124,83 @@ 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:
 
 ```pycon
->>> conversation = sm.create_conversation(llm_model="gpt-4o-mini", llm_provider="openai")
+>>> conv = sm.create_conversation(llm_model="gpt-4o-mini", llm_provider="openai")
 
 >>> # Add a message to the conversation
->>> conversation.add_message("user", "Hi there, how are you?")
+>>> conv.add_message("user", "Hi there, how are you?")
 
->>> conversation.send()
+>>> conv.send()
 <Message role=assistant text="Hello! I'm just a computer program, so I don't have feelings, but I'm here and ready to help you. How can I assist you today?">
 ```
 
-To continue the conversation, you can call `conversation.send()` again, which returns the next message in the conversation:
+To continue the conversation, you can call `conv.send()` again, which returns the next message in the conversation:
 
 ```pycon
->>> conversation.add_message("user", "What is the meaning of life?")
->>> conversation.send()
+>>> conv.add_message("user", "What is the meaning of life?")
+>>> conv.send()
 <Message role=assistant text="The meaning of life is a profound philosophical question that has been explored by cultures, religions, and philosophers for centuries. Different people and belief systems offer varying interpretations:\n\n1. **Religious Perspectives:** Many religions propose that the meaning of life is to fulfill a divine purpose, serve God, or reach an afterlife. For example, Christianity often emphasizes love, faith, and service to God and others as central to life’s meaning.\n\n2. **Philosophical Views:** Philosophers offer diverse answers. Existentialists like Jean-Paul Sartre argue that life has no inherent meaning, and it is up to individuals to create their own purpose. Others, like Aristotle, suggest that achieving eudaimonia (flourishing or happiness) through virtuous living is the key to a meaningful life.\n\n3. **Scientific and Secular Approaches:** Some people find meaning through understanding the natural world, contributing to human knowledge, or through personal accomplishments and happiness. They may view life’s meaning as a product of connection, legacy, or the pursuit of knowledge and creativity.\n\n4. **Personal Perspective:** For many, the meaning of life is deeply personal, involving their relationships, passions, and goals. These individuals define life’s purpose through experiences, connections, and the impact they have on others and the world.\n\nUltimately, the meaning of life is a subjective question, with each person finding their own answers based on their beliefs, experiences, and reflections.">
 ```
 
+### Stop Repeating Yourself
+
+You can use the `Session` class to set default parameters for all calls:
+
+```python
+# Create a session with defaults
+gpt_4o_mini = sm.Session(llm_provider="openai", llm_model="gpt-4o-mini")
+
+# Now all calls use these defaults
+response = gpt_4o_mini.generate_text("Hello!")
+conversation = gpt_4o_mini.create_conversation()
+```
+
+This maintains the simplicity of the original API while reducing repetition. The session object also supports overriding defaults on a per-call basis:
+
+```python
+response = gpt_4o_mini.generate_text(
+    "Complex task here",
+    llm_model="gpt-4"
+)
+```
+
 ### Basic Memory Plugin
 
 Harnessing the power of Python, you can easily create your own plugins to add additional functionality to your conversations:
 
 ```python
-import simplemind as sm
-
-
-class SimpleMemoryPlugin:
+class SimpleMemoryPlugin(sm.BasePlugin):
     def __init__(self):
         self.memories = [
             "the earth has fictionally beeen destroyed.",
@@ -122,7 +210,7 @@ class SimpleMemoryPlugin:
     def yield_memories(self):
         return (m for m in self.memories)
 
-    def send_hook(self, conversation: sm.Conversation):
+    def pre_send_hook(self, conversation: sm.Conversation):
         for m in self.yield_memories():
             conversation.add_message(role="system", text=m)
 
@@ -133,9 +221,10 @@ conversation.add_plugin(SimpleMemoryPlugin())
 
 conversation.add_message(
     role="user",
-    text="Write a poem about the moon",
+    text="Please write a poem about the moon",
 )
 ```
+
 ```pycon
 >>> conversation.send()
 In the vast expanse where stars do play,
@@ -169,9 +258,20 @@ A reminder that in tales and fun,
 The universe is never done.
 ```
 
+Simple, yet effective.
+
+### Logging
+
+Simplemind uses [Logfire](https://pydantic.dev/logfire) for logging. To enable logging, call `sm.enable_logfire()`.
+
+### More Examples
+
 Please see the [examples](examples) directory for executable examples.
 
+---
+
 ## Contributing
+
 We welcome contributions of all kinds. Feel free to open issues for bug reports or feature requests, and submit pull requests to make SimpleMind even better.
 
 To get started:
@@ -182,11 +282,9 @@ To get started:
 4. Submit a pull request.
 
 ## License
-SimpleMind is licensed under the Apache 2.0 License.
+
+Simplemind is licensed under the Apache 2.0 License.
 
 ## Acknowledgements
-SimpleMind is inspired by the philosophy of "code for humans" and aims to make working with AI models accessible to all. Special thanks to the open-source community for their contributions and inspiration.
 
----------------
-
-SimpleMind: Keep it simple, keep it human.
+Simplemind is inspired by the philosophy of "code for humans" and aims to make working with AI models accessible to all. Special thanks to the open-source community for their contributions and inspiration.

docker-compose.yaml

@@ -1,10 +0,0 @@
-services:
-  simplemind:
-    build:
-      context: .
-      dockerfile: Dockerfile
-    volumes:
-      - ./simplemind:/src/simplemind
-      - ./build.py:/src/build.py
-    env_file:
-      - .env

docs/Makefile

@@ -0,0 +1,20 @@
+# Minimal makefile for Sphinx documentation
+#
+
+# You can set these variables from the command line, and also
+# from the environment for the first two.
+SPHINXOPTS    ?=
+SPHINXBUILD   ?= sphinx-build
+SOURCEDIR     = .
+BUILDDIR      = _build
+
+# Put it first so that "make" without argument is like "make help".
+help:
+	@$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)
+
+.PHONY: help Makefile
+
+# Catch-all target: route all unknown targets to Sphinx using the new
+# "make mode" option.  $(O) is meant as a shortcut for $(SPHINXOPTS).
+%: Makefile
+	@$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)

docs/conf.py

@@ -0,0 +1,34 @@
+# Configuration file for the Sphinx documentation builder.
+#
+# For the full list of built-in configuration values, see the documentation:
+# https://www.sphinx-doc.org/en/master/usage/configuration.html
+
+# -- Project information -----------------------------------------------------
+# https://www.sphinx-doc.org/en/master/usage/configuration.html#project-information
+
+import os
+import sys
+
+sys.path.insert(0, os.path.abspath(".."))
+
+import simplemind
+
+project = "simplemind"
+copyright = "2024 Kenneth Reitz"
+author = "Kenneth Reitz"
+release = "v0.2.2"
+
+# -- General configuration ---------------------------------------------------
+# https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration
+
+extensions = ["sphinx.ext.autodoc"]
+
+templates_path = ["_templates"]
+exclude_patterns = ["_build", "Thumbs.db", ".DS_Store"]
+
+
+# -- Options for HTML output -------------------------------------------------
+# https://www.sphinx-doc.org/en/master/usage/configuration.html#options-for-html-output
+
+html_theme = "alabaster"
+html_static_path = ["_static"]

docs/index.rst

@@ -0,0 +1,236 @@
+.. simplemind documentation master file, created by
+   sphinx-quickstart on Wed Oct 30 08:08:14 2024.
+   You can adapt this file completely to your liking, but it should at least
+   contain the root `toctree` directive.
+
+SimpleMind: AI for Humans™
+==========================
+
+**SimpleMind** is a versatile Python library designed to simplify interactions with various AI models. It provides a consistent and user-friendly interface to numerous AI providers, enabling developers to seamlessly integrate powerful AI capabilities into their applications without the overhead of managing multiple APIs and configurations.
+
+Features
+--------
+
+- **Unified Interface**: Interact with multiple AI providers using a single, consistent API
+- **Plugin Architecture**: Extend functionality with custom plugins for tasks like memory management and sentiment analysis
+- **Structured Data Support**: Generate and manipulate structured data using Pydantic models
+- **Human-Centered Design**: Prioritizes readability and ease of use, making AI integration accessible to all developers
+- **Minimal Configuration**: Quickly get started without extensive setup or configuration
+
+Supported Providers
+------------------
+
+SimpleMind supports a variety of AI providers:
+
+- `OpenAI's GPT <https://openai.com/gpt>`_
+- `Anthropic's Claude <https://www.anthropic.com/claude>`_
+- `xAI's Grok <https://x.ai/>`_
+- `Groq's Groq <https://groq.com/>`_
+- `Ollama <https://ollama.com>`_
+
+Installation
+-----------
+
+Install SimpleMind using pip:
+
+.. code-block:: shell
+
+    $ pip install simplemind
+
+Quickstart
+----------
+
+1. Set your API keys as environment variables:
+
+   .. code-block:: bash
+
+       $ export OPENAI_API_KEY="sk-..."
+       $ export ANTHROPIC_API_KEY="..."
+       $ export XAI_API_KEY="..."
+       $ export GROQ_API_KEY="..."
+
+This is the only required configuration.
+
+2. Import and use SimpleMind:
+
+   .. code-block:: python
+
+       import simplemind as sm
+
+       # Generate text using the default provider (OpenAI)
+       response = sm.generate_text("Write a poem about the moon.", llm_model="gpt-4o-mini")
+       print(response)
+
+Things to know:
+
+- The primary function for generating text is ``generate_text()``, which is used in the example above.
+- To generate structured data, use ``generate_data()``, which most providers support. This is extremely useful.
+- The third function, ``create_conversation()``, is used to engage in conversations with AI models.
+
+All of these functions accept an ``llm_model`` and ``llm_provider`` parameter, which allows you to specify the AI model to use. If not provided, the default model for the given provider will be used.
+
+
+Usage Examples
+--------------
+
+Here are some examples demonstrating SimpleMind's key features. From generating creative text and structured data to engaging in conversations and extending functionality with plugins, these examples showcase the library's versatility and ease of use.
+
+Feel free to adapt these examples to your specific use cases!
+
+
+Text Generation
+~~~~~~~~~~~~~~~
+
+This example generates a poem about the moon using the ``gpt-4o-mini`` model.
+
+.. code-block:: python
+
+    import simplemind as sm
+
+    poem = sm.generate_text("Write a poem about the moon.", llm_model="gpt-4o-mini")
+    print(poem)
+
+Structured Data Generation
+~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+This example generates a poem about love using the ``gpt-4o-mini`` model.
+
+.. code-block:: python
+
+    from pydantic import BaseModel
+
+    class Poem(BaseModel):
+        title: str
+        content: str
+
+    poem = sm.generate_data(
+        prompt="Write a poem about love",
+        llm_model="gpt-4o-mini",
+        response_model=Poem,
+    )
+    print(poem)
+
+Conversational AI
+~~~~~~~~~~~~~~~~~
+
+This example engages in a conversation with the ``gpt-4o-mini`` model.
+
+
+.. code-block:: python
+
+    conversation = sm.create_conversation(llm_model="gpt-4o-mini")
+    conversation.add_message("user", "Hi there, how are you?")
+    response = conversation.send()
+    print(response.text)
+
+Plugins
+~~~~~~~
+
+This example adds a simple custom memory plugin to the conversation.
+
+.. code-block:: python
+
+    class SimpleMemoryPlugin:
+        def __init__(self):
+            self.memories = ["the moon is made of cheese."]
+
+        def send_hook(self, conversation):
+            for memory in self.memories:
+                conversation.add_message(role="system", text=memory)
+
+    conversation = sm.create_conversation()
+    conversation.add_plugin(SimpleMemoryPlugin())
+    conversation.add_message("user", "Write a poem about the moon")
+    print(conversation.send().text)
+
+Plugin Development
+~~~~~~~~~~~~~~~~~~
+
+Plugins in SimpleMind follow a simple hook-based architecture. The ``send_hook`` method shown above is just one of several hooks available. Here's a more detailed example showing the complete plugin interface:
+
+.. code-block:: python
+
+    from simplemind.plugins import BasePlugin
+
+    class CustomPlugin(BasePlugin):
+        def __init__(self):
+            self.conversation_history = []
+
+        def initialize_hook(self, conversation):
+            """Called when the plugin is first added to a conversation."""
+            print("Plugin initialized!")
+
+        def pre_send_hook(self, conversation):
+            """Called before the conversation is sent to the AI provider."""
+            # Add any system messages or modify the conversation
+            conversation.add_message("system", "Remember to be helpful.")
+
+        def send_hook(self, conversation):
+            """Called during the send process."""
+            # Add messages or modify the conversation
+            self.conversation_history.append(conversation.messages)
+
+        def post_send_hook(self, conversation, response):
+            """Called after receiving a response from the AI provider."""
+            # Process or modify the response
+            return response
+
+        def cleanup_hook(self):
+            """Called when the plugin is removed or the conversation ends."""
+            self.conversation_history.clear()
+
+All plugins should inherit from ``BasePlugin``, which provides default no-op implementations of these hooks. You only need to implement the hooks you want to use. Here's a simpler example:
+
+.. code-block:: python
+
+    from simplemind.plugins import BasePlugin
+
+    class LoggingPlugin(BasePlugin):
+        def pre_send_hook(self, conversation):
+            print(f"Sending conversation with {len(conversation.messages)} messages")
+
+        def post_send_hook(self, conversation, response):
+            print(f"Received response: {response.text[:50]}...")
+            return response
+
+    conversation = sm.create_conversation()
+    conversation.add_plugin(LoggingPlugin())
+    conversation.add_message("user", "Hello!")
+    response = conversation.send()
+
+Plugins can be used to implement features like:
+
+- Conversation logging
+- Memory management
+- Response filtering
+- Token counting
+- Custom prompt engineering
+- Analytics and monitoring
+
+Multiple plugins can be added to a single conversation, and they will be executed in the order they were added.
+
+
+Contributing
+-----------
+
+1. Fork the Repository
+2. Create a New Branch
+3. Make Your Changes
+4. Submit a Pull Request
+
+Please review our `Code of Conduct <LICENSE>`_ before contributing.
+
+License
+-------
+
+SimpleMind is licensed under the `Apache 2.0 License <LICENSE>`_.
+
+.. toctree::
+   :maxdepth: 2
+   :caption: Contents:
+
+   installation
+   usage
+   api
+   contributing
+   changelog

docs/make.bat

@@ -0,0 +1,35 @@
+@ECHO OFF
+
+pushd %~dp0
+
+REM Command file for Sphinx documentation
+
+if "%SPHINXBUILD%" == "" (
+	set SPHINXBUILD=sphinx-build
+)
+set SOURCEDIR=.
+set BUILDDIR=_build
+
+%SPHINXBUILD% >NUL 2>NUL
+if errorlevel 9009 (
+	echo.
+	echo.The 'sphinx-build' command was not found. Make sure you have Sphinx
+	echo.installed, then set the SPHINXBUILD environment variable to point
+	echo.to the full path of the 'sphinx-build' executable. Alternatively you
+	echo.may add the Sphinx directory to PATH.
+	echo.
+	echo.If you don't have Sphinx installed, grab it from
+	echo.https://www.sphinx-doc.org/
+	exit /b 1
+)
+
+if "%1" == "" goto help
+
+%SPHINXBUILD% -M %1 %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O%
+goto end
+
+:help
+%SPHINXBUILD% -M help %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O%
+
+:end
+popd

examples/_context.py

@@ -5,6 +5,7 @@ import sys
 sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
 
 
+import simplemind
 import simplemind as sm
 
-__all__ = ["sm"]
+__all__ = ["simplemind", "sm"]

examples/bible_verses.py

@@ -0,0 +1,137 @@
+from _context import simplemind as sm
+from pydantic import BaseModel
+from rich.console import Console
+from rich.panel import Panel
+from rich.text import Text
+
+console = Console()
+
+gpt_4o_mini = sm.Session(llm_provider="openai")
+claude_sonnet = sm.Session(llm_provider="anthropic")
+
+
+class BibleVerse(BaseModel):
+    book: str
+    chapter: int
+    verse: int
+    text: str
+    translation: str
+
+
+class BiblePassage(BaseModel):
+    book: str
+    chapter: int
+    verses: list[BibleVerse]
+    translation: str
+
+
+class CrossReference(BaseModel):
+    passage: BiblePassage
+    notes: list[str]
+    origin_verse: BibleVerse
+    ai_perspective: str
+    anthropic_perspective: str
+
+
+def get_passage(book: str, chapter: int, translation: str = "ESV") -> BiblePassage:
+    passage = gpt_4o_mini.generate_data(
+        prompt=f"""Return {book} chapter {chapter} from the {translation} translation.
+Format each verse as plain text without any special characters or formatting.
+For example:
+- "Love is patient, love is kind."
+- "It does not envy, it does not boast"
+
+Return only the biblical text, formatted as a BiblePassage object.""",
+        response_model=BiblePassage,
+        max_tokens=8000,
+    )
+    return passage
+
+
+def get_cross_reference(passage: BiblePassage) -> CrossReference:
+    verses_text = "\n".join([f"Verse {v.verse}: {v.text}" for v in passage.verses])
+
+    # Get main cross-reference from OpenAI
+    ref = gpt_4o_mini.generate_data(
+        prompt=f"""Find a thematically related Bible passage that connects with this text:
+{verses_text}
+
+Return a CrossReference object with:
+1. A related passage (using plain text without special characters)
+2. A list of clear, specific notes explaining the thematic connections
+3. The original passage included
+4. An AI perspective that provides a thoughtful, modern interpretation of how these passages relate to contemporary life and universal human experiences""",
+        response_model=CrossReference,
+    )
+
+    # Get Anthropic's perspective separately
+    anthropic_insight = claude_sonnet.generate_text(
+        prompt=f"""Analyze these biblical passages from a philosophical and ethical perspective:
+
+Original passage:
+{verses_text}
+
+Cross-reference passage:
+{' '.join([f'Verse {v.verse}: {v.text}' for v in ref.passage.verses])}
+
+Provide a thoughtful analysis focusing on the philosophical and ethical implications of these passages, drawing from your training in ethics and philosophy.
+Return your response as a plain string.""",
+    )
+
+    # Add Anthropic's perspective to the reference object
+    ref.anthropic_perspective = anthropic_insight
+    return ref
+
+
+def pretty_print_reference(ref: CrossReference):
+    # Create origin passage panel
+    origin_text = Text()
+    origin_text.append(
+        f"{ref.origin_verse.book} {ref.origin_verse.chapter}\n",
+        style="bold blue",
+    )
+    origin_text.append(f"{ref.origin_verse.verse}. ", style="blue")
+    origin_text.append(f"{ref.origin_verse.text}\n", style="italic")
+    origin_text.append(f"\n({ref.origin_verse.translation})", style="dim")
+
+    origin_panel = Panel(origin_text, title="Original Passage", border_style="blue")
+
+    # Create cross reference panel
+    ref_text = Text()
+    ref_text.append(
+        f"{ref.passage.book} {ref.passage.chapter}\n",
+        style="bold green",
+    )
+    for verse in ref.passage.verses:
+        ref_text.append(f"{verse.verse}. ", style="green")
+        ref_text.append(f"{verse.text}\n", style="italic")
+    ref_text.append(f"\n({ref.passage.translation})", style="dim")
+
+    ref_panel = Panel(ref_text, title="Cross Reference", border_style="green")
+
+    # Create notes panel with bullet points
+    notes_text = Text()
+    for note in ref.notes:
+        notes_text.append("• ", style="yellow")
+        notes_text.append(f"{note}\n")
+
+    notes_panel = Panel(notes_text, title="Thematic Connections", border_style="yellow")
+
+    # Add new AI perspective panel
+    ai_text = Text()
+    ai_text.append(ref.ai_perspective)
+
+    ai_panel = Panel(ai_text, title="AI Perspective", border_style="magenta")
+
+    # Print all panels
+    console.print(origin_panel)
+    console.print(ref_panel)
+    console.print(notes_panel)
+    console.print(ai_panel)
+
+
+if __name__ == "__main__":
+    # Get 1 Corinthians 13 (The Love Chapter)
+    passage = get_passage("1 Corinthians", 13)
+    ref = get_cross_reference(passage)
+    pretty_print_reference(ref)

examples/contextual_memory.py

@@ -1,15 +1,20 @@
-from _context import sm
-
-from pydantic import BaseModel
-import openai
-import faiss
-import numpy as np
 import os
 import pickle
 
+import faiss
+import numpy as np
+import openai
+from _context import sm
+from pydantic import BaseModel
 
-class ContextualMemoryPlugin:
-    def __init__(self, api_key: str, memory_file: str = "memories.pkl", embedding_model: str = "text-embedding-ada-002"):
+
+class ContextualMemoryPlugin(sm.BasePlugin):
+    def __init__(
+        self,
+        api_key: str,
+        memory_file: str = "memories.pkl",
+        embedding_model: str = "text-embedding-ada-002",
+    ):
         openai.api_key = api_key
         self.memory_file = memory_file
         self.embedding_model = embedding_model
@@ -35,29 +40,29 @@ class ContextualMemoryPlugin:
     def build_faiss_index(self):
         if self.embeddings:
             self.index = faiss.IndexFlatL2(len(self.embeddings[0]))
-            self.index.add(np.array(self.embeddings).astype('float32'))
+            self.index.add(np.array(self.embeddings).astype("float32"))
         else:
             self.index = faiss.IndexFlatL2(1536)
 
     def get_embedding(self, text: str) -> list:
         response = openai.Embedding.create(input=text, model=self.embedding_model)
-        return response['data'][0]['embedding']
+        return response["data"][0]["embedding"]
 
     def add_memory(self, memory: str):
         embedding = self.get_embedding(memory)
         self.memories.append(memory)
         self.embeddings.append(embedding)
-        self.index.add(np.array([embedding]).astype('float32'))
+        self.index.add(np.array([embedding]).astype("float32"))
         self.save_memories()
 
     def retrieve_memories(self, query: str, top_k: int = 3) -> list:
         if not self.index or len(self.embeddings) == 0:
             return []
         query_embedding = self.get_embedding(query)
-        D, I = self.index.search(np.array([query_embedding]).astype('float32'), top_k)
+        D, I = self.index.search(np.array([query_embedding]).astype("float32"), top_k)
         return [self.memories[i] for i in I[0] if i < len(self.memories)]
 
-    def send_hook(self, conversation: sm.Conversation):
+    def pre_send_hook(self, conversation: sm.Conversation):
         # Retrieve relevant memories based on the latest user message
         if conversation.messages:
             last_user_message = conversation.messages[-1].text
@@ -69,13 +74,16 @@ class ContextualMemoryPlugin:
         # Optionally, add the AI's response to memories
         self.add_memory(response)
 
+
 # Example Usage
 
+
 # Define a Pydantic model if needed
 class Story(BaseModel):
     title: str
     content: str
 
+
 # Initialize the conversation with the ContextualMemoryPlugin
 memory_plugin = ContextualMemoryPlugin(api_key=sm.settings.OPENAI_API_KEY)
 

examples/cooking_recipe_example.py

@@ -0,0 +1,66 @@
+from pydantic import BaseModel
+import simplemind as sm
+
+
+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]
+    
+    def __str__(self) -> str:
+        output = f"\n=== {self.name.upper()} ===\n\n"
+        
+        output += "INGREDIENTS:\n"
+        for ing in self.ingredients:
+            output += f"• {ing.quantity} {ing.unit} {ing.name}\n"
+        
+        output += "\nINSTRUCTIONS:\n"
+        for step in self.instructions:
+            output += f"{step.step_number}. {step.instruction}\n"
+        
+        return output
+    
+
+recipe = sm.generate_data(
+    "Write a recipe for chocolate chip cookies",
+    llm_model="gpt-4o-mini",
+    llm_provider="openai",
+    response_model=Recipe,
+)
+
+print(recipe)
+# Expected output is something like this:
+#
+# === CHOCOLATE CHIP COOKIES ===
+#
+# INGREDIENTS:
+# • 2.25 cups all-purpose flour
+# • 1.0 teaspoon baking soda
+# • 0.5 teaspoon salt
+# • 1.0 cup unsalted butter
+# • 0.75 cup sugar
+# • 0.75 cup brown sugar
+# • 1.0 teaspoon vanilla extract
+# • 2.0 large eggs
+# • 2.0 cups semi-sweet chocolate chips
+#
+# INSTRUCTIONS:
+# 1. Preheat your oven to 350°F (175°C).
+# 2. In a small bowl, combine flour, baking soda, and salt; set aside.
+# 3. In a large bowl, cream together the butter, sugar, and brown sugar until smooth.
+# 4. Beat in the vanilla extract and eggs, one at a time.
+# 5. Gradually blend in the flour mixture until just combined.
+# 6. Stir in the chocolate chips.
+# 7. Drop by rounded tablespoon onto ungreased cookie sheets.
+# 8. Bake for 9 to 11 minutes, or until edges are golden.
+# 9. Let cool on the cookie sheet for a few minutes before transferring to wire racks to cool completely.
+

examples/discussion.py

@@ -0,0 +1,132 @@
+import time
+from typing import List, Tuple
+
+from rich.console import Console
+from rich.markdown import Markdown
+
+from _context import sm
+
+
+class MultiAIConversation:
+    """Orchestrates conversations between multiple AI models."""
+
+    MODEL_SESSIONS = {
+        "Llama3.2": sm.Session(
+            llm_provider="ollama",
+            llm_model="llama3.2",
+        ),
+        "Claude-3.5-Sonnet": sm.Session(
+            llm_provider="anthropic",
+            llm_model="claude-3-5-sonnet-20241022",
+        ),
+        "GPT-4o": sm.Session(
+            llm_provider="openai",
+            llm_model="gpt-4o",
+        ),
+        "Grok-Beta": sm.Session(
+            llm_provider="xai",
+            llm_model="grok-beta",
+        ),
+    }
+
+    def __init__(self, topic: str, turns_per_model: int = 1, max_rounds: int = 5):
+
+        self.topic = topic
+        self.turns_per_model = turns_per_model
+        self.max_rounds = max_rounds
+        self.conversation_history: List[Tuple[str, str]] = []
+        self.console = Console()
+
+    def _format_system_prompt(self, ai_name: str) -> str:
+        """Creates a system prompt for each AI model."""
+        return f"""You are {ai_name}. You are participating in a thoughtful discussion with other AI models about {self.topic}.
+
+Rules:
+1. Be concise but insightful (keep responses under 100 words)
+2. Build upon previous points made in the conversation
+3. Ask questions to deepen the discussion when appropriate
+4. Stay on topic while maintaining your unique perspective
+5. Be respectful of other viewpoints while maintaining your distinct voice
+
+Current discussion topic: {self.topic}"""
+
+    def _create_conversation(
+        self, session: sm.Session, ai_name: str
+    ) -> sm.Conversation:
+        """Creates a new conversation with appropriate context for an AI model."""
+        conv = session.create_conversation()
+
+        # Add system prompt
+        conv.add_message(role="user", text=self._format_system_prompt(ai_name))
+
+        # Add conversation history
+        for speaker, message in self.conversation_history[-3:]:  # Last 3 messages
+            conv.add_message(role="user", text=f"{speaker} said: {message}")
+
+        return conv
+
+    def _print_response(self, ai_name: str, response: str):
+        """Pretty prints an AI response using Rich."""
+        self.console.print(f"\n[bold blue]{ai_name}[/bold blue]:")
+        self.console.print(Markdown(response))
+        # Store in history
+        self.conversation_history.append((ai_name, response))
+
+    def run_conversation(self):
+        """Runs the multi-AI conversation."""
+
+        # Initialize the conversation
+        initial_prompt = (
+            f"Let's have a thoughtful discussion about {self.topic}. "
+            "Please share your initial thoughts in 2-3 sentences."
+        )
+
+        for round_num in range(self.max_rounds):
+            self.console.print(f"\n[bold green]Round {round_num + 1}[/bold green]")
+
+            for model_name, session in self.MODEL_SESSIONS.items():
+                for turn in range(self.turns_per_model):
+                    conversation = self._create_conversation(session, model_name)
+
+                    # Add the prompt
+                    prompt = (
+                        initial_prompt
+                        if round_num == 0 and turn == 0
+                        else (
+                            f"Continue the discussion about {self.topic}, "
+                            "responding to the previous points made."
+                        )
+                    )
+
+                    conversation.add_message(role="user", text=prompt)
+
+                    # Get and print response
+                    response = conversation.send()
+                    self._print_response(model_name, response.text)
+
+                    # Small delay to prevent rate limiting
+                    time.sleep(1)
+
+            # Optional: Add a separator between rounds
+            self.console.print("\n" + "-" * 50)
+
+
+def have_ai_discussion(topic: str, turns_per_model: int = 1, max_rounds: int = 3):
+    """Convenience function to start an AI discussion."""
+    debate = MultiAIConversation(
+        topic=topic, turns_per_model=turns_per_model, max_rounds=max_rounds
+    )
+
+    print(f"\nStarting AI discussion on: {topic}")
+    print("=" * 50)
+
+    debate.run_conversation()
+
+
+# Example usage
+if __name__ == "__main__":
+    # Example topics
+    topic = "The future of human-AI collaboration in creative fields",
+
+    # Run a discussion on the first topic
+    have_ai_discussion(topic=topic, turns_per_model=1, max_rounds=3)

examples/generate_data.py

@@ -1,29 +1,43 @@
-from typing import List
-
-from pydantic import BaseModel
+from typing import Iterator, List
 
 from _context import sm
+from pydantic import BaseModel
 
 
 class Movie(BaseModel):
     title: str
     year: int
 
+
 class MovieCharecter(BaseModel):
     name: str
     actor: str
 
+
 class MovieQuote(BaseModel):
     quote: str
     movie: Movie
     charecter: MovieCharecter
 
+
 class QuotesList(BaseModel):
     quotes: List[MovieQuote]
-    theme: str
 
 
-quotes = sm.generate_data(llm_provider="openai", llm_model="gpt-4o-mini", prompt="Generate 20 quotes from famous movies", response_model=QuotesList)
+def gen_quotes(n: int = 10) -> Iterator[MovieQuote]:
+    """Generate a list of quotes from famous movies."""
 
-for quote in quotes.quotes:
-    print(f"{quote.charecter.name} from {quote.movie.title} ({quote.movie.year}): {quote.quote!r}")
+    for q in sm.generate_data(
+        llm_provider="openai",
+        llm_model="gpt-4o-mini",
+        prompt=f"Generate {n} quotes from famous movies",
+        response_model=QuotesList,
+    ).quotes:
+        yield q
+
+
+if __name__ == "__main__":
+    for quote in gen_quotes(n=20):
+        print(
+            f"{quote.charecter.name} from {quote.movie.title} ({quote.movie.year}): {quote.quote!r}"
+        )

examples/generate_stream_text.py

@@ -0,0 +1,7 @@
+from _context import sm
+
+# Defaults to the default provider (openai)
+r = sm.generate_text("Write a poem about the moon", llm_provider="gemini", stream=True)
+
+for chunk in r:
+    print(chunk, end="", flush=True)

examples/math_plugin.py

@@ -1,9 +1,11 @@
 from _context import sm
 
 
-class MathPlugin:
-    def send_hook(self, conversation: sm.Conversation):
+class MathPlugin(sm.BasePlugin):
+    def pre_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.",
                 )
 
 

examples/medicine_data.py

@@ -0,0 +1,94 @@
+from pydantic import BaseModel
+from _context import simplemind as sm
+from rich.console import Console
+from rich.panel import Panel
+from rich.table import Table
+
+
+class SideEffect(BaseModel):
+    effect: str
+    severity: str  # mild, moderate, severe
+    frequency: str  # common, uncommon, rare
+
+
+class Medication(BaseModel):
+    brand_name: str
+    generic_name: str
+    drug_class: str
+    half_life: str
+    common_uses: list[str]
+    side_effects: list[SideEffect]
+    typical_dosage: str
+    warnings: list[str]
+
+
+class MedicationList(BaseModel):
+    root: list[Medication]
+
+
+# Create a session with your preferred model
+session = sm.Session(llm_provider="openai", llm_model="gpt-4o-mini")
+
+
+# Update the prompt to use an f-string with a parameter
+def get_medication_prompt(medications: list[str]) -> str:
+    return f"""
+Provide detailed medical information about {', '.join(medications)}.
+Include their generic names, drug classes, half-lives, common uses, side effects (with severity and frequency),
+typical dosages, and important warnings.
+Return the information as separate medication entries.
+"""
+
+
+# Example usage
+medications_to_lookup = ["Abilify (aripiprazole)", "Trileptal (oxcarbazepine)"]
+prompt = get_medication_prompt(medications_to_lookup)
+
+# Generate structured data for medications
+medications = session.generate_data(prompt=prompt, response_model=MedicationList)
+
+# Create a Rich console
+console = Console()
+
+# Replace the print section with Rich formatting
+for med in medications.root:
+    # Create a table for the medication details
+    table = Table(show_header=False, box=None)
+    table.add_row("[bold cyan]Generic Name:[/]", med.generic_name)
+    table.add_row("[bold cyan]Drug Class:[/]", med.drug_class)
+    table.add_row("[bold cyan]Half Life:[/]", med.half_life)
+
+    # Create a nested table for common uses
+    uses_table = Table(show_header=False, box=None, padding=(0, 2))
+    for use in med.common_uses:
+        uses_table.add_row("•", use)
+
+    # Create a nested table for side effects
+    effects_table = Table(show_header=False, box=None, padding=(0, 2))
+    for effect in med.side_effects:
+        severity_color = {"mild": "green", "moderate": "yellow", "severe": "red"}.get(
+            effect.severity.lower(), "white"
+        )
+        effects_table.add_row(
+            "•",
+            effect.effect,
+            f"[{severity_color}]{effect.severity}[/]",
+            f"({effect.frequency})",
+        )
+
+    # Create a nested table for warnings
+    warnings_table = Table(show_header=False, box=None, padding=(0, 2))
+    for warning in med.warnings:
+        warnings_table.add_row("•", f"[red]{warning}[/]")
+
+    # Add the nested tables to the main table
+    table.add_row("[bold cyan]Common Uses:[/]", uses_table)
+    table.add_row("[bold cyan]Side Effects:[/]", effects_table)
+    table.add_row("[bold cyan]Typical Dosage:[/]", med.typical_dosage)
+    table.add_row("[bold cyan]Warnings:[/]", warnings_table)
+
+    # Create and print a panel for each medication
+    console.print(
+        Panel(table, title=f"[bold blue]{med.brand_name}[/]", border_style="blue")
+    )
+    console.print()  # Add a blank line between medications

examples/requirements.txt

@@ -2,3 +2,4 @@ numpy
 openai
 pydantic
 faiss-cpu
+rich

examples/sentiment_analysis.py

@@ -1,8 +1,9 @@
-from _context import sm
-
-from pydantic import BaseModel
 from typing import Literal
 
+from _context import sm
+from pydantic import BaseModel
+
+# Note: you should probably be using textblob for this.
 
 class SentimentAnalysis(BaseModel):
     sentiment: Literal["positive", "negative", "neutral"]

examples/simple_logging_plugin.py

@@ -0,0 +1,31 @@
+import simplemind as sm
+
+
+class LoggingPlugin(sm.BasePlugin):
+    def pre_send_hook(self, conversation):
+        print(f"Sending conversation with {len(conversation.messages)} messages")
+
+    def add_message_hook(self, conversation, message):
+        print(f"Adding message to conversation: {message.text}")
+
+    def cleanup_hook(self, conversation):
+        print(f"Cleaning up conversation with {len(conversation.messages)} messages")
+
+    def initialize_hook(self, conversation):
+        print("Initializing conversation")
+
+    def post_send_hook(self, conversation, response):
+        print(f"Received response: {response.text}")
+
+
+with sm.create_conversation() as conversation:
+    # Add the logging plugin.
+    conversation.add_plugin(LoggingPlugin())
+
+    # Add a message to the conversation.
+    conversation.add_message("user", "Hello!", meta={})
+
+    # Send the conversation.
+    response = conversation.send()
+
+print(f"Response: {response.text}")

examples/simple_memory.py

@@ -11,9 +11,9 @@ class SimpleMemoryPlugin:
     def yield_memories(self):
         return (m for m in self.memories)
 
-    def send_hook(self, conversation: sm.Conversation):
+    def initialize_hook(self, conversation: sm.Conversation):
         for m in self.yield_memories():
-            conversation.add_message(role="system", text=m)
+            conversation.prepend_system_message(text=m)
 
 
 conversation = sm.create_conversation(llm_model="grok-beta", llm_provider="xai")
@@ -21,7 +21,7 @@ conversation.add_plugin(SimpleMemoryPlugin())
 
 conversation.add_message(
     role="user",
-    text="Write a poem about the moon",
+    text="Please write a poem about the moon",
 )
 
 r = conversation.send()

examples/translate_text.py

@@ -1,9 +1,13 @@
 from _context import sm
 
-conversation = sm.create_conversation(llm_model="gpt-4o", llm_provider="openai")
 
-conversation.add_message(
-    "user", "Translate the following text to French: 'Hello, world!'"
-)
+def translate_to_french(text: str) -> str:
+    conversation = sm.create_conversation(llm_model="gpt-4o", llm_provider="openai")
 
-print(conversation.send().text)
+    conversation.add_message(
+        "user", f"Translate the following text to French: {text!r}"
+    )
+    return conversation.send().text
+
+
+print(translate_to_french("an omlette with cheese"))

examples/two_llms_talking.py

@@ -0,0 +1,59 @@
+import time
+
+import simplemind as sm
+
+
+class ConversationPlugin(sm.BasePlugin):
+    def post_send_hook(self, conversation, response):
+        # Print the LLM model and the response text.
+        print(f"{conversation.llm_model}:\n{response.text.strip()}\n\n------------\n")
+
+
+def have_conversation(rounds: int = 3):
+    # Create two conversations - one for each AI
+    with (
+        sm.create_conversation(
+            llm_model="claude-3-5-sonnet-20241022", llm_provider="anthropic"
+        ) as claude_conv,
+        sm.create_conversation(
+            llm_model="llama3.2", llm_provider="ollama"
+        ) as llama_conv,
+    ):
+
+        # Add our plugin to both
+        plugin = ConversationPlugin()
+        claude_conv.add_plugin(plugin)
+        llama_conv.add_plugin(plugin)
+
+        # Start the conversation
+        prompt = "What do you think about the future of artificial intelligence? Please keep your response brief."
+        claude_conv.add_message("user", prompt, meta={})
+        claude_response = claude_conv.send()
+
+        # Have them discuss back and forth
+        for _ in range(rounds):
+            # Llama responds to Claude
+            llama_conv.add_message(
+                "user",
+                f"Respond to this statement from another AI: {claude_response.text}",
+                meta={},
+            )
+            llama_response = llama_conv.send()
+
+            time.sleep(1)  # Add a small delay between responses
+
+            # Claude responds to Llama
+            claude_conv.add_message(
+                "user",
+                f"Respond to this statement from another AI: {llama_response.text}",
+                meta={},
+            )
+            claude_response = claude_conv.send()
+
+            time.sleep(1)
+
+
+if __name__ == "__main__":
+    print("Starting AI conversation...\n")
+    have_conversation()
+    print("\nConversation ended.")

pyproject.toml

@@ -1,10 +1,27 @@
 [project]
 name = "simplemind"
-version = "0.1.0"
+version = "0.2.2"
 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.11"
-dependencies = ["pydantic", "pydantic-settings", "instructor", "openai", "anthropic", "groq"]
+requires-python = ">=3.10"
+dependencies = ["pydantic", "pydantic-settings", "instructor", "logfire"]
+
+[project.optional-dependencies]
+full = [
+    "openai",
+    "anthropic",
+    "ollama",
+    "groq",
+    "google-generativeai",
+    "botocore",
+    "boto3"
+]
+openai = ["openai"]
+anthropic = ["anthropic"]
+ollama = ["ollama", "openai"]
+groq = ["groq"]
+gemini = ["google-generativeai"]
+amazon = ["boto3", "botocore", "anthropic"]
 
 [build-system]
 requires = ["hatchling"]

simplemind/__init__.py

@@ -1,40 +1,144 @@
-from .models import Conversation
-from .utils import find_provider
+from typing import Generator, List, Type
+
+from .models import BaseModel, BasePlugin, Conversation
 from .settings import settings
+from .utils import find_provider
 
 
-def create_conversation(llm_model=None, llm_provider=None):
+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 | None = None,
+        **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: str | None = None,
+    llm_provider: str | None = None,
+    plugins: List[BasePlugin] | None = None,
+    **kwargs,
+) -> Conversation:
     """Create a new conversation."""
 
-    return Conversation(
-        llm_model=llm_model, llm_provider=llm_provider or settings.DEFAULT_LLM_PROVIDER
+    # Create the conversation.
+    conv = Conversation(
+        llm_model=llm_model,
+        llm_provider=llm_provider or settings.DEFAULT_LLM_PROVIDER,
     )
 
+    # Add plugins to the conversation.
+    for plugin in plugins or []:
+        conv.add_plugin(plugin)
 
-def generate_data(prompt, *, llm_model=None, llm_provider=None, response_model=None):
+    return conv
+
+
+def generate_data(
+    prompt: str,
+    *,
+    llm_model: str | None = None,
+    llm_provider: str | None = None,
+    response_model: Type[BaseModel],
+    **kwargs,
+) -> BaseModel:
     """Generate structured data from a given prompt."""
 
+    # Find the provider.
     provider = find_provider(llm_provider or settings.DEFAULT_LLM_PROVIDER)
 
+    # Generate the data.
     return provider.structured_response(
         prompt=prompt,
         llm_model=llm_model,
         response_model=response_model,
+        **kwargs,
     )
 
 
-def generate_text(prompt, *, llm_model=None, llm_provider=None, **kwargs):
+def generate_text(
+    prompt: str,
+    *,
+    llm_model: str | None = None,
+    llm_provider: str | None = None,
+    stream: bool = False,
+    **kwargs,
+) -> str:
     """Generate text from a given prompt."""
+
+    # Find the provider.
     provider = find_provider(llm_provider or settings.DEFAULT_LLM_PROVIDER)
 
-    return provider.generate_text(prompt=prompt, llm_model=llm_model, **kwargs)
+    # Generate the text.
+    if stream:
+        if not provider.supports_streaming:
+            raise ValueError(f"{provider} does not support streaming.")
 
+        return provider.generate_stream_text(
+            prompt=prompt, llm_model=llm_model, **kwargs
+        )
+    else:
+        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
 
 __all__ = [
-    "Conversation",
     "create_conversation",
     "find_provider",
     "generate_data",
     "generate_text",
     "settings",
+    "BasePlugin",
+    "Session",
+    "Plugin",
+    "enable_logfire",
 ]

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

simplemind/models.py

@@ -1,18 +1,18 @@
 import uuid
-from abc import ABC, abstractmethod
 from datetime import datetime
+from types import TracebackType
 from typing import Any, Dict, List, Literal, Optional
 
-
 from pydantic import BaseModel, Field
 
 from .utils import find_provider
 
-
 MESSAGE_ROLE = Literal["system", "user", "assistant"]
 
 
 class SMBaseModel(BaseModel):
+    """The base SimpleMind model class."""
+
     date_created: datetime = Field(default_factory=datetime.now)
 
     def __str__(self):
@@ -22,16 +22,36 @@ class SMBaseModel(BaseModel):
         return str(self)
 
 
-class BasePlugin(ABC):
+class BasePlugin(SMBaseModel):
     """The base conversation plugin class."""
 
-    @abstractmethod
-    def send_hook(self, conversation: "Conversation"):
-        """Send a hook to the plugin."""
+    # Plugin metadata.
+    meta: Dict[str, Any] = {}
+
+    def initialize_hook(self, conversation: "Conversation") -> Any:
+        """Initialize a hook for the plugin."""
+        raise NotImplementedError
+
+    def cleanup_hook(self, conversation: "Conversation") -> Any:
+        """Cleanup a hook for the plugin."""
+        raise NotImplementedError
+
+    def add_message_hook(self, conversation: "Conversation", message: "Message") -> Any:
+        """Add a message hook for the plugin."""
+        raise NotImplementedError
+
+    def pre_send_hook(self, conversation: "Conversation") -> Any:
+        """Pre-send hook for the plugin."""
+        raise NotImplementedError
+
+    def post_send_hook(self, conversation: "Conversation", response: "Message") -> Any:
+        """Post-send hook for the plugin."""
         raise NotImplementedError
 
 
 class Message(SMBaseModel):
+    """A message in a conversation."""
+
     role: MESSAGE_ROLE
     text: str
     meta: Dict[str, Any] = {}
@@ -43,7 +63,16 @@ class Message(SMBaseModel):
         return f"<Message role={self.role} text={self.text!r}>"
 
     @classmethod
-    def from_raw_response(cls, *, text: str, raw):
+    def from_raw_response(cls, *, text: str, raw: Any) -> "Message":
+        """Create a Message instance from a raw response.
+
+        Args:
+            text (str): The message text.
+            raw (Any): The raw response data.
+
+        Returns:
+            Message: A new Message instance.
+        """
         self = cls()
         self.text = text
         self.raw = raw
@@ -51,40 +80,114 @@ class Message(SMBaseModel):
 
 
 class Conversation(SMBaseModel):
+    """A conversation between a user and an assistant."""
+
     id: str = Field(default_factory=lambda: str(uuid.uuid4()))
     messages: List[Message] = []
     llm_model: Optional[str] = None
     llm_provider: Optional[str] = None
-    plugins: List[Any] = []
+    plugins: List[BasePlugin] = []
 
     def __str__(self):
         return f"<Conversation id={self.id!r}>"
 
+    def __enter__(self):
+        # Execute all initialize hooks.
+        for plugin in self.plugins:
+            if hasattr(plugin, "initialize_hook"):
+                try:
+                    plugin.initialize_hook(self)
+                except NotImplementedError:
+                    pass
+
+        return self
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException],
+        exc_value: BaseException,
+        traceback: TracebackType,
+    ) -> None:
+        """Execute all cleanup hooks."""
+        for plugin in self.plugins:
+            if hasattr(plugin, "cleanup_hook"):
+                try:
+                    plugin.cleanup_hook(self)
+                except NotImplementedError:
+                    pass
+
+    def prepend_system_message(self, text: str, meta: Dict[str, Any] | None = None):
+        """Prepend a system message to the conversation."""
+        self.messages = [
+            Message(role="system", text=text, meta=meta or {})
+        ] + self.messages
+
     def add_message(
-        self, role: MESSAGE_ROLE, text: str, meta: Optional[Dict[str, Any]] = None
+        self,
+        role: MESSAGE_ROLE = "user",
+        text: str | None = None,
+        *,
+        meta: Optional[Dict[str, Any]] = None,
     ):
         """Add a new message to the conversation."""
+
+        assert text is not None
+
+        # Ensure meta is a dict.
         if meta is None:
             meta = {}
+
+        # Execute all add-message hooks.
+        for plugin in self.plugins:
+            if hasattr(plugin, "add_message_hook"):
+                try:
+                    plugin.add_message_hook(
+                        self, Message(role=role, text=text, meta=meta)
+                    )
+                except NotImplementedError:
+                    pass
+
+        # Add the message to the conversation.
         self.messages.append(Message(role=role, text=text, meta=meta))
 
     def send(
-        self, llm_model: Optional[str] = None, llm_provider: Optional[str] = None
+        self,
+        llm_model: str | None = None,
+        llm_provider: str | None = None,
     ) -> Message:
         """Send the conversation to the LLM."""
-        for plugin in self.plugins:
-            plugin.send_hook(self)
 
+        # TODO: llm_model and llm_provider should override the conversation's.
+
+        # Execute all pre send hooks.
+        for plugin in self.plugins:
+            if hasattr(plugin, "pre_send_hook"):
+                try:
+                    plugin.pre_send_hook(self)
+                except NotImplementedError:
+                    pass
+
+        # Find the provider and send the conversation.
         provider = find_provider(llm_provider or self.llm_provider)
         response = provider.send_conversation(self)
 
+        # Execute all post-send hooks.
+        for plugin in self.plugins:
+            if hasattr(plugin, "post_send_hook"):
+                try:
+                    plugin.post_send_hook(self, response)
+                except NotImplementedError:
+                    pass
+
+        # Add the response to the conversation.
         self.add_message(role="assistant", text=response.text, meta=response.meta)
+
         return response
 
-    def get_last_message(self, role: MESSAGE_ROLE) -> Optional[Message]:
+    def get_last_message(self, role: MESSAGE_ROLE) -> Message | None:
         """Get the last message with the given role."""
         return next((m for m in reversed(self.messages) if m.role == role), None)
 
-    def add_plugin(self, plugin: Any):
+    def add_plugin(self, plugin: BasePlugin) -> None:
         """Add a plugin to the conversation."""
         self.plugins.append(plugin)

simplemind/providers/__init__.py

@@ -1,9 +1,12 @@
 from typing import List, Type
 
-from simplemind.providers._base import BaseProvider
-from simplemind.providers.anthropic import Anthropic
-from simplemind.providers.groq import Groq
-from simplemind.providers.openai import OpenAI
-from simplemind.providers.xai import XAI
+from ._base import BaseProvider
+from .anthropic import Anthropic
+from .gemini import Gemini
+from .groq import Groq
+from .ollama import Ollama
+from .openai import OpenAI
+from .xai import XAI
+from .amazon import Amazon
 
-providers: List[Type[BaseProvider]] = [Anthropic, Groq, OpenAI, XAI]
+providers: List[Type[BaseProvider]] = [Anthropic, Gemini, Groq, OpenAI, Ollama, XAI, Amazon]

simplemind/providers/_base.py

@@ -1,6 +1,14 @@
 from abc import ABC, abstractmethod
+from functools import cached_property
+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)
 
 
 class BaseProvider(ABC):
@@ -8,14 +16,16 @@ class BaseProvider(ABC):
 
     NAME: str
     DEFAULT_MODEL: str
+    supports_streaming: bool = False
+    supports_structured_responses: bool = True
 
-    @property
+    @cached_property
     @abstractmethod
-    def client(self):
+    def client(self) -> Any:
         """The instructor client for the provider."""
         raise NotImplementedError
 
-    @property
+    @cached_property
     @abstractmethod
     def structured_client(self) -> Instructor:
         """The structured client for the provider."""
@@ -27,11 +37,11 @@ class BaseProvider(ABC):
         raise NotImplementedError
 
     @abstractmethod
-    def structured_response(self, prompt: str, response_model, **kwargs):
+    def structured_response(self, prompt: str, response_model: Type[T], **kwargs) -> T:
         """Get a structured response."""
         raise NotImplementedError
 
     @abstractmethod
-    def generate_text(self, prompt: str, **kwargs) -> str:
+    def generate_text(self, prompt: str, *, stream: bool = False, **kwargs) -> str:
         """Generate text from a prompt."""
         raise NotImplementedError

simplemind/providers/amazon.py

@@ -0,0 +1,116 @@
+from typing import Type, TypeVar
+from functools import cached_property
+
+import instructor
+from pydantic import BaseModel
+
+from ._base import BaseProvider
+from ..settings import settings
+
+T = TypeVar("T", bound=BaseModel)
+
+PROVIDER_NAME = "amazon"
+DEFAULT_MODEL = "anthropic.claude-3-sonnet-20240229-v1:0"
+DEFAULT_MAX_TOKENS = 5_000
+
+
+class Amazon(BaseProvider):
+    NAME = PROVIDER_NAME
+    DEFAULT_MODEL = DEFAULT_MODEL
+    supports_streaming = True
+
+    def __init__(self, profile_name: str | None = None):
+        self.profile_name = profile_name or settings.AMAZON_PROFILE_NAME
+
+    @cached_property
+    def client(self):
+        """The AnthropicBedrock client."""
+        import anthropic
+
+        if not self.profile_name:
+            raise ValueError("Profile name is not provided")
+
+        return anthropic.AnthropicBedrock(aws_profile=self.profile_name)
+
+    @cached_property
+    def structured_client(self):
+        """A client patched with Instructor."""
+
+        return instructor.from_anthropic(self.client)
+
+    def send_conversation(self, conversation: "Conversation", **kwargs):
+        """Send a conversation to the OpenAI API."""
+
+        from ..models import Message
+
+        messages = [
+            {"role": msg.role, "content": msg.text} for msg in conversation.messages
+        ]
+
+        response = self.client.chat.completions.create(
+            model=conversation.llm_model or DEFAULT_MODEL, messages=messages, **kwargs
+        )
+
+        # Get the response content from the OpenAI response
+        assistant_message = response.choices[0].message
+
+        # Create and return a properly formatted Message instance
+        return Message(
+            role="assistant",
+            text=assistant_message.content or "",
+            raw=response,
+            llm_model=conversation.llm_model or DEFAULT_MODEL,
+            llm_provider=PROVIDER_NAME,
+        )
+
+    def structured_response(
+        self, prompt, response_model: Type[T], *, llm_model: str | None = None, **kwargs
+    ) -> T:
+        # Ensure messages are provided in kwargs
+        messages = [
+            {"role": "user", "content": prompt},
+        ]
+
+        response = self.structured_client.chat.completions.create(
+            messages=messages,
+            model=llm_model or self.DEFAULT_MODEL,
+            response_model=response_model,
+            max_tokens=DEFAULT_MAX_TOKENS,
+            **kwargs,
+        )
+        return response
+
+    def generate_text(self, prompt, *, llm_model, **kwargs):
+        messages = [
+            {"role": "user", "content": prompt},
+        ]
+
+        response = self.client.messages.create(
+            model=llm_model or self.DEFAULT_MODEL,
+            messages=messages,
+            max_tokens=DEFAULT_MAX_TOKENS,
+            **kwargs,
+        )
+
+        return response.content[0].text
+
+    def generate_stream_text(self, prompt, *, llm_model, **kwargs):
+        """Generate streaming text using the Amazon API."""
+
+        # Prepare the messages.
+        messages = [
+            {"role": "user", "content": prompt},
+        ]
+
+        # Send the request to the API.
+        response = self.client.messages.create(
+            model=llm_model or self.DEFAULT_MODEL,
+            messages=messages,
+            stream=True,
+            **kwargs,
+        )
+
+        # Yield the text chunks.
+        for chunk in response:
+            if chunk.text:
+                yield chunk.text

simplemind/providers/anthropic.py

@@ -1,36 +1,55 @@
-from typing import Union
+from functools import cached_property
+from typing import TYPE_CHECKING, Type, TypeVar
 
-import anthropic
 import instructor
+from pydantic import BaseModel
 
-from ._base import BaseProvider
+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 = "anthropic"
 DEFAULT_MODEL = "claude-3-5-sonnet-20241022"
-DEFAULT_MAX_TOKENS = 1000
+DEFAULT_MAX_TOKENS = 1_000
+DEFAULT_KWARGS = {"max_tokens": DEFAULT_MAX_TOKENS}
 
 
 class Anthropic(BaseProvider):
     NAME = PROVIDER_NAME
     DEFAULT_MODEL = DEFAULT_MODEL
+    DEFAULT_KWARGS = DEFAULT_KWARGS
+    supports_streaming = True
 
-    def __init__(self, api_key: Union[str, None] = None):
+    def __init__(self, api_key: str | None = None):
         self.api_key = api_key or settings.get_api_key(PROVIDER_NAME)
 
-    @property
+    @cached_property
     def client(self):
         """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)
 
-    @property
+    @cached_property
     def structured_client(self):
         """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
 
@@ -39,10 +58,9 @@ class Anthropic(BaseProvider):
         ]
 
         response = self.client.messages.create(
-            model=conversation.llm_model or DEFAULT_MODEL,
+            model=conversation.llm_model or self.DEFAULT_MODEL,
             messages=messages,
-            max_tokens=DEFAULT_MAX_TOKENS,
-            **kwargs,
+            **{**self.DEFAULT_KWARGS, **kwargs},
         )
 
         # Get the response content from the Anthropic response
@@ -53,26 +71,57 @@ class Anthropic(BaseProvider):
             role="assistant",
             text=assistant_message,
             raw=response,
-            llm_model=conversation.llm_model or DEFAULT_MODEL,
+            llm_model=conversation.llm_model or self.DEFAULT_MODEL,
             llm_provider=PROVIDER_NAME,
         )
 
-    def structured_response(self, model, response_model, **kwargs):
-        response = self.structured_client.messages.create(
-            model=model, response_model=response_model, **kwargs
-        )
-        return response
+    @logger
+    def structured_response(
+        self, response_model: Type[T], *, llm_model: str | None = None, **kwargs
+    ) -> T:
+        model = llm_model or self.DEFAULT_MODEL
 
-    def generate_text(self, prompt, *, llm_model, **kwargs):
+        # Extract the prompt from kwargs if it exists
+        prompt = kwargs.pop("prompt", kwargs.pop("messages", ""))
+
+        # Format the messages properly
+        messages = [{"role": "user", "content": prompt}]
+
+        response = self.structured_client.messages.create(
+            model=model,
+            messages=messages,  # Add the messages parameter
+            response_model=response_model,
+            **{**self.DEFAULT_KWARGS, **kwargs},
+        )
+        return response_model.model_validate(response)
+
+    @logger
+    def generate_text(self, prompt: str, *, llm_model: str, **kwargs):
         messages = [
             {"role": "user", "content": prompt},
         ]
 
         response = self.client.messages.create(
-            model=llm_model,
+            model=llm_model or self.DEFAULT_MODEL,
             messages=messages,
-            max_tokens=DEFAULT_MAX_TOKENS,
-            **kwargs,
+            **{**self.DEFAULT_KWARGS, **kwargs},
         )
 
         return response.content[0].text
+
+    @logger
+    def generate_stream_text(self, prompt: str, *, llm_model: str, **kwargs):
+        # Prepare the messages.
+        messages = [
+            {"role": "user", "content": prompt},
+        ]
+
+        # Make the request.
+        with self.client.messages.stream(
+            model=llm_model or self.DEFAULT_MODEL,
+            messages=messages,
+            **{**self.DEFAULT_KWARGS, **kwargs},
+        ) as stream:
+            # Yield each chunk of text from the stream.
+            for chunk in stream.text_stream:
+                yield chunk

simplemind/providers/gemini.py

@@ -0,0 +1,124 @@
+# TODO: this is a placeholder file for the Gemini provider
+# IT is not currently working as desired.
+
+from functools import cached_property
+from typing import TYPE_CHECKING, Type, TypeVar
+
+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)
+
+
+PROVIDER_NAME = "gemini"
+DEFAULT_MODEL = "models/gemini-1.5-flash-latest"
+
+
+class Gemini(BaseProvider):
+    NAME = PROVIDER_NAME
+    DEFAULT_MODEL = DEFAULT_MODEL
+    supports_streaming = True
+
+    def __init__(self, api_key: str | None = None):
+        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):
+        """The raw Gemini client."""
+        if not self.api_key:
+            raise ValueError("Gemini API key is required")
+        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
+    def structured_client(self):
+        """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
+
+        # Convert messages to Gemini's format
+        chat = self.client.start_chat()
+
+        # Send all previous messages to establish context
+        for msg in conversation.messages[:-1]:  # All messages except the last one
+            chat.send_message(msg.text)
+
+        # Send the final message and get response
+        try:
+            response = chat.send_message(conversation.messages[-1].text)
+        except Exception as e:
+            raise RuntimeError(f"Failed to send conversation to Gemini API: {e}") from e
+
+        # Create and return a properly formatted Message instance
+        return Message(
+            role="assistant",
+            text=response.text,
+            raw=response,
+            llm_model=self.model_name,
+            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."""
+        # 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(
+                messages=[{"role": "user", "content": prompt}],
+                response_model=response_model,
+                **kwargs,
+            )
+        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_model.model_validate(response)
+
+    @logger
+    def generate_text(self, prompt: str, **kwargs) -> str:
+        """Generate text using the Gemini API."""
+        kwargs.pop("llm_model")
+        try:
+            response = self.client.generate_content(prompt, **kwargs)
+        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 response.text
+
+    @logger
+    def generate_stream_text(self, prompt: str, **kwargs) -> str:
+        """Generate streaming text using the Gemini API."""
+        kwargs.pop("llm_model", None)
+        try:
+            response = self.client.generate_content(prompt, stream=True, **kwargs)
+            for chunk in response:
+                if chunk.text:
+                    yield chunk.text
+        except Exception as e:
+            raise RuntimeError(
+                f"Failed to generate streaming text with Gemini API: {e}"
+            ) from e

simplemind/providers/groq.py

@@ -1,34 +1,53 @@
-from typing import Union
+from functools import cached_property
+from typing import TYPE_CHECKING, Type, TypeVar
 
-import groq
 import instructor
+from pydantic import BaseModel
 
-from ._base import BaseProvider
+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 = "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
+    supports_streaming = True
 
-    def __init__(self, api_key: Union[str, None] = None):
+    def __init__(self, api_key: str | None = None):
         self.api_key = api_key or settings.get_api_key(PROVIDER_NAME)
 
-    @property
+    @cached_property
     def client(self):
         """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)
 
-    @property
+    @cached_property
     def structured_client(self):
         """A client patched with Instructor."""
         return instructor.from_groq(self.client)
 
+    @logger
     def send_conversation(
         self,
         conversation: "Conversation",
@@ -42,9 +61,9 @@ class Groq(BaseProvider):
         ]
 
         response = self.client.chat.completions.create(
-            model=conversation.llm_model or DEFAULT_MODEL,
+            model=conversation.llm_model or self.DEFAULT_MODEL,
             messages=messages,
-            **kwargs,
+            **{**self.DEFAULT_KWARGS, **kwargs},
         )
 
         # Get the response content from the Groq response
@@ -55,11 +74,12 @@ class Groq(BaseProvider):
             role="assistant",
             text=assistant_message.content or "",
             raw=response,
-            llm_model=conversation.llm_model or DEFAULT_MODEL,
+            llm_model=conversation.llm_model or self.DEFAULT_MODEL,
             llm_provider=PROVIDER_NAME,
         )
 
-    def structured_response(self, prompt: str, response_model, **kwargs):
+    @logger
+    def structured_response(self, prompt: str, response_model: Type[T], **kwargs) -> T:
         # Ensure messages are provided in kwargs
         messages = [
             {"role": "user", "content": prompt},
@@ -68,25 +88,56 @@ class Groq(BaseProvider):
         response = self.structured_client.chat.completions.create(
             messages=messages,
             response_model=response_model,
-            **kwargs,
+            model=kwargs.pop("llm_model", self.DEFAULT_MODEL),
+            **{**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},
         ]
 
-        response = self.structured_client.chat.completions.create(
+        response = self.client.chat.completions.create(
             messages=messages,
-            model=llm_model,
-            **kwargs,
+            model=llm_model or self.DEFAULT_MODEL,
+            **{**self.DEFAULT_KWARGS, **kwargs},
         )
 
-        return response.choices[0].message.content
+        return str(response.choices[0].message.content)
+
+    @logger
+    def generate_stream_text(
+        self,
+        prompt: str,
+        *,
+        llm_model: str | None = None,
+        **kwargs,
+    ) -> str:
+        """Generate streaming text using the Groq API."""
+        messages = [
+            {"role": "user", "content": prompt},
+        ]
+
+        response = self.client.chat.completions.create(
+            messages=messages,
+            model=llm_model or self.DEFAULT_MODEL,
+            stream=True,
+            **{**self.DEFAULT_KWARGS, **kwargs},
+        )
+
+        try:
+            for chunk in response:
+                if chunk.choices and chunk.choices[0].delta.content:
+                    yield chunk.choices[0].delta.content
+        except Exception as e:
+            raise RuntimeError(
+                f"Failed to generate streaming text with Groq API: {e}"
+            ) from e

simplemind/providers/ollama.py

@@ -0,0 +1,137 @@
+from functools import cached_property
+from typing import TYPE_CHECKING, Type, TypeVar
+
+import instructor
+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
+    supports_streaming = True
+
+    def __init__(self, host_url: str | None = None):
+        self.host_url = host_url or settings.OLLAMA_HOST_URL
+
+    @cached_property
+    def client(self):
+        """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
+    def structured_client(self) -> instructor.Instructor:
+        """A client patched with Instructor."""
+        return instructor.from_openai(
+            OpenAI(
+                base_url=f"{self.host_url}/v1",
+                api_key="ollama",
+            ),
+            mode=instructor.Mode.JSON,
+        )
+
+    @logger
+    def send_conversation(self, conversation: "Conversation", **kwargs) -> "Message":
+        """Send a conversation to the Ollama API."""
+        from ..models import Message
+
+        messages = [
+            {"role": msg.role, "content": msg.text} for msg in conversation.messages
+        ]
+        response = self.client.chat(
+            model=conversation.llm_model or DEFAULT_MODEL,
+            messages=messages,
+            **{**self.DEFAULT_KWARGS, **kwargs},
+        )
+        assistant_message = response.get("message")
+
+        # Create and return a properly formatted Message instance
+        return Message(
+            role="assistant",
+            text=assistant_message.get("content"),
+            raw=response,
+            llm_model=conversation.llm_model or self.DEFAULT_MODEL,
+            llm_provider=PROVIDER_NAME,
+        )
+
+    @logger
+    def structured_response(
+        self,
+        prompt: str,
+        response_model: Type[T],
+        *,
+        llm_model: str | None = None,
+        **kwargs,
+    ) -> T:
+        """Get a structured response from the Ollama API."""
+        messages = [
+            {"role": "user", "content": prompt},
+        ]
+
+        response = self.structured_client.chat.completions.create(
+            messages=messages,
+            model=llm_model or self.DEFAULT_MODEL,
+            response_model=response_model,
+            **{**self.DEFAULT_KWARGS, **kwargs},
+        )
+        return response_model.model_validate(response)
+
+    @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,
+            **{**self.DEFAULT_KWARGS, **kwargs},
+        )
+
+        return response.get("message", {}).get("content", "")
+
+    @logger
+    def generate_stream_text(self, prompt: str, *, llm_model: str, **kwargs) -> str:
+        # Prepare the messages.
+        messages = [
+            {"role": "user", "content": prompt},
+        ]
+
+        response = self.client.chat(
+            messages=messages,
+            model=llm_model or self.DEFAULT_MODEL,
+            stream=True,
+            **{**self.DEFAULT_KWARGS, **kwargs},
+        )
+
+        # Iterate over the response and yield the content.
+        for chunk in response:
+            yield chunk["message"]["content"]

simplemind/providers/openai.py

@@ -1,35 +1,54 @@
-from typing import Union
+from functools import cached_property
+from typing import TYPE_CHECKING, Type, TypeVar
 
 import instructor
-import openai as oa
+from pydantic import BaseModel
 
-from ._base import BaseProvider
+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
+    supports_streaming = True
+    
 
-    def __init__(self, api_key: Union[str, None] = None):
+    def __init__(self, api_key: str | None = None):
         self.api_key = api_key or settings.get_api_key(PROVIDER_NAME)
 
-    @property
+    @cached_property
     def client(self):
         """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)
 
-    @property
+    @cached_property
     def structured_client(self):
         """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
 
@@ -38,7 +57,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
@@ -53,24 +74,59 @@ class OpenAI(BaseProvider):
             llm_provider=PROVIDER_NAME,
         )
 
-    def structured_response(self, prompt, response_model, *, llm_model: str, **kwargs):
+    @logger
+    def structured_response(
+        self,
+        prompt: str,
+        response_model: Type[T],
+        *,
+        llm_model: str | None = None,
+        **kwargs,
+    ) -> T:
+        """Get a structured response from the OpenAI API."""
         # Ensure messages are provided in kwargs
         messages = [
             {"role": "user", "content": prompt},
         ]
-
         response = self.structured_client.chat.completions.create(
-            messages=messages, model=llm_model, response_model=response_model, **kwargs
+            messages=messages,
+            model=llm_model or self.DEFAULT_MODEL,
+            response_model=response_model,
+            **{**self.DEFAULT_KWARGS, **kwargs},
         )
-        return response
+        return response_model.model_validate(response)
 
-    def generate_text(self, prompt, *, llm_model, **kwargs):
+    @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, **kwargs
+            messages=messages,
+            model=llm_model or self.DEFAULT_MODEL,
+            **{**self.DEFAULT_KWARGS, **kwargs},
+        )
+        return response.choices[0].message.content
+
+    @logger
+    def generate_stream_text(
+        self, prompt: str, *, llm_model: str | None = None, **kwargs
+    ):
+        """Generate streaming text using the OpenAI API.
+
+        Yields chunks of text as they are generated by the model.
+        """
+        messages = [
+            {"role": "user", "content": prompt},
+        ]
+        response = self.client.chat.completions.create(
+            messages=messages,
+            model=llm_model or self.DEFAULT_MODEL,
+            stream=True,  # Enable streaming
+            **{**self.DEFAULT_KWARGS, **kwargs},
         )
 
-        return response.choices[0].message.content
+        for chunk in response:
+            if chunk.choices[0].delta.content is not None:
+                yield chunk.choices[0].delta.content

simplemind/providers/xai.py

@@ -1,40 +1,59 @@
-from typing import Union
+from functools import cached_property
+from typing import TYPE_CHECKING, Type, TypeVar
 
 import instructor
-import openai as oa
+from pydantic import BaseModel
 
-from ._base import BaseProvider
+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
+    supports_streaming = True
+    supports_structured_responses = False
 
-    def __init__(self, api_key: Union[str, None] = None):
+    def __init__(self, api_key: str | None = None):
         self.api_key = api_key or settings.get_api_key(PROVIDER_NAME)
 
-    @property
+    @cached_property
     def client(self):
         """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,
         )
 
-    @property
+    @cached_property
     def structured_client(self):
         """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
 
@@ -43,9 +62,9 @@ class XAI(BaseProvider):
         ]
 
         response = self.client.chat.completions.create(
-            model=conversation.llm_model or DEFAULT_MODEL,
+            model=conversation.llm_model or self.DEFAULT_MODEL,
             messages=messages,
-            **kwargs,
+            **{**self.DEFAULT_KWARGS, **kwargs},
         )
 
         # Get the response content from the OpenAI response
@@ -56,22 +75,48 @@ class XAI(BaseProvider):
             role="assistant",
             text=assistant_message.content,
             raw=response,
-            llm_model=conversation.llm_model or DEFAULT_MODEL,
+            llm_model=conversation.llm_model or self.DEFAULT_MODEL,
             llm_provider=PROVIDER_NAME,
         )
 
-    def structured_response(self, prompt: str, response_model, *, llm_model):
+    @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, *, llm_model, **kwargs):
+    @logger
+    def generate_text(self, prompt: str, *, llm_model: str, **kwargs) -> str:
+        # Prepare the messages.
         messages = [
             {"role": "user", "content": prompt},
         ]
 
+        # Make the request.
         response = self.client.chat.completions.create(
             messages=messages,
-            model=llm_model,
-            **kwargs,
+            model=llm_model or self.DEFAULT_MODEL,
+            **{**self.DEFAULT_KWARGS, **kwargs},
         )
 
-        return response.choices[0].message.content
+        # Return the response content.
+        return str(response.choices[0].message.content)
+
+    @logger
+    def generate_stream_text(self, prompt: str, *, llm_model: str, **kwargs) -> str:
+        # Prepare the messages.
+        messages = [
+            {"role": "user", "content": prompt},
+        ]
+
+        # Make the request.
+        response = self.client.chat.completions.create(
+            messages=messages,
+            model=llm_model or self.DEFAULT_MODEL,
+            stream=True,
+            **{**self.DEFAULT_KWARGS, **kwargs},
+        )
+
+        # Iterate over the response and yield the content.
+        for chunk in response:
+            yield chunk.choices[0].delta.content

simplemind/settings.py

@@ -4,20 +4,61 @@ from pydantic import Field, SecretStr, field_validator
 from pydantic_settings import BaseSettings, SettingsConfigDict
 
 
+class LoggingConfig(BaseSettings):
+    """The class that holds all the logging settings for the application."""
+
+    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."""
 
+    AMAZON_PROFILE_NAME: Optional[str] = Field(
+        "default", description="AWS Named Profile"
+    )
     ANTHROPIC_API_KEY: Optional[SecretStr] = Field(
         None, description="API key for Anthropic"
     )
     GROQ_API_KEY: Optional[SecretStr] = Field(None, description="API key for Groq")
+    GEMINI_API_KEY: Optional[SecretStr] = Field(None, description="API key for Gemini")
     OPENAI_API_KEY: Optional[SecretStr] = Field(None, description="API key for OpenAI")
+    OLLAMA_HOST_URL: Optional[str] = Field(
+        "http://127.0.0.1:11434", description="Fully qualified host URL for Ollama"
+    )
     XAI_API_KEY: Optional[SecretStr] = Field(None, description="API key for xAI")
     DEFAULT_LLM_PROVIDER: str = Field("openai", description="The default LLM provider")
 
     model_config = SettingsConfigDict(
         env_file=".env", env_file_encoding="utf-8", case_sensitive=True, extra="ignore"
     )
+    logging: LoggingConfig = LoggingConfig()
 
     @field_validator("*", mode="before")
     @classmethod

simplemind/utils.py

@@ -1,13 +1,38 @@
-from typing import Union
+import difflib
 
-from .providers import providers
+from .providers import BaseProvider, providers
+
+_PROVIDER_NAMES = [provider.NAME.lower() for provider in providers]
 
 
-def find_provider(provider_name: Union[str, None]):
-    """Find a provider by name."""
-    if provider_name:
-        for provider_class in providers:
-            if provider_class.NAME.lower() == provider_name.lower():
-                # Instantiate the provider
-                return provider_class()
-    raise ValueError(f"Provider {provider_name} not found")
+def find_provider(provider_name: str | None) -> BaseProvider:
+    """
+    Find and instantiate a provider by name.
+
+    Parameters:
+        provider_name (Union[str, None]): The name of the provider to find.
+
+    Returns:
+        An instance of the provider class if found.
+
+    Raises:
+        ValueError: If the provider is not specified or is not found, with a suggestion for the closest match.
+    """
+    if provider_name is None:
+        raise ValueError("No provider specified.")
+
+    # Find the provider by name.
+    for provider_class in providers:
+        if provider_class.NAME.lower() == provider_name.lower():
+            # Instantiate the provider
+            return provider_class()
+
+    # Find the closest match
+    provider_found = difflib.get_close_matches(
+        provider_name.lower(), _PROVIDER_NAMES, n=1
+    )
+    if provider_found:
+        raise ValueError(
+            f"Provider {provider_name!r} not found. Did you mean {provider_found[0]!r}?"
+        )
+    raise ValueError(f"Provider {provider_name} not found.")

tests/conftest.py

@@ -0,0 +1,15 @@
+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__))))
+
+from simplemind import Session
+
+
+@pytest.fixture
+def sm():
+    """Fixture that provides a simplemind Session instance with default settings."""
+    return Session()

tests/test_basics.py

@@ -0,0 +1,2 @@
+def test_basic_math():
+    assert 1 + 1 == 2

tests/test_conversations.py

@@ -0,0 +1,28 @@
+import pytest
+
+from simplemind.providers import Anthropic, Gemini, OpenAI, Groq, Ollama, Amazon
+
+import simplemind as sm
+
+
+@pytest.mark.parametrize(
+    "provider_cls",
+    [
+        Anthropic,
+        Gemini,
+        OpenAI,
+        Groq,
+        Ollama,
+        # Amazon
+    ],
+)
+def test_generate_data(provider_cls):
+    conv = sm.create_conversation(
+        llm_model=provider_cls.DEFAULT_MODEL, llm_provider=provider_cls.NAME
+    )
+
+    conv.add_message(text="hey")
+    data = conv.send()
+
+    assert isinstance(data.text, str)
+    assert len(data.text) > 0

tests/test_generate_data.py

@@ -0,0 +1,30 @@
+import pytest
+
+from simplemind.providers import Anthropic, Gemini, OpenAI, Groq, Ollama, Amazon
+
+from pydantic import BaseModel
+
+
+class ResponseModel(BaseModel):
+    result: int
+
+
+@pytest.mark.parametrize(
+    "provider_cls",
+    [
+        Anthropic,
+        Gemini,
+        OpenAI,
+        Groq,
+        Ollama,
+        # Amazon
+    ],
+)
+def test_generate_data(provider_cls):
+    provider = provider_cls()
+    prompt = "What is 2+2?"
+
+    data = provider.structured_response(prompt=prompt, response_model=ResponseModel)
+
+    assert isinstance(data, ResponseModel)
+    assert isinstance(data.result, int)

tests/test_generate_text.py

@@ -0,0 +1,24 @@
+import pytest
+
+from simplemind.providers import Anthropic, Gemini, OpenAI, Groq, Ollama, Amazon
+
+
+@pytest.mark.parametrize(
+    "provider_cls",
+    [
+        Anthropic,
+        Gemini,
+        OpenAI,
+        Groq,
+        Ollama,
+        # Amazon,
+    ],
+)
+def test_generate_text(provider_cls):
+    provider = provider_cls()
+    prompt = "What is 2+2?"
+
+    response = provider.generate_text(prompt=prompt, llm_model=provider.DEFAULT_MODEL)
+
+    assert isinstance(response, str)
+    assert len(response) > 0