data.py

import threading
import logging
import os
import pandas as pd
from typing import Dict
import psutil
import google.generativeai as genai
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_openai import ChatOpenAI
from neo4j import GraphDatabase
from datetime import datetime
import sys
import time

# accessing and defining the llm served in the cluster
from langchain_community.chat_models import ChatOpenAI

NGROK_URL = "https://071bb12c21bd.ngrok-free.app" 
VLLM_MODEL_NAME = "hugging-quants/Meta-Llama-3.1-8B-Instruct-AWQ-INT4"

llm1 = ChatOpenAI(
    openai_api_key="EMPTY",
    openai_api_base=f"{NGROK_URL}/v1",
    model_name=VLLM_MODEL_NAME,
    temperature=0.0,
)

# graph db credentials and gemini api key management 

NEO4J_URI = "neo4j+s://83c91252.databases.neo4j.io"
NEO4J_USER = "neo4j"
NEO4J_PASS = "TKUo104Fqm6BtDNJ2eeb6RRM8xIBS80kdJKLOFv_CkI"

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[logging.StreamHandler()]
)
logger = logging.getLogger(__name__)

API_KEYS = [
    "AIzaSyC4qW9fQAbiCS30xY-Tq086AoRB8lUfxes",
    "AIzaSyAxa91t80DFbtZG0Z8pyHlGG2X0v-sYSRc",
    "AIzaSyDGc7l82E1sa4GarTfGiTXC6dLD7Crk8oo",
    "AIzaSyC463t1T7I6AkKFt9z0eP85KpBVJsGljHc",
    "AIzaSyAHoi9xbYAThtjXlyF_IKFtruoWYoUCjJQ"
]


model_key_counters: Dict[str, int] = {
    "gemini-2.5-flash": 0,
    "gemini-2.0-flash-lite": 0,
    "gemini-2.0-flash": 0,
    "gemini-2.5-flash-lite-preview-06-17": 0,
    "models/embedding-001": 0,
    "models/text-embedding-004": 0
}
counter_lock = threading.Lock()

# Usage and Error Tracking
key_usage: Dict[tuple, int] = {(model, key): 0 for model in model_key_counters for key in API_KEYS}
key_errors: Dict[tuple, int] = {(model, key): 0 for model in model_key_counters for key in API_KEYS}
MAX_ERRORS_PER_KEY = 50

# Lock for genai API calls
genai_lock = threading.Lock()

# Chat Model Cache
model_cache: Dict[str, ChatGoogleGenerativeAI | genai.GenerativeModel] = {}

def select_key_for_model(model_name: str) -> str:
    """Select the next API key for the model using round-robin, skipping unreliable keys."""
    with counter_lock:
        available_keys = [k for k in API_KEYS if key_errors[(model_name, k)] < MAX_ERRORS_PER_KEY]
        if not available_keys:
            logger.warning(f"All keys for {model_name} exhausted. Resetting errors.")
            for k in API_KEYS:
                key_errors[(model_name, k)] = 0
            available_keys = API_KEYS
        index = model_key_counters[model_name] % len(available_keys)
        selected_key = available_keys[index]
        model_key_counters[model_name] = (index + 1) % len(available_keys)
        key_usage[(model_name, selected_key)] += 1
        logger.info(f"Key {selected_key[:10]}... selected for {model_name}, usage: {key_usage[(model_name, selected_key)]}")
        return selected_key

def get_model(model_name: str, use_chat: bool) -> ChatGoogleGenerativeAI | genai.GenerativeModel:
    """Get or create a ChatGoogleGenerativeAI model with a selected key."""
    key = select_key_for_model(model_name)
    cache_key = f"{model_name}_{key}_{'chat' if use_chat else 'genai'}"
    if cache_key not in model_cache:
        try:
            genai.configure(api_key=key)
            if use_chat:
                model = ChatGoogleGenerativeAI(
                    model=model_name,
                    temperature=0.7,
                    google_api_key=key
                )
            else:
                model = genai.GenerativeModel(model_name)

            model_cache[cache_key] = model
            logger.info(f"Created Model for {model_name} with key {key[:10]}...")
        except Exception as e:
            key_errors[(model_name, key)] += 1
            logger.error(f"Error creating model for {model_name}: {str(e)}")
            if key_errors[(model_name, key)] >= MAX_ERRORS_PER_KEY:
                logger.warning(f"Key {key[:10]}... unreliable for {model_name}")
            return get_model(model_name)  # Retry
    return model_cache[cache_key]

def call_genai_embedding_api(model_name: str, *args, **kwargs):
    """Call genai API with selected key, using a lock."""
    key = select_key_for_model(model_name)
    with genai_lock:
        try:
            genai.configure(api_key=key)
            return genai.embed_content(*args, **kwargs)
        except Exception as e:
            key_errors[(model_name, key)] += 1
            logger.error(f"API call failed for {model_name}: {str(e)}")
            if key_errors[(model_name, key)] >= MAX_ERRORS_PER_KEY:
                logger.warning(f"Key {key[:10]}... unreliable for {model_name}")
            raise

def reset_key_metrics():
    """Reset usage and error counts for all (model, key) pairs."""
    with counter_lock:
        for (model, key) in key_usage:
            key_usage[(model, key)] = 0
            key_errors[(model, key)] = 0
        logger.info("Reset all API key usage and error counts for all models.")

# Schedule metrics reset every 8 hours
def schedule_reset():
    while True:
        try:
            time.sleep(8 * 60 * 60)  # Reset every 8 hours
            reset_key_metrics()
        except Exception as e:
            logging.error(f"Error in reset_key_metrics: {e}")

# CPU monitoring every 20 seconds
def monitor_cpu():
    while True:
        try:
            cpu_usage = psutil.cpu_percent(interval=20)  # Combine interval with sleep
            logging.info(f"CPU Utilization: {cpu_usage}%")
        except Exception as e:
            logging.error(f"Error in CPU monitoring: {e}")

# Start threads
reset_thread = threading.Thread(target=schedule_reset, daemon=True)
monitor_thread = threading.Thread(target=monitor_cpu, daemon=True)
reset_thread.start()
monitor_thread.start()

genai.configure(api_key=os.getenv("GEMINI_API_KEY_1"))

_driver = None
def get_driver():
    global _driver
    if _driver is None:
        _driver = GraphDatabase.driver(NEO4J_URI, auth=(NEO4J_USER,NEO4J_PASS))
    return _driver



def debug_print(category, message):
    """
    Print a debug message with timestamp and category, flushed immediately.
    
    Args:
        category (str): Category of the debug message (e.g., NODE, AGENT, TOOL)
        message (str): The message to print
    """
    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    print(f"[{timestamp}] [{category}] {message}", flush=True, file=sys.stdout)


# Data Loading 
eligibility_df = pd.read_csv("cards_eligibility_updated.csv")
df = pd.read_csv("credit_card_data_updated.csv")
# print(df.head())
card_descriptions = dict(zip(df["name"], df["description"]))
features = ['Fuel Surcharge Waiver', 'Insurance', 'Shopping Benefits', 'Airport Lounge Access', 'Co-Branded', 'Daily Spends (Grocery)', 'Dining Benefits', 'Domestic Travel Benefits', 'Entertainment', 'General Reward Points', 'Movie Benefits', 'Rupay Network Support', 'Student', 'UPI Transaction Support', 'Welcome Bonus', 'International Travel Benefits', 'premium', 'Flight Discounts', 'Hotel Benefits', 'Travel Benefits', 'Railway Benefits', 'Railway Lounge', 'Utility', 'Beginners (Entry Level)', 'E-commerce Platform Benefits', 'Air Miles', 'Jewellery Spends', 'Concierge Services', 'Food Delivery Benefits', 'Lifestyle & Luxury Perks', 'Spa Access Benefits', 'Golf Access & Perks', 'Super Premium', 'Frequent Flyer Benefits', 'Health Benefits', 'Rent Payment Benefits', 'Education', 'Lifetime Free', 'Roadside Assistance', 'EMI Conversion Options', 'No Forex Markup Fee', 'Secured FD Based', 'Cashback', 'Fuel Benefits', 'Business']
df_all_cards = pd.read_csv("credit_card_data_updated.csv") 
all_card_names = df_all_cards["name"].tolist()
all_card_lookup = dict(zip(df_all_cards["name"], df_all_cards["description"]))


# Eligibility Lookup 
eligibility_lookup = {}
for _, row in eligibility_df.iterrows():
    name = row["Name"].strip()
    eligibility_lookup[name] = {
        "bank": row["Bank"],
        "age_range": f"{row['Minimum Age']} to {row['Maximum Age']}",
        "income": f"{row['Minimum Income (LPA)']} LPA",
        "credit_score": row["Minimum Credit Score"],
        "joining_fee": row["Joining fee"],
        "annual_fee": row["Annual fee"],
        "issuer_link": row.get("Issuer Link", "").strip()
    }