app.py

import gradio as gr
from llm_agent import AgriculturalAgent
import os

agent = None


def initialize_agent(api_key: str) -> str:
    """
    Initialize the agent with the provided API key.
    
    Args:
        api_key: Groq API key
        
    Returns:
        Status message
    """
    global agent
    if not api_key or not api_key.strip():
        return "Please enter a valid API key."
    
    try:
        agent = AgriculturalAgent(api_key=api_key.strip())
        return "Agent initialized successfully! Dataset loaded and ready. You can now start asking questions."
    except Exception as e:
        return f"Error initializing agent: {str(e)}"


def chat_with_agent(message: str, history: list, api_key: str) -> tuple:
    """
    Handle chat messages with the agent.
    
    Args:
        message: User's message
        history: Chat history (Gradio format: list of tuples (user_msg, assistant_msg))
        api_key: Groq API key
        
    Returns:
        Tuple of (empty string, updated history in Gradio format)
    """
    global agent
    
    # Ensure history is a list
    if not isinstance(history, list):
        history = []
    
    # Convert history to tuple format (compatible with Gradio 4.0.0+)
    converted_history = []
    for item in history:
        if isinstance(item, (tuple, list)) and len(item) == 2:
            # Already in tuple format (user_msg, assistant_msg)
            converted_history.append((str(item[0]) if item[0] else "", str(item[1]) if item[1] else ""))
        elif isinstance(item, dict) and "role" in item and "content" in item:
            # Convert dict format to tuple format
            # This handles dicts but we need to pair them properly
            # For now, just skip dicts and rebuild from tuples
            pass
    
    history = converted_history
    
    # Skip if message is empty
    if not message or not message.strip():
        return "", history
    
    print(f"\n[APP] Received message: {message[:100]}...") 
    
    if agent is None:
        if not api_key or not api_key.strip():
            new_history = list(history)
            new_history.append((str(message), "Please initialize the agent with an API key first."))
            return "", new_history
        print("[APP] Initializing agent...")
        init_msg = initialize_agent(api_key)
        if "Error" in init_msg:
            new_history = list(history)
            new_history.append((str(message), str(init_msg)))
            return "", new_history
    
    if agent is None:
        new_history = list(history)
        new_history.append((str(message), "Please initialize the agent with an API key first."))
        return "", new_history
    
    print("[APP] Sending message to agent...")
    
    try:
        response = agent.chat(message)
        print(f"[APP] Received response from agent: {response[:100]}...")
        
        # Ensure both message and response are strings
        if not isinstance(response, str):
            response = str(response)
        if not isinstance(message, str):
            message = str(message)
        
        # Append as tuple format (Gradio 4.0.0+ expects tuples)
        new_history = list(history)
        new_history.append((message, response))
        return "", new_history
    except Exception as e:
        error_msg = f"Error: {str(e)}"
        print(f"[APP] ERROR: {error_msg}")
        new_history = list(history)
        new_history.append((str(message), error_msg))
        return "", new_history


def reset_chat():
    """Reset the chat conversation."""
    global agent
    if agent:
        agent.reset_conversation()
    return []  


try:
    theme = gr.themes.Soft()
    blocks_kwargs = {"title": "Agricultural Research AI Agent", "theme": theme}
except (AttributeError, ImportError):
    blocks_kwargs = {"title": "Agricultural Research AI Agent"}

with gr.Blocks(**blocks_kwargs) as demo:
    gr.Markdown("""
    # 🌾 Agricultural Research AI Agent
    
    This AI agent helps you explore and find information from a comprehensive collection of 
    **45,232 agricultural research publications** from CGIAR (Consultative Group on International Agricultural Research).
    
    ## Features
    - 🔍 Search for research documents on specific agricultural topics
    - 📚 Browse documents by topic (crop management, pest control, climate adaptation, etc.)
    - 📄 Get detailed information about specific research papers
    - 💡 Get insights and answers based on the latest agricultural research
    
    ## How to Use
    1. Enter your Groq API key in the field below (get a free key at https://console.groq.com)
    2. Click "Initialize Agent" to start
    3. Ask questions about agricultural topics, search for documents, or browse research papers
    
    **Example questions:**
    - "What research is available on rice cultivation?"
    - "Find documents about pest control in agriculture"
    - "Tell me about climate adaptation strategies for small-scale farmers"
    - "What does the dataset contain?"
    
    ---
    """)
    
    with gr.Row():
        with gr.Column(scale=3):
            api_key_input = gr.Textbox(
                label="Groq API Key",
                placeholder="Enter your Groq API key here...",
                type="password",
                info="Your API key is not stored and is only used for this session. Get a free API key at https://console.groq.com"
            )
            init_btn = gr.Button("Initialize Agent", variant="primary")
            init_status = gr.Textbox(label="Status", interactive=False)
        
        with gr.Column(scale=1):
            reset_btn = gr.Button("Reset Chat", variant="secondary")
    
    chatbot = gr.Chatbot(
        label="Chat with Agricultural Research Agent",
        height=500
    )
    
    with gr.Row():
        msg = gr.Textbox(
            label="Your Question",
            placeholder="Ask about agricultural research, search for documents, or browse topics...",
            scale=4,
            lines=2
        )
        submit_btn = gr.Button("Send", variant="primary", scale=1)
    
    init_btn.click(
        fn=initialize_agent,
        inputs=[api_key_input],
        outputs=[init_status]
    )
    
    submit_btn.click(
        fn=chat_with_agent,
        inputs=[msg, chatbot, api_key_input],
        outputs=[msg, chatbot],
        api_name="chat_with_agent_submit"
    )
    
    msg.submit(
        fn=chat_with_agent,
        inputs=[msg, chatbot, api_key_input],
        outputs=[msg, chatbot],
        api_name="chat_with_agent_enter"
    )
    
    reset_btn.click(
        fn=reset_chat,
        outputs=[chatbot]
    )
    
    gr.Markdown("""
    ---
    **Dataset Information:**
    - Source: [CGIAR/gardian-ai-ready-docs](https://huggingface.co/datasets/CGIAR/gardian-ai-ready-docs) on HuggingFace
    - Contains 45,232 structured agricultural research publications
    - Topics: Crop management, pest control, climate adaptation, farming systems, and more
    
    **Note:** This application uses Groq's free API for the LLM agent. Get your free API key at [console.groq.com](https://console.groq.com)
    """)

if __name__ == "__main__":
    demo.launch(share=False, server_name="0.0.0.0", server_port=7860)