From 28518e6c170b21f6067028e478b111ea9c4937d0 Mon Sep 17 00:00:00 2001
From: Adam Wilson <lightbrok3r@gmail.com>
Date: Wed, 30 Jul 2025 07:17:48 -0600
Subject: [PATCH] fixes for basic prompt template

---
 .../phi-3-mini-4k-instruct.01-basic.json      |  4 +--
 .../nlp/text_generation_completion_service.py | 36 ++++++++++++-------
 .../test_prompt_template_repository.py        |  4 +--
 3 files changed, 28 insertions(+), 16 deletions(-)

diff --git a/infrastructure/prompt_templates/phi-3-mini-4k-instruct.01-basic.json b/infrastructure/prompt_templates/phi-3-mini-4k-instruct.01-basic.json
index d967211da..0e526e837 100644
--- a/infrastructure/prompt_templates/phi-3-mini-4k-instruct.01-basic.json
+++ b/infrastructure/prompt_templates/phi-3-mini-4k-instruct.01-basic.json
@@ -1,14 +1,14 @@
 {
     "name": null,
     "input_variables": [
-        "question"
+        "input"
     ],
     "optional_variables": [],
     "output_parser": null,
     "partial_variables": {},
     "metadata": null,
     "tags": null,
-    "template": "<|system|>\n            Act as a conversational agent to respond to the end user's prompt. \n            <|user|>\n            Question: {question}<|end|>\n            <|assistant|>\n            ",
+    "template": "<|system|>\n            Act as a conversational agent to respond to the end user's prompt. \n            <|user|>\n            Question: {input}<|end|>\n            <|assistant|>\n            ",
     "template_format": "f-string",
     "validate_template": false,
     "_type": "prompt"
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 1fdb1f3c1..d95a6510b 100644
--- a/src/text_generation/services/nlp/text_generation_completion_service.py
+++ b/src/text_generation/services/nlp/text_generation_completion_service.py
@@ -1,7 +1,8 @@
-from langchain.prompts import PromptTemplate
+from langchain.prompts import StringPromptTemplate
 from langchain_core.output_parsers import StrOutputParser
 from langchain_core.runnables import RunnablePassthrough, RunnableConfig
 from langchain_huggingface import HuggingFacePipeline
+from langchain_core.prompt_values import PromptValue
 
 from src.text_generation.common.constants import Constants
 from src.text_generation.domain.alternate_completion_result import AlternateCompletionResult
@@ -152,14 +153,31 @@ class TextGenerationCompletionService(
     def _handle_without_guidelines(self, user_prompt: str) -> TextGenerationCompletionResult:
         """Handle: CoT=False, RAG=False"""
         try:
-            chain = self._create_chain_without_guidelines()
+            prompt_template: StringPromptTemplate = self.prompt_template_service.get(
+                id=self.constants.PromptTemplateIds.PHI_3_MINI_4K_INSTRUCT__01_BASIC
+            )
+
+            if prompt_template is None:
+                raise ValueError(f"Prompt template not found for ID: {self.constants.PromptTemplateIds.PHI_3_MINI_4K_INSTRUCT__01_BASIC}")
+            
+            chain = self._create_chain_without_guidelines(prompt_template)
             llm_config = self.llm_configuration_introspection_service.get_config(chain)
 
+            prompt_value: PromptValue = prompt_template.format_prompt(input=user_prompt)
+            prompt_dict = {
+                "messages": [
+                    {"role": msg.type, "content": msg.content, "additional_kwargs": msg.additional_kwargs}
+                    for msg in prompt_value.to_messages()
+                ],
+                "string_representation": prompt_value.to_string(),
+            }
+
             result = TextGenerationCompletionResult(
                 original_result=OriginalCompletionResult(
                     user_prompt=user_prompt,
-                    completion_text=chain.invoke(user_prompt),
-                    llm_config=llm_config
+                    completion_text=chain.invoke({ self.constants.INPUT_VARIABLE_TOKEN: user_prompt }),
+                    llm_config=llm_config,
+                    full_prompt=prompt_dict
             ))
             return self._process_completion_result(result)
         except Exception as e:
@@ -198,16 +216,10 @@ class TextGenerationCompletionService(
         self._use_reflexion_guardrails = True
         return self
 
-    def _create_chain_without_guidelines(self):
-        prompt_template = self.prompt_template_service.get(
-            id=self.constants.PromptTemplateIds.PHI_3_MINI_4K_INSTRUCT__01_BASIC
-        )
-
-        if prompt_template is None:
-            raise ValueError(f"Prompt template not found for ID: {self.constants.PromptTemplateIds.PHI_3_MINI_4K_INSTRUCT__01_BASIC}")
+    def _create_chain_without_guidelines(self, prompt_template):
     
         return (
-            { "question": RunnablePassthrough() }
+            { f"{self.constants.INPUT_VARIABLE_TOKEN}": RunnablePassthrough() }
             | prompt_template
             | self.foundation_model_pipeline
             | StrOutputParser()
diff --git a/tests/integration/test_prompt_template_repository.py b/tests/integration/test_prompt_template_repository.py
index 2c4436771..68b99cbb2 100644
--- a/tests/integration/test_prompt_template_repository.py
+++ b/tests/integration/test_prompt_template_repository.py
@@ -13,10 +13,10 @@ def test_1_save_basic_template(constants, prompt_template_service):
         template=f"""{constants.PHI_3_SYSTEM_START_TOKEN}
             Act as a conversational agent to respond to the end user's prompt. 
             {constants.PHI_3_USER_START_TOKEN}
-            Question: {{question}}{constants.PHI_3_END_TOKEN}
+            Question: {{{constants.INPUT_VARIABLE_TOKEN}}}{constants.PHI_3_END_TOKEN}
             {constants.PHI_3_ASSISTANT_START_TOKEN}
             """,
-        input_variables=["question"]
+        input_variables=[constants.INPUT_VARIABLE_TOKEN]
     )
     id = constants.PromptTemplateIds.PHI_3_MINI_4K_INSTRUCT__01_BASIC