diff --git a/demos/hr_agent/main.py b/demos/hr_agent/main.py index 5bd65b62..a773488a 100644 --- a/demos/hr_agent/main.py +++ b/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) diff --git a/demos/hr_agent/run_demo.sh b/demos/hr_agent/run_demo.sh index f6097b09..fc32430c 100644 --- a/demos/hr_agent/run_demo.sh +++ b/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