From 9c00212ce3060aee432e753e78fa9c3a6521bf22 Mon Sep 17 00:00:00 2001
From: Erni Durdevic <erni.durdevic@databricks.com>
Date: Fri, 11 Oct 2024 23:35:38 +0200
Subject: [PATCH] Added save for agent config (#35)

* Added save for agent config

* Removed default agent config file
---
 .../03_agent_proof_of_concept.py              |  9 +++++
 agent_app_sample_code/agents/agent_config.py  |  9 +++++
 .../agents/generated_configs/agent.yaml       | 40 -------------------
 3 files changed, 18 insertions(+), 40 deletions(-)
 delete mode 100644 agent_app_sample_code/agents/generated_configs/agent.yaml

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
diff --git a/agent_app_sample_code/agents/generated_configs/agent.yaml b/agent_app_sample_code/agents/generated_configs/agent.yaml
deleted file mode 100644
index e09fa95..0000000
--- a/agent_app_sample_code/agents/generated_configs/agent.yaml
+++ /dev/null
@@ -1,40 +0,0 @@
-input_example:
-  messages:
-  - content: What is RAG?
-    role: user
-llm_config:
-  llm_endpoint_name: databricks-meta-llama-3-1-70b-instruct
-  llm_parameters:
-    max_tokens: 1500
-    temperature: 0.01
-  llm_system_prompt_template: You are a helpful assistant that answers questions by
-    calling tools.  Provide responses ONLY based on the information from tools that
-    are explictly specified to you.  If you do not have a relevant tool for a question,
-    respond with 'Sorry, I'm not trained to answer that question'.
-retriever_tool:
-  chunk_template: 'Passage text: {chunk_text}
-
-    Passage metadata: {metadata}
-
-
-    '
-  parameters:
-    num_results: 5
-    query_type: ann
-  prompt_template: 'Use the following pieces of retrieved context to answer the question.
-
-    Only use the passages from context that are relevant to the query to answer the
-    question, ignore the irrelevant passages.  When responding, cite your source,
-    referring to the passage by the columns in the passage''s metadata.
-
-
-    Context: {context}'
-  tool_description_prompt: Search for documents that are relevant to a user's query
-    about the [REPLACE WITH DESCRIPTION OF YOUR DOCS].
-  vector_search_index: prithvikannan_catalog.cookbook.my_agent_app_chunked_docs_index
-  vector_search_schema:
-    additional_metadata_columns: []
-    chunk_text: content_chunked
-    document_uri: doc_uri
-    primary_key: chunk_id
-  vector_search_threshold: 0.1