Update main.py

main.py

@@ -5,8 +5,7 @@ from fastapi import FastAPI, UploadFile, File, Header, HTTPException, Body
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
 from spitch import Spitch
-from langchain.prompts import PromptTemplate
-from langchain_huggingface import HuggingFaceEndpoint
+from transformers import pipeline
 from langdetect import detect, DetectorFactory
 from smebuilder_vector import retriever
 
@@ -24,9 +23,10 @@ if not SPITCH_API_KEY:
 os.environ["SPITCH_API_KEY"] = SPITCH_API_KEY
 spitch_client = Spitch()
 
-# HuggingFace LLM
-llm = HuggingFaceEndpoint(
-    repo_id=HF_MODEL,
+# ----------------- HUGGINGFACE PIPELINE -----------------
+llm_pipeline = pipeline(
+    task="text-generation",
+    model=HF_MODEL,
     temperature=0.7,
     top_p=0.9,
     do_sample=True,
@@ -35,7 +35,7 @@ llm = HuggingFaceEndpoint(
 )
 
 # ----------------- FASTAPI -----------------
-app = FastAPI(title="DevAssist AI Backend (FastAPI + LangChain)")
+app = FastAPI(title="DevAssist AI Backend (FastAPI + HuggingFace Pipeline)")
 
 app.add_middleware(
     CORSMiddleware,
@@ -45,7 +45,7 @@ app.add_middleware(
     allow_headers=["Authorization", "Content-Type"],
 )
 
-# ----------------- PROMPT TEMPLATES -----------------
+# ----------------- PROMPTS -----------------
 chat_template = """You are DevAssist, an AI coding assistant.
 
 Guidelines:
@@ -94,12 +94,6 @@ Context: {context}
 Output:
 """
 
-# ----------------- CHAINS -----------------
-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):
     question: str
@@ -117,20 +111,35 @@ def check_auth(authorization: str | None):
     if token != PROJECT_API_KEY:
         raise HTTPException(status_code=403, detail="Invalid token")
 
-# ----------------- HELPER FUNCTIONS -----------------
-def run_llm(prompt_text: str):
-    """
-    Directly run HuggingFaceEndpoint with string input.
-    Returns text or error dict.
-    """
+# ----------------- DEBUG LOGGING -----------------
+DEBUG_LOG_FILE = "llm_debug.log"
+
+def run_pipeline(prompt_text: str):
     try:
-        output = llm(prompt_text)
-        if not output.strip():
+        output_list = llm_pipeline(prompt_text, max_new_tokens=2048, do_sample=True)
+        text = output_list[0]['generated_text'].strip()
+
+        # Debug logging
+        with open(DEBUG_LOG_FILE, "a", encoding="utf-8") as f:
+            f.write("=== PROMPT START ===\n")
+            f.write(prompt_text + "\n")
+            f.write("--- MODEL OUTPUT ---\n")
+            f.write(text + "\n")
+            f.write("=== PROMPT END ===\n\n")
+
+        if not text:
             return {"success": False, "error": "⚠️ LLM returned empty output", "prompt": prompt_text}
-        return output.strip()
+        return text
     except Exception:
+        with open(DEBUG_LOG_FILE, "a", encoding="utf-8") as f:
+            f.write("=== PROMPT START ===\n")
+            f.write(prompt_text + "\n")
+            f.write("--- EXCEPTION ---\n")
+            f.write(traceback.format_exc() + "\n")
+            f.write("=== PROMPT END ===\n\n")
         return {"success": False, "error": "⚠️ LLM error", "details": traceback.format_exc(), "prompt": prompt_text}
 
+# ----------------- AUDIO PROCESSING -----------------
 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:
@@ -173,16 +182,16 @@ def root():
 @app.post("/chat")
 def chat(req: ChatRequest, authorization: str | None = Header(None)):
     check_auth(authorization)
-    prompt_text = chat_chain.format(question=req.question)
-    result = run_llm(prompt_text)
+    prompt_text = chat_template.format(question=req.question)
+    result = run_pipeline(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)
-    prompt_text = stt_chain.format(speech=translation)
-    result = run_llm(prompt_text)
+    prompt_text = stt_chat_template.format(speech=translation)
+    result = run_pipeline(prompt_text)
     return {
         "transcription": transcription,
         "detected_language": detected_lang,
@@ -193,8 +202,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)
-    prompt_text = autodoc_chain.format(code=req.code)
-    result = run_llm(prompt_text)
+    prompt_text = autodoc_template.format(code=req.code)
+    result = run_pipeline(prompt_text)
     return result if isinstance(result, dict) else {"documentation": result}
 
 @app.post("/sme/generate")
@@ -204,8 +213,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"
-        prompt_text = sme_chain.format(user_prompt=user_prompt, context=context)
-        result = run_llm(prompt_text)
+        prompt_text = sme_template.format(user_prompt=user_prompt, context=context)
+        result = run_pipeline(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()}
@@ -217,8 +226,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"
-        prompt_text = sme_chain.format(user_prompt=translation, context=context)
-        result = run_llm(prompt_text)
+        prompt_text = sme_template.format(user_prompt=translation, context=context)
+        result = run_pipeline(prompt_text)
         return {
             "success": True,
             "transcription": transcription,