app.py2

# This block contains the full combined script.
# It includes all the code from the previous successful steps, ordered correctly.

# Combined Imports
import os
import gradio as gr
from huggingface_hub import InferenceClient
import torch
import re
import warnings
import time
import json
from transformers import AutoTokenizer, AutoModelForCausalLM, GenerationConfig, BitsAndBytesConfig
from sentence_transformers import SentenceTransformer, util, CrossEncoder
import gspread
# from google.colab import auth
from google.auth import default
from tqdm import tqdm
from duckduckgo_search import DDGS
import spacy
from datetime import date, timedelta, datetime # Import datetime
from dateutil.relativedelta import relativedelta
import traceback
import base64
import pytz # Import pytz for timezone handling


# Suppress warnings
warnings.filterwarnings("ignore", category=UserWarning)

# Define global variables and load secrets
HF_TOKEN = os.getenv("HF_TOKEN")
SHEET_ID = "19ipxC2vHYhpXCefpxpIkpeYdI3a1Ku2kYwecgUULIw"
GOOGLE_BASE64_CREDENTIALS = os.getenv("GOOGLE_BASE64_CREDENTIALS")

# Initialize InferenceClient
client = InferenceClient("google/gemma-2-9b-it", token=HF_TOKEN)


# Load spacy model for sentence splitting
nlp = None
try:
    nlp = spacy.load("en_core_web_sm")
    print("SpaCy model 'en_core_web_sm' loaded.")
except OSError:
    print("SpaCy model 'en_core_web_sm' not found. Downloading...")
    try:
        os.system("python -m spacy download en_core_web_sm")
        nlp = spacy.load("en_core_web_sm")
        print("SpaCy model 'en_core_web_sm' downloaded and loaded.")
    except Exception as e:
        print(f"Failed to download or load SpaCy model: {e}")


# Load SentenceTransformer for RAG/business info retrieval
embedder = None
try:
    print("Attempting to load Sentence Transformer (sentence-transformers/paraphrase-MiniLM-L6-v2)...")
    embedder = SentenceTransformer("sentence-transformers/paraphrase-MiniLM-L6-v2")
    print("Sentence Transformer loaded.")
except Exception as e:
     print(f"Error loading Sentence Transformer: {e}")


# Load a Cross-Encoder model for re-ranking retrieved documents
reranker = None
try:
    print("Attempting to load Cross-Encoder Reranker (cross-encoder/ms-marco-MiniLM-L6-v2)...")
    reranker = CrossEncoder('cross-encoder/ms-marco-MiniLM-L6-v2')
    print("Cross-Encoder Reranker loaded.")
except Exception as e:
    print(f"Error loading Cross-Encoder Reranker: {e}")
    print("Please ensure the model identifier 'cross-encoder/ms-marco-MiniLM-L6-v2' is correct and accessible on Hugging Face Hub.")
    print(traceback.format_exc())
    reranker = None


# Google Sheets Authentication
gc = None # Global variable for gspread client
def authenticate_google_sheets():
    """Authenticates with Google Sheets using base64 encoded credentials."""
    global gc
    print("Authenticating Google Account...")
    if not GOOGLE_BASE64_CREDENTIALS:
        print("Error: GOOGLE_BASE64_CREDENTIALS secret not found.")
        return False

    try:
        # Decode the base64 credentials
        credentials_json = base64.b64decode(GOOGLE_BASE64_CREDENTIALS).decode('utf-8')
        credentials = json.loads(credentials_json)

        # Authenticate using service account from dictionary
        gc = gspread.service_account_from_dict(credentials)
        print("Google Sheets authentication successful via service account.")
        return True
    except Exception as e:
        print(f"Google Sheets authentication failed: {e}")
        print("Please ensure your GOOGLE_BASE64_CREDENTIALS secret is correctly set and contains valid service account credentials.")
        print(traceback.format_exc())
        return False

# Google Sheets Data Loading and Embedding
data = [] # Global variable to store loaded data
descriptions_for_embedding = []
embeddings = torch.tensor([])
business_info_available = False # Flag to indicate if business info was loaded successfully

def load_business_info():
    """Loads business information from Google Sheet and creates embeddings."""
    global data, descriptions_for_embedding, embeddings, business_info_available
    business_info_available = False # Reset flag

    if gc is None:
        print("Skipping Google Sheet loading: Google Sheets client not authenticated.")
        return

    if not SHEET_ID:
         print("Error: SHEET_ID not set.")
         return

    try:
        sheet = gc.open_by_key(SHEET_ID).sheet1
        print(f"Successfully opened Google Sheet with ID: {SHEET_ID}")
        data_records = sheet.get_all_records()

        if not data_records:
            print(f"Warning: No data records found in Google Sheet with ID: {SHEET_ID}")
            data = []
            descriptions_for_embedding = []
        else:
            # Filter out rows missing 'Service' or 'Description'
            filtered_data = [row for row in data_records if row.get('Service') and row.get('Description')]
            if not filtered_data:
                print("Warning: Filtered data is empty after checking for 'Service' and 'Description'.")
                data = []
                descriptions_for_embedding = []
            else:
                data = filtered_data
                # Use BOTH Service and Description for embedding
                descriptions_for_embedding = [f"Service: {row['Service']}. Description: {row['Description']}" for row in data]

                # Only encode if descriptions_for_embedding are found and embedder is available
                if descriptions_for_embedding and embedder is not None:
                    print("Encoding descriptions...")
                    try:
                        embeddings = embedder.encode(descriptions_for_embedding, convert_to_tensor=True)
                        print("Encoding complete.")
                        business_info_available = True # Set flag if successful
                    except Exception as e:
                        print(f"Error during description encoding: {e}")
                        embeddings = torch.tensor([]) # Ensure embeddings is an empty tensor on error
                        business_info_available = False # Encoding failed
                else:
                    print("Skipping encoding descriptions: No descriptions found or embedder not available.")
                    embeddings = torch.tensor([]) # Ensure embeddings is an empty tensor
                    business_info_available = False # Cannot use RAG without descriptions or embedder

        print(f"Loaded {len(descriptions_for_embedding)} entries from Google Sheet for embedding/RAG.")
        if not business_info_available:
            print("Business information retrieval (RAG) is NOT available.")

    except gspread.exceptions.SpreadsheetNotFound:
        print(f"Error: Google Sheet with ID '{SHEET_ID}' not found.")
        print("Please check the SHEET_ID and ensure your authenticated Google Account has access to this sheet.")
        business_info_available = False # Sheet not found
    except Exception as e:
        print(f"An error occurred while accessing the Google Sheet: {e}")
        print(traceback.format_exc())
        business_info_available = False # Other sheet access error

# Business Info Retrieval (RAG)
def retrieve_business_info(query: str, top_n: int = 3) -> list:
    """
    Retrieves relevant business information from loaded data based on a query.

    Args:
        query: The user's query string.
        top_n: The number of top relevant entries to retrieve.

    Returns:
        A list of dictionaries, where each dictionary is a relevant row from the
        Google Sheet data. Returns an empty list if RAG is not available or
        no relevant information is found.
    """
    global data
    if not business_info_available or embedder is None or not descriptions_for_embedding or not data:
        print("Business information retrieval is not available or data is empty.")
        return []

    try:
        query_embedding = embedder.encode(query, convert_to_tensor=True)
        cosine_scores = util.cos_sim(query_embedding, embeddings)[0]
        top_results_indices = torch.topk(cosine_scores, k=min(top_n, len(data)))[1].tolist()
        top_results = [data[i] for i in top_results_indices]

        if reranker is not None and top_results:
            print("Re-ranking top results...")
            rerank_pairs = [(query, descriptions_for_embedding[i]) for i in top_results_indices]
            rerank_scores = reranker.predict(rerank_pairs)
            reranked_indices = sorted(range(len(rerank_scores)), key=lambda i: rerank_scores[i], reverse=True)
            reranked_results = [top_results[i] for i in reranked_indices]
            print("Re-ranking complete.")
            return reranked_results
        else:
            return top_results

    except Exception as e:
        print(f"Error during business information retrieval: {e}")
        print(traceback.format_exc())
        return []

# Function to perform DuckDuckGo Search and return results with URLs
def perform_duckduckgo_search(query: str, max_results: int = 5):
    """
    Performs a search using DuckDuckGo and returns a list of dictionaries.
    Includes a delay to avoid rate limits.
    Returns an empty list and prints an error if search fails.
    """
    print(f"Executing Tool: perform_duckduckgo_search with query='{query}')")
    search_results_list = []
    try:
        time.sleep(1)

        with DDGS() as ddgs:
            if not query or len(query.split()) < 2:
                 print(f"Skipping search for short query: '{query}'")
                 return []

            results_generator = ddgs.text(query, max_results=max_results)
            results_found = False
            for r in results_generator:
                search_results_list.append(r)
                results_found = True

            if not results_found and max_results > 0:
                 print(f"DuckDuckGo search for '{query}' returned no results.")


    except Exception as e:
        print(f"Error during Duckduckgo search for '{query}': {e}")
        return []

    return search_results_list

# Function to perform date calculation if needed
def perform_date_calculation(query: str):
    """
    Analyzes query for date calculation requests and performs the calculation.
    Returns a dict describing the calculation and result, or None.
    Handles formats like 'X days ago', 'X days from now', 'X weeks ago', 'X weeks from now', 'what is today's date'.
    Uses dateutil for slightly more flexibility (though core logic remains simple).
    """
    print(f"Executing Tool: perform_date_calculation with query='{query}')")
    query_lower = query.lower()

    # Use datetime.now(timezone).date() with an explicit timezone
    try:
        # Using UTC as a default timezone, you can change this based on your needs
        tz = pytz.timezone("UTC")
        today = datetime.now(tz).date()
        print(f"[DEBUG] Current system date used: {today}")
    except Exception as e:
        print(f"Error getting timezone-aware date: {e}")
        # Fallback to date.today() if timezone handling fails
        today = date.today()
        print(f"[DEBUG] Falling back to date.today(): {today}")


    result_date = None
    calculation_description = None

    if re.search(r"\btoday'?s date\b|what is today'?s date\b|what day is it\b", query_lower):
        result_date = today
        calculation_description = f"The current date is: {today.strftime('%Y-%m-%d')}"
        print(f"Identified query for today's date.")
        return {"query": query, "description": calculation_description, "result": result_date.strftime('%Y-%m-%d'), "success": True}

    match = re.search(r"(\d+)\s+(day|week|month|year)s?\s+(ago|from now)", query_lower)

    if match:
        value = int(match.group(1))
        unit = match.group(2)
        direction = match.group(3)

        try:
            if unit == 'day':
                delta = timedelta(days=value)
            elif unit == 'week':
                delta = timedelta(weeks=value)
            elif unit == 'month':
                 delta = relativedelta(months=value)
            elif unit == 'year':
                 delta = relativedelta(years=value)
            else:
                 desc = f"Could not understand the time unit '{unit}' in '{query}'."
                 print(desc)
                 return {"query": query, "description": desc, "result": None, "success": False, "error": desc}

            if direction == 'ago':
                result_date = today - delta
                calculation_description = f"Calculating date {value} {unit}s ago from {today.strftime('%Y-%m-%d')}: {result_date.strftime('%Y-%m-%d')}"
            elif direction == 'from now':
                result_date = today + delta
                calculation_description = f"Calculating date {value} {unit}s from now from {today.strftime('%Y-%m-%d')}: {result_date.strftime('%Y-%m-%d')}"

            print(f"Performed date calculation: {calculation_description}")
            return {"query": query, "description": calculation_description, "result": result_date.strftime('%Y-%m-%d'), "success": True}

        except OverflowError:
             desc = f"Date calculation overflow for query: {query}"
             print(f"Date calculation overflow for query: {query}")
             return {"query": query, "description": desc, "result": None, "success": False, "error": desc}
        except Exception as e:
             desc = f"An error occurred during date calculation for query '{query}': {e}"
             print(desc)
             return {"query": query, "description": desc, "result": None, "success": False, "error": str(e)}

    desc = "No specific date calculation pattern recognized."
    print(f"No specific date calculation pattern found in query: '{query}'")
    return {"query": query, "description": desc, "result": None, "success": False}

# Function to identify questions using spaCy
def identify_questions(message: str) -> list[str]:
    """
    Identifies potential questions in the user's message using spaCy.

    Args:
        message: The user's input string.

    Returns:
        A list of strings, where each string is an identified question.
    """
    if nlp is None:
        print("SpaCy model not loaded, cannot identify questions.")
        return [message]

    doc = nlp(message)
    questions = []
    for sent in doc.sents:
        if sent.text.strip().endswith('?'):
            questions.append(sent.text.strip())

    if not questions and message.strip():
        print("No specific questions identified, treating entire message as query.")
        questions.append(message.strip())

    return questions

# Function to classify questions
def classify_question(question: str) -> dict:
    """
    Classifies a question as 'business', 'general', or 'date' based on RAG results
    or date calculation pattern.

    Args:
        question: The question string.

    Returns:
        A dictionary containing the question and its classification,
        including calculation result for 'date' questions.
    """
    date_calculation_result = perform_date_calculation(question)
    if date_calculation_result and date_calculation_result['success']:
        print(f"Question '{question}' classified as 'date'. Date calculation pattern found.")
        return {'question': question, 'type': 'date', 'retrieved_info': [], 'calculation_result': date_calculation_result['result']}

    retrieved_info = retrieve_business_info(question, top_n=1)

    if retrieved_info:
        print(f"Question '{question}' classified as 'business'. Retrieved info found.")
        return {'question': question, 'type': 'business', 'retrieved_info': retrieved_info}
    else:
        print(f"Question '{question}' classified as 'general'. No relevant info found.")
        return {'question': question, 'type': 'general', 'retrieved_info': []}


def create_model_prompt(system_message: str, processed_questions_with_context: list, history: list) -> list:
    """
    Creates a detailed prompt for the model based on the system message,
    classified questions with context, and conversation history.

    Args:
        system_message: The base system message (now includes persona).
        processed_questions_with_context: List of dictionaries with questions,
                                           types, retrieved info, and calculation results.
        history: Conversation history.

    Returns:
        A list of messages in ChatML format to be sent to the model.
    """
    messages = [{"role": "system", "content": system_message}]

    for user_msg, bot_msg in history:
        if user_msg:
            messages.append({"role": "user", "content": user_msg})
        if bot_msg:
            messages.append({"role": "assistant", "content": bot_msg})

    messages.append({"role": "user", "content": "The user has asked one or more questions. Process each question individually based on its type (business, general, or date). For business questions, use the provided context. For general questions, use your internal knowledge. FOR DATE QUESTIONS, USE THE PROVIDED DATE *DIRECTLY* in your answer. Combine your answers into a single, coherent response, clearly addressing each question. If you cannot answer a question, state that you are unable and explain why if possible. Ensure your final response directly addresses the user's queries without unnecessary preamble or repetition of the original full query. You have access to a date calculation tool whose results are provided for 'date' type questions."})


    for item in processed_questions_with_context:
        question = item['question']
        question_type = item['type']
        retrieved_info = item.get('retrieved_info', [])
        calculation_result = item.get('calculation_result', None)

        prompt_content = f"**Question:** {question}\n**Question Type:** {question_type}\n"

        if question_type == 'business' and retrieved_info:
            context_message = "Relevant Business Information (use this to answer if applicable):\n"
            for i, info in enumerate(retrieved_info):
                context_message += f"--- Business Info Entry {i+1} ---\n"
                for key, value in info.items():
                     context_message += f"{key}: {str(value)}\n"
                context_message += "---\n"
            prompt_content += context_message
            print(f"Added business context for question: '{question}'")
        elif question_type == 'business' and not retrieved_info and business_info_available:
             prompt_content += "No highly relevant business information found for this specific question in the loaded data. Please answer based on general knowledge if possible, or state that business information is not available for this query."
             print(f"No relevant business context found for business question: '{question}'")
        elif question_type == 'business' and not business_info_available:
             prompt_content += "Business information is not available. Please answer based on general knowledge if possible."
             print(f"Business information not available for business question: '{question}'")
        elif question_type == 'date' and calculation_result is not None:
            prompt_content += f"Calculated Date: {calculation_result}\nUse this date to answer the question."
            print(f"Added calculated date result for question: '{question}')")
        elif question_type == 'general':
             prompt_content += "This is a general knowledge question. Answer it using your internal knowledge."
             print(f"Handling general question: '{question}'")
        else:
             prompt_content += "Unable to process this question with available tools or information."
             print(f"Unable to process question: '{question}' (Type: {question_type})")


        messages.append({"role": "user", "content": prompt_content})

    messages.append({"role": "user", "content": "Now, synthesize the answers to all the above questions into a single, coherent, and polite response. Address each question clearly. Do not repeat the original full query. If you could not answer a specific question, mention that."})


    return messages


# Chat handler function
def respond(
    message: str,
    history: list[tuple[str, str]],
    system_message: str,
    max_tokens: int,
    temperature: float,
    top_p: float,
):
    identified_questions = identify_questions(message)
    print(f"Identified questions for processing: {identified_questions}")

    processed_questions_with_context = []
    for question in identified_questions:
        classification_result = classify_question(question)
        question_type = classification_result['type']
        retrieved_info = classification_result.get('retrieved_info', [])
        calculation_result = classification_result.get('calculation_result', None)

        if question_type == 'business':
            enhanced_retrieved_info = retrieve_business_info(question, top_n=5)
            print(f"Enhanced RAG results for business question '{question}': {len(enhanced_retrieved_info)} items retrieved.")
            processed_questions_with_context.append({
                'question': question,
                'type': question_type,
                'retrieved_info': enhanced_retrieved_info,
                'calculation_result': None
            })
        elif question_type == 'date':
             print(f"Date calculation result for '{question}': {calculation_result}")
             processed_questions_with_context.append({
                 'question': question,
                 'type': question_type,
                 'retrieved_info': [],
                 'calculation_result': calculation_result
             })
        else:
            processed_questions_with_context.append({
                'question': question,
                'type': question_type,
                'retrieved_info': [],
                'calculation_result': None
            })


    print(f"Processed questions with enhanced context/calculation results: {processed_questions_with_context}")

    messages = create_model_prompt(system_message, processed_questions_with_context, history)

    response = ""
    try:
        result = client.chat_completion(
            messages=messages,
            max_tokens=max_tokens,
            stream=False,
            temperature=temperature,
            top_p=top_p,
        )
        response = result.choices[0].message.content or ""
        print(f"Generated combined response: {response[:200]}...")
        yield response
    except Exception as e:
        print(f"Error generating combined response: {e}")
        print(traceback.format_exc())
        yield f"An error occurred while generating the combined response: {e}"


# Gradio interface
print(f"RAG functionality available: {business_info_available}")

demo = gr.ChatInterface(
    fn=respond,
    additional_inputs=[
        gr.Textbox(value="You are FutureAi, a helpful AI developed by the Futuresony team in October 2024. Use the provided business information as if it is from your company.", label="System message"),
        gr.Slider(1, 2048, value=512, step=1, label="Max new tokens"),
        gr.Slider(0.1, 4.0, value=0.7, step=0.1, label="Temperature"),
        gr.Slider(0.1, 1.0, value=0.95, step=0.05, label="Top‑p (nucleus sampling)"),
    ],
    title="Gemma‑2‑9B‑IT Chat with RAG",
    description="Chat with Google Gemma‑2‑9B‑IT via Hugging Face Inference API, with business info retrieved from Google Sheets.",
)

demo.queue()

if __name__ == "__main__":
    if authenticate_google_sheets():
        load_business_info()
    else:
        print("Google Sheets authentication failed. RAG functionality will not be available.")

    print(f"RAG functionality available: {business_info_available}")

    demo.launch()