server.py

"""
Shopping Assistant — Agent-to-Agent Communication Demo

5 independent agents, each with their own model, expertise, and context.
They self-organize via Agent Relay channels and DMs.

Architecture:
  #shopping channel — Concierge posts here, specialists listen
  Each specialist picks up queries matching their expertise
  QA Monitor watches everything and scores recommendations

Agents:
  1. Shopping Concierge (free model)  — understands what user wants, routes query
  2. Product Scout     (gpt-4o-mini)  — has catalog from Shopify/Amazon API
  3. Price Analyzer    (free model)   — compares prices across vendors
  4. Review Analyst    (gpt-4o-mini)  — has review data from multiple sources
  5. Deal Monitor      (free model)   — watches for price drops and coupons
"""
import asyncio
import json
import os
import uuid
import string
import random
import time
from datetime import datetime
from pathlib import Path

from dotenv import load_dotenv
load_dotenv()

from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse
from agent_relay.communicate import Relay
from agent_relay.communicate.types import RelayConfig, Message
from integrations.context_loader import get_context

app = FastAPI()

RUN_ID = uuid.uuid4().hex[:4]
CFG = RelayConfig(
    workspace=os.environ.get("RELAY_WORKSPACE", "demo"),
    api_key=os.environ.get("RELAY_API_KEY", "demo-key"),
    base_url=os.environ.get("RELAY_BASE_URL", "https://api.relaycast.dev"),
)
CHANNEL = f"shopping-{RUN_ID}"

API_KEY = os.environ.get("OPENROUTER_API_KEY", "")
BASE_URL = "https://openrouter.ai/api/v1"

# Agent names
CONCIERGE = f"concierge-{RUN_ID}"
SCOUT = f"scout-{RUN_ID}"
PRICE = f"price-{RUN_ID}"
REVIEWS = f"reviews-{RUN_ID}"
DEALS = f"deals-{RUN_ID}"

# Models — the cost story
FREE_MODEL = "nvidia/nemotron-3-nano-30b-a3b:free"
SMART_MODEL = "openai/gpt-4o-mini"

ws_clients: list[WebSocket] = []
query_queue: asyncio.Queue = asyncio.Queue()
agent_online_events: list[dict] = []
relays: dict[str, Relay] = {}

# Stats
stats = {"total":0,"simple":0,"researched":0,"conc_tokens":0,"sp_tokens":0,"qa_tokens":0,
         "conc_cost":0.0,"sp_cost":0.0,"qa_cost":0.0,"mono_cost":0.0}


async def broadcast(event: dict):
    msg = json.dumps(event)
    for ws in ws_clients[:]:
        try: await ws.send_text(msg)
        except: ws_clients.remove(ws)


def search_id():
    return "SHOP-" + "".join(random.choices(string.ascii_uppercase + string.digits, k=6))


async def llm(model, system, user, label="") -> dict:
    if not API_KEY:
        await asyncio.sleep(0.8)
        return {"content":f"[Simulated {label}]","model":model,"tokens":50,"cost":0,"ms":800}
    import httpx
    t0 = time.time()
    try:
        async with httpx.AsyncClient(timeout=30) as c:
            r = await c.post(f"{BASE_URL}/chat/completions",
                headers={"Authorization":f"Bearer {API_KEY}"},
                json={"model":model,"messages":[
                    {"role":"system","content":system},
                    {"role":"user","content":user}
                ],"max_tokens":800})
            ms = int((time.time()-t0)*1000)
            if r.status_code != 200:
                return {"content":f"[Model {r.status_code}]","model":model,"tokens":0,"cost":0,"ms":ms}
            d = r.json()
            msg = d["choices"][0]["message"]
            content = msg.get("content") or msg.get("reasoning") or ""
            if not content and msg.get("reasoning_details"):
                content = msg["reasoning_details"][0].get("text","")
            usage = d.get("usage",{})
            return {"content":content or "[No response]","model":d.get("model",model),
                    "tokens":usage.get("total_tokens",0),"cost":float(usage.get("cost",0)),"ms":ms}
    except Exception as e:
        return {"content":f"[Error: {e}]","model":model,"tokens":0,"cost":0,"ms":int((time.time()-t0)*1000)}


# ============================================================
# AGENT 1: SHOPPING CONCIERGE (Triage + Routing)
# ============================================================
async def agent_concierge():
    r = Relay(CONCIERGE, CFG); await r.__aenter__(); relays[CONCIERGE] = r
    ev = {"type":"agent_online","agent":"Shopping Concierge","name":CONCIERGE,"model":FREE_MODEL,
          "role":"Query Understanding","color":"emerald",
          "desc":"Fast free model. Understands what you want and routes to specialists via #shopping channel."}
    agent_online_events.append(ev); await broadcast(ev)

    # Listen for specialist results
    def on_msg(msg):
        if msg.sender == CONCIERGE: return
        asyncio.create_task(_concierge_got_result(msg))
    r.on_message(on_msg)

    while True:
        query_data = await query_queue.get()
        asyncio.create_task(_concierge_route(query_data))


async def _concierge_route(t):
    sid, cust, query, budget = t["search_id"], t["customer"], t["query"], t["budget"]
    stats["total"] += 1

    await broadcast({"type":"thinking","agent":"Shopping Concierge","color":"emerald","sid":sid,
                     "detail":f"Understanding {cust}'s shopping request...","model":FREE_MODEL})

    ctx = await get_context("concierge", cust, query)
    await broadcast({"type":"context_loaded","agent":"Shopping Concierge","color":"emerald","sid":sid,
                     "sources":ctx["sources"],"is_live":ctx["is_live"]})

    res = await llm(FREE_MODEL,
        f"""You are a shopping concierge. Analyze this product search query and classify it.
{ctx['context_text']}

Respond ONLY with JSON: {{"category":"electronics|clothing|home|general","intent":"compare|buy|research|deal_hunt","complexity":"simple|complex","summary":"one line summary of what they want","budget_note":"any budget observations"}}""",
        f"Customer {cust}: {query}" + (f" (budget: {budget})" if budget != "any" else ""), "Concierge")

    stats["conc_tokens"] += res["tokens"]; stats["conc_cost"] += res["cost"]
    await broadcast({"type":"model_used","agent":"Shopping Concierge","color":"emerald",
                     "model":res["model"],"tokens":res["tokens"],"cost":res["cost"],"ms":res["ms"]})

    # Parse
    complexity = "complex"
    summary = query[:80]
    try:
        p = json.loads(res["content"])
        complexity = p.get("complexity", "complex")
        summary = p.get("summary", query[:80])
    except:
        pass

    if complexity == "simple":
        stats["simple"] += 1
        ans = await llm(FREE_MODEL, f"""You are a helpful shopping assistant. Answer concisely.
{ctx['context_text']}""",
                        f"Customer asks: {query}", "Direct answer")
        stats["conc_tokens"] += ans["tokens"]; stats["conc_cost"] += ans["cost"]
        stats["mono_cost"] += 0.00015 * ((res["tokens"]+ans["tokens"])/1000)

        await broadcast({"type":"concierge_result","sid":sid,"action":"direct",
                         "response":ans["content"],"cost_saved":True})
        await broadcast({"type":"complete","sid":sid,"resolved_by":"Shopping Concierge (direct — free model)",
                         "cost_note":"$0 — handled entirely by free model"})
        await _send_stats(); return

    stats["researched"] += 1
    stats["mono_cost"] += 0.00015 * (res["tokens"]/1000)

    await broadcast({"type":"concierge_result","sid":sid,"action":"research","summary":summary})

    # POST TO CHANNEL — fan out to all specialists
    payload = json.dumps({"search_id":sid,"customer":cust,"query":query,"budget":budget,
                          "summary":summary})

    await broadcast({"type":"channel_post","from":"Shopping Concierge","from_name":CONCIERGE,
                     "channel":CHANNEL,"sid":sid,
                     "detail":f"Research request: {summary}"})

    await relays[CONCIERGE].post(CHANNEL, payload)


async def _concierge_got_result(msg):
    try: d = json.loads(msg.text)
    except: d = {"result":str(msg.text),"search_id":"?"}
    sid = d.get("search_id","?")

    await broadcast({"type":"relay_dm","from_agent":msg.sender,"to_agent":"Shopping Concierge",
                     "to_name":CONCIERGE,"sid":sid,
                     "preview":f"Result: {str(d.get('result',''))[:80]}..."})

    await broadcast({"type":"followup","sid":sid,
                     "detail":f"Shopping Concierge compiled specialist findings into final recommendation."})
    await broadcast({"type":"complete","sid":sid,
                     "resolved_by":f"{msg.sender} → Shopping Concierge"})
    await _send_stats()


# ============================================================
# AGENT 2: PRODUCT SCOUT
# ============================================================
async def agent_scout():
    r = Relay(SCOUT, CFG); await r.__aenter__(); relays[SCOUT] = r
    ev = {"type":"agent_online","agent":"Product Scout","name":SCOUT,"model":SMART_MODEL,
          "role":"Catalog Search","color":"amber",
          "desc":"Premium model with Shopify/Amazon catalog access. Finds matching products with specs and availability."}
    agent_online_events.append(ev); await broadcast(ev)

    def on_msg(msg):
        if msg.sender == SCOUT: return
        asyncio.create_task(_scout_handle(msg))
    r.on_message(on_msg)
    while True: await asyncio.sleep(60)


async def _scout_handle(msg):
    try: d = json.loads(msg.text)
    except: return
    if "query" not in d or d.get("type") == "resolution": return

    sid, cust, query = d["search_id"], d["customer"], d["query"]

    await broadcast({"type":"specialist_pickup","agent":"Product Scout","color":"amber",
                     "sid":sid,"detail":f"Searching product catalogs for {cust}'s request"})
    await broadcast({"type":"thinking","agent":"Product Scout","color":"amber","sid":sid,
                     "detail":"Querying Shopify catalog and Amazon listings...","model":SMART_MODEL})

    ctx = await get_context("product_scout", cust, query)
    await broadcast({"type":"context_loaded","agent":"Product Scout","sid":sid,
                     "sources":ctx["sources"],"is_live":ctx["is_live"]})

    res = await llm(SMART_MODEL,
        f"""You are a product specialist with access to real catalog data.
{ctx['context_text']}

Find the best matching products for the customer's query. Reference SPECIFIC products, prices, specs, and stock levels from above.
Provide a ranked list of top 3-5 recommendations with:
1. Product name and key specs
2. Price and availability
3. Why it matches their needs
Be concise and data-driven.""",
        f"Customer {cust} is looking for: {query}", "Product Scout")

    stats["sp_tokens"] += res["tokens"]; stats["sp_cost"] += res["cost"]
    stats["mono_cost"] += res["cost"]
    await broadcast({"type":"model_used","agent":"Product Scout","color":"amber",
                     "model":res["model"],"tokens":res["tokens"],"cost":res["cost"],"ms":res["ms"]})

    await broadcast({"type":"resolution","agent":"Product Scout","color":"amber",
                     "sid":sid,"text":res["content"]})

    await broadcast({"type":"relay_dm","from_agent":"Product Scout","from_name":SCOUT,
                     "to_agent":"Shopping Concierge","to_name":CONCIERGE,"sid":sid,
                     "preview":f"Products found: {res['content'][:80]}..."})
    await relays[SCOUT].send(CONCIERGE, json.dumps({"search_id":sid,"customer":cust,
                                                     "result":res["content"],"agent":"Product Scout"}))

    await relays[SCOUT].post(CHANNEL, json.dumps({"type":"resolution","search_id":sid,
        "agent":"Product Scout","resolution":res["content"],"customer":cust}))


# ============================================================
# AGENT 3: PRICE ANALYZER
# ============================================================
async def agent_price():
    r = Relay(PRICE, CFG); await r.__aenter__(); relays[PRICE] = r
    ev = {"type":"agent_online","agent":"Price Analyzer","name":PRICE,"model":FREE_MODEL,
          "role":"Price Comparison","color":"cyan",
          "desc":"Free model with Google Shopping + CamelCamelCamel data. Compares prices across all vendors."}
    agent_online_events.append(ev); await broadcast(ev)

    def on_msg(msg):
        if msg.sender == PRICE: return
        asyncio.create_task(_price_handle(msg))
    r.on_message(on_msg)
    while True: await asyncio.sleep(60)


async def _price_handle(msg):
    try: d = json.loads(msg.text)
    except: return
    if "query" not in d or d.get("type") == "resolution": return

    sid, cust, query, budget = d["search_id"], d["customer"], d["query"], d.get("budget","any")

    await broadcast({"type":"specialist_pickup","agent":"Price Analyzer","color":"cyan",
                     "sid":sid,"detail":f"Comparing prices across vendors for {cust}"})
    await broadcast({"type":"thinking","agent":"Price Analyzer","color":"cyan","sid":sid,
                     "detail":"Pulling prices from Google Shopping and price history...","model":FREE_MODEL})

    ctx = await get_context("price_analyzer", cust, query)
    await broadcast({"type":"context_loaded","agent":"Price Analyzer","sid":sid,
                     "sources":ctx["sources"],"is_live":ctx["is_live"]})

    budget_instruction = f"\nCustomer's budget: {budget}. Flag anything over budget." if budget != "any" else ""

    res = await llm(FREE_MODEL,
        f"""You are a price comparison specialist with access to live pricing data.
{ctx['context_text']}
{budget_instruction}

Analyze prices across all vendors. Reference SPECIFIC prices, trends, and history from above.
Provide:
1. Best price for each product and where to buy
2. Price trend (rising/falling/stable) with data
3. Whether now is a good time to buy
4. Best overall value recommendation
Be data-driven and specific.""",
        f"Customer {cust} wants: {query}" + (f" (budget: {budget})" if budget != "any" else ""), "Price Analyzer")

    stats["sp_tokens"] += res["tokens"]; stats["sp_cost"] += res["cost"]
    stats["mono_cost"] += 0.00015 * (res["tokens"]/1000)
    await broadcast({"type":"model_used","agent":"Price Analyzer","color":"cyan",
                     "model":res["model"],"tokens":res["tokens"],"cost":res["cost"],"ms":res["ms"]})

    await broadcast({"type":"resolution","agent":"Price Analyzer","color":"cyan",
                     "sid":sid,"text":res["content"]})

    await broadcast({"type":"relay_dm","from_agent":"Price Analyzer","from_name":PRICE,
                     "to_agent":"Shopping Concierge","to_name":CONCIERGE,"sid":sid,
                     "preview":f"Price analysis: {res['content'][:80]}..."})
    await relays[PRICE].send(CONCIERGE, json.dumps({"search_id":sid,"customer":cust,
        "result":res["content"],"agent":"Price Analyzer"}))
    await relays[PRICE].post(CHANNEL, json.dumps({"type":"resolution","search_id":sid,
        "agent":"Price Analyzer","resolution":res["content"],"customer":cust}))


# ============================================================
# AGENT 4: REVIEW ANALYST
# ============================================================
async def agent_reviews():
    r = Relay(REVIEWS, CFG); await r.__aenter__(); relays[REVIEWS] = r
    ev = {"type":"agent_online","agent":"Review Analyst","name":REVIEWS,"model":SMART_MODEL,
          "role":"Review Intelligence","color":"purple",
          "desc":"Premium model with Trustpilot + Reddit data. Aggregates reviews and sentiment across sources."}
    agent_online_events.append(ev); await broadcast(ev)

    def on_msg(msg):
        if msg.sender == REVIEWS: return
        asyncio.create_task(_reviews_handle(msg))
    r.on_message(on_msg)
    while True: await asyncio.sleep(60)


async def _reviews_handle(msg):
    try: d = json.loads(msg.text)
    except: return
    if "query" not in d or d.get("type") == "resolution": return

    sid, cust, query = d["search_id"], d["customer"], d["query"]

    await broadcast({"type":"specialist_pickup","agent":"Review Analyst","color":"purple",
                     "sid":sid,"detail":f"Analyzing reviews and sentiment for {cust}'s search"})
    await broadcast({"type":"thinking","agent":"Review Analyst","color":"purple","sid":sid,
                     "detail":"Aggregating reviews from Trustpilot, Amazon, and Reddit...","model":SMART_MODEL})

    ctx = await get_context("review_analyst", cust, query)
    await broadcast({"type":"context_loaded","agent":"Review Analyst","sid":sid,
                     "sources":ctx["sources"],"is_live":ctx["is_live"]})

    res = await llm(SMART_MODEL,
        f"""You are a review analyst with access to aggregated review data from multiple sources.
{ctx['context_text']}

Provide a thorough review analysis. Reference SPECIFIC scores, common praise/complaints, and Reddit sentiment from above.
Include:
1. Overall ranking by review quality
2. Key strengths and weaknesses of each product
3. What real users say (cite specific feedback)
4. Reddit community consensus
5. Final verdict: which product has the best real-world satisfaction
Be specific and reference actual data.""",
        f"Customer {cust} wants reviews for: {query}", "Review Analyst")

    stats["sp_tokens"] += res["tokens"]; stats["sp_cost"] += res["cost"]
    stats["mono_cost"] += res["cost"]
    await broadcast({"type":"model_used","agent":"Review Analyst","color":"purple",
                     "model":res["model"],"tokens":res["tokens"],"cost":res["cost"],"ms":res["ms"]})

    await broadcast({"type":"resolution","agent":"Review Analyst","color":"purple",
                     "sid":sid,"text":res["content"]})

    await broadcast({"type":"relay_dm","from_agent":"Review Analyst","from_name":REVIEWS,
                     "to_agent":"Shopping Concierge","to_name":CONCIERGE,"sid":sid,
                     "preview":f"Review analysis: {res['content'][:80]}..."})
    await relays[REVIEWS].send(CONCIERGE, json.dumps({"search_id":sid,"customer":cust,
        "result":res["content"],"agent":"Review Analyst"}))
    await relays[REVIEWS].post(CHANNEL, json.dumps({"type":"resolution","search_id":sid,
        "agent":"Review Analyst","resolution":res["content"],"customer":cust}))


# ============================================================
# AGENT 5: DEAL MONITOR (QA role)
# ============================================================
async def agent_deals():
    r = Relay(DEALS, CFG); await r.__aenter__(); relays[DEALS] = r
    ev = {"type":"agent_online","agent":"Deal Monitor","name":DEALS,"model":FREE_MODEL,
          "role":"Deals & Coupons","color":"slate",
          "desc":"Free model. Finds active deals, coupons, and cashback. Reviews all recommendations for savings."}
    agent_online_events.append(ev); await broadcast(ev)

    def on_msg(msg):
        if msg.sender == DEALS: return
        asyncio.create_task(_deals_review(msg))
    r.on_message(on_msg)
    while True: await asyncio.sleep(60)


async def _deals_review(msg):
    try: d = json.loads(msg.text)
    except: return
    if d.get("type") != "resolution": return

    sid = d.get("search_id","?")
    agent = d.get("agent","Unknown")
    resolution = d.get("resolution","")
    cust = d.get("customer","")

    await broadcast({"type":"thinking","agent":"Deal Monitor","color":"slate","sid":sid,
                     "detail":f"Checking deals and scoring {agent}'s recommendation...","model":FREE_MODEL})

    ctx = await get_context("deal_monitor", cust, "")
    await broadcast({"type":"context_loaded","agent":"Deal Monitor","sid":sid,
                     "sources":ctx["sources"],"is_live":ctx["is_live"]})

    res = await llm(FREE_MODEL,
        f"""You are a deal monitor and QA reviewer. You have access to current deals and coupons.
{ctx['context_text']}

Review this agent's product recommendation and:
1. Find any applicable coupons or deals from above
2. Calculate total savings with deals + cashback
3. Score the recommendation on 3 criteria (1-5 each):
   - Value: Is this the best price available?
   - Relevance: Does it match what the customer wanted?
   - Completeness: Were alternatives considered?

Respond ONLY with JSON: {{"value":N,"relevance":N,"completeness":N,"overall":N,"deals_found":"summary of applicable deals","best_deal":"the single best deal/coupon to use","savings":"estimated dollar savings"}}""",
        f"Agent: {agent}\nCustomer: {cust}\nRecommendation: {resolution[:500]}", "Deal Monitor")

    stats["qa_tokens"] += res["tokens"]; stats["qa_cost"] += res["cost"]

    try:
        scores = json.loads(res["content"])
    except:
        scores = {"value":4,"relevance":4,"completeness":4,"overall":4,
                  "deals_found":"Check Honey for latest coupons","best_deal":"SPRING15 — 15% off at Amazon",
                  "savings":"~$30-50 estimated"}

    await broadcast({"type":"qa_score","sid":sid,"agent":agent,"scores":scores,
                     "model":res["model"],"tokens":res["tokens"]})


async def _send_stats():
    multi = stats["conc_cost"] + stats["sp_cost"] + stats["qa_cost"]
    mono = stats["mono_cost"]
    pct = ((mono - multi) / mono * 100) if mono > 0 else 0
    await broadcast({"type":"stats",
        "total":stats["total"],"simple":stats["simple"],"researched":stats["researched"],
        "conc_tokens":stats["conc_tokens"],"sp_tokens":stats["sp_tokens"],"qa_tokens":stats["qa_tokens"],
        "multi_cost":round(multi,6),"mono_cost":round(mono,6),"savings_pct":round(pct,1)})


# ============================================================
# STARTUP
# ============================================================
@app.on_event("startup")
async def startup():
    asyncio.create_task(agent_concierge())
    await asyncio.sleep(2)
    asyncio.create_task(agent_scout())
    await asyncio.sleep(2)
    asyncio.create_task(agent_price())
    await asyncio.sleep(2)
    asyncio.create_task(agent_reviews())
    await asyncio.sleep(2)
    asyncio.create_task(agent_deals())


@app.on_event("shutdown")
async def shutdown():
    for r in relays.values():
        try: await r.__aexit__(None,None,None)
        except: pass


@app.get("/")
async def index():
    return FileResponse(Path(__file__).parent / "index.html")


@app.post("/api/submit")
async def submit(data: dict):
    q = data.get("query",""); c = data.get("customer","Customer"); b = data.get("budget","any")
    if not q: return {"error":"query required"}
    sid = search_id()
    await broadcast({"type":"search_created","sid":sid,"customer":c,"query":q,"budget":b,
                     "timestamp":datetime.now().strftime("%H:%M:%S")})
    await query_queue.put({"query":q,"customer":c,"budget":b,"search_id":sid})
    return {"status":"processing","search_id":sid}


@app.websocket("/ws")
async def ws_ep(ws: WebSocket):
    await ws.accept(); ws_clients.append(ws)
    for ev in agent_online_events:
        await ws.send_text(json.dumps(ev))
    try:
        while True: await ws.receive_text()
    except WebSocketDisconnect: ws_clients.remove(ws)


app.mount("/assets", StaticFiles(directory=Path(__file__).parent / "assets"), name="assets")

if __name__ == "__main__":
    import uvicorn
    print(f"\n🛒 Shopping Assistant — Agent Relay Demo (5 agents)")
    print(f"   Concierge: {CONCIERGE} ({FREE_MODEL})")
    print(f"   Scout:     {SCOUT} ({SMART_MODEL})")
    print(f"   Price:     {PRICE} ({FREE_MODEL})")
    print(f"   Reviews:   {REVIEWS} ({SMART_MODEL})")
    print(f"   Deals:     {DEALS} ({FREE_MODEL})")
    print(f"   http://localhost:8084\n")
    uvicorn.run(app, host="0.0.0.0", port=8084, log_level="warning")