added option to energy wizard to return openai chat object instead of…

… text

elm/wizard.py

@@ -218,7 +218,8 @@ def chat(self, query,
              convo=False,
              token_budget=None,
              new_info_threshold=0.7,
-             print_references=False):
+             print_references=False,
+             return_chat_obj=False):
         """Answers a query by doing a semantic search of relevant text with
         embeddings and then sending engineered query to the LLM.
 
@@ -248,6 +249,8 @@ def chat(self, query,
         print_references : bool
             Flag to print references if EnergyWizard is initialized with a
             valid ref_col.
+        return_chat_obj : bool
+            Flag to only return the ChatCompletion from OpenAI API.
 
         Returns
         -------
@@ -279,6 +282,9 @@ def chat(self, query,
 
         response = openai.ChatCompletion.create(**kwargs)
 
+        if return_chat_obj:
+            return response, query, references
+
         if stream:
             for chunk in response:
                 chunk_msg = chunk['choices'][0]['delta']