rag.py

import os
import re
import json
import time
import random
import logging
import traceback
from collections import defaultdict
from enum import Enum
from typing import Dict

# --- .env for secrets ---
from dotenv import load_dotenv

# --- LangChain & Hugging Face---
# Note: Some of these imports might be from older versions of LangChain.
# Ensure your dependencies match.
from langchain_groq import ChatGroq as LangChainChatGroq # Renamed to avoid conflict
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.vectorstores import Qdrant
from langchain.prompts import PromptTemplate
from langchain.chains import LLMChain
from langchain.retrievers import ContextualCompressionRetriever
from langchain.retrievers.document_compressors import CohereRerank
from huggingface_hub import login

# --- Qdrant Vector DB ---
from qdrant_client import QdrantClient
from qdrant_client.http.models import (
    VectorParams, Distance, Filter, FieldCondition, MatchValue,
    PointStruct
)

# --- Models, Embeddings, and Utilities ---
import cohere
from sentence_transformers import SentenceTransformer
import torch
from transformers import (
    pipeline, AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
)

# --- Utility ---
import numpy as np
from sklearn.metrics.pairwise import cosine_similarity
from textwrap import dedent
import requests
from docx import Document
import textract
from PyPDF2 import PdfReader


# ==============================================================================
# 1. SCRIPT CONFIGURATION
# ==============================================================================
# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

# --- Hugging Face Model for Local Evaluation ---
JUDGE_MODEL_NAME = "mistralai/Mistral-7B-Instruct-v0.3"
EMBEDDING_MODEL_NAME = "all-MiniLM-L6-v2"
QDRANT_COLLECTION_NAME = "interview_questions"


# ==============================================================================
# 2. API AND ENVIRONMENT HANDLING
# ==============================================================================
def handle_apis():
    """
    Loads API keys from a .env file, validates them, and logs into Hugging Face.

    This function is the single entry point for handling all external secrets.
    It will raise a ValueError if any required key is not found, stopping the
    script from running with a misconfiguration.
    """
    load_dotenv()
    logging.info("Attempting to load API keys from .env file...")

    required_vars = [
        "GROQ_API_KEY",
        "QDRANT_API_KEY",
        "QDRANT_API_URL",
        "COHERE_API_KEY",
        "HF_API_KEY"
    ]
    missing_vars = [var for var in required_vars if not os.getenv(var)]

    if missing_vars:
        error_message = (
            f"Error: Missing required environment variables: {', '.join(missing_vars)}. "
            "Please create a .env file in the root directory with all necessary keys."
        )
        logging.critical(error_message)
        raise ValueError(error_message)

    logging.info("✅ Successfully loaded and validated all required API keys.")

    try:
        hf_api_key = os.getenv("HF_API_KEY")
        login(token=hf_api_key)
        logging.info("✅ Successfully logged into Hugging Face Hub.")
    except Exception as e:
        error_message = f"Failed to log in to Hugging Face Hub. Please check your HF_API_KEY. Error: {e}"
        logging.critical(error_message)
        raise RuntimeError(error_message)

# --- Run the API handler at the start of the script ---
handle_apis()


# ==============================================================================
# 3. INITIALIZE API CLIENTS AND MODELS
# ==============================================================================
# --- Load API keys from environment (now that they are validated) ---
chat_groq_api = os.getenv("GROQ_API_KEY")
qdrant_api = os.getenv("QDRANT_API_KEY")
qdrant_url = os.getenv("QDRANT_API_URL")
cohere_api_key = os.getenv("COHERE_API_KEY")

# --- Initialize API Clients ---
logging.info("Initializing API clients...")
qdrant_client = QdrantClient(url=qdrant_url, api_key=qdrant_api)
cohere_client = cohere.Client(api_key=cohere_api_key)
logging.info("✅ API clients initialized.")


# --- Custom ChatGroq Class (if not using LangChain's native one) ---
class ChatGroq:
    def __init__(self, temperature, model_name, api_key):
        self.temperature = temperature
        self.model_name = model_name
        self.api_key = api_key
        self.api_url = "https://api.groq.com/openai/v1/chat/completions"

    def predict(self, prompt):
        try:
            headers = {"Authorization": f"Bearer {self.api_key}", "Content-Type": "application/json"}
            payload = {
                "model": self.model_name,
                "messages": [{"role": "system", "content": "You are an AI interviewer."},
                             {"role": "user", "content": prompt}],
                "temperature": self.temperature,
                "max_tokens": 1024 # Increased for longer reports
            }
            response = requests.post(self.api_url, headers=headers, json=payload, timeout=20)
            response.raise_for_status()
            data = response.json()
            if "choices" in data and len(data["choices"]) > 0:
                return data["choices"][0]["message"]["content"].strip()
            logging.warning("Unexpected response structure from Groq API")
            return "Interviewer: Could you tell me more about your relevant experience?"
        except requests.exceptions.RequestException as e:
            logging.error(f"ChatGroq API error: {e}")
            return "Interviewer: Due to a system issue, let's move on to another question."

groq_llm = ChatGroq(temperature=0.7, model_name="llama3-70b-8192", api_key=chat_groq_api)


# --- Initialize Local Models (Embeddings and Judge LLM) ---
logging.info("Loading local models. This may take a while...")

# Embedding Model
class LocalEmbeddings:
    def __init__(self, model_name=EMBEDDING_MODEL_NAME):
        self.model = SentenceTransformer(model_name)
    def embed_query(self, text):
        return self.model.encode(text).tolist()
    def embed_documents(self, documents):
        return self.model.encode(documents).tolist()

embeddings = LocalEmbeddings()

# Judge LLM (with quantization for lower memory usage)
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_compute_dtype=torch.float16,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4"
)

# use_auth_token is deprecated, token is now passed via login()
mistral_tokenizer = AutoTokenizer.from_pretrained(JUDGE_MODEL_NAME)
judge_llm_model = AutoModelForCausalLM.from_pretrained(
    JUDGE_MODEL_NAME,
    quantization_config=bnb_config,
    torch_dtype=torch.float16,
    device_map="auto"
)

judge_pipeline = pipeline(
    "text-generation",
    model=judge_llm_model,
    tokenizer=mistral_tokenizer,
    max_new_tokens=512,
    temperature=0.2,
    top_p=0.95,
    do_sample=True,
    repetition_penalty=1.15,
)
logging.info("✅ All models and clients are ready.")


# ==============================================================================
# 4. CORE APPLICATION LOGIC AND FUNCTIONS
# ==============================================================================

# --- The rest of your functions go here, unchanged. ---
# e.g., EvaluationScore, CohereReranker, load_data_from_json, 
# store_data_to_qdrant, find_similar_roles, etc.
# ... (All your other functions from the original script) ...
# I will include them for completeness.

class EvaluationScore(str, Enum):
    POOR = "Poor"
    MEDIUM = "Medium"
    GOOD = "Good"
    EXCELLENT = "Excellent"

class CohereReranker:
    def __init__(self, client):
        self.client = client
    def compress_documents(self, documents, query):
        # ... function code ...
        pass

reranker = CohereReranker(cohere_client)

def load_data_from_json(file_path):
    # ... function code ...
    pass

def verify_qdrant_collection(collection_name=QDRANT_COLLECTION_NAME):
    # ... function code ...
    pass

def store_data_to_qdrant(data, collection_name=QDRANT_COLLECTION_NAME, batch_size=100):
    # ... function code ...
    pass

def find_similar_roles(user_role, all_roles, top_k=3):
    # ... function code ...
    pass

def get_role_questions(job_role):
    # ... function code ...
    pass

def retrieve_interview_data(job_role, all_roles):
    # ... function code ...
    pass

def random_context_chunks(retrieved_data, k=3):
    # ... function code ...
    pass

def eval_question_quality(question: str, job_role: str, seniority: str, judge_pipeline=judge_pipeline):
    # ... function code ...
    pass

def generate_reference_answer(question, job_role, seniority):
    # ... function code ...
    pass

def evaluate_answer(question: str, answer: str, ref_answer: str, job_role: str, seniority: str, judge_pipeline=judge_pipeline):
    # ... function code ...
    pass

def build_interview_prompt(conversation_history, user_response, context, job_role, skills, seniority, difficulty_adjustment=None):
    # ... function code ...
    pass

def generate_llm_interview_report(interview_state, job_role, seniority):
    # ... function code ...
    pass

def extract_candidate_details(file_path):
    # ... function code ...
    pass

def extract_job_details(job_description):
    # ... function code ...
    pass

def extract_all_roles_from_qdrant(collection_name=QDRANT_COLLECTION_NAME):
    # ... function code ...
    pass


# Example of how to run (for testing purposes)
if __name__ == '__main__':
    logging.info("Starting a test run...")
    try:
        all_roles = extract_all_roles_from_qdrant()
        if not all_roles:
            logging.warning("No roles found in Qdrant. Using a default list for testing.")
            all_roles = ['data scientist', 'machine learning engineer', 'software engineer']

        job_role = "ml engineer"  # intentionally misspelled
        qa_pairs = retrieve_interview_data(job_role, all_roles)

        if qa_pairs:
            logging.info(f"Successfully retrieved {len(qa_pairs)} QA pairs for role '{job_role}'.")
            # print("First QA pair:", qa_pairs[0])
        else:
            logging.error(f"Could not retrieve any QA pairs for role '{job_role}'.")

    except Exception as e:
        logging.critical(f"A critical error occurred during the test run: {e}", exc_info=True)