import json import logging import time import uuid from fastapi import FastAPI, File, UploadFile, Form from pydantic import BaseModel from typing import List, Any, Optional from pathlib import Path from app.routing.router import route_query from app.generation.context_builder import build_context # ── Structured JSON logging setup ──────────────────────────────────────────── # Each log line is a JSON object — queryable in CloudWatch Logs Insights. # Fields: timestamp, level, logger, message, plus any extra kwargs passed # to logger.info(..., extra={...}). class _JsonFormatter(logging.Formatter): def format(self, record: logging.LogRecord) -> str: payload = { "timestamp": self.formatTime(record, "%Y-%m-%dT%H:%M:%S"), "level": record.levelname, "logger": record.name, "message": record.getMessage(), } # Attach any extra fields (query_id, latency_ms, cache_hit, etc.) for key, val in record.__dict__.items(): if key not in ( "name", "msg", "args", "levelname", "levelno", "pathname", "filename", "module", "exc_info", "exc_text", "stack_info", "lineno", "funcName", "created", "msecs", "relativeCreated", "thread", "threadName", "processName", "process", "message", "taskName", ): payload[key] = val return json.dumps(payload, default=str) def _configure_json_logging(): formatter = _JsonFormatter() root = logging.getLogger() if not root.handlers: handler = logging.StreamHandler() handler.setFormatter(formatter) root.addHandler(handler) else: for h in root.handlers: h.setFormatter(formatter) root.setLevel(logging.INFO) _configure_json_logging() logger = logging.getLogger(__name__) # ---------- App ---------- app = FastAPI( title="GST Legal RAG API", version="1.0", description="In-house GST knowledge assistant", ) # Readiness gate: ALB health check returns 503 until models are loaded. # Prevents live traffic from hitting a task that isn't ready to serve yet. _warmup_complete: bool = False @app.on_event("startup") async def _startup_tasks(): """Seed feed store and pre-warm AI models in background.""" import asyncio as _asyncio # 1. Seed in-memory event log try: from app.api.documents import get_activity_feed from app.feed_store import _event_log items = get_activity_feed() for item in reversed(items): _event_log.append(item) logger.info(f"Feed store seeded with {len(items)} events from filesystem") except Exception as e: logger.warning(f"Feed store seed failed (non-fatal): {e}") # 2. Pre-load embedding model + FAISS index so the first /ask-sync request # doesn't pay the ~15-20s cold-start cost and timeout at API Gateway. async def _warmup(): global _warmup_complete try: from app.dependencies import preload_all_models await _asyncio.to_thread(preload_all_models) logger.info("Startup model warmup complete") except Exception as e: logger.warning(f"Startup warmup failed (non-fatal): {e}") finally: _warmup_complete = True # always unblock health check, even on error _asyncio.ensure_future(_warmup()) @app.get("/") async def root(): return { "message": "GST Legal RAG API is running.", "docs": "/docs", "health": "/api/health" } from fastapi.middleware.cors import CORSMiddleware import os # CORS: Use ALLOWED_ORIGINS env var in production (comma-separated), # falls back to localhost dev origins when not set. _default_origins = "http://localhost:5173,http://127.0.0.1:5173,http://localhost:5174,http://127.0.0.1:5174,https://gst-rag-95li.vercel.app,https://main.d1q7i80dk455hq.amplifyapp.com,https://letatec.com,https://www.letatec.com" ALLOWED_ORIGINS = [ origin.strip() for origin in os.getenv("ALLOWED_ORIGINS", _default_origins).split(",") if origin.strip() ] app.add_middleware( CORSMiddleware, allow_origins=ALLOWED_ORIGINS, allow_credentials=True, allow_methods=["GET", "POST", "OPTIONS"], allow_headers=["Content-Type", "Authorization", "X-Request-ID"], expose_headers=["X-Request-ID"], max_age=600, ) # ── Security headers middleware ─────────────────────────────────────────────── from starlette.middleware.base import BaseHTTPMiddleware from starlette.requests import Request as StarletteRequest class SecurityHeadersMiddleware(BaseHTTPMiddleware): async def dispatch(self, request: StarletteRequest, call_next): response = await call_next(request) # Prevent browsers from MIME-sniffing the content type response.headers["X-Content-Type-Options"] = "nosniff" # Block clickjacking — skip for document view so the frontend iframe can embed PDFs if not request.url.path.startswith("/api/documents/view"): response.headers["X-Frame-Options"] = "DENY" # Force HTTPS for 1 year, include subdomains response.headers["Strict-Transport-Security"] = "max-age=31536000; includeSubDomains; preload" # Only send origin in Referer header, no full URL response.headers["Referrer-Policy"] = "strict-origin-when-cross-origin" # Disable browser features not needed by the API response.headers["Permissions-Policy"] = "geolocation=(), camera=(), microphone=(), payment=()" # Prevent caching of API responses (contains legal/confidential data) if not request.url.path.startswith("/api/documents/view"): response.headers["Cache-Control"] = "no-store, no-cache, must-revalidate, private" # Remove server fingerprinting headers for h in ("server", "x-powered-by"): if h in response.headers: del response.headers[h] return response app.add_middleware(SecurityHeadersMiddleware) # ── Request size limit middleware ───────────────────────────────────────────── class RequestSizeLimitMiddleware(BaseHTTPMiddleware): MAX_BODY_BYTES = 5 * 1024 * 1024 # 5 MB async def dispatch(self, request: StarletteRequest, call_next): if request.method in ("POST", "PUT", "PATCH"): content_length = request.headers.get("content-length") if content_length and int(content_length) > self.MAX_BODY_BYTES: from starlette.responses import JSONResponse return JSONResponse({"detail": "Request body too large"}, status_code=413) return await call_next(request) app.add_middleware(RequestSizeLimitMiddleware) from fastapi import Request # ── Request logging middleware ──────────────────────────────────────────────── # Emits one JSON log line per request with fields queryable in CloudWatch: # query_id, method, path, status_code, latency_ms # The /ask endpoint adds extra fields (cache_hit, domain, etc.) via its own logs. _request_logger = logging.getLogger("leta.request") @app.middleware("http") async def log_requests(request: Request, call_next): # Full query_id: LETA-YYYYMMDD-HHMMSS-XXXXXXXX (CloudWatch-searchable prefix) import datetime as _dt_mw _ts = _dt_mw.datetime.now().strftime("%Y%m%d-%H%M%S") _rand = str(uuid.uuid4()).replace("-", "")[:6].upper() query_id = f"LETA-{_ts}-{_rand}" request.state.query_id = query_id t0 = time.monotonic() response = await call_next(request) latency_ms = round((time.monotonic() - t0) * 1000, 1) _request_logger.info( f"{request.method} {request.url.path} {response.status_code}", extra={ "query_id": query_id, "method": request.method, "path": request.url.path, "status_code": response.status_code, "latency_ms": latency_ms, "client_ip": request.client.host if request.client else "unknown", }, ) response.headers["X-Request-ID"] = query_id return response # Rate limiting — user-ID keyed, Redis-backed when available import base64 as _b64 import json as _json_ratelimit from slowapi import Limiter, _rate_limit_exceeded_handler from slowapi.util import get_remote_address from slowapi.errors import RateLimitExceeded def _rate_limit_key(request: Request) -> str: """Prefer JWT user-ID over IP so limits survive proxy/CDN hops.""" auth = request.headers.get("Authorization", "") if auth.startswith("Bearer "): try: parts = auth.split(".") if len(parts) == 3: padded = parts[1] + "=" * (-len(parts[1]) % 4) payload = _json_ratelimit.loads(_b64.b64decode(padded)) if "sub" in payload: return f"user:{payload['sub']}" except Exception: pass forwarded = request.headers.get("X-Forwarded-For", "") if forwarded: return forwarded.split(",")[0].strip() return get_remote_address(request) def _build_limiter() -> Limiter: _redis_url = os.getenv("REDIS_URL", "redis://localhost:6379/0") try: import redis as _r _r.from_url(_redis_url, socket_connect_timeout=1).ping() lim = Limiter(key_func=_rate_limit_key, storage_uri=_redis_url) logger.info("Rate limiter: Redis-backed (distributed)") return lim except Exception as _e: logger.warning(f"Rate limiter: in-memory fallback (Redis unavailable: {_e})") return Limiter(key_func=_rate_limit_key) limiter = _build_limiter() app.state.limiter = limiter app.add_exception_handler(RateLimitExceeded, _rate_limit_exceeded_handler) from app.api import documents app.include_router(documents.router, prefix="/api/documents", tags=["Documents"]) from app.api import advisory app.include_router(advisory.router, prefix="/api/advisory", tags=["Advisory"]) from app.api import sessions app.include_router(sessions.router, prefix="/api/sessions", tags=["Sessions"]) from app.api import auth app.include_router(auth.router, prefix="/api/auth", tags=["Auth"]) from app.api import templates app.include_router(templates.router, prefix="/api/templates", tags=["Templates"]) from app.api import admin app.include_router(admin.router, prefix="/api/admin", tags=["Admin"]) from app.api import feed app.include_router(feed.router, prefix="/api/feed", tags=["Feed"]) from app.api import transcribe app.include_router(transcribe.router, prefix="/api", tags=["Transcribe"]) from app.api import payments app.include_router(payments.router, tags=["Payments"]) from app.api import debug as _debug_api app.include_router(_debug_api.router, tags=["Debug"]) # ---------- Request / Response ---------- class QuestionRequest(BaseModel): question: str session_id: Optional[str] = None intent: Optional[str] = "general" class Source(BaseModel): source: str page: int class AnswerResponse(BaseModel): answer: str confidence: float intent: str sources: List[Source] reasoning: Optional[Any] = None # ---------- Lazy Load Retriever ---------- from app.dependencies import get_retriever from app.database import get_session_collection from datetime import datetime # ---------- Endpoint ---------- @app.get("/api/health") async def health_check(): from starlette.responses import JSONResponse # Return 503 while models are still loading so ALB withholds live traffic # until this task is genuinely ready to serve requests. if not _warmup_complete: return JSONResponse( status_code=503, content={ "status": "warming_up", "timestamp": datetime.now().isoformat(), "message": "Model warmup in progress — retry in 30s", }, headers={"Retry-After": "30"}, ) health_status = { "status": "active", "timestamp": datetime.now().isoformat(), "systems": { "api": "ok", "database": "unknown", "retriever": "unknown" } } # Check Database try: from app.database import get_db db = get_db() if db is not None: # Simple ping db.command('ping') health_status["systems"]["database"] = "connected" except Exception as e: health_status["systems"]["database"] = f"error: {str(e)}" health_status["status"] = "degraded" # Check Retriever try: retriever = get_retriever() if retriever is not None: health_status["systems"]["retriever"] = "initialized" except Exception as e: health_status["systems"]["retriever"] = f"error: {str(e)}" health_status["status"] = "degraded" return health_status async def stream_and_save(generator, session_id, user_query, chunks=None, context="", truth_rules_text="", query_vec=None, is_draft_session: bool = False): """ Wrapper that streams the LLM response AND runs post-generation accuracy layers before saving to DB. Post-stream pipeline: 1. Citation Validator — cross-checks cited sections against chunks 2. Answer Verifier — LLM second-pass for logical consistency 3. Hallucination Guard — flags ungrounded numbers 4. Template Matcher — surfaces relevant litigation templates """ import logging _logger = logging.getLogger("stream_and_save") import asyncio import json full_answer = "" unique_sources = [] # 1. Immediately emit metadata about retrieved sources (Perplexity-style transparency) if chunks: try: unique_sources = [] seen_src = set() import urllib.parse as _urlparse for c in chunks: _dedup_rel = c.get("rel_path", "") or c.get("metadata", {}).get("rel_path", "") or c.get("source", "") src_key = (_dedup_rel, c.get("page", 0)) if src_key not in seen_src: seen_src.add(src_key) _raw_src = c.get("source", "") or c.get("metadata", {}).get("source", "") _rel_path = c.get("rel_path", "") or c.get("metadata", {}).get("rel_path", "") # Use rel_path for display name — it holds the real filename. _basename = os.path.basename(_rel_path) if _rel_path else os.path.basename(_raw_src) _enc_path = _urlparse.quote(_rel_path.replace("\\", "/"), safe="") if _rel_path else "" _enc_name = _urlparse.quote(_basename, safe="") _url = ( f"/api/documents/view_by_path?path={_enc_path}" if _enc_path else f"/api/documents/view?category=all&filename={_enc_name}" ) _snippet = (c.get("text") or "").strip() unique_sources.append({ "title": _basename or "Document", "page": c.get("page", 1), "url": _url, "rel_path": _rel_path, "score": float(c.get("_rerank_score", 0)), "snippet": _snippet[:800] if _snippet else "", }) if len(unique_sources) >= 20: # collect more, then sort + cap break # Sort by relevance score — the reranker already encodes legal # authority (30% weight) + semantic similarity (50%) + topic (20%). # Most relevant doc wins regardless of type; linkifyLegalRefs no # longer uses sources[0] as a fallback so AAR-as-first-link is gone. unique_sources.sort(key=lambda s: s.get("score", 0), reverse=True) unique_sources = unique_sources[:8] metadata_payload = { "type": "metadata", "sources": unique_sources } yield f"__METADATA__:{json.dumps(metadata_payload)}__END_METADATA__" except Exception as me: _logger.warning(f"Metadata emission failed: {me}") try: for chunk in generator: full_answer += chunk yield chunk if chunks and full_answer.strip(): # Use the session-level draft flag passed in — do not re-detect from # user_query alone, which breaks on follow-up/correction messages. is_draft = is_draft_session or any(kw in user_query.lower() for kw in [ "draft","notice","reply","appeal","submission","advisory","scn", "show cause","drc-01","drc 01","asmt-10","our understanding", "gst implications","provide opinion","our comments","tax position", "advise on","legal opinion","our client is","we are engaged in", ]) if not is_draft: # --- Parallel Accuracy Pipeline --- async def run_citation_validator(): try: from app.generation.citation_validator import CitationValidator return CitationValidator.validate_citations(full_answer, chunks) except Exception as e: _logger.warning(f"Citation validator error: {e}") return full_answer async def run_answer_verifier(): try: from app.generation.answer_verifier import verify_answer return await asyncio.to_thread(verify_answer, user_query, full_answer, chunks) except Exception as e: _logger.warning(f"Answer verifier error: {e}") return None async def run_authority_verifier(): """Priority 10: Answer Verification Agent. Checks if the LLM answer cited every mandatory governing authority predicted by the authority taxonomy. Runs in ~1ms (string matching, no LLM call). Logs gaps for monitoring and future regression tracking.""" try: from app.dependencies import get_retriever from app.retrieval.query_refiner import verify_answer_authority_coverage _ret = get_retriever() _tax = getattr(_ret, "_last_taxonomy", {}) _cov = getattr(_ret, "_last_coverage", {}) if _tax.get("confidence", 0) > 0 and (_tax.get("sections") or _tax.get("circulars")): _av = verify_answer_authority_coverage(user_query, full_answer, _tax, _cov) if _av["verdict"] != "pass": _logger.warning( f"[AUTHORITY_VERIFY] verdict={_av['verdict']} | " f"topics={_tax.get('topics')} | " f"cited={_av['cited']} | " f"missing={_av['missing']} | " f"note={_av['note']}" ) return None except Exception as e: _logger.warning(f"Authority verifier error: {e}") return None async def run_hallucination_guard(): try: from app.generation.hallucination_guard import check_hallucinated_numbers return check_hallucinated_numbers(full_answer, context, truth_rules_text, chunks) except Exception as e: _logger.warning(f"Hallucination guard error: {e}") return "" async def run_template_matcher(): try: from app.retrieval.template_matcher import search_templates, format_template_suggestions matched = await asyncio.to_thread(search_templates, user_query, top_k=3) return format_template_suggestions(matched) except Exception as e: _logger.warning(f"Template matcher error: {e}") return "" try: results = await asyncio.wait_for( asyncio.gather( run_citation_validator(), run_hallucination_guard(), run_authority_verifier(), ), timeout=4.0, ) except asyncio.TimeoutError: _logger.warning("Post-generation validators timed out — skipping") results = (full_answer, "", None) _, hallu_warn, _ = results # Log internally — validator output is NEVER streamed to the user. # Citation validator already logs its own warnings. if hallu_warn: _logger.warning(f"Hallucination guard: {hallu_warn[:300]}") # template_block intentionally not streamed — surfaced via /api/templates instead except Exception as e: _logger.error(f"Error in stream_and_save: {e}", exc_info=True) yield "\n[System: An error occurred processing your request. Please try again.]" finally: # DB save — wrapped in to_thread so the sync pymongo write doesn't # block the event loop while the HTTP response is still open. if session_id and full_answer.strip(): collection = get_session_collection() if collection is not None: try: await asyncio.to_thread( collection.update_one, {"session_id": session_id}, { "$push": {"messages": { "role": "assistant", "content": full_answer, "sources": unique_sources, "timestamp": datetime.now(), }}, "$set": {"updated_at": datetime.now()}, }, ) except Exception as dbe: _logger.warning(f"DB save error (non-fatal): {dbe}") # Confidence-gated cache store — also off the event loop if full_answer.strip() and query_vec is not None: try: from app.cache import cache_store from app.generation.confidence import estimate_confidence confidence = await asyncio.to_thread(estimate_confidence, full_answer, chunks or []) await asyncio.to_thread(cache_store, user_query, query_vec, full_answer, confidence, unique_sources if chunks else []) _logger.debug(f"Cache store attempt | confidence={confidence:.2f}") except Exception as ce: _logger.warning(f"Cache store error (non-fatal): {ce}") _ask_logger = logging.getLogger("leta.ask") @app.post("/ask") @limiter.limit("30/minute") async def ask_question(request: Request, req: QuestionRequest): if not _warmup_complete: from starlette.responses import JSONResponse as _jr return _jr(status_code=503, content={"detail": "Service warming up — retry in 30s"}, headers={"Retry-After": "30"}) question = req.question.strip() session_id = req.session_id query_id = getattr(request.state, "query_id", str(uuid.uuid4())[:8]) t0 = time.monotonic() # IMMEDIATE SAVE: Save User Question First if session_id: collection = get_session_collection() if collection is not None: collection.update_one( {"session_id": session_id}, {"$push": {"messages": {"role": "user", "content": question, "timestamp": datetime.now()}}} ) async def rag_pipeline_orchestrator(): import asyncio as _asyncio from app.routing.router import route_query from app.generation.synthesizer import _estimate_complexity from app.generation.calculation_engine import detect_and_calculate, format_for_context from app.generation.context_compressor import compress_context from app.cache import cache_lookup from app.retrieval.retriever import embed_query from app.retrieval.retrieval_trace import RetrievalTrace, store_trace # ── RetrievalTrace — created once per query, threaded through entire pipeline ── _trace = RetrievalTrace(query_id=query_id, query=question) # ── Stage 1: Instant intent classification (pure keyword, ~1ms) ────────── route = route_query(question) domain_paths = route.get("domain_paths", []) _complexity = _estimate_complexity(question) _q = question.lower() _DRAFT_KW = [ "draft","notice","reply","appeal","submission","advisory","scn","show cause", "drc-01","drc 01","asmt-10","asmt 10","drc-07","drc 07","drc-03","drc 03", "write a letter","write letter","prepare reply","representation", "response to notice","respond to","our understanding","gst implications", "gst implication","provide opinion","provide advisory","our comments", "tax position","gst treatment of","advise on","legal opinion", "our client is","we are engaged in","facts of the case", ] # Definition/explanation queries must NEVER route to DRAFTING_PROMPT regardless # of session history. A plain knowledge query ("define X", "what is X", # "provide definition of X") has no missing facts and should never trigger # the check-3 clarification flow. _NEVER_DRAFT_PATTERNS = [ r'\b(what\s+is|what\s+are|define|definition\s+of|explain|meaning\s+of)\b', r'\b(provide|give|state|share)\s+(the\s+)?(definition|meaning|explanation|rate|provision)', r'\b(relevant\s+circular|applicable\s+circular|circular\s+on)\b', r'\b(full\s+form|abbreviation|what\s+does\s+.+stand\s+for)\b', # Correction / continuation signals — never reroute to advisory/drafting mode r'you\s+got\s+me\s+(all\s+)?wrong', r'(that\'s|that\s+is)\s+(wrong|incorrect|not\s+right|off\s+route|off\s+track)', r'completely\s+(switched|diverted|changed)\s+to', r'stick\s+to\s+(our|the|this)\s+conversation', r'take\s+it\s+forward|taking\s+(it|this)\s+forward', r'continue\s+(the|our|this)\s+(discussion|conversation|analysis|thread|topic)', r'\bi\s+(already|have\s+already)\s+(said|told|replied|mentioned)', r'\bi\s+replied\s+to\s+you', r'as\s+(i|we)\s+(said|mentioned|discussed|told|replied)', r'not\s+what\s+(i|we)\s+(asked|said|meant)', ] import re as _re _is_never_draft = any(_re.search(p, _q) for p in _NEVER_DRAFT_PATTERNS) _is_draft_early = (not _is_never_draft) and any(k in _q for k in _DRAFT_KW) if "rule 42" in _q or "rule42" in _q: _init_msg = "Computing Rule 42 ITC Reversal..." elif "rule 43" in _q or "rule43" in _q: _init_msg = "Computing Rule 43 Capital Goods Reversal..." elif any(k in _q for k in ["itc", "input tax credit", "section 16", "section 17"]): _init_msg = "Analyzing ITC Eligibility Provisions..." elif _is_draft_early: _init_msg = "Initializing Advisory & Drafting Engine..." elif any(k in _q for k in ["refund", "export", "lut"]): _init_msg = "Checking Refund & Export Provisions..." elif any(k in _q for k in ["interest", "section 50", "penalty"]): _init_msg = "Computing Interest & Penalty Exposure..." else: _init_msg = "Initializing Statutory Analyzer..." yield f"__STATUS__:{json.dumps({'msg': _init_msg})}__END_STATUS__" # ── Stage 2: PARALLEL — session history + query embedding ───────────────── # Both are independent I/O-bound tasks; run them simultaneously. _HISTORY_DRAFT_KW = [ "advisory","our understanding","gst implications","gst implication", "provide opinion","our comments","tax position","gst treatment", "advise on","legal opinion","our client","we are engaged", "facts of the case","draft","notice","reply","appeal","scn","show cause", "drc-01","drc-07","asmt-10","representation", "our comments from gst perspective","i've reviewed your", "issue raised","section invoked","to draft a strong reply", ] # Phrases that indicate the user is referencing a previous session _CROSS_SESSION_KW = [ "remember", "that chat", "previous session", "last time", "we discussed", "earlier session", "you mentioned", "from before", "that consultation", "we talked about", "as discussed", "recall when", "in our previous", "that case we", "previous chat", "last session", ] _is_cross_session_ref = any(k in question.lower() for k in _CROSS_SESSION_KW) def _fetch_cross_session_context(): """Search the user's other sessions for relevant content when they reference a past chat.""" if not _is_cross_session_ref: return "" _coll = get_session_collection() if _coll is None: return "" # Get the user_id from the current session _sess_doc = _coll.find_one({"session_id": session_id}, {"user_id": 1}) if session_id else None if not _sess_doc: return "" _user_id = _sess_doc.get("user_id", "") if not _user_id: return "" # Build search terms: strip cross-session keywords and use remaining words _search_q = question.lower() for kw in _CROSS_SESSION_KW: _search_q = _search_q.replace(kw, " ") _search_terms = [w for w in _search_q.split() if len(w) > 3][:6] if not _search_terms: return "" _regex = "|".join(_search_terms) try: _matches = list(_coll.find( {"user_id": _user_id, "session_id": {"$ne": session_id}, "messages.content": {"$regex": _regex, "$options": "i"}}, {"_id": 0, "session_id": 1, "title": 1, "messages": 1, "updated_at": 1} ).sort("updated_at", -1).limit(3)) except Exception: return "" if not _matches: return "" _ctx_parts = [] for _m in _matches: # Find the most relevant message pair (user question + LETA answer) _msgs = _m.get("messages", []) _best_pair = "" for _i, _msg in enumerate(_msgs): _content = _msg.get("content", "") if any(t in _content.lower() for t in _search_terms): # Include user question + LETA response pair _u_msg = _msgs[_i - 1]["content"][:300] if _i > 0 and _msgs[_i-1]["role"] == "user" else "" _a_msg = _content[:500] if _u_msg: _best_pair = f"USER: {_u_msg}\nLETA: {_a_msg}" else: _best_pair = f"LETA: {_a_msg}" break if _best_pair: _title = _m.get("title", "Previous Session") _date = _m.get("updated_at", "") _date_str = _date.strftime("%b %d, %Y") if hasattr(_date, "strftime") else str(_date)[:10] _ctx_parts.append(f'[MEMORY — "{_title}" ({_date_str})]:\n{_best_pair}') return "\n\n".join(_ctx_parts) def _fetch_history_sync(): if not session_id: return "", False _coll = get_session_collection() if _coll is None: return "", False _sess = _coll.find_one({"session_id": session_id}) if not _sess or "messages" not in _sess: return "", False # IMPORTANT: include the current user message (already saved by the # immediate save above) in the history window. Previously [:-1] excluded # it, so CHECK 1 in the drafting prompt could not see the user's reply # to LETA's questions and would ask again. Using [-7:] shows the model # the full sequence: LETA-questions → USER-answers → produce draft. _recent = _sess["messages"][-7:] _hist = "".join(f"{m['role'].upper()}: {m['content']}\n" for m in _recent) # Draft detection: check only the 2 most recent messages so a single # old advisory turn doesn't permanently lock the whole session into # drafting mode for every follow-up question. _last_2_text = " ".join( m.get("content", "") for m in _sess["messages"][-2:] ).lower() _is_d = any(k in _last_2_text for k in _HISTORY_DRAFT_KW) return _hist, _is_d yield f"__STATUS__:{json.dumps({'msg': 'Scanning Semantic Cache...'})}__END_STATUS__" (history_context, _session_is_draft), query_vec, cross_session_context = await _asyncio.gather( _asyncio.to_thread(_fetch_history_sync), _asyncio.to_thread(embed_query, question), _asyncio.to_thread(_fetch_cross_session_context), ) _is_draft = _is_draft_early or _session_is_draft # ── Stage 3: Cache lookup ───────────────────────────────────────────────── cached_answer = await _asyncio.to_thread(cache_lookup, question, query_vec) if cached_answer: cached_text, cached_sources = cached_answer _ask_logger.info("Cache HIT", extra={"query_id": query_id, "cache_hit": True}) # Log minimal trace for cache hits so query_id is still searchable try: _trace.record_preprocessing( original_query=question, refined_query=question, domain_route=domain_paths, complexity_score=_complexity, response_mode="cache_hit", ) _trace.answer = {"cache_hit": True, "query_id": query_id} store_trace(_trace) except Exception: pass yield f"__STATUS__:{json.dumps({'msg': 'Cache Hit — Instant Retrieval Complete.'})}__END_STATUS__" if cached_sources: yield f"__METADATA__:{json.dumps({'type': 'metadata', 'sources': cached_sources})}__END_METADATA__" yield cached_text return # ── Stage 4: Pre-computed calculations (sync, instant) ──────────────────── calc_result = detect_and_calculate(question) calc_block = format_for_context(calc_result) if calc_result else "" if calc_result: yield f"__STATUS__:{json.dumps({'msg': 'Pre-computing Statutory Formula...'})}__END_STATUS__" # ── Stage 5: Retrieval — strategy depends on query type ─────────────────── # Follow-up query enrichment: if query is short (<8 words) and history # exists, enrich the retrieval query with context from the last exchange # so "what about penalties?" searches with full legal context. retriever = get_retriever() _retrieval_top_k = 30 if _is_draft else (25 if _complexity >= 0.60 else 20) domain_label = ", ".join(domain_paths[:2]) if domain_paths else "All Databases" def _enrich_for_retrieval(q: str, hist: str) -> str: if hist and len(q.split()) < 8: _last_lines = [l for l in hist.split('\n') if l.startswith('ASSISTANT:')] if _last_lines: _ctx = _last_lines[-1].replace('ASSISTANT:', '').strip()[:100] return f"{q} [context: {_ctx}]" return q _retrieval_q = _enrich_for_retrieval(question, history_context) yield f"__STATUS__:{json.dumps({'msg': f'Searching Statutory Database ({domain_label})...'})}__END_STATUS__" if _is_draft: # Draft: rule-based sub-queries, no LLM expansion call needed _statute_q = _retrieval_q + " section rule act provisions conditions eligibility liability" _caselaw_q = _retrieval_q + " high court supreme court judgment held ruling decision AAR" _circular_q = _retrieval_q + " CBIC circular notification clarification instruction" _adv = { "queries": [_retrieval_q, _statute_q, _caselaw_q, _circular_q], "hyde_document": "", "topic": "General", "subtopic": None, } # Record preprocessing before handing off to retriever _trace.record_preprocessing( original_query=question, refined_query=_retrieval_q, sub_queries=[_statute_q, _caselaw_q, _circular_q], topic="General", domain_route=domain_paths, complexity_score=_complexity, response_mode="draft", ) chunks = await _asyncio.to_thread( retriever.search, _retrieval_q, _retrieval_top_k, route["use_sources"], _adv, domain_paths, True, False, _trace, # skip_rerank=False, trace=_trace ) else: # All non-draft queries: run query expansion + fast retrieval IN PARALLEL. # Expansion uses 4 angle-specific sub-queries (statutory/circular/notification/factual) # + a corpus-authentic HyDE document. Parallel execution means the LLM expansion # adds zero latency beyond what retrieval already takes. yield f"__STATUS__:{json.dumps({'msg': 'Expanding Query for Precision Retrieval...'})}__END_STATUS__" def _fast_retrieve(): # Trace is NOT passed here — fast retrieve is skip_rerank=True # and supplement_and_rerank runs the full reranking pass with trace. return retriever.search( _retrieval_q, _retrieval_top_k, route["use_sources"], None, domain_paths, False, skip_rerank=True, ) def _expand(): from app.retrieval.query_refiner import generate_advanced_queries return generate_advanced_queries(_retrieval_q) try: fast_chunks, advanced_queries = await _asyncio.gather( _asyncio.to_thread(_fast_retrieve), _asyncio.to_thread(_expand), ) # Record preprocessing now that we have expanded queries _adv_qs = advanced_queries or {} _trace.record_preprocessing( original_query=question, refined_query=_retrieval_q, sub_queries=_adv_qs.get("queries", []), hyde_doc=_adv_qs.get("hyde_document", ""), topic=_adv_qs.get("topic", "General"), subtopic=_adv_qs.get("subtopic"), domain_route=domain_paths, complexity_score=_complexity, response_mode=( "detailed" if _complexity >= 0.60 else "standard" if _complexity >= 0.25 else "brief" ), ) # Supplement the fast pool with expanded-query FAISS results, # then do ONE FlashRank + LegalReranker pass on the merged pool. chunks = await _asyncio.to_thread( retriever.supplement_and_rerank, fast_chunks, advanced_queries, _retrieval_q, _retrieval_top_k, _trace, # trace ) except Exception as _retrieval_exc: import traceback as _tb logger.error( "[CRASH] rag_pipeline_orchestrator retrieval stage failed | " f"query_id={query_id} | {_retrieval_exc!r}\n{_tb.format_exc()}" ) yield f"__STATUS__:{json.dumps({'msg': 'Retrieval engine error. Please try again.'})}__END_STATUS__" yield "\n\n⚠ **LETA encountered a retrieval error.** The statutory database search failed for this query. Please try again — if the problem persists, try rephrasing your question." return # ── Stage 6: Context assembly (pure Python, no blocking I/O) ───────────── citation_block = build_context(chunks, is_draft=_is_draft) compressed_block = compress_context(chunks, question, is_draft=_is_draft) # Detect if this is a follow-up / correction so the model gets a strong # continuation signal and doesn't restart from first principles. _FOLLOWUP_KW = [ "you got me wrong", "you are wrong", "that's wrong", "that is wrong", "off route", "off track", "completely switched", "completely diverted", "stick to our conversation", "stick to the conversation", "take it forward", "taking it forward", "continue the discussion", "as i said", "as i mentioned", "i already said", "i already told", "i replied", "i have said", "i told you", "you mentioned", "building on", "following up", "based on what you said", "not what i asked", "not what i meant", "you missed", ] _is_followup = bool(history_context) and ( len(question.split()) < 25 or any(k in question.lower() for k in _FOLLOWUP_KW) ) _history_label = ( "⚠ ACTIVE CONVERSATION — CONTINUE FROM HERE. Do NOT restart with basics already covered. " "Directly continue the specific legal discussion in progress.\n" if _is_followup else "CHAT HISTORY" ) full_rag_context = ( (f"--- MEMORY FROM PREVIOUS SESSIONS ---\n{cross_session_context}\n--- END MEMORY ---\n\n" if cross_session_context else "") + (f"--- {_history_label} ---\n{history_context}\n--- END HISTORY ---\n\n" if history_context else "") + citation_block + (f"\n\n{calc_block}" if calc_block else "") + "\n\n--- COMPRESSED STATUTORY EXCERPTS (for quick reference) ---\n\n" + compressed_block ) from app.generation.rules_engine import rules_engine truth_rules_text = rules_engine.get_all_rules_as_text() # ── Stage 7: Streaming LLM synthesis ────────────────────────────────────── if _is_draft: _synth_msg = "Synthesizing Advisory & Drafting..." elif _complexity >= 0.60: _synth_msg = "Synthesizing Legal Position & Precedents..." else: _synth_msg = "Synthesizing Sovereign Legal Position..." yield f"__STATUS__:{json.dumps({'msg': _synth_msg})}__END_STATUS__" from app.generation.synthesizer import synthesize_answer_stream response_stream = synthesize_answer_stream( question, full_rag_context, session_is_draft=_is_draft ) # ── Log and store the retrieval trace ──────────────────────────────── # This fires before streaming starts — captures everything up to synthesis. # Answer metadata (latency, model) is captured later in stream_and_save. try: _trace.answer = { "query_id": query_id, "cache_hit": False, "complexity": round(_complexity, 3), "mode": ("draft" if _is_draft else "detailed" if _complexity >= 0.60 else "standard" if _complexity >= 0.25 else "brief"), } store_trace(_trace) # Single structured JSON log line — filterable in CloudWatch Logs Insights: # filter trace_type = "retrieval_trace" # filter query_id = "LETA-20260810-..." import logging as _log_mod _tlog = _log_mod.getLogger("leta.trace") _tlog.info( "retrieval_trace", extra=_trace.to_log_dict(), ) except Exception as _te: logger.debug(f"Trace logging failed (non-fatal): {_te}") try: async for chunk in stream_and_save( response_stream, session_id, question, chunks=chunks, context=citation_block, truth_rules_text=truth_rules_text, query_vec=query_vec, is_draft_session=_is_draft, ): yield chunk except Exception as _synth_exc: import traceback as _tb logger.error( "[CRASH] rag_pipeline_orchestrator synthesis stage failed | " f"query_id={query_id} | {_synth_exc!r}\n{_tb.format_exc()}" ) yield "\n\n⚠ **LETA encountered an error while generating the response.** The statutory sources were retrieved successfully but the synthesis step failed. Please try asking again — your question is valid and the answer exists in our database." from fastapi.responses import StreamingResponse return StreamingResponse(rag_pipeline_orchestrator(), media_type="text/event-stream") @app.post("/ask-sync") @limiter.limit("20/minute") async def ask_question_sync(request: Request, req: QuestionRequest): """Non-streaming version of /ask — returns complete JSON response. Required for AWS API Gateway HTTP_PROXY compatibility (no SSE streaming support).""" if not _warmup_complete: from starlette.responses import JSONResponse as _jr return _jr(status_code=503, content={"detail": "Service warming up — retry in 30s"}, headers={"Retry-After": "30"}) import asyncio as _asyncio import urllib.parse as _urlparse from fastapi.responses import JSONResponse as _JSONResponse from app.routing.router import route_query as _route_query from app.generation.synthesizer import _estimate_complexity, synthesize_answer_stream as _synth_stream from app.generation.calculation_engine import detect_and_calculate, format_for_context from app.generation.context_compressor import compress_context from app.cache import cache_lookup from app.retrieval.retriever import embed_query question = req.question.strip() session_id = req.session_id if session_id: collection = get_session_collection() if collection is not None: collection.update_one( {"session_id": session_id}, {"$push": {"messages": {"role": "user", "content": question, "timestamp": datetime.now()}}} ) route = _route_query(question) domain_paths = route.get("domain_paths", []) _q = question.lower() _complexity = _estimate_complexity(question) history_context = "" _session_is_draft = False if session_id: collection = get_session_collection() if collection is not None: _sess = collection.find_one({"session_id": session_id}) if _sess and "messages" in _sess: recent = _sess["messages"][:-1][-6:] for msg in recent: history_context += f"{msg['role'].upper()}: {msg['content']}\n" _DRAFT_HIST_KW = [ "advisory", "our understanding", "gst implications", "provide opinion", "our comments", "tax position", "gst treatment", "advise on", "legal opinion", "draft", "notice", "reply", "appeal", "scn", "show cause", "drc-01", ] _session_is_draft = any(k in history_context.lower() for k in _DRAFT_HIST_KW) query_vec = embed_query(question) cached_answer = cache_lookup(question, query_vec) if cached_answer: cached_text, cached_sources = cached_answer return _JSONResponse({"answer": cached_text, "sources": cached_sources or []}) calc_result = detect_and_calculate(question) calc_block = format_for_context(calc_result) if calc_result else "" _DRAFT_KW = [ "draft", "notice", "reply", "appeal", "submission", "advisory", "scn", "show cause", "drc-01", "drc 01", "asmt-10", "our understanding", "gst implications", "provide opinion", "our comments", "tax position", "advise on", ] _is_draft = _session_is_draft or any(k in _q for k in _DRAFT_KW) # Sync mode: use rule-based 4-angle multi-query (avoids extra LLM latency on the sync path). # Covers statutory, circular, notification, and factual angles — same structure as the # LLM expansion but keyword-driven so it completes in microseconds. refined_q = question if _is_draft: advanced_queries = { "queries": [ refined_q, refined_q + " section rule act provisions conditions eligibility liability", refined_q + " high court supreme court judgment held ruling decision AAR", refined_q + " CBIC circular notification clarification instruction", ], "hyde_document": "", "topic": "General", "subtopic": None, } else: advanced_queries = { "queries": [ refined_q, refined_q + " section CGST Act rule proviso conditions eligibility", refined_q + " CBIC Circular clarification instruction", refined_q + " GST Notification Central Tax Rate exemption 2017 2018 2019 2020 2021 2022 2023 2024 2025", ], "hyde_document": "", "topic": "General", "subtopic": None, } retriever = get_retriever() # CrossEncoder reranking enabled: ms-marco-MiniLM-L-6-v2 runs in ~200-400ms # for 50 candidate pairs on CPU — well within the 29s ALB timeout. # The old NLI DeBERTa model was both slow AND broken (returned 3-class arrays). _top_k = 20 if _is_draft else (18 if _complexity >= 0.60 else 15) chunks = retriever.search( query=refined_q, top_k=_top_k, allowed_sources=route["use_sources"], advanced_queries=advanced_queries, domain_paths=domain_paths, is_draft=_is_draft, skip_rerank=False, ) citation_block = build_context(chunks, is_draft=_is_draft) compressed_block = compress_context(chunks, question, is_draft=_is_draft) full_rag_context = ( (f"--- CHAT HISTORY ---\n{history_context}\n--- END HISTORY ---\n\n" if history_context else "") + citation_block + (f"\n\n{calc_block}" if calc_block else "") + "\n\n--- COMPRESSED STATUTORY EXCERPTS (for quick reference) ---\n\n" + compressed_block ) # Draft/advisory queries: Haiku + 4000 tokens (~23-27s) — fits API Gateway's 29s limit. # Non-draft Q&A: Sonnet (6000-12000 tokens — Quick Take + Key Extracts + Detailed Advisory). answer = await _asyncio.to_thread( lambda: "".join(_synth_stream(question, full_rag_context, session_is_draft=_is_draft, force_haiku=_is_draft)) ) unique_sources: list = [] seen_src: set = set() for c in chunks: _dedup_rel = c.get("rel_path", "") or c.get("metadata", {}).get("rel_path", "") or c.get("source", "") src_key = (_dedup_rel, c.get("page", 0)) if src_key not in seen_src: seen_src.add(src_key) _raw_src = c.get("source", "") or c.get("metadata", {}).get("source", "") _rel_path = c.get("rel_path", "") or c.get("metadata", {}).get("rel_path", "") _basename = os.path.basename(_rel_path) if _rel_path else os.path.basename(_raw_src) _enc_path = _urlparse.quote(_rel_path.replace("\\", "/"), safe="") if _rel_path else "" _enc_name = _urlparse.quote(_basename, safe="") _url = ( f"/api/documents/view_by_path?path={_enc_path}" if _enc_path else f"/api/documents/view?category=all&filename={_enc_name}" ) _snippet = (c.get("text") or "").strip() unique_sources.append({ "title": _basename or "Document", "page": c.get("page", 1), "url": _url, "rel_path": _rel_path, "score": float(c.get("_rerank_score", 0)), "snippet": _snippet[:800] if _snippet else "", }) if len(unique_sources) >= 8: break if session_id and answer.strip(): collection = get_session_collection() if collection is not None: collection.update_one( {"session_id": session_id}, { "$push": {"messages": {"role": "assistant", "content": answer, "timestamp": datetime.now()}}, "$set": {"updated_at": datetime.now()}, } ) return _JSONResponse({"answer": answer, "sources": unique_sources}) @app.post("/ask-with-file") @limiter.limit("20/minute") async def ask_question_with_file( request: Request, file: UploadFile = File(...), question: str = Form(...), session_id: Optional[str] = Form(None), ): if not _warmup_complete: from starlette.responses import JSONResponse as _jr return _jr(status_code=503, content={"detail": "Service warming up — retry in 30s"}, headers={"Retry-After": "30"}) question_text = question.strip() # 1. Read the file file_bytes = await file.read() filename = (file.filename or "").lower() extracted_text = "" # 2. Parse based on extension if filename.endswith(".pdf"): from app.ingestion.pdf_scanned import extract_text_from_scanned_pdf import tempfile import os # Need to save to temp file for fitz to open with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp: tmp.write(file_bytes) tmp_path = tmp.name try: pages = extract_text_from_scanned_pdf(tmp_path) extracted_text = "\n".join(p["text"] for p in pages) finally: os.remove(tmp_path) elif filename.endswith((".png", ".jpg", ".jpeg")): from app.ingestion.pdf_scanned import extract_text_from_image extracted_text = extract_text_from_image(file_bytes) elif filename.endswith(".txt"): extracted_text = file_bytes.decode("utf-8", errors="replace") elif filename.endswith(".docx"): try: import docx as _docx from io import BytesIO as _BytesIO _doc = _docx.Document(_BytesIO(file_bytes)) _paras = [p.text for p in _doc.paragraphs if p.text.strip()] # Also extract table cell text for _tbl in _doc.tables: for _row in _tbl.rows: for _cell in _row.cells: if _cell.text.strip(): _paras.append(_cell.text.strip()) extracted_text = "\n".join(_paras) except Exception as _e: extracted_text = f"[DOCX extraction error: {_e}]" else: extracted_text = "[Unsupported file format. Please upload PDF, DOCX, TXT, PNG, JPG, or JPEG.]" # Save User Question if session_id: collection = get_session_collection() if collection is not None: collection.update_one( {"session_id": session_id}, {"$push": {"messages": {"role": "user", "content": question_text, "timestamp": datetime.now()}}} ) # Fetch history history_context = "" if session_id: collection = get_session_collection() if collection is not None: session = collection.find_one({"session_id": session_id}) if session and "messages" in session: recent = session["messages"][:-1][-6:] for msg in recent: history_context += f"{msg['role'].upper()}: {msg['content']}\n" # Route + optional query expansion (blocked Haiku call — skip for simple queries) route = route_query(question_text) from app.generation.synthesizer import _estimate_complexity _file_complexity = _estimate_complexity(question_text) if _file_complexity >= 0.80: from app.retrieval.query_refiner import generate_advanced_queries advanced_queries = generate_advanced_queries(question_text) refined_q = advanced_queries.get("queries", [question_text])[0] else: advanced_queries = {"queries": [question_text], "hyde_document": "", "topic": "General", "subtopic": None} refined_q = question_text retriever = get_retriever() chunks = retriever.search( query=refined_q, top_k=15, allowed_sources=route["use_sources"], advanced_queries=advanced_queries, domain_paths=route.get("domain_paths", []), ) from app.generation.context_builder import build_context from app.generation.context_compressor import compress_context rag_context = ( build_context(chunks) + "\n\n--- COMPRESSED STATUTORY EXCERPTS (for quick reference) ---\n\n" + compress_context(chunks, question_text) ) # Combine File Content with RAG Context file_context = "" if extracted_text.strip(): file_context = f"\n--- UPLOADED FILE CONTENT ({file.filename}) ---\n{extracted_text}\n--- END UPLOADED FILE ---\n\n" full_rag_context = file_context + rag_context if history_context: full_rag_context = f"--- CHAT HISTORY ---\n{history_context}\n--- END HISTORY ---\n\n" + full_rag_context # Build truth rules text for hallucination guard from app.generation.rules_engine import rules_engine truth_rules_text = rules_engine.get_all_rules_as_text() # Generate Stream from app.generation.synthesizer import synthesize_answer_stream response_stream = synthesize_answer_stream(question_text, full_rag_context) from fastapi.responses import StreamingResponse wrapped_stream = stream_and_save( response_stream, session_id, question_text, chunks=chunks, context=rag_context, truth_rules_text=truth_rules_text, ) return StreamingResponse(wrapped_stream, media_type="text/event-stream") # ---------- Feedback Endpoint ---------- class FeedbackRequest(BaseModel): session_id: Optional[str] = None question: str answer_preview: str rating: int # 1 = thumbs up, -1 = thumbs down comment: Optional[str] = None @app.post("/feedback") @limiter.limit("60/minute") async def submit_feedback(request: Request, req: FeedbackRequest): """Stores user ratings (thumbs up/down) for quality monitoring and future fine-tuning.""" import logging _fb_logger = logging.getLogger("feedback") try: from app.database import get_db db = get_db() if db is not None: db["feedback"].insert_one({ "session_id": req.session_id, "question": req.question, "answer_preview": req.answer_preview[:200], "rating": req.rating, "comment": req.comment, "user": "anonymous", "timestamp": datetime.now(), }) _fb_logger.info(f"Feedback recorded | rating={req.rating} | q={req.question[:60]}") return {"status": "recorded", "rating": req.rating} except Exception as e: _fb_logger.error(f"Feedback save error: {e}") return {"status": "error", "detail": str(e)} # ---------- PDF Reporting Endpoint ---------- from fastapi.responses import Response from app.generation.pdf_report import PDFReportGenerator import os pdf_gen = PDFReportGenerator(output_dir="data/generated_reports") class PDFRequest(BaseModel): title: str content: str @app.post("/generate-pdf") def create_pdf(req: PDFRequest): # Use the class to generate PDF # We construct a filename import hashlib title_hash = hashlib.md5(req.title.encode()).hexdigest()[:8] filename = f"Report_{title_hash}.pdf" pdf_path = pdf_gen.generate_report(req.content, filename) if not os.path.exists(pdf_path): return Response(status_code=500, content="Error generating PDF") with open(pdf_path, "rb") as f: pdf_bytes = f.read() return Response( content=pdf_bytes, media_type="application/pdf", headers={"Content-Disposition": "attachment; filename=advisory_report.pdf"} ) # ---------- Frontend Serving ---------- from fastapi.staticfiles import StaticFiles from fastapi.responses import FileResponse # We look for the dist directory which should exist in the multi-stage docker container frontend_dist_path = Path(__file__).resolve().parent.parent.parent / "frontend" / "dist" if frontend_dist_path.exists() and frontend_dist_path.is_dir(): # Mount assets directly assets_path = frontend_dist_path / "assets" if assets_path.exists(): app.mount("/assets", StaticFiles(directory=assets_path), name="assets") # Catch-all route for SPA @app.get("/{full_path:path}") async def serve_frontend(full_path: str): # Ignore API routes if full_path.startswith("api/"): return Response(status_code=404) filepath = frontend_dist_path / full_path if filepath.exists() and filepath.is_file(): return FileResponse(filepath) # Fallback to index.html for React Router index_path = frontend_dist_path / "index.html" if index_path.exists(): return FileResponse(index_path) return Response(status_code=404, content="Frontend not found")