Building Robust AI-Powered Agents for Third Parties API: From Basic Setup to Scalable, Validated Architecture

Integrating AI with Notion can be transformative — imagine an AI assistant that understands your inputs and updates your Notion workspace with data, ideas, or even haikus! Starting with a basic setup is effective, but as functionality grows, it’s crucial to build a scalable, flexible, and validated system. This guide takes you through setting up a simple haiku generator, then elevates it into a robust, modular, and validated agent that can adapt to changing requirements.

Code on this GitHub repo

— -

Part 1: Setting Up a Naive Integration

We begin with a simple AI-driven tool that allows the user to input a haiku, which is then inserted into a Notion page. This is straightforward and effective for a single-purpose, small-scale application.

How the Naive Setup Works

1. API Access Setup:

• Use OpenAI’s API to generate haikus and Notion’s API to save them.
• Store API keys securely in environment variables.

# Load environment variables
load_dotenv()

# OpenAI and Notion API setup
client = OpenAI(api_key=os.getenv('OPENAI_API_KEY'))
notion_api_key = os.getenv('NOTION_API_KEY')
notion_page_id = os.getenv('NOTION_PAGE_ID')
openai_model = os.getenv('OPENAI_MODEL')

headers = {
    "Authorization": f"Bearer {notion_api_key}",
    "Content-Type": "application/json",
    "Notion-Version": "2022-06-28"
}

2. Basic Tool Function:

• Define Tool Function: We create a single function to insert haikus into Notion.

def insert_haiku(haiku: str) -> str:
    """
    Inserts a haiku as a new page in Notion under a specified parent page.

    Parameters
    ----------
    haiku : str
        The haiku content to be inserted into Notion.

    Returns
    -------
    str
        The result of the API call. If successful, a confirmation message
        is returned; otherwise, an error message is returned.
    """
    
    haiku_body = {
        "parent": {"page_id": notion_page_id},
        "properties": {
            "title": {  # Correct title property for Notion API
                "title": [
                    {
                        "type": "text",
                        "text": {
                            "content": haiku
                        }
                    }
                ]
            }
        }
    }

    response = requests.post('https://api.notion.com/v1/pages', headers=headers, json=haiku_body)

    if response.status_code == 200:
        return "Haiku successfully inserted in Notion!"
    else:
        return f"Failed to insert haiku: {response.status_code}, {response.text}"

3. Prompt and Tool Invocation:

• Prompt and Tool Invocation: The function prompts OpenAI for user input, checks if it should invoke the tool, and inserts the haiku if requested.

def prompt_ai(messages: list) -> str:
    """
    Handles user messages and determines if a tool needs to be called by
    the OpenAI model.

    Parameters
    ----------
    messages : list
        A list of conversation messages between the user and the assistant.

    Returns
    -------
    str
        The assistant's response after processing the message or invoking
        a function.
    """
    # First, prompt the AI with the latest user message
    completion = client.chat.completions.create(
        model=openai_model,  # Make sure to specify the correct model
        messages=messages,  # No need for a separate 'prompt' parameter
        tools=get_tools()
    )

    response_message = completion.choices[0].message
    tool_calls = response_message.tool_calls

    # Second, see if the AI decided it needs to invoke a tool
    if tool_calls:
        # If the AI decided to invoke a tool, invoke it
        available_functions = {
            "insert_haiku": insert_haiku,
        }

        # Add the tool request to the list of messages so the AI knows later it invoked the tool
        messages.append(response_message)

        # Next, for each tool the AI wanted to call, call it and add the tool result to the list of messages
        for tool_call in tool_calls:
            function_name = tool_call.function.name
            function_to_call = available_functions[function_name]
            function_args = json.loads(tool_call.function.arguments)
            function_response = function_to_call(**function_args)

            messages.append({
                "tool_call_id": tool_call.id,
                "role": "tool",
                "name": function_name,
                "content": function_response
            })

        # Call the AI again so it can produce a response with the result of calling the tool(s)
        second_response = client.chat.completions.create(
            model=openai_model,
            messages=messages,
        )

        return second_response.choices[0].message.content

    return response_message.content

Why Upgrade? Limitations of the Naive Approach

While this simple setup is a great starting point, it has drawbacks that limit scalability:

• No Input Validation: There’s no structure to ensure correct input data, so erroneous input could break the function.
• Single-Purpose Design: It’s restricted to creating haikus, making it difficult to add other features.
• Locked to One Provider: Tied to OpenAI, so adapting to other providers would require significant changes.

Without a more modular setup, every new feature could turn into a major overhaul. To scale effectively, we need a more structured approach.

— -

Part 2: Upgrading with LiteLLM, Instructor, and UV

Provider Abstraction with LiteLLM

• Flexibility: LiteLLM decouples the assistant from OpenAI, allowing you to switch providers by changing configuration settings, not code.
• Cleaner Code: Centralized language model calls mean your main assistant logic focuses on tasks, not provider details.

Input Validation with Instructor and Pydantic

• Structured Data Models: Define Pydantic models for each action (e.g., HaikuRequest, WeatherRequest) to validate inputs, reducing runtime errors and ensuring data quality.
• Dynamic Tool Selection: Instructor automatically chooses the right tool based on validated inputs, so adding new tools (like checking weather or logging tasks) is easy.

Smooth Execution and Orchestration with UV

• Environment Management: Run your project with consistent dependencies using a single uv run command.
• Easy Deployment and Reproducibility: Ensures everyone on your team has the same environment, minimizing setup time and dependency issues.

1. Validated Configuration

We leverage `dataclass` to simplify the Agent configuration

@dataclass
class AgentConfig:
    """Configuration for the AI Assistant.

    Attributes
    ----------
    openai_api_key : str
        The OpenAI API key for accessing language models.
    notion_api_key : str
        The Notion API key for accessing the Notion API.
    notion_page_id : str
        The Notion page ID where data will be inserted.
    openai_model : str
        The OpenAI model to use for completions.
    """
    openai_api_key: str = os.getenv('OPENAI_API_KEY')
    notion_api_key: str = os.getenv('NOTION_API_KEY')
    notion_page_id: str = os.getenv('NOTION_PAGE_ID')
    openai_model: str = os.getenv('OPENAI_MODEL', 'gpt-3.5-turbo')  # Default model

2. Abstracting the Provider with `LiteLLM`

With `LiteLLM`, we decouple our code from OpenAI’s API, replacing direct calls with an interface that can interact with different language model providers and enabling lazy initialization for performance.

    @property
    def instructor_client(self):
        """Lazily initializes and returns the instructor client.

        Returns
        -------
        InstructorClient
            The client for interacting with Instructor API.
        """
        if self._instructor_client is None:
            self._instructor_client = instructor.from_litellm(litellm.completion)
        return self._instructor_client

- Provider Flexibility: The assistant is no longer tied to OpenAI. If new language model providers (like Anthropic, Cohere, or custom models) become preferable, we can switch providers by changing the configuration instead of rewriting the code.
- Code Readability and Maintainability: By centralizing language model calls, `LiteLLM` keeps the main assistant logic clean and focused on task handling, not API specifics.

This abstraction lays the groundwork for a truly modular AI system that can evolve as technology and business needs change, a significant advantage over a single-provider setup.

3. Validating Input with `Instructor` and `Pydantic`

Validation is vital for scaling, especially as the assistant takes on more complex actions. Here’s how we approach validation:

• Define Action Models with `Pydantic`: Instead of allowing free-form strings, we define structured data models for each action (e.g., a `HaikuRequest` model and a `WeatherRequest` model). These models specify expected inputs and allow automated checks, ensuring that only valid data reaches each tool function.
• Dynamic Tool Handling with `Instructor`: `Instructor` automates tool selection, allowing the AI to choose the appropriate tool based on the user’s input and the models defined. Each action model has its schema, so `Instructor` can dynamically identify which tool to invoke based on input validation, making it easy to add new actions (like checking the weather or adding to-do items) without disrupting existing functions.

# Action models
# Instructor leverages pydantic to extract the arguments 
# to pass to the actions from the User's prompt
class HaikuRequest(BaseModel):
    """Model for a request to create a haiku.

    Attributes
    ----------
    text : str
        The text of the haiku to insert into Notion.
    title : str, optional
        The title for the haiku in Notion, defaults to "Haiku".
    """
    text: str = Field(..., description="The haiku text to insert into Notion")
    title: str = Field("Haiku", description="Title for the haiku in Notion")

class WeatherRequest(BaseModel):
    """Model for a request to retrieve weather data.

    Attributes
    ----------
    location : str
        The location for which to retrieve the weather data.
    """
    location: str = Field(..., description="The location to retrieve the weather for")

# Parent model to support multiple actions
class ActionModel(BaseModel):
    """Parent model to handle multiple actions in a single request.

    Attributes
    ----------
    actions : list of Union[HaikuRequest, WeatherRequest]
        List of requested actions to be performed by the assistant.
    """
    actions: List[Union[HaikuRequest, WeatherRequest]] = Field(
        ..., description="List of requested actions"
    )

• Input Validation: Ensures that only valid data reaches each tool, improving robustness. If a user input doesn’t meet the model’s criteria, `Instructor` will detect this early, reducing the chance of runtime errors.
• Modularity: Each tool (haiku creation, weather check, etc.) has its own model and handler function, which are isolated from each other. This makes it easy to add, modify, or remove tools as needed.
• Error Handling: Errors are caught at the model validation level, making debugging easier and improving error feedback for end users.

The resultant is a greatly simplified Agent implementation:

class AIAssistant:
    """AI Assistant that interacts with OpenAI's API and handles multiple actions.

    Attributes
    ----------
    config : AgentConfig
        The configuration object containing API keys and model settings.
    actions_dispatch : dict
        A dictionary mapping action models to their handler functions.
    """

    def __init__(self, config: AgentConfig):
        self.config = config
        self._instructor_client = None
        # we define the list of actions accessible to the AI Agent
        self.actions_dispatch: Dict[type, Callable[[BaseModel], str]] = {
            HaikuRequest: lambda action: handle_haiku_request(action, self.config.notion_api_key, self.config.notion_page_id),
            WeatherRequest: handle_weather_request,
        }
    
    @property
    def instructor_client(self):
        """Lazily initializes and returns the instructor client.

        Returns
        -------
        InstructorClient
            The client for interacting with Instructor API.
        """
        if self._instructor_client is None:
            self._instructor_client = instructor.from_litellm(litellm.completion)
        return self._instructor_client

    def prompt_ai(self, messages: List[Dict[str, Any]]) -> List[str]:
        """Processes user messages and invokes actions based on AI responses.

        Parameters
        ----------
        messages : list of dict
            The conversation history between the user and the assistant.

        Returns
        -------
        list of str
            List of responses from the assistant after processing actions.
        """
        completion = self.instructor_client.chat.completions.create(
            model=self.config.openai_model,
            messages=messages,
            response_model=ActionModel  # Accepts multiple actions
        )

        results = []
        for action in completion.actions:
            action_type = type(action)
            if action_type in self.actions_dispatch:
                results.append(self.actions_dispatch[action_type](action))
            else:
                results.append(f"No handler found for action type: {action_type.__name__}")
        return results

4. `UV` for Dependency Management and Execution

While `LiteLLM` and `Instructor` enhance modularity and validation, managing dependencies and executing the script can still be challenging. This is where `UV` shines. For more details, I covered UV in a 2 parts story: UV part 1 and UV part 2.

Advantages of Using `UV`:

• Simplified Dependency Management: With `UV`, you add dependency extremely fast.
• Portable and Isolated Environments: `UV` enables the script to run in a container-like environment, ensuring that dependencies don’t conflict with other projects and are version-controlled.
• Easy Execution and Reproducibility: Running `uv run` automatically sets up and activates the environment, ensuring consistency across machines and making it easy for new team members to onboard without complex setup steps.

By incorporating `UV`, our AI assistant becomes easier to deploy, maintain, and scale across different environments and team members.

Why This Approach is Ideal: Scalability, Validation, and Ease of Management

Upgrading from a basic setup to one that leverages LiteLLM, Instructor, and UV, you gain significant advantages in structure, flexibility, and usability.

Key Benefits:

1. Scalability: Modular tools, validation, and abstraction enable you to easily add new models or providers without affecting existing features.
2. Robustness: Input validation and dynamic tool handling guarantee data integrity and protect against errors and unexpected user input.
3. Simplified Management: UV ensures consistent, isolated environments that are simple to set up, manage, and reproduce, making your assistant production-ready and suitable for team-wide deployment.

Conclusion

Building a powerful AI-powered assistant for Notion (or any third-party API) involves a balance of simplicity and sophistication. By integrating LiteLLM for provider abstraction, Instructor with Pydantic for validation, and UV for streamlined execution and dependency management, you can create a robust and scalable agent. This structured approach not only future-proofs your AI applications but also empowers you to confidently expand functionality while maintaining reliability and reproducibility.