Update main.py

main.py

@@ -1,5 +1,6 @@
 import os
 import tempfile
+import traceback
 from fastapi import FastAPI, UploadFile, File, Header, HTTPException, Body
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
@@ -21,11 +22,10 @@ PROJECT_API_KEY = os.getenv("PROJECT_API_KEY")
 if not SPITCH_API_KEY:
     raise RuntimeError("Set SPITCH_API_KEY in environment before starting.")
 
-# Init Spitch
 os.environ["SPITCH_API_KEY"] = SPITCH_API_KEY
 spitch_client = Spitch()
 
-# HuggingFace LLM (removed task="conversational" to avoid StopIteration bug)
+# HuggingFace LLM
 llm = HuggingFaceEndpoint(
     repo_id=HF_MODEL,
     temperature=0.7,
@@ -35,10 +35,9 @@ llm = HuggingFaceEndpoint(
     max_new_tokens=2048
 )
 
-# FastAPI app
+# ----------------- FASTAPI -----------------
 app = FastAPI(title="DevAssist AI Backend (FastAPI + LangChain)")
 
-# CORS
 app.add_middleware(
     CORSMiddleware,
     allow_origins=[FRONTEND_ORIGIN] if FRONTEND_ORIGIN != "*" else ["*"],
@@ -111,8 +110,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:
         return
     if not authorization or not authorization.startswith("Bearer "):
         raise HTTPException(status_code=401, detail="Missing bearer token")
@@ -120,38 +118,41 @@ def check_auth(authorization: str | None):
     if token != PROJECT_API_KEY:
         raise HTTPException(status_code=403, detail="Invalid token")
 
-# ----------------- ENDPOINTS -----------------
-@app.get("/")
-def root():
-    return {"status": "✅ DevAssist AI Backend running"}
-
-@app.post("/chat")
-def chat(req: ChatRequest, authorization: str | None = Header(None)):
-    check_auth(authorization)
+# ----------------- HELPER FUNCTIONS -----------------
+def run_chain(chain, input_dict: dict):
+    """
+    Safely run a LangChain PromptTemplate | HuggingFaceEndpoint chain.
+    Converts output to string and captures errors.
+    """
     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
+        # Render template
+        if hasattr(chain, "prompt"):
+            prompt_text = chain.prompt.format(**input_dict)
+        else:
+            prompt_text = str(input_dict)
 
-@app.post("/stt")
-async def stt_audio(file: UploadFile = File(...), lang_hint: str | None = None, authorization: str | None = Header(None)):
-    check_auth(authorization)
+        # Generate using HuggingFaceEndpoint (expects str input)
+        output = chain.llm.generate([{"role": "user", "content": prompt_text}])
+        return output.generations[0][0].text.strip()
+    except Exception:
+        return {"success": False, "error": "⚠️ LLM error", "details": traceback.format_exc()}
 
+async def process_audio(file: UploadFile, lang_hint: str | None = None):
     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
 
+    with open(tmp_path, "rb") as f:
+        audio_bytes = f.read()
+
     try:
         if lang_hint:
-            resp = spitch_client.speech.transcribe(language=lang_hint, content=open(tmp_path, "rb").read())
+            resp = spitch_client.speech.transcribe(language=lang_hint, content=audio_bytes)
         else:
-            resp = spitch_client.speech.transcribe(content=open(tmp_path, "rb").read())
+            resp = spitch_client.speech.transcribe(content=audio_bytes)
     except Exception:
-        resp = spitch_client.speech.transcribe(language="en", content=open(tmp_path, "rb").read())
+        resp = spitch_client.speech.transcribe(language="en", content=audio_bytes)
 
     transcription = getattr(resp, "text", "") or (resp.get("text", "") if isinstance(resp, dict) else "")
     detected_lang = "en"
@@ -168,19 +169,36 @@ async def stt_audio(file: UploadFile = File(...), lang_hint: str | None = None,
         except Exception:
             translation = transcription
 
-    reply = stt_chain.invoke({"speech": translation})
+    return transcription, detected_lang, translation
+
+# ----------------- ENDPOINTS -----------------
+@app.get("/")
+def root():
+    return {"status": "✅ DevAssist AI Backend running"}
+
+@app.post("/chat")
+def chat(req: ChatRequest, authorization: str | None = Header(None)):
+    check_auth(authorization)
+    result = run_chain(chat_chain, {"question": req.question})
+    return result if isinstance(result, dict) else {"reply": result}
+
+@app.post("/stt")
+async def stt_audio(file: UploadFile = File(...), lang_hint: str | None = None, authorization: str | None = Header(None)):
+    check_auth(authorization)
+    transcription, detected_lang, translation = await process_audio(file, lang_hint)
+    result = run_chain(stt_chain, {"speech": translation})
     return {
         "transcription": transcription,
         "detected_language": detected_lang,
         "translation": translation,
-        "reply": reply.strip() if isinstance(reply, str) else str(reply)
+        "reply": result if isinstance(result, str) else result.get("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)}
+    result = run_chain(autodoc_chain, {"code": req.code})
+    return result if isinstance(result, dict) else {"documentation": result}
 
 @app.post("/sme/generate")
 async def sme_generate(payload: dict = Body(...), authorization: str | None = Header(None)):
@@ -189,56 +207,28 @@ async def sme_generate(payload: dict = Body(...), authorization: str | None = He
         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 Exception as e:
-        return {"success": False, "error": f"⚠️ LLM error: {str(e)}"}
+        result = run_chain(sme_chain, {"user_prompt": user_prompt, "context": context})
+        return {"success": True, "data": result if isinstance(result, str) else result.get("reply", "")}
+    except Exception:
+        return {"success": False, "error": "⚠️ LLM error", "details": traceback.format_exc()}
 
 @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())
-    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
-
+    transcription, detected_lang, translation = await process_audio(file, lang_hint)
     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})
+        result = run_chain(sme_chain, {"user_prompt": translation, "context": context})
         return {
             "success": True,
             "transcription": transcription,
             "detected_language": detected_lang,
             "translation": translation,
-            "sme_site": sme_response
+            "sme_site": result if isinstance(result, str) else result.get("reply", "")
         }
-    except Exception as e:
-        return {"success": False, "error": f"⚠️ LLM error: {str(e)}"}
+    except Exception:
+        return {"success": False, "error": "⚠️ LLM error", "details": traceback.format_exc()}
 
 # ----------------- MAIN -----------------
 if __name__ == "__main__":