hot fix on energy wizard ref print

NREL · Nov 2, 2023 · c7026b1 · c7026b1
1 parent 32b6b13
commit c7026b1
Showing 1 changed file with 11 additions and 5 deletions.
diff --git a/elm/wizard.py b/elm/wizard.py
@@ -228,6 +228,9 @@ def chat(self, query,
             Question being asked of EnergyWizard
         debug : bool
             Flag to return extra diagnostics on the engineered question.
+        stream : bool
+            Flag to print subsequent chunks of the response in a streaming
+            fashion
         temperature : float
             GPT model temperature, a measure of response entropy from 0 to 1. 0
             is more reliable and nearly deterministic; 1 will give the model
@@ -286,14 +289,17 @@ def chat(self, query,
         else:
             response_message = response["choices"][0]["message"]["content"]
 
-        if stream and print_references and any(references):
-            print('\n\nThe model was provided with the following documents to '
-                  'support its answer:')
-            print(' - ' + '\n - '.join(references))
-
         self.messages.append({'role': 'assistant',
                               'content': response_message})
 
+        if any(references) and print_references:
+            ref_msg = ('\n\nThe model was provided with the '
+                       'following documents to support its answer:')
+            ref_msg += '\n - ' + '\n - '.join(references)
+            response_message += ref_msg
+            if stream:
+                print(ref_msg)
+
         if debug:
             return response_message, query, references
         else: