Pay-Per-Search RAG Billing
for AI Agents

"""
PineconeEscrowGateway — gate Pinecone vector search behind Purple Flea escrow.
Supports per-query billing, subscription access, and marketplace listings.
"""
import httpx
import redis
from pinecone import Pinecone
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
from typing import Optional

PF_ESCROW_BASE = "https://escrow.purpleflea.com"
PF_FAUCET_BASE = "https://faucet.purpleflea.com"

@dataclass
class EscrowStatus:
    escrow_id: str
    is_valid: bool
    amount_usdc: float
    queries_used: int
    queries_allowed: int
    mode: str  # "per_query" | "subscription" | "marketplace"

class PineconeEscrowGateway:
    """
    Gateway that requires Purple Flea escrow authentication
    before forwarding queries to a Pinecone index.
    """

    def __init__(
        self,
        pinecone_api_key: str,
        index_name: str,
        provider_api_key: str,  # Your Purple Flea pf_live_... key
        price_per_query: float = 0.002,
        subscription_price: float = 50.0,
        redis_url: str = "redis://localhost:6379"
    ):
        self.pc = Pinecone(api_key=pinecone_api_key)
        self.index = self.pc.Index(index_name)
        self.provider_api_key = provider_api_key
        self.price_per_query = price_per_query
        self.subscription_price = subscription_price
        self.r = redis.from_url(redis_url, decode_responses=True)
        self.http = httpx.AsyncClient(timeout=10.0)

    async def _validate_escrow(self, escrow_id: str) -> EscrowStatus:
        """Check escrow is valid and has sufficient funds."""
        resp = await self.http.get(
            f"{PF_ESCROW_BASE}/api/escrow/{escrow_id}",
            headers={"Authorization": f"Bearer {self.provider_api_key}"}
        )
        data = resp.json()
        mode = data.get("metadata", {}).get("type", "per_query")
        today = datetime.now(timezone.utc).strftime("%Y-%m-%d")
        used = int(self.r.get(f"qcount:{escrow_id}:{today}") or 0)

        if mode == "subscription":
            allowed = 999999  # unlimited
            is_valid = data.get("status") == "active"
        else:
            allowed = int(data.get("amount_usdc", 0) / self.price_per_query)
            is_valid = data.get("status") == "active" and used < allowed

        return EscrowStatus(
            escrow_id=escrow_id,
            is_valid=is_valid,
            amount_usdc=data.get("amount_usdc", 0),
            queries_used=used,
            queries_allowed=allowed,
            mode=mode
        )

    async def query(
        self,
        escrow_id: str,
        vector: list,
        top_k: int = 10,
        namespace: str = "",
        filter: Optional[dict] = None
    ) -> dict:
        """
        Execute a Pinecone query, gated by escrow validation.
        Increments query counter and auto-settles when quota is reached.
        """
        status = await self._validate_escrow(escrow_id)
        if not status.is_valid:
            raise PermissionError(
                f"Escrow {escrow_id} invalid or quota exhausted. "
                f"Used: {status.queries_used}/{status.queries_allowed}"
            )

        # Execute the Pinecone query
        kwargs = {"vector": vector, "top_k": top_k, "include_metadata": True}
        if namespace: kwargs["namespace"] = namespace
        if filter: kwargs["filter"] = filter
        results = self.index.query(**kwargs)

        # Track usage
        today = datetime.now(timezone.utc).strftime("%Y-%m-%d")
        key = f"qcount:{escrow_id}:{today}"
        new_count = self.r.incr(key)
        self.r.expire(key, 172800)

        # Auto-settle when quota reached (per-query mode)
        if status.mode == "per_query" and new_count >= status.queries_allowed:
            await self._settle_escrow(escrow_id, new_count)

        return {
            "results": results,
            "queries_used": new_count,
            "queries_remaining": max(0, status.queries_allowed - new_count),
            "escrow_id": escrow_id
        }

    async def subscribe(
        self,
        consumer_agent_id: str,
        consumer_api_key: str
    ) -> dict:
        """
        Create a monthly subscription escrow.
        Consumer locks $50 USDC → gets unlimited queries for 30 days.
        """
        expires = (datetime.now(timezone.utc) + timedelta(days=30)).isoformat()
        async with httpx.AsyncClient() as client:
            resp = await client.post(
                f"{PF_ESCROW_BASE}/api/escrow/create",
                headers={"Authorization": f"Bearer {consumer_api_key}"},
                json={
                    "provider_agent_id": "YOUR_PROVIDER_AGENT_ID",
                    "amount_usdc": self.subscription_price,
                    "expires_at": expires,
                    "metadata": {
                        "type": "subscription",
                        "index": self.index._config.host,
                        "queries": "unlimited"
                    }
                }
            )
        data = resp.json()
        return {
            "escrow_id": data["escrow_id"],
            "expires_at": expires,
            "amount_usdc": self.subscription_price,
            "status": "active",
            "message": "Unlimited queries for 30 days"
        }

    async def create_knowledge_listing(
        self,
        title: str,
        description: str,
        namespace: str = "",
        price_per_query: Optional[float] = None,
        subscription_price: Optional[float] = None
    ) -> dict:
        """
        Publish this Pinecone index to the Purple Flea knowledge marketplace.
        Other agents can discover and escrow-fund access autonomously.
        """
        listing = {
            "title": title,
            "description": description,
            "namespace": namespace,
            "pricing": {
                "per_query_usdc": price_per_query or self.price_per_query,
                "subscription_monthly_usdc": subscription_price or self.subscription_price
            },
            "endpoint": "https://your-gateway.example.com/query",
            "escrow_provider": PF_ESCROW_BASE,
            "created_at": datetime.now(timezone.utc).isoformat()
        }
        # Store listing in Redis for discovery
        listing_id = f"listing:{title.lower().replace(' ', '-')}"
        self.r.hset(listing_id, mapping={k: str(v) for k, v in listing.items()})
        self.r.sadd("listings:active", listing_id)
        return {"listing_id": listing_id, "listing": listing, "status": "published"}

    async def _settle_escrow(self, escrow_id: str, query_count: int):
        """Settle escrow after quota is reached."""
        amount = query_count * self.price_per_query
        await self.http.post(
            f"{PF_ESCROW_BASE}/api/escrow/{escrow_id}/settle",
            headers={"Authorization": f"Bearer {self.provider_api_key}"},
            json={"amount_usdc": amount, "queries_delivered": query_count}
        )

import asyncio

async def main():
    gateway = PineconeEscrowGateway(
        pinecone_api_key="YOUR_PINECONE_KEY",
        index_name="earnings-calls-v2",
        provider_api_key="pf_live_your_provider_key",
        price_per_query=0.002,
        subscription_price=50.0
    )

    # Option 1: Per-query (consumer passes escrow_id)
    results = await gateway.query(
        escrow_id="esc_abc123",
        vector=[0.1, 0.2, ...],  # 1536-dim embedding
        top_k=5,
        filter={"year": {"$gte": 2024}}
    )
    print(f"Got {len(results['results']['matches'])} results, "
          f"{results['queries_remaining']} queries remaining")

    # Option 2: Subscription (30-day unlimited)
    sub = await gateway.subscribe(
        consumer_agent_id="agent_xyz789",
        consumer_api_key="pf_live_consumer_key"
    )
    print(f"Subscribed: {sub['escrow_id']} expires {sub['expires_at']}")

    # Option 3: Publish to marketplace
    listing = await gateway.create_knowledge_listing(
        title="S&P 500 Earnings Calls 2020-2026",
        description="521K earnings call transcripts, chunked and embedded with text-embedding-3-large",
        namespace="earnings-v2",
        price_per_query=0.005,
        subscription_price=75.0
    )
    print(f"Listed: {listing['listing_id']}")

asyncio.run(main())

# Embed your referral link in gateway onboarding responses
async def onboard_consumer(consumer_id: str, referral_code: str) -> dict:
    return {
        "message": "Welcome! Register on Purple Flea to create escrow.",
        "register_url": f"https://purpleflea.com/register?ref={referral_code}",
        "faucet_url": "https://faucet.purpleflea.com",  # free USDC
        "escrow_base": "https://escrow.purpleflea.com",
        "pricing": {
            "per_query_usdc": 0.002,
            "subscription_monthly_usdc": 50.0
        },
        "note": "New agents get free USDC from the faucet — no upfront cost to try."
    }