Updated

main.py

@@ -4,78 +4,111 @@ from transformers import pipeline
 from langchain.llms import HuggingFacePipeline
 from langchain.chains import LLMChain
 from langchain.prompts import PromptTemplate
+from langchain.memory import ConversationBufferMemory
 import torch
 import logging
 
-# ===== CONFIG =====
+# ===============================================
+# CONFIGURATION
+# ===============================================
 API_SECRET = "techdisciplesai404"
 MODEL_NAME = "google/flan-t5-large"
 DEVICE = 0 if torch.cuda.is_available() else -1
 
-# ===== LOGGING =====
+# ===============================================
+# LOGGING SETUP
+# ===============================================
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger("TechDisciplesAI")
 
-# ===== INITIALIZE APP =====
-app = FastAPI(title="TechDisciples AI", version="2.0")
+# ===============================================
+# FASTAPI APP
+# ===============================================
+app = FastAPI(title="Tech Disciples AI (LangChain Conversational)", version="3.0")
 
-# ===== MODEL SETUP =====
+# ===============================================
+# LOAD MODEL USING PIPELINE + LANGCHAIN
+# ===============================================
 try:
     logger.info(f"🚀 Loading model: {MODEL_NAME}")
-    pipe = pipeline(
+
+    hf_pipeline = pipeline(
         "text2text-generation",
         model=MODEL_NAME,
         device=DEVICE,
         max_new_tokens=256,
         temperature=0.3,
-        do_sample=True
+        do_sample=True,
+        top_p=0.9
     )
 
-    llm = HuggingFacePipeline(pipeline=pipe)
+    llm = HuggingFacePipeline(pipeline=hf_pipeline)
     logger.info("✅ Model loaded successfully.")
+
 except Exception as e:
     logger.error(f"❌ Failed to load model: {e}")
     llm = None
 
-# ===== PROMPT TEMPLATE =====
+# ===============================================
+# MEMORY SYSTEM
+# ===============================================
+memory = ConversationBufferMemory(memory_key="conversation_history")
+
+# ===============================================
+# PROMPT TEMPLATE
+# ===============================================
 prompt_template = """
-You are a Christian conversational AI named TechDisciples AI. 
-Answer the question naturally and clearly, providing biblical or inspirational insight where possible.
+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.
 
-Question: {query}
+Conversation so far:
+{conversation_history}
 
-Response:
+User: {query}
+Tech Disciples AI:
 """
 
-prompt = PromptTemplate(template=prompt_template, input_variables=["query"])
-chain = LLMChain(prompt=prompt, llm=llm)
-
-# ===== REQUEST MODEL =====
+prompt = PromptTemplate(
+    template=prompt_template,
+    input_variables=["conversation_history", "query"]
+)
+
+# ===============================================
+# LLM CHAIN (with memory)
+# ===============================================
+chain = LLMChain(
+    prompt=prompt,
+    llm=llm,
+    memory=memory
+)
+
+# ===============================================
+# REQUEST MODEL
+# ===============================================
 class QueryInput(BaseModel):
     query: str
+    session_id: str | None = "default"  # optional: could be user/session-based
 
+# ===============================================
+# ROUTES
+# ===============================================
+@app.get("/")
+async def root():
+    return {"message": "✅ Tech Disciples AI (LangChain Memory) is running."}
 
-# ===== ROUTES =====
 @app.post("/ai-chat")
 async def ai_chat(data: QueryInput, x_api_key: str = Header(None)):
+    # --- Authentication ---
     if x_api_key != API_SECRET:
         raise HTTPException(status_code=403, detail="Forbidden: Invalid API key")
 
-    if llm is None:
+    if not llm:
         raise HTTPException(status_code=500, detail="Model not initialized")
 
-    user_query = data.query.strip()
-    if not user_query:
-        raise HTTPException(status_code=400, detail="Query cannot be empty")
-
+    # --- Process Query ---
     try:
-        response = chain.run(query=user_query)
-        return {"reply": response.strip(), "tone_used": "neutral"}
+        response = chain.run(query=data.query.strip())
+        return {"reply": response.strip()}
     except Exception as e:
-        logger.error(f"⚠️ Generation error: {e}")
+        logger.error(f"⚠️ Model error: {e}")
         raise HTTPException(status_code=500, detail="Model failed to respond")
-
-
-@app.get("/")
-async def root():
-    return {"message": "✅ TechDisciples AI (LangChain) is running."}