remove some duplication

Author: bcallender

src/fenic/api/mcp/tool_generation.py

@@ -413,7 +413,7 @@ def _auto_generate_sql_tool(
         raise ConfigurationError("Cannot create SQL tool: no datasets provided.")
 
     def analyze_func(
-        full_sql: Annotated[str, "Full SELECT SQL. Refer to DataFrames by name in braces, e.g., {orders}."]
+        full_sql: Annotated[str, "Full SELECT SQL. Refer to DataFrames by name in braces, e.g., `SELECT * FROM {orders}`. JOINs between the provided datasets are allowed. SQL dialect: DuckDB. DDL/DML, CTEs, subqueries, UNION, and multiple top-level queries are not allowed"]
     ) -> LogicalPlan:
         return session.sql(full_sql.strip(), **{spec.table_name: spec.df for spec in datasets})._logical_plan
 
@@ -789,6 +789,7 @@ def _auto_generate_core_tools(
         tool_name=f"{tool_group_name} - Read",
         tool_description="\n\n".join([
             "Read rows from a single dataset. Use to sample data, or to execute simple queries over the data that do not require filtering or grouping.",
+            "Use `include_columns` and `exclude_columns` to filter columns by name -- this is important to conserve token usage. Use the `Profile` tool to understand the columns and their sizes.",
             "Available datasets:",
             group_desc,
         ]),

src/fenic/core/mcp/_server.py

@@ -15,6 +15,7 @@
 from functools import wraps
 from typing import Any, Callable, Dict, List, Optional, Union
 
+import polars as pl
 from pydantic import BaseModel, ConfigDict
 from typing_extensions import Annotated, Literal
 
@@ -48,7 +49,8 @@ class MCPResultSet(BaseModel):
 
     table_schema: Optional[List[Dict[str, Any]]]
     rows: Union[List[Dict[str, Any]], str]
-    row_count: int
+    returned_result_count: int
+    total_result_count: int
 
 MCPTransport = Literal["http", "stdio"]
 
@@ -135,6 +137,24 @@ def http_app(self, **kwargs):
         """Create a Starlette ASGI app for the MCP server."""
         return self.mcp.http_app(**kwargs)
 
+    def _handle_result_set(self, pl_df: pl.DataFrame, effective_limit: Optional[int], table_format: TableFormat) -> MCPResultSet:
+        """Handle the result set from a logical plan."""
+        original_result_count = len(pl_df)
+        if effective_limit and original_result_count > effective_limit:
+            pl_df = pl_df.limit(effective_limit)
+        rows_list = pl_df.to_dicts()
+        schema_fields = [{"name": name, "type": str(dtype)} for name, dtype in pl_df.schema.items()]
+        result_set = MCPResultSet(
+            table_schema=schema_fields,
+            rows=rows_list,
+            returned_result_count=len(rows_list),
+            total_result_count=original_result_count,
+        )
+        if table_format == "markdown":
+            result_set.rows = _render_markdown_preview(rows_list)
+            result_set.table_schema = None
+        return result_set
+
     def _build_parameterized_tool(self, tool: ParameterizedToolDefinition):
         """Build a keyword-argument tool function with per-field schema for FastMCP.
 
@@ -162,24 +182,13 @@ async def tool_fn_wrapper(*args, **kwargs) -> MCPResultSet:
                 bound_plan = bind_parameters(tool._parameterized_view, payload, tool.params)
                 async with self._collect_semaphore:
                     pl_df, metrics = await asyncio.to_thread(
-                        lambda: self.session_state.execution.collect(bound_plan, n=effective_limit)
+                        lambda: self.session_state.execution.collect(bound_plan)
                     )
                     logger.info(f"Completed query for {tool.name}")
                     logger.info(metrics.get_summary())
                     logger.debug(f"Query Details: {params_obj.model_dump_json()}")
 
-                rows_list = pl_df.to_dicts()
-                schema_fields = [{"name": name, "type": str(dtype)} for name, dtype in pl_df.schema.items()]
-                result_set = MCPResultSet(
-                    table_schema=schema_fields,
-                    rows=rows_list,
-                    row_count=len(rows_list),
-                )
-                if table_format == "markdown":
-                    result_set.rows = _render_markdown_preview(rows_list)
-                    result_set.table_schema = None
-
-                return result_set
+                return self._handle_result_set(pl_df, effective_limit, table_format)
             except Exception as e:
                 from fastmcp.exceptions import ToolError
                 raise ToolError(f"Fenic server failed to execute tool {tool.name}. Underlying error: {e}") from e
@@ -263,19 +272,13 @@ async def wrapper(*args, **kwargs) -> MCPResultSet:
                 # collections with a semaphore to protect the backend executor.
                 async with self._collect_semaphore:
                     pl_df, metrics = await asyncio.to_thread(
-                        lambda: self.session_state.execution.collect(bound_plan, n=effective_limit)
+                        lambda: self.session_state.execution.collect(bound_plan)
                     )
                     logger.info(f"Completed query for {tool.name}")
                     logger.info(metrics.get_summary())
                     logger.debug(f"Query Details: {args if args else kwargs}")
-                rows_list = pl_df.to_dicts()
-                schema_fields = [{"name": name, "type": str(dtype)} for name, dtype in pl_df.schema.items()]
-                out = MCPResultSet(table_schema=schema_fields, rows=rows_list, row_count=len(rows_list))
-                if table_format == "markdown":
-                    out.rows = _render_markdown_preview(rows_list)
-                    out.table_schema = None
-
-                return out
+
+                return self._handle_result_set(pl_df, effective_limit, table_format)
             except Exception as e:
                 from fastmcp.exceptions import ToolError
                 raise ToolError(f"Fenic server failed to execute tool {tool.name}. Underlying error: {e}") from e