Refactor local and production deployment pipelines to integrate chat interface visualization

Author: safoinme

llm-complete-guide/pipelines/local_deployment.py

@@ -8,6 +8,3 @@
 def local_deployment():
     bento = bento_builder()
     bento_deployment(bento)
-    create_chat_interface()
-    
-    #vllm_model_deployer_step()

llm-complete-guide/pipelines/prod_deployment.py

@@ -17,6 +17,7 @@
 
 from steps.bento_dockerizer import bento_dockerizer
 from steps.k8s_deployment import k8s_deployment
+from steps.visualize_chat import create_chat_interface
 from zenml import pipeline
 
 
@@ -28,4 +29,5 @@ def production_deployment(
     This is a pipeline deploys trained model for future inference.
     """
     bento_model_image = bento_dockerizer()
-    k8s_deployment(bento_model_image)
+    deployment_info = k8s_deployment(bento_model_image)
+    create_chat_interface(deployment_info)

llm-complete-guide/steps/visualize_chat.py

@@ -1,9 +1,12 @@
-from typing import Optional
-from zenml import pipeline, step
+from typing import Optional, Dict, Any
+from typing_extensions import Annotated
+from zenml import log_artifact_metadata, pipeline, step
 from zenml.types import HTMLString
 
 @step(enable_cache=False)
-def create_chat_interface() -> HTMLString:
+def create_chat_interface(
+        deployment_info: Dict[str, Any],
+    ) -> Annotated[HTMLString, "chat_bot"]:
     html = """
     <div id="zenml-chat-container" class="w-full max-w-4xl mx-auto">
         <style>