Update app.py

app.py

@@ -1,148 +1,179 @@
-from fastapi import FastAPI, Request, Depends, HTTPException
-from fastapi.responses import JSONResponse
-from datetime import datetime
-import uuid
 import os
-from typing import Any, Dict
+import time
+from fastapi import FastAPI, Request, HTTPException, Depends, status
+from fastapi.security import APIKeyHeader
+from typing import Dict, Any
+from datetime import datetime
 
-from models import Message, ExtractedIntelligence
+# LangGraph and Model Imports
+from langgraph.checkpoint.memory import MemorySaver
+from langgraph.checkpoint.base import BaseCheckpointSaver
 from agent import create_honeypot_graph, final_callback
+from models import HoneypotRequest, HoneypotResponse, AgentState, ExtractedIntelligence, Message
+
+# --- Configuration ---
+API_KEY_NAME = "x-api-key"
+API_KEY = os.environ.get("HONEYPOT_API_KEY", "sk_test_123456789") # Default for local testing
+api_key_header = APIKeyHeader(name=API_KEY_NAME, auto_error=False)
+
+# --- Initialization ---
+app = FastAPI(
+    title="Agentic Honey-Pot API",
+    description="REST API for Scam Detection and Intelligence Extraction using LangGraph and Qwen 2.5 3B.",
+    version="1.0.0"
+)
+
+# Initialize LangGraph Checkpointer (Use MemorySaver for simplicity, replace with a database for production)
+# For Hugging Face Spaces, a persistent volume or database is recommended for checkpointer.
+# For this example, we use MemorySaver, which will reset on Space restart.
+checkpointer: BaseCheckpointSaver = MemorySaver()
+honeypot_app = create_honeypot_graph(checkpointer)
+
+# --- Dependency for API Key Validation ---
+
+async def get_api_key(api_key_header: str = Depends(api_key_header)):
+    if api_key_header is None or api_key_header != API_KEY:
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Invalid API Key or missing 'x-api-key' header.",
+        )
+    return api_key_header
 
-app = FastAPI()
-
-# -----------------------------
-# API KEY AUTH
-# -----------------------------
-HONEYPOT_API_KEY = os.environ.get("HONEYPOT_API_KEY", "sk_test_123456789")
-
-def get_api_key(request: Request):
-    api_key = request.headers.get("x-api-key")
-    if not api_key or api_key != HONEYPOT_API_KEY:
-        raise HTTPException(status_code=401, detail="Invalid API Key")
-    return api_key
-
-# -----------------------------
-# INIT GRAPH
-# -----------------------------
-graph = create_honeypot_graph()
-
-# -----------------------------
-# UTILS
-# -----------------------------
-def normalize_timestamp(ts: Any) -> str:
-    """
-    Convert numeric or invalid timestamps into ISO string
-    """
-    if ts is None:
-        return datetime.utcnow().isoformat() + "Z"
-
-    if isinstance(ts, (int, float)):
-        return datetime.utcfromtimestamp(ts / 1000).isoformat() + "Z"
-
-    return str(ts)
+# --- API Endpoints ---
 
-def safe_message_from_payload(payload: Any) -> Message:
+@app.post("/api/honeypot-detection", response_model=HoneypotResponse)
+async def honeypot_detection(
+    request_data: HoneypotRequest,
+    api_key: str = Depends(get_api_key)
+) -> Dict[str, Any]:
     """
-    Convert ANY payload into a valid Message object
+    Accepts an incoming message event, runs the LangGraph agent, and returns the response.
     """
-    if isinstance(payload, dict):
-        sender = payload.get("sender", "scammer")
-        text = str(payload.get("text", "Hello"))
-        timestamp = normalize_timestamp(payload.get("timestamp"))
+    session_id = request_data.sessionId
+    
+    # 1. Load or Initialize State
+    # LangGraph uses the thread_id for checkpointing
+    config = {"configurable": {"thread_id": session_id}}
+    
+    # Check if a checkpoint exists for this session
+    # Get the latest checkpoint using the graph's state loader
+    checkpoint = honeypot_app.get_state(config)
+
+    if checkpoint:
+        # Load existing state and append the new message
+        # checkpoint.values contains the state dictionary saved at the last checkpoint
+        current_state_dict = checkpoint.values if hasattr(checkpoint, "values") else checkpoint.dict["values"]
+        # Add default for new fields if missing
+        if "callbackSent" not in current_state_dict:
+            current_state_dict["callbackSent"] = False
+        if "agentNotes" not in current_state_dict:
+            # Initialize agentNotes if missing to avoid KeyError in state machine
+            current_state_dict["agentNotes"] = ""
+        if "extractedIntelligence" not in current_state_dict:
+            # Initialize extractedIntelligence if missing
+            current_state_dict["extractedIntelligence"] = ExtractedIntelligence()
+        # Ensure conversation history exists
+        if "conversationHistory" not in current_state_dict:
+            current_state_dict["conversationHistory"] = []
+        if "totalMessagesExchanged" not in current_state_dict:
+            current_state_dict["totalMessagesExchanged"] = 0
+        # Ensure sessionId is present
+        if "sessionId" not in current_state_dict:
+            current_state_dict["sessionId"] = session_id
+        current_state = AgentState(**current_state_dict)
+        
+        # Append the new message from the scammer
+        current_state["conversationHistory"].append(request_data.message)
+        current_state["totalMessagesExchanged"] += 1
+        
+        # The LangGraph will continue from the last node
+        input_state = current_state
+        start_time = time.time()
+        
     else:
-        sender = "scammer"
-        text = str(payload)
-        timestamp = normalize_timestamp(None)
-
-    return Message(sender=sender, text=text, timestamp=timestamp)
+        # New conversation: Initialize the state
+        initial_history = request_data.conversationHistory + [request_data.message]
+        
+        input_state = AgentState(
+            sessionId=session_id,
+            conversationHistory=initial_history,
+            scamDetected=False,
+            extractedIntelligence=ExtractedIntelligence(),
+            agentNotes="New session started. ",
+            totalMessagesExchanged=len(initial_history),
+            should_continue_engagement=False,
+            agent_response_text="",
+            callbackSent=False
+        )
+        start_time = time.time()
 
-# -----------------------------
-# MAIN ENDPOINT (TESTER SAFE)
-# -----------------------------
-@app.post("/api/honeypot-detection")
-async def honeypot_detection(
-    request: Request,
-    api_key: str = Depends(get_api_key)
-):
+    # 2. Invoke LangGraph
     try:
-        body = await request.json()
-    except Exception:
-        body = {}
-
-    # SESSION ID
-    session_id = body.get("sessionId") or str(uuid.uuid4())
-
-    # MESSAGE
-    msg_payload = body.get("message", "Hello")
-    new_message = safe_message_from_payload(msg_payload)
-
-    # CONVERSATION HISTORY
-    history_payload = body.get("conversationHistory", [])
-    conversation_history = []
-
-    if isinstance(history_payload, list):
-        for item in history_payload:
-            conversation_history.append(safe_message_from_payload(item))
-
-    conversation_history.append(new_message)
-
-    # INITIAL STATE
-    state: Dict[str, Any] = {
-        "sessionId": session_id,
-        "conversationHistory": conversation_history,
-        "scamDetected": False,
-        "agentNotes": "",
-        "extractedIntelligence": ExtractedIntelligence(),
-        "callbackSent": False,
-    }
+        # Invoke the graph with the updated state
+        final_state_dict = honeypot_app.invoke(input_state, config=config)
+        final_state = AgentState(**final_state_dict)
+        
+        end_time = time.time()
+        engagement_duration = end_time - start_time
+        
+        # 3. Prepare API Response
+        response_data = {
+            "status": "success",
+            "scamDetected": final_state["scamDetected"],
+            "engagementMetrics": {
+                "engagementDurationSeconds": round(engagement_duration, 2),
+                "totalMessagesExchanged": final_state["totalMessagesExchanged"]
+            },
+            "extractedIntelligence": final_state["extractedIntelligence"].model_dump(),
+            "agentNotes": final_state["agentNotes"]
+        }
+        
+        # The LangGraph handles engagement end and final callback internally.
+        
+        return response_data
 
-    # RUN AGENT
-    try:
-        result = graph.invoke(state)
     except Exception as e:
-        return JSONResponse(
-            status_code=200,
-            content={
-                "status": "success",
-                "scamDetected": False,
-                "agentNotes": f"Agent fallback triggered: {str(e)}",
-                "extractedIntelligence": ExtractedIntelligence().model_dump(),
-                "engagementMetrics": {
-                    "totalMessagesExchanged": len(conversation_history),
-                    "engagementDurationSeconds": 0
-                }
-            }
+        print(f"An error occurred during LangGraph invocation: {e}")
+        raise HTTPException(
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            detail=f"Internal server error during agent processing: {str(e)}",
         )
 
-    # RESPONSE (PDF-COMPLIANT)
-    return {
-        "status": "success",
-        "scamDetected": result.get("scamDetected", False),
-        "engagementMetrics": {
-            "totalMessagesExchanged": len(conversation_history),
-            "engagementDurationSeconds": 5
-        },
-        "extractedIntelligence": result.get("extractedIntelligence").model_dump(),
-        "agentNotes": result.get("agentNotes", "")
-    }
-
-# -----------------------------
-# MANUAL CALLBACK TEST
-# -----------------------------
 @app.post("/api/trigger-final-callback")
-async def trigger_callback(
-    session_id: str,
-    api_key: str = Depends(get_api_key)
-):
-    state = {
-        "sessionId": session_id,
-        "scamDetected": True,
-        "extractedIntelligence": ExtractedIntelligence(),
-        "agentNotes": "Manual callback trigger",
-        "callbackSent": False,
-        "totalMessagesExchanged": 1
-    }
-
-    final_callback(state)
-
-    return {"status": "callback triggered"}
+async def trigger_callback(session_id: str, api_key: str = Depends(get_api_key)):
+    """
+    Manually triggers the final result callback for a specific session.
+    This is useful for testing or for an external system to signal the end of engagement.
+    """
+    config = {"configurable": {"thread_id": session_id}}
+    checkpoint = honeypot_app.get_state(config)
+    
+    if not checkpoint:
+        raise HTTPException(status_code=404, detail=f"Session ID {session_id} not found.")
+        
+    current_state_dict = checkpoint.values if hasattr(checkpoint, "values") else checkpoint.dict["values"]
+    # Ensure mandatory fields exist for older checkpoints
+    if "callbackSent" not in current_state_dict:
+        current_state_dict["callbackSent"] = False
+    if "agentNotes" not in current_state_dict:
+        current_state_dict["agentNotes"] = ""
+    if "extractedIntelligence" not in current_state_dict:
+        current_state_dict["extractedIntelligence"] = ExtractedIntelligence()
+    if "conversationHistory" not in current_state_dict:
+        current_state_dict["conversationHistory"] = []
+    if "totalMessagesExchanged" not in current_state_dict:
+        current_state_dict["totalMessagesExchanged"] = len(current_state_dict.get("conversationHistory", []))
+    # Ensure sessionId is present
+    if "sessionId" not in current_state_dict:
+        current_state_dict["sessionId"] = session_id
+    current_state = AgentState(**current_state_dict)
+    
+    # Trigger the final callback if not already sent
+    if not current_state.get("callbackSent", False):
+        final_callback(current_state)
+    
+    return {"status": "success", "message": f"Final callback triggered for session {session_id}."}
+
+@app.get("/")
+async def root():
+    return {"message": "Agentic Honey-Pot API is running. Use /api/honeypot-detection endpoint."}