"""LLM-friendly output formatting for database results.""" from dataclasses import dataclass from mcp_database_universal.config import DatabaseConfig from mcp_database_universal.engines.base import DBInfo, TableInfo, TableDetail, TableStats, QueryResult TYPE_TRANSLATIONS = { "INTEGER": "integer", "INT": "integer", "BIGINT": "large integer", "SMALLINT": "small integer", "TINYINT": "tiny integer", "REAL": "decimal number", "FLOAT": "decimal number", "DOUBLE": "decimal number", "DOUBLE PRECISION": "decimal number", "NUMERIC": "precise decimal", "DECIMAL": "precise decimal", "TEXT": "text", "VARCHAR": "text", "CHAR": "text", "BOOLEAN": "true/false", "BOOL": "true/false", "DATETIME": "date and time", "DATE": "date only", "TIMESTAMP": "date and time", "TIMESTAMP WITH TIME ZONE": "date and time (timezone)", "TIMESTAMP WITHOUT TIME ZONE": "date and time", "BLOB": "binary data", "BYTEA": "binary data", "JSON": "JSON data", "JSONB": "JSON data (optimized)", "UUID": "unique identifier", "ARRAY": "list of values", "SERIAL": "auto-increment integer", "BIGSERIAL": "auto-increment large integer", } class LLMFormatter: def __init__(self, config: DatabaseConfig): self.config = config def translate_type(self, raw_type: str) -> str: upper = raw_type.upper().strip() if "(" in upper: base = upper.split("(")[0].strip() size = upper.split("(")[1].rstrip(")") translated = TYPE_TRANSLATIONS.get(base, base.lower()) return f"{translated} (max {size} chars)" if "text" in translated else translated return TYPE_TRANSLATIONS.get(upper, raw_type.lower()) def format_number(self, n) -> str: if n is None: return "(empty)" if isinstance(n, float): return f"{n:,.2f}" if isinstance(n, int): return f"{n:,}" return str(n) def format_value(self, v) -> str: if v is None: return "(empty)" if isinstance(v, float): return f"{v:,.2f}" if isinstance(v, int): return f"{v:,}" return str(v) def truncate_for_llm(self, text: str) -> tuple[str, bool]: max_bytes = self.config.max_output_bytes if len(text.encode("utf-8")) <= max_bytes: return text, False lines = text.split("\n") truncated_lines = [] byte_count = 0 for line in lines: line_bytes = len(line.encode("utf-8")) + 1 if byte_count + line_bytes > max_bytes - 200: truncated_lines.append("... (output truncated)") break truncated_lines.append(line) byte_count += line_bytes return "\n".join(truncated_lines), True def format_db_info(self, info: DBInfo) -> str: lines = [ f"## Database Connection OK", f"- **Engine:** {info.engine}", f"- **Version:** {info.version}", f"- **Name:** {info.name}", f"- **Size:** {info.size_approx}", ] text = "\n".join(lines) text, _ = self.truncate_for_llm(text) return text def format_table_list(self, tables: list[TableInfo]) -> str: if not tables: return "## Tables\nNo tables found in this database." lines = [ f"## Tables ({len(tables)} total)", "", "| Table | Rows | Columns | FK In | FK Out |", "|-------|------|---------|-------|--------|", ] for t in tables: lines.append( f"| {t.name} | {self.format_number(t.row_count)} | " f"{t.column_count} | {t.foreign_keys_in} | {t.foreign_keys_out} |" ) text = "\n".join(lines) text, _ = self.truncate_for_llm(text) return text def format_table_detail(self, detail: TableDetail) -> str: lines = [ f"## Table: {detail.name}", "", "### Columns", "| # | Name | Type | Nullable | Key |", "|---|------|------|----------|-----|", ] for i, col in enumerate(detail.columns, 1): key = "" if col.is_primary_key: key = "PK" elif col.is_foreign_key: key = f"FK -> {col.foreign_key_table}.{col.foreign_key_column}" lines.append( f"| {i} | {col.name} | {self.translate_type(col.type)} | " f"{'yes' if col.nullable else 'no'} | {key} |" ) if detail.foreign_keys: lines.extend(["", "### Relationships"]) for fk in detail.foreign_keys: lines.append( f"- `{detail.name}.{fk.column}` -> " f"`{fk.references_table}.{fk.references_column}`" ) if detail.indexes: lines.extend(["", "### Indexes"]) for idx in detail.indexes: unique = " UNIQUE" if idx.unique else "" lines.append(f"- `{idx.name}` ON ({', '.join(idx.columns)}){unique}") if detail.sample_data: lines.extend(["", "### Sample Data (first 5 rows)"]) if detail.sample_data: headers = list(detail.sample_data[0].keys()) lines.append("| " + " | ".join(headers) + " |") lines.append("| " + " | ".join(["---"] * len(headers)) + " |") for row in detail.sample_data: vals = [self.format_value(row.get(h)) for h in headers] lines.append("| " + " | ".join(vals) + " |") if detail.stats: lines.extend(["", "### Statistics"]) lines.append(f"- **Row count:** {self.format_number(detail.stats.row_count)}") lines.append(f"- **Avg row size:** {detail.stats.avg_row_size}") lines.append(f"- **Total size:** {detail.stats.total_size}") if detail.stats.null_counts: non_zero = {k: v for k, v in detail.stats.null_counts.items() if v > 0} if non_zero: lines.append("- **NULL counts:**") for col, count in non_zero.items(): lines.append(f" - {col}: {self.format_number(count)}") text = "\n".join(lines) text, _ = self.truncate_for_llm(text) return text def format_query_result(self, result: QueryResult) -> str: lines = [ "## Query Results", f"**SQL:** `{result.sql}`", f"**Rows:** {result.row_count}{' (truncated)' if result.truncated else ''}", f"**Time:** {result.execution_time_ms}ms", ] if result.warning: lines.append(f"**Warning:** {result.warning}") if result.columns and result.rows: lines.append("") lines.append("| " + " | ".join(result.columns) + " |") lines.append("| " + " | ".join(["---"] * len(result.columns)) + " |") for row in result.rows: vals = [self.format_value(row.get(c)) for c in result.columns] lines.append("| " + " | ".join(vals) + " |") elif not result.rows: lines.append("\n*No rows returned.*") text = "\n".join(lines) text, _ = self.truncate_for_llm(text) return text def format_profile(self, profile_data: dict) -> str: lines = ["## Database Profile", ""] if "db_info" in profile_data: info = profile_data["db_info"] lines.extend([ f"**Engine:** {info.engine} {info.version}", f"**Size:** {info.size_approx}", f"**Tables:** {profile_data.get('table_count', 'unknown')}", ]) if "tables" in profile_data: lines.extend(["", "### Overview", "| Table | Rows | Columns | FK In | FK Out |", "|-------|------|---------|-------|--------|"]) for t in profile_data["tables"]: lines.append( f"| {t.name} | {self.format_number(t.row_count)} | " f"{t.column_count} | {t.foreign_keys_in} | {t.foreign_keys_out} |" ) if "relationships" in profile_data and profile_data["relationships"]: lines.extend(["", f"### Relationships ({len(profile_data['relationships'])})"]) for rel in profile_data["relationships"]: explicit = "explicit" if rel.is_explicit else "inferred" lines.append(f"- {rel.from_table}.{rel.from_column} -> {rel.to_table}.{rel.to_column} ({rel.cardinality}, {explicit})") if "junction_tables" in profile_data and profile_data["junction_tables"]: lines.extend(["", "### Junction Tables (many-to-many)"]) for jt in profile_data["junction_tables"]: lines.append(f"- `{jt}`") if "table_stats" in profile_data: lines.extend(["", "### Table Statistics"]) for table_name, stats in profile_data["table_stats"].items(): lines.append(f"\n#### {table_name}") lines.append(f"- Rows: {self.format_number(stats.row_count)}") lines.append(f"- Size: {stats.total_size}") if stats.null_counts: non_zero = {k: v for k, v in stats.null_counts.items() if v > 0} if non_zero: lines.append("- NULL columns:") for col, count in non_zero.items(): lines.append(f" - {col}: {self.format_number(count)}") text = "\n".join(lines) text, _ = self.truncate_for_llm(text) return text def format_schema_graph(self, mermaid: str) -> str: lines = [ "## Schema Graph", "", "```mermaid", mermaid, "```", ] return "\n".join(lines)