From ece3ac1d3c44de50cff1a12d3d3bf0d22221e6f0 Mon Sep 17 00:00:00 2001
From: Laine <lainelminor@gmail.com>
Date: Thu, 13 Mar 2025 14:53:03 -0400
Subject: [PATCH] Add the category tag to the goals and example env file, and
 filter the results based on tags in list_agents

---
 .env.example               |  7 +++++--
 models/tool_definitions.py |  1 +
 tools/goal_registry.py     |  4 ++++
 tools/list_agents.py       | 28 ++++++++++++++++++++++------
 4 files changed, 32 insertions(+), 8 deletions(-)

diff --git a/.env.example b/.env.example
index 4701495..ad5cffd 100644
--- a/.env.example
+++ b/.env.example
@@ -35,8 +35,11 @@ OPENAI_API_KEY=sk-proj-...
 # Uncomment if using API key (not needed for local dev server)
 # TEMPORAL_API_KEY=abcdef1234567890
 
-# Agent Goal Configuration
-# AGENT_GOAL=goal_event_flight_invoice  # (default) or goal_match_train_invoice
+# Set starting goal of agent
+AGENT_GOAL=goal_choose_agent_type  # (default) 
+
+#Choose which category(ies) of goals you want to be listed by the Agent - options are system (always included), hr, travel, or all. 
+GOAL_CATEGORIES=hr,travel # default is all
 
 # Set if the UI should force a user confirmation step or not
 SHOW_CONFIRM=True
\ No newline at end of file
diff --git a/models/tool_definitions.py b/models/tool_definitions.py
index 76d9ab3..59cf27b 100644
--- a/models/tool_definitions.py
+++ b/models/tool_definitions.py
@@ -18,6 +18,7 @@ class ToolDefinition:
 @dataclass
 class AgentGoal:
     id: str
+    category_tag: str
     agent_name: str
     agent_friendly_description: str
     tools: List[ToolDefinition]
diff --git a/tools/goal_registry.py b/tools/goal_registry.py
index 6ffa293..055fee3 100644
--- a/tools/goal_registry.py
+++ b/tools/goal_registry.py
@@ -6,6 +6,7 @@ starter_prompt_generic = "Welcome me, give me a description of what you can do,
 
 goal_choose_agent_type = AgentGoal(
     id = "goal_choose_agent_type",
+    category_tag="system",
     agent_name="Choose Agent",
     agent_friendly_description="Choose the type of agent to assist you today.",
     tools=[
@@ -33,6 +34,7 @@ goal_choose_agent_type = AgentGoal(
 
 goal_match_train_invoice = AgentGoal(
     id = "goal_match_train_invoice",
+    category_tag="travel",
     agent_name="UK Premier League Match Trip Booking",
     agent_friendly_description="Book a trip to a city in the UK around the dates of a premier league match.",
     tools=[
@@ -80,6 +82,7 @@ goal_match_train_invoice = AgentGoal(
 
 goal_event_flight_invoice = AgentGoal(
     id = "goal_event_flight_invoice",
+    category_tag="travel",
     agent_name="Australia and New Zealand Event Flight Booking",
     agent_friendly_description="Book a trip to a city in Australia or New Zealand around the dates of events in that city.",    
     tools=[
@@ -119,6 +122,7 @@ goal_event_flight_invoice = AgentGoal(
 # This goal uses the data/employee_pto_data.json file as dummy data.
 goal_hr_schedule_pto = AgentGoal(
     id = "goal_hr_schedule_pto",
+    category_tag="hr",
     agent_name="Schedule PTO",
     agent_friendly_description="Schedule PTO based on your available PTO.",   
     tools=[
diff --git a/tools/list_agents.py b/tools/list_agents.py
index 1fc56f3..ba9aa20 100644
--- a/tools/list_agents.py
+++ b/tools/list_agents.py
@@ -1,16 +1,32 @@
+import os
 import tools.goal_registry as goals
 
 def list_agents(args: dict) -> dict:
 
+    goal_categories_start = os.getenv("GOAL_CATEGORIES")
+    if goal_categories_start is None:
+        goal_categories = ["all"] # default to 'all' categories
+    else:
+        goal_categories_start.strip().lower() # handle extra spaces or non-lowercase
+        goal_categories = goal_categories_start.split(",")
+
+    # always show goals labeled as "system," like the goal chooser
+    if "system" not in goal_categories:
+        goal_categories.append("system")
+
     agents = []
     if goals.goal_list is not None:
         for goal in goals.goal_list:
-            agents.append(
-                {
-                    "agent_name": goal.agent_name,
-                    "goal_id": goal.id,
-                    "agent_description": goal.agent_friendly_description,
-                }
+            # add to list if either
+            #   - all
+            #   - current goal's tag is in goal_categories
+            if "all" in goal_categories or goal.category_tag in goal_categories:
+                agents.append(
+                    {
+                        "agent_name": goal.agent_name,
+                        "goal_id": goal.id,
+                        "agent_description": goal.agent_friendly_description,
+                    }
             )
     return {
         "agents": agents,