Modified the response structure

app.py

@@ -1,9 +1,10 @@
-# app.py (simplified generate endpoint)
+# app.py (refined with clean metadata)
 import os
 import logging
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
 import tempfile
+from typing import List, Dict
 
 # --- Use writable temp dir for Hugging Face caches ---
 TMP_CACHE = os.environ.get("HF_CACHE_DIR", os.path.join(tempfile.gettempdir(), "hf_cache"))
@@ -17,7 +18,7 @@ os.environ["HF_HOME"] = TMP_CACHE
 os.environ["HF_DATASETS_CACHE"] = TMP_CACHE
 os.environ["HF_METRICS_CACHE"] = TMP_CACHE
 
-app = FastAPI(title="DirectEd LoRA API (simplified)")
+app = FastAPI(title="DirectEd LoRA API with metadata")
 
 @app.get("/health")
 def health():
@@ -25,15 +26,24 @@ def health():
 
 @app.get("/")
 def root():
-    return {"Status": "AI backend is running"}
+    return {"status": "AI backend is running"}
 
 class PromptRequest(BaseModel):
     prompt: str
 
+class Source(BaseModel):
+    name: str
+    url: str
+
+class ResponseWithMetadata(BaseModel):
+    answer: str
+    sources: List[Source] = []
+
 pipe = None
 
 @app.on_event("startup")
 def load_model():
+    """Load base + LoRA adapter model at startup."""
     global pipe
     try:
         from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
@@ -53,19 +63,52 @@ def load_model():
         model = PeftModel.from_pretrained(base_model, ADAPTER_REPO)
         model.eval()
 
-        pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, device_map="auto")
+        pipe = pipeline(
+            "text-generation",
+            model=model,
+            tokenizer=tokenizer,
+            device_map="auto",
+        )
         logging.info("Model and adapter loaded successfully.")
     except Exception as e:
         logging.exception("Failed to load model at startup: %s", e)
         pipe = None
 
-@app.post("/generate")
+def parse_response(raw_text: str) -> ResponseWithMetadata:
+    """Extract answer and sources from raw model output."""
+    import re
+    from collections import OrderedDict
+
+    # Attempt to extract sources if present (looking for URLs)
+    source_pattern = r"(https?://[^\s]+)"
+    urls = re.findall(source_pattern, raw_text)
+
+    # Deduplicate and create simple source list
+    seen = set()
+    sources: List[Source] = []
+    for url in urls:
+        if url not in seen:
+            seen.add(url)
+            sources.append(Source(name="Reference", url=url))
+
+    # Remove sources from the text to keep answer clean
+    answer = re.sub(source_pattern, "", raw_text).strip()
+
+    return ResponseWithMetadata(answer=answer, sources=sources)
+
+@app.post("/generate", response_model=ResponseWithMetadata)
 def generate(req: PromptRequest):
+    """Generate a concise response with optional metadata."""
     if pipe is None:
         raise HTTPException(status_code=503, detail="Model not loaded. Check logs.")
     try:
         output = pipe(req.prompt, max_new_tokens=150, do_sample=True)
-        return {"response": output[0]["generated_text"]}
+        full_text = output[0].get("generated_text", "").strip()
+        if not full_text:
+            raise HTTPException(status_code=500, detail="Model returned empty response.")
+
+        return parse_response(full_text)
+
     except Exception as e:
         logging.exception("Generation failed: %s", e)
-        raise HTTPException(status_code=500, detail=str(e))
+        raise HTTPException(status_code=500, detail=f"Generation failed: {e}")

requirements.txt

@@ -5,4 +5,5 @@ accelerate
 bitsandbytes
 fastapi
 uvicorn
-bitsandbytes
+bitsandbytes
+requests