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('/') 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()