add logging

src/llm/llm.py

@@ -2,7 +2,9 @@
 RAG implementation with local Phi-3-mini-4k-instruct-onnx and embeddings
 """
 
+import logging
 import os
+import sys
 from typing import List
 
 # LangChain imports
@@ -26,6 +28,13 @@ from transformers import AutoTokenizer, pipeline
 
 class Phi3LanguageModel:
 
+    def __init__(self):
+        logger = logging.getLogger()
+        logger.setLevel(logging.DEBUG)
+        handler = logging.StreamHandler(sys.stdout)
+        logger.addHandler(handler)
+        self.logger = logger
+
     def extract_assistant_response(self, text):
         if "<|assistant|>" in text:
             return text.split("<|assistant|>")[-1].strip()
@@ -36,7 +45,7 @@ class Phi3LanguageModel:
         # Set up paths to the local model
         base_dir = os.path.dirname(os.path.abspath(__file__))
         model_path = os.path.join(base_dir, "cpu_and_mobile", "cpu-int4-rtn-block-32-acc-level-4")
-        print(f"Loading Phi-3 model from: {model_path}")
+        self.logger.debug(f"Loading Phi-3 model from: {model_path}")
 
         # Load the tokenizer and model
         tokenizer = AutoTokenizer.from_pretrained(
@@ -87,12 +96,12 @@ class Phi3LanguageModel:
         
         try:
             # Get response from the chain
-            print(f'===Prompt: {user_input}\n\n')
+            self.logger.debug(f'===Prompt: {user_input}\n\n')
             response = chain.invoke(user_input)
             # Print the answer
-            print(f'===Response: {response}\n\n')
+            self.logger.debug(f'===Response: {response}\n\n')
             return response
         except Exception as e:
-            print(f"Failed: {e}")
+            self.logger.error(f"Failed: {e}")
             return e