app.py

import gradio as gr
from langchain_community.vectorstores import Chroma
from langchain_core.documents import Document
from langchain_huggingface import HuggingFaceEmbeddings
from openai import OpenAI
import json
import os

# -------------------------------
# CONFIGURATION
# -------------------------------
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY", "sk-your-key-here")

# -------------------------------
# LOAD DATA AND BUILD VECTOR STORE
# -------------------------------
def clean_metadata(metadata):
    """Convert list values to comma-separated strings for ChromaDB compatibility"""
    cleaned = {}
    for key, value in metadata.items():
        if isinstance(value, list):
            cleaned[key] = ", ".join(str(v) for v in value)
        elif isinstance(value, (str, int, float, bool)) or value is None:
            cleaned[key] = value
        else:
            cleaned[key] = str(value)
    return cleaned

print("Loading documents...")
docs = []
with open("helpwildlife_rag.jsonl", "r", encoding="utf-8") as f:
    for line in f:
        entry = json.loads(line)
        metadata = entry.get("metadata", {})
        docs.append(Document(
            page_content=entry["text"],
            metadata=clean_metadata(metadata)
        ))
print(f"✓ Loaded {len(docs)} documents")

print("Loading embedding model...")
embeddings = HuggingFaceEmbeddings(
    model_name="sentence-transformers/all-MiniLM-L6-v2"
)

print("Building vector store...")
try:
    vectorstore = Chroma.from_documents(docs, embedding=embeddings)
except TypeError:
    vectorstore = Chroma.from_documents(docs, embedding_function=embeddings)

retriever = vectorstore.as_retriever(search_kwargs={"k": 3})
print("✓ Vector store created")

# Initialize OpenAI client
client = OpenAI(api_key=OPENAI_API_KEY)
print("✓ OpenAI client ready")

# Improved, more conversational prompt
SYSTEM_PROMPT = """You are a compassionate wildlife advice assistant having a conversation with someone concerned about an animal.

CRITICAL FORMATTING RULES:
- Write in short, conversational paragraphs (2-3 sentences each)
- NEVER use asterisks (*) or bullet points
- Use numbered lists ONLY when giving step-by-step instructions
- Keep your initial response brief and focused
- Ask 1-2 clarifying questions to better understand the situation

Your approach:
1. First, acknowledge their concern warmly
2. Ask specific questions to assess the situation (e.g., "Is the animal bleeding or visibly injured?", "Can you see the parents nearby?", "Is it in immediate danger from traffic or predators?")
3. Based on the context provided, give concise, practical first steps
4. Always emphasize: contact a local wildlife rescue or 24-hour animal hospital for the safest outcome

Key principles:
- Never suggest killing or harming any animal
- Discourage people from handling wildlife themselves
- Note that regular vets may euthanize wild animals - wildlife rescues are better
- All animals deserve care and respect - no animal is a "pest"
- Use empathy and clarity

Use ONLY the information in the CONTEXT section below. If the context doesn't have enough information, say so and suggest contacting a wildlife rescue."""

# -------------------------------
# CHAT FUNCTION WITH HISTORY
# -------------------------------
def chat_wildlife(message, history):
    """Process a wildlife question with conversation history"""
    if not message.strip():
        return history + [["", "Please tell me about the animal you've found, and I'll help you figure out what to do."]]
    
    try:
        # Retrieve relevant documents
        retrieved_docs = retriever.invoke(message)
        
        # Format context from retrieved documents
        context = "\n\n".join([doc.page_content for doc in retrieved_docs])
        
        # Build conversation history for context
        messages = [{"role": "system", "content": SYSTEM_PROMPT}]
        
        # Add conversation history
        for user_msg, assistant_msg in history:
            if user_msg:
                messages.append({"role": "user", "content": user_msg})
            if assistant_msg:
                messages.append({"role": "assistant", "content": assistant_msg})
        
        # Add current message with context
        current_message = f"""CONTEXT (from HelpWildlife guidance):
{context}

QUESTION: {message}"""
        
        messages.append({"role": "user", "content": current_message})
        
        # Call OpenAI API
        response = client.chat.completions.create(
            model="gpt-4o-mini",
            messages=messages,
            temperature=0.3,
            max_tokens=300
        )
        
        answer = response.choices[0].message.content
        history.append([message, answer])
        return history
        
    except Exception as e:
        error_msg = f"I'm having trouble connecting right now. Please check that the OpenAI API key is set correctly, or try again in a moment."
        history.append([message, error_msg])
        return history

# -------------------------------
# CUSTOM CSS - CHATGPT STYLE
# -------------------------------
custom_css = """
/* Import ChatGPT-like font */
@import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;500;600&display=swap');

* {
    font-family: 'Inter', -apple-system, BlinkMacSystemFont, 'Segoe UI', sans-serif !important;
}

.gradio-container {
    max-width: 900px !important;
    margin: 0 auto !important;
}

/* Main chat area styling */
.chatbot {
    border-radius: 12px !important;
    border: 1px solid #e5e7eb !important;
}

/* Message bubbles */
.message-wrap {
    padding: 20px !important;
}

.message {
    font-size: 15px !important;
    line-height: 1.6 !important;
    color: #374151 !important;
}

/* User message */
.user {
    background: #f7f7f8 !important;
    border-radius: 12px !important;
    padding: 12px 16px !important;
}

/* Assistant message */
.bot {
    background: transparent !important;
    padding: 12px 16px !important;
}

/* Input box */
.input-box {
    border-radius: 12px !important;
    border: 1px solid #d1d5db !important;
    font-size: 15px !important;
    padding: 12px 16px !important;
}

/* Buttons */
button {
    border-radius: 8px !important;
    font-weight: 500 !important;
    font-size: 14px !important;
    transition: all 0.2s !important;
}

button:hover {
    transform: translateY(-1px) !important;
    box-shadow: 0 4px 12px rgba(0,0,0,0.1) !important;
}

/* Primary button (Send) */
.primary {
    background: #10a37f !important;
    color: white !important;
    border: none !important;
}

.primary:hover {
    background: #0d8a6a !important;
}

/* Secondary buttons */
.secondary {
    background: white !important;
    border: 1px solid #d1d5db !important;
    color: #374151 !important;
}

/* Example buttons */
.sm {
    background: white !important;
    border: 1px solid #d1d5db !important;
    color: #374151 !important;
    padding: 8px 16px !important;
}

/* Headers */
h1, h2, h3 {
    font-weight: 600 !important;
    color: #111827 !important;
}

h1 {
    font-size: 32px !important;
    margin-bottom: 8px !important;
}

h3 {
    font-size: 16px !important;
    font-weight: 500 !important;
    color: #6b7280 !important;
    margin-top: 24px !important;
    margin-bottom: 12px !important;
}

/* Markdown text */
.markdown {
    color: #374151 !important;
    font-size: 15px !important;
    line-height: 1.6 !important;
}
"""

# -------------------------------
# GRADIO INTERFACE
# -------------------------------
with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        # 🦔 Wildlife Rescue Assistant
        
        Found an animal that might need help? I'm here to guide you with compassionate, evidence-based advice.
        
        **⚠️ For emergencies:** Contact your local wildlife rescue or 24-hour animal hospital immediately.
        """
    )
    
    chatbot = gr.Chatbot(
        height=500,
        show_label=False,
        avatar_images=(None, "🦔"),
        bubble_full_width=False
    )
    
    with gr.Row():
        with gr.Column(scale=9):
            msg = gr.Textbox(
                placeholder="Describe what you've found (e.g., 'I found a baby bird on the ground')...",
                show_label=False,
                container=False
            )
        with gr.Column(scale=1, min_width=80):
            submit = gr.Button("Send", variant="primary")
    
    gr.Markdown(
        """
        ### Common situations:
        """
    )
    
    with gr.Row():
        example1 = gr.Button("🐦 Baby bird on ground", size="sm")
        example2 = gr.Button("🦔 Hedgehog in daytime", size="sm")
        example3 = gr.Button("🦊 Injured fox", size="sm")
    
    clear = gr.Button("🔄 Start new conversation", variant="secondary")
    
    # Event handlers
    def respond(message, chat_history):
        return chat_wildlife(message, chat_history), ""
    
    msg.submit(respond, [msg, chatbot], [chatbot, msg])
    submit.click(respond, [msg, chatbot], [chatbot, msg])
    
    example1.click(lambda: ("I found a baby bird on the ground", []), None, [msg, chatbot])
    example2.click(lambda: ("I found a hedgehog out during the day", []), None, [msg, chatbot])
    example3.click(lambda: ("There's an injured fox in my garden", []), None, [msg, chatbot])
    
    clear.click(lambda: [], None, chatbot)
    
    gr.Markdown(
        """
        ---
        *This tool provides general guidance based on wildlife rescue best practices. 
        Always consult with a wildlife rescue professional for specific situations.*
        """
    )

if __name__ == "__main__":
    demo.launch()