Update main.py

main.py

@@ -1,14 +1,17 @@
+# main.py
 import os
 import tempfile
 import logging
 import traceback
-from fastapi import FastAPI, UploadFile, File, Header, HTTPException, Body
+from fastapi import FastAPI, UploadFile, File, Header, HTTPException, Body, Request
 from fastapi.responses import JSONResponse
 from pydantic import BaseModel
 from transformers import pipeline
 from langdetect import detect, DetectorFactory
 from PIL import Image
-from smebuilder_vector import retriever  # Your vector retrieval module
+import io
+from smebuilder_vector import retriever  # your existing retriever module
+import spitch
 
 # ==============================
 # Logging Setup
@@ -16,10 +19,13 @@ from smebuilder_vector import retriever  # Your vector retrieval module
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger("DevAssist")
 
+# Debug log file for prompts + outputs
+DEBUG_LOG_FILE = os.getenv("LLM_DEBUG_LOG", "llm_debug.log")
+
 # ==============================
 # App Init
 # ==============================
-app = FastAPI(title="DevAssist AI Backend")
+app = FastAPI(title="DevAssist / CuraAI Backend")
 
 # ==============================
 # Config
@@ -27,17 +33,25 @@ app = FastAPI(title="DevAssist AI Backend")
 DetectorFactory.seed = 0
 PROJECT_API_KEY = os.getenv("PROJECT_API_KEY")
 SPITCH_API_KEY = os.getenv("SPITCH_API_KEY")
+
+# Models chosen per task (public/reasonable defaults)
 HF_MODELS = {
-    "chat": "bigcode/starcoderbase",
-    "autodoc": "Salesforce/codegen-2B-mono",
-    "sme": "deepseek-ai/deepseek-coder-1.3b-instruct"
+    "chat": os.getenv("CHAT_MODEL", "bigcode/starcoderbase"),                     # coding assistant
+    "autodoc": os.getenv("AUTODOC_MODEL", "Salesforce/codegen-2B-mono"),         # code -> docs
+    "sme": os.getenv("SME_MODEL", "deepseek-ai/deepseek-coder-1.3b-instruct"),   # frontend generation
+    "image_caption": os.getenv("IMAGE_CAPTION_MODEL", "Salesforce/blip-image-captioning-base")
 }
 
 if not SPITCH_API_KEY:
     raise RuntimeError("Set SPITCH_API_KEY in environment before starting.")
 
+# Initialize Spitch client once
+spitch_client = spitch.Spitch()
+# Optionally set env var for Spitch API if required by client library
+os.environ["SPITCH_API_KEY"] = SPITCH_API_KEY
+
 # ==============================
-# Auth Check
+# Authentication helper
 # ==============================
 def check_auth(authorization: str | None):
     if not PROJECT_API_KEY:
@@ -49,15 +63,15 @@ def check_auth(authorization: str | None):
         raise HTTPException(status_code=403, detail="Invalid token")
 
 # ==============================
-# Global Exception Handler
+# Global exception handler
 # ==============================
 @app.exception_handler(Exception)
-async def global_exception_handler(request, exc: Exception):
-    logger.error(f"Unhandled error: {exc}")
+async def global_exception_handler(request: Request, exc: Exception):
+    logger.error(f"Unhandled error: {exc}", exc_info=True)
     return JSONResponse(status_code=500, content={"error": str(exc)})
 
 # ==============================
-# Request Models
+# Request models
 # ==============================
 class ChatRequest(BaseModel):
     question: str
@@ -68,53 +82,122 @@ class AutoDocRequest(BaseModel):
 class SMERequest(BaseModel):
     user_prompt: str
 
+# For simple vector search API
+class VectorRequest(BaseModel):
+    query: str
+
 # ==============================
-# Pipeline Loader
+# Pipeline loader with fallback
 # ==============================
 def load_pipeline(task: str, model_name: str, fallback: str = None):
+    """
+    Load a HuggingFace pipeline with a fallback option.
+    Keep the load minimal (no device_map here — set in env for production).
+    """
     try:
+        logger.info(f"Loading pipeline task={task} model={model_name}")
         return pipeline(task, model=model_name)
     except Exception as e:
-        logger.warning(f"Failed to load {model_name}: {e}")
+        logger.warning(f"Failed to load {model_name} for task={task}: {e}")
         if fallback:
-            logger.info(f"Falling back to {fallback}")
+            logger.info(f"Falling back to {fallback} for task={task}")
             return pipeline(task, model=fallback)
-        raise e
+        raise
 
 # ==============================
-# Pipelines
+# Pipelines (load on startup)
 # ==============================
-chat_pipe = load_pipeline("text-generation", HF_MODELS["chat"], "gpt2")
-autodoc_pipe = load_pipeline("text-generation", HF_MODELS["autodoc"], "gpt2")
-sme_pipe = load_pipeline("text-generation", HF_MODELS["sme"], "gpt2")
+# text-generation pipelines for chat/autodoc/sme
+chat_pipe = load_pipeline("text-generation", HF_MODELS["chat"], fallback="gpt2")
+autodoc_pipe = load_pipeline("text-generation", HF_MODELS["autodoc"], fallback="gpt2")
+sme_pipe = load_pipeline("text-generation", HF_MODELS["sme"], fallback="gpt2")
+
+# image caption / image-to-text pipeline for crop/vision tasks
+image_caption_pipe = load_pipeline("image-to-text", HF_MODELS["image_caption"], fallback="Salesforce/blip-image-captioning-base")
 
 # ==============================
-# Helper Functions
+# Helper / wrapper functions
 # ==============================
-def run_pipeline(pipe, prompt: str):
+def debug_log_prompt(prompt: str, output: str, tag: str = "LLM"):
+    try:
+        with open(DEBUG_LOG_FILE, "a", encoding="utf-8") as fh:
+            fh.write(f"=== {tag} PROMPT START ===\n")
+            fh.write(prompt + "\n")
+            fh.write("--- MODEL OUTPUT ---\n")
+            fh.write(output + "\n")
+            fh.write(f"=== {tag} PROMPT END ===\n\n")
+    except Exception:
+        logger.exception("Failed to write debug log")
+
+def run_pipeline(pipe, prompt: str, max_new_tokens: int = 1024):
+    """
+    Run a text-generation pipeline and return text or structured error.
+    Logs prompt + output to debug file.
+    """
     try:
-        output_list = pipe(prompt, max_new_tokens=1024, do_sample=True)
-        text = output_list[0].get("generated_text", "").strip() if isinstance(output_list, list) else str(output_list)
+        # call pipeline (many models return list with 'generated_text')
+        output_list = pipe(prompt, max_new_tokens=max_new_tokens, do_sample=True)
+        text = ""
+        if isinstance(output_list, list) and len(output_list) > 0:
+            # handle generators that include 'generated_text'
+            first = output_list[0]
+            if isinstance(first, dict) and "generated_text" in first:
+                text = first["generated_text"]
+            else:
+                text = str(first)
+        else:
+            text = str(output_list)
+
+        text = text.strip()
+        debug_log_prompt(prompt, text, tag="TEXT-GEN")
+        logger.info("Prompt executed successfully")
 
-        # Log prompt + output
-        logger.info(f"Prompt:\n{prompt}\n--- Output:\n{text}\n--- End")
         if not text:
             return {"success": False, "error": "⚠️ LLM returned empty output", "prompt": prompt}
         return text
     except Exception as e:
-        logger.error(f"Pipeline error: {e}")
-        return {"success": False, "error": f"⚠️ LLM error: {str(e)}", "prompt": prompt, "trace": traceback.format_exc()}
+        logger.error("Pipeline execution error", exc_info=True)
+        trace = traceback.format_exc()
+        debug_log_prompt(prompt, f"EXCEPTION:\n{trace}", tag="TEXT-GEN")
+        return {"success": False, "error": f"⚠️ LLM error: {str(e)}", "trace": trace, "prompt": prompt}
+
+def run_image_to_text(pipe, image_bytes: bytes, prompt: str):
+    """
+    Run image-to-text pipelines (image captioning / multimodal).
+    Returns generated_text or error structure.
+    """
+    try:
+        image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
+        output_list = pipe(image, prompt=prompt)
+        text = ""
+        if isinstance(output_list, list) and len(output_list) > 0 and isinstance(output_list[0], dict):
+            text = output_list[0].get("generated_text", "")
+        else:
+            text = str(output_list)
+        text = text.strip()
+        debug_log_prompt(prompt, text, tag="IMG-TO-TEXT")
+        if not text:
+            return {"success": False, "error": "⚠️ Vision model returned empty output", "prompt": prompt}
+        return text
+    except Exception as e:
+        logger.exception("Image-to-text pipeline error")
+        trace = traceback.format_exc()
+        debug_log_prompt(prompt, f"EXCEPTION:\n{trace}", tag="IMG-TO-TEXT")
+        return {"success": False, "error": f"⚠️ Vision model error: {str(e)}", "trace": trace, "prompt": prompt}
 
 # ==============================
-# Audio Processing Helper
+# Audio processing (Spitch) helper
 # ==============================
 async def process_audio(file: UploadFile, lang_hint: str | None = None):
-    import spitch
-    spitch_client = spitch.Spitch()
+    """
+    Save audio temporarily, transcribe via Spitch client, detect language and optionally translate to English.
+    Returns (transcription, detected_lang, translation)
+    """
     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()
 
@@ -124,6 +207,7 @@ async def process_audio(file: UploadFile, lang_hint: str | None = None):
         else:
             resp = spitch_client.speech.transcribe(content=audio_bytes)
     except Exception:
+        # fallback to english if Spitch fails with the given hint
         resp = spitch_client.speech.transcribe(language="en", content=audio_bytes)
 
     transcription = getattr(resp, "text", "") or (resp.get("text", "") if isinstance(resp, dict) else "")
@@ -131,13 +215,13 @@ async def process_audio(file: UploadFile, lang_hint: str | None = None):
     try:
         detected_lang = detect(transcription) if transcription.strip() else "en"
     except Exception:
-        pass
+        detected_lang = "en"
 
     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", "")
+            translation = getattr(translation_resp, "text", "") or translation_resp.get("text", "") or transcription
         except Exception:
             translation = transcription
 
@@ -148,43 +232,72 @@ async def process_audio(file: UploadFile, lang_hint: str | None = None):
 # ==============================
 @app.get("/")
 async def root_endpoint():
-    return {"status": "✅ DevAssist AI Backend running"}
+    return {"status": "✅ DevAssist / CuraAI Backend running"}
 
+# ----- Chat: coding assistant -----
 @app.post("/chat")
 async def chat_endpoint(req: ChatRequest, authorization: str | None = Header(None)):
     check_auth(authorization)
-    prompt = f"You are a professional coding assistant. Answer clearly:\nQuestion: {req.question}\nAnswer:"
-    result = run_pipeline(chat_pipe, prompt)
+    # prompt template tuned for coding Q&A
+    prompt = (
+        "You are DevAssist — a helpful, concise coding assistant. "
+        f"Answer clearly with code samples if relevant.\n\nQuestion:\n{req.question}\n\nAnswer:"
+    )
+    result = run_pipeline(chat_pipe, prompt, max_new_tokens=512)
     return result if isinstance(result, dict) else {"reply": result}
 
+# ----- Autodoc: code -> documentation -----
 @app.post("/autodoc")
 async def autodoc_endpoint(req: AutoDocRequest, authorization: str | None = Header(None)):
     check_auth(authorization)
-    prompt = f"Generate professional documentation for the following code in Markdown:\n{req.code}\nDocumentation:"
-    result = run_pipeline(autodoc_pipe, prompt)
+    prompt = (
+        "You are DevAssist DocBot. Produce professional Markdown documentation for the provided code.\n\n"
+        f"Code:\n{req.code}\n\nDocumentation:"
+    )
+    result = run_pipeline(autodoc_pipe, prompt, max_new_tokens=512)
     return result if isinstance(result, dict) else {"documentation": result}
 
+# ----- SME: production-ready frontend generation (with retriever context) -----
 @app.post("/sme/generate")
 async def sme_generate_endpoint(req: SMERequest, authorization: str | None = Header(None)):
     check_auth(authorization)
     try:
-        context_docs = retriever.get_relevant_documents(req.user_prompt)
+        # Use retriever for context injection (keep old method for compatibility)
+        try:
+            context_docs = retriever.get_relevant_documents(req.user_prompt)
+        except AttributeError:
+            # if newer retriever API uses .invoke
+            context_docs = retriever.invoke(req.user_prompt)
+
         context = "\n".join([doc.page_content for doc in context_docs]) if context_docs else "No extra context"
-        prompt = f"Generate production-ready frontend code based on this prompt:\n{req.user_prompt}\nContext:\n{context}\nOutput:"
-        result = run_pipeline(sme_pipe, prompt)
+        prompt = (
+            "You are a senior full-stack engineer. "
+            "Generate production-ready frontend code (index.html, styles.css, script.js) "
+            f"based on the prompt:\n{req.user_prompt}\n\nContext:\n{context}\n\nOutput:"
+        )
+        result = run_pipeline(sme_pipe, prompt, max_new_tokens=1500)
         return {"success": True, "data": result if isinstance(result, str) else result.get("reply", "")}
     except Exception as e:
+        logger.exception("SME generate endpoint error")
         return {"success": False, "error": f"⚠️ LLM error: {str(e)}", "trace": traceback.format_exc()}
 
+# ----- SME Speech generate: STT -> SME -----
 @app.post("/sme/speech-generate")
 async def sme_speech_endpoint(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)
     try:
-        context_docs = retriever.get_relevant_documents(translation)
+        try:
+            context_docs = retriever.get_relevant_documents(translation)
+        except AttributeError:
+            context_docs = retriever.invoke(translation)
+
         context = "\n".join([doc.page_content for doc in context_docs]) if context_docs else "No extra context"
-        prompt = f"Generate production-ready frontend code based on this prompt:\n{translation}\nContext:\n{context}\nOutput:"
-        result = run_pipeline(sme_pipe, prompt)
+        prompt = (
+            "You are a senior full-stack engineer. Generate production-ready frontend code "
+            f"based on the prompt:\n{translation}\n\nContext:\n{context}\n\nOutput:"
+        )
+        result = run_pipeline(sme_pipe, prompt, max_new_tokens=1500)
         return {
             "success": True,
             "transcription": transcription,
@@ -193,11 +306,46 @@ async def sme_speech_endpoint(file: UploadFile = File(...), lang_hint: str | Non
             "sme_site": result if isinstance(result, str) else result.get("reply", "")
         }
     except Exception as e:
+        logger.exception("SME speech-generate error")
         return {"success": False, "error": f"⚠️ LLM error: {str(e)}", "trace": traceback.format_exc()}
 
+# ----- Vision/crop doctor style endpoint (image + text -> diagnosis / explanation) -----
+@app.post("/vision/diagnose")
+async def vision_diagnose(symptoms: str = Header(...), image: UploadFile = File(...), authorization: str | None = Header(None)):
+    """
+    Use an image-to-text model (BLIP) to analyze an image + farmer description, then produce
+    a simple diagnosis & treatment plan. Returns a string or error object.
+    """
+    check_auth(authorization)
+    image_bytes = await image.read()
+    prompt = (
+        f"Farmer reports: {symptoms}. Analyze this plant image, diagnose the likely disease, "
+        "provide simple treatment steps and short prevention advice in plain language."
+    )
+    result = run_image_to_text(image_caption_pipe, image_bytes, prompt)
+    return {"diagnosis": result} if isinstance(result, str) else result
+
+# ----- Vector search wrapper endpoint -----
+@app.post("/vector-search")
+async def vector_search(req: VectorRequest, authorization: str | None = Header(None)):
+    check_auth(authorization)
+    try:
+        # call your existing vector query function in smebuilder_vector (query_vector)
+        try:
+            results = retriever.get_relevant_documents(req.query)
+        except AttributeError:
+            # fallback to invoke if retriever API differs
+            results = retriever.invoke(req.query)
+        # normalize a simple list response
+        brief = [{"page_content": getattr(r, "page_content", str(r)), "meta": getattr(r, "metadata", {})} for r in results]
+        return {"results": brief}
+    except Exception as e:
+        logger.exception("Vector search error")
+        return {"error": f"Vector search error: {str(e)}", "trace": traceback.format_exc()}
+
 # ==============================
 # Run App
 # ==============================
 if __name__ == "__main__":
     import uvicorn
-    uvicorn.run("main:app", host="0.0.0.0", port=7860, reload=False)
+    uvicorn.run("main:app", host="0.0.0.0", port=int(os.getenv("PORT", "7860")), reload=False)