utils.py

# # utils.py
# import os
# import json
# from langchain_groq import ChatGroq
# from langchain_text_splitters import RecursiveCharacterTextSplitter
# # from langchain.schema import Document
# # from langchain.chains import RetrievalQA
# # from langchain_huggingface import HuggingFaceEmbeddings
# # from langchain_community.vectorstores import Chroma
# # from langchain.prompts import PromptTemplate

# from langchain_core.documents import Document
# from langchain_chroma import Chroma
# from langchain_huggingface import HuggingFaceEmbeddings
# from langchain_core.prompts import PromptTemplate
# # from langchain.chains import RetrievalQA
# from langchain_classic.chains import create_retrieval_chain
# from langchain_classic.chains.combine_documents import create_stuff_documents_chain
# from configure import USER_DATA_PATH, RAG_BASE_DIRECTORY, RAG_CATEGORIES
# import shutil
# from dotenv import load_dotenv
# from pymongo import MongoClient
# import certifi
# import re

# load_dotenv()


# def get_mongo_collection():
#     CONNECTION_STRING = os.getenv("CONNECTION_STRING")
#     DB_NAME = os.getenv("DB_NAME")
#     COLLECTION_NAME = os.getenv("COLLECTION_NAME")
#     try:
#         # Connect with certifi to avoid SSL errors
#         client = MongoClient(CONNECTION_STRING, tlsCAFile=certifi.where())
#         db = client[DB_NAME]
#         return db[COLLECTION_NAME]
#     except Exception as e:
#         print(f"Error connecting to Mongo: {e}")
#         return None
    

# def LLMChunking():
#     pass




# # LLM setup
# llm = ChatGroq(
#     api_key="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",
#     model="llama-3.3-70b-versatile",
#     temperature=0,
#     max_tokens=4000
# )

# # Text splitter
# text_splitter = RecursiveCharacterTextSplitter(
#     separators=["\n\n", "\n", ".", " ", ""],
#     chunk_size=500,
#     chunk_overlap=100,
#     length_function=len
# )
# # text_splitter = LLMChunking()

# # Embeddings
# embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
# # embeddings = None

# # Load user data
# def load_user_data():
#     try:
#         if os.path.exists(USER_DATA_PATH) and os.path.getsize(USER_DATA_PATH) > 0:
#             with open(USER_DATA_PATH, 'r') as f:
#                 return json.load(f)
#     except Exception:
#         pass
#     return {"users": {}, "user_info": {}}

# # Save user data
# def save_user_data(data):
#     with open(USER_DATA_PATH, 'w') as f:
#         json.dump(data, f, indent=2)

# # Initialize RAG with per-category vector stores and QA chains
# def initialize_rag():
#     """
#     Initialize RAG vector stores and chains using modern LangChain (LCEL).
#     Args:
#         llm: The initialized ChatGroq (or other) LLM object.
#         embeddings: The initialized HuggingFaceEmbeddings object.
#         text_splitter: The initialized RecursiveCharacterTextSplitter object.
#     """
#     vector_stores = {}
#     qa_chains = {}
    
#     # 1. Define Prompt Template (Modern LCEL Format)
#     # Note: Modern chains typically look for "context" and "input" variables.
#     base_prompt_template = """You are a {category} wellness expert. Provide helpful advice with specific actions:

# 1. Start with a brief empathetic response to the user's concern
# 2. Offer 1-3 actionable suggestions with brief explanations
# 3. End with an open-ended question to continue conversation

# Guidelines:
# - Keep responses conversational and supportive
# - Avoid clinical jargon
# - Focus on practical, implementable advice
# - Maintain hopeful and encouraging tone

# Context:
# {context}

# Question: {input}
# """
    
#     for category in RAG_CATEGORIES:
#         persist_dir = f"./chroma_db_{category}"
#         vector_store = None 

#         # --- 2. Check/Load Existing Vector Store ---
#         if os.path.exists(persist_dir):
#             print(f"Found existing vector store for {category}. Attempting to load...")
#             try:
#                 vector_store = Chroma(
#                     persist_directory=persist_dir,
#                     embedding_function=embeddings  # UPDATED: 'embedding_function', not 'embedding'
#                 )
#                 vector_stores[category] = vector_store
#             except Exception as e:
#                 print(f"Error loading existing store {persist_dir}: {e}")
#                 print("Will delete and attempt to re-build.")
#                 shutil.rmtree(persist_dir)
        
#         # --- 3. Create Vector Store if needed ---
#         if vector_store is None: 
#             print(f"No valid vector store for {category} found. Creating new one...")
            
#             dir_path = os.path.join(RAG_BASE_DIRECTORY, category)
#             docs = []
            
#             if os.path.exists(dir_path):
#                 for filename in os.listdir(dir_path):
#                     if filename.endswith('.txt'):
#                         file_path = os.path.join(dir_path, filename)
#                         try:
#                             with open(file_path, 'r', encoding='utf-8') as f:
#                                 text = f.read()
                            
#                             chunks = text_splitter.split_text(text)
#                             for chunk in chunks:
#                                 if chunk.strip():
#                                     metadata = {
#                                         "source": filename,
#                                         "category": category
#                                     }
#                                     docs.append(Document(
#                                         page_content=chunk.strip(),
#                                         metadata=metadata
#                                     ))
#                         except Exception as e:
#                             print(f"Error processing {file_path}: {e}")
            
#             if docs:
#                 # UPDATED: Use 'embedding_function' instead of 'embedding'
#                 # UPDATED: Removed .persist() call (Auto-persists in new version)
#                 vector_store = Chroma.from_documents(
#                     documents=docs,
#                     embedding=embeddings, 
#                     persist_directory=persist_dir
#                 )
#                 vector_stores[category] = vector_store
#                 print(f"Created new vector store for {category} with {len(docs)} documents.")
#             else:
#                 print(f"No documents found for {category}. Skipping QA chain setup.")
#                 continue 

#         # --- 4. Create QA Chain (LCEL Style) ---
#         if vector_store:
#             # A. Create the Prompt
#             # We inject the specific category into the template string immediately
#             category_specific_template = base_prompt_template.replace("{category}", category)
            
#             prompt = PromptTemplate(
#                 template=category_specific_template,
#                 input_variables=["context", "input"] # LCEL standard variables
#             )

#             # B. Create the Document Chain (LLM + Prompt)
#             question_answer_chain = create_stuff_documents_chain(llm, prompt)

#             # C. Create the Retrieval Chain (Retriever + Document Chain)
#             retriever = vector_store.as_retriever(search_kwargs={"k": 5})
#             rag_chain = create_retrieval_chain(retriever, question_answer_chain)

#             qa_chains[category] = rag_chain
#             print(f"Initialized {category} QA chain.")

#     return qa_chains



# # Classify question to category
# def classify_question_category(question):
#     prompt = f"""
#     Classify this question into one category: {', '.join(RAG_CATEGORIES)}
#     Question: {question}
#     Respond with only the category name.
#     """
#     response = llm.invoke(prompt)
#     print(response)
#     return response.content.strip()

# # Get RAG response using category-specific QA chain
# def get_rag_response(question, qa_chains):
#     # Classify question
#     category = classify_question_category(question)
    
#     if category not in qa_chains:
#         # Fallback to first available chain
#         category = list(qa_chains.keys())[0]
    
#     # Get response
#     result = qa_chains[category].invoke({"input": question})
#     return result['answer']


# # Classify user input
# def classify_input(user_input):
#     prompt = f"""
#     Classify the following user input into one of these categories:
#     1. "question" - If the user is asking a factual question that could be answered with knowledge
#     2. "general" - If the user is just chatting or expressing feelings

#     User Input: {user_input}

#     Respond with only one word: either "question" or "general"
#     """
#     response = llm.invoke(prompt)
#     return response.content.strip().lower()



# def parse_weird_json(text_data):
#     fixed_json_string = re.sub(r'\]\s*\[', ', ', text_data.strip())
    
#     # Step B: Load it as standard JSON
#     try:
#         data_list = json.loads(fixed_json_string)
#         return data_list
#     except json.JSONDecodeError as e:
#         print(f"❌ JSON Parsing Error: {e}")
#         return []







# utils.py
import os
import json
from langchain_groq import ChatGroq
from langchain_text_splitters import RecursiveCharacterTextSplitter
# from langchain.schema import Document
# from langchain.chains import RetrievalQA
# from langchain_huggingface import HuggingFaceEmbeddings
# from langchain_community.vectorstores import Chroma
# from langchain.prompts import PromptTemplate

from langchain_core.documents import Document
from langchain_chroma import Chroma
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_core.prompts import PromptTemplate
# from langchain.chains import RetrievalQA
from langchain_classic.chains import create_retrieval_chain
from langchain_classic.chains.combine_documents import create_stuff_documents_chain
from configure import USER_DATA_PATH, RAG_BASE_DIRECTORY, RAG_CATEGORIES
import shutil
from dotenv import load_dotenv
from pymongo import MongoClient
import certifi
import re

load_dotenv()
GROQ_API_KEY = os.getenv("GROQ_API_KEY")

def get_mongo_collection():
    CONNECTION_STRING = os.getenv("CONNECTION_STRING")
    DB_NAME = os.getenv("DB_NAME")
    COLLECTION_NAME = os.getenv("COLLECTION_NAME")
    try:
        # Connect with certifi to avoid SSL errors
        client = MongoClient(CONNECTION_STRING, tlsCAFile=certifi.where())
        db = client[DB_NAME]
        return db[COLLECTION_NAME]
    except Exception as e:
        print(f"Error connecting to Mongo: {e}")
        return None
    

def LLMChunking():
    pass




# LLM setup
llm = ChatGroq(
    api_key=GROQ_API_KEY,
    model="llama-3.3-70b-versatile",
    temperature=0.7,
    max_tokens=500,
    model_kwargs={
        "top_p": 0.9,
        "presence_penalty": 0.5,
        "frequency_penalty": 0.4
    }
)


# Text splitter
text_splitter = RecursiveCharacterTextSplitter(
    separators=["\n\n", "\n", ".", " ", ""],
    chunk_size=500,
    chunk_overlap=100,
    length_function=len
)
# text_splitter = LLMChunking()

# Embeddings
embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
# embeddings = None

# Load user data
def load_user_data():
    try:
        if os.path.exists(USER_DATA_PATH) and os.path.getsize(USER_DATA_PATH) > 0:
            with open(USER_DATA_PATH, 'r') as f:
                return json.load(f)
    except Exception:
        pass
    return {"users": {}, "user_info": {}}

# Save user data
def save_user_data(data):
    with open(USER_DATA_PATH, 'w') as f:
        json.dump(data, f, indent=2)

# Initialize RAG with per-category vector stores and QA chains
def initialize_rag():
    """
    Initialize RAG vector stores and chains using modern LangChain (LCEL).
    Args:
        llm: The initialized ChatGroq (or other) LLM object.
        embeddings: The initialized HuggingFaceEmbeddings object.
        text_splitter: The initialized RecursiveCharacterTextSplitter object.
    """
    vector_stores = {}
    qa_chains = {}
    retrievers = {}
    
    # 1. Define Prompt Template (Modern LCEL Format)
    # Note: Modern chains typically look for "context" and "input" variables.
    base_prompt_template = """You are a {category} wellness expert. Provide helpful advice with specific actions:

1. Start with a brief empathetic response to the user's concern
2. Offer 1-3 actionable suggestions with brief explanations
3. End with an open-ended question to continue conversation

Guidelines:
- Keep responses conversational and supportive
- Avoid clinical jargon
- Focus on practical, implementable advice
- Maintain hopeful and encouraging tone

Context:
{context}

Question: {input}
"""
    
    for category in RAG_CATEGORIES:
        persist_dir = f"./chroma_db_{category}"
        vector_store = None 

        # --- 2. Check/Load Existing Vector Store ---
        if os.path.exists(persist_dir):
            print(f"Found existing vector store for {category}. Attempting to load...")
            try:
                vector_store = Chroma(
                    persist_directory=persist_dir,
                    embedding_function=embeddings  # UPDATED: 'embedding_function', not 'embedding'
                )
                vector_stores[category] = vector_store
            except Exception as e:
                print(f"Error loading existing store {persist_dir}: {e}")
                print("Will delete and attempt to re-build.")
                shutil.rmtree(persist_dir)
        
        # --- 3. Create Vector Store if needed ---
        if vector_store is None: 
            print(f"No valid vector store for {category} found. Creating new one...")
            
            dir_path = os.path.join(RAG_BASE_DIRECTORY, category)
            docs = []
            
            if os.path.exists(dir_path):
                for filename in os.listdir(dir_path):
                    if filename.endswith('.txt'):
                        file_path = os.path.join(dir_path, filename)
                        try:
                            with open(file_path, 'r', encoding='utf-8') as f:
                                text = f.read()
                            
                            chunks = text_splitter.split_text(text)
                            for chunk in chunks:
                                if chunk.strip():
                                    metadata = {
                                        "source": filename,
                                        "category": category
                                    }
                                    docs.append(Document(
                                        page_content=chunk.strip(),
                                        metadata=metadata
                                    ))
                        except Exception as e:
                            print(f"Error processing {file_path}: {e}")
            
            if docs:
                # UPDATED: Use 'embedding_function' instead of 'embedding'
                # UPDATED: Removed .persist() call (Auto-persists in new version)
                vector_store = Chroma.from_documents(
                    documents=docs,
                    embedding=embeddings, 
                    persist_directory=persist_dir
                )
                vector_stores[category] = vector_store
                print(f"Created new vector store for {category} with {len(docs)} documents.")
            else:
                print(f"No documents found for {category}. Skipping QA chain setup.")
                continue 

        # --- 4. Create QA Chain (LCEL Style) ---
        if vector_store:
            # A. Create the Prompt
            # We inject the specific category into the template string immediately
            category_specific_template = base_prompt_template.replace("{category}", category)
            
            prompt = PromptTemplate(
                template=category_specific_template,
                input_variables=["context", "input"] # LCEL standard variables
            )

            # B. Create the Document Chain (LLM + Prompt)
            question_answer_chain = create_stuff_documents_chain(llm, prompt)

            # C. Create the Retrieval Chain (Retriever + Document Chain)
            retriever = vector_store.as_retriever(search_kwargs={"k": 5})
            rag_chain = create_retrieval_chain(retriever, question_answer_chain)

            qa_chains[category] = rag_chain
            retrievers[category] = retriever
            print(f"Initialized {category} QA chain.")

    return qa_chains, retrievers



# Classify question to category
def classify_question_category(question):
    prompt = f"""
    Classify this question into one category: {', '.join(RAG_CATEGORIES)}
    Question: {question}
    Respond with only the category name.
    """
    response = llm.invoke(prompt)
    print(response)
    return response.content.strip()

# Get RAG response using category-specific QA chain
# def get_rag_response(question, qa_chains):
#     # Classify question
#     category = classify_question_category(question)
    
#     if category not in qa_chains:
#         # Fallback to first available chain
#         category = list(qa_chains.keys())[0]
    
#     # Get response
#     # result = qa_chains[category].invoke({"input": question})
#     retriever = qa_chains[category].retriever
#     docs = retriever.invoke(question)

#     # return result['answer']
#     return [doc.page_content for doc in docs]



def get_rag_response(question, retrievers_dict):
    category = classify_question_category(question)
    
    if category not in retrievers_dict:
        category = list(retrievers_dict.keys())[0]
    
    docs = retrievers_dict[category].invoke(question)
    
    return "\n\n".join([doc.page_content for doc in docs])

# Classify user input
def classify_input(user_input):
    prompt = f"""
    Classify the following user input into one of these categories:
    1. "question" - If the user is asking a factual question that could be answered with knowledge
    2. "general" - If the user is just chatting or expressing feelings

    User Input: {user_input}

    Respond with only one word: either "question" or "general"
    """
    response = llm.invoke(prompt)
    return response.content.strip().lower()



def parse_weird_json(text_data):
    fixed_json_string = re.sub(r'\]\s*\[', ', ', text_data.strip())
    
    # Step B: Load it as standard JSON
    try:
        data_list = json.loads(fixed_json_string)
        return data_list
    except json.JSONDecodeError as e:
        print(f"❌ JSON Parsing Error: {e}")
        return []