diff --git a/demos/network_agent/arch_config.yaml b/demos/network_agent/arch_config.yaml
index 250a1037..f46dd204 100644
--- a/demos/network_agent/arch_config.yaml
+++ b/demos/network_agent/arch_config.yaml
@@ -17,36 +17,32 @@ system_prompt: |
   You are a network assistant that helps operators with a better understanding of network traffic flow and perform actions on networking operations. No advice on manufacturers or purchasing decisions.
 
 prompt_targets:
-    - name: device_summary
-      description: Retrieve network statistics for specific devices within a time range
-      endpoint:
-        name: app_server
-        path: /agent/device_summary
-        http_method: POST
-      parameters:
-        - name: device_ids
-          type: list
-          description: A list of device identifiers (IDs) to retrieve statistics for.
-          required: true  # device_ids are required to get device statistics
-        - name: days
-          type: int
-          description: The number of days for which to gather device statistics.
-          default: "7"
-    - name: reboot_devices
-      description: Reboot a list of devices
-      endpoint:
-        name: app_server
-        path: /agent/device_reboot
-        http_method: POST
-      parameters:
-        - name: device_ids
-          type: list
-          description: A list of device identifiers (IDs).
-          required: true
-        - name: days
-          type: int
-          description: A list of device identifiers (IDs)
-          default: "7"
+  - name: device_summary
+    description: Retrieve network statistics for specific devices within a time range
+    endpoint:
+      name: app_server
+      path: /agent/device_summary
+      http_method: POST
+    parameters:
+      - name: device_ids
+        type: list
+        description: A list of device identifiers (IDs) to retrieve statistics for.
+        required: true # device_ids are required to get device statistics
+      - name: days
+        type: int
+        description: The number of days for which to gather device statistics.
+        default: "7"
+  - name: reboot_devices
+    description: Reboot a list of devices
+    endpoint:
+      name: app_server
+      path: /agent/device_reboot
+      http_method: POST
+    parameters:
+      - name: device_ids
+        type: list
+        description: A list of device identifiers (IDs).
+        required: true
 
 # Arch creates a round-robin load balancing between different endpoints, managed via the cluster subsystem.
 endpoints:
@@ -54,6 +50,6 @@ endpoints:
     # value could be ip address or a hostname with port
     # this could also be a list of endpoints for load balancing
     # for example endpoint: [ ip1:port, ip2:port ]
-    endpoint: host.docker.internal:18080
+    endpoint: host.docker.internal:18083
     # max time to wait for a connection to be established
     connect_timeout: 0.005s
diff --git a/demos/network_agent/docker-compose.yaml b/demos/network_agent/docker-compose.yaml
index 5d9fb0cb..f8235331 100644
--- a/demos/network_agent/docker-compose.yaml
+++ b/demos/network_agent/docker-compose.yaml
@@ -2,14 +2,19 @@ services:
   api_server:
     build:
       context: .
+      dockerfile: Dockerfile
+    ports:
+      - "18083:80"
+
+  chatbot_ui:
+    build:
+      context: ../shared/chatbot_ui
+      dockerfile: Dockerfile
+    ports:
+      - "18080:8080"
     environment:
       - CHAT_COMPLETION_ENDPOINT=http://host.docker.internal:10000/v1
+    extra_hosts:
+      - "host.docker.internal:host-gateway"
     volumes:
       - ./arch_config.yaml:/app/arch_config.yaml
-      - ../shared/chatbot_ui/common.py:/app/common.py
-    ports:
-      - "18080:80"
-    healthcheck:
-        test: ["CMD", "curl" ,"http://localhost:80/healthz"]
-        interval: 5s
-        retries: 20
diff --git a/demos/network_agent/main.py b/demos/network_agent/main.py
index 746353b7..bc1a56e4 100644
--- a/demos/network_agent/main.py
+++ b/demos/network_agent/main.py
@@ -1,11 +1,9 @@
-from openai import OpenAI
-from fastapi import FastAPI, HTTPException
-from pydantic import BaseModel, Field
-from typing import List, Optional
-from common import create_gradio_app
-import gradio as gr
 import os
+from typing import List, Optional
 
+from fastapi import FastAPI, HTTPException
+from openai import OpenAI
+from pydantic import BaseModel, Field
 
 app = FastAPI()
 demo_description = """This demo illustrates how **Arch** can be used to perform function calling with network-related tasks.
@@ -100,10 +98,3 @@ client = OpenAI(
     api_key="--",
     base_url=CHAT_COMPLETION_ENDPOINT,
 )
-
-gr.mount_gradio_app(
-    app, create_gradio_app(demo_description, client), path="/agent/chat"
-)
-
-if __name__ == "__main__":
-    app.run(debug=True)
diff --git a/demos/network_agent/utils.py b/demos/network_agent/utils.py
index 18b782c5..5fdf59d0 100644
--- a/demos/network_agent/utils.py
+++ b/demos/network_agent/utils.py
@@ -1,10 +1,11 @@
-import pandas as pd
-import random
-from datetime import datetime, timedelta, timezone
-import re
 import logging
-from dateparser import parse
+import random
+import re
 import sqlite3
+from datetime import datetime, timedelta, timezone
+
+import pandas as pd
+from dateparser import parse
 
 logging.basicConfig(
     level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"