search_personalization/agentic_memory/demo_runner.py

"""
Demo Runner — Orchestrates the 3-agent pipeline for the agentic memory demo.

Architecture:
  - Long-term memory: user preferences fetched from OpenSearch agentic memory
  - Query Understanding Agent: enriches query, searches catalog, returns results directly
  - Ranking Agent: receives results directly, reranks with personalization
  - Working memory write: async fire-and-forget for observability/tracing

Usage:
    python -m src.agentic_memory.demo_runner
    python -m src.agentic_memory.demo_runner --query "black boots for a concert"
    python -m src.agentic_memory.demo_runner --interactive
"""

import json
import time
import uuid
import threading
from datetime import datetime
from typing import Optional

from search_personalization.agentic_memory.config import PERSONAS
from search_personalization.agentic_memory.memory_tools import get_user_profile
from search_personalization.agentic_memory.agents.query_agent import invoke_query_agent
from search_personalization.agentic_memory.memory_tools import write_session_memory, write_history, read_session_memory


DEFAULT_QUERY = "backpack"


def run_pipeline_for_persona(
    query: str,
    persona_name: str,
    session_id: Optional[str] = None,
    conversation_history: Optional[str] = None,
) -> dict:
    """
    Run the 3-step pipeline for a single persona:
      1. Profile Cache (instant)
      2. Query Understanding Agent (enrich + search)
      3. Ranking Agent (rerank with direct results)
    """
    persona_id = PERSONAS[persona_name]
    if not session_id:
        session_id = f"sess-{uuid.uuid4().hex[:8]}"

    print(f"\n{'='*60}")
    print(f"  PERSONA: {persona_name.upper()} ({persona_id})")
    print(f"  Query: \"{query}\"")
    print(f"  Session: {session_id}")
    print(f"{'='*60}")

    pipeline_start = time.time()
    pipeline_trace = {
        "persona_id": persona_id,
        "persona_name": persona_name,
        "query": query,
        "session_id": session_id,
        "timestamp": datetime.utcnow().isoformat(),
        "agents": {},
    }

    # --- Step 1: Profile + session in parallel ---
    from concurrent.futures import ThreadPoolExecutor

    start = time.time()
    with ThreadPoolExecutor(max_workers=2) as pool:
        profile_future = pool.submit(get_user_profile, persona_id)
        if conversation_history:
            session_context = conversation_history
            profile_raw = profile_future.result()
        else:
            session_future = pool.submit(read_session_memory, persona_id, session_id)
            profile_raw = profile_future.result()
            session_context = session_future.result()

    # Format USER_PREFERENCE memories into a profile string for the agent
    profile_data = json.loads(profile_raw)
    memories = profile_data.get("memories", [])
    preference_statements = [
        m.get("_source", {}).get("memory", "")
        for m in memories
        if m.get("_source", {}).get("strategy_type") == "USER_PREFERENCE"
    ]
    profile_lines = [f"- {s}" for s in preference_statements if s]
    profile = (
        f"Profile based on {len(preference_statements)} behavioral observations:\n"
        + "\n".join(profile_lines)
    ) if profile_lines else ""

    memory_duration = time.time() - start
    print(f"\n  [memory] Profile + session loaded ({memory_duration*1000:.1f}ms)")
    pipeline_trace["agents"]["memory"] = {
        "duration_ms": round(memory_duration * 1000, 1),
        "source": "opensearch-long-term-memory",
        "preference_records": len(preference_statements),
    }

    # --- Step 2: Query Understanding Agent (enrich only) ---
    print(f"  [agent] Query Understanding Agent — enrich...")
    start = time.time()
    query_result = invoke_query_agent(query, persona_id, session_id, profile=profile or "", session_context=session_context)
    query_duration = time.time() - start
    pipeline_trace["agents"]["query"] = {
        "duration_ms": int(query_duration * 1000),
        "output": query_result,
    }
    print(f"          Done ({query_duration:.1f}s)")

    # --- Step 2b: Embed enriched query + knn search ---
    print(f"  [search] Embed + catalog search...")
    start = time.time()
    try:
        enrichment = json.loads(query_result)
    except (json.JSONDecodeError, TypeError):
        # Agent output may contain surrounding text — extract JSON object
        enrichment = None
        try:
            _start = query_result.find("{")
            _end = query_result.rfind("}") + 1
            if _start >= 0 and _end > _start:
                enrichment = json.loads(query_result[_start:_end])
        except (json.JSONDecodeError, TypeError):
            pass
        if not enrichment:
            enrichment = {"enriched_query": query, "inferred_attributes": {}}

    attrs = enrichment.get("inferred_attributes", {})
    enriched_query = enrichment.get("enriched_query", query)

    # Let the agent handle all semantic resolution — no deterministic style matching.
    # Style is a semantic field: "shoes" should surface boots, sneakers, sandals, etc.
    # Category and gender_affinity are hard kNN pre-filters.
    # max_price is only applied as a kNN filter when the user explicitly mentions
    # price in their original query. The agent may infer max_price from profile
    # (useful for reranking/explanation) but we don't use it as a hard filter
    # unless the user asked for it.
    style_filter = None
    category_filter = attrs.get("category")
    max_price_filter = attrs.get("max_price")
    # Ensure max_price is numeric if present
    if max_price_filter is not None:
        try:
            max_price_filter = float(max_price_filter)
        except (ValueError, TypeError):
            max_price_filter = None

    # Only use max_price as a kNN pre-filter if the user explicitly stated price
    import re
    _price_match = re.search(
        r'\$\s*(\d+(?:\.\d+)?)|(?:under|below|less\s+than|cheaper\s+than|max)\s+\$?\s*(\d+(?:\.\d+)?)|(\d+(?:\.\d+)?)\s*dollars',
        query, re.IGNORECASE
    )
    if _price_match:
        # Extract the price value from whichever group matched
        _explicit_price = next((g for g in _price_match.groups() if g is not None), None)
        if _explicit_price:
            knn_price_filter = float(_explicit_price)
        else:
            knn_price_filter = max_price_filter
    else:
        knn_price_filter = None

    from search_personalization.agentic_memory.memory_tools import search_product_catalog, rerank_results
    gender_filter = attrs.get("gender_affinity")
    search_results = search_product_catalog(
        query=enriched_query,
        category=category_filter,
        max_price=knn_price_filter,
        gender_affinity=gender_filter,
    )

    search_duration = time.time() - start
    print(f"          Done ({search_duration:.1f}s)")

    # Extract the actual OpenSearch query from the search results for display
    try:
        _search_parsed = json.loads(search_results) if isinstance(search_results, str) else search_results
        actual_opensearch_query = _search_parsed.get("_opensearch_query")
    except (json.JSONDecodeError, TypeError):
        actual_opensearch_query = None

    # --- Step 3: Rerank + filter aversions (no LLM needed) ---
    print(f"  [rerank] Reranking + filtering...")
    start = time.time()
    # Build a rerank query that blends user intent with style/use-context signals.
    # The search step already used the enriched query (with profile preferences),
    # so reranking surfaces results that best match what the user actually asked for
    # plus their functional context.
    rerank_parts = [query]
    use_ctx = attrs.get("use_context")
    preferred_materials = attrs.get("preferred_materials", [])
    if use_ctx:
        rerank_parts.append(use_ctx)
    if preferred_materials:
        rerank_parts.append(", ".join(preferred_materials[:2]))
    rerank_query = ", ".join(rerank_parts)
    reranked_raw = rerank_results(query=rerank_query, documents=search_results, top_n=10)

    # Aversions are handled semantically — the query agent expresses them as positive
    # opposites in the enriched_query (e.g., "flashy" → "subtle understated") so the
    # vector search naturally favors non-aversive products. No post-filtering needed.
    raw_aversions = attrs.get("aversions", [])
    reranked = json.loads(reranked_raw) if isinstance(reranked_raw, str) else reranked_raw
    results_list = reranked.get("reranked_results", [])

    # Price post-filter: apply after reranking so semantic relevance isn't starved.
    # If too few results survive the price filter, keep the full set and let the
    # explanation agent inform the user about the price gap.
    price_filter_relaxed = False
    # Only apply price filtering if the user explicitly mentioned price in their query
    import re
    _user_mentioned_price = bool(re.search(r'(\$\d|under|below|less than|cheap|budget|affordable|\d+\s*dollars)', query.lower()))
    if max_price_filter and results_list and _user_mentioned_price:
        price_filtered = [r for r in results_list if (r.get("price") or 0) <= max_price_filter]
        if len(price_filtered) >= 3:
            results_list = price_filtered
        else:
            # Not enough results within budget — keep full set, flag for explanation
            price_filter_relaxed = True
            print(f"  [price] Only {len(price_filtered)} results under ${max_price_filter:.0f}, showing all with explanation")

    # Build final output
    final_results = {
        "results": [
            {"rank": i + 1, "product_id": r.get("id"), "name": r.get("name"), "price": r.get("price"), "score": r.get("rerank_score")}
            for i, r in enumerate(results_list[:7])
        ],
        "enriched_query": enriched_query,
        "personalization_summary": (
            f"Use context: {use_ctx or 'none'}, "
            f"materials: {preferred_materials or 'none'}, "
            f"filtered aversions: {raw_aversions}"
        ) if (raw_aversions or use_ctx or preferred_materials) else "No personalization applied",
    }
    ranking_result = json.dumps(final_results, indent=2, default=str)
    ranking_duration = time.time() - start
    pipeline_trace["agents"]["ranking"] = {
        "duration_ms": int(ranking_duration * 1000),
        "output": ranking_result,
    }
    print(f"          Done ({ranking_duration:.1f}s)")

    # Combine for working memory
    combined_payload = json.dumps({**enrichment, "search_results": results_list}, default=str)

    # --- Async: write working memory for observability (fire-and-forget) ---
    def _write_working_memory_async():
        try:
            from search_personalization.agentic_memory.memory_tools import write_working_memory
            write_working_memory(persona_id, session_id, combined_payload)
        except Exception:
            pass

    threading.Thread(target=_write_working_memory_async, daemon=True).start()

    # --- Async: Results Explanation Agent (runs in parallel, doesn't block) ---
    from concurrent.futures import Future
    explanation_future: Future = Future()

    def _run_explanation_async():
        try:
            from search_personalization.agentic_memory.agents.results_explanation_agent import (
                invoke_results_explanation_agent,
            )
            explanation = invoke_results_explanation_agent(
                query=query,
                results=results_list[:5],
                price_filter_relaxed=price_filter_relaxed,
                original_max_price=max_price_filter if price_filter_relaxed else None,
                inferred_attributes=attrs,
            )
            explanation_future.set_result(explanation)
        except Exception:
            explanation_future.set_result(None)

    explanation_thread = threading.Thread(target=_run_explanation_async, daemon=True)
    explanation_thread.start()

    pipeline_trace["total_duration_ms"] = int((time.time() - pipeline_start) * 1000)
    pipeline_trace["opensearch_query"] = actual_opensearch_query
    print(f"\n  Total pipeline: {pipeline_trace['total_duration_ms']}ms")

    # Price relaxation explanation is synchronous (no LLM needed, just a string).
    # LLM-based explanations run fully async — the future is returned in the trace
    # so the caller can resolve it without blocking the pipeline.
    if price_filter_relaxed and max_price_filter:
        pipeline_trace["agents"]["explanation"] = {
            "output": f"Not enough options under ${max_price_filter:.0f} — showing best matches across all price ranges."
        }
    else:
        pipeline_trace["agents"]["explanation"] = {"output": None, "_future": explanation_future}

    # --- Post-pipeline: session memory write (async fire-and-forget) ---
    def _write_session_memory_async():
        try:
            top_results = ", ".join(
                r.get("name", "") for r in results_list[:3]
            )
            session_content = (
                f"User searched: \"{query}\". "
                f"Category: {attrs.get('category') or 'general'}. "
                f"Top results: {top_results}."
            )
            write_session_memory(persona_id, session_id, session_content)
        except Exception:
            pass

    threading.Thread(target=_write_session_memory_async, daemon=True).start()

    # History write can remain async (not needed for multi-turn)
    def _write_history_async():
        try:
            write_history(persona_id, session_id,
                f"Session {session_id} | Query: \"{query}\" | "
                f"Persona: {persona_name} | Timestamp: {pipeline_trace['timestamp']}")
        except Exception:
            pass

    threading.Thread(target=_write_history_async, daemon=True).start()

    return pipeline_trace


def run_comparison_demo(query: str):
    """Run the same query through both personas side by side."""
    print("\n" + "=" * 70)
    print("  AGENTIC MEMORY DEMO — PERSONALIZED SEARCH")
    print("=" * 70)
    print(f"\n  Query: \"{query}\"")
    print(f"  Personas: Sarah (classic/professional style) vs Alex (casual/streetwear style)")
    print(f"  Architecture: Cache → Query Understanding → Ranking (3 agents)")

    demo_session_prefix = f"demo-{uuid.uuid4().hex[:6]}"

    sarah_trace = run_pipeline_for_persona(
        query=query, persona_name="sarah", session_id=f"{demo_session_prefix}-sarah",
    )
    alex_trace = run_pipeline_for_persona(
        query=query, persona_name="alex", session_id=f"{demo_session_prefix}-alex",
    )

    print("\n" + "=" * 70)
    print("  COMPARISON SUMMARY")
    print("=" * 70)
    print(f"\n  Same query: \"{query}\"")
    print(f"\n  Sarah's results:")
    print(f"    {_extract_summary(sarah_trace['agents'].get('ranking', {}).get('output', ''))}")
    print(f"\n  Alex's results:")
    print(f"    {_extract_summary(alex_trace['agents'].get('ranking', {}).get('output', ''))}")
    print(f"\n  Key insight: Same query — memory + decoupled ranking is the difference.")
    print("=" * 70)

    return {"sarah": sarah_trace, "alex": alex_trace}


def _extract_summary(output: str) -> str:
    if not output:
        return "(no output)"
    clean = output.replace("\n", " ").strip()
    return clean[:300] + "..." if len(clean) > 300 else clean


def run_interactive():
    """Interactive mode."""
    print("\n" + "=" * 70)
    print("  AGENTIC MEMORY DEMO — INTERACTIVE MODE")
    print("=" * 70)
    print("\n  Type a query to run through both personas")
    print("  'sarah: <query>' or 'alex: <query>' for one persona")
    print("  'quit' to stop\n")

    session_ids = {
        "sarah": f"interactive-{uuid.uuid4().hex[:6]}-sarah",
        "alex": f"interactive-{uuid.uuid4().hex[:6]}-alex",
    }

    while True:
        try:
            user_input = input("Query: ").strip()
        except (EOFError, KeyboardInterrupt):
            print("\nGoodbye!")
            break

        if not user_input or user_input.lower() in ("quit", "exit"):
            break

        if user_input.lower().startswith("sarah:"):
            run_pipeline_for_persona(user_input[6:].strip(), "sarah", session_ids["sarah"])
        elif user_input.lower().startswith("alex:"):
            run_pipeline_for_persona(user_input[5:].strip(), "alex", session_ids["alex"])
        else:
            run_comparison_demo(user_input)
        print()


def main():
    import argparse
    parser = argparse.ArgumentParser(description="Run the agentic memory demo pipeline")
    parser.add_argument("--query", type=str, default=DEFAULT_QUERY)
    parser.add_argument("--interactive", action="store_true")
    parser.add_argument("--persona", type=str, choices=["sarah", "alex"])
    args = parser.parse_args()

    # Refresh schema cache and start daily scheduler (7am)
    from search_personalization.agentic_memory.schema_cache import refresh_cache, start_scheduler
    try:
        print("  Refreshing schema cache...")
        refresh_cache()
    except Exception as e:
        print(f"  Schema cache refresh failed (will use defaults): {e}")
    start_scheduler(hour=7, minute=0)

    if args.interactive:
        run_interactive()
    elif args.persona:
        run_pipeline_for_persona(query=args.query, persona_name=args.persona)
    else:
        run_comparison_demo(args.query)


if __name__ == "__main__":
    main()