feat: introduce inferring of log levels

fal-ai · Dec 10, 2024 · 0f68ae9 · 0f68ae9
1 parent 6d20281
commit 0f68ae9
Show file tree

Hide file tree

Showing 3 changed files with 35 additions and 6 deletions.
diff --git a/src/isolate/connections/_local/_base.py b/src/isolate/connections/_local/_base.py
@@ -115,7 +115,9 @@ def start_process(
         python_executable = get_executable_path(self.environment_path, "python")
         with logged_io(
             partial(
-                self.handle_agent_log, source=LogSource.USER, level=LogLevel.STDOUT
+                self.handle_agent_log,
+                source=LogSource.USER,
+                level=None,  # Will be inferred
             ),
             partial(
                 self.handle_agent_log, source=LogSource.USER, level=LogLevel.STDERR
@@ -173,7 +175,30 @@ def get_python_cmd(
         """Return the command to run the agent process with."""
         raise NotImplementedError
 
-    def handle_agent_log(self, line: str, level: LogLevel, source: LogSource) -> None:
+    def infer_log_level(self, line: str, source: LogSource) -> LogLevel:
+        """Infer the log level of the given line."""
+        if source in (LogSource.BUILDER, LogSource.BRIDGE):
+            return LogLevel.TRACE
+        else:
+            if "[error]" in line.lower():
+                return LogLevel.ERROR
+            elif "[warning]" in line.lower():
+                return LogLevel.WARNING
+            elif "[warn]" in line.lower():
+                return LogLevel.WARNING
+            elif "[info]" in line.lower():
+                return LogLevel.INFO
+            elif "[debug]" in line.lower():
+                return LogLevel.DEBUG
+            elif "[trace]" in line.lower():
+                return LogLevel.TRACE
+
+        # Default
+        return LogLevel.INFO
+
+    def handle_agent_log(
+        self, line: str, *, level: LogLevel | None, source: LogSource
+    ) -> None:
         """Handle a log line emitted by the agent process. The level will be either
         STDOUT or STDERR."""
         raise NotImplementedError
diff --git a/src/isolate/connections/grpc/_base.py b/src/isolate/connections/grpc/_base.py
@@ -147,5 +147,7 @@ def get_python_cmd(
             str(log_fd),
         ]
 
-    def handle_agent_log(self, line: str, level: LogLevel, source: LogSource) -> None:
-        self.log(line, level=level, source=source)
+    def handle_agent_log(
+        self, line: str, *, level: LogLevel | None, source: LogSource
+    ) -> None:
+        self.log(line, level=level or self.infer_log_level(line, source), source=source)
diff --git a/src/isolate/connections/ipc/_base.py b/src/isolate/connections/ipc/_base.py
@@ -219,5 +219,7 @@ def get_python_cmd(
             str(log_fd),
         ]
 
-    def handle_agent_log(self, line: str, level: LogLevel, source: LogSource) -> None:
-        self.log(line, level=level, source=source)
+    def handle_agent_log(
+        self, line: str, *, level: LogLevel | None, source: LogSource
+    ) -> None:
+        self.log(line, level=level or self.infer_log_level(line, source), source=source)