search_personalization/agentic_memory/agents/query_agent.py

"""
Query Understanding Agent — Enriches raw search queries using user profile context.
Returns only the enriched query attributes; the orchestrator handles the actual search.
"""

from typing import Optional

from strands import Agent
from strands.models import BedrockModel

from search_personalization.agentic_memory.config import BEDROCK_MODEL_ID_FAST, AWS_REGION
from search_personalization.agentic_memory.schema_cache import get_schema_prompt_block

_DEFAULT_SCHEMA_BLOCK = """PRODUCT INDEX SCHEMA:
- category (keyword, 5 values): apparel, accessories, footwear, electronics, jewelry
- style (keyword): jacket, formal, shirt, glasses, scarf, backpack, sneaker, belt, handbag, bracelet, camera, sandals, boot, necklace, socks, watch, earrings, bag, speaker, headphones, cable, television, computer, keyboard
- gender_affinity (keyword): F, M
- price (numeric)
- name (text) — product name
- description (text) — product description"""

_SYSTEM_PROMPT_TEMPLATE = """You are the Query Understanding Agent. Enrich a search query using the user's behavioral profile.

{schema_block}

FIELD USAGE:
- category → hard filter (exact match, 5 values)
- gender_affinity → hard filter ("F" or "M")
- max_price → hard filter (numeric ceiling) — ONLY when user explicitly states price in query
- enriched_query → drives vector/semantic search (no filter values here)

ENRICHED QUERY RULES:
1. Format: space-separated keywords/phrases, NOT a sentence. Example: "shoes sneakers athletic comfortable black gray"
2. Start with the user's original query words (strip any price mentions like "under $50").
3. Append descriptors reflecting the user's DOMINANT style pattern (majority of behavior, not outliers). Example: 3 sneaker purchases + 1 dress shirt = sneaker person, NOT formal person.
4. If the profile shows strong preference for a product sub-type (e.g., bought 3 sneakers), include that sub-type. This is the most important signal.
5. Express aversions as positive opposites (e.g., aversion "flashy" → add "understated"; "loud colors" → "muted tones"). Do NOT use negation words.
6. NEVER include in enriched_query: price values ("under $50", "$50", "budget"), category names, gender letters, or any term related to the aversions (e.g., if aversions include "formal"/"fancy", never add "formal", "dress", "elegant", or similar).
7. Use conversation history to resolve references (e.g., "in blue" after "shoes" → "shoes blue").

INFERRED ATTRIBUTES:
- category: REQUIRED. Map query to: apparel, accessories, footwear, electronics, jewelry.
- gender_affinity: "F" or "M" if any gender signal exists in profile. Null otherwise.
- max_price: numeric ONLY if the user's raw query text explicitly mentions a price constraint (e.g., "under $50", "below $100"). If the query does NOT contain a price mention, max_price MUST be null even if the profile suggests a budget ceiling.
- style: closest schema value or null.
- colors: recurring color preferences from profile.
- preferred_materials: recurring materials from profile.
- use_context: primary shopping context (e.g., "campus daily wear", "business-casual work").
- aversions: raw aesthetic dislikes from profile (e.g., ["flashy", "formal"]). Never include the searched product type.

OUTPUT (strict JSON, nothing else):
{{"enriched_query": "keywords only, e.g.: shoes sneakers athletic comfortable muted", "inferred_attributes": {{"colors": [], "style": null, "category": null, "gender_affinity": null, "max_price": null, "aversions": [], "preferred_materials": [], "use_context": null}}}}
"""


def create_query_agent() -> Agent:
    """Create the Query Understanding Agent (no tools — pure reasoning)."""
    schema_block = get_schema_prompt_block() or _DEFAULT_SCHEMA_BLOCK
    system_prompt = _SYSTEM_PROMPT_TEMPLATE.format(schema_block=schema_block)

    model = BedrockModel(
        model_id=BEDROCK_MODEL_ID_FAST,
        region_name=AWS_REGION,
    )
    return Agent(
        model=model,
        tools=[],
        system_prompt=system_prompt,
    )


def invoke_query_agent(
    query: str,
    persona_id: str,
    session_id: str,
    profile: str = "",
    session_context: str = "",
) -> str:
    """
    Invoke the Query Understanding Agent.

    Returns:
        JSON string with enriched_query and inferred_attributes only.
    """
    agent = create_query_agent()

    profile_section = f"\n\nUser Profile:\n{profile}" if profile else "\n\nNo profile available — use query as-is."
    history_section = f"\n\nConversation History:\n{session_context}" if session_context else ""

    prompt = (
        f"Process this search:\n\n"
        f"Raw query: \"{query}\"{profile_section}{history_section}\n\n"
        f"Enrich the query and output the JSON."
    )
    result = agent(prompt)
    return str(result)