Added save for agent config

databricks · Oct 8, 2024 · 0ab578e · 0ab578e
1 parent 4669f46
commit 0ab578e
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 0 deletions.
diff --git a/agent_app_sample_code/03_agent_proof_of_concept.py b/agent_app_sample_code/03_agent_proof_of_concept.py
@@ -91,6 +91,15 @@
 
 # COMMAND ----------
 
+# MAGIC %md
+# MAGIC ## Save the agent config
+
+# COMMAND ----------
+
+save_agent_config(agent_config.dict(), './agents/generated_configs/agent.yaml')
+
+# COMMAND ----------
+
 # MAGIC %md
 # MAGIC ## Set the MLflow experiement name
 # MAGIC

diff --git a/agent_app_sample_code/agents/agent_config.py b/agent_app_sample_code/agents/agent_config.py
@@ -13,6 +13,8 @@
 
 from pydantic import BaseModel
 from typing import Literal, Any, List
+import yaml
+import os
 
 # COMMAND ----------
 
@@ -86,3 +88,10 @@ class AgentConfig(BaseModel):
 
 def validate_agent_config(config: dict) -> None:
   AgentConfig.parse_obj(config)
+
+def save_agent_config(config: dict, file_path: str) -> None:
+  # Ensure the directory exists
+  os.makedirs(os.path.dirname(file_path), exist_ok=True)
+
+  with open(file_path, 'w') as file:
+      yaml.dump(config, file)