DataPilot MCP Server

Instructions

Implementation Reference

• src/main.py:275-308 (handler)

  MCP tool handler and registration for 'natural_language_to_sql'. Fetches relevant table schemas from Snowflake, passes to OpenAI client for SQL generation, and returns the generated SQL query.
  @mcp.tool() async def natural_language_to_sql(request: NaturalLanguageRequest, ctx: Context) -> str: """Convert natural language question to SQL query using AI""" await ctx.info(f"Converting natural language to SQL: {request.question}") try: openai = await get_openai_client() # Get schema context for better SQL generation table_schemas = [] if request.database or request.schema: snowflake = await get_snowflake_client() tables = await snowflake.list_tables(request.database, request.schema) for table in tables[:5]: # Limit to first 5 tables for context columns = await snowflake.describe_table( table.table_name, request.database or table.database_name, request.schema or table.schema_name ) table_schemas.append({ 'table_name': table.table_name, 'columns': [{'name': col.column_name, 'type': col.data_type} for col in columns] }) sql_query = await openai.natural_language_to_sql(request, table_schemas) await ctx.info("Successfully generated SQL query") return sql_query except Exception as e: logger.error(f"Error converting to SQL: {str(e)}") await ctx.error(f"Failed to convert to SQL: {str(e)}") raise
• src/openai_client.py:17-74 (helper)

  Core helper function in OpenAIClient that generates SQL query from natural language using OpenAI chat completions API. Builds schema context and crafts specialized prompt for Snowflake SQL.
  async def natural_language_to_sql( self, request: NaturalLanguageRequest, table_schemas: Optional[List[Dict[str, Any]]] = None ) -> str: """Convert natural language question to SQL query""" # Build context about available tables and schemas context = self._build_schema_context(table_schemas) system_prompt = f""" You are a Snowflake SQL expert. Convert natural language questions to SQL queries. Guidelines: - Use proper Snowflake SQL syntax - Include appropriate WHERE clauses for filters - Use proper JOIN syntax when needed - Always include LIMIT clause for safety (default 100) - Use uppercase for SQL keywords - Be precise with column names and table names - Handle date/time functions appropriately for Snowflake Available schema context: {context} Return only the SQL query without any explanation or markdown formatting. """ user_prompt = f""" Question: {request.question} Additional context: {request.context or 'None'} Database: {request.database or 'Use current database'} Schema: {request.schema or 'Use current schema'} """ try: response = await self.client.chat.completions.create( model=self.model, messages=[ {"role": "system", "content": system_prompt}, {"role": "user", "content": user_prompt} ], temperature=0.1, max_tokens=500 ) sql_query = response.choices[0].message.content.strip() # Remove any markdown formatting sql_query = sql_query.replace('```sql', '').replace('```', '').strip() logger.info(f"Generated SQL query: {sql_query}") return sql_query except Exception as e: logger.error(f"Error generating SQL: {str(e)}") raise Exception(f"Failed to generate SQL query: {str(e)}")
• src/models.py:56-62 (schema)

  Pydantic model defining the input schema for the natural_language_to_sql tool, including question, optional context, database, and schema fields.
  class NaturalLanguageRequest(BaseModel): """Request for natural language to SQL conversion""" question: str = Field(..., description="Natural language question about the data") context: Optional[str] = Field(None, description="Additional context about tables/schema") database: Optional[str] = Field(None, description="Database to query") schema: Optional[str] = Field(None, description="Schema to query")