Update main.py

main.py

@@ -8,8 +8,7 @@ from spitch import Spitch
 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 for context injection
+from smebuilder_vector import retriever
 
 # ----------------- CONFIG -----------------
 DetectorFactory.seed = 0
@@ -96,10 +95,10 @@ Output:
 """
 
 # ----------------- CHAINS -----------------
-chat_chain = PromptTemplate(input_variables=["question"], template=chat_template) | llm
-stt_chain = PromptTemplate(input_variables=["speech"], template=stt_chat_template) | llm
-autodoc_chain = PromptTemplate(input_variables=["code"], template=autodoc_template) | llm
-sme_chain = PromptTemplate(input_variables=["user_prompt", "context"], template=sme_template) | llm
+chat_chain = PromptTemplate(input_variables=["question"], template=chat_template)
+stt_chain = PromptTemplate(input_variables=["speech"], template=stt_chat_template)
+autodoc_chain = PromptTemplate(input_variables=["code"], template=autodoc_template)
+sme_chain = PromptTemplate(input_variables=["user_prompt", "context"], template=sme_template)
 
 # ----------------- REQUEST MODELS -----------------
 class ChatRequest(BaseModel):
@@ -119,26 +118,16 @@ def check_auth(authorization: str | None):
         raise HTTPException(status_code=403, detail="Invalid token")
 
 # ----------------- HELPER FUNCTIONS -----------------
-def run_chain(chain, input_dict: dict):
+def run_llm(prompt_text: str):
     """
-    Safely run a LangChain PromptTemplate | HuggingFaceEndpoint chain.
-    Returns non-empty string, or detailed error info for debugging.
+    Directly run HuggingFaceEndpoint with string input.
+    Returns text or error dict.
     """
     try:
-        # Render template
-        prompt_text = chain.prompt.format(**input_dict) if hasattr(chain, "prompt") else str(input_dict)
-
-        # Generate using HuggingFaceEndpoint
-        output = chain.llm.generate([{"role": "user", "content": prompt_text}])
-
-        # Get text safely
-        text = getattr(output.generations[0][0], "text", "") or ""
-        text = text.strip()
-
-        if not text:
+        output = llm(prompt_text)
+        if not output.strip():
             return {"success": False, "error": "⚠️ LLM returned empty output", "prompt": prompt_text}
-
-        return text
+        return output.strip()
     except Exception:
         return {"success": False, "error": "⚠️ LLM error", "details": traceback.format_exc(), "prompt": prompt_text}
 
@@ -184,14 +173,16 @@ def root():
 @app.post("/chat")
 def chat(req: ChatRequest, authorization: str | None = Header(None)):
     check_auth(authorization)
-    result = run_chain(chat_chain, {"question": req.question})
+    prompt_text = chat_chain.format(question=req.question)
+    result = run_llm(prompt_text)
     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})
+    prompt_text = stt_chain.format(speech=translation)
+    result = run_llm(prompt_text)
     return {
         "transcription": transcription,
         "detected_language": detected_lang,
@@ -202,7 +193,8 @@ async def stt_audio(file: UploadFile = File(...), lang_hint: str | None = None,
 @app.post("/autodoc")
 def autodoc(req: AutoDocRequest, authorization: str | None = Header(None)):
     check_auth(authorization)
-    result = run_chain(autodoc_chain, {"code": req.code})
+    prompt_text = autodoc_chain.format(code=req.code)
+    result = run_llm(prompt_text)
     return result if isinstance(result, dict) else {"documentation": result}
 
 @app.post("/sme/generate")
@@ -212,7 +204,8 @@ 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"
-        result = run_chain(sme_chain, {"user_prompt": user_prompt, "context": context})
+        prompt_text = sme_chain.format(user_prompt=user_prompt, context=context)
+        result = run_llm(prompt_text)
         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()}
@@ -224,7 +217,8 @@ async def sme_speech_generate(file: UploadFile = File(...), lang_hint: str | Non
     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"
-        result = run_chain(sme_chain, {"user_prompt": translation, "context": context})
+        prompt_text = sme_chain.format(user_prompt=translation, context=context)
+        result = run_llm(prompt_text)
         return {
             "success": True,
             "transcription": transcription,