diff --git a/api/main.py b/api/main.py index f518325..7a05ba9 100644 --- a/api/main.py +++ b/api/main.py @@ -1,3 +1,4 @@ +import os from fastapi import FastAPI from typing import Optional from temporalio.client import Client @@ -11,7 +12,7 @@ from workflows.agent_goal_workflow import AgentGoalWorkflow from models.data_types import CombinedInput, AgentGoalWorkflowParams from tools.goal_registry import goal_list from fastapi.middleware.cors import CORSMiddleware -from shared.config import get_temporal_client, TEMPORAL_TASK_QUEUE, AGENT_GOAL +from shared.config import get_temporal_client, TEMPORAL_TASK_QUEUE app = FastAPI() temporal_client: Optional[Client] = None @@ -22,13 +23,10 @@ load_dotenv() def get_initial_agent_goal(): """Get the agent goal from environment variables.""" - if AGENT_GOAL is not None: - for listed_goal in goal_list: - if listed_goal.id == AGENT_GOAL: - return listed_goal - else: - #if no goal is set in the config file, default to choosing an agent - return goal_list.get("goal_choose_agent_type") + env_goal = os.getenv("AGENT_GOAL", "goal_choose_agent_type") #if no goal is set in the env file, default to choosing an agent + for listed_goal in goal_list: + if listed_goal.id == env_goal: + return listed_goal @app.on_event("startup") diff --git a/shared/config.py b/shared/config.py index cb4b9da..9590634 100644 --- a/shared/config.py +++ b/shared/config.py @@ -16,11 +16,6 @@ TEMPORAL_TLS_CERT = os.getenv("TEMPORAL_TLS_CERT", "") TEMPORAL_TLS_KEY = os.getenv("TEMPORAL_TLS_KEY", "") TEMPORAL_API_KEY = os.getenv("TEMPORAL_API_KEY", "") -#Starting agent goal - 1st goal is always to help user pick a next goal -AGENT_GOAL = "goal_choose_agent_type" -#AGENT_GOAL = "goal_event_flight_invoice" - - async def get_temporal_client() -> Client: """ Creates a Temporal client based on environment configuration.