server.py

from fastapi import FastAPI, HTTPException, File, UploadFile
from pydantic import BaseModel
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
import io
import os
import pypdf
import traceback
from core.agent import NyayaAgent
from tools.search_tool import search_indian_kanoon, get_act_details, search_legal_documents
from tools.legal_tools import draft_rti_query, draft_legal_notice, analyze_case_summary
from tools.compliance import audit_privacy_policy
from tools.law_bridge import ipc_to_bns_lookup
from google.genai import types

# Initialize Agent and App
agent = NyayaAgent()
app = FastAPI(title="Nyaya Agent")

# CORS
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

class ChatRequest(BaseModel):
    message: str
    persona: str = "explorer"

# API ROUTES MUST COME BEFORE STATIC MOUNT
@app.post("/api/chat")
async def chat(request: ChatRequest):
    try:
        custom_instruction = f"{agent.system_instruction}\n\nCURRENT USER PERSONA: {request.persona.upper()}. "
        if request.persona == "lawyer":
            custom_instruction += "Provide highly technical legal analysis with specific sections and procedural details."
        elif request.persona == "student":
            custom_instruction += "Provide academic analysis, citing landmark cases and legal theory."
        else:
            custom_instruction += "Use simple language, avoid jargon, and explain basic legal concepts clearly."
            
        import datetime
        current_date_str = datetime.date.today().strftime("%B %d, %Y")
        custom_instruction += f"\n\nCURRENT DATE: {current_date_str}. When drafting any document, use this exact date for any 'Date:' fields UNLESS the user explicitly provides a different date in their instructions."
            
        custom_instruction += "\n\nCRITICAL: If the user asks to draft a document (Legal Notice, RTI, etc.), you MUST FIRST check if they have provided all necessary details. IF NOT, you MUST NOT generate the draft immediately. Instead, you MUST dictate the exact fields the UI form should render by replying EXACTLY with the marker '===REQUEST_NOTICE_FORM===' followed IMMEDIATELY by a JSON array of strings representing the needed fields. Provide absolutely no other text. Example: ===REQUEST_NOTICE_FORM===\n[\"Sender's Details (Name, Address)\", \"Receiver's Details (Name, Address)\", \"Marriage Details (Date, Place)\", \"Grounds for Divorce\", \"Relief Sought\"]\n\nONLY IF the user has already provided the details (or if they are submitting the form), you MUST output the actual full draft text in your response. Do NOT just give procedural advice. You MUST wrap the drafted text EXACTLY between the markers '===DRAFT BEGIN===' and '===DRAFT END==='.\n\nIF DRAFTING A LEGAL NOTICE: You MUST strictly follow this exact standard Indian format and use DOUBLE NEWLINES between paragraphs:\n\n===DRAFT BEGIN===\nREGISTERED A.D. / SPEED POST\n\nDate: [Date]\n\nTo,\n[Receiver Name]\n[Receiver Address]\n\nSubject: Legal Notice for [Subject]\n\nDear Sir/Madam,\n\nUnder instructions from and on behalf of my client [Client Name], resident of [Client Address], I hereby serve upon you the following legal notice:\n\n1. [State facts sequentially...]\n\n2. [State legal basis...]\n\nI therefore call upon you to [Specific Relief] within 15 days of receipt of this notice, failing which my client shall be constrained to initiate appropriate civil/criminal proceedings against you at your cost and risk.\n\nYours faithfully,\n\n[Advocate Name/Sender]\n===DRAFT END===\n\nIF DRAFTING AN RTI: You MUST strictly follow this exact standard format and use DOUBLE NEWLINES:\n\n===DRAFT BEGIN===\nTo,\nThe Public Information Officer,\n[Department Name]\n[Address]\n\nSubject: Application for information under Section 6(1) of the Right to Information Act, 2005.\n\nSir/Madam,\n\nI, [Sender Name], am a citizen of India and I request you to provide the following information:\n\n1. [Specific Query 1]\n2. [Specific Query 2]\n\nThe period to which the information relates: [Time Period]\n\nI have attached the application fee of Rs. 10/- vide [IPO/DD/Online Receipt Details].\n\nSincerely,\n\n[Sender Name]\n[Sender Address]\n===DRAFT END==="



        tools = [
            search_indian_kanoon, 
            get_act_details, 
            search_legal_documents,
            draft_rti_query,
            draft_legal_notice,
            analyze_case_summary,
            audit_privacy_policy,
            ipc_to_bns_lookup
        ]
        
        config = types.GenerateContentConfig(
            system_instruction=custom_instruction,
            tools=tools,
            temperature=0.2,
        )
        
        response_text = agent.generate_response(request.message, config=config)
        return {"response": response_text}
    except Exception as e:
        print(f"ERROR IN CHAT: {str(e)}", flush=True)
        traceback.print_exc()
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/api/upload")
async def upload_file(file: UploadFile = File(...)):
    try:
        content = await file.read()
        text = ""
        if file.filename.endswith(".pdf"):
            reader = pypdf.PdfReader(io.BytesIO(content))
            for page in reader.pages:
                text += page.extract_text() + "\n"
        else:
            text = content.decode("utf-8")
        return {"filename": file.filename, "text": text}
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error processing file: {str(e)}")

@app.get("/api/status")
async def status():
    return {"status": "Satyameva Jayate - Online"}

# MOUNT STATIC FILES LAST
app.mount("/", StaticFiles(directory="frontend", html=True), name="frontend")

if __name__ == "__main__":
    import uvicorn
    # Hugging Face uses 7860
    port = int(os.getenv("PORT", 7860))
    uvicorn.run(app, host="0.0.0.0", port=port)