Merge pull request #15 from tomchapin/feature/fixing-logging

Moving logger into a more granular setup

OS/01/device.py

@@ -1,7 +1,6 @@
 import asyncio
 import threading
 import os
-import logging
 import pyaudio
 from starlette.websockets import WebSocket
 from queue import Queue
@@ -23,8 +22,9 @@ from interpreter import interpreter # Just for code execution. Maybe we should l
 from utils.kernel import put_kernel_messages_into_queue
 from stt import stt_wav
 
-# Configure logging
-logging.basicConfig(format='%(message)s', level=logging.getLevelName(os.getenv('DEBUG_LEVEL', 'INFO').upper()))
+from utils.logs import setup_logging
+from utils.logs import logger
+setup_logging()
 
 # Configuration for Audio Recording
 CHUNK = 1024  # Record in chunks of 1024 samples
@@ -55,7 +55,7 @@ def record_audio():
 
     """Record audio from the microphone and add it to the queue."""
     stream = p.open(format=FORMAT, channels=CHANNELS, rate=RATE, input=True, frames_per_buffer=CHUNK)
-    logging.info("Recording started...")
+    logger.info("Recording started...")
     global RECORDING
 
     # Create a temporary WAV file to store the audio data
@@ -73,7 +73,7 @@ def record_audio():
     wav_file.close()
     stream.stop_stream()
     stream.close()
-    logging.info("Recording stopped.")
+    logger.info("Recording stopped.")
 
     duration = wav_file.getnframes() / RATE
     if duration < 0.3:
@@ -124,7 +124,7 @@ def on_release(key):
     if key == keyboard.Key.space:
         toggle_recording(False)
     elif key == keyboard.Key.esc:
-        logging.info("Exiting...")
+        logger.info("Exiting...")
         os._exit(0)
 
 import asyncio
@@ -141,7 +141,7 @@ async def websocket_communication(WS_URL):
     while True:
         try:
             async with websockets.connect(WS_URL) as websocket:
-                logging.info("Press the spacebar to start/stop recording. Press ESC to exit.")
+                logger.info("Press the spacebar to start/stop recording. Press ESC to exit.")
                 asyncio.create_task(message_sender(websocket))
 
                 initial_message = {"role": None, "type": None, "format": None, "content": None} 
@@ -150,16 +150,18 @@ async def websocket_communication(WS_URL):
                 while True:
                     message = await websocket.recv()
 
-                    logging.info(f"Got this message from the server: {type(message)} {message}")
+                    logger.debug(f"Got this message from the server: {type(message)} {message}")
 
                     if type(message) == str:
                         message = json.loads(message)
 
                     if message.get("end"):
-                        logging.info(f"Complete message from the server: {message_so_far}")
+                        logger.debug(f"Complete message from the server: {message_so_far}")
+                        logger.info("\n")
                         message_so_far = initial_message
 
                     if "content" in message:
+                        print(message['content'], end="", flush=True)
                         if any(message_so_far[key] != message[key] for key in message_so_far if key != "content"):
                             message_so_far = message
                         else:
@@ -187,7 +189,7 @@ async def websocket_communication(WS_URL):
   
         except:
             # traceback.print_exc()
-            logging.info(f"Connecting to `{WS_URL}`...")
+            logger.info(f"Connecting to `{WS_URL}`...")
             await asyncio.sleep(2)
             
 

OS/01/server.py

@@ -4,7 +4,6 @@ import json
 import time
 import queue
 import os
-import logging
 import traceback
 from queue import Queue
 from threading import Thread
@@ -23,8 +22,10 @@ from utils.kernel import put_kernel_messages_into_queue
 from i import configure_interpreter
 from interpreter import interpreter
 
-# Configure logging
-logging.basicConfig(format='%(message)s', level=logging.getLevelName(os.getenv('DEBUG_LEVEL', 'INFO').upper()))
+from utils.logs import setup_logging
+from utils.logs import logger
+setup_logging()
+
 
 app = FastAPI()
 
@@ -68,10 +69,10 @@ if os.getenv('CODE_RUNNER') == "device":
                 to_device.put({"role": "assistant", "type": "code", "format": "python", "end": True})
             
             # Stream the response
-            logging.info("Waiting for the device to respond...")
+            logger.info("Waiting for the device to respond...")
             while True:
                 chunk = from_computer.get()
-                logging.info(f"Server received from device: {chunk}")
+                logger.info(f"Server received from device: {chunk}")
                 if "end" in chunk:
                     break
                 yield chunk
@@ -98,7 +99,7 @@ async def websocket_endpoint(websocket: WebSocket):
         await asyncio.gather(receive_task, send_task)
     except Exception as e:
         traceback.print_exc()
-        logging.info(f"Connection lost. Error: {e}")
+        logger.info(f"Connection lost. Error: {e}")
 
 async def receive_messages(websocket: WebSocket):
     while True:
@@ -113,7 +114,7 @@ async def receive_messages(websocket: WebSocket):
 async def send_messages(websocket: WebSocket):
     while True:
         message = await to_device.get()
-        logging.debug(f"Sending to the device: {type(message)} {message}")
+        logger.debug(f"Sending to the device: {type(message)} {message}")
         await websocket.send_json(message)
 
 async def listener():
@@ -163,7 +164,7 @@ async def listener():
         
         for chunk in interpreter.chat(messages, stream=True, display=False):
 
-            logging.debug("Got chunk:", chunk)
+            logger.debug("Got chunk:", chunk)
 
             # Send it to the user
             await to_device.put(chunk)
@@ -199,7 +200,7 @@ async def listener():
                 with open(conversation_history_path, 'w') as file:
                     json.dump(interpreter.messages, file, indent=4)
 
-                logging.info("New user message recieved. Breaking.")
+                logger.info("New user message recieved. Breaking.")
                 break
 
             # Also check if there's any new computer messages
@@ -208,7 +209,7 @@ async def listener():
                 with open(conversation_history_path, 'w') as file:
                     json.dump(interpreter.messages, file, indent=4)
 
-                logging.info("New computer message recieved. Breaking.")
+                logger.info("New computer message recieved. Breaking.")
                 break
         else:
             with open(conversation_history_path, 'w') as file:
@@ -243,7 +244,7 @@ if __name__ == "__main__":
         if not server_url:
             raise ValueError("The environment variable SERVER_URL is not set. Please set it to proceed.")
         parsed_url = urllib.parse.urlparse(server_url)
-        logging.info("Starting `server.py`...")
+        logger.info("Starting `server.py`...")
 
         config = Config(app, host=parsed_url.hostname, port=parsed_url.port, lifespan='on')
         server = Server(config)

OS/01/start.sh

@@ -21,8 +21,8 @@ export STT_RUNNER=device # If server, audio will be sent over websocket.
 export SERVER_EXPOSE_PUBLICALLY=False
 
 # Debug level
-# export DEBUG_LEVEL=DEBUG
-export DEBUG_LEVEL="INFO"
+# export LOG_LEVEL=DEBUG
+export LOG_LEVEL="INFO"
 
 ### SETUP
 

OS/01/stt.py

@@ -4,7 +4,6 @@ Defines a function which takes a path to an audio file and turns it into text.
 
 from datetime import datetime
 import os
-import logging
 import contextlib
 import tempfile
 import ffmpeg
@@ -12,8 +11,9 @@ import subprocess
 import openai
 from openai import OpenAI
 
-# Configure logging
-logging.basicConfig(format='%(message)s', level=logging.getLevelName(os.getenv('DEBUG_LEVEL', 'INFO').upper()))
+from utils.logs import setup_logging
+from utils.logs import logger
+setup_logging()
 
 client = OpenAI()
 
@@ -85,10 +85,10 @@ def stt_wav(wav_file_path: str):
                 response_format="text"
             )
         except openai.BadRequestError as e:
-            logging.info(f"openai.BadRequestError: {e}")
+            logger.info(f"openai.BadRequestError: {e}")
             return None
 
-        logging.info(f"Transcription result: {transcript}")
+        logger.info(f"Transcription result: {transcript}")
         return transcript
     else:
         temp_dir = tempfile.gettempdir()

OS/01/utils/kernel.py

@@ -1,11 +1,10 @@
 import asyncio
 import subprocess
 import platform
-import os
-import logging
 
-# Configure logging
-logging.basicConfig(format='%(message)s', level=logging.getLevelName(os.getenv('DEBUG_LEVEL', 'INFO').upper()))
+from utils.logs import setup_logging
+from utils.logs import logger
+setup_logging()
 
 def get_kernel_messages():
     """
@@ -21,7 +20,7 @@ def get_kernel_messages():
         with open('/var/log/dmesg', 'r') as file:
             return file.read()
     else:
-        logging.info("Unsupported platform.")
+        logger.info("Unsupported platform.")
 
 def custom_filter(message):
     # Check for {TO_INTERPRETER{ message here }TO_INTERPRETER} pattern

OS/01/utils/logs.py

@@ -0,0 +1,22 @@
+import os
+import logging
+
+logger: logging.Logger = logging.getLogger("01")
+root_logger: logging.Logger = logging.getLogger()
+
+
+def _basic_config() -> None:
+    logging.basicConfig(
+        format="%(message)s"
+    )
+
+
+def setup_logging() -> None:
+    env = os.environ.get("LOG_LEVEL", "").upper()
+    if env == "DEBUG":
+        _basic_config()
+        logger.setLevel(logging.DEBUG)
+        root_logger.setLevel(logging.DEBUG)
+    elif env == "INFO":
+        _basic_config()
+        logger.setLevel(logging.INFO)