Config.yaml Support for Pebblo (daxa-ai#137)

* Updated Test Cases

* Fixed EntityClassifier init

* Updated outputdir path

---------

Co-authored-by: kunal.cd <kunal@clouddefense.io>

README.md

@@ -55,6 +55,14 @@ pebblo
 
 Pebblo daemon now listens to `localhost:8000` to accept Gen-AI application data snippets for inspection and reporting.
 
+#### Pebblo Optional Flags
+
+- `--config <file>`: Specifies a custom configuration file in yaml format.
+
+```bash
+pebblo --config config.yaml
+````
+
 ## Pebblo Safe DataLoader for Langchain
 
 `Pebblo Safe DataLoader` currently supports Langchain framework.

pebblo/app/config/config.py

@@ -0,0 +1,70 @@
+import yaml
+
+from pydantic import BaseSettings, Field
+import pathlib
+
+# Default config value
+dir_path = pathlib.Path().absolute()
+
+
+# Port BaseModel
+class PortConfig(BaseSettings):
+    host: str = Field(default='localhost')
+    port: int = Field(default=8000)
+
+
+# Report BaseModel
+class ReportConfig(BaseSettings):
+    format: str = Field(default='pdf')
+    outputDir: str = Field(dir_path)
+
+
+# Logging BaseModel
+class LoggingConfig(BaseSettings):
+    level: str = Field(default='info')
+
+
+# ConfigFile BaseModel
+class Config(BaseSettings):
+    daemon:  PortConfig
+    reports: ReportConfig
+    logging: LoggingConfig
+
+
+def load_config(path) -> Config:
+    try:
+        # If Path does not exist in command, set default config value
+        conf_obj = Config(
+            daemon=PortConfig(
+                host='localhost',
+                port=8000
+            ),
+            reports=ReportConfig(
+                format='pdf',
+                outputDir='~/.pebblo'
+            ),
+            logging=LoggingConfig(
+                level='info'
+            )
+        )
+        if not path:
+            # Setting Default config details
+            return conf_obj.dict()
+
+        # If Path exist, set config value
+        else:
+            con_file = path
+            try:
+                with open(con_file, "r") as output:
+                    cred_json = yaml.safe_load(output)
+                    parsed_config = Config.parse_obj(cred_json)
+                    config_dict = parsed_config.dict()
+                    return config_dict
+            except IOError as err:
+                print(f"no credentials file found at {con_file}")
+                return conf_obj.dict()
+
+    except Exception as err:
+        print(f'Error while loading config details, err: {err}')
+
+

pebblo/app/config/config.yaml

@@ -0,0 +1,8 @@
+daemon:
+  port: 8000
+  host: localhost
+logging:
+  level: info
+reports:
+  format: pdf
+  outputDir:  ~/.pebblo

pebblo/app/config/service.py

@@ -0,0 +1,30 @@
+from fastapi import FastAPI
+
+import uvicorn
+import asyncio
+from pebblo.app.routers.routers import router_instance
+
+
+class Service:
+    def __init__(self, config_details):
+        # Initialise app instance
+        self.app = FastAPI()
+        # Register the router instance with the main app
+        self.app.include_router(router_instance.router)
+        # Fetching Details from Config File
+        self.config_details = config_details
+        self.port = self.config_details.get('daemon', {}).get('port', 8000)
+        self.host = self.config_details.get('daemon', {}).get('host', '0.0.0.0')
+        self.log_level = self.config_details.get('logging', {}).get('level', 'info')
+
+    async def create_main_api_server(self):
+        # Add config Details to Uvicorn
+        config = uvicorn.Config(app=self.app, host=self.host, port=self.port, log_level=self.log_level)
+        server = uvicorn.Server(config)
+        await server.serve()
+
+    def start(self):
+        asyncio.run(self.create_main_api_server())
+
+
+

pebblo/app/daemon.py

@@ -1,22 +1,28 @@
-import uvicorn
-from fastapi import FastAPI
-from pebblo.app.routers.routers import router_instance
+from pebblo.app.config.config import load_config
+import argparse
 
-from pebblo.topic_classifier.topic_classifier import TopicClassifier
-from pebblo.entity_classifier.entity_classifier import EntityClassifier
+
+config_details = {}
 
 
 def start():
+    global config_details
+    # For loading config file details
+    parser = argparse.ArgumentParser(description="Pebblo  CLI")
+    parser.add_argument('--config', type=str, help="Config file path")
+    args = parser.parse_args()
+    path = args.config
+    config_details = load_config(path)
+
+    # LazyLoading TopicClassifier and EntityClassifier to avoid circular dependency
+    from pebblo.topic_classifier.topic_classifier import TopicClassifier
+    from pebblo.entity_classifier.entity_classifier import EntityClassifier
     # Init TopicClassifier(This step downloads the models and put in cache)
     _ = TopicClassifier()
     # Init EntityClassifier(This step downloads all necessary training models)
     _ = EntityClassifier()
 
-    # Initialise app instance
-    app = FastAPI()
-
-    # Register the router instance with the main app
-    app.include_router(router_instance.router)
-
-    # running local server
-    uvicorn.run(app, host="localhost", port=8000, log_level="info")
+    # Starting Uvicorn Service Using config details
+    from pebblo.app.config.service import Service
+    svc = Service(config_details)
+    svc.start()

pebblo/app/enums/enums.py

@@ -2,14 +2,16 @@
 These are all enums related to Inspector.
 """
 from enum import Enum
+import os
+from pebblo.app.daemon import config_details
 
 
 class CacheDir(Enum):
     metadata_folder = "/metadata"
     metadata_file_path = f"{metadata_folder}/metadata.json"
     report_file_name = "report.json"
     pdf_report_file_name = "pebblo_report.pdf"
-    home_dir = ".pebblo"
+    home_dir = config_details.get('reports', {}).get('outputDir', '~/.pebblo')
 
 
 class ReportConstants(Enum):

pebblo/app/libs/logger.py

@@ -2,16 +2,17 @@
 Module to handle logging functionality
 """
 import logging
-import os
+from pebblo.app.daemon import config_details
 
 
 def get_logger():
     """Get object of logger"""
+    logger_level = config_details.get('logging', {}).get('level', 'info').upper()
     logger_obj = logging.getLogger("Pebblo Logger")
     formatter = logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s")
     console_handle = logging.StreamHandler()
     console_handle.setFormatter(formatter)
-    logger_obj.setLevel(os.environ.get('LOG_LEVEL', 'INFO'))
+    logger_obj.setLevel(logger_level)
     logger_obj.propagate = False
     if not logger_obj.handlers:
         logger_obj.addHandler(console_handle)

tests/app/test_daemon.py

@@ -6,6 +6,7 @@
 from starlette.testclient import TestClient
 
 from pebblo.app.routers.routers import router_instance
+from pebblo.app import daemon
 
 app = FastAPI()
 app.include_router(router_instance.router)
@@ -15,20 +16,20 @@
 
 @pytest.fixture(scope="module")
 def mocked_objects():
-    with patch('pebblo.app.daemon.TopicClassifier') as topic_classifier, \
-            patch('pebblo.app.daemon.EntityClassifier') as entity_classifier:
+    with (patch.object(daemon, 'start', 'TopicClassifier') as topic_classifier,
+          patch.object(daemon, 'start', 'EntityClassifier') as entity_classifier):
         yield topic_classifier, entity_classifier
 
 
 @pytest.fixture(scope="module")
 def topic_classifier():
-    with patch('pebblo.app.daemon.TopicClassifier') as topic_classifier:
+    with patch.object(daemon, 'start', 'TopicClassifier') as topic_classifier:
         yield topic_classifier
 
 
 @pytest.fixture(scope="module")
 def entity_classifier():
-    with patch('pebblo.app.daemon.EntityClassifier') as entity_classifier:
+    with patch.object(daemon, 'start', 'EntityClassifier') as entity_classifier:
         yield entity_classifier