From 9dad13d05b87ccc2af9c2eea75ce8458906d7fd1 Mon Sep 17 00:00:00 2001
From: Alex Dhillon <12418963+adhillo@users.noreply.github.com>
Date: Sun, 16 Jul 2023 17:22:46 -0400
Subject: [PATCH] simplify pretty print (#575)

---
 ...w_to_call_functions_with_chat_models.ipynb | 35 ++++++++++++-------
 1 file changed, 22 insertions(+), 13 deletions(-)

diff --git a/examples/How_to_call_functions_with_chat_models.ipynb b/examples/How_to_call_functions_with_chat_models.ipynb
index ea4f567..b434fca 100644
--- a/examples/How_to_call_functions_with_chat_models.ipynb
+++ b/examples/How_to_call_functions_with_chat_models.ipynb
@@ -1,6 +1,7 @@
 {
  "cells": [
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "3e67f200",
    "metadata": {},
@@ -22,6 +23,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "64c85e26",
    "metadata": {},
@@ -65,6 +67,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "69ee6a93",
    "metadata": {},
@@ -119,28 +122,22 @@
     "        \"assistant\": \"blue\",\n",
     "        \"function\": \"magenta\",\n",
     "    }\n",
-    "    formatted_messages = []\n",
+    "    \n",
     "    for message in messages:\n",
     "        if message[\"role\"] == \"system\":\n",
-    "            formatted_messages.append(f\"system: {message['content']}\\n\")\n",
+    "            print(colored(f\"system: {message['content']}\\n\", role_to_color[message[\"role\"]]))\n",
     "        elif message[\"role\"] == \"user\":\n",
-    "            formatted_messages.append(f\"user: {message['content']}\\n\")\n",
+    "            print(colored(f\"user: {message['content']}\\n\", role_to_color[message[\"role\"]]))\n",
     "        elif message[\"role\"] == \"assistant\" and message.get(\"function_call\"):\n",
-    "            formatted_messages.append(f\"assistant: {message['function_call']}\\n\")\n",
+    "            print(colored(f\"assistant: {message['function_call']}\\n\", role_to_color[message[\"role\"]]))\n",
     "        elif message[\"role\"] == \"assistant\" and not message.get(\"function_call\"):\n",
-    "            formatted_messages.append(f\"assistant: {message['content']}\\n\")\n",
+    "            print(colored(f\"assistant: {message['content']}\\n\", role_to_color[message[\"role\"]]))\n",
     "        elif message[\"role\"] == \"function\":\n",
-    "            formatted_messages.append(f\"function ({message['name']}): {message['content']}\\n\")\n",
-    "    for formatted_message in formatted_messages:\n",
-    "        print(\n",
-    "            colored(\n",
-    "                formatted_message,\n",
-    "                role_to_color[messages[formatted_messages.index(formatted_message)][\"role\"]],\n",
-    "            )\n",
-    "        )"
+    "            print(colored(f\"function ({message['name']}): {message['content']}\\n\", role_to_color[message[\"role\"]]))\n"
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "29d4e02b",
    "metadata": {},
@@ -204,6 +201,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "bfc39899",
    "metadata": {},
@@ -242,6 +240,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "4c999375",
    "metadata": {},
@@ -280,6 +279,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "c14d4762",
    "metadata": {},
@@ -318,6 +318,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "6172ddac",
    "metadata": {},
@@ -356,6 +357,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "4b758a0a",
    "metadata": {},
@@ -364,6 +366,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "412f79ba",
    "metadata": {},
@@ -434,6 +437,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "3bd70e48",
    "metadata": {},
@@ -469,6 +473,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "b4482aee",
    "metadata": {},
@@ -481,6 +486,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "f7654fef",
    "metadata": {},
@@ -546,6 +552,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "77e6e5ea",
    "metadata": {},
@@ -570,6 +577,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "ae73c9ee",
    "metadata": {},
@@ -608,6 +616,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "da08c121",
    "metadata": {},