Upload mcp_database_universal/formatters/llm.py with huggingface_hub

mcp_database_universal/formatters/llm.py

@@ -0,0 +1,262 @@
+"""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)