DataPilot MCP Server

Instructions

Implementation Reference

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

  MCP tool handler for natural_language_to_sql. Gathers relevant table schemas from Snowflake and delegates to OpenAI client for SQL generation.

  @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/models.py:56-62 (schema)

  Pydantic BaseModel defining the input parameters for the natural_language_to_sql tool.

  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")

• src/openai_client.py:17-74 (helper)

  OpenAIClient method implementing the core NL-to-SQL logic using chat completions with schema-aware prompting.

  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)}")