Updated

main.py

@@ -1,9 +1,5 @@
 import os
-import json
 import tempfile
-import traceback
-from typing import Optional
-
 from fastapi import FastAPI, UploadFile, File, Header, HTTPException, Body
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
@@ -12,15 +8,15 @@ from langchain.prompts import PromptTemplate
 from langchain_huggingface import HuggingFaceEndpoint
 from langdetect import detect, DetectorFactory
 from huggingface_hub.utils import HfHubHTTPError
-from smebuilder_vector import retriever  # retriever that exposes .get_relevant_documents(...)
+from smebuilder_vector import retriever  # Retriever for context injection
 
+# ----------------- CONFIG -----------------
 DetectorFactory.seed = 0
 
-# ----------------- CONFIG -----------------
 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", "")
+PROJECT_API_KEY = os.getenv("PROJECT_API_KEY", "super-secret-123")  # Default if not set
 
 if not SPITCH_API_KEY:
     raise RuntimeError("Set SPITCH_API_KEY in environment before starting.")
@@ -30,14 +26,13 @@ os.environ["SPITCH_API_KEY"] = SPITCH_API_KEY
 spitch_client = Spitch()
 
 # HuggingFace LLM
-# NOTE: pass generation params explicitly (pydantic validation requires explicit params)
 llm = HuggingFaceEndpoint(
     repo_id=HF_MODEL,
     temperature=0.7,
     top_p=0.9,
     do_sample=True,
     repetition_penalty=1.1,
-    max_new_tokens=2048,
+    max_new_tokens=2048
 )
 
 # FastAPI app
@@ -57,11 +52,10 @@ chat_template = """You are DevAssist, an AI coding assistant.
 
 Guidelines:
 - Always format responses in Markdown.
-- Do NOT use inline formatting like bold (**), italics (*), or underlines.
-- Use plain text section headers (e.g., Explanation:, Steps to Fix:, Fixed Code:).
-- Use bullet points (-, 1.) for steps.
-- Use fenced code blocks (```python ... ```) for code.
-- Be friendly yet professional; explain step by step.
+- Use section headers: Explanation:, Steps:, Fixed Code:
+- Use bullet points for steps.
+- Use fenced code blocks for code.
+- Be friendly yet professional.
 
 Question: {question}
 
@@ -69,16 +63,16 @@ Answer:
 """
 
 stt_chat_template = """You are DevAssist, an AI coding assistant.
-- The input is transcribed speech. Interpret it as a dev question.
-- Provide clear answers with code examples (use markdown triple backticks).
-- If input is unclear, ask a clarifying question.
+The input is transcribed speech. Interpret it as a developer question.
+Provide clear answers with code examples.
+If unclear, ask for clarification.
 
 Spoken Question: {speech}
 Answer:
 """
 
 autodoc_template = """You are DevAssist DocBot.
-- Read the code and produce professional documentation in markdown.
+Read the code and produce professional documentation in markdown.
 
 Code: {code}
 Documentation:
@@ -86,24 +80,21 @@ Documentation:
 
 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.
+Your job is to generate **production-ready code** for websites and apps.
 
 Guidelines:
 - Always return three separate files: index.html, styles.css, and script.js
-- HTML must be semantic, responsive, and mobile-first (include <meta name="viewport">)
-- CSS should use Flexbox/Grid and include hover/transition effects
-- JavaScript should add interactivity (e.g. button actions, basic animations, toggles)
-- Include a hero section, a feature grid, testimonials, and footer
-- Use realistic content (avoid lorem ipsum), sensible copy, and accessible markup
-- Return only valid JSON with the keys: "files" -> { "index.html": "...", "styles.css": "...", "script.js": "..." }
+- HTML must be semantic, responsive, and mobile-first
+- CSS should use Flexbox/Grid with hover/transition effects
+- 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": "..." } }
 
-User Prompt:
-{user_prompt}
+Prompt: {user_prompt}
+Context: {context}
 
-Context:
-{context}
-
-Return:
+Output:
 """
 
 # ----------------- CHAINS -----------------
@@ -120,9 +111,9 @@ class AutoDocRequest(BaseModel):
     code: str
 
 # ----------------- AUTH -----------------
-def check_auth(authorization: Optional[str] = None):
-    if not PROJECT_API_KEY:
-        # No API key enforced in this environment
+def check_auth(authorization: str | None):
+    """Validate Bearer token against PROJECT_API_KEY"""
+    if not PROJECT_API_KEY:  # If not set, skip auth
         return
     if not authorization or not authorization.startswith("Bearer "):
         raise HTTPException(status_code=401, detail="Missing bearer token")
@@ -130,29 +121,13 @@ def check_auth(authorization: Optional[str] = None):
     if token != PROJECT_API_KEY:
         raise HTTPException(status_code=403, detail="Invalid token")
 
-# ----------------- HELPERS -----------------
-def try_parse_json(maybe_str: str):
-    """Try to parse JSON; if fails, return None."""
-    try:
-        return json.loads(maybe_str)
-    except Exception:
-        # attempt to find a JSON substring
-        import re
-        m = re.search(r"\{[\s\S]*\}\s*$", maybe_str.strip())
-        if m:
-            try:
-                return json.loads(m.group(0))
-            except Exception:
-                return None
-        return None
-
 # ----------------- ENDPOINTS -----------------
 @app.get("/")
 def root():
-    return {"status": "DevAssist AI Backend running"}
+    return {"status": "✅ DevAssist AI Backend running"}
 
 @app.post("/chat")
-def chat(req: ChatRequest, authorization: Optional[str] = Header(None)):
+def chat(req: ChatRequest, authorization: str | None = Header(None)):
     check_auth(authorization)
     try:
         answer = chat_chain.invoke({"question": req.question})
@@ -163,123 +138,72 @@ def chat(req: ChatRequest, authorization: Optional[str] = Header(None)):
         raise e
 
 @app.post("/stt")
-async def stt_audio(
-    file: UploadFile = File(...),
-    lang_hint: Optional[str] = None,
-    authorization: Optional[str] = Header(None),
-):
+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"
 
-    # create temp file
+    suffix = os.path.splitext(file.filename)[1] or ".wav"
     with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tf:
-        content = await file.read()
-        tf.write(content)
+        tf.write(await file.read())
         tmp_path = tf.name
 
     try:
-        # transcribe
         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:
-        # fallback to english transcription if something fails
         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:
-        detected_lang = "en"
+        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", "") if isinstance(translation_resp, dict) else translation)
+            translation = getattr(translation_resp, "text", "") or translation_resp.get("text", "")
         except Exception:
             translation = transcription
 
-    # call the STT chain (LLM)
-    try:
-        reply = stt_chain.invoke({"speech": translation})
-    except Exception as e:
-        # on LLM problems return transcription anyway
-        reply = f"(LLM error) Transcription: {translation}"
-
-    # cleanup temp file to avoid storage bloat
-    try:
-        os.remove(tmp_path)
-    except Exception:
-        pass
-
+    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": reply.strip() if isinstance(reply, str) else str(reply)
     }
 
 @app.post("/autodoc")
-def autodoc(req: AutoDocRequest, authorization: Optional[str] = Header(None)):
+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)}
 
 @app.post("/sme/generate")
-async def sme_generate(payload: dict = Body(...), authorization: Optional[str] = Header(None)):
-    """
-    Payload expected: { "user_prompt": "Create ...", (optionally) "force_simple": true }
-    Returns: success, data (if success) or error
-    """
+async def sme_generate(payload: dict = Body(...), authorization: str | None = Header(None)):
     check_auth(authorization)
-    user_prompt = payload.get("user_prompt", "")
-    if not user_prompt or not user_prompt.strip():
-        raise HTTPException(status_code=400, detail="user_prompt is required")
-
-    # Get context from retriever (if available)
     try:
-        context_docs = retriever.get_relevant_documents(user_prompt) if retriever else []
-        context = "\n\n".join([getattr(d, "page_content", str(d)) for d in context_docs]) if context_docs else "No extra context"
-    except Exception:
-        context = "No extra context"
-
-    # Invoke SME chain
-    try:
-        raw = sme_chain.invoke({"user_prompt": user_prompt, "context": context})
-        # Try to parse returned JSON
-        parsed = None
-        if isinstance(raw, str):
-            parsed = try_parse_json(raw)
-        elif isinstance(raw, dict):
-            parsed = raw
-
-        if parsed:
-            return {"success": True, "data": parsed}
-        else:
-            # If model didn't return strict JSON, return helpful error + raw output so frontend can show it
-            return {"success": False, "error": "LLM did not return valid JSON", "raw": raw}
+        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
-    except Exception as e:
-        # Debug info for devs (but don't leak sensitive internals in production)
-        return {"success": False, "error": "SME generation failed", "details": str(e), "trace": traceback.format_exc()}
 
 @app.post("/sme/speech-generate")
-async def sme_speech_generate(
-    file: UploadFile = File(...),
-    lang_hint: Optional[str] = None,
-    authorization: Optional[str] = Header(None),
-):
+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"
 
+    suffix = os.path.splitext(file.filename)[1] or ".wav"
     with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tf:
-        content = await file.read()
-        tf.write(content)
+        tf.write(await file.read())
         tmp_path = tf.name
 
     try:
@@ -291,72 +215,35 @@ async def sme_speech_generate(
         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:
-        detected_lang = "en"
+        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", "") if isinstance(translation_resp, dict) else translation)
+            translation = getattr(translation_resp, "text", "") or translation_resp.get("text", "")
         except Exception:
             translation = transcription
 
-    # Get context docs for the transcribed prompt
     try:
-        context_docs = retriever.get_relevant_documents(translation) if retriever else []
-        context = "\n\n".join([getattr(d, "page_content", str(d)) for d in context_docs]) if context_docs else "No extra context"
-    except Exception:
-        context = "No extra context"
-
-    # Invoke SME chain
-    try:
-        raw = sme_chain.invoke({"user_prompt": translation, "context": context})
-        parsed = None
-        if isinstance(raw, str):
-            parsed = try_parse_json(raw)
-        elif isinstance(raw, dict):
-            parsed = raw
-
-        # cleanup tmp file
-        try:
-            os.remove(tmp_path)
-        except Exception:
-            pass
-
-        if parsed:
-            return {
-                "success": True,
-                "transcription": transcription,
-                "detected_language": detected_lang,
-                "translation": translation,
-                "sme_site": parsed,
-            }
-        else:
-            return {
-                "success": False,
-                "error": "LLM did not return valid JSON",
-                "raw": raw,
-                "transcription": transcription,
-                "detected_language": detected_lang,
-                "translation": translation,
-            }
+        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:
-        try:
-            os.remove(tmp_path)
-        except Exception:
-            pass
         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
-    except Exception as e:
-        try:
-            os.remove(tmp_path)
-        except Exception:
-            pass
-        return {"success": False, "error": "SME generation failed", "details": str(e), "trace": traceback.format_exc()}
 
 # ----------------- MAIN -----------------
 if __name__ == "__main__":