From 296669ac4999cbaf3b2d15b0c60b57ca0f445896 Mon Sep 17 00:00:00 2001
From: Adam Wilson <lightbrok3r@gmail.com>
Date: Sat, 19 Jul 2025 07:59:54 -0600
Subject: [PATCH] add env var for prompt template dir; error handling

---
 run.sh                                                       | 2 ++
 src/text_generation/adapters/prompt_template_repository.py   | 5 +++--
 .../services/nlp/text_generation_completion_service.py       | 4 ++++
 3 files changed, 9 insertions(+), 2 deletions(-)

diff --git a/run.sh b/run.sh
index c584a1414..8277689d2 100755
--- a/run.sh
+++ b/run.sh
@@ -30,11 +30,13 @@ git lfs install
 pip install -r ./requirements.txt
 
 # environment variables
+export PROMPT_TEMPLATES_DIR="./infrastructure/prompt_templates"
 export MODEL_BASE_DIR="./infrastructure/foundation_model"
 export MODEL_CPU_DIR="cpu_and_mobile/cpu-int4-rtn-block-32-acc-level-4"
 MODEL_DATA_FILENAME="phi3-mini-4k-instruct-cpu-int4-rtn-block-32-acc-level-4.onnx.data"
 MODEL_DATA_FILEPATH="$MODEL_BASE_DIR/$MODEL_CPU_DIR/$MODEL_DATA_FILENAME"
 
+
 # get foundation model dependencies from HuggingFace / Microsoft
 if [ ! -f "$MODEL_DATA_FILEPATH" ]; then
     echo "Downloading foundation model..."
diff --git a/src/text_generation/adapters/prompt_template_repository.py b/src/text_generation/adapters/prompt_template_repository.py
index 807a7edc4..91fba2864 100644
--- a/src/text_generation/adapters/prompt_template_repository.py
+++ b/src/text_generation/adapters/prompt_template_repository.py
@@ -13,10 +13,11 @@ class PromptTemplateRepository(AbstractPromptTemplateRepository):
         return os.path.join(self.templates_dir, template_filename)
 
     def get(self, id: str) -> PromptTemplate:
+        path = self._create_path_from_id(id)
         try:
-            return load_prompt(self._create_path_from_id(id))
+            return load_prompt(path)
         except Exception as e:
-            print(e)
+            print(f'Failed to load template from path "{path}":\n{e}')
             return None
     
     def add(self, id: str, prompt_template: PromptTemplate) -> None:
diff --git a/src/text_generation/services/nlp/text_generation_completion_service.py b/src/text_generation/services/nlp/text_generation_completion_service.py
index ed74c7f44..65e056970 100644
--- a/src/text_generation/services/nlp/text_generation_completion_service.py
+++ b/src/text_generation/services/nlp/text_generation_completion_service.py
@@ -186,6 +186,10 @@ class TextGenerationCompletionService(
         prompt_template = self.prompt_template_service.get(
             id=self.constants.PromptTemplateIds.PHI_3_MINI_4K_INSTRUCT_BASIC
         )
+
+        if prompt_template is None:
+            raise ValueError(f"Prompt template not found for ID: {self.constants.PromptTemplateIds.PHI_3_MINI_4K_INSTRUCT_BASIC}")
+    
         return (
             { "question": RunnablePassthrough() }
             | prompt_template