Update main.py

main.py

@@ -1,58 +1,33 @@
-# main.py
 import os
-import tempfile
 import logging
-import traceback
-from fastapi import FastAPI, UploadFile, File, Header, HTTPException, Body, Request
+import io
+from fastapi import FastAPI, Request, Header, HTTPException, UploadFile, File
 from fastapi.responses import JSONResponse
 from pydantic import BaseModel
 from transformers import pipeline
-from langdetect import detect, DetectorFactory
 from PIL import Image
-import io
-from smebuilder_vector import retriever  # your existing retriever module
-import spitch
+from vector import query_vector
 
 # ==============================
 # Logging Setup
 # ==============================
 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 / CuraAI Backend")
+logger = logging.getLogger("AgriCopilot")
 
 # ==============================
-# Config
+# App Initialization
 # ==============================
-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": 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.")
+app = FastAPI(title="AgriCopilot AI API", version="2.0")
 
-# 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
+@app.get("/")
+async def root():
+    return {"status": "AgriCopilot AI Backend is running smoothly ✅"}
 
 # ==============================
-# Authentication helper
+# AUTH CONFIGURATION
 # ==============================
+PROJECT_API_KEY = os.getenv("PROJECT_API_KEY", "agricopilot404")
+
 def check_auth(authorization: str | None):
     if not PROJECT_API_KEY:
         return
@@ -63,289 +38,125 @@ def check_auth(authorization: str | None):
         raise HTTPException(status_code=403, detail="Invalid token")
 
 # ==============================
-# Global exception handler
+# Exception Handling
 # ==============================
 @app.exception_handler(Exception)
 async def global_exception_handler(request: Request, exc: Exception):
-    logger.error(f"Unhandled error: {exc}", exc_info=True)
+    logger.error(f"Unhandled error: {exc}")
     return JSONResponse(status_code=500, content={"error": str(exc)})
 
 # ==============================
-# Request models
+# Request Models
 # ==============================
 class ChatRequest(BaseModel):
-    question: str
+    query: str
 
-class AutoDocRequest(BaseModel):
-    code: str
+class DisasterRequest(BaseModel):
+    report: str
 
-class SMERequest(BaseModel):
-    user_prompt: str
+class MarketRequest(BaseModel):
+    product: str
 
-# For simple vector search API
 class VectorRequest(BaseModel):
     query: str
 
 # ==============================
-# Pipeline loader with fallback
+# Load Hugging Face Pipelines
 # ==============================
-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} for task={task}: {e}")
-        if fallback:
-            logger.info(f"Falling back to {fallback} for task={task}")
-            return pipeline(task, model=fallback)
-        raise
+HF_TOKEN = os.getenv("HUGGINGFACEHUB_API_TOKEN")
 
-# ==============================
-# Pipelines (load on startup)
-# ==============================
-# 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")
+if not HF_TOKEN:
+    logger.warning("⚠️ No Hugging Face token found. Gated models may fail.")
+else:
+    logger.info("✅ Hugging Face token loaded successfully.")
+
+# General text-generation model for chat, disaster, market endpoints
+default_model = "meta-llama/Llama-3.1-8B-Instruct"
+vision_model = "meta-llama/Llama-3.2-11B-Vision-Instruct"
 
-# 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")
+chat_pipe = pipeline("text-generation", model=default_model, token=HF_TOKEN)
+disaster_pipe = pipeline("text-generation", model=default_model, token=HF_TOKEN)
+market_pipe = pipeline("text-generation", model=default_model, token=HF_TOKEN)
+
+# Multimodal crop diagnostic model
+try:
+    crop_pipe = pipeline("image-text-to-text", model=vision_model, token=HF_TOKEN)
+except Exception as e:
+    logger.warning(f"Crop model load failed: {e}")
+    crop_pipe = None
 
 # ==============================
-# Helper / wrapper functions
+# Helper Functions
 # ==============================
-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.
-    """
+def run_conversational(pipe, prompt: str):
     try:
-        # 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")
-
-        if not text:
-            return {"success": False, "error": "⚠️ LLM returned empty output", "prompt": prompt}
-        return text
+        output = pipe(prompt, max_new_tokens=200)
+        if isinstance(output, list) and len(output) > 0:
+            return output[0].get("generated_text", str(output))
+        return str(output)
     except Exception as e:
-        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}
+        logger.error(f"Pipeline error: {e}")
+        return f"⚠️ Model error: {str(e)}"
 
-def run_image_to_text(pipe, image_bytes: bytes, prompt: str):
+def run_crop_doctor(image_bytes: bytes, symptoms: str):
     """
-    Run image-to-text pipelines (image captioning / multimodal).
-    Returns generated_text or error structure.
+    Diagnose crop issues using Meta's multimodal LLaMA Vision model.
     """
+    if not crop_pipe:
+        return "⚠️ Crop analysis temporarily unavailable (model not loaded)."
     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
+        prompt = (
+            f"The farmer reports: {symptoms}. "
+            "Analyze the image and diagnose the likely crop disease. "
+            "Then explain it simply and recommend possible treatment steps."
+        )
+        output = crop_pipe(image, prompt)
+        if isinstance(output, list) and len(output) > 0:
+            return output[0].get("generated_text", str(output))
+        return str(output)
     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}
+        logger.error(f"Crop Doctor pipeline error: {e}")
+        return f"⚠️ Unexpected model error: {str(e)}"
 
 # ==============================
-# Audio processing (Spitch) helper
+# API ROUTES
 # ==============================
-async def process_audio(file: UploadFile, lang_hint: str | None = None):
-    """
-    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()
-
-    try:
-        if lang_hint:
-            resp = spitch_client.speech.transcribe(language=lang_hint, content=audio_bytes)
-        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 "")
-    detected_lang = "en"
-    try:
-        detected_lang = detect(transcription) if transcription.strip() else "en"
-    except Exception:
-        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", "") or transcription
-        except Exception:
-            translation = transcription
-
-    return transcription, detected_lang, translation
-
-# ==============================
-# Endpoints
-# ==============================
-@app.get("/")
-async def root_endpoint():
-    return {"status": "✅ DevAssist / CuraAI Backend running"}
-
-# ----- Chat: coding assistant -----
-@app.post("/chat")
-async def chat_endpoint(req: ChatRequest, authorization: str | None = Header(None)):
+@app.post("/multilingual-chat")
+async def multilingual_chat(req: ChatRequest, authorization: str | None = Header(None)):
     check_auth(authorization)
-    # 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}
+    reply = run_conversational(chat_pipe, req.query)
+    return {"reply": reply}
 
-# ----- Autodoc: code -> documentation -----
-@app.post("/autodoc")
-async def autodoc_endpoint(req: AutoDocRequest, authorization: str | None = Header(None)):
+@app.post("/disaster-summarizer")
+async def disaster_summarizer(req: DisasterRequest, authorization: str | None = Header(None)):
     check_auth(authorization)
-    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}
+    summary = run_conversational(disaster_pipe, req.report)
+    return {"summary": summary}
 
-# ----- SME: production-ready frontend generation (with retriever context) -----
-@app.post("/sme/generate")
-async def sme_generate_endpoint(req: SMERequest, authorization: str | None = Header(None)):
+@app.post("/marketplace")
+async def marketplace(req: MarketRequest, authorization: str | None = Header(None)):
     check_auth(authorization)
-    try:
-        # 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 = (
-            "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()}
+    recommendation = run_conversational(market_pipe, req.product)
+    return {"recommendation": recommendation}
 
-# ----- 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:
-        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 = (
-            "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,
-            "detected_language": detected_lang,
-            "translation": translation,
-            "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}
+        results = query_vector(req.query)
+        return {"results": results}
     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=int(os.getenv("PORT", "7860")), reload=False)
+        logger.error(f"Vector search error: {e}")
+        return {"error": f"Vector search error: {str(e)}"}
+
+@app.post("/crop-doctor")
+async def crop_doctor(
+    symptoms: str = Header(...),
+    image: UploadFile = File(...),
+    authorization: str | None = Header(None)
+):
+    check_auth(authorization)
+    image_bytes = await image.read()
+    diagnosis = run_crop_doctor(image_bytes, symptoms)
+    return {"diagnosis": diagnosis}