app.py

import os
import subprocess
import uuid
import datetime
from flask import Flask, request, jsonify, send_from_directory
from langchain_community.document_loaders import TextLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_community.vectorstores import Chroma
from langchain_community.llms import CTransformers
from langchain.chains import ConversationalRetrievalChain
from langchain.memory import ConversationBufferMemory

# ---------- CONFIGURATION ----------
COMBINED_DATASET_FILE = "combined_dataset.txt"
CHUNK_SIZE = 300
CHUNK_OVERLAP = 30
API_TOKEN = os.environ.get('API_TOKEN', 'supersecret')  # Token for auth

DATA_DIR = os.environ.get('DATA_DIR', '.')
VECTORSTORE_DIR = os.path.join(DATA_DIR, 'vectorstore')
MODELS_DIR = os.path.join(DATA_DIR, 'models')
HF_CACHE_DIR = os.path.join(DATA_DIR, '.cache')
CHAT_LOG_FILE = os.path.join(DATA_DIR, 'chat_logs.txt')

COMBINED_DATASET_FILE_PATH = os.path.join(DATA_DIR, COMBINED_DATASET_FILE)
# ------------------------------------

app = Flask(__name__, static_folder='frontend/build', static_url_path='')

vectorstore_instance = None
chat_sessions = {}
loaded_models = {}

def log_chat(session_id, user_msg, bot_msg):
    with open(CHAT_LOG_FILE, 'a') as f:
        f.write(f"{datetime.datetime.now()} | Session {session_id}\n")
        f.write(f"You: {user_msg}\n")
        f.write(f"Bot: {bot_msg}\n\n")

def prepare_vectorstore():
    global vectorstore_instance
    os.makedirs(VECTORSTORE_DIR, exist_ok=True)

    if not os.listdir(VECTORSTORE_DIR):
        print("📄 Vector store not found or empty. Creating...")

        if not os.path.exists(COMBINED_DATASET_FILE_PATH):
            raise FileNotFoundError(f"Dataset file missing at {COMBINED_DATASET_FILE_PATH}")

        loader = TextLoader(COMBINED_DATASET_FILE_PATH)
        documents = loader.load()
        splitter = RecursiveCharacterTextSplitter(chunk_size=CHUNK_SIZE, chunk_overlap=CHUNK_OVERLAP)
        chunks = splitter.split_documents(documents)

        embeddings = HuggingFaceEmbeddings(
            model_name="sentence-transformers/all-MiniLM-L6-v2",
            cache_folder=HF_CACHE_DIR
        )
        vectorstore_instance = Chroma.from_documents(chunks, embedding=embeddings, persist_directory=VECTORSTORE_DIR)
        print("✅ Vector store created.")
    else:
        embeddings = HuggingFaceEmbeddings(
            model_name="sentence-transformers/all-MiniLM-L6-v2",
            cache_folder=HF_CACHE_DIR
        )
        vectorstore_instance = Chroma(persist_directory=VECTORSTORE_DIR, embedding_function=embeddings)
        print("✅ Vector store loaded.")

def load_model(model_name):
    global loaded_models

    if model_name in loaded_models:
        return loaded_models[model_name]

    models_config = {
        "TinyLLaMA (1.1B)": {
            "file": "tinyllama-1.1b-chat-v1.0.Q4_K_M.gguf",
            "url": "https://huggingface.co/TheBloke/TinyLlama-1.1B-Chat-v1.0-GGUF/resolve/main/tinyllama-1.1b-chat-v1.0.Q4_K_M.gguf"
        },
        "LLaMA 2 (7B)": {
            "file": "llama-2-7b-chat.Q4_K_M.gguf",
            "url": "https://huggingface.co/TheBloke/Llama-2-7B-Chat-GGUF/resolve/main/llama-2-7b-chat.Q4_K_M.gguf"
        },
        "CodeLLaMA (7B)": {
            "file": "codellama-7b.Q4_K_M.gguf",
            "url": "https://huggingface.co/TheBloke/CodeLlama-7B-GGUF/resolve/main/codellama-7b.Q4_K_M.gguf"
        }
    }

    os.makedirs(MODELS_DIR, exist_ok=True)
    model_info = models_config.get(model_name)
    if not model_info:
        raise ValueError(f"Model '{model_name}' not configured.")

    model_path = os.path.join(MODELS_DIR, model_info["file"])

    if not os.path.exists(model_path):
        subprocess.run(["wget", model_info["url"], "-O", model_path], check=True)

    llm = CTransformers(
        model=model_path,
        model_type="llama",
        config={
            'top_p': 0.9,
            'repetition_penalty': 1.2,
            'max_new_tokens': 512,
            'temperature': 0.7,
            'context_length': 1024
        }
    )
    loaded_models[model_name] = llm
    return llm

@app.route('/')
def serve_index():
    return send_from_directory(app.static_folder, 'index.html')

@app.route('/<path:path>')
def serve_static(path):
    if not os.path.exists(os.path.join(app.static_folder, path)):
        return "Not Found", 404
    return send_from_directory(app.static_folder, path)

@app.route('/api/chat', methods=['POST'])
def chat_with_bot_api():
    token = request.headers.get('Authorization')
    if token != f"Bearer {API_TOKEN}":
        return jsonify({'error': 'Unauthorized'}), 401

    data = request.get_json()
    message = data.get('message')
    model_choice = data.get('model_choice')
    session_id = data.get('session_id', str(uuid.uuid4()))

    if not message or not model_choice:
        return jsonify({'error': 'Message and model_choice are required'}), 400

    current_llm = load_model(model_choice)

    if session_id not in chat_sessions:
        if vectorstore_instance is None:
            prepare_vectorstore()

        retriever = vectorstore_instance.as_retriever()
        chain = ConversationalRetrievalChain.from_llm(
            llm=current_llm,
            retriever=retriever,
            memory=ConversationBufferMemory(memory_key="chat_history", return_messages=True)
        )
        chat_sessions[session_id] = {"chain": chain, "history": [], "model_choice": model_choice}
    else:
        session_data = chat_sessions[session_id]
        if session_data.get("model_choice") != model_choice:
            retriever = vectorstore_instance.as_retriever()
            chain = ConversationalRetrievalChain.from_llm(
                llm=current_llm,
                retriever=retriever,
                memory=ConversationBufferMemory(memory_key="chat_history", return_messages=True)
            )
            chat_sessions[session_id]["chain"] = chain
            chat_sessions[session_id]["model_choice"] = model_choice

    chain = chat_sessions[session_id]["chain"]
    try:
        result = chain({"question": message})
        answer = result["answer"]

        chat_sessions[session_id]["history"].append(("You", message))
        chat_sessions[session_id]["history"].append(("Bot", answer))

        log_chat(session_id, message, answer)

        return jsonify({
            'answer': answer,
            'chat_history': chat_sessions[session_id]["history"]
        })
    except Exception as e:
        return jsonify({'error': f'Processing error: {str(e)}'}), 500

@app.route('/api/reset', methods=['POST'])
def reset_session_api():
    token = request.headers.get('Authorization')
    if token != f"Bearer {API_TOKEN}":
        return jsonify({'error': 'Unauthorized'}), 401

    data = request.get_json()
    session_id = data.get('session_id')

    if session_id in chat_sessions:
        chat_sessions.pop(session_id, None)
    return jsonify({'status': 'success', 'message': f'Session {session_id} reset'})

with app.app_context():
    if os.path.exists(COMBINED_DATASET_FILE_PATH):
        prepare_vectorstore()
    try:
        load_model("TinyLLaMA (1.1B)")
    except Exception as e:
        print(f"Model preload failed: {e}")

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=int(os.environ.get('PORT', 7860)), debug=False)



    import sqlite3
from flask import g

DATABASE = os.path.join(DATA_DIR, "users.db")

def get_db():
    db = getattr(g, "_database", None)
    if db is None:
        db = g._database = sqlite3.connect(DATABASE)
        db.row_factory = sqlite3.Row
    return db

@app.teardown_appcontext
def close_connection(exception):
    db = getattr(g, "_database", None)
    if db is not None:
        db.close()

def init_db():
    with app.app_context():
        db = get_db()
        db.execute("""
        CREATE TABLE IF NOT EXISTS users (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            username TEXT UNIQUE NOT NULL,
            password TEXT NOT NULL
        )
        """)
        db.commit()

import hashlib

def hash_password(password):
    return hashlib.sha256(password.encode()).hexdigest()

@app.route('/api/register', methods=['POST'])
def register():
    data = request.get_json()
    username = data.get("username")
    password = data.get("password")
    if not username or not password:
        return {"error": "Username and password required"}, 400

    hashed = hash_password(password)
    db = get_db()
    try:
        db.execute("INSERT INTO users (username, password) VALUES (?, ?)", (username, hashed))
        db.commit()
        return {"message": "User registered successfully"}
    except sqlite3.IntegrityError:
        return {"error": "Username already exists"}, 409

@app.route('/api/login', methods=['POST'])
def login():
    data = request.get_json()
    username = data.get("username")
    password = data.get("password")
    if not username or not password:
        return {"error": "Username and password required"}, 400

    hashed = hash_password(password)
    db = get_db()
    user = db.execute("SELECT * FROM users WHERE username = ? AND password = ?", (username, hashed)).fetchone()
    if user:
        return {"message": "Login successful"}
    else:
        return {"error": "Invalid credentials"}, 401

# Call init_db() once on startup to ensure DB is ready
init_db()