NVIDIA_API_KEY=nvapi-...
DATABRICKS_HOST=https://<your-workspace>.azuredatabricks.net
DATABRICKS_TOKEN=dapi...
MLFLOW_TRACKING_URI=databricks
MLFLOW_REGISTRY_URI=databricks-uc
MLFLOW_EXPERIMENT_ID=<experiment-id-from-databricks-ui>
UPSTASH_REDIS_REST_URL=https://...
UPSTASH_REDIS_REST_TOKEN=...

!pip install \
    "ipykernel>=7.3.0" \
    "langchain>=1.3.11" \
    "langchain-nvidia-ai-endpoints>=1.4.2" \
    "mlflow>=3.14.0" \
    "pydantic>=2.13.4" \
    "upstash-redis>=1.7.0"

import os
import mlflow
from langchain.chat_models import init_chat_model

# --- NVIDIA API config ---
API_KEY = os.environ["NVIDIA_API_KEY"]
BASE_URL = "https://integrate.api.nvidia.com/v1"
CHAT_MODEL_NAME = "openai/gpt-oss-120b"

# --- MLflow tracing config ---
# Old (local) tracking setup — kept for reference / fallback
# TRACKING_URI = "http://localhost:5909"
# mlflow.set_tracking_uri(TRACKING_URI)
# mlflow.set_experiment("supply-chain-disruption-poc")

# Databricks MLflow tracking (credentials loaded from .env)
from dotenv import load_dotenv
load_dotenv()

os.environ["DATABRICKS_TOKEN"] = os.getenv("DATABRICKS_TOKEN", "")
os.environ["DATABRICKS_HOST"] = os.getenv("DATABRICKS_HOST", "")

TRACKING_URI = os.getenv("MLFLOW_TRACKING_URI", "databricks")
REGISTRY_URI = os.getenv("MLFLOW_REGISTRY_URI", "databricks-uc")
EXPERIMENT_ID = os.getenv("MLFLOW_EXPERIMENT_ID", "")

mlflow.set_tracking_uri(TRACKING_URI)
mlflow.set_registry_uri(REGISTRY_URI)
if EXPERIMENT_ID:
    mlflow.set_experiment(experiment_id=EXPERIMENT_ID)
else:
    mlflow.set_experiment("supply-chain-disruption-poc")

# Auto-trace ALL LangChain calls (LLM, tools, agents, chains)
mlflow.langchain.autolog()

print("✓ MLflow tracing enabled")
print(f"  Tracking URI: {TRACKING_URI}")
print(f"  Registry URI: {REGISTRY_URI}")

llm = init_chat_model(
    CHAT_MODEL_NAME,
    model_provider="nvidia",
    base_url=BASE_URL,
    api_key=API_KEY,
)

response = llm.invoke("Reply with just the word 'pong'.")
print(f"Smoke test: {response.content}")

from pydantic import BaseModel, Field
from typing import Literal
import time


class ShipmentStatus(BaseModel):
    """Current state of a shipment in the supply chain."""
    shipment_id: str
    status: Literal["in_transit", "delayed", "rerouted", "delivered", "lost"]
    reason: str = ""
    source_agent_id: str
    timestamp: float = Field(default_factory=time.time)
    ttl_seconds: int = 3600
    version: int = 1


class EventLog(BaseModel):
    """A discrete event in the supply chain."""
    event_id: str
    event_type: Literal["disruption_detected", "reroute_planned", "customer_notified", "delivery_confirmed"]
    shipment_id: str
    details: str = ""
    source_agent_id: str
    timestamp: float = Field(default_factory=time.time)
    ttl_seconds: int = 86400
    version: int = 1


print("✓ Schemas defined: ShipmentStatus, EventLog")

from dotenv import load_dotenv
from upstash_redis import Redis

load_dotenv()

redis_client = Redis(
    url=os.environ["UPSTASH_REDIS_REST_URL"],
    token=os.environ["UPSTASH_REDIS_REST_TOKEN"],
)

import json as _json


class SharedEpistemicMemory:
    """Upstash Redis-backed Shared Epistemic Memory store.

    Agents read/write through this object exclusively — no direct agent-to-agent calls.
    """

    def __init__(self, client: Redis):
        self._client = client

    def write(self, key: str, entry: BaseModel) -> None:
        self._client.set(key, entry.model_dump_json())

    def read(self, key: str) -> BaseModel | None:
        raw = self._client.get(key)
        if raw is None:
            return None
        if isinstance(raw, bytes):
            raw = raw.decode("utf-8")
        entry = self._deserialize(key, raw)
        if entry is None:
            return None
        if not self.is_fresh(entry):
            return None
        return entry

    def update(self, key: str, entry: BaseModel, expected_version: int) -> BaseModel:
        current = self.read(key)
        if current is None:
            raise KeyError(f"Key {key!r} not found")
        if current.version != expected_version:
            raise ValueError(
                f"Version conflict on {key!r}: expected {expected_version}, got {current.version}"
            )
        updated = entry.model_copy(update={"version": expected_version + 1})
        self._client.set(key, updated.model_dump_json())
        return updated

    def delete(self, key: str) -> None:
        self._client.delete(key)

    def list_keys(self) -> list[str]:
        keys = self._client.keys("*")
        return [k.decode("utf-8") if isinstance(k, bytes) else k for k in (keys or [])]

    def is_fresh(self, entry: BaseModel) -> bool:
        age = time.time() - entry.timestamp
        return age < entry.ttl_seconds

    def snapshot(self) -> dict:
        result = {}
        for key in self.list_keys():
            raw = self._client.get(key)
            if raw is None:
                continue
            if isinstance(raw, bytes):
                raw = raw.decode("utf-8")
            entry = self._deserialize(key, raw)
            if entry is not None:
                result[key] = entry.model_dump()
        return result

    def _deserialize(self, key: str, raw: str) -> BaseModel | None:
        try:
            data = _json.loads(raw)
        except (_json.JSONDecodeError, TypeError):
            return None
        if key.startswith("shipment:"):
            return ShipmentStatus.model_validate(data)
        if key.startswith("evt-"):
            return EventLog.model_validate(data)
        for model_cls in (ShipmentStatus, EventLog):
            try:
                return model_cls.model_validate(data)
            except Exception:
                continue
        return None


memory = SharedEpistemicMemory(redis_client)
print("✓ SharedEpistemicMemory initialized (Upstash Redis backend)")

from langchain.tools import tool
import uuid


@tool
def log_event(
    event_type: str,
    shipment_id: str,
    details: str,
    source_agent_id: str,
) -> str:
    """Log a discrete event to the shared memory.

    Args:
        event_type: One of 'disruption_detected', 'reroute_planned',
                    'customer_notified', 'delivery_confirmed'.
        shipment_id: The shipment this event relates to.
        details: Free-text description of what happened.
        source_agent_id: ID of the agent logging this event.

    Returns:
        The event_id of the newly created event.
    """
    event = EventLog(
        event_id=f"evt-{uuid.uuid4().hex[:8]}",
        event_type=event_type,
        shipment_id=shipment_id,
        details=details,
        source_agent_id=source_agent_id,
    )
    memory.write(event.event_id, event)
    return f"Logged event {event.event_id} ({event_type}) for shipment {shipment_id}"


@tool
def update_shipment_status(
    shipment_id: str,
    status: str,
    reason: str,
    source_agent_id: str,
) -> str:
    """Update a shipment's status in the shared memory.

    Args:
        shipment_id: The shipment to update.
        status: One of 'in_transit', 'delayed', 'rerouted', 'delivered', 'lost'.
        reason: Why the status changed.
        source_agent_id: ID of the agent making the update.

    Returns:
        Confirmation message with the new version number.
    """
    key = f"shipment:{shipment_id}"
    existing = memory.read(key)

    if existing is None:
        entry = ShipmentStatus(
            shipment_id=shipment_id,
            status=status,
            reason=reason,
            source_agent_id=source_agent_id,
        )
        memory.write(key, entry)
        return f"Created shipment {shipment_id} with status={status} (v{entry.version})"

    updated = memory.update(
        key,
        existing.model_copy(update={"status": status, "reason": reason, "source_agent_id": source_agent_id}),
        expected_version=existing.version,
    )
    return f"Updated shipment {shipment_id} to status={status} (v{updated.version})"


@tool
def read_memory(key: str) -> str:
    """Read an entry from the shared memory by its key.

    Args:
        key: The memory key (e.g. 'shipment:SHP-001' or 'evt-abc12345').

    Returns:
        JSON string of the entry, or 'NOT_FOUND' if missing/stale.
    """
    entry = memory.read(key)
    if entry is None:
        return "NOT_FOUND"
    return entry.model_dump_json()


print("✓ Tools defined: log_event, update_shipment_status, read_memory")

from langchain.agents import create_agent


monitoring_agent = create_agent(
    llm,
    tools=[log_event],
    system_prompt=(
        "You are the MonitoringAgent. Your job is to detect supply chain disruptions "
        "(storms, port closures, road blocks) and log them as events. "
        "Use the log_event tool with event_type='disruption_detected'. "
        "Always pass source_agent_id='monitoring-agent'. "
        "Do NOT update shipment statuses — that's the LogisticsAgent's job."
    ),
)

logistics_agent = create_agent(
    llm,
    tools=[read_memory, update_shipment_status],
    system_prompt=(
        "You are the LogisticsAgent. Your job is to read disruption events from "
        "shared memory and update affected shipment statuses. "
        "First call read_memory to discover recent events, then call "
        "update_shipment_status with status='delayed' or 'rerouted'. "
        "Always pass source_agent_id='logistics-agent'. "
        "Do NOT notify customers — that's the CustomerNotificationAgent's job."
    ),
)

customer_agent = create_agent(
    llm,
    tools=[read_memory, log_event],
    system_prompt=(
        "You are the CustomerNotificationAgent. Your job is to read shipment "
        "statuses from shared memory and log customer notification events for "
        "any shipment that is delayed, rerouted, or lost. "
        "First call read_memory to discover shipment statuses, then call "
        "log_event with event_type='customer_notified'. "
        "Always pass source_agent_id='customer-notification-agent'."
    ),
)

print("✓ Three agents created:")
print("  - MonitoringAgent (logs disruptions)")
print("  - LogisticsAgent (updates shipment statuses)")
print("  - CustomerNotificationAgent (logs notifications)")

import json


def run_agent(agent, prompt: str, label: str):
    print(f"\n{'='*60}")
    print(f"  {label}")
    print(f"{'='*60}")
    result = agent.invoke(
        {"messages": [{"role": "user", "content": prompt}]},
    )
    final_msg = result["messages"][-1]
    print(f"\n{label} response:\n  {final_msg.content}")
    return result


run_agent(
    monitoring_agent,
    "A severe storm has closed the I-80 corridor in Nebraska. Shipment SHP-001 "
    "is currently in transit on this route. Log this disruption.",
    "MonitoringAgent",
)

run_agent(
    logistics_agent,
    "Check shared memory for recent disruption events. If any shipment is "
    "affected, update its status accordingly.",
    "LogisticsAgent",
)

run_agent(
    customer_agent,
    "Check shared memory for shipment statuses. For any shipment that is "
    "delayed, rerouted, or lost, log a customer_notified event.",
    "CustomerNotificationAgent",
)

print(f"\n{'='*60}")
print("  FINAL SHARED MEMORY STATE")
print(f"{'='*60}")
print(json.dumps(memory.snapshot(), indent=2, default=str))

short_entry = EventLog(
    event_id="evt-stale-test",
    event_type="disruption_detected",
    shipment_id="SHP-002",
    details="Test entry with 2s TTL",
    source_agent_id="test",
    ttl_seconds=2,
)
memory.write("evt-stale-test", short_entry)

fresh = memory.read("evt-stale-test")
print(f"Immediate read: {'FOUND' if fresh else 'NOT_FOUND'} (expected: FOUND)")

print("Sleeping 3 seconds...")
time.sleep(3)

stale = memory.read("evt-stale-test")
print(f"After 3s read:   {'FOUND' if stale else 'NOT_FOUND'} (expected: NOT_FOUND)")

print("\n✓ Staleness tracking works — stale entries are treated as missing.")