Updated

main.py

@@ -1,7 +1,3 @@
-# ===============================================
-# TechDisciples AI Backend — Stable + LangChain Memory
-# ===============================================
-
 from fastapi import FastAPI, HTTPException, Header
 from pydantic import BaseModel
 import torch
@@ -11,49 +7,36 @@ import logging
 from langchain.chains import LLMChain
 from langchain.prompts import PromptTemplate
 from langchain.memory import ConversationBufferMemory
-from langchain.llms.base import LLM
+from langchain.llms.base import LLM  # For custom LLM wrappers
 
 # Transformers pipeline
 from transformers import pipeline
 
-# ==============================
-# Logging Setup
-# ==============================
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger("TechDisciplesAI")
-
-# ==============================
-# FastAPI App Init
-# ==============================
-app = FastAPI(title="TechDisciples AI Backend")
 
-@app.get("/")
-async def root():
-    return {"status": "✅ TechDisciples AI Backend is running and stable."}
-
-# ==============================
-# Auth Config
-# ==============================
+# ===============================================
+# CONFIGURATION
+# ===============================================
 API_SECRET = "techdisciplesai404"
-
-def check_auth(x_api_key: str | None):
-    if not x_api_key or x_api_key != API_SECRET:
-        raise HTTPException(status_code=403, detail="Forbidden: Invalid API key")
-
-# ==============================
-# Request Schema
-# ==============================
-class QueryInput(BaseModel):
-    query: str
-
-# ==============================
-# Hugging Face Pipeline
-# ==============================
 MODEL_NAME = "google/flan-t5-large"
 DEVICE = 0 if torch.cuda.is_available() else -1
 
+# ===============================================
+# LOGGING SETUP
+# ===============================================
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger("TechDisciplesAI")
+
+# ===============================================
+# FASTAPI APP
+# ===============================================
+app = FastAPI(title="Tech Disciples AI (LangChain Conversational)", version="3.0")
+
+# ===============================================
+# LOAD MODEL USING PIPELINE + LANGCHAIN
+# ===============================================
 try:
     logger.info(f"🚀 Loading model: {MODEL_NAME}")
+
     hf_pipeline = pipeline(
         "text2text-generation",
         model=MODEL_NAME,
@@ -63,40 +46,23 @@ try:
         do_sample=True,
         top_p=0.9
     )
+
+    # You can later wrap hf_pipeline in a custom LLM class compatible with LLMChain
+    llm = hf_pipeline
     logger.info("✅ Model loaded successfully.")
+
 except Exception as e:
     logger.error(f"❌ Failed to load model: {e}")
-    hf_pipeline = None
-
-# ==============================
-# Minimal LangChain LLM Wrapper
-# ==============================
-class HuggingFaceLLM(LLM):
-    """Wraps a Hugging Face pipeline for LangChain compatibility."""
-    def __init__(self, pipeline):
-        self.pipeline = pipeline
-
-    @property
-    def _llm_type(self):
-        return "huggingface_pipeline"
-
-    def _call(self, prompt: str, stop=None):
-        try:
-            output = self.pipeline(prompt)
-            if isinstance(output, list) and len(output) > 0:
-                return output[0].get("generated_text", "")
-            return str(output)
-        except Exception as e:
-            logger.error(f"⚠️ LLM wrapper error: {e}")
-            return f"⚠️ Model error: {str(e)}"
-
-llm = HuggingFaceLLM(hf_pipeline) if hf_pipeline else None
-
-# ==============================
-# LangChain Memory + Prompt
-# ==============================
+    llm = None
+
+# ===============================================
+# MEMORY SYSTEM
+# ===============================================
 memory = ConversationBufferMemory(memory_key="conversation_history")
 
+# ===============================================
+# PROMPT TEMPLATE
+# ===============================================
 prompt_template = """
 You are Tech Disciples AI — a spiritually aware, intelligent, and kind conversational assistant.
 You offer thoughtful, biblical, and insightful answers with grace, empathy, and calm intelligence.
@@ -113,18 +79,33 @@ prompt = PromptTemplate(
     input_variables=["conversation_history", "query"]
 )
 
+# ===============================================
+# LLM CHAIN (with memory)
+# ===============================================
 chain = LLMChain(
-    llm=llm,
     prompt=prompt,
+    llm=llm,
     memory=memory
 )
 
-# ==============================
-# Endpoint
-# ==============================
+# ===============================================
+# REQUEST MODEL
+# ===============================================
+class QueryInput(BaseModel):
+    query: str
+    session_id: str | None = "default"
+
+# ===============================================
+# ROUTES
+# ===============================================
+@app.get("/")
+async def root():
+    return {"message": "✅ Tech Disciples AI (LangChain Memory) is running."}
+
 @app.post("/ai-chat")
 async def ai_chat(data: QueryInput, x_api_key: str = Header(None)):
-    check_auth(x_api_key)
+    if x_api_key != API_SECRET:
+        raise HTTPException(status_code=403, detail="Forbidden: Invalid API key")
 
     if not llm:
         raise HTTPException(status_code=500, detail="Model not initialized")