import json import re from pathlib import Path import faiss import torch from fastapi import FastAPI, File, Form, HTTPException, UploadFile from pydantic import BaseModel from sentence_transformers import SentenceTransformer from transformers import AutoTokenizer, AutoModelForSeq2SeqLM from peft import PeftModel from app.indexing import ( append_chunks_to_index, build_chunks_from_document, existing_sources, list_indexed_documents, save_raw_document, ) from app.paths import index_path, lora_path, metadata_path # ----------------------------- # API Initialization # ----------------------------- app = FastAPI( title="Enterprise Code Comment Generator", description="CodeT5 + LoRA + RAG for ERPNext Accounting", version="1.0", ) # ----------------------------- # Request Schema # ----------------------------- class CodeRequest(BaseModel): code: str def _safe_filename(filename: str) -> str: name = Path(filename).name if not name or name.startswith("."): raise HTTPException(status_code=400, detail="Invalid filename") if not name.lower().endswith(".md"): raise HTTPException(status_code=400, detail="Only Markdown (.md) files are allowed") return name def _safe_category(category: str) -> str: if not category or not re.fullmatch(r"[a-zA-Z0-9_-]+", category): raise HTTPException(status_code=400, detail="Invalid category") return category # ----------------------------- # Global Variables for Models # ----------------------------- tokenizer = None model = None embedder = None index = None metadata = None BASE_MODEL = "Salesforce/codet5-small" # ----------------------------- # Load Models (Startup Event) # ----------------------------- @app.on_event("startup") async def load_models() -> None: """ Lazily load all heavy models and indexes when the app starts. Mirrors the setup logic from the RAG notebook. """ global tokenizer, model, embedder, index, metadata print("Loading CodeT5 model...") tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL) base_model = AutoModelForSeq2SeqLM.from_pretrained(BASE_MODEL) model = PeftModel.from_pretrained( base_model, str(lora_path()), is_trainable=False, ) model.eval() print("CodeT5 loaded") print("Loading SentenceBERT...") embedder = SentenceTransformer("all-MiniLM-L6-v2") print("SentenceBERT loaded") print("Loading FAISS index and metadata...") print(f"DATA_DIR={index_path().parent}") index = faiss.read_index(str(index_path())) with open(metadata_path(), "r", encoding="utf-8") as f: metadata = json.load(f) print("FAISS index loaded") # ----------------------------- # Comment Generation # ----------------------------- def generate_comment(code: str) -> str: """ Generate a natural-language summary/comment for the given code. """ prompt = "summarize the purpose of this function: " + code inputs = tokenizer( prompt, return_tensors="pt", truncation=True, ) with torch.no_grad(): outputs = model.generate( **inputs.to(model.device), max_new_tokens=80, min_length=30, num_beams=6, repetition_penalty=1.25, length_penalty=1.35, no_repeat_ngram_size=3, early_stopping=True, ) comment = tokenizer.decode(outputs[0], skip_special_tokens=True) return comment.strip() # ----------------------------- # Context Retrieval (RAG-style) # ----------------------------- def retrieve_context(code: str, top_k: int = 5) -> str: """ Retrieve the most relevant documentation chunk for the given code. Mirrors the logic used in the RAG notebook: - Extract function name from the code - Use the function name as the semantic query - Re-rank FAISS candidates using simple keyword overlap """ # Extract function name fn_match = re.search(r"def\s+([a-zA-Z0-9_]+)", code) fn_name = fn_match.group(1).lower() if fn_match else "" keywords = fn_name.split("_") if fn_name else [] # Encode query using the function name (more focused than raw code) query_embedding = embedder.encode( [fn_name or code], convert_to_numpy=True, ) # Retrieve FAISS candidates _, I = index.search(query_embedding, k=top_k) candidates = [metadata[i] for i in I[0]] # Hybrid re-ranking: boost chunks that contain function-name keywords scored = [] for chunk in candidates: text = chunk["text"].lower() keyword_score = sum(1 for k in keywords if k and k in text) scored.append((keyword_score, chunk)) scored.sort(key=lambda x: x[0], reverse=True) best_chunk = scored[0][1] if scored else candidates[0] return best_chunk["text"] # ----------------------------- # Context Summarization # ----------------------------- def summarize_context(context: str) -> str: """ Clean and lightly summarize the retrieved context, similar to the notebook: - Strip markdown artifacts (#, \-, backslashes) - Collapse newlines - Keep only the first couple of meaningful sentences """ clean = context.replace("#", "") clean = clean.replace("\\-", "") clean = clean.replace("\\", "") clean = re.sub(r"\n+", " ", clean) sentences = re.split(r'(?<=[.!?])\s+', clean.strip()) meaningful = [s.strip() for s in sentences if len(s.strip()) > 25] return " ".join(meaningful[:2]) if meaningful else clean.strip() # ----------------------------- # API Endpoint # ----------------------------- @app.post("/generate_comment") def generate_comment_api(request: CodeRequest): code = request.code comment = generate_comment(code) raw_context = retrieve_context(code) context = summarize_context(raw_context) return { "comment": comment, "context": context, } @app.get("/documents") def list_documents(): """ List all documents currently indexed in the FAISS metadata. """ return list_indexed_documents(metadata, index.ntotal) @app.post("/documents") async def add_document( file: UploadFile = File(..., description="Markdown file (.md only)"), category: str = Form( default="accounting", description="Subfolder under data/enterprise_accounting_docs/", ), ): """ Upload a Markdown file and append its chunks to the FAISS index and metadata. Existing documents and vectors are preserved; duplicate sources are rejected. """ if not file.filename: raise HTTPException(status_code=400, detail="Missing filename") filename = _safe_filename(file.filename) category = _safe_category(category) source = f"{category}/{filename}" raw = await file.read() try: content = raw.decode("utf-8") except UnicodeDecodeError as exc: raise HTTPException( status_code=400, detail="File must be valid UTF-8 text", ) from exc if not content.strip(): raise HTTPException(status_code=400, detail="Uploaded file is empty") if source in existing_sources(metadata): raise HTTPException( status_code=409, detail=f"Document already indexed: {source}", ) new_chunks = build_chunks_from_document(source, content) if not new_chunks: raise HTTPException( status_code=400, detail="No indexable chunks produced (document empty or too short)", ) save_raw_document(category, filename, content) vectors_added = append_chunks_to_index( embedder, index, metadata, new_chunks, persist=True, ) return { "message": "Document indexed successfully", "source": source, "chunks_added": len(new_chunks), "vectors_added": vectors_added, "total_chunks": len(metadata), "total_vectors": index.ntotal, }