assistant: Improve JSON handling in `/fetch` command (#12864)

Marshall Bowers created 2 years ago

This PR improves the `/fetch` command with better support for URLs that
return JSON content.

JSON response bodies will now be pretty-printed and placed within a
Markdown code block:

<img width="690" alt="Screenshot 2024-06-10 at 3 39 52 PM"
src="https://github.com/zed-industries/zed/assets/1486634/4a7c1cb7-9f5b-4a63-9e8e-5168bf9a6625">

Release Notes:

- Improved the handling of JSON response bodies in the `/fetch` command
in the Assistant.

Change summary

crates/assistant/src/slash_command/fetch_command.rs | 67 +++++++++++---
1 file changed, 50 insertions(+), 17 deletions(-)

Detailed changes

crates/assistant/src/slash_command/fetch_command.rs

@@ -11,6 +11,13 @@ use language::LspAdapterDelegate;
 use ui::{prelude::*, ButtonLike, ElevationIndex};
 use workspace::Workspace;
 
+#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Clone, Copy)]
+enum ContentType {
+    Html,
+    Plaintext,
+    Json,
+}
+
 pub(crate) struct FetchSlashCommand;
 
 impl FetchSlashCommand {
@@ -37,24 +44,50 @@ impl FetchSlashCommand {
             );
         }
 
-        let mut handlers: Vec<Box<dyn HandleTag>> = vec![
-            Box::new(markdown::ParagraphHandler),
-            Box::new(markdown::HeadingHandler),
-            Box::new(markdown::ListHandler),
-            Box::new(markdown::TableHandler::new()),
-            Box::new(markdown::StyledTextHandler),
-        ];
-        if url.contains("wikipedia.org") {
-            use html_to_markdown::structure::wikipedia;
-
-            handlers.push(Box::new(wikipedia::WikipediaChromeRemover));
-            handlers.push(Box::new(wikipedia::WikipediaInfoboxHandler));
-            handlers.push(Box::new(wikipedia::WikipediaCodeHandler::new()));
-        } else {
-            handlers.push(Box::new(markdown::CodeHandler));
-        }
+        let Some(content_type) = response.headers().get("content-type") else {
+            bail!("missing Content-Type header");
+        };
+        let content_type = content_type
+            .to_str()
+            .context("invalid Content-Type header")?;
+        let content_type = match content_type {
+            "text/html" => ContentType::Html,
+            "text/plain" => ContentType::Plaintext,
+            "application/json" => ContentType::Json,
+            _ => ContentType::Html,
+        };
 
-        convert_html_to_markdown(&body[..], handlers)
+        match content_type {
+            ContentType::Html => {
+                let mut handlers: Vec<Box<dyn HandleTag>> = vec![
+                    Box::new(markdown::ParagraphHandler),
+                    Box::new(markdown::HeadingHandler),
+                    Box::new(markdown::ListHandler),
+                    Box::new(markdown::TableHandler::new()),
+                    Box::new(markdown::StyledTextHandler),
+                ];
+                if url.contains("wikipedia.org") {
+                    use html_to_markdown::structure::wikipedia;
+
+                    handlers.push(Box::new(wikipedia::WikipediaChromeRemover));
+                    handlers.push(Box::new(wikipedia::WikipediaInfoboxHandler));
+                    handlers.push(Box::new(wikipedia::WikipediaCodeHandler::new()));
+                } else {
+                    handlers.push(Box::new(markdown::CodeHandler));
+                }
+
+                convert_html_to_markdown(&body[..], handlers)
+            }
+            ContentType::Plaintext => Ok(std::str::from_utf8(&body)?.to_owned()),
+            ContentType::Json => {
+                let json: serde_json::Value = serde_json::from_slice(&body)?;
+
+                Ok(format!(
+                    "```json\n{}\n```",
+                    serde_json::to_string_pretty(&json)?
+                ))
+            }
+        }
     }
 }