From 6929e825086cb076e4aac4b2b8b4620a21f42133 Mon Sep 17 00:00:00 2001
From: david qiu <david@qiu.dev>
Date: Wed, 5 Apr 2023 16:43:59 -0700
Subject: [PATCH] Magic tweaks (#31)

* suppress warning when using OpenAIChat models

* do not call IPython.display.display() when returning output

* add raw format mode
---
 packages/jupyter-ai/jupyter_ai/magics.py | 22 ++++++++++++++++------
 1 file changed, 16 insertions(+), 6 deletions(-)

diff --git a/packages/jupyter-ai/jupyter_ai/magics.py b/packages/jupyter-ai/jupyter_ai/magics.py
index 476d80414..bc0125f08 100644
--- a/packages/jupyter-ai/jupyter_ai/magics.py
+++ b/packages/jupyter-ai/jupyter_ai/magics.py
@@ -1,13 +1,15 @@
 import os
 import json
+import warnings
 from typing import Optional
-from importlib_metadata import entry_points
 
-from jupyter_ai.providers import BaseProvider
+from importlib_metadata import entry_points
 from IPython import get_ipython
 from IPython.core.magic import Magics, magics_class, line_cell_magic
 from IPython.core.magic_arguments import magic_arguments, argument, parse_argstring
-from IPython.display import display, HTML, Markdown, Math, JSON
+from IPython.display import HTML, Markdown, Math, JSON
+
+from jupyter_ai.providers import BaseProvider
 
 
 MODEL_ID_ALIASES = {
@@ -23,6 +25,7 @@
     "math": Math,
     "md": Markdown,
     "json": JSON,
+    "raw": None
 }
 
 class FormatDict(dict):
@@ -41,6 +44,11 @@ def __init__(self, shell):
         super(AiMagics, self).__init__(shell)
         self.transcript_openai = []
 
+        # suppress warning when using old OpenAIChat provider
+        warnings.filterwarnings("ignore", message="You are trying to use a chat model. This way of initializing it is "
+            "no longer supported. Instead, please use: "
+            "`from langchain.chat_models import ChatOpenAI`")
+
         # load model providers from entry point
         self.providers = {}
         eps = entry_points()
@@ -97,7 +105,7 @@ def _get_provider(self, provider_id: Optional[str]) -> BaseProvider:
                 optionally prefixed with the ID of the model provider, delimited
                 by a colon.""")
     @argument('-f', '--format',
-                choices=["markdown", "html", "json", "math", "md"],
+                choices=["markdown", "html", "json", "math", "md", "raw"],
                 nargs="?",
                 default="markdown",
                 help="""IPython display to use when rendering output. [default="markdown"]""")
@@ -124,7 +132,7 @@ def ai(self, line, cell=None):
         provider_id, local_model_id = self._decompose_model_id(args.model_id)
         Provider = self._get_provider(provider_id)
         if Provider is None:
-            return display(f"Cannot determine model provider from model ID {args.model_id}.")
+            return f"Cannot determine model provider from model ID {args.model_id}."
 
         # if `--reset` is specified, reset transcript and return early
         if (provider_id == "openai-chat" and args.reset):
@@ -162,10 +170,12 @@ def ai(self, line, cell=None):
 
         # build output display
         DisplayClass = DISPLAYS_BY_FORMAT[args.format]
+        if DisplayClass is None:
+            return output
         if args.format == 'json':
             # JSON display expects a dict, not a JSON string
             output = json.loads(output)
         output_display = DisplayClass(output)
 
         # finally, display output display
-        return display(output_display)
+        return output_display