app.py

# app.py

import os
import pandas as pd
import chardet
import logging
import gradio as gr
import json
import hashlib
import numpy as np
from typing import Optional, List, Tuple, ClassVar, Dict

from sentence_transformers import SentenceTransformer, util, CrossEncoder
from langchain.llms.base import LLM
import google.generativeai as genai

# Import smolagents components
from smolagents import CodeAgent, LiteLLMModel, DuckDuckGoSearchTool, ManagedAgent

###############################################################################
# 1) Logging Setup
###############################################################################
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("Daily Wellness AI")

###############################################################################
# 2) API Key Handling and LiteLLMModel Instantiation
###############################################################################
def clean_api_key(key: str) -> str:
    """Remove non-ASCII characters and strip whitespace from the API key."""
    return ''.join(c for c in key if ord(c) < 128).strip()

gemini_api_key = os.environ.get("GEMINI_API_KEY")
if not gemini_api_key:
    logger.error("GEMINI_API_KEY environment variable not set.")
    raise EnvironmentError("Please set the GEMINI_API_KEY environment variable.")

gemini_api_key = clean_api_key(gemini_api_key)
logger.info("GEMINI API Key loaded successfully.")

# Instantiate the model using LiteLLMModel
llm = LiteLLMModel(model_id="gemini/gemini-pro", api_key=gemini_api_key)

###############################################################################
# 3) CSV Loading and Processing
###############################################################################
def load_csv(file_path: str):
    try:
        if not os.path.isfile(file_path):
            logger.error(f"CSV file does not exist: {file_path}")
            return [], []
        
        with open(file_path, 'rb') as f:
            result = chardet.detect(f.read())
            encoding = result['encoding']

        data = pd.read_csv(file_path, encoding=encoding)
        if 'Question' not in data.columns or 'Answers' not in data.columns:
            raise ValueError("CSV must contain 'Question' and 'Answers' columns.")
        data = data.dropna(subset=['Question', 'Answers'])

        logger.info(f"Loaded {len(data)} entries from {file_path}")
        return data['Question'].tolist(), data['Answers'].tolist()
    except Exception as e:
        logger.error(f"Error loading CSV: {e}")
        return [], []

csv_file_path = "AIChatbot.csv"
corpus_questions, corpus_answers = load_csv(csv_file_path)
if not corpus_questions:
    raise ValueError("Failed to load the knowledge base.")

###############################################################################
# 4) Sentence Embeddings & Cross-Encoder
###############################################################################
embedding_model_name = "sentence-transformers/multi-qa-mpnet-base-dot-v1"
try:
    embedding_model = SentenceTransformer(embedding_model_name)
    logger.info(f"Loaded embedding model: {embedding_model_name}")
except Exception as e:
    logger.error(f"Failed to load embedding model: {e}")
    raise e

try:
    question_embeddings = embedding_model.encode(corpus_questions, convert_to_tensor=True)
    logger.info("Encoded question embeddings successfully.")
except Exception as e:
    logger.error(f"Failed to encode question embeddings: {e}")
    raise e

cross_encoder_name = "cross-encoder/ms-marco-MiniLM-L-6-v2"
try:
    cross_encoder = CrossEncoder(cross_encoder_name)
    logger.info(f"Loaded cross-encoder model: {cross_encoder_name}")
except Exception as e:
    logger.error(f"Failed to load cross-encoder model: {e}")
    raise e

###############################################################################
# 5) Retrieval + Re-Ranking
###############################################################################
class EmbeddingRetriever:
    def __init__(self, questions, answers, embeddings, model, cross_encoder):
        self.questions = questions
        self.answers = answers
        self.embeddings = embeddings
        self.model = model
        self.cross_encoder = cross_encoder

    def retrieve(self, query: str, top_k: int = 3):
        try:
            query_embedding = self.model.encode(query, convert_to_tensor=True)
            scores = util.pytorch_cos_sim(query_embedding, self.embeddings)[0].cpu().tolist()
            scored_data = sorted(zip(self.questions, self.answers, scores), key=lambda x: x[2], reverse=True)[:top_k]

            cross_inputs = [[query, candidate[0]] for candidate in scored_data]
            cross_scores = self.cross_encoder.predict(cross_inputs)

            reranked = sorted(zip(scored_data, cross_scores), key=lambda x: x[1], reverse=True)
            final_retrieved = [(entry[0][1], entry[1]) for entry in reranked]
            logger.debug(f"Retrieved and reranked answers: {final_retrieved}")
            return final_retrieved
        except Exception as e:
            logger.error(f"Error during retrieval: {e}")
            logger.debug("Exception details:", exc_info=True)
            return []

retriever = EmbeddingRetriever(corpus_questions, corpus_answers, question_embeddings, embedding_model, cross_encoder)

###############################################################################
# 6) Sanity Check Tool
###############################################################################
class QuestionSanityChecker:
    def __init__(self, llm):
        self.llm = llm

    def is_relevant(self, question: str) -> bool:
        prompt = (
            f"You are an assistant that determines whether a question is relevant to daily wellness.\n\n"
            f"Question: {question}\n\n"
            f"Is the above question relevant to daily wellness? Respond with 'Yes' or 'No' only."
        )
        try:
            response = self.llm(prompt)
            is_yes = 'yes' in response.lower()
            is_no = 'no' in response.lower()
            logger.debug(f"Sanity check response: '{response}', interpreted as is_yes={is_yes}, is_no={is_no}")
            if is_yes and not is_no:
                return True
            elif is_no and not is_yes:
                return False
            else:
                logger.warning(f"Sanity check ambiguous response: '{response}'. Defaulting to 'No'.")
                return False
        except Exception as e:
            logger.error(f"Error in sanity check: {e}")
            logger.debug("Exception details:", exc_info=True)
            return False

sanity_checker = QuestionSanityChecker(llm)

###############################################################################
# 7) smolagents Integration: GROQ Model and Web Search
###############################################################################
# Initialize the smolagents' LiteLLMModel with GROQ model (already instantiated as llm if needed elsewhere)

# Instantiate the DuckDuckGo search tool
search_tool = DuckDuckGoSearchTool()

# Create the web agent with the search tool
web_agent = CodeAgent(
    tools=[search_tool],
    model=llm  # Use the direct model for web queries if applicable
)

# Define the managed web agent
managed_web_agent = ManagedAgent(
    agent=web_agent,
    name="web_search",
    description="Runs a web search for you. Provide your query as an argument."
)

# Create the manager agent with managed web agent and additional tools if needed
manager_agent = CodeAgent(
    tools=[],  # Add additional tools here if required
    model=llm,
    managed_agents=[managed_web_agent]
)

###############################################################################
# 8) Answer Expansion
###############################################################################
class AnswerExpander:
    def __init__(self, llm):
        self.llm = llm

    def expand(self, query: str, retrieved_answers: List[str], detail: bool = False) -> str:
        try:
            reference_block = "\n".join(
                f"- {idx+1}) {ans}" for idx, ans in enumerate(retrieved_answers, start=1)
            )

            detail_instructions = (
                "Provide a thorough, in-depth explanation, adding relevant tips and context, "
                "while remaining creative and brand-aligned. "  
                if detail else
                "Provide a concise response in no more than 4 sentences."
            )

            prompt = (
                f"You are Daily Wellness AI, a friendly wellness expert. Below are multiple "
                f"potential answers retrieved from a local knowledge base. You have a user question.\n\n"
                f"Question: {query}\n\n"
                f"Retrieved Answers:\n{reference_block}\n\n"
                f"Please synthesize these references into a single cohesive, creative, and brand-aligned response. "
                f"{detail_instructions} "
                f"End with a short inspirational note.\n\n"
                "Disclaimer: This is general wellness information, not a substitute for professional medical advice."
            )

            logger.debug(f"Generated prompt for answer expansion: {prompt}")
            response = self.llm(prompt)
            logger.debug(f"Expanded answer: {response}")
            return response.strip()
        except Exception as e:
            logger.error(f"Error expanding answer: {e}")
            logger.debug("Exception details:", exc_info=True)
            return "Sorry, an error occurred while generating a response."

answer_expander = AnswerExpander(llm)

###############################################################################
# 9) Persistent Cache (ADDED)
###############################################################################
CACHE_FILE = "query_cache.json"
SIMILARITY_THRESHOLD_CACHE = 0.8

def load_cache() -> Dict:
    if os.path.isfile(CACHE_FILE):
        try:
            with open(CACHE_FILE, "r", encoding="utf-8") as f:
                return json.load(f)
        except Exception as e:
            logger.error(f"Failed to load cache file: {e}")
            return {}
    return {}

def save_cache(cache_data: Dict):
    try:
        with open(CACHE_FILE, "w", encoding="utf-8") as f:
            json.dump(cache_data, f, ensure_ascii=False, indent=2)
    except Exception as e:
        logger.error(f"Failed to save cache file: {e}")

def compute_hash(text: str) -> str:
    return hashlib.md5(text.encode("utf-8")).hexdigest()

cache_store = load_cache()

###############################################################################
# 9.1) Utility to attempt cached retrieval (ADDED)
###############################################################################
def get_cached_answer(query: str) -> Optional[str]:
    if not cache_store:
        return None

    query_embedding = embedding_model.encode(query, convert_to_tensor=True)

    best_score = 0.0
    best_answer = None

    for cached_q, cache_data in cache_store.items():
        stored_embedding = np.array(cache_data["embedding"], dtype=np.float32)
        score = util.pytorch_cos_sim(query_embedding, stored_embedding)[0].item()
        if score > best_score:
            best_score = score
            best_answer = cache_data["answer"]

    if best_score >= SIMILARITY_THRESHOLD_CACHE:
        logger.info(f"Cache hit! Similarity: {best_score:.2f}, returning cached answer.")
        return best_answer
    return None

def store_in_cache(query: str, answer: str):
    query_embedding = embedding_model.encode(query, convert_to_tensor=True).cpu().tolist()
    cache_key = compute_hash(query)
    cache_store[cache_key] = {
        "query": query,
        "answer": answer,
        "embedding": query_embedding
    }
    save_cache(cache_store)

###############################################################################
# 10) Query Handling
###############################################################################
def handle_query(query: str, detail: bool = False) -> str:
    if not query or not isinstance(query, str) or len(query.strip()) == 0:
        return "Please provide a valid question."

    try:
        is_relevant = sanity_checker.is_relevant(query)
        if not is_relevant:
            return "Your question seems out of context or not related to daily wellness. Please ask a wellness-related question."

        retrieved = retriever.retrieve(query)
        cached_answer = get_cached_answer(query)

        if not retrieved:
            if cached_answer:
                logger.info("No relevant entries found in knowledge base. Returning cached answer.")
                return cached_answer
            return "I'm sorry, I couldn't find an answer to your question."

        top_score = retrieved[0][1]
        similarity_threshold = 0.3

        if top_score < similarity_threshold:
            logger.info("Similarity score below threshold. Performing web search.")
            web_search_response = manager_agent.run(query)
            logger.debug(f"Web search response: {web_search_response}")

            if cached_answer:
                blend_prompt = (
                    f"Combine the following previous answer with the new web results to create a more creative and accurate response. "
                    f"Do not include any of the previous prompt or instructions in your response. "
                    f"Add positivity and conclude with a short inspirational note.\n\n"
                    f"Previous Answer:\n{cached_answer}\n\n"
                    f"Web Results:\n{web_search_response}"
                )
                final_answer = llm(blend_prompt).strip()
            else:
                final_answer = (
                    f"**Daily Wellness AI**\n\n"
                    f"{web_search_response}\n\n"
                    "Disclaimer: This information is retrieved from the web and is not a substitute for professional medical advice.\n\n"
                    "Wishing you a calm and wonderful day!"
                )

            store_in_cache(query, final_answer)
            return final_answer

        responses = [ans for ans, score in retrieved]

        if cached_answer:
            blend_prompt = (
                f"Combine the previous answer with the newly retrieved answers to enhance creativity and accuracy. "
                f"Do not include any of the previous prompt or instructions in your response. "
                f"Add new insights, creativity, and conclude with a short inspirational note.\n\n"
                f"Previous Answer:\n{cached_answer}\n\n"
                f"New Retrieved Answers:\n" + "\n".join(f"- {r}" for r in responses)
            )
            final_answer = llm(blend_prompt).strip()
        else:
            final_answer = answer_expander.expand(query, responses, detail=detail)

        store_in_cache(query, final_answer)
        return final_answer

    except Exception as e:
        logger.error(f"Error handling query: {e}")
        logger.debug("Exception details:", exc_info=True)
        return "An error occurred while processing your request."

###############################################################################
# 11) Gradio Interface
###############################################################################
def gradio_interface(query: str, detail: bool):
    try:
        response = handle_query(query, detail=detail)
        formatted_response = response
        return formatted_response
    except Exception as e:
        logger.error(f"Error in Gradio interface: {e}")
        logger.debug("Exception details:", exc_info=True)
        return "**An error occurred while processing your request. Please try again later.**"

interface = gr.Interface(
    fn=gradio_interface,
    inputs=[
        gr.Textbox(
            lines=2,
            placeholder="e.g., What is box breathing?",
            label="Ask Daily Wellness AI"
        ),
        gr.Checkbox(
            label="In-Depth Answer?",
            value=False,
            info="Check for a longer, more detailed response."
        )
    ],
    outputs=gr.Markdown(label="Answer from Daily Wellness AI"),
    title="Daily Wellness AI",
    description="Ask wellness-related questions and receive synthesized, creative answers. Optionally request a more in-depth response.",
    theme="default",
    examples=[
        ["What is box breathing and how does it help reduce anxiety?", True],
        ["Provide a daily wellness schedule incorporating box breathing techniques.", False],
        ["What are some tips for maintaining good posture while working at a desk?", True],
        ["Who is the CEO of Hugging Face?", False]
    ],
    allow_flagging="never"
)

###############################################################################
# 12) Launch Gradio
###############################################################################
if __name__ == "__main__":
    try:
        interface.launch(server_name="0.0.0.0", server_port=7860, debug=False, share=True)
    except Exception as e:
        logger.error(f"Failed to launch Gradio interface: {e}")
        logger.debug("Exception details:", exc_info=True)