Update main.py

main.py

@@ -1,8 +1,6 @@
 import os
 import tempfile
-import logging
-from fastapi import FastAPI, UploadFile, File, Header, HTTPException, Body, Request
-from fastapi.responses import JSONResponse
+from fastapi import FastAPI, UploadFile, File, Header, HTTPException, Body
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
 from spitch import Spitch
@@ -12,17 +10,13 @@ from langdetect import detect, DetectorFactory
 from huggingface_hub.utils import HfHubHTTPError
 from smebuilder_vector import retriever  # Retriever for context injection
 
-# ----------------- LOGGING -----------------
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger("DevAssist")
-
 # ----------------- CONFIG -----------------
 DetectorFactory.seed = 0
 
 SPITCH_API_KEY = os.getenv("SPITCH_API_KEY")
 HF_MODEL = os.getenv("HF_MODEL", "deepseek-ai/deepseek-coder-1.3b-instruct")
 FRONTEND_ORIGIN = os.getenv("ALLOWED_ORIGIN", "*")
-PROJECT_API_KEY = os.getenv("PROJECT_API_KEY")  # default if not set
+PROJECT_API_KEY = os.getenv("PROJECT_API_KEY")
 
 if not SPITCH_API_KEY:
     raise RuntimeError("Set SPITCH_API_KEY in environment before starting.")
@@ -41,31 +35,7 @@ llm = HuggingFaceEndpoint(
     max_new_tokens=2048
 )
 
-# ----------------- HELPERS -----------------
-def run_llm_model(chain, payload: dict):
-    """
-    Safely run HuggingFace model through LangChain chain.
-    Handles string, dict, and list responses without crashing.
-    """
-    try:
-        result = chain.invoke(payload)
-        logger.info(f"HF raw response: {result}")
-
-        if isinstance(result, str):
-            return result.strip()
-
-        if isinstance(result, dict) and "generated_text" in result:
-            return result["generated_text"].strip()
-
-        if isinstance(result, list) and len(result) > 0 and "generated_text" in result[0]:
-            return result[0]["generated_text"].strip()
-
-        return str(result).strip()
-    except Exception as e:
-        logger.error(f"LLM execution failed: {e}")
-        return f"⚠️ LLM error: {str(e)}"
-
-# ----------------- FASTAPI -----------------
+# FastAPI app
 app = FastAPI(title="DevAssist AI Backend (FastAPI + LangChain)")
 
 # CORS
@@ -77,13 +47,7 @@ app.add_middleware(
     allow_headers=["Authorization", "Content-Type"],
 )
 
-# Global exception handler
-@app.exception_handler(Exception)
-async def global_exception_handler(request: Request, exc: Exception):
-    logger.error(f"Unhandled error: {exc}")
-    return JSONResponse(status_code=500, content={"error": str(exc)})
-
-# ----------------- PROMPTS -----------------
+# ----------------- PROMPT TEMPLATES -----------------
 chat_template = """You are DevAssist, an AI coding assistant.
 
 Guidelines:
@@ -114,6 +78,7 @@ Code: {code}
 Documentation:
 """
 
+# 🔥 Fixed SME template with escaped curly braces
 sme_template = """
 You are a senior full-stack engineer specializing in modern front-end development.
 Your job is to generate **production-ready code** for websites and apps.
@@ -125,7 +90,7 @@ Guidelines:
 - JavaScript must add interactivity (animations, toggles, button actions)
 - Include hero, feature grid, testimonials, and footer
 - Use realistic content (no lorem ipsum, no placeholders)
-- Return ONLY valid JSON: { "files": { "index.html": "...", "styles.css": "...", "script.js": "..." } }
+- Return ONLY valid JSON: {{ "files": {{ "index.html": "...", "styles.css": "...", "script.js": "..." }} }}
 
 Prompt: {user_prompt}
 Context: {context}
@@ -139,7 +104,7 @@ stt_chain = PromptTemplate(input_variables=["speech"], template=stt_chat_templat
 autodoc_chain = PromptTemplate(input_variables=["code"], template=autodoc_template) | llm
 sme_chain = PromptTemplate(input_variables=["user_prompt", "context"], template=sme_template) | llm
 
-# ----------------- MODELS -----------------
+# ----------------- REQUEST MODELS -----------------
 class ChatRequest(BaseModel):
     question: str
 
@@ -148,7 +113,8 @@ class AutoDocRequest(BaseModel):
 
 # ----------------- AUTH -----------------
 def check_auth(authorization: str | None):
-    if not PROJECT_API_KEY:  # if no key set, skip
+    """Validate Bearer token against PROJECT_API_KEY"""
+    if not PROJECT_API_KEY:
         return
     if not authorization or not authorization.startswith("Bearer "):
         raise HTTPException(status_code=401, detail="Missing bearer token")
@@ -164,18 +130,28 @@ def root():
 @app.post("/chat")
 def chat(req: ChatRequest, authorization: str | None = Header(None)):
     check_auth(authorization)
-    return {"reply": run_llm_model(chat_chain, {"question": req.question})}
+    try:
+        answer = chat_chain.invoke({"question": req.question})
+        return {"reply": answer.strip() if isinstance(answer, str) else str(answer)}
+    except HfHubHTTPError as e:
+        if "exceeded" in str(e).lower() or "quota" in str(e).lower():
+            return {"reply": "⚠️ Daily token limit reached. Try again in 24 hours."}
+        raise e
 
 @app.post("/stt")
 async def stt_audio(file: UploadFile = File(...), lang_hint: str | None = None, authorization: str | None = Header(None)):
     check_auth(authorization)
+
     suffix = os.path.splitext(file.filename)[1] or ".wav"
     with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tf:
         tf.write(await file.read())
         tmp_path = tf.name
 
     try:
-        resp = spitch_client.speech.transcribe(language=lang_hint or "en", content=open(tmp_path, "rb").read())
+        if lang_hint:
+            resp = spitch_client.speech.transcribe(language=lang_hint, content=open(tmp_path, "rb").read())
+        else:
+            resp = spitch_client.speech.transcribe(content=open(tmp_path, "rb").read())
     except Exception:
         resp = spitch_client.speech.transcribe(language="en", content=open(tmp_path, "rb").read())
 
@@ -194,56 +170,33 @@ async def stt_audio(file: UploadFile = File(...), lang_hint: str | None = None,
         except Exception:
             translation = transcription
 
-    reply = run_llm_model(stt_chain, {"speech": translation})
-    return {"transcription": transcription, "detected_language": detected_lang, "translation": translation, "reply": reply}
+    reply = stt_chain.invoke({"speech": translation})
+    return {
+        "transcription": transcription,
+        "detected_language": detected_lang,
+        "translation": translation,
+        "reply": reply.strip() if isinstance(reply, str) else str(reply)
+    }
 
 @app.post("/autodoc")
 def autodoc(req: AutoDocRequest, authorization: str | None = Header(None)):
     check_auth(authorization)
-    docs = run_llm_model(autodoc_chain, {"code": req.code})
-    return {"documentation": docs}
+    docs = autodoc_chain.invoke({"code": req.code})
+    return {"documentation": docs.strip() if isinstance(docs, str) else str(docs)}
 
 @app.post("/sme/generate")
 async def sme_generate(payload: dict = Body(...), authorization: str | None = Header(None)):
     check_auth(authorization)
-    user_prompt = payload.get("user_prompt", "")
-    context_docs = retriever.get_relevant_documents(user_prompt)
-    context = "\n".join([doc.page_content for doc in context_docs]) if context_docs else "No extra context"
-    response = run_llm_model(sme_chain, {"user_prompt": user_prompt, "context": context})
-    return {"success": True, "data": response}
-
-@app.post("/sme/speech-generate")
-async def sme_speech_generate(file: UploadFile = File(...), lang_hint: str | None = None, authorization: str | None = Header(None)):
-    check_auth(authorization)
-    suffix = os.path.splitext(file.filename)[1] or ".wav"
-    with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tf:
-        tf.write(await file.read())
-        tmp_path = tf.name
-
-    try:
-        resp = spitch_client.speech.transcribe(language=lang_hint or "en", content=open(tmp_path, "rb").read())
-    except Exception:
-        resp = spitch_client.speech.transcribe(language="en", content=open(tmp_path, "rb").read())
-
-    transcription = getattr(resp, "text", "") or (resp.get("text", "") if isinstance(resp, dict) else "")
-    detected_lang = "en"
     try:
-        detected_lang = detect(transcription) if transcription.strip() else "en"
-    except Exception:
-        pass
-
-    translation = transcription
-    if detected_lang != "en":
-        try:
-            translation_resp = spitch_client.text.translate(text=transcription, source=detected_lang, target="en")
-            translation = getattr(translation_resp, "text", "") or translation_resp.get("text", "")
-        except Exception:
-            translation = transcription
-
-    context_docs = retriever.get_relevant_documents(translation)
-    context = "\n".join([doc.page_content for doc in context_docs]) if context_docs else "No extra context"
-    sme_response = run_llm_model(sme_chain, {"user_prompt": translation, "context": context})
-    return {"success": True, "transcription": transcription, "detected_language": detected_lang, "translation": translation, "sme_site": sme_response}
+        user_prompt = payload.get("user_prompt", "")
+        context_docs = retriever.get_relevant_documents(user_prompt)
+        context = "\n".join([doc.page_content for doc in context_docs]) if context_docs else "No extra context"
+        response = sme_chain.invoke({"user_prompt": user_prompt, "context": context})
+        return {"success": True, "data": response}
+    except HfHubHTTPError as e:
+        if "exceeded" in str(e).lower() or "quota" in str(e).lower():
+            return {"success": False, "error": "⚠️ Token quota for today has been used. Please come back in 24 hours."}
+        raise e
 
 # ----------------- MAIN -----------------
 if __name__ == "__main__":