ADDED FEATURES

app.py

@@ -1,102 +1,149 @@
 # app.py
 
 import os
-import getpass
 import pandas as pd
 import chardet
 import logging
 import gradio as gr
+from typing import Optional, List, Tuple, ClassVar, Dict
 
 from sentence_transformers import SentenceTransformer, util, CrossEncoder
-from langchain_community.retrievers import BM25Retriever
-from smolagents import (
-    CodeAgent,
-    HfApiModel,
-    DuckDuckGoSearchTool,
-    Tool,
-    ManagedAgent,
-    LiteLLMModel
-)
+from langchain.llms.base import LLM
+import google.generativeai as genai
+
+###############################################################################
+# 1) Logging Setup
+###############################################################################
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger("Daily Wellness AI")
+
+###############################################################################
+# 2) API Key Handling and Enhanced GeminiLLM Class
+###############################################################################
+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()
+
+# Load the GEMINI API key from environment variables
+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.")
+
+# Configure Google Generative AI
+try:
+    genai.configure(api_key=gemini_api_key)
+    logger.info("Configured Google Generative AI with provided API key.")
+except Exception as e:
+    logger.error(f"Failed to configure Google Generative AI: {e}")
+    raise e
+
+class GeminiLLM(LLM):
+    model_name: ClassVar[str] = "gemini-2.0-flash-exp"
+    temperature: float = 0.7
+    top_p: float = 0.95
+    top_k: int = 40
+    max_tokens: int = 2048
+
+    @property
+    def _llm_type(self) -> str:
+        return "custom_gemini"
+
+    def _call(self, prompt: str, stop: Optional[List[str]] = None) -> str:
+        generation_config = {
+            "temperature": self.temperature,
+            "top_p": self.top_p,
+            "top_k": self.top_k,
+            "max_output_tokens": self.max_tokens,
+        }
+
+        try:
+            logger.debug(f"Initializing GenerativeModel with config: {generation_config}")
+            model = genai.GenerativeModel(
+                model_name=self.model_name,
+                generation_config=generation_config,
+            )
+            logger.debug("GenerativeModel initialized successfully.")
+
+            chat_session = model.start_chat(history=[])
+            logger.debug("Chat session started.")
+
+            response = chat_session.send_message(prompt)
+            logger.debug(f"Prompt sent to model: {prompt}")
+            logger.debug(f"Raw response received: {response.text}")
+
+            return response.text
+        except Exception as e:
+            logger.error(f"Error generating response with GeminiLLM: {e}")
+            logger.debug("Exception details:", exc_info=True)
+            raise e
+
+# Instantiate the GeminiLLM globally
+llm = GeminiLLM()
 
-# --------------------------------------------------------------------------------
-# Set up logging
-# --------------------------------------------------------------------------------
-logging.basicConfig(level=logging.DEBUG)
-logger = logging.getLogger("Daily Wellness AI Guru")
-
-# --------------------------------------------------------------------------------
-# Ensure Hugging Face API Token
-# --------------------------------------------------------------------------------
-# In a Hugging Face Space, you can set HF_API_TOKEN as a secret variable.
-# If it's not set, you could prompt for it locally, but in Spaces,
-# you typically wouldn't do getpass. We'll leave the logic here as fallback.
-if 'HF_API_TOKEN' not in os.environ or not os.environ['HF_API_TOKEN']:
-    os.environ['HF_API_TOKEN'] = getpass.getpass('Enter your Hugging Face API Token: ')
-else:
-    print("HF_API_TOKEN is already set.")
-
-# --------------------------------------------------------------------------------
-# CSV Loading and Processing
-# --------------------------------------------------------------------------------
-def load_csv(file_path):
-    """
-    Load and process a CSV file into two lists: questions and answers.
-    """
+###############################################################################
+# 3) CSV Loading and Processing
+###############################################################################
+def load_csv(file_path: str):
     try:
-        # Detect the encoding of the file
+        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']
 
-        # Load the CSV using the detected encoding
         data = pd.read_csv(file_path, encoding=encoding)
-
-        # Validate that the required columns are present
         if 'Question' not in data.columns or 'Answers' not in data.columns:
-            raise ValueError("The CSV file must contain 'Question' and 'Answers' columns.")
-
-        # Drop any rows with missing values in 'Question' or 'Answers'
+            raise ValueError("CSV must contain 'Question' and 'Answers' columns.")
         data = data.dropna(subset=['Question', 'Answers'])
 
-        # Extract questions and answers
-        questions = data['Question'].tolist()
-        answers = data['Answers'].tolist()
-
-        logger.info(f"Loaded {len(questions)} questions and {len(answers)} answers from {file_path}")
-        return questions, 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 file: {e}")
+        logger.error(f"Error loading CSV: {e}")
         return [], []
 
-# --------------------------------------------------------------------------------
-# Load the AIChatbot.csv file
-# --------------------------------------------------------------------------------
-file_path = "AIChatbot.csv"  # Ensure this file is in the same directory as app.py
-corpus_questions, corpus_answers = load_csv(file_path)
-
+# Path to your CSV file (ensure 'AIChatbot.csv' is in the repository)
+csv_file_path = "AIChatbot.csv"
+corpus_questions, corpus_answers = load_csv(csv_file_path)
 if not corpus_questions:
-    raise ValueError(f"Failed to load questions from {file_path}.")
+    raise ValueError("Failed to load the knowledge base.")
 
-# --------------------------------------------------------------------------------
-# Embedding Model
-# --------------------------------------------------------------------------------
+###############################################################################
+# 4) Sentence Embeddings & Cross-Encoder
+###############################################################################
 embedding_model_name = "sentence-transformers/multi-qa-mpnet-base-dot-v1"
-embedding_model = SentenceTransformer(embedding_model_name)
-logger.info(f"Loaded sentence embedding model: {embedding_model_name}")
-
-# Encode Questions (for retrieval)
-question_embeddings = embedding_model.encode(corpus_questions, convert_to_tensor=True)
-
-# --------------------------------------------------------------------------------
-# Cross-Encoder for Re-Ranking
-# --------------------------------------------------------------------------------
-cross_encoder_model_name = "cross-encoder/ms-marco-MiniLM-L-6-v2"
-cross_encoder = CrossEncoder(cross_encoder_model_name)
-logger.info(f"Loaded cross-encoder model: {cross_encoder_model_name}")
-
-# --------------------------------------------------------------------------------
-# Retrieval + Re-ranking Class
-# --------------------------------------------------------------------------------
+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
@@ -105,204 +152,121 @@ class EmbeddingRetriever:
         self.model = model
         self.cross_encoder = cross_encoder
 
-    def retrieve(self, query, top_k=3):
-        # Compute query embedding
-        query_embedding = self.model.encode(query, convert_to_tensor=True)
-        scores = util.pytorch_cos_sim(query_embedding, self.embeddings)[0].cpu().tolist()
-
-        # Combine data
-        scored_data = list(zip(self.questions, self.answers, scores))
-        # Sort by best scores
-        scored_data = sorted(scored_data, key=lambda x: x[2], reverse=True)
-        # Take top_k
-        top_candidates = scored_data[:top_k]
-
-        # Cross-encode re-rank
-        cross_inputs = [[query, candidate[0]] for candidate in top_candidates]
-        cross_scores = self.cross_encoder.predict(cross_inputs)
-
-        reranked = sorted(
-            zip(top_candidates, cross_scores),
-            key=lambda x: x[1],
-            reverse=True
-        )
+    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]
 
-        # The best candidate
-        best_candidate = reranked[0][0]  # (question, answer, score)
-        best_answer = best_candidate[1]
-        return best_answer
-
-retriever = EmbeddingRetriever(
-    questions=corpus_questions,
-    answers=corpus_answers,
-    embeddings=question_embeddings,
-    model=embedding_model,
-    cross_encoder=cross_encoder
-)
+            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 []
 
-# --------------------------------------------------------------------------------
-# Simple Answer Expander (Without custom sampling parameters)
-# --------------------------------------------------------------------------------
+retriever = EmbeddingRetriever(corpus_questions, corpus_answers, question_embeddings, embedding_model, cross_encoder)
+
+###############################################################################
+# 6) Answer Expansion
+###############################################################################
 class AnswerExpander:
-    def __init__(self, model: HfApiModel):
-        self.model = model
+    def __init__(self, llm: GeminiLLM):
+        self.llm = llm
 
-    def expand(self, question: str, short_answer: str) -> str:
-        """
-        Prompt the LLM to provide a more creative, brand-aligned answer.
-        """
-        prompt = (
-            "You are Daily Wellness AI, a friendly and creative wellness expert. "
-            "The user has a question about well-being. Provide an encouraging, day-to-day "
-            "wellness perspective. Be gentle, uplifting, and brand-aligned.\n\n"
-            f"Question: {question}\n"
-            f"Current short answer: {short_answer}\n\n"
-            "Please rephrase and expand with more detail, wellness tips, daily-life "
-            "applications, and an optimistic tone. Keep it informal, friendly, and end "
-            "with a short inspirational note.\n"
-        )
+    def expand(self, query: str, retrieved_answers: List[str]) -> str:
         try:
-            expanded_answer = self.model.run(prompt)
-            return expanded_answer.strip()
+            reference_block = "\n".join(f"- {idx+1}) {ans}" for idx, ans in enumerate(retrieved_answers, start=1))
+            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"
+                "Please synthesize these references into a single cohesive, creative, "
+                "and brand-aligned response. Add practical tips and positivity, and 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._call(prompt)
+            logger.debug(f"Expanded answer: {response}")
+            return response.strip()
         except Exception as e:
-            logger.error(f"Failed to expand answer: {e}")
-            return short_answer
-
-# NOTE: We are using a basic HfApiModel here (no custom sampling).
-expander_model = HfApiModel()
-answer_expander = AnswerExpander(expander_model)
-
-# --------------------------------------------------------------------------------
-# Enhanced Retriever Tool
-# --------------------------------------------------------------------------------
-from smolagents import Tool
-
-class RetrieverTool(Tool):
-    name = "retriever_tool"
-    description = "Uses semantic search + cross-encoder re-ranking to retrieve the best answer."
-    inputs = {
-        "query": {
-            "type": "string",
-            "description": "User query for retrieving relevant information.",
-        }
-    }
-    output_type = "string"
-
-    def __init__(self, retriever, expander):
-        super().__init__()
-        self.retriever = retriever
-        self.expander = expander
-
-    def forward(self, query):
-        best_answer = self.retriever.retrieve(query, top_k=3)
-        if best_answer:
-            # If short, expand it
-            if len(best_answer.strip()) < 80:
-                logger.info("Answer is short. Expanding with LLM.")
-                best_answer = self.expander.expand(query, best_answer)
-            return best_answer
-        return "No relevant information found."
-
-retriever_tool = RetrieverTool(retriever, answer_expander)
-
-# --------------------------------------------------------------------------------
-# DuckDuckGo (Web) Fallback
-# --------------------------------------------------------------------------------
-search_tool = DuckDuckGoSearchTool()
-
-# --------------------------------------------------------------------------------
-# Managed Agents
-# --------------------------------------------------------------------------------
-from smolagents import ManagedAgent, CodeAgent, LiteLLMModel
-
-retriever_agent = ManagedAgent(
-    agent=CodeAgent(tools=[retriever_tool], model=LiteLLMModel("groq/llama3-8b-8192")),
-    name="retriever_agent",
-    description="Retrieves answers from the local knowledge base (CSV file)."
-)
+            logger.error(f"Error expanding answer: {e}")
+            logger.debug("Exception details:", exc_info=True)
+            return "Sorry, an error occurred while generating a response."
 
-web_agent = ManagedAgent(
-    agent=CodeAgent(tools=[search_tool], model=HfApiModel()),
-    name="web_search_agent",
-    description="Performs web searches if the local knowledge base doesn't have an answer."
-)
+answer_expander = AnswerExpander(llm)
 
-# --------------------------------------------------------------------------------
-# Manager Agent to Orchestrate
-# --------------------------------------------------------------------------------
-manager_agent = CodeAgent(
-    tools=[],
-    model=HfApiModel(),
-    managed_agents=[retriever_agent, web_agent],
-    verbose=True
-)
+###############################################################################
+# 7) Query Handling
+###############################################################################
+def handle_query(query: str) -> str:
+    if not query or not isinstance(query, str) or len(query.strip()) == 0:
+        return "Please provide a valid question."
 
-# --------------------------------------------------------------------------------
-# Gradio Interface
-# --------------------------------------------------------------------------------
-def gradio_interface(query):
     try:
-        logger.info(f"User query: {query}")
-
-        # 1) Query local knowledge base
-        retriever_response = retriever_tool.forward(query)
-        if retriever_response != "No relevant information found.":
-            logger.info("Provided answer from local DB (possibly expanded).")
-            return (
-                f"Hello! This is **Daily Wellness AI**.\n\n"
-                f"{retriever_response}\n\n"
-                "Disclaimer: This is general wellness information, "
-                "not a substitute for professional medical advice.\n\n"
-                "Wishing you a calm and wonderful day!"
-            )
-
-        # 2) Fallback to Web if no relevant local info
-        logger.info("Falling back to web search.")
-        web_response = web_agent.run(query)
-        if web_response:
-            logger.info("Response retrieved from the web.")
-            return (
-                f"Hello! This is **Daily Wellness AI**.\n\n"
-                f"{web_response.strip()}\n\n"
-                "Disclaimer: This is general wellness information, "
-                "not a substitute for professional medical advice.\n\n"
-                "Wishing you a calm and wonderful day!"
-            )
-
-        # 3) Default fallback
-        logger.info("No response found from any source.")
-        return (
-            "Hello! This is **Daily Wellness AI**.\n\n"
-            "I'm sorry, I couldn't find an answer to your question. "
-            "Please try rephrasing or ask something else.\n\n"
-            "Take care, and have a wonderful day!"
+        retrieved = retriever.retrieve(query)
+        if not retrieved:
+            return "I'm sorry, I couldn't find an answer to your question."
+        responses = [ans[0] for ans in retrieved]
+        expanded_answer = answer_expander.expand(query, responses)
+        return expanded_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."
+
+###############################################################################
+# 8) Gradio Interface
+###############################################################################
+def gradio_interface(query: str):
+    try:
+        response = handle_query(query)
+        formatted_response = (
+            f"**Daily Wellness AI**\n\n"
+            f"{response}\n\n"
+            "Disclaimer: This is general wellness information, "
+            "not a substitute for professional medical advice.\n\n"
+            "Wishing you a calm and wonderful day!"
         )
+        return formatted_response
     except Exception as e:
-        logger.error(f"Error processing query: {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.**"
 
-# --------------------------------------------------------------------------------
-# Launch Gradio App
-# --------------------------------------------------------------------------------
 interface = gr.Interface(
     fn=gradio_interface,
     inputs=gr.Textbox(
-        label="Ask Daily Wellness AI",
-        placeholder="e.g., What is box breathing?"
+        lines=2,
+        placeholder="e.g., What is box breathing?",
+        label="Ask Daily Wellness AI"
     ),
     outputs=gr.Markdown(label="Answer from Daily Wellness AI"),
-    title="Daily Wellness AI Guru Chatbot",
-    description=(
-        "Ask wellness-related questions to get detailed, creative answers from "
-        "our knowledge base—expanded by an LLM if needed—or from the web. "
-        "We aim to bring calm and positivity to your day."
-    ),
-    theme="compact"
+    title="Daily Wellness AI",
+    description="Ask wellness-related questions and receive synthesized, creative answers.",
+    theme="default",
+    examples=[
+        "What is box breathing and how does it help reduce anxiety?",
+        "Provide a daily wellness schedule incorporating box breathing techniques.",
+        "What are some tips for maintaining good posture while working at a desk?"
+    ],
+    allow_flagging="never"
 )
 
-def main():
-    interface.launch(server_name="0.0.0.0", server_port=7860, debug=True)
-
-# If running in a local environment, we can also just call main()
+###############################################################################
+# 9) Launch Gradio
+###############################################################################
 if __name__ == "__main__":
-    main()
+    try:
+        # For Hugging Face Spaces, set share=False
+        interface.launch(server_name="0.0.0.0", server_port=7860, debug=False)
+    except Exception as e:
+        logger.error(f"Failed to launch Gradio interface: {e}")
+        logger.debug("Exception details:", exc_info=True)