app.py

import os
import uuid
import logging
from datetime import datetime, timedelta
from contextlib import asynccontextmanager

from fastapi import FastAPI, HTTPException, Request, Depends, Response, Cookie
from fastapi.responses import HTMLResponse, FileResponse, JSONResponse
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from pydantic_settings import BaseSettings
from dotenv import load_dotenv

from upstash_redis.asyncio import Redis
from slowapi import Limiter
from slowapi.errors import RateLimitExceeded
from slowapi.util import get_remote_address
from slowapi.middleware import SlowAPIMiddleware

from openai import OpenAI
from langchain_community.embeddings import OpenAIEmbeddings
from langchain_community.vectorstores import Chroma
from langchain_community.chat_models import ChatOpenAI
from langchain.chains import LLMChain
from langchain.prompts import PromptTemplate

# ─── SETTINGS ────────────────────────────────────────────────────────────────────
class Settings(BaseSettings):
    OPENAI_API_KEY: str
    UPSTASH_REDIS_REST_URL: str
    UPSTASH_REDIS_REST_TOKEN: str
    VECTOR_DB_PATH: str = "./chroma_db"
    TOP_K: int = 5
    SESSION_TIMEOUT_MIN: int = 30
    RATE_LIMIT: str = "60/minute"

    class Config:
        env_file = ".env"
        extra = "ignore"  # Add this line to ignore extra variables

settings = Settings()
load_dotenv()

# ─── LOGGING ─────────────────────────────────────────────────────────────────────
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s %(levelname)s %(name)s %(message)s'
)
logger = logging.getLogger("legal-bot")

# ─── LIFESPAN MANAGEMENT ─────────────────────────────────────────────────────────
@asynccontextmanager
async def lifespan(app: FastAPI):
    global redis
    redis = Redis(
        url=settings.UPSTASH_REDIS_REST_URL,
        token=settings.UPSTASH_REDIS_REST_TOKEN
    )
    logger.info("Upstash Redis connection established")
    yield
    await redis.close()
    logger.info("Upstash Redis connection closed")

# ─── FASTAPI APP ────────────────────────────────────────────────────────────────
app = FastAPI(
    title="Irish Legal AI Bot",
    description="RAG‑driven Irish legal assistant",
    lifespan=lifespan
)

# CORS
app.add_middleware(
    CORSMiddleware,
    allow_origins=["http://localhost:8000"],
    allow_methods=["GET", "POST"],
    allow_headers=["*"],
    allow_credentials=True,
)

# Rate limiting
limiter = Limiter(key_func=get_remote_address)
app.state.limiter = limiter
app.add_middleware(SlowAPIMiddleware)

# ─── SECURITY & MODERATION ───────────────────────────────────────────────────────
openai_client = OpenAI(api_key=settings.OPENAI_API_KEY)

async def moderate_content(text: str) -> bool:
    try:
        resp = openai_client.moderations.create(input=text)
        return not resp.results[0].flagged
    except Exception as e:
        logger.error(f"Moderation error: {e}")
        return False

# ─── SESSION MANAGEMENT ──────────────────────────────────────────────────────────
class SessionData(BaseModel):
    session_id: str
    created_at: datetime
    last_activity: datetime
    history: list

async def get_session(session_id: str = Cookie(default=None), response: Response = None) -> SessionData:
    if session_id:
        raw = await redis.get(session_id)
        if raw:
            data = SessionData.parse_raw(raw)
            # Update last activity
            data.last_activity = datetime.utcnow()
            await redis.setex(
                session_id, 
                settings.SESSION_TIMEOUT_MIN * 60, 
                data.json()
            )
            return data
    
    # Create new session
    new_id = str(uuid.uuid4())
    data = SessionData(
        session_id=new_id,
        created_at=datetime.utcnow(),
        last_activity=datetime.utcnow(),
        history=[]
    )
    await redis.setex(
        new_id,
        settings.SESSION_TIMEOUT_MIN * 60,
        data.json()
    )
    response.set_cookie(key="session_id", value=new_id, httponly=True, secure=True)
    return data

# ─── VECTOR & LLM SETUP ─────────────────────────────────────────────────────────
embeddings = OpenAIEmbeddings(openai_api_key=settings.OPENAI_API_KEY)
vectordb = Chroma(embedding_function=embeddings, persist_directory=settings.VECTOR_DB_PATH)
LEGAL_PROMPT = PromptTemplate(
    input_variables=["context","question","history"],
    template=(
        "As an Irish legal expert, provide a precise, concise answer using ONLY the context below."
        "\n1. Direct answer (1-2 sentences)\n2. Key legal basis (cite sources)\n3. Practical implications"
        "\n\nContext:\n{context}\n\nHistory:\n{history}\n\nQuestion: {question}\n\nAnswer:" )
)
POLISH_PROMPT = PromptTemplate(
    input_variables=["raw_answer","question"],
    template=(
        "Enhance this Irish legal answer with current figures/fines (2024), recent amendments, and practical next steps."
        " Keep response under 150 words.\n\nOriginal:\n{raw_answer}\n\nQuestion: {question}\n\nEnhanced Answer:" )
)
legal_chain = LLMChain(llm=ChatOpenAI(temperature=0, openai_api_key=settings.OPENAI_API_KEY, model="gpt-4-turbo"), prompt=LEGAL_PROMPT)
polish_chain = LLMChain(llm=ChatOpenAI(temperature=0.3, openai_api_key=settings.OPENAI_API_KEY, model="gpt-4-turbo"), prompt=POLISH_PROMPT)

# ─── HELPERS ───────────────────────────────────────────────────────────────────
def retrieve_context(query: str):
    docs = vectordb.similarity_search_with_score(query, k=settings.TOP_K)
    snippets = [f"[Source {i+1} | Relevance: {score:.2f}] {doc.page_content.strip()}" for i,(doc,score) in enumerate(docs)]
    sources = [f"Source {i+1}" for i in range(len(docs))]
    return "\n\n".join(snippets), sources

# ─── MODELS ─────────────────────────────────────────────────────────────────────
class QueryRequest(BaseModel):
    query: str

class QueryResponse(BaseModel):
    answer: str
    session_id: str
    sources: list

class SessionStatusResponse(BaseModel):
    status: str  # "active", "expired", or "new"
    ttl: int     # seconds until expiration (-2 = expired, -1 = no expiration)
    session_id: str | None
    created_at: datetime | None
    last_activity: datetime | None
    history_count: int | None

# ─── ROUTES ─────────────────────────────────────────────────────────────────────
@app.get("/", response_class=HTMLResponse)
async def root():
    return FileResponse("frontend/index.html")

@app.post("/query", response_model=QueryResponse)
@limiter.limit(settings.RATE_LIMIT)
async def handle_query(
    request: Request,    
    req: QueryRequest,
    session: SessionData = Depends(get_session),
    response: Response = None
):
    if not await moderate_content(req.query):
        raise HTTPException(400, "Content policy violation")

    context, sources = retrieve_context(req.query)
    history = session.history[-3:] if session.history else []

    raw = legal_chain.run({"context": context, "question": req.query, "history": history})
    polished = polish_chain.run({"raw_answer": raw, "question": req.query})
    if not await moderate_content(polished):
        polished = "Restricted content."

    # Update session
    session.history.append({"q": req.query, "a": polished, "timestamp": datetime.utcnow().isoformat()})
    if len(session.history) > 5:
        session.history.pop(0)
    
    # Save with TTL refresh
    await redis.setex(
        session.session_id,
        settings.SESSION_TIMEOUT_MIN * 60,
        session.json()
    )

    return QueryResponse(answer=polished, session_id=session.session_id, sources=sources)

@app.get("/session/status", response_model=SessionStatusResponse)
async def get_session_status(session_id: str = Cookie(default=None)):
    if not session_id:
        return SessionStatusResponse(
            status="new",
            ttl=-2,
            session_id=None,
            created_at=None,
            last_activity=None,
            history_count=None
        )
    
    ttl = await redis.ttl(session_id)
    if ttl < 0:  # Key doesn't exist or has no TTL
        return SessionStatusResponse(
            status="expired",
            ttl=-2,
            session_id=session_id,
            created_at=None,
            last_activity=None,
            history_count=None
        )
    
    raw = await redis.get(session_id)
    if not raw:
        return SessionStatusResponse(
            status="expired",
            ttl=-2,
            session_id=session_id,
            created_at=None,
            last_activity=None,
            history_count=None
        )
    
    data = SessionData.parse_raw(raw)
    return SessionStatusResponse(
        status="active",
        ttl=ttl,
        session_id=session_id,
        created_at=data.created_at,
        last_activity=data.last_activity,
        history_count=len(data.history)
    )

# ─── SERVER LAUNCH ──────────────────────────────────────────────────────────────
if __name__ == "__main__":
    import uvicorn
    uvicorn.run("app:app", host="0.0.0.0", port=7860, log_level="info")