from fastapi import FastAPI, HTTPException from fastapi.middleware.cors import CORSMiddleware from pydantic import BaseModel import numpy as np import json import faiss import re from sentence_transformers import SentenceTransformer, CrossEncoder from groq import Groq import os from typing import List, Dict, Optional import logging import httpx logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) app = FastAPI( title="LexNepal AI API", description="Advanced Legal Intelligence API for Nepal Legal Code", version="1.0.0", docs_url="/", redoc_url="/redoc" ) app.add_middleware( CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) class QueryRequest(BaseModel): query: str max_sources: Optional[int] = 10 class Source(BaseModel): law: str section: str section_title: str text: str rel_score: float class QueryResponse(BaseModel): answer: str sources: List[Source] query: str total_candidates: int class StatsResponse(BaseModel): total_provisions: int total_laws: int vector_dimensions: int embedding_model: str reranking_model: str llm_model: str class HealthResponse(BaseModel): status: str models_loaded: bool message: Optional[str] = None _bi_encoder = None _cross_encoder = None _groq_client = None _index = None _metadata = None def get_bi_encoder(): global _bi_encoder if _bi_encoder is None: logger.info("Loading bi-encoder (MPNet)...") _bi_encoder = SentenceTransformer("all-mpnet-base-v2") logger.info("✅ Bi-encoder loaded successfully") return _bi_encoder def get_cross_encoder(): global _cross_encoder if _cross_encoder is None: logger.info("Loading cross-encoder...") _cross_encoder = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2") logger.info("✅ Cross-encoder loaded successfully") return _cross_encoder def get_groq_client(): global _groq_client if _groq_client is None: logger.info("Initializing Groq client...") # Get API key from environment ONLY (no fallback) groq_api_key = os.getenv("GROQ_API_KEY") if not groq_api_key: logger.error("❌ GROQ_API_KEY not found in environment") raise HTTPException( status_code=503, detail="GROQ_API_KEY not configured. Please set it in Hugging Face Space secrets." ) try: _groq_client = Groq(api_key=groq_api_key) logger.info("✅ Groq client initialized successfully") except Exception as e: logger.error(f"❌ Failed to initialize Groq client: {e}") raise HTTPException( status_code=503, detail=f"Failed to initialize Groq client: {str(e)}" ) return _groq_client def get_index(): global _index if _index is None: logger.info("Loading embeddings and creating FAISS index...") try: embeddings = np.load("final_legal_embeddings.npy") logger.info(f"Embeddings shape: {embeddings.shape}") _index = faiss.IndexFlatL2(embeddings.shape[1]) _index.add(embeddings.astype('float32')) logger.info(f"✅ FAISS index created with {embeddings.shape[0]} vectors") except FileNotFoundError: logger.error("❌ Embeddings file not found") raise HTTPException( status_code=503, detail="Embeddings file not found. Please upload final_legal_embeddings.npy" ) return _index def get_metadata(): global _metadata if _metadata is None: logger.info("Loading metadata...") try: with open("final_legal_laws_metadata.json", "r", encoding="utf-8") as f: _metadata = json.load(f) logger.info(f"✅ Loaded {len(_metadata)} legal provisions") except FileNotFoundError: logger.error("❌ Metadata file not found") raise HTTPException( status_code=503, detail="Metadata file not found. Please upload final_legal_laws_metadata.json" ) return _metadata def get_premium_context(query: str, max_sources: int = 10) -> List[Dict]: try: bi_encoder = get_bi_encoder() cross_encoder = get_cross_encoder() index = get_index() metadata = get_metadata() # Stage 1: Encode query query_embedding = bi_encoder.encode([query], convert_to_numpy=True) # Stage 2: Dense retrieval _, indices = index.search(query_embedding.astype('float32'), 25) candidates = [] seen = set() for i in indices[0]: if i != -1 and i < len(metadata): candidates.append(metadata[i].copy()) seen.add(i) # Stage 3: Keyword boosting numbers = re.findall(r'\d+', query) if numbers: for i, item in enumerate(metadata): if any(str(item.get('section', '')) == n for n in numbers): if i not in seen: candidates.append(item.copy()) seen.add(i) # Stage 4: Cross-encoder reranking if candidates: pairs = [ [query, f"{c.get('law', '')} {c.get('section_title', '')} {c.get('text', '')}"] for c in candidates ] scores = cross_encoder.predict(pairs) for i, c in enumerate(candidates): c['rel_score'] = float(scores[i]) candidates = sorted(candidates, key=lambda x: x['rel_score'], reverse=True)[:max_sources] logger.info(f"Retrieved {len(candidates)} relevant candidates") return candidates except Exception as e: logger.error(f"Error in context retrieval: {str(e)}") raise HTTPException(status_code=500, detail=f"Context retrieval error: {str(e)}") @app.get("/health", response_model=HealthResponse) async def health_check(): """Health check endpoint""" try: metadata = get_metadata() models_loaded = True message = f"API is healthy. {len(metadata)} provisions loaded." except Exception as e: models_loaded = False message = f"Error: {str(e)}" return { "status": "healthy" if models_loaded else "unhealthy", "models_loaded": models_loaded, "message": message } @app.get("/stats", response_model=StatsResponse) async def get_statistics(): """Get database statistics""" try: metadata = get_metadata() unique_laws = len(set(d.get('law', '') for d in metadata)) return { "total_provisions": len(metadata), "total_laws": unique_laws, "vector_dimensions": 768, "embedding_model": "all-mpnet-base-v2", "reranking_model": "ms-marco-MiniLM-L-6-v2", "llm_model": "llama-3.3-70b-versatile" } except Exception as e: logger.error(f"Error getting stats: {str(e)}") raise HTTPException(status_code=503, detail=str(e)) @app.post("/query", response_model=QueryResponse) async def process_legal_query(request: QueryRequest): """Process legal query with RAG pipeline""" # Validation if not request.query.strip(): raise HTTPException(status_code=400, detail="Query cannot be empty") if len(request.query) > 1000: raise HTTPException(status_code=400, detail="Query too long (max 1000 characters)") try: logger.info(f"Processing query: {request.query[:100]}...") # Get relevant context candidates = get_premium_context(request.query, request.max_sources) if not candidates: return { "answer": "No relevant legal provisions found in the database for your query. Please try rephrasing or consult a legal professional.", "sources": [], "query": request.query, "total_candidates": 0 } # Build context string context_str = "\n\n".join([ f"[{d['law']} Section {d['section']}]: {d['text']}" for d in candidates ]) # System prompt system_prompt = """You are an Elite Legal Advisor specializing in Nepal law. OPERATIONAL MANDATE: 1. Answer STRICTLY from provided legal text 2. If information is absent, state: "No specific provision found in current database" 3. Always cite exact Law name and Section number 4. Use formal, authoritative legal language 5. NEVER hallucinate or infer beyond provided text 6. Maintain zero-tolerance policy for speculation When citing, use format: "According to [Law Name], Section [Number]..." Provide clear, structured answers with proper legal citations.""" # Generate response using Groq logger.info("Generating LLM response...") groq_client = get_groq_client() response = groq_client.chat.completions.create( model="llama-3.3-70b-versatile", messages=[ {"role": "system", "content": system_prompt}, {"role": "user", "content": f"Legal Context:\n{context_str}\n\nQuery: {request.query}"} ], temperature=0, max_tokens=1500 ) answer = response.choices[0].message.content # Format sources sources = [ Source( law=d['law'], section=str(d['section']), section_title=d['section_title'], text=d['text'], rel_score=d['rel_score'] ) for d in candidates ] logger.info(f"✅ Query processed successfully with {len(sources)} sources") return { "answer": answer, "sources": sources, "query": request.query, "total_candidates": len(candidates) } except HTTPException: raise except Exception as e: logger.error(f"Error processing query: {str(e)}") raise HTTPException(status_code=500, detail=f"Query processing error: {str(e)}") @app.get("/") async def root(): """Root endpoint - API info""" return { "message": "🇳🇵 LexNepal AI API is running", "version": "1.0.0", "description": "Advanced Legal Intelligence for Nepal Legal Code", "endpoints": { "docs": "/ (Swagger UI)", "health": "/health (GET)", "stats": "/stats (GET)", "query": "/query (POST)" }, "technology": "RAG with Hybrid Retrieval + Cross-Encoder Reranking", "support": "https://huggingface.co/spaces/yamraj047/lexnepal-api" } if __name__ == "__main__": import uvicorn uvicorn.run(app, host="0.0.0.0", port=7860)