diff --git a/docs/build_docs.py b/docs/build_docs.py
index da4555c09..de8764974 100644
--- a/docs/build_docs.py
+++ b/docs/build_docs.py
@@ -26,16 +26,21 @@
 import textwrap
 import typing
 
-# For showing the conditional imports and types in `content_types.py`
-typing.TYPE_CHECKING = True
 
 from absl import app
 from absl import flags
 
 import google
-from google import generativeai as palm
 from google.ai import generativelanguage as glm
 
+import grpc 
+# For showing the conditional imports and types in `content_types.py`
+# grpc must be imported first.
+typing.TYPE_CHECKING = True
+from google import generativeai as palm
+
+
+
 from tensorflow_docs.api_generator import generate_lib
 from tensorflow_docs.api_generator import public_api