search_personalization/ubi_tools.py

"""
Custom Strands tools for querying UBI (User Behavior Insights) data
from OpenSearch ubi_events and ubi_queries indices.
"""

import json
from typing import Optional
from strands import tool
from search_personalization.data_loader import get_opensearch_client


def _get_client():
    """Get a shared OpenSearch client."""
    return get_opensearch_client()


@tool
def query_ubi_events(
    action_name: Optional[str] = None,
    user_id: Optional[str] = None,
    query_id: Optional[str] = None,
    session_id: Optional[str] = None,
    time_range_minutes: int = 1440,
    size: int = 50,
    custom_query: Optional[str] = None,
) -> str:
    """Search UBI events from the ubi_events OpenSearch index. Events track user interactions
    like searches, product views, add-to-cart, hovers, checkouts, and filter usage.

    Args:
        action_name: Filter by event type (e.g. search, add_to_cart, view_product_details, hover, checkout, filter_applied, user_feedback)
        user_id: Filter by user ID
        query_id: Filter by query ID to see all events tied to a specific search
        session_id: Filter by session ID
        time_range_minutes: How far back to look in minutes (default 1440 = 24h)
        size: Max number of results to return (default 50)
        custom_query: Raw OpenSearch query JSON string to use instead of the built-in filters
    """
    client = _get_client()

    if custom_query:
        body = json.loads(custom_query)
    else:
        must_clauses = []
        if action_name:
            must_clauses.append({"match": {"action_name": action_name}})
        if user_id:
            must_clauses.append({"match": {"user_id": user_id}})
        if query_id:
            must_clauses.append({"match": {"query_id": query_id}})
        if session_id:
            must_clauses.append({"match": {"session_id": session_id}})
        must_clauses.append({
            "range": {
                "timestamp": {
                    "gte": f"now-{time_range_minutes}m",
                    "lte": "now"
                }
            }
        })

        body = {
            "query": {"bool": {"must": must_clauses}} if must_clauses else {"match_all": {}},
            "size": size,
            "sort": [{"timestamp": {"order": "desc"}}],
        }

    resp = client.search(index="ubi_events", body=body)
    hits = resp["hits"]["hits"]
    total = resp["hits"]["total"]["value"]

    results = []
    for h in hits:
        results.append(h["_source"])

    return json.dumps({"total_hits": total, "returned": len(results), "events": results}, indent=2, default=str)


@tool
def query_ubi_queries(
    user_query_text: Optional[str] = None,
    query_id: Optional[str] = None,
    user_id: Optional[str] = None,
    time_range_minutes: int = 1440,
    size: int = 50,
    custom_query: Optional[str] = None,
) -> str:
    """Search UBI queries from the ubi_queries OpenSearch index. Queries track what users searched for,
    including the search text, query_id, client_id, and the list of result hit IDs returned.

    Args:
        user_query_text: Filter by search text (fuzzy match)
        query_id: Filter by specific query ID
        user_id: Filter by user ID
        time_range_minutes: How far back to look in minutes (default 1440 = 24h)
        size: Max number of results to return (default 50)
        custom_query: Raw OpenSearch query JSON string to use instead of the built-in filters
    """
    client = _get_client()

    if custom_query:
        body = json.loads(custom_query)
    else:
        must_clauses = []
        if user_query_text:
            must_clauses.append({"match": {"user_query": user_query_text}})
        if query_id:
            must_clauses.append({"match": {"query_id": query_id}})
        if user_id:
            must_clauses.append({"match": {"user_id": user_id}})
        must_clauses.append({
            "range": {
                "timestamp": {
                    "gte": f"now-{time_range_minutes}m",
                    "lte": "now"
                }
            }
        })

        body = {
            "query": {"bool": {"must": must_clauses}} if must_clauses else {"match_all": {}},
            "size": size,
            "sort": [{"timestamp": {"order": "desc"}}],
        }

    resp = client.search(index="ubi_queries", body=body)
    hits = resp["hits"]["hits"]
    total = resp["hits"]["total"]["value"]

    results = []
    for h in hits:
        results.append(h["_source"])

    return json.dumps({"total_hits": total, "returned": len(results), "queries": results}, indent=2, default=str)


@tool
def aggregate_ubi_events(
    agg_field: str,
    time_range_minutes: int = 1440,
    agg_size: int = 20,
    action_name: Optional[str] = None,
) -> str:
    """Run aggregations on UBI events to get counts and distributions.
    Useful for understanding which actions are most common, which products get the most interaction, etc.

    Args:
        agg_field: The field to aggregate on (e.g. action_name, user_id, event_attributes.object.object_id, session_id)
        time_range_minutes: How far back to look in minutes (default 1440 = 24h)
        agg_size: Number of top buckets to return (default 20)
        action_name: Optional filter to only aggregate events of a specific action type
    """
    client = _get_client()

    filters = [{"range": {"timestamp": {"gte": f"now-{time_range_minutes}m", "lte": "now"}}}]
    if action_name:
        filters.append({"match": {"action_name": action_name}})

    body = {
        "size": 0,
        "query": {"bool": {"must": filters}},
        "aggs": {
            "field_distribution": {
                "terms": {"field": agg_field, "size": agg_size}
            }
        },
    }

    resp = client.search(index="ubi_events", body=body)
    buckets = resp["aggregations"]["field_distribution"]["buckets"]
    total_docs = resp["hits"]["total"]["value"]

    return json.dumps({
        "total_events_in_range": total_docs,
        "aggregation_field": agg_field,
        "buckets": buckets,
    }, indent=2, default=str)


@tool
def aggregate_ubi_queries(
    agg_field: str,
    time_range_minutes: int = 1440,
    agg_size: int = 20,
) -> str:
    """Run aggregations on UBI queries to understand search patterns.
    Useful for finding top search terms, most active users, query volume over time, etc.

    Args:
        agg_field: The field to aggregate on (e.g. user_query.keyword, user_id, client_id, application)
        time_range_minutes: How far back to look in minutes (default 1440 = 24h)
        agg_size: Number of top buckets to return (default 20)
    """
    client = _get_client()

    body = {
        "size": 0,
        "query": {
            "range": {
                "timestamp": {
                    "gte": f"now-{time_range_minutes}m",
                    "lte": "now"
                }
            }
        },
        "aggs": {
            "field_distribution": {
                "terms": {"field": agg_field, "size": agg_size}
            }
        },
    }

    resp = client.search(index="ubi_queries", body=body)
    buckets = resp["aggregations"]["field_distribution"]["buckets"]
    total_docs = resp["hits"]["total"]["value"]

    return json.dumps({
        "total_queries_in_range": total_docs,
        "aggregation_field": agg_field,
        "buckets": buckets,
    }, indent=2, default=str)


@tool
def get_ubi_index_mappings(index_name: str = "ubi_events") -> str:
    """Get the field mappings for a UBI index. Useful for discovering available fields
    before writing queries or aggregations.

    Args:
        index_name: The index to inspect — either ubi_events or ubi_queries
    """
    client = _get_client()
    mappings = client.indices.get_mapping(index=index_name)
    return json.dumps(mappings, indent=2, default=str)


@tool
def compute_user_preferences(
    user_id: str,
    time_range_minutes: int = 10080,
) -> str:
    """Analyze a user's filter history from UBI events and compute preference weights.
    Aggregates price and category filter events to determine the user's preferred price range,
    favorite categories, and price sensitivity. Results are saved to the user_preferences index.

    Args:
        user_id: The user ID to compute preferences for
        time_range_minutes: How far back to look in minutes (default 10080 = 7 days)
    """
    from search_personalization.user_preferences import save_user_preferences

    client = _get_client()

    # Query all filter events for this user
    body = {
        "size": 0,
        "query": {
            "bool": {
                "must": [
                    {"term": {"user_id": user_id}},
                    {"term": {"action_name": "filter"}},
                    {"range": {"timestamp": {"gte": f"now-{time_range_minutes}m"}}},
                ]
            }
        },
        "aggs": {
            "price_filters": {
                "filter": {"term": {"event_attributes.filter_type": "price"}},
                "aggs": {
                    "avg_max": {"avg": {"field": "event_attributes.max_price"}},
                    "avg_min": {"avg": {"field": "event_attributes.min_price"}},
                    "count": {"value_count": {"field": "event_attributes.filter_type"}},
                },
            },
            "category_filters": {
                "filter": {"term": {"event_attributes.filter_type": "category"}},
                "aggs": {
                    "top_categories": {
                        "terms": {"field": "event_attributes.filter_value.keyword", "size": 5}
                    }
                },
            },
            "total_filter_events": {"value_count": {"field": "action_name"}},
        },
    }

    resp = client.search(index="ubi_events", body=body)
    aggs = resp["aggregations"]

    total_events = aggs["total_filter_events"]["value"]
    price_agg = aggs["price_filters"]
    cat_agg = aggs["category_filters"]

    price_count = price_agg["count"]["value"]
    avg_max = price_agg["avg_max"]["value"]
    avg_min = price_agg["avg_min"]["value"]

    # Price sensitivity = proportion of filter events that are price-related
    price_sensitivity = price_count / total_events if total_events > 0 else 0.0

    top_categories = [b["key"] for b in cat_agg["top_categories"]["buckets"]]

    preferences = {
        "preferred_max_price": avg_max if avg_max else 2000.0,
        "preferred_min_price": avg_min if avg_min else 1.0,
        "preferred_categories": top_categories,
        "price_sensitivity": round(price_sensitivity, 3),
        "total_filter_events": int(total_events),
    }

    save_user_preferences(user_id, preferences)

    return json.dumps({"user_id": user_id, "preferences": preferences, "saved": True}, indent=2, default=str)