Updated app

app.py

@@ -1,10 +1,9 @@
-# app.py (refined with clean metadata)
+# app.py
 import os
 import logging
+import tempfile
 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"))
@@ -18,8 +17,17 @@ 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 with metadata")
+app = FastAPI(title="DirectEd LoRA API (concise)")
+
+# ---------------------
+# Request Model
+# ---------------------
+class PromptRequest(BaseModel):
+    prompt: str
 
+# ---------------------
+# Health & Root
+# ---------------------
 @app.get("/health")
 def health():
     return {"ok": True}
@@ -28,22 +36,13 @@ def health():
 def root():
     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] = []
-
+# ---------------------
+# Load Model on Startup
+# ---------------------
 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
@@ -63,52 +62,38 @@ 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
 
-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)
+# ---------------------
+# Generate Endpoint
+# ---------------------
+@app.post("/generate")
 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)
-        full_text = output[0].get("generated_text", "").strip()
-        if not full_text:
-            raise HTTPException(status_code=500, detail="Model returned empty response.")
+        # Limit tokens to avoid huge outputs
+        max_tokens = 200
+
+        output = pipe(req.prompt, max_new_tokens=max_tokens, do_sample=True)
+        text = output[0].get("generated_text", "").strip()
+
+        # Remove repeated context if present
+        if text.startswith(req.prompt):
+            text = text[len(req.prompt):].strip()
+
+        if not text:
+            logging.warning("Model returned empty response for prompt: %s", req.prompt)
+            text = "No response generated by the model."
 
-        return parse_response(full_text)
+        return {"response": text}
 
     except Exception as e:
-        logging.exception("Generation failed: %s", e)
+        logging.exception("Generation failed for prompt '%s': %s", req.prompt, e)
         raise HTTPException(status_code=500, detail=f"Generation failed: {e}")