runbooks 0.7.0__py3-none-any.whl → 0.7.6__py3-none-any.whl

jupyter-agent/utils.py

@@ -1,409 +0,0 @@
-"""
-utils.py
-
-This module provides helper functions to:
-- Generate and update Jupyter Notebook structures.
-- Execute code in a sandbox environment.
-- Parse and convert execution results to notebook cell outputs.
-- Export notebooks to HTML.
-"""
-
-import json
-from pathlib import Path
-from typing import Any, Dict, List, Tuple
-
-import nbformat
-from e2b_code_interpreter import Sandbox
-from huggingface_hub import InferenceClient
-from nbconvert import HTMLExporter
-from nbformat.v4 import new_code_cell, new_markdown_cell, new_notebook
-from traitlets.config import Config
-from transformers import AutoTokenizer
-
-## --- Global Configuration and Template Loading ---
-
-## Create a configuration for nbconvert and set up the HTML exporter
-config = Config()
-html_exporter = HTMLExporter(config=config, template_name="classic")
-
-## Load the Jinja template for the LLaMA model
-TEMPLATE_PATH = Path("llama3_template.jinja")
-try:
-    with TEMPLATE_PATH.open("r", encoding="utf-8") as f:
-        llama_template = f.read()
-except FileNotFoundError:
-    raise FileNotFoundError(f"Template file {TEMPLATE_PATH} not found.")
-
-MAX_TURNS = 4
-
-## --- Code Execution Functions ---
-
-
-def parse_exec_result_nb(execution: Any) -> List[Dict[str, Any]]:
-    """
-    Convert an E2B execution object to a list of Jupyter notebook cell outputs format.
-
-    :param execution: Execution object from the sandbox.
-    :return: List of output dictionaries.
-    """
-    outputs: List[Dict[str, Any]] = []
-
-    if execution.logs.stdout:
-        outputs.append(
-            {
-                "output_type": "stream",
-                "name": "stdout",
-                "text": "".join(execution.logs.stdout),
-            }
-        )
-
-    if execution.logs.stderr:
-        outputs.append(
-            {
-                "output_type": "stream",
-                "name": "stderr",
-                "text": "".join(execution.logs.stderr),
-            }
-        )
-
-    if execution.error:
-        outputs.append(
-            {
-                "output_type": "error",
-                "ename": execution.error.name,
-                "evalue": execution.error.value,
-                "traceback": [line for line in execution.error.traceback.split("\n")],
-            }
-        )
-
-    for result in execution.results:
-        output = {
-            "output_type": (
-                "execute_result" if result.is_main_result else "display_data"
-            ),
-            "metadata": {},
-            "data": {},
-        }
-
-        if result.text:
-            output["data"]["text/plain"] = [result.text]  # Array for text/plain
-        if result.html:
-            output["data"]["text/html"] = result.html
-        if result.png:
-            output["data"]["image/png"] = result.png
-        if result.svg:
-            output["data"]["image/svg+xml"] = result.svg
-        if result.jpeg:
-            output["data"]["image/jpeg"] = result.jpeg
-        if result.pdf:
-            output["data"]["application/pdf"] = result.pdf
-        if result.latex:
-            output["data"]["text/latex"] = result.latex
-        if result.json:
-            output["data"]["application/json"] = result.json
-        if result.javascript:
-            output["data"]["application/javascript"] = result.javascript
-
-        if result.is_main_result and execution.execution_count is not None:
-            output["execution_count"] = execution.execution_count
-
-        if output["data"]:
-            outputs.append(output)
-
-    return outputs
-
-
-## HTML and CSS templates for notebook cells
-system_template = """\
-<details>
-  <summary style="display: flex; align-items: center;">
-    <div class="alert alert-block alert-info" style="margin: 0; width: 100%;">
-      <b>System: <span class="arrow">▶</span></b>
-    </div>
-  </summary>
-  <div class="alert alert-block alert-info">
-    {}
-  </div>
-</details>
-
-<style>
-details > summary .arrow {{
-  display: inline-block;
-  transition: transform 0.2s;
-}}
-details[open] > summary .arrow {{
-  transform: rotate(90deg);
-}}
-</style>
-"""
-
-user_template = """<div class="alert alert-block alert-success">
-<b>User:</b> {}
-</div>
-"""
-
-header_message = """<p align="center">
-  <img src="cloudops-agent.png" alt="Jupyter Agent" />
-</p>
-
-
-<p style="text-align:center;">Let a LLM agent write and execute code inside a notebook!</p>"""
-
-bad_html_bad = """input[type="file"] {
-  display: block;
-}"""
-
-
-## --- Notebook Creation and Update Functions ---
-
-
-def create_base_notebook(messages: List[Dict[str, Any]]) -> Tuple[Dict[str, Any], int]:
-    """
-    Create the base Jupyter Notebook structure with initial cells.
-
-    :param messages: List of conversation messages.
-    :return: A tuple of the notebook data dictionary and the current code cell counter.
-    """
-    base_notebook = {
-        "metadata": {
-            "kernel_info": {"name": "python3"},
-            "language_info": {
-                "name": "python",
-                "version": "3.12",
-            },
-        },
-        "nbformat": 4,
-        "nbformat_minor": 0,
-        "cells": [],
-    }
-    base_notebook["cells"].append(
-        {"cell_type": "markdown", "metadata": {}, "source": header_message}
-    )
-
-    if len(messages) == 0:
-        base_notebook["cells"].append(
-            {
-                "cell_type": "code",
-                "execution_count": None,
-                "metadata": {},
-                "source": "",
-                "outputs": [],
-            }
-        )
-
-    code_cell_counter = 0
-
-    for message in messages:
-        if message["role"] == "system":
-            text = system_template.format(message["content"].replace("\n", "<br>"))
-            base_notebook["cells"].append(
-                {"cell_type": "markdown", "metadata": {}, "source": text}
-            )
-        elif message["role"] == "user":
-            text = user_template.format(message["content"].replace("\n", "<br>"))
-            base_notebook["cells"].append(
-                {"cell_type": "markdown", "metadata": {}, "source": text}
-            )
-
-        elif message["role"] == "assistant" and "tool_calls" in message:
-            base_notebook["cells"].append(
-                {
-                    "cell_type": "code",
-                    "execution_count": None,
-                    "metadata": {},
-                    "source": message["content"],
-                    "outputs": [],
-                }
-            )
-
-        elif message["role"] == "ipython":
-            code_cell_counter += 1
-            base_notebook["cells"][-1]["outputs"] = message["nbformat"]
-            base_notebook["cells"][-1]["execution_count"] = code_cell_counter
-
-        elif message["role"] == "assistant" and "tool_calls" not in message:
-            base_notebook["cells"].append(
-                {"cell_type": "markdown", "metadata": {}, "source": message["content"]}
-            )
-
-        else:
-            raise ValueError(message)
-
-    return base_notebook, code_cell_counter
-
-
-def execute_code(sbx: Sandbox, code: str) -> Tuple[str, Any]:
-    """
-    Execute the given code in the provided sandbox.
-
-    :param sbx: Sandbox instance to run the code.
-    :param code: Code to execute.
-    :return: Tuple of aggregated output string and the raw execution object.
-    """
-    execution = sbx.run_code(code, on_stdout=lambda data: print("stdout:", data))
-    output = ""
-    if len(execution.logs.stdout) > 0:
-        output += "\n".join(execution.logs.stdout)
-    if len(execution.logs.stderr) > 0:
-        output += "\n".join(execution.logs.stderr)
-    if execution.error is not None:
-        output += execution.error.traceback
-    return output, execution
-
-
-def parse_exec_result_llm(execution: Any) -> str:
-    """
-    Parse the execution results and return a single concatenated output string.
-
-    :param execution: Execution object from the sandbox.
-    :return: Concatenated string of output messages.
-    """
-    output = ""
-    if len(execution.logs.stdout) > 0:
-        output += "\n".join(execution.logs.stdout)
-    if len(execution.logs.stderr) > 0:
-        output += "\n".join(execution.logs.stderr)
-    if execution.error is not None:
-        output += execution.error.traceback
-    return output
-
-
-def update_notebook_display(notebook_data):
-    notebook = nbformat.from_dict(notebook_data)
-    notebook_body, _ = html_exporter.from_notebook_node(notebook)
-    notebook_body = notebook_body.replace(bad_html_bad, "")
-    return notebook_body
-
-
-## --- Interactive Notebook Generation ---
-
-
-def run_interactive_notebook(
-    client: InferenceClient,
-    model: str,
-    tokenizer: Any,
-    messages: List[Dict[str, Any]],
-    sbx: Sandbox,
-    max_new_tokens: int = 512,
-) -> Any:
-    """
-    Generator function that iteratively builds and updates the Jupyter Notebook.
-
-    :param client: Hugging Face InferenceClient for text generation.
-    :param model: Model identifier.
-    :param tokenizer: Tokenizer corresponding to the model.
-    :param messages: List of conversation messages.
-    :param sbx: Sandbox instance for executing code.
-    :param max_new_tokens: Maximum tokens to generate per turn.
-    :yield: Tuple containing updated notebook HTML, notebook data, and messages.
-    """
-    notebook_data, code_cell_counter = create_base_notebook(messages)
-    turns = 0
-
-    # code_cell_counter = 0
-    while turns <= MAX_TURNS:
-        turns += 1
-        input_tokens = tokenizer.apply_chat_template(
-            messages,
-            chat_template=llama_template,
-            builtin_tools=["code_interpreter"],
-            add_generation_prompt=True,
-        )
-        model_input = tokenizer.decode(input_tokens)
-
-        print(f"Model input:\n{model_input}\n{'='*80}")
-
-        response_stream = client.text_generation(
-            model=model,
-            prompt=model_input,
-            details=True,
-            stream=True,
-            do_sample=True,
-            repetition_penalty=1.1,
-            temperature=0.8,
-            max_new_tokens=max_new_tokens,
-        )
-
-        assistant_response = ""
-        tokens = []
-
-        code_cell = False
-        for i, chunk in enumerate(response_stream):
-            if not chunk.token.special:
-                content = chunk.token.text
-            else:
-                content = ""
-            tokens.append(chunk.token.text)
-            assistant_response += content
-
-            if len(tokens) == 1:
-                create_cell = True
-                code_cell = "<|python_tag|>" in tokens[0]
-                if code_cell:
-                    code_cell_counter += 1
-            else:
-                create_cell = False
-
-            ## Update notebook cells in real-time
-            if create_cell:
-                if "<|python_tag|>" in tokens[0]:
-                    notebook_data["cells"].append(
-                        {
-                            "cell_type": "code",
-                            "execution_count": None,
-                            "metadata": {},
-                            "source": assistant_response,
-                            "outputs": [],
-                        }
-                    )
-                else:
-                    notebook_data["cells"].append(
-                        {
-                            "cell_type": "markdown",
-                            "metadata": {},
-                            "source": assistant_response,
-                        }
-                    )
-            else:
-                notebook_data["cells"][-1]["source"] = assistant_response
-            if i % 16 == 0:
-                yield update_notebook_display(notebook_data), notebook_data, messages
-        yield update_notebook_display(notebook_data), notebook_data, messages
-
-        ## If a code cell was generated, execute the code
-        if code_cell:
-            notebook_data["cells"][-1]["execution_count"] = code_cell_counter
-
-            exec_result, execution = execute_code(sbx, assistant_response)
-            messages.append(
-                {
-                    "role": "assistant",
-                    "content": assistant_response,
-                    "tool_calls": [
-                        {
-                            "type": "function",
-                            "function": {
-                                "name": "code_interpreter",
-                                "arguments": {"code": assistant_response},
-                            },
-                        }
-                    ],
-                }
-            )
-            messages.append(
-                {
-                    "role": "ipython",
-                    "content": parse_exec_result_llm(execution),
-                    "nbformat": parse_exec_result_nb(execution),
-                }
-            )
-
-            ## Update the last code cell with execution results
-            notebook_data["cells"][-1]["outputs"] = parse_exec_result_nb(execution)
-            update_notebook_display(notebook_data)
-        else:
-            messages.append({"role": "assistant", "content": assistant_response})
-            if tokens[-1] == "<|eot_id|>":
-                break
-
-    yield update_notebook_display(notebook_data), notebook_data, messages

runbooks/aws/__init__.py

@@ -1,58 +0,0 @@
-## src/runbooks/aws/__init__.py
-"""AWS Runbooks Initialization Module."""
-
-import importlib
-import os
-import sys
-
-from runbooks.utils.logger import configure_logger
-
-logger = configure_logger(__name__)
-
-
-def discover_scripts():
-    """
-    Dynamically discovers and lists all AWS scripts in this package.
-
-    Returns:
-        dict: A mapping of script names to their main functions.
-    """
-    scripts = {}
-    aws_path = os.path.dirname(__file__)
-    for filename in os.listdir(aws_path):
-        if filename.endswith(".py") and filename != "__init__.py":
-            module_name = f"runbooks.aws.{filename[:-3]}"
-            try:
-                module = importlib.import_module(module_name)
-                if hasattr(module, "main"):
-                    scripts[filename[:-3]] = module.main
-            except Exception as e:
-                logger.error(f"Error importing {module_name}: {e}")
-    return scripts
-
-
-def run_script(script_name, *args):
-    """
-    Executes the given script by name.
-
-    Args:
-        script_name (str): The name of the script to execute.
-        *args: Additional arguments to pass to the script.
-    """
-    scripts = discover_scripts()
-    if script_name in scripts:
-        try:
-            scripts[script_name](*args)
-        except Exception as e:
-            logger.error(f"Error executing script {script_name}: {e}")
-    else:
-        logger.error(f"Script {script_name} not found.")
-        sys.exit(1)
-
-
-if __name__ == "__main__":
-    if len(sys.argv) < 2:
-        logger.error("Usage: python -m runbooks.aws <script_name> [<args>]")
-        sys.exit(1)
-
-    run_script(sys.argv[1], *sys.argv[2:])

runbooks/aws/dynamodb_operations.py

@@ -1,231 +0,0 @@
-#!/usr/bin/env python3
-
-"""
-DynamoDB Operations: Put Item, Delete Item, and Batch Write.
-
-This script supports the following functionalities:
-1. Insert or update a single item (Put Item).
-2. Retrieve and delete a single item (Delete Item).
-3. Batch insert multiple items efficiently (Batch Write).
-
-Designed for usage in Python, Docker, and AWS Lambda environments.
-
-Author: nnthanh101@gmail.com
-Date: 2025-01-09
-Version: 1.0.0
-"""
-
-import json
-import os
-from typing import Dict, List
-
-import boto3
-from botocore.exceptions import BotoCoreError, ClientError
-
-from runbooks.utils.logger import configure_logger
-
-## ✅ Configure Logger
-logger = configure_logger(__name__)
-
-# ==============================
-# CONFIGURATION VARIABLES
-# ==============================
-AWS_REGION = os.getenv("AWS_REGION", "us-east-1")
-TABLE_NAME = os.getenv("TABLE_NAME", "employees")
-MAX_BATCH_ITEMS = int(os.getenv("MAX_BATCH_ITEMS", 100))
-
-
-# ==============================
-# AWS CLIENT INITIALIZATION
-# ==============================
-try:
-    dynamodb = boto3.resource("dynamodb", region_name=AWS_REGION)
-    table = dynamodb.Table(TABLE_NAME)
-    logger.info(f"✅ DynamoDB Table '{TABLE_NAME}' initialized successfully.")
-except Exception as e:
-    logger.error(f"❌ Failed to initialize DynamoDB table: {e}")
-    raise
-
-
-# ==============================
-# FUNCTION: PUT ITEM
-# ==============================
-def put_item(emp_id: str, name: str, salary: int) -> None:
-    """
-    Inserts or updates a single item in DynamoDB.
-
-    Args:
-        emp_id (str): Employee ID.
-        name (str): Employee name.
-        salary (int): Employee salary.
-
-    Raises:
-        Exception: If item insertion fails.
-    """
-    try:
-        logger.info(f"🚀 Inserting/Updating item in table '{TABLE_NAME}'...")
-        table.put_item(Item={"emp_id": emp_id, "name": name, "salary": salary})
-        logger.info(f"✅ Item added successfully: emp_id={emp_id}, name={name}, salary={salary}")
-
-    except ClientError as e:
-        logger.error(f"❌ AWS Client Error: {e}")
-        raise
-
-    except Exception as e:
-        logger.error(f"❌ Unexpected Error: {e}")
-        raise
-
-
-# ==============================
-# FUNCTION: DELETE ITEM
-# ==============================
-def delete_item(emp_id: str) -> Dict:
-    """
-    Retrieves and deletes a single item from DynamoDB.
-
-    Args:
-        emp_id (str): Employee ID.
-
-    Returns:
-        Dict: Deleted item details.
-
-    Raises:
-        Exception: If retrieval or deletion fails.
-    """
-    try:
-        ## ✅ 1. Retrieve the item
-        logger.info(f"🔍 Retrieving item with emp_id={emp_id}...")
-        response = table.get_item(Key={"emp_id": emp_id})
-
-        if "Item" not in response:
-            raise ValueError(f"Item with emp_id={emp_id} not found.")
-        item = response["Item"]
-        logger.info(f"✅ Item retrieved: {item}")
-
-        ## ✅ 2. Delete the item
-        logger.info(f"🗑️ Deleting item with emp_id={emp_id}...")
-        table.delete_item(Key={"emp_id": emp_id})
-        logger.info(f"✅ Item deleted successfully: emp_id={emp_id}")
-
-        return item
-
-    except ClientError as e:
-        logger.error(f"❌ AWS Client Error: {e}")
-        raise
-
-    except BotoCoreError as e:
-        logger.error(f"❌ BotoCore Error: {e}")
-        raise
-
-    except Exception as e:
-        logger.error(f"❌ Unexpected Error: {e}")
-        raise
-
-
-# ==============================
-# FUNCTION: BATCH WRITE ITEMS
-# ==============================
-def batch_write_items(batch_size: int = MAX_BATCH_ITEMS) -> None:
-    """
-    Inserts multiple items into DynamoDB using batch writer.
-
-    Args:
-        batch_size (int): Number of items to write in a batch.
-
-    Raises:
-        Exception: If batch write fails.
-    """
-    try:
-        logger.info(f"🚀 Starting batch write with {batch_size} items...")
-        with table.batch_writer() as batch:
-            for i in range(batch_size):
-                batch.put_item(
-                    Item={
-                        "emp_id": str(i),
-                        "name": f"Name-{i}",
-                        "salary": 50000 + i * 100,  ## Incremental salary
-                    }
-                )
-        logger.info(f"✅ Batch write completed successfully with {batch_size} items.")
-
-    except ClientError as e:
-        logger.error(f"❌ AWS Client Error: {e}")
-        raise
-
-    except BotoCoreError as e:
-        logger.error(f"❌ BotoCore Error: {e}")
-        raise
-
-    except Exception as e:
-        logger.error(f"❌ Unexpected Error: {e}")
-        raise
-
-
-# ==============================
-# MAIN FUNCTION (CLI/DOCKER)
-# ==============================
-def main():
-    """
-    Main function for CLI/Docker execution.
-    """
-    try:
-        ## Use-Case 1: Put Item
-        put_item(emp_id="2", name="John Doe", salary=75000)
-
-        ## Use-Case 2: Delete Item
-        delete_item(emp_id="2")
-
-        ## Use-Case 3: Batch Write Items
-        batch_write_items(batch_size=MAX_BATCH_ITEMS)
-
-    except Exception as e:
-        logger.error(f"❌ Error in main execution: {e}")
-        raise
-
-
-# ==============================
-# AWS LAMBDA HANDLER
-# ==============================
-def lambda_handler(event, context):
-    """
-    AWS Lambda handler for DynamoDB operations.
-
-    Args:
-        event (dict): AWS Lambda event with action details.
-        context: AWS Lambda context object.
-
-    Returns:
-        dict: Status code and message.
-    """
-    try:
-        action = event.get("action")
-        emp_id = event.get("emp_id")
-        name = event.get("name")
-        salary = event.get("salary", 0)
-        batch_size = int(event.get("batch_size", MAX_BATCH_ITEMS))
-
-        if action == "put":
-            put_item(emp_id, name, salary)
-            return {"statusCode": 200, "body": f"Item {emp_id} inserted."}
-
-        elif action == "delete":
-            item = delete_item(emp_id)
-            return {"statusCode": 200, "body": f"Item {item} deleted."}
-
-        elif action == "batch_write":
-            batch_write_items(batch_size)
-            return {"statusCode": 200, "body": "Batch write completed."}
-
-        else:
-            raise ValueError("Invalid action. Use 'put', 'delete', or 'batch_write'.")
-
-    except Exception as e:
-        logger.error(f"❌ Lambda Error: {e}")
-        return {"statusCode": 500, "body": str(e)}
-
-
-# ==============================
-# SCRIPT ENTRY POINT
-# ==============================
-if __name__ == "__main__":
-    main()