Add logging to Flo library

.gitignore

@@ -11,4 +11,5 @@ dataset
 bin
 .DS_Store
 *.sql
-*.png
+*.png
+*.log

examples/simple_blogging_team.py

@@ -33,13 +33,13 @@
 """
 
 llm = ChatOpenAI(temperature=0, model_name='gpt-4o-mini')
-session = FloSession(llm).register_tool(
+session = FloSession(llm, log_level="INFO").register_tool(
     name="TavilySearchResults", 
     tool=TavilySearchResults()
 ).register_tool(
     name="DummyTool", 
     tool=TavilySearchResults(description="Tool is a dummy tool, dont use this")
 )
-flo: Flo = Flo.build(session, yaml=yaml_data)
+flo: Flo = Flo.build(session, yaml=yaml_data, log_level="INFO")
 data = flo.invoke(input_prompt)
 print((data['messages'][-1]).content)

flo_ai/__init__.py

@@ -5,4 +5,5 @@
 from flo_ai.router.flo_linear import FloLinear
 from flo_ai.state.flo_session import FloSession
 from flo_ai.retrievers.flo_retriever import FloRagBuilder
-
+from flo_ai.common.flo_logger import get_logger
+from flo_ai.common.flo_langchain_logger import FloLangchainLogger

flo_ai/common/README.md

@@ -0,0 +1,86 @@
+# Understanding Log Levels in FloAI
+
+FloAI uses standard Python logging levels to indicate the severity of logged messages. Here are the common levels used (from least to most severe):
+
+- **DEBUG**: Detailed information for debugging purposes.
+- **INFO**: Informational messages about the normal operation of the system.
+- **WARNING**: Potential issues or unexpected conditions.
+- **ERROR**: Errors that may have caused the system to malfunction.
+- **CRITICAL**: Critical errors that have caused the system to crash.
+
+By adjusting the log level, you can control how much information is logged and the verbosity of the output.
+
+## Controlling Log Levels
+
+FloAI provides multiple ways to control log levels:
+
+### 1. Environment Variables for Log Level Control
+
+Export environment variables to set the log level for specific components before running the application:
+
+- `FLO_LOG_LEVEL_COMMON`: Controls the level for the "CommonLogs" logger.
+  - CommonLogs: General-purpose logging used across the entire FloAI system. It captures broad, system-wide events and information.
+
+- `FLO_LOG_LEVEL_BUILDER`: Controls the level for the "BuilderLogs" logger.
+  - BuilderLogs: Specific to the process of building and configuring FloAI instances. It logs information about YAML parsing, component creation, and FloAI structure setup.
+
+- `FLO_LOG_LEVEL_SESSION`: Controls the level for the "SessionLogs" logger. 
+  - SessionLogs: Dedicated to logging session-specific information. It captures events and data related to individual FloAI sessions, including session creation, tool registration, and session-level operations.
+
+These loggers allow for granular control over logging output in different parts of the FloAI system. By adjusting their levels independently, you can focus on debugging or monitoring specific aspects of FloAI's operation.
+
+Example:
+
+```bash
+export FLO_LOG_LEVEL_COMMON=DEBUG
+export FLO_LOG_LEVEL_BUILDER=INFO
+export FLO_LOG_LEVEL_SESSION=WARNING
+```
+
+### 2. FloSession Creation
+
+When creating a FloSession object, you can specify the desired log level:
+
+```python
+session = FloSession(llm, log_level="DEBUG")
+```
+
+This session will log messages at DEBUG level and above.
+
+### 3. Flo Instance Creation
+
+Similar to FloSession, you can set the log level when creating a Flo instance:
+
+```python
+flo = Flo.build(session, yaml_config, log_level="DEBUG")
+```
+
+This Flo instance will inherit the specified log level.
+
+### 4. Global Log Level Change (Runtime)
+
+You can dynamically change the global log level at runtime using the set_global_log_level function from flo_ai.common.logging_config:
+
+```python
+from flo_ai.common.logging_config import set_global_log_level
+
+set_global_log_level("DEBUG")  # Set the global log level to DEBUG
+```
+
+This will affect all logging throughout the application.
+
+### 5. Specific Logger Level Change (Runtime)
+
+If you need to adjust the level for a specific logger, use the set_log_level method of the FloLogger class:
+
+```python
+from flo_ai.common.logging_config import FloLogger
+
+FloLogger.set_log_level("COMMON", "DEBUG")  # Set COMMON logger to DEBUG level
+```
+
+## Best Practices
+
+- **Environment variables**: Use these for setting default levels without modifying code, useful in different deployment environments.
+- **Object creation**: This approach allows setting specific levels for individual sessions or Flo instances.
+- **Runtime changes**: Use these methods for dynamic adjustments during program execution.

flo_ai/common/flo_langchain_logger.py

@@ -0,0 +1,53 @@
+from typing import Any, Dict, List, Union
+from langchain.callbacks.base import BaseCallbackHandler
+from langchain.schema import AgentAction, AgentFinish, LLMResult
+from flo_ai.common.flo_logger import get_logger, FloLogConfig
+from typing import Optional
+
+class FloLangchainLogger(BaseCallbackHandler):
+
+    def __init__(self, 
+                 session_id: str,
+                 logger_name: Optional[str] = "FloLangChainLogger", 
+                 log_level: Optional[str] = "WARN"):
+        self.logger = get_logger(FloLogConfig(logger_name, log_level))
+        self.session_id = session_id
+
+    def on_llm_start(self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any) -> None:
+        self.logger.info(f"Session ID: {self.session_id}: onLLMStart: {prompts}")
+
+    def on_llm_new_token(self, token: str, **kwargs: Any) -> None:
+        self.logger.debug(f"Session ID: {self.session_id}: onNewToken: {token}")
+
+    def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
+        self.logger.info(f"Session ID: {self.session_id}: onLLMEnd: {response.generations}")
+
+    def on_llm_error(self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any) -> None:
+        self.logger.error(f"Session ID: {self.session_id}: onLLMError: {error}")
+
+    def on_chain_start(self, serialized: Dict[str, Any], inputs: Dict[str, Any], **kwargs: Any) -> None:
+        self.logger.info(f"Session ID: {self.session_id}: onChainStart: {inputs}")
+
+    def on_chain_end(self, outputs: Dict[str, Any], **kwargs: Any) -> None:
+        self.logger.info(f"Session ID: {self.session_id}: onChainEnd: {outputs}")
+
+    def on_chain_error(self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any) -> None:
+        self.logger.error(f"Session ID: {self.session_id}: onChainError: {error}")
+
+    def on_tool_start(self, serialized: Dict[str, Any], input_str: str, **kwargs: Any) -> None:
+        self.logger.info(f"Session ID: {self.session_id}: onToolStart: {input_str}")
+
+    def on_tool_end(self, output: str, **kwargs: Any) -> None:
+        self.logger.info(f"Session ID: {self.session_id}: onToolEnd: {output}")
+
+    def on_tool_error(self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any) -> None:
+        self.logger.error(f"Session ID: {self.session_id}: onToolError: {error}")
+
+    def on_text(self, text: str, **kwargs: Any) -> None:
+        self.logger.info(f"Session ID: {self.session_id}: onText: {text}")
+
+    def on_agent_action(self, action: AgentAction, **kwargs: Any) -> Any:
+        self.logger.info(f"Session ID: {self.session_id}: onAgentAction: {action.tool} - {action.tool_input}")
+
+    def on_agent_finish(self, finish: AgentFinish, **kwargs: Any) -> None:
+        self.logger.info(f"Session ID: {self.session_id}: onAgentFinish: {finish.return_values}")

flo_ai/common/flo_logger.py

@@ -0,0 +1,54 @@
+import os
+import logging
+from logging.handlers import RotatingFileHandler
+from dataclasses import dataclass
+
+@dataclass
+class FloLogConfig:
+    name: str
+    level: str = "INFO"
+    file_path: str = None
+    max_bytes: int = 1048576
+
+class FloLoggerUtil(logging.Logger):
+    def __init__(self, config: FloLogConfig):
+        super().__init__(config.name, config.level)
+        self.setLevel(config.level)
+
+        formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
+
+        console_handler = logging.StreamHandler()
+        console_handler.setFormatter(formatter)
+        self.addHandler(console_handler)
+
+        if config.file_path:
+            file_handler = RotatingFileHandler(config.file_path, maxBytes=config.max_bytes)
+            file_handler.setFormatter(formatter)
+            self.addHandler(file_handler)
+
+class FloLogger:
+    _loggers = {}
+
+    @classmethod
+    def get_logger(cls, config: FloLogConfig) -> FloLoggerUtil:
+        if config.name not in cls._loggers:
+            level = config.level or os.environ.get(f"FLO_LOG_LEVEL_{config.name.upper()}", "INFO")
+            config.level = level
+            cls._loggers[config.name] = FloLoggerUtil(config)
+        return cls._loggers[config.name]
+
+    @classmethod
+    def set_log_level(cls, name: str, level: str):
+        if name in cls._loggers:
+            cls._loggers[name].setLevel(level)
+
+def get_logger(config: FloLogConfig) -> FloLoggerUtil:
+    return FloLogger.get_logger(config)
+
+common_logger = get_logger(FloLogConfig("COMMON"))
+builder_logger = get_logger(FloLogConfig("BUILDER"))
+session_logger = get_logger(FloLogConfig("SESSION"))
+
+def set_global_log_level(level: str):
+    for name in ["COMMON", "BUILDER", "SESSION"]:
+        FloLogger.set_log_level(name, level)

flo_ai/core.py

@@ -1,30 +1,40 @@
 from flo_ai.yaml.flo_team_builder import to_supervised_team
 from flo_ai.builders.yaml_builder import build_supervised_team, FloRoutedTeamConfig
-from typing import (
-    Any,
-    Iterator,
-    Union
-)
+from typing import Any, Iterator, Union
 from flo_ai.state.flo_session import FloSession
 from flo_ai.models.flo_executable import ExecutableFlo
+from flo_ai.common.flo_logger import common_logger, builder_logger, set_global_log_level
 
 class Flo:
 
     def __init__(self,
                  session: FloSession,
-                 config: FloRoutedTeamConfig) -> None:
+                 config: FloRoutedTeamConfig,
+                 log_level: str = "INFO") -> None:
         self.config = config
+        self.session = session
         self.runnable: ExecutableFlo = build_supervised_team(session, config)
 
+        set_global_log_level(log_level)
+        self.logger = common_logger
+        self.langchain_logger = session.langchain_logger
+        self.logger.info(f"Flo instance created for session {session.session_id}")
+
     def stream(self, query, config = None) -> Iterator[Union[dict[str, Any], Any]]:
+        self.logger.info(f"Streaming query for session {self.session.session_id}: {query}")
         return self.runnable.stream(query, config)
 
     def invoke(self, query, config = None) -> Iterator[Union[dict[str, Any], Any]]:
+        config = {
+         'callbacks' : [self.session.langchain_logger]
+        }
+        self.logger.info(f"Invoking query for session {self.session.session_id}: {query}")
         return self.runnable.invoke(query, config)
 
     @staticmethod
-    def build(session: FloSession, yaml: str):
-        return Flo(session, to_supervised_team(yaml))
+    def build(session: FloSession, yaml: str, log_level: str = "INFO"):
+        builder_logger.info("Building Flo instance from YAML")
+        return Flo(session, to_supervised_team(yaml), log_level)
 
     def draw(self, xray=True):
         return self.runnable.draw(xray)

flo_ai/state/flo_session.py

@@ -1,21 +1,43 @@
+import uuid
 from langchain_core.language_models import BaseLanguageModel
 from langchain_core.tools import BaseTool
+from flo_ai.common.flo_logger import session_logger, FloLogger, FloLogConfig
+from flo_ai.common.flo_langchain_logger import FloLangchainLogger
+from typing import Optional
 
 class FloSession:
-    def __init__(self, llm: BaseLanguageModel, loop_size: int = 2, max_loop: int = 3) -> None:
+
+    def __init__(self, 
+                 llm: BaseLanguageModel, 
+                 loop_size: int = 2, 
+                 max_loop: int = 3, 
+                 log_level: Optional[str] = "INFO",
+                 custom_langchainlog_handler: Optional[FloLangchainLogger] = None) -> None:
+        self.session_id = str(uuid.uuid4())
         self.llm = llm
         self.tools = dict()
         self.counter = dict()
         self.navigation: list[str] = list()
         self.pattern_series = dict()
         self.loop_size: int = loop_size
         self.max_loop: int = max_loop
+
+        self.init_logger(log_level)
+        self.logger = session_logger
+        self.logger.info(f"New FloSession created with ID: {self.session_id}")
+        self.langchain_logger = custom_langchainlog_handler or FloLangchainLogger(self.session_id, log_level=log_level, logger_name=f"FloLangChainLogger-{self.session_id}")
+        self.langchain_logger.set_session_id(self.session_id)
+
+    def init_logger(self, log_level: str):
+        FloLogger.set_log_level("SESSION", log_level)
 
     def register_tool(self, name: str, tool: BaseTool):
         self.tools[name] = tool
+        self.logger.info(f"Tool '{name}' registered for session {self.session_id}")
         return self
 
     def append(self, node: str) -> int:
+        self.logger.debug(f"Appending node: {node}")
         self.counter[node] = self.counter.get(node, 0) + 1
         if node in self.navigation:
             last_known_index = len(self.navigation) - 1 - self.navigation[::-1].index(node)
@@ -29,6 +51,7 @@ def append(self, node: str) -> int:
         self.navigation.append(node)
 
     def is_looping(self, node) -> bool:
+        self.logger.debug(f"Checking if node {node} is looping")
         patterns = self.pattern_series[node] if node in self.pattern_series else []
         if len(patterns) < self.max_loop:
             return False