HR agent demo (#206)

* commiting my hr_agent branch

* updating the HR agent config

* pushing to remote

* fix hr agent

* committing to merge with main

* updating to merge from main

* updating the demo and model-server-tests to pull from poetry

* updating the poetry.lock files

* updating based on feedback

* updated sysmte prompt for hr_agent

---------

Co-authored-by: Salman Paracha <salmanparacha@MacBook-Pro-261.local>
Co-authored-by: Adil Hafeez <adil@katanemo.com>

demos/hr_agent/Dockerfile

@@ -0,0 +1,20 @@
+FROM python:3.10 AS base
+
+FROM base AS builder
+
+WORKDIR /src
+
+COPY requirements.txt /src/
+COPY workforce_data.json /src/
+RUN pip install --prefix=/runtime --force-reinstall -r requirements.txt
+
+COPY . /src
+
+FROM python:3.10-slim AS output
+
+COPY --from=builder /runtime /usr/local
+
+COPY . /app
+WORKDIR /app
+
+CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "80", "--log-level", "info"]

demos/hr_agent/arch_config.yaml

@@ -0,0 +1,63 @@
+version: v0.1
+listener:
+  address: 127.0.0.1
+  port: 8080 #If you configure port 443, you'll need to update the listener with tls_certificates
+  message_format: huggingface
+
+# Centralized way to manage LLMs, manage keys, retry logic, failover and limits in a central way
+llm_providers:
+  - name: OpenAI
+    provider: openai
+    access_key: OPENAI_API_KEY
+    model: gpt-4o
+    default: true
+
+# Arch creates a round-robin load balancing between different endpoints, managed via the cluster subsystem.
+endpoints:
+  app_server:
+    # 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:18083
+    # max time to wait for a connection to be established
+    connect_timeout: 0.005s
+
+# default system prompt used by all prompt targets
+system_prompt: |
+  You are a Workforce assistant that helps on workforce planning and HR decision makers with reporting and workfoce planning. NOTHING ELSE. When you get data in json format, offer some summary but don't be too verbose.
+
+prompt_targets:
+    - name: hr_qa
+      endpoint:
+        name: app_server
+        path: /agent/hr_qa
+      description: Handle general Q/A related to HR.
+      default: true
+    - name: workforce
+      description: Get workforce data like headcount and satisfacton levels by region and staffing type
+      endpoint:
+        name: app_server
+        path: /agent/workforce
+      parameters:
+        - name: staffing_type
+          type: str
+          description: Staffing type like contract, fte or agency
+          required: true
+        - name: region
+          type: str
+          required: true
+          description: Geographical region for which you want workforce data like asia, europe, americas.
+        - name: point_in_time
+          type: int
+          required: false
+          description: the point in time for which to retrieve data. For e.g 0 days ago, 30 days ago, etc.
+    - name: slack_message
+      endpoint:
+        name: app_server
+        path: /agent/slack_message
+      description: sends a slack message on a channel
+      parameters:
+        - name: slack_message
+          type: string
+          required: true
+          description: the message that should be sent to a slack channel

demos/hr_agent/docker-compose.yaml

@@ -0,0 +1,23 @@
+services:
+  api_server:
+    build:
+      context: .
+    environment:
+      - SLACK_BOT_TOKEN=${SLACK_BOT_TOKEN:-None}
+    ports:
+      - "18083:80"
+    healthcheck:
+        test: ["CMD", "curl" ,"http://localhost:80/healthz"]
+        interval: 5s
+        retries: 20
+
+  chatbot_ui:
+    build:
+      context: ../../chatbot_ui
+    ports:
+      - "18080:8080"
+    environment:
+      - OPENAI_API_KEY=${OPENAI_API_KEY:?error}
+      - CHAT_COMPLETION_ENDPOINT=http://host.docker.internal:10000/v1
+    extra_hosts:
+      - "host.docker.internal:host-gateway"

demos/hr_agent/main.py

@@ -0,0 +1,106 @@
+import os
+import json
+import pandas as pd
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel, Field
+from typing import Optional
+from enum import Enum
+from slack_sdk import WebClient
+from slack_sdk.errors import SlackApiError
+
+app = FastAPI()
+workforce_data_df = None
+
+with open("workforce_data.json") as file:
+    workforce_data = json.load(file)
+    workforce_data_df = pd.json_normalize(
+        workforce_data, record_path=["regions"], meta=["point_in_time", "satisfaction"]
+    )
+
+
+# Define the request model
+class WorkforceRequset(BaseModel):
+    region: str
+    staffing_type: str
+    point_in_time: Optional[int] = None
+
+
+class SlackRequest(BaseModel):
+    slack_message: str
+
+
+class WorkforceResponse(BaseModel):
+    region: str
+    staffing_type: str
+    headcount: int
+    satisfaction: float
+
+
+# Post method for device summary
+@app.post("/agent/workforce")
+def get_workforce(request: WorkforceRequset):
+    """
+    Endpoint to workforce data by region, staffing type at a given point in time.
+    """
+    region = request.region.lower()
+    staffing_type = request.staffing_type.lower()
+    point_in_time = request.point_in_time if request.point_in_time else 0
+
+    response = {
+        "region": region,
+        "staffing_type": f"Staffing agency: {staffing_type}",
+        "headcount": f"Headcount: {int(workforce_data_df[(workforce_data_df['region']==region) & (workforce_data_df['point_in_time']==point_in_time)][staffing_type].values[0])}",
+        "satisfaction": f"Satisifaction: {float(workforce_data_df[(workforce_data_df['region']==region) & (workforce_data_df['point_in_time']==point_in_time)]['satisfaction'].values[0])}",
+    }
+    return response
+
+
+@app.post("/agent/slack_message")
+def send_slack_message(request: SlackRequest):
+    """
+    Endpoint that sends slack message
+    """
+    slack_message = request.slack_message
+
+    # Load the bot token from an environment variable or replace it directly
+    slack_token = os.getenv(
+        "SLACK_BOT_TOKEN"
+    )  # Replace with your token if needed: 'xoxb-your-token'
+
+    if slack_token is None:
+        print(f"Message for slack: {slack_message}")
+    else:
+        client = WebClient(token=slack_token)
+        channel = "hr_agent_demo"
+        try:
+            # Send the message
+            response = client.chat_postMessage(channel=channel, text=slack_message)
+            return f"Message sent to {channel}: {response['message']['text']}"
+        except SlackApiError as e:
+            print(f"Error sending message: {e.response['error']}")
+
+
+@app.post("/agent/hr_qa")
+async def general_hr_qa():
+    """
+    This method handles Q/A related to general issues in HR.
+    It forwards the conversation to the OpenAI client via a local proxy and returns the response.
+    """
+    return {
+        "choices": [
+            {
+                "message": {
+                    "role": "assistant",
+                    "content": "I am a helpful HR agent, and I can help you plan for workforce related questions",
+                },
+                "finish_reason": "completed",
+                "index": 0,
+            }
+        ],
+        "model": "hr_agent",
+        "usage": {"completion_tokens": 0},
+    }
+
+
+if __name__ == "__main__":
+    app.run(debug=True)

demos/hr_agent/requirements.txt

@@ -0,0 +1,6 @@
+fastapi
+uvicorn
+pydantic
+slack-sdk
+typing
+pandas

demos/hr_agent/run_demo.sh

@@ -0,0 +1,46 @@
+#!/bin/bash
+
+# Function to start the demo
+start_demo() {
+  # Step 1: Check if .env file exists
+  if [ -f ".env" ]; then
+    echo ".env file already exists. Skipping creation."
+  else
+    # Step 2: Create `.env` file and set OpenAI key
+    if [ -z "$OPENAI_API_KEY" ]; then
+      echo "Error: OPENAI_API_KEY environment variable is not set for the demo."
+      exit 1
+    fi
+
+    echo "Creating .env file..."
+    echo "OPENAI_API_KEY=$OPENAI_API_KEY" > .env
+    echo ".env file created with OPENAI_API_KEY."
+  fi
+
+  # Step 3: Start Arch
+  echo "Starting Arch with arch_config.yaml..."
+  archgw up arch_config.yaml
+
+  # Step 4: Start Network Agent
+  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 HR Agent using Docker Compose..."
+  docker compose down
+
+  # Step 2: Stop Arch
+  echo "Stopping Arch..."
+  archgw down
+}
+
+# Main script logic
+if [ "$1" == "down" ]; then
+  stop_demo
+else
+  # Default action is to bring the demo up
+  start_demo
+fi

demos/hr_agent/workforce_data.json

@@ -0,0 +1,29 @@
+[
+    {
+        "point_in_time": 0,
+        "regions": [
+            { "region": "asia", "contract": 100, "fte": 150, "agency": 2000 },
+            { "region": "europe", "contract": 80, "fte": 120, "agency": 2500 },
+            { "region": "americas", "contract": 90, "fte": 200, "agency": 3100 }
+        ],
+        "satisfaction": 3.5
+    },
+    {
+        "point_in_time": 30,
+        "regions": [
+            { "region": "asia", "contract": 110, "fte": 155, "agency": 1000 },
+            { "region": "europe", "contract": 85, "fte": 130, "agency": 1600 },
+            { "region": "americas", "contract": 95, "fte": 210, "agency": 3100 }
+        ],
+        "satisfaction": 4.0
+    },
+    {
+        "point_in_time": 60,
+        "regions": [
+            { "region": "asia", "contract": 115, "fte": 160, "agency": 500 },
+            { "region": "europe", "contract": 90, "fte": 140, "agency": 700 },
+            { "region": "americas", "contract": 100, "fte": 220, "agency": 1200 }
+        ],
+        "satisfaction": 4.7
+    }
+]