From 0ab578e11831700349b1b99b8c09f7bd5efc336c Mon Sep 17 00:00:00 2001
From: Erni Durdevic <erni.durdevic@databricks.com>
Date: Tue, 8 Oct 2024 11:18:15 +0200
Subject: [PATCH] Added save for agent config

---
 agent_app_sample_code/03_agent_proof_of_concept.py | 9 +++++++++
 agent_app_sample_code/agents/agent_config.py       | 9 +++++++++
 2 files changed, 18 insertions(+)

diff --git a/agent_app_sample_code/03_agent_proof_of_concept.py b/agent_app_sample_code/03_agent_proof_of_concept.py
index 7b389a2..2f488ee 100644
--- 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
index 0f803b3..e90ecac 100644
--- 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)
\ No newline at end of file