Update main.py

main.py

@@ -1,6 +1,8 @@
 import os
 import tempfile
-from fastapi import FastAPI, UploadFile, File, Header, HTTPException, Body
+import logging
+from fastapi import FastAPI, UploadFile, File, Header, HTTPException, Body, Request
+from fastapi.responses import JSONResponse
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
 from spitch import Spitch
@@ -10,13 +12,17 @@ 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", "super-secret-123")  # Default if not set
+PROJECT_API_KEY = os.getenv("PROJECT_API_KEY")  # default if not set
 
 if not SPITCH_API_KEY:
     raise RuntimeError("Set SPITCH_API_KEY in environment before starting.")
@@ -35,7 +41,31 @@ llm = HuggingFaceEndpoint(
     max_new_tokens=2048
 )
 
-# FastAPI app
+# ----------------- 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 -----------------
 app = FastAPI(title="DevAssist AI Backend (FastAPI + LangChain)")
 
 # CORS
@@ -47,7 +77,13 @@ app.add_middleware(
     allow_headers=["Authorization", "Content-Type"],
 )
 
-# ----------------- PROMPT TEMPLATES -----------------
+# 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 -----------------
 chat_template = """You are DevAssist, an AI coding assistant.
 
 Guidelines:
@@ -103,7 +139,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
 
-# ----------------- REQUEST MODELS -----------------
+# ----------------- MODELS -----------------
 class ChatRequest(BaseModel):
     question: str
 
@@ -112,8 +148,7 @@ class AutoDocRequest(BaseModel):
 
 # ----------------- AUTH -----------------
 def check_auth(authorization: str | None):
-    """Validate Bearer token against PROJECT_API_KEY"""
-    if not PROJECT_API_KEY:  # If not set, skip auth
+    if not PROJECT_API_KEY:  # if no key set, skip
         return
     if not authorization or not authorization.startswith("Bearer "):
         raise HTTPException(status_code=401, detail="Missing bearer token")
@@ -129,28 +164,18 @@ def root():
 @app.post("/chat")
 def chat(req: ChatRequest, authorization: str | None = Header(None)):
     check_auth(authorization)
-    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
+    return {"reply": run_llm_model(chat_chain, {"question": req.question})}
 
 @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:
-        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())
+        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())
 
@@ -169,48 +194,34 @@ async def stt_audio(file: UploadFile = File(...), lang_hint: str | None = None,
         except Exception:
             translation = transcription
 
-    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)
-    }
+    reply = run_llm_model(stt_chain, {"speech": translation})
+    return {"transcription": transcription, "detected_language": detected_lang, "translation": translation, "reply": reply}
 
 @app.post("/autodoc")
 def autodoc(req: AutoDocRequest, authorization: str | None = Header(None)):
     check_auth(authorization)
-    docs = autodoc_chain.invoke({"code": req.code})
-    return {"documentation": docs.strip() if isinstance(docs, str) else str(docs)}
+    docs = run_llm_model(autodoc_chain, {"code": req.code})
+    return {"documentation": docs}
 
 @app.post("/sme/generate")
 async def sme_generate(payload: dict = Body(...), authorization: str | None = Header(None)):
     check_auth(authorization)
-    try:
-        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
+    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:
-        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())
+        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())
 
@@ -229,21 +240,10 @@ async def sme_speech_generate(file: UploadFile = File(...), lang_hint: str | Non
         except Exception:
             translation = transcription
 
-    try:
-        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 = sme_chain.invoke({"user_prompt": translation, "context": context})
-        return {
-            "success": True,
-            "transcription": transcription,
-            "detected_language": detected_lang,
-            "translation": translation,
-            "sme_site": sme_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
+    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}
 
 # ----------------- MAIN -----------------
 if __name__ == "__main__":