app.py

import streamlit as st
import pandas as pd
import chromadb
from sentence_transformers import SentenceTransformer
from transformers import pipeline, AutoModelForQuestionAnswering, AutoTokenizer
from PIL import Image
from io import BytesIO
import requests

# --- 1. Load Recipes Dataset ---
@st.cache_data
def load_recipes():
    try:
        recipes_df = pd.read_csv("recipes.csv")
        recipes_df = recipes_df.rename(columns={"recipe_name": "title", "directions": "instructions"})
        recipes_df = recipes_df[['title', 'ingredients', 'instructions', 'img_src']]
        recipes_df.fillna("", inplace=True)
        recipes_df["ingredients"] = recipes_df["ingredients"].str.lower().str.replace(r'[^\w\s]', '', regex=True)
        recipes_df["combined_text"] = recipes_df["title"] + " " + recipes_df["ingredients"]
        return recipes_df
    except Exception as e:
        st.error(f"⚠ Error loading recipes: {e}")
        return pd.DataFrame()

recipes_df = load_recipes()

# --- 2. Load SentenceTransformer Model ---
@st.cache_resource
def load_embedding_model():
    return SentenceTransformer("all-mpnet-base-v2")

embedding_model = load_embedding_model()

# --- 3. Initialize ChromaDB ---
chroma_client = chromadb.PersistentClient(path="./chroma_db")
collection = chroma_client.get_or_create_collection(name="recipe_collection")

# --- 4. Generate & Store Embeddings ---
def get_sentence_transformer_embeddings(text):
    return embedding_model.encode(text).tolist()

try:
    existing_data = collection.get()
    existing_ids = set(existing_data["ids"]) if existing_data and "ids" in existing_data else set()
except Exception as e:
    st.error(f"⚠ ChromaDB Error: {e}")
    existing_ids = set()

for index, row in recipes_df.iterrows():
    recipe_id = str(index)
    if recipe_id in existing_ids:
        continue
    embedding = get_sentence_transformer_embeddings(row["combined_text"])
    if embedding:
        collection.add(embeddings=[embedding], documents=[row["combined_text"]], ids=[recipe_id])

# --- 5. Retrieve Similar Recipes ---
def retrieve_recipes(query, top_k=3):
    query_embedding = get_sentence_transformer_embeddings(query)
    results = collection.query(query_embeddings=[query_embedding], n_results=top_k)
    
    if results and "ids" in results and results["ids"]:  # Check existence before accessing
        recipe_indices = [int(id) for id in results["ids"][0] if id.isdigit()]
        return recipes_df.iloc[recipe_indices] if recipe_indices else None
    return None

# --- 6. Load a Compatible LLM for Q&A ---
@st.cache_resource
def load_llm_model():
    tokenizer = AutoTokenizer.from_pretrained("deepset/roberta-base-squad2")  # Better Q&A model
    model = AutoModelForQuestionAnswering.from_pretrained("deepset/roberta-base-squad2")
    return pipeline("question-answering", model=model, tokenizer=tokenizer)

llm_model = load_llm_model()

def answer_question(query, context=""):
    greetings = ["hi", "hello", "hii", "hey", "greetings", "how are you", "what's up", "how's it going"]
    
    # Normalize input
    query_cleaned = query.lower().strip()
    
    # Handle greetings first
    if query_cleaned in greetings:
        return "Hello! I'm here to assist with recipes and food-related questions. 🍽️ What would you like to know?"

    # Handle gibberish (random characters, very long non-words)
    if len(query_cleaned) > 12 and not any(char.isdigit() for char in query_cleaned):  
        return "I couldn’t understand your question. Could you rephrase it? 🤔"

    # Attempt to retrieve a recipe-related context
    related_recipes = retrieve_recipes(query, top_k=1)
    if related_recipes is None or related_recipes.empty:
        return "I specialize in recipes! 🍽️ Feel free to ask me about ingredients, cooking methods, or meal ideas. 😊"

    # If a valid recipe is found, use its instructions as context
    context = related_recipes.iloc[0]['instructions']

    # Generate an answer using the LLM
    response = llm_model(question=query, context=context)

    # Ensure we only return meaningful answers
    answer = response.get("answer", "").strip()
    if not answer:
        return "I'm not sure, but I can help with recipes! What would you like to cook?"

    return answer

# --- 8. Few-Shot Classification Function ---
def classify_with_few_shot(query):
    prompt = """
    Classify the following query as one of three types: 
    1. Greeting 
    2. Recipe Search 
    3. Non-Recipe Query
    
    Example Queries:
    - "Hi" -> Greeting
    - "How do I make lasagna?" -> Recipe Search
    - "What is the capital of France?" -> Non-Recipe Query
    
    Query: "{query}"
    Classification:
    """
    
    full_prompt = prompt.format(query=query)
    
    # Use your existing LLM model (or another one) for the few-shot classification
    response = llm_model(question=full_prompt, context="")  # Adjust to how your model processes prompts
    
    classification = response.get("answer", "").strip()  # Assuming the model returns the classification directly
    
    # Check the output and map it to a valid intent
    if "Greeting" in classification:
        return "Greeting"
    elif "Recipe Search" in classification:
        return "Recipe Search"
    elif "Non-Recipe Query" in classification:
        return "Non-Recipe Query"
    else:
        return "Unclassified"
        
# --- 9. Display Image Function ---
def display_image(image_url, recipe_name):
    try:
        if not isinstance(image_url, str) or not image_url.startswith("http"):
            raise ValueError("Invalid or missing image URL")
        response = requests.get(image_url, timeout=5)
        response.raise_for_status()
        image = Image.open(BytesIO(response.content))
        st.image(image, caption=recipe_name, use_container_width=True)
    except requests.exceptions.RequestException as e:
        st.warning(f"⚠ Image fetch error: {e}")
        placeholder_url = "https://via.placeholder.com/300?text=No+Image"
        st.image(placeholder_url, caption=recipe_name, use_container_width=True)

# --- 10. Streamlit UI --- 
st.title("🍽️ AI Recipe & Q&A Assistant")

user_query = st.text_input("Enter your question or recipe search query:", "", key="main_query_input")

if "retrieved_recipes" not in st.session_state:
    st.session_state["retrieved_recipes"] = None

if st.button("Ask AI"):
    if user_query:
        # Classify query using few-shot prompting
        intent = classify_with_few_shot(user_query)  # Updated to use classify_with_few_shot()

        if intent == "Greeting":
            st.subheader("🤖 AI Answer:")
            st.write(answer_question(user_query))

        elif intent == "Recipe Search":
            retrieved_recipes = retrieve_recipes(user_query)
            if retrieved_recipes is not None and not retrieved_recipes.empty:
                st.session_state["retrieved_recipes"] = retrieved_recipes
                st.subheader("🍴 Found Recipes:")
                for index, recipe in retrieved_recipes.iterrows():
                    st.markdown(f"### {recipe['title']}")
                    st.write(f"**Ingredients:** {recipe['ingredients']}")
                    st.write(f"**Instructions:** {recipe['instructions']}")
                    display_image(recipe.get('img_src', ''), recipe['title'])
            else:
                st.warning("⚠️ No relevant recipes found.")
        elif intent == "Non-Recipe Query":
            st.subheader("🤖 AI Answer:")
            st.write("I'm specialized in recipes! Feel free to ask me anything food-related. 😊")
        else:
            st.warning("❌ Unable to classify the query.")