feat(data-model-discovery-agent): fix lint issues

Author: sanughosh-google

agent-app/app/agent.py

@@ -28,6 +28,7 @@
     mosaic_rag_agent_presales,
 )
 from .sub_agents.compliance_and_security_baseline_agent import compliance_agent
+from .sub_agents.data_model_discovery_agent import data_model_discovery_agent
 from .sub_agents.detailed_architecture_design_agent import (
     detailed_architecture_design_agent,
 )
@@ -38,8 +39,6 @@
 from .sub_agents.recommendation_agent import recommendation_agent
 from .sub_agents.strategy_recommender_agent import strategy_recommender_agent
 from .sub_agents.tech_stack_profiler_agent import tech_stack_profiler
-from .sub_agents.data_model_discovery_agent import data_model_discovery_agent
-
 from .tools import (
     transfer_to_capability_mapper_agent_tool,
     transfer_to_discovery_agent_tool,

agent-app/app/sub_agents/data_model_discovery_agent/agent.py

@@ -1,13 +1,15 @@
+import logging
+
 from google.adk.agents.llm_agent import LlmAgent
 from google.adk.agents.readonly_context import ReadonlyContext
+
+from app.config import MODEL
+
+from .sub_agents.data_profiling_agent.agent import data_profiling_agent
 from .sub_agents.database_cred_agent.agent import database_cred_agent
-from .sub_agents.schema_introspection_agent.agent import schema_introspection_agent
 from .sub_agents.qa_agent.agent import qa_agent
-from .sub_agents.data_profiling_agent.agent import data_profiling_agent
 from .sub_agents.reporting_agent.agent import reporting_agent
-
-import logging
-import json
+from .sub_agents.schema_introspection_agent.agent import schema_introspection_agent
 
 logger = logging.getLogger(__name__)
 logging.basicConfig(level=logging.INFO)
@@ -200,7 +202,7 @@ def root_agent_instruction(ctx: ReadonlyContext) -> str:
 
 
 data_model_discovery_agent = LlmAgent(
-    model="gemini-2.5-flash",
+    model=MODEL,
     name="data_model_discovery_agent",
     description=(
         "A helpful root agent that orchestrates sub-agents to introspect and profile legacy databases."

agent-app/app/sub_agents/data_model_discovery_agent/sub_agents/data_profiling_agent/agent.py

@@ -1,9 +1,11 @@
 from google.adk.agents.llm_agent import LlmAgent
+
+from app.config import MODEL
+
 from .tools import profile_schema_data
-from ..qa_agent.agent import qa_agent
 
 data_profiling_agent = LlmAgent(
-    model="gemini-2.5-flash",
+    model=MODEL,
     name="data_profiling_agent",
     description="Profiles data quality for the selected schema and then calls QA agent to summarize.",
     instruction="""

agent-app/app/sub_agents/data_model_discovery_agent/sub_agents/data_profiling_agent/tools.py

@@ -1,20 +1,22 @@
 import logging
-from typing import Dict, Any
-from google.adk.tools import ToolContext
-import psycopg2
+from typing import Any
+
 import mysql.connector
+import psycopg2
 import pyodbc
+from google.adk.tools import ToolContext
+
 from .utils import (
-    postgres_profiling_utils,
-    mysql_profiling_utils,
     mssql_profiling_utils,
+    mysql_profiling_utils,
+    postgres_profiling_utils,
 )
 
 logger = logging.getLogger(__name__)
 logging.basicConfig(level=logging.INFO)
 
 
-def _get_db_connection(metadata: Dict[str, Any], password: str) -> Any:
+def _get_db_connection(metadata: dict[str, Any], password: str) -> Any:
     db_type = metadata.get("db_type")
     host = metadata.get("host")
     port = int(metadata.get("port"))
@@ -39,8 +41,8 @@ def _get_db_connection(metadata: Dict[str, Any], password: str) -> Any:
 
 
 async def profile_schema_data(
-    tool_context: ToolContext, args: Dict[str, Any]
-) -> Dict[str, Any]:
+    tool_context: ToolContext, args: dict[str, Any]
+) -> dict[str, Any]:
     """
     Profiles the data in the selected schema based on the schema structure.
     Calculates nullability, cardinality, orphan records, and type anomalies.
@@ -104,7 +106,7 @@ async def profile_schema_data(
     except Exception as e:
         logger.error(f"Error during data profiling: {e}", exc_info=True)
         return {
-            "error": f"Failed to profile data for {db_type} ({schema_name}): {str(e)}"
+            "error": f"Failed to profile data for {db_type} ({schema_name}): {e!s}"
         }
     finally:
         if conn:

agent-app/app/sub_agents/data_model_discovery_agent/sub_agents/data_profiling_agent/utils/mssql_profiling_utils.py

@@ -1,20 +1,18 @@
 import logging
-from typing import Dict, Any, List
-import pyodbc
-from decimal import Decimal
+from typing import Any
 
 logger = logging.getLogger(__name__)
 
 
-def _execute_query(conn: Any, query: str) -> List[Dict[str, Any]]:
+def _execute_query(conn: Any, query: str) -> list[dict[str, Any]]:
     """Executes a SQL query and returns results as a list of dicts for SQL Server."""
     cursor = conn.cursor()
     try:
         cursor.execute(query)
         if cursor.description:
             columns = [column[0] for column in cursor.description]
             rows = cursor.fetchall()
-            return [dict(zip(columns, row)) for row in rows]
+            return [dict(zip(columns, row, strict=False)) for row in rows]
         return []
     finally:
         cursor.close()
@@ -23,9 +21,9 @@ def _execute_query(conn: Any, query: str) -> List[Dict[str, Any]]:
 def profile_mssql_data(
     conn: Any,
     schema_name: str,
-    schema_structure: Dict[str, Any],
+    schema_structure: dict[str, Any],
     sample_size: int = 10000,
-) -> Dict[str, Any]:
+) -> dict[str, Any]:
     profile_results = {
         "nullability": {},
         "cardinality": {},

agent-app/app/sub_agents/data_model_discovery_agent/sub_agents/data_profiling_agent/utils/mysql_profiling_utils.py

@@ -1,10 +1,10 @@
 import logging
-from typing import Dict, Any, List
+from typing import Any
 
 logger = logging.getLogger(__name__)
 
 
-def _execute_query(conn: Any, query: str) -> List[Dict[str, Any]]:
+def _execute_query(conn: Any, query: str) -> list[dict[str, Any]]:
     cursor = conn.cursor(dictionary=True)
     try:
         cursor.execute(query)
@@ -16,9 +16,9 @@ def _execute_query(conn: Any, query: str) -> List[Dict[str, Any]]:
 def profile_mysql_data(
     conn: Any,
     schema_name: str,
-    schema_structure: Dict[str, Any],
+    schema_structure: dict[str, Any],
     sample_size: int = 10000,
-) -> Dict[str, Any]:
+) -> dict[str, Any]:
     try:
         conn.database = schema_name
     except Exception as e:

agent-app/app/sub_agents/data_model_discovery_agent/sub_agents/data_profiling_agent/utils/postgres_profiling_utils.py

@@ -1,12 +1,10 @@
 import logging
-from typing import Dict, Any, List
-import psycopg2
-from decimal import Decimal
+from typing import Any
 
 logger = logging.getLogger(__name__)
 
 
-def _execute_query(conn: Any, query: str) -> List[Dict[str, Any]]:
+def _execute_query(conn: Any, query: str) -> list[dict[str, Any]]:
     """Executes a SQL query and returns results as a list of dicts for PostgreSQL."""
     cursor = conn.cursor()
     try:
@@ -15,7 +13,7 @@ def _execute_query(conn: Any, query: str) -> List[Dict[str, Any]]:
         if cursor.description:
             columns = [desc[0] for desc in cursor.description]
             rows = cursor.fetchall()
-            return [dict(zip(columns, row)) for row in rows]
+            return [dict(zip(columns, row, strict=False)) for row in rows]
         return []
     finally:
         cursor.close()
@@ -24,9 +22,9 @@ def _execute_query(conn: Any, query: str) -> List[Dict[str, Any]]:
 def profile_postgres_data(
     conn: Any,
     schema_name: str,
-    schema_structure: Dict[str, Any],
+    schema_structure: dict[str, Any],
     sample_size: int = 10000,
-) -> Dict[str, Any]:
+) -> dict[str, Any]:
     profile_results = {
         "nullability": {},
         "cardinality": {},

agent-app/app/sub_agents/data_model_discovery_agent/sub_agents/database_cred_agent/agent.py

@@ -1,9 +1,11 @@
 from google.adk.agents.llm_agent import LlmAgent
-from .tools import validate_db_connection
 
+from app.config import MODEL
+
+from .tools import validate_db_connection
 
 database_cred_agent = LlmAgent(
-    model="gemini-2.5-flash",
+    model=MODEL,
     name="database_cred_agent",
     description="A helpful assistant that collects and validates database connection details, and lists available schemas.",
     instruction="""

agent-app/app/sub_agents/data_model_discovery_agent/sub_agents/database_cred_agent/tools.py

@@ -1,16 +1,16 @@
-from google.adk.tools import ToolContext
 import logging
-from typing import Dict, Any, List
+from typing import Any
 
-import psycopg2
 import mysql.connector
+import psycopg2
 import pyodbc
+from google.adk.tools import ToolContext
 
 logger = logging.getLogger(__name__)
 logging.basicConfig(level=logging.INFO)
 
 
-def _get_schemas(conn: Any, db_type: str) -> List[str]:
+def _get_schemas(conn: Any, db_type: str) -> list[str]:
     """Fetches list of schemas/databases based on db type."""
     schemas = []
     cursor = conn.cursor()
@@ -52,8 +52,8 @@ def _get_schemas(conn: Any, db_type: str) -> List[str]:
 
 
 async def validate_db_connection(
-    connection_details: Dict[str, Any], tool_context: ToolContext
-) -> Dict[str, Any]:
+    connection_details: dict[str, Any], tool_context: ToolContext
+) -> dict[str, Any]:
     """Validates a database connection for PostgreSQL, MySQL, or MSSQL,
     fetches available schemas, and saves metadata to session memory.
 

agent-app/app/sub_agents/data_model_discovery_agent/sub_agents/qa_agent/agent.py

@@ -1,15 +1,17 @@
 import json
 from decimal import Decimal
+
 from google.adk.agents.llm_agent import LlmAgent
 from google.adk.agents.readonly_context import ReadonlyContext
 
+from app.config import MODEL
+
 
 def json_encoder_default(obj):
     if isinstance(obj, Decimal):
         return str(obj)
     raise TypeError(f"Object of type {obj.__class__.__name__} is not JSON serializable")
 
-
 def qa_agent_instruction(ctx: ReadonlyContext) -> str:
     """Builds the QA agent's instruction for schema and data profiling queries."""
 
@@ -30,9 +32,7 @@ def qa_agent_instruction(ctx: ReadonlyContext) -> str:
         """
 
     try:
-        schema_json = json.dumps(
-            schema_structure, indent=2, default=json_encoder_default
-        )
+        schema_json = json.dumps(schema_structure, indent=2, default=json_encoder_default)
     except Exception as e:
         schema_json = f"Error serializing schema structure: {e}"
 
@@ -45,15 +45,11 @@ def qa_agent_instruction(ctx: ReadonlyContext) -> str:
                 "Nullability": data_profile.get("nullability", "Not available"),
                 "Cardinality": data_profile.get("cardinality", "Not available"),
                 "Orphan Records": data_profile.get("orphan_records", "Not available"),
-                "Type Anomalies": data_profile.get("type_anomalies", "Not available"),
+                "Type Anomalies": data_profile.get("type_anomalies", "Not available")
             }
-            profile_message = json.dumps(
-                profile_summary, indent=2, default=json_encoder_default
-            )
+            profile_message = json.dumps(profile_summary, indent=2, default=json_encoder_default)
         except Exception:
-            profile_message = (
-                "Data profiling results exist but could not be summarized."
-            )
+            profile_message = "Data profiling results exist but could not be summarized."
     else:
         profile_message = (
             "Data profiling has not been run yet. "
@@ -97,11 +93,10 @@ def qa_agent_instruction(ctx: ReadonlyContext) -> str:
     Always respond in clear, human-readable sentences. If profiling data is missing, offer to run profiling on a sample of up to 10,000 rows to provide a summary.
     """
 
-
 qa_agent = LlmAgent(
-    model="gemini-2.5-flash",
-    name="qa_agent",
-    description="Answers natural language questions about the discovered database schema structure and data profiling results.",
+    model=MODEL,
+    name='qa_agent',
+    description='Answers natural language questions about the discovered database schema structure and data profiling results.',
     instruction=qa_agent_instruction,
-    tools=[],
+    tools=[]
 )