prompts.py

# prompts.py

def get_classification_prompt(user_input: str) -> str:
    return f"""
    You are an intent classifier.
    
    Classify the user input into EXACTLY ONE of the following categories.
    
    IMPORTANT RULES:
    - Choose "query_similarity_search" ONLY if the user clearly names
      a topic, entity, document, or subject to search for.
    - If the input is vague, contextual, conversational, or does NOT
      specify what to search for, it is NOT a similarity search.
    - Questions like "what is this?", "what can I do here?",
      "explain this", or unclear references MUST be "query_non_db_chat".
    
    Categories:
    
    1. query_similarity_search
       - User explicitly asks to find information ABOUT a named topic,
         document, concept, or entity.
    
    2. query_sql_generation
       - User asks for aggregation, calculations, filtering,
         summaries, or analysis over stored data, or just general questions over the content/data etc included
    
    3. add_table
       - User explicitly asks to create or add a database table.
    
    4. query_non_db_chat
       - Vague, conversational, unclear, or non-database-related input.
       - Includes questions without a clear search target.
    
    User Input:
    {user_input}
    
    Return ONLY the category name.
    """


def get_table_filter_prompt(user_input: str, all_tables: str) -> str:
    return f"""
    User Query: {user_input}
    Available Tables: {all_tables}
    
    Select the tables that are likely to contain information RELEVANT to the user's query.
    If the query is generic, select the most important core tables (like 'nodes' or 'edges').
    
    Return a JSON list of strings, e.g. ["table1", "table2"].
    """

def get_sql_generation_prompt(user_input: str, formatted_table_names: str, context_data: str) -> str:
    return f"""
    You are a BigQuery SQL expert.
    
    User Question: {user_input}
    
    Relevant Tables (Fully Qualified):
    {formatted_table_names}
    
    Context (Schemas & Metadata):
    {context_data}
    
    Generate a valid BigQuery SQL query to answer the question.
    Use the fully qualified table names provided.
    
    CRITICAL RULES:
    1. First, think step-by-step (Chain-of-Thought) about which tables are needed and how they join.
    2. Then, write the SQL.
    3. Use Standard SQL syntax for BigQuery.
    4. Use `LIMIT n` instead of `TOP n`.
    5. Return ONLY the raw SQL string. Do NOT use markdown code blocks (```sql ... ```).
    6. Ensure column names exist in the provided schema.
    7. Pay attention to 'mode': 'REPEATED' in the schema. These are ARRAYs and require UNNEST() to query effectively if filtering by value.
    8. Use the provided table metadata (row counts, etc.) to optimize your query.

    COMMON COLUMN MAPPINGS (Use these if applicable):
    - "Filename", "Source File", "File" -> `file_id`
    - "Date", "Timestamp", "Created" -> `ingested_at`
    - "Text", "Body", "Document" -> `content`
    """

def get_query_expansion_prompt(user_input: str) -> str:
    return f"""
    You are a Search Query Optimizer.
    Your goal is to improve the retrieval of relevant documents by expanding the user's query.

    User Input: "{user_input}"

    Generate 3 distinct search variations using these strategies:
    1. **Decomposition**: Break complex questions into simpler keyword phrases.
    2. **Synonyms**: Use professional or technical synonyms for key terms.
    3. **Hypothetical Answer**: What key phrases would appear in a document that answers this?

    Return ONLY a JSON list of strings. Example: ["variation 1", "variation 2", "variation 3"]
    """

def get_natural_answer_prompt(user_input: str, sql_query: str, query_result: str) -> str:
    return f"""
    You are a helpful and knowledgeable data assistant. 
    You have just executed a SQL query to answer the user's question.

    User's Question: "{user_input}"
    
    Data Retrieved (Result of SQL Query):
    {query_result}

    Instructions:
    1. Synthesize the data into a natural, friendly response.
    2. Do not mention "SQL", "rows", or "query results" explicitly unless necessary for clarity. 
    3. Speak as if you analyzed the data yourself. 
    4. If the data corresponds to a specific file or item, mention it naturally.
    5. Be concise but complete.
    """

def get_upload_instructions_text() -> str:
    return """
    To add a table or ingest data, please use the **Ingest** command in the CLI.
    
    Example:
    `python cli.py ingest --chunk-size 1000 --use-docai`
    
    Or ensure your files are in `data_dir` and ask me to "ingest data" if configured.
    """

def get_platform_help_prompt(user_input: str) -> str:
    return f"""
    You are the **BigQuery AI Toolbox** Platform Assistant.
    Your specific role is to help the user understand how to use this platform, explain its features, and offer best practices.

    **Platform Overview:**
    - **Purpose**: Ingest unstructured data (PDFs, Images, CSVs) into BigQuery, auto-extract content, generate embeddings, and enable RAG + SQL Analytics.
    - **Core Features**:
      - **Ingestion**: Supports PDF/Image (via DocAI) and CSV. Chunks content and stores in `KB` table.
      - **Search**: "Find X" performs vector similarity search.
      - **Analytics**: "Count Y" or "How many..." generates SQL queries.
      - **Security**: Data is stored in your personal BigQuery dataset.

    **Instructions:**
    1. Answer the user's question **only** if it relates to the platform, its usage, or best practices.
    2. **DO NOT** attempt to answer questions about specific documents, files, or data in the Knowledge Base (you do not have access to them in this mode).
    3. If the user input is nonsense, gibberish (e.g. "asdfgh"), or completely irrelevant, respond with a friendly follow-up question like: "I'm not sure I understood that correctly. Did you want to search your knowledge base, analyze data, or learn how to ingest new files? I'm here to help!"
    4. If the user asks about their data (e.g. "What is in file X?"), politely guide them to use a search command (e.g. "You can ask 'Find info about X'").
    5. If the user asks general world knowledge questions (e.g. "What is an iPod?"), politely redirect them to how they could *ingest* information about that topic into the platform, or answer very briefly and pivot back to the platform.
    6. Be helpful, professional, and concise.

    User Question: "{user_input}"
    """

def get_query_rewrite_prompt(user_input: str, history_text: str) -> str:
    return f"""
    You are a Query Transformation AI.
    Your job is to rewrite the User's latest input into a standalone, fully contextualized query based on the Conversation History.
    
    Conversation History:
    {history_text}
    
    User Input: {user_input}
    
    Instructions:
    1. If the User Input is a follow-up (e.g., "what about for X?", "and the price?"), rewrite it to include the missing context from history.
    2. If the User Input is standalone and clear, return it exactly as is.
    3. Do NOT answer the question. Only REWRITE it.
    4. Output ONLY the rewritten string.
    """