Update main.py

main.py

@@ -1,107 +1,106 @@
-import uvicorn
-from fastapi import FastAPI, HTTPException
-from fastapi.middleware.cors import CORSMiddleware
-from pydantic import BaseModel
-from typing import List, Optional
-from graph import app_graph
-
-app = FastAPI(title="Clause.ai Backend")
-
-# --- CORS SETUP (LOCKED DOWN) ---
-origins = [
-    "https://clause-ai-nbu8.vercel.app"  
-]
-
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=origins,        # Only allow your specific frontend
-    allow_credentials=True,
-    allow_methods=["*"],          
-    allow_headers=["*"],          
-)
-
-# --- API MODELS ---
-class DraftRequest(BaseModel):
-    query: str
-
-class AgentResponse(BaseModel):
-    status: str
-    phase: str
-    message: str
-    missing_info: List[str] = []
-    draft: Optional[str] = None
-    reference: Optional[str] = None
-
-@app.get("/")
-def home():
-    return {"status": "Clause.ai Brain is Online"}
-
-@app.post("/draft", response_model=AgentResponse)
-async def generate_clause(request: DraftRequest):
-    try:
-        # Initialize the state
-        initial_state = {
-            "query": request.query,
-            "messages": [],
-            "context": "",
-            "reference_clause": "",
-            "final_draft": "",
-            "phase": "",
-            "missing_info": [],
-            "clarification_question": ""
-        }
-        
-        # Run the LangGraph Agent
-        result = app_graph.invoke(initial_state)
-        
-        phase = result.get("phase", "")
-        
-        # --- SCENARIO 1: Guardrail stopped (general question/greeting/injection) ---
-        if phase == "stopped":
-            return {
-                "status": "general_response",
-                "phase": "stopped",
-                "message": result.get("final_draft", ""),
-                "missing_info": [],
-                "draft": None,
-                "reference": None
-            }
-        
-        # --- SCENARIO 2: Triage needs clarification ---
-        if phase == "planning":
-            return {
-                "status": "needs_info",
-                "phase": "planning",
-                "message": result.get("clarification_question", "Please provide more details."),
-                "missing_info": result.get("missing_info", []),
-                "draft": None,
-                "reference": None
-            }
-        
-        # --- SCENARIO 3: Draft completed successfully ---
-        if phase == "drafting" or result.get("final_draft"):
-            return {
-                "status": "success",
-                "phase": "drafting",
-                "message": "Draft generated successfully.",
-                "missing_info": [],
-                "draft": result.get("final_draft", ""),
-                "reference": result.get("reference_clause", "")
-            }
-        
-        # --- FALLBACK: Unknown state ---
-        return {
-            "status": "error",
-            "phase": "unknown",
-            "message": "Unable to process your request. Please try again.",
-            "missing_info": [],
-            "draft": None,
-            "reference": None
-        }
-
-    except Exception as e:
-        print(f"❌ Error: {e}")
-        raise HTTPException(status_code=500, detail=str(e))
-
-if __name__ == "__main__":
+import uvicorn
+from fastapi import FastAPI, HTTPException
+from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel
+from typing import List, Optional
+from graph import app_graph
+
+app = FastAPI(title="Clause.ai Backend")
+
+# --- CORS SETUP (LOCKED DOWN) ---
+origins = [
+    "https://clause-ai-nbu8.vercel.app"  
+]
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],  # Allow ANYONE to connect
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# --- API MODELS ---
+class DraftRequest(BaseModel):
+    query: str
+
+class AgentResponse(BaseModel):
+    status: str
+    phase: str
+    message: str
+    missing_info: List[str] = []
+    draft: Optional[str] = None
+    reference: Optional[str] = None
+
+@app.get("/")
+def home():
+    return {"status": "Clause.ai Brain is Online"}
+
+@app.post("/draft", response_model=AgentResponse)
+async def generate_clause(request: DraftRequest):
+    try:
+        # Initialize the state
+        initial_state = {
+            "query": request.query,
+            "messages": [],
+            "context": "",
+            "reference_clause": "",
+            "final_draft": "",
+            "phase": "",
+            "missing_info": [],
+            "clarification_question": ""
+        }
+        
+        # Run the LangGraph Agent
+        result = app_graph.invoke(initial_state)
+        
+        phase = result.get("phase", "")
+        
+        # --- SCENARIO 1: Guardrail stopped (general question/greeting/injection) ---
+        if phase == "stopped":
+            return {
+                "status": "general_response",
+                "phase": "stopped",
+                "message": result.get("final_draft", ""),
+                "missing_info": [],
+                "draft": None,
+                "reference": None
+            }
+        
+        # --- SCENARIO 2: Triage needs clarification ---
+        if phase == "planning":
+            return {
+                "status": "needs_info",
+                "phase": "planning",
+                "message": result.get("clarification_question", "Please provide more details."),
+                "missing_info": result.get("missing_info", []),
+                "draft": None,
+                "reference": None
+            }
+        
+        # --- SCENARIO 3: Draft completed successfully ---
+        if phase == "drafting" or result.get("final_draft"):
+            return {
+                "status": "success",
+                "phase": "drafting",
+                "message": "Draft generated successfully.",
+                "missing_info": [],
+                "draft": result.get("final_draft", ""),
+                "reference": result.get("reference_clause", "")
+            }
+        
+        # --- FALLBACK: Unknown state ---
+        return {
+            "status": "error",
+            "phase": "unknown",
+            "message": "Unable to process your request. Please try again.",
+            "missing_info": [],
+            "draft": None,
+            "reference": None
+        }
+
+    except Exception as e:
+        print(f"❌ Error: {e}")
+        raise HTTPException(status_code=500, detail=str(e))
+
+if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=8000)