pushing to remote

demos/hr_agent/main.py

@@ -6,48 +6,55 @@ import re
 
 app = FastAPI()
 
+
 class StaffingType(Enum):
-    CONTRACT = "contract"
     FTE = "fte"
     AGENCY = "agency"
+    CONTRACT = "contract"
+
+
+class RegionType(Enum):
+    ASIA = "asia"
+    EUROPE = "europe"
+    AMERICAS = "americas"
+
 
 # Define the request model
 class HeadcountRequest(BaseModel):
-    region: str
+    region: RegionType
     staffing_type: str
 
+
 class HeadcountResponseSummary(BaseModel):
     region: str
     headcount: int
     staffing_type: str
 
+
+HEADCOUNT = {
+    ASIA: {CONTRACT: 100, FTE: 150, AGENCY: 2000},
+    EUROPE: {CONTRACT: 80, FTE: 120, AGENCY: 2500},
+    AMERICAS: {CONTRACT: 90, FTE: 200, AGENCY: 3000},
+}
+
+
 # Post method for device summary
 @app.post("/agent/headcount")
 def get_headcount(request: HeadcountRequest):
     """
     Endpoint to headcount data by region, staffing type over time range
     """
-    staffing_type_value = request.staffing_type
-
-    if re.match(r"(?i)contract", staffing_type_value):  # Case-insensitive regex match
-        headcount = 500
-    elif re.match(r"(?i)fte", staffing_type_value):
-        headcount = 1000
-    elif re.match(r"(?i)agency", staffing_type_value):
-        headcount = 4000
-    else:
-        raise HTTPException(
-            status_code=400, detail="staffing_type parameter is invalid."
-        )
+    headcount = HEADCOUNT[request.region][request.staffing_type]
 
     response = {
-            "region": request.region,
-            "staffing_type": f"Staffing agency: {staffing_type_value}",
-            "headcount" : f"Headcount: {headcount}"
-        }
+        "region": request.region.value,
+        "staffing_type": f"Staffing agency: {staffing_type}",
+        "headcount": f"Headcount: {headcount}",
+    }
 
     return response
 
+
 @app.post("/agent/hr_qa")
 async def general_hr_qa():
     """
@@ -69,5 +76,6 @@ async def general_hr_qa():
         "usage": {"completion_tokens": 0},
     }
 
+
 if __name__ == "__main__":
     app.run(debug=True)

demos/hr_agent/run_demo.sh

@@ -22,14 +22,14 @@ start_demo() {
   archgw up arch_config.yaml
 
   # Step 4: Start Network Agent
-  echo "Starting Network Agent using Docker Compose..."
+  echo "Starting HR Agent using Docker Compose..."
   docker compose up -d  # Run in detached mode
 }
 
 # Function to stop the demo
 stop_demo() {
   # Step 1: Stop Docker Compose services
-  echo "Stopping Network Agent using Docker Compose..."
+  echo "Stopping HR Agent using Docker Compose..."
   docker compose down
 
   # Step 2: Stop Arch