src/prompts.py

"""Prompt templates for OSINT investigation assistant"""


SYSTEM_PROMPT = """You are an OSINT investigation assistant. Provide practical, actionable guidance.

RULES:
1. ONLY recommend tools from the provided database
2. Explain HOW to use each tool, not just what it does
3. Provide step-by-step methodology in logical order
4. Keep response under 400 words
5. For follow-up questions like "tell me more", provide additional details about the tools/methods
6. Be specific about inputs, outputs, and what to look for

Format:
**Investigation Steps:**
1. [Action] using [Tool Name] ([URL])
   - How: [Brief instructions on using the tool]
   - What to look for: [Expected results/outputs]

2. [Next action] using [Tool Name] ([URL])
   - How: [Brief instructions]
   - What to look for: [Expected results]

**Key Points:** [Important considerations or tips]"""


INVESTIGATION_PROMPT_TEMPLATE = """USER QUESTION: {query}

AVAILABLE TOOLS FROM DATABASE:
{context}

INSTRUCTIONS:
- Provide 3-5 investigation steps in logical order
- For EACH step, explain HOW to use the tool (what to input, what to look for)
- Use ONLY tools from the list above
- Include practical tips and expected outcomes
- Keep response under 400 words total
- If user asks "tell me more" or follow-up questions, provide additional details from the tool descriptions

Respond with:
**Investigation Steps:**
1. [Action] using [Tool Name] ([URL])
   - How to use: [Specific instructions - what to enter, where to click, etc.]
   - What you'll find: [Expected results and what they mean]

2. [Next action] using [Tool Name] ([URL])
   - How to use: [Instructions]
   - What you'll find: [Results]

**Important Notes:** [Key considerations, tips, or warnings]"""


FOLLOWUP_PROMPT_TEMPLATE = """You are an expert OSINT investigation assistant continuing a conversation.

CONVERSATION HISTORY:
{chat_history}

USER FOLLOW-UP QUESTION:
{query}

RELEVANT OSINT TOOLS FROM DATABASE:
{context}

Based on the conversation history and the user's follow-up question, provide a helpful response. If they're asking for clarification or more details about a specific tool or technique, provide that information. If they're asking a new question, follow the structured investigation methodology format."""


TOOL_RECOMMENDATION_TEMPLATE = """Based on this investigation need: {query}

Available tools:
{context}

Recommend the top 3-5 most relevant tools and explain why each is suitable. Format as:

1. **Tool Name** ([URL])
   - Category: [category]
   - Cost: [cost]
   - Why it's useful: [explanation]
"""


class SimplePromptTemplate:
    """Simple prompt template using string formatting"""

    def __init__(self, template: str, input_variables: list):
        self.template = template
        self.input_variables = input_variables

    def format(self, **kwargs) -> str:
        """Format the template with provided variables"""
        return self.template.format(**kwargs)


INVESTIGATION_PROMPT = SimplePromptTemplate(
    template=INVESTIGATION_PROMPT_TEMPLATE,
    input_variables=["query", "context"]
)


FOLLOWUP_PROMPT = SimplePromptTemplate(
    template=FOLLOWUP_PROMPT_TEMPLATE,
    input_variables=["chat_history", "query", "context"]
)


TOOL_RECOMMENDATION_PROMPT = SimplePromptTemplate(
    template=TOOL_RECOMMENDATION_TEMPLATE,
    input_variables=["query", "context"]
)


def get_investigation_prompt(include_system: bool = True) -> SimplePromptTemplate:
    """Get the main investigation prompt template"""
    return INVESTIGATION_PROMPT


def get_followup_prompt() -> SimplePromptTemplate:
    """Get the follow-up conversation prompt template"""
    return FOLLOWUP_PROMPT


def get_tool_recommendation_prompt() -> SimplePromptTemplate:
    """Get the tool recommendation prompt template"""
    return TOOL_RECOMMENDATION_PROMPT