miscellaneous/stable2.py

# ==============================================================================
# GRADIO APP WITH LANGCHAIN AGENT (DIAGNOSIS + REMEDY)
# ==============================================================================
"""
This script launches a multi-step agentic application for maize health.
The workflow is as follows:
1.  A fine-tuned Gemma-3N vision model (the 'Diagnoser Tool') analyzes an
    uploaded image to identify the plant's condition.
2.  The diagnosis text is passed to a LangChain agent.
3.  The agent uses a 'Remedy Retriever Tool' (a RAG pipeline) to search a
    knowledge base for relevant treatments.
4.  The agent synthesizes the diagnosis and the retrieved remedy into a
    comprehensive, helpful response.
"""

# --- Step 0: Essential Imports ---
import gradio as gr
import torch
from PIL import Image
from unsloth import FastVisionModel
from transformers import AutoProcessor
import os

# LangChain and RAG components
from langchain_community.vectorstores import FAISS
from langchain_community.document_loaders import TextLoader
from langchain_huggingface import HuggingFaceEmbeddings
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnablePassthrough
from langchain_core.output_parsers import StrOutputParser
from langchain_huggingface import ChatHuggingFace, HuggingFaceEndpoint

print("✅ All libraries imported successfully.")

# --- Step 1: Global Setup - Vision Model (The 'Diagnoser') ---
# This expensive setup runs only ONCE when the application starts.

print("Performing initial setup for the Vision Model...")
BASE_MODEL_NAME = "unsloth/gemma-3n-E2B-it-unsloth-bnb-4bit"

# !!! IMPORTANT: Replace this with the model repo ID you created on the Hub !!!
ADAPTER_PATH = "surfiniaburger/maize-health-diagnosis-adapter" # <--- YOUR HF ADAPTER REPO

VISION_MODEL = None
PROCESSOR = None

try:
    VISION_MODEL, PROCESSOR = FastVisionModel.from_pretrained(
        model_name=BASE_MODEL_NAME,
        max_seq_length=2048,
        load_in_4bit=True,
        dtype=None,
    )
    FastVisionModel.for_inference(VISION_MODEL)
    VISION_MODEL.load_adapter(ADAPTER_PATH)
    print("✅ Vision model and adapter loaded successfully!")

except Exception as e:
    print(f"❌ CRITICAL ERROR during vision model loading: {e}")

# --- Step 2: Global Setup - RAG Knowledge Base (The 'Remedy Retriever') ---
# We create an in-memory vector store with remedy information. This also runs only once.

print("Building the RAG knowledge base for remedies...")
try:
    # In a real enterprise app, this text would come from a database or file system.
    remedy_knowledge_text = """
    # Maize Health Guide
    ## Phosphorus Deficiency / Purple Leaf Disease
    **Symptoms:** Plants are often stunted, dark green or purplish. This condition is sometimes called Purple leaf disease. The purplish discoloration is most evident on the tips of leaves of young plants.
    **Cause:** Lack of available phosphorus in the soil, often due to incorrect pH or cold soil temperatures.
    **Local Remedy:** Apply a high-phosphorus starter fertilizer at planting time. In organic systems, bone meal or rock phosphate are excellent sources. A foliar spray of a water-soluable phosphate fertilizer can be applied directly to the leaves.
    ## Nitrogen Deficiency
    **Symptoms:** General yellowing (chlorosis) of the plant, starting with the lower, older leaves in a distinct "V" shape.
    **Cause:** Insufficient nitrogen in the soil.
    **Local Remedy:** Side-dress with a nitrogen-rich fertilizer like urea or composted manure.
    ## Leaf Spot
    **Symptoms:** Small, circular to oval spots with tan centers and dark borders on the leaves. This is a common fungal disease.
    **Local Remedy:** Use resistant maize varieties if available. Apply appropriate fungicides if the infection is severe, following label directions. Improve air circulation and destroy infected crop debris after harvest.
    ## Leaf Blight
    **Symptoms:** Large, irregular, grayish-green or tan lesions on the leaves. These lesions can grow and merge, leading to significant leaf death.
    **Cause:** Fungal pathogens that thrive in warm, humid conditions.
    **Local Remedy:** The primary defense is planting resistant hybrids. Fungicide applications may be necessary and should be timed based on disease scouting. Crop rotation and tilling under infected residue can help reduce the pathogen for the next season.
    ## Healthy Plant
    **Symptoms:** Vigorous growth with lush, uniformly dark green leaves.
    **Local Remedy:** No remedy needed. Maintain good agricultural practices.
    """

    
    # Create a temporary file to load the text
    with open("knowledge.txt", "w") as f:
        f.write(remedy_knowledge_text)
    
    loader = TextLoader("knowledge.txt")
    documents = loader.load()

    # Split documents into chunks
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
    docs = text_splitter.split_documents(documents)

    # Create embeddings model
    embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")

    # Create FAISS vector store from documents
    db = FAISS.from_documents(docs, embeddings)
    
    # Create the retriever object
    RETRIEVER = db.as_retriever(search_kwargs={"k": 2}) # Retrieves the top 2 most relevant chunks
    
    print("✅ RAG knowledge base and retriever created successfully!")

except Exception as e:
    print(f"❌ CRITICAL ERROR during RAG setup: {e}")

# --- Step 3: Global Setup - The LangChain Agent (The 'Synthesizer') ---
# The agent's "brain" is a powerful LLM from the Hugging Face Hub.
# You must set your HF_TOKEN in the Space's secrets for this to work.

print("Initializing the LangChain agent...")
try:
    # *** THE FIX IS APPLIED HERE ***
    # 1. Explicitly get the token from the environment secrets.
    hf_token = os.environ.get("HF_TOKEN")
    if not hf_token:
        # This provides a clear error in the logs if the secret is missing.
        raise ValueError("HF_TOKEN secret not found. Please add it in your Space settings.")
    else:
        print("✅ HF_TOKEN secret found successfully.")


    prompt = ChatPromptTemplate.from_messages(
        [
            ("system", "You are an expert agricultural assistant. Your goal is to provide a clear, actionable guide based on a plant's diagnosis and retrieved context. Synthesize the information into a helpful, well-structured response."),
            ("human", "My plant has the following diagnosis: {diagnosis}\n\nUsing this information, please provide a remedy based on the following retrieved context:\n\nCONTEXT:\n{context}")
        ]        
    )
    print(f"DEBUG: Created prompt of type: {type(prompt)}")

    base_llm = HuggingFaceEndpoint(
        repo_id="HuggingFaceH4/zephyr-7b-beta",
        huggingfacehub_api_token=hf_token,
        max_new_tokens=512,
        temperature=0.2,
    )
    
    # 2. Wrap the base LLM with ChatHuggingFace, passing the base_llm as the required 'llm' field.
    chat_model = ChatHuggingFace(llm=base_llm)
    # This is the full LCEL (LangChain Expression Language) chain.
    # It defines the entire multi-step workflow.
    AGENT_CHAIN = (
        {"context": RETRIEVER, "diagnosis": RunnablePassthrough()}
        | prompt | chat_model | StrOutputParser()
    )
    print("✅ LangChain agent initialized successfully!")

except Exception as e:
    print(f"❌ CRITICAL ERROR during LangChain agent initialization: {e}")
    AGENT_CHAIN = None

# --- Step 4: Define the Core Functions for the Gradio App ---

def diagnose_plant_from_image(uploaded_image: Image.Image) -> str:
    """
    This function takes a PIL Image and runs ONLY the vision model diagnosis.
    It's the first step in our agentic workflow.
    """
    if VISION_MODEL is None or PROCESSOR is None or uploaded_image is None:
        return "ERROR: Vision model not loaded."

    image = uploaded_image.convert("RGB")
    messages = [
        {"role": "user", "content": [{"type": "text", "text": "What is the condition of this maize plant? Provide only the name of the condition."}, {"type": "image", "image": image}]}
    ]
    text_prompt = PROCESSOR.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    inputs = PROCESSOR(text=text_prompt, images=image, return_tensors="pt").to(VISION_MODEL.device)

    with torch.inference_mode():
        outputs = VISION_MODEL.generate(**inputs, max_new_tokens=48, use_cache=True)
    
    response = PROCESSOR.batch_decode(outputs, skip_special_tokens=True)[0]
    answer_start_index = response.rfind("model\n")
    if answer_start_index != -1:
        return response[answer_start_index + len("model\n"):].strip()
    return "Could not parse diagnosis."


def run_full_analysis(uploaded_image: Image.Image) -> str:
    """
    This is the main function for the Gradio interface. It orchestrates the entire
    diagnosis-to-remedy workflow.
    """
    if uploaded_image is None:
        return "Please upload an image of a maize plant."

    # Step 1: Get the initial diagnosis from the vision model.
    print("Workflow Step 1: Running diagnosis...")
    diagnosis = diagnose_plant_from_image(uploaded_image)
    print(f"Diagnosis received: {diagnosis}")

    if "ERROR" in diagnosis or "Could not parse" in diagnosis:
        return f"Sorry, I couldn't identify the condition from the image. Raw output: {diagnosis}"

    # Step 2: Invoke the LangChain agent with the diagnosis to get the remedy.
    print("Workflow Step 2: Invoking agent for remedy retrieval...")
    if AGENT_CHAIN:
        try:
            retrieved_docs = RETRIEVER.invoke(diagnosis)
            print(f"DEBUG: Retrieved docs for '{diagnosis}': {retrieved_docs}")
            response = AGENT_CHAIN.invoke(diagnosis)
            print("Agent invocation successful.")
            return response
        except Exception as e:
            print(f"Agent invocation failed: {e}")
            return "The diagnosis was successful, but I failed to retrieve a remedy. Please check the logs."
    else:
        return "ERROR: The main analysis agent is not available."


# --- Step 5: Build and Launch the Gradio Interface ---

print("Building Gradio interface...")


demo = gr.Interface(
    fn=run_full_analysis,
    inputs=gr.Image(type="pil", label="Upload Maize Plant Image"),
    outputs=gr.Markdown(label="Full Analysis and Remedy", value="The agent's report will appear here..."),
    title="🌽 Maize Health Automated Diagnosis & Remedy Agent",
    description="**This is an advanced, multi-step agent.** Upload an image of a maize plant. The AI will first **diagnose** the condition using a fine-tuned vision model, then **retrieve** a detailed treatment plan from its knowledge base.",
    article="Built with Unsloth, LangChain, and Gradio. Powered by Gemma-3N and Zephyr-7B.",
    allow_flagging="never",
)

print("Launching Gradio app...")
demo.launch(share=True)