chapter_05_ai_agent_workflows.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Chapter 05: AI Agent Workflows\n",
    "\n",
    "In our previous discussions, we explored the utility of structured outputs and AI agents in the context of chat applications. Now, we take a leap forward by combining these foundational elements with workflow orchestration, leveraging LLMs. This enables us to build highly sophisticated workflows, integrating multiple steps and diverse components such as functions, data transformations, machine learning models, vision models, and more. The result is an output in a predefined structure that can be integrated into business systems like databases, ERP systems, and other organizational infrastructure.  \n",
    "\n",
    "The integration of LLMs into these workflows offers a unique advantage: the ability to handle unstructured data and convert it into structured formats that can be effortlessly pushed into various systems. This capability addresses a pivotal challenge faced by businesses today - efficiently managing and utilizing the vast amounts of unstructured data generated daily.  \n",
    "\n",
    "Moreover, these workflows can be automated to react to specific triggers, such as the addition of a new object in a system or the receipt of an alert. Such automation not only enhances operational efficiency but also ensures timely and accurate decision-making. The output generated from these workflows can either be written directly into the systems or subjected to a human-in-the-loop approval process, ensuring the perfect balance between automation and human oversight."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import annotations\n",
    "import sys\n",
    "from typing import Any, Callable, Literal\n",
    "from functools import reduce\n",
    "from pydantic import BaseModel, Field\n",
    "from loguru import logger\n",
    "from language_models.agent import Agent, OutputType, PromptingStrategy\n",
    "from language_models.models.llm import OpenAILanguageModel\n",
    "from language_models.proxy_client import ProxyClient\n",
    "from language_models.settings import settings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "logger.remove()\n",
    "logger.add(sys.stderr, format=\"{message}\", level=\"INFO\")\n",
    "\n",
    "proxy_client = ProxyClient(\n",
    "    client_id=settings.CLIENT_ID,\n",
    "    client_secret=settings.CLIENT_SECRET,\n",
    "    auth_url=settings.AUTH_URL,\n",
    "    api_base=settings.API_BASE,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "llm = OpenAILanguageModel(\n",
    "    proxy_client=proxy_client,\n",
    "    model=\"gpt-4\",\n",
    "    max_tokens=250,\n",
    "    temperature=0.2,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In our simple workflow design, we have integrated distinct components:\n",
    "- **`WorkflowFunctionStep`:** Manages specific function invocations.\n",
    "- **`WorkflowLLMStep`:** Leverages LLMs for tasks.\n",
    "- **`WorkflowTransformationStep`:** Focuses on data transformations, offering mechanisms like map, filter, and reduce to efficiently process and restructure data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class WorkflowStepOutput(BaseModel):\n",
    "    \"\"\"Class that represents the output of a step.\"\"\"\n",
    "\n",
    "    inputs: dict[str, Any]\n",
    "    output: (\n",
    "        str\n",
    "        | int\n",
    "        | float\n",
    "        | dict[str, Any]\n",
    "        | BaseModel\n",
    "        | list[str]\n",
    "        | list[int]\n",
    "        | list[float]\n",
    "        | list[dict[str, Any]]\n",
    "        | list[BaseModel]\n",
    "        | None\n",
    "    )\n",
    "\n",
    "class WorkflowFunctionStep(BaseModel):\n",
    "    \"\"\"Class that implements a function step.\n",
    "\n",
    "    Attributes:\n",
    "        name: The name of the step.\n",
    "        inputs: The Pydantic model that represents the input arguments.\n",
    "        function: The function that will be invoked when calling this step.\n",
    "    \"\"\"\n",
    "\n",
    "    name: str\n",
    "    inputs: type[BaseModel]\n",
    "    function: Callable[[Any], Any]\n",
    "\n",
    "    def invoke(self, inputs: dict[str, Any], verbose: bool) -> WorkflowStepOutput:\n",
    "        if verbose:\n",
    "            logger.opt(colors=True).info(f\"<b><fg #2D72D2>Use Function</fg #2D72D2></b>: {self.name}\")\n",
    "\n",
    "        inputs = {key: value for key, value in inputs.items() if key in self.inputs.model_fields}\n",
    "        if verbose:\n",
    "            logger.opt(colors=True).info(f\"<b><fg #EC9A3C>Inputs</fg #EC9A3C></b>: {inputs}\")\n",
    "\n",
    "        output = self.function(**inputs)\n",
    "        if verbose:\n",
    "            logger.opt(colors=True).info(f\"<b><fg #EC9A3C>Output</fg #EC9A3C></b>: {output}\")\n",
    "\n",
    "        return WorkflowStepOutput(inputs=inputs, output=output)\n",
    "\n",
    "class WorkflowLLMStep(BaseModel):\n",
    "    \"\"\"Class that implements an agent step.\n",
    "\n",
    "    Attributes:\n",
    "        name: The name of the step.\n",
    "        agent: The agent that will be invoked when calling this step.\n",
    "    \"\"\"\n",
    "\n",
    "    name: str\n",
    "    agent: Agent\n",
    "\n",
    "    def invoke(self, inputs: dict[str, Any], verbose: bool) -> WorkflowStepOutput:\n",
    "        if verbose:\n",
    "            logger.opt(colors=True).info(f\"<b><fg #2D72D2>Use LLM</fg #2D72D2></b>: {self.name}\")\n",
    "\n",
    "        inputs = {variable: inputs.get(variable) for variable in self.agent.prompt_variables}\n",
    "        if verbose:\n",
    "            logger.opt(colors=True).info(f\"<b><fg #EC9A3C>Inputs</fg #EC9A3C></b>: {inputs}\")\n",
    "\n",
    "        prompt = self.agent.prompt.format(**inputs)\n",
    "        if verbose:\n",
    "            logger.opt(colors=True).info(f\"<b><fg #738091>Prompt</fg #738091></b>: {prompt}\")\n",
    "\n",
    "        output = self.agent.invoke(inputs)\n",
    "        if verbose:\n",
    "            logger.opt(colors=True).info(f\"<b><fg #EC9A3C>Output</fg #EC9A3C></b>: {output.final_answer}\")\n",
    "\n",
    "        return WorkflowStepOutput(inputs=inputs, output=output.final_answer)\n",
    "\n",
    "class WorkflowTransformationStep(BaseModel):\n",
    "    \"\"\"Class that implements a transformation step.\n",
    "\n",
    "    Attributes:\n",
    "        name: The name of the step.\n",
    "        input_field: The name of the field values to transform.\n",
    "        transformation: The transformation to apply (can be map, filter, reduce).\n",
    "        function: The function used for the transformation.\n",
    "    \"\"\"\n",
    "\n",
    "    name: str\n",
    "    input_field: str\n",
    "    transformation: Literal[\"map\", \"filter\", \"reduce\"]\n",
    "    function: Callable[[Any], Any]\n",
    "\n",
    "    def invoke(self, inputs: dict[str, Any], verbose: bool) -> WorkflowStepOutput:\n",
    "        if verbose:\n",
    "            logger.opt(colors=True).info(f\"<b><fg #2D72D2>Use Transformation</fg #2D72D2></b>: {self.name}\")\n",
    "\n",
    "        values = inputs[self.input_field]\n",
    "        inputs = {self.input_field: values}\n",
    "        if verbose:\n",
    "            logger.opt(colors=True).info(f\"<b><fg #EC9A3C>Inputs</fg #EC9A3C></b>: {inputs}\")\n",
    "\n",
    "        if self.transformation == \"map\":\n",
    "            transformed_values = map(self.function, values)\n",
    "            output = list(transformed_values) if isinstance(values, list) else dict(transformed_values)\n",
    "        elif self.transformation == \"filter\":\n",
    "            transformed_values = filter(self.function, values)\n",
    "            output = list(transformed_values) if isinstance(values, list) else dict(transformed_values)\n",
    "        else:\n",
    "            output = reduce(self.function, values)\n",
    "\n",
    "        if verbose:\n",
    "            logger.opt(colors=True).info(f\"<b><fg #EC9A3C>Output</fg #EC9A3C></b>: {output}\")\n",
    "\n",
    "        return WorkflowStepOutput(inputs=inputs, output=output)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we introduce several components to execute a workflow comprising various steps. It's noteworthy that when creating sophisticated workflows, some steps may not depend on the previous ones. To improve efficiency, we could identify processes that can run in parallel. However, for simplicity, we will skip this optimization in our current setup."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "class WorkflowStateManager(BaseModel):\n",
    "    \"\"\"Class that implements a state manager.\"\"\"\n",
    "\n",
    "    state: dict[str, Any]\n",
    "\n",
    "    def update(self, name: str, step: WorkflowStepOutput) -> None:\n",
    "        \"\"\"Updates the state values.\"\"\"\n",
    "        self.state[name] = step.output\n",
    "\n",
    "class WorkflowOutput(BaseModel):\n",
    "    \"\"\"Class that represents the workflow output.\"\"\"\n",
    "\n",
    "    inputs: dict[str, Any]\n",
    "    output: (\n",
    "        str\n",
    "        | int\n",
    "        | float\n",
    "        | dict[str, Any]\n",
    "        | BaseModel\n",
    "        | list[str]\n",
    "        | list[int]\n",
    "        | list[float]\n",
    "        | list[dict[str, Any]]\n",
    "        | list[BaseModel]\n",
    "        | None\n",
    "    )\n",
    "\n",
    "class Workflow(BaseModel):\n",
    "    \"\"\"Class that implements a workflow.\n",
    "\n",
    "    Attributes:\n",
    "        name: The name of the workflow.\n",
    "        description: The description of what the workflow does.\n",
    "        steps: The steps of the workflow.\n",
    "        inputs: The workflow inputs.\n",
    "        output: The name of the step value to output.\n",
    "    \"\"\"\n",
    "\n",
    "    name: str\n",
    "    description: str\n",
    "    steps: list[WorkflowLLMStep | WorkflowFunctionStep | WorkflowTransformationStep]\n",
    "    inputs: type[BaseModel]\n",
    "    output: str\n",
    "    verbose: bool = True\n",
    "\n",
    "    def invoke(self, inputs: dict[str, Any]) -> WorkflowOutput:\n",
    "        \"\"\"Runs the workflow.\"\"\"\n",
    "        _ = self.inputs.model_validate(inputs)\n",
    "        state_manager = WorkflowStateManager(state=inputs)\n",
    "        for step in self.steps:\n",
    "            output = step.invoke(state_manager.state, self.verbose)\n",
    "            state_manager.update(step.name, output)\n",
    "\n",
    "        output = state_manager.state.get(self.output)\n",
    "        if self.verbose:\n",
    "            logger.opt(colors=True).success(f\"<b><fg #32A467>Workflow Output</fg #32A467></b>: {output}\")\n",
    "\n",
    "        return WorkflowOutput(inputs=inputs, output=output)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this workflow, each step's output is automatically stored in a global state, making it accessible for subsequent steps. This approach simplifies the process of chaining tasks together, as each step can easily reference the results of previous steps by simply using the step's name.\n",
    "\n",
    "For example, after the LLM step completes the task of extracting numbers from the text, the resulting list of numbers is stored with the identifier \"numbers.\" This allows the function step, which sorts these numbers, to simply refer to \"numbers\" to access the data it needs. The sorted list produced by the function step is then stored with its step name \"sort,\" making it available for the transformation step."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "system_prompt = \"\"\"You are an AI assistant designed to help users with a variety of tasks.\n",
    "\n",
    "Extract all numbers from the user's input text.\"\"\"\n",
    "\n",
    "agent = Agent.create(\n",
    "    llm=llm,\n",
    "    system_prompt=system_prompt,\n",
    "    prompt=\"{prompt}\",\n",
    "    prompt_variables=[\"prompt\"],\n",
    "    output_type=OutputType.ARRAY_INTEGER,\n",
    "    prompting_strategy=PromptingStrategy.SINGLE_COMPLETION,\n",
    "    verbose=True,\n",
    ")\n",
    "\n",
    "agent_step = WorkflowLLMStep(name=\"numbers\", agent=agent)\n",
    "\n",
    "class Function(BaseModel):\n",
    "    numbers: list[int]\n",
    "\n",
    "function_step = WorkflowFunctionStep(name=\"sort\", inputs=Function, function=lambda numbers: sorted(numbers))\n",
    "\n",
    "filter_step = WorkflowTransformationStep(name=\"numbers_greater_10\", input_field=\"sort\", transformation=\"filter\", function=lambda number: number > 10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Prompt(BaseModel):\n",
    "    prompt: str = Field(description=\"The user prompt\")\n",
    "\n",
    "workflow = Workflow(\n",
    "    name=\"Find numbers greater than 10\",\n",
    "    description=\"Extracts numbers from a given text\",\n",
    "    steps=[agent_step, function_step, filter_step],\n",
    "    inputs=Prompt,\n",
    "    output=\"numbers_greater_10\",\n",
    "    verbose=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The workflow is driven by a prompt containing a narrative about a hiking trip. Upon execution, the workflow seamlessly processes the input, notably extracting key numerical data points and refining them as specified. These workflows will be more useful later on."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[1m\u001b[38;2;45;114;210mUse LLM\u001b[0m\u001b[1m\u001b[0m: numbers\n",
      "\u001b[1m\u001b[38;2;236;154;60mInputs\u001b[0m\u001b[1m\u001b[0m: {'prompt': \"Last weekend, six of us went on a 15-kilometer hike, starting at 7 AM.\\n\\nBy noon, we had covered 10 kilometers and reached Mount Elbert's 4,401-meter summit by 2 PM, with a temperature of 12°C.\\n\\nWe camped 5 kilometers away by 6 PM with 12 others and returned home by 5 PM the next day.\"}\n",
      "\u001b[1m\u001b[38;2;115;128;145mPrompt\u001b[0m\u001b[1m\u001b[0m: Last weekend, six of us went on a 15-kilometer hike, starting at 7 AM.\n",
      "\n",
      "By noon, we had covered 10 kilometers and reached Mount Elbert's 4,401-meter summit by 2 PM, with a temperature of 12°C.\n",
      "\n",
      "We camped 5 kilometers away by 6 PM with 12 others and returned home by 5 PM the next day.\n",
      "\u001b[1m\u001b[38;2;50;164;103mFinal Answer\u001b[0m\u001b[1m\u001b[0m: [6, 15, 7, 10, 4401, 2, 12, 5, 6, 12, 5]\n",
      "\u001b[1m\u001b[38;2;236;154;60mOutput\u001b[0m\u001b[1m\u001b[0m: [6, 15, 7, 10, 4401, 2, 12, 5, 6, 12, 5]\n",
      "\u001b[1m\u001b[38;2;45;114;210mUse Function\u001b[0m\u001b[1m\u001b[0m: sort\n",
      "\u001b[1m\u001b[38;2;236;154;60mInputs\u001b[0m\u001b[1m\u001b[0m: {'numbers': [6, 15, 7, 10, 4401, 2, 12, 5, 6, 12, 5]}\n",
      "\u001b[1m\u001b[38;2;236;154;60mOutput\u001b[0m\u001b[1m\u001b[0m: [2, 5, 5, 6, 6, 7, 10, 12, 12, 15, 4401]\n",
      "\u001b[1m\u001b[38;2;45;114;210mUse Transformation\u001b[0m\u001b[1m\u001b[0m: numbers_greater_10\n",
      "\u001b[1m\u001b[38;2;236;154;60mInputs\u001b[0m\u001b[1m\u001b[0m: {'sort': [2, 5, 5, 6, 6, 7, 10, 12, 12, 15, 4401]}\n",
      "\u001b[1m\u001b[38;2;236;154;60mOutput\u001b[0m\u001b[1m\u001b[0m: [12, 12, 15, 4401]\n",
      "\u001b[1m\u001b[38;2;50;164;103mWorkflow Output\u001b[0m\u001b[1m\u001b[0m: [12, 12, 15, 4401]\n"
     ]
    }
   ],
   "source": [
    "prompt = \"\"\"Last weekend, six of us went on a 15-kilometer hike, starting at 7 AM.\n",
    "\n",
    "By noon, we had covered 10 kilometers and reached Mount Elbert's 4,401-meter summit by 2 PM, with a temperature of 12°C.\n",
    "\n",
    "We camped 5 kilometers away by 6 PM with 12 others and returned home by 5 PM the next day.\"\"\"\n",
    "\n",
    "output = workflow.invoke({\"prompt\": prompt})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[12, 12, 15, 4401]\n"
     ]
    }
   ],
   "source": [
    "print(output.output)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "llm-powered-ai-agents",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}