Update app.py

app.py

@@ -5,6 +5,7 @@ import datetime
 import requests
 from checks.failed_check import  create_failed_gradio_ui
 from checks.endpoint_check import check_public_endpoint
+from checks.health_check import check_model_endpoint, should_launch_ui
 
 import os
 import pytz # Had to give it permission in Code agent  
@@ -15,65 +16,66 @@ from agents.model import load_huggingface_model
 from agents.prompts import load_prompts
 
 from ContentGradio import ContentAgentUI
+   
+    # Try a direct test of your endpoint
+    test_prompt = "Hello, how are you?"
+    response = requests.post(
+        endpoint_uri,
+        json={"inputs": test_prompt},
+        headers={"Authorization": f"Bearer {os.getenv('HF_TOKEN')}"}
+    )
+    print(response.json())
 
+    
 
-# Prechecks 
-
-# Get the URI for the endpoint 
-endpoint_uri = load_huggingface_model()
-
-# Test the endpoint
-
-status_info = check_public_endpoint(endpoint_uri)
+   
+def initialize_agent(endpoint_uri: str):
 
- 
-print(status_info["status"])
-print(status_info["status_code"])
- 
+     # Load prompts
+    combined_prompts = load_prompts()
+    combined_tools = load_tools() ## from the tools folder (don't remove final answer)
 
-if status_info["status"] is False or status_info["status_code"] == "503":
-     interface = create_failed_gradio_ui(status_info)
-     # Launch the UI
-     interface.launch(show_error=True)
-else: 
     
+    """Initialize and configure the CodeAgent"""
     model = HfApiModel(
         max_tokens=2096,
         temperature=0.5,
-        model_id= endpoint_uri,
+        model_id=endpoint_uri,
         custom_role_conversions=None,
-     )
-
-    # Load prompts
-    combined_prompts = load_prompts()
-    combined_tools = load_tools() ## from the tools folder (don't remove final answer)
-
-    # Try a direct test of your endpoint
-    test_prompt = "Hello, how are you?"
-    response = requests.post(
-        endpoint_uri,
-        json={"inputs": test_prompt},
-        headers={"Authorization": f"Bearer {os.getenv('HF_TOKEN')}"}
     )
-    print(response.json())
-
-    agent = CodeAgent(
+    
+    return CodeAgent(
         model=model,
-        tools=combined_tools, 
+        tools=combined_tools(),
         max_steps=6,
         verbosity_level=3,
-        grammar=None,
-        planning_interval=None,
         name="content_agent",
-        description="Evaluates whether text is polite or impolite. ",
-        prompt_templates=combined_prompts,
+        description="Evaluates whether text is polite or impolite.",
+        prompt_templates=combined_prompts(),
         additional_authorized_imports=["pytz"]
-     
-    )
+    ) 
+
+
+def main():
+    # Load endpoint and check health
 
+
+
+    # Prechecks 
+    endpoint_uri = load_huggingface_model() # Get the URI for the endpoint 
+    is_healthy, status_info = check_model_endpoint(endpoint_uri) # Test the endpoint
+    
+    if not is_healthy:
+        interface = create_failed_gradio_ui(status_info)
+        interface.launch(show_error=True)
+        return
+    
+    # Initialize and run the agent
+    agent = initialize_agent(endpoint_uri)
     # Create an instance of the ContentAgentUI class
     ui = ContentAgentUI()
+    ui.pass_through_agent(agent)  # Pass through the agent
     
-    # Pass through the agent
-    ui.pass_through_agent(agent)
 
+if __name__ == "__main__":
+    main()