style(data source): format code

Author: SSSuperDan

alias/src/alias/agent/agents/data_source/_data_profiler_factory.py

@@ -1,4 +1,6 @@
 # -*- coding: utf-8 -*-
+# pylint: disable=R1702,R0912,R0915
+
 import os
 import json
 from abc import ABC, abstractmethod
@@ -167,9 +169,6 @@ def _call_model(
 
 
 class StructuredDataProfiler(BaseDataProfiler):
-    def __init__(self, api_key, path, source_type):
-        super().__init__(api_key, path, source_type)
-
     def _generate_content(self, prompt: str, data: Any) -> str:
         return prompt.format(data=data)
 
@@ -403,9 +402,6 @@ def _extract_irregular_table(
 
 
 class RelationalDatabaseProfiler(StructuredDataProfiler):
-    def __init__(self, api_key, path, source_type):
-        super().__init__(api_key, path, source_type)
-
     def _read_data(self):
         """
         Extracts metadata (schema) for all tables in a relational db.
@@ -432,7 +428,7 @@ def _read_data(self):
             connection = engine.connect()
         except Exception as e:
             print(f"Connection to {self.path} failed: {e}")
-            raise Exception(f"Failed to connect to database: {e}")
+            raise ConnectionError(f"Failed to connect to database: {e}") from e
 
         # Use DSN as the db identifier (can parsed cleaner)
         database_name = self.path

alias/src/alias/agent/agents/data_source/data_source.py

@@ -173,7 +173,10 @@ async def prepare(self, toolkit: AliasToolkit):
             )
 
     def get_coarse_desc(self):
-        return f"{self.source_desc}. {self.source_access_desc}: {self._general_profile()}"
+        return (
+            f"{self.source_desc}. {self.source_access_desc}: "
+            + f"{self._general_profile()}"
+        )
 
     def _get_profile(self, sandbox) -> Optional[Dict[str, Any]]:
         """Run type-specific profiling."""