Updated

main.py

@@ -1,18 +1,22 @@
+# ===============================================
+# Tech Disciples AI Backend — Main.py
+# ===============================================
+
 from fastapi import FastAPI, HTTPException, Header
 from pydantic import BaseModel
 import torch
 import logging
 
-# LangChain imports (stable >=1.0)
-from langchain.chains import LLMChain
+# LangChain ≥1.0 imports
 from langchain.prompts import PromptTemplate
 from langchain.memory import ConversationBufferMemory
-from langchain.llms.base import LLM  # For custom LLM wrappers
+from langchain.chains import LLMChain
+from langchain.llms.base import LLM
+from typing import Optional, List
 
 # Transformers pipeline
 from transformers import pipeline
 
-
 # ===============================================
 # CONFIGURATION
 # ===============================================
@@ -32,7 +36,7 @@ logger = logging.getLogger("TechDisciplesAI")
 app = FastAPI(title="Tech Disciples AI (LangChain Conversational)", version="3.0")
 
 # ===============================================
-# LOAD MODEL USING PIPELINE + LANGCHAIN
+# HUGGING FACE PIPELINE
 # ===============================================
 try:
     logger.info(f"🚀 Loading model: {MODEL_NAME}")
@@ -47,13 +51,31 @@ try:
         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.")
+    logger.info("✅ Hugging Face pipeline loaded successfully.")
 
 except Exception as e:
-    logger.error(f"❌ Failed to load model: {e}")
-    llm = None
+    logger.error(f"❌ Failed to load Hugging Face pipeline: {e}")
+    hf_pipeline = None
+
+# ===============================================
+# HUGGING FACE LLM WRAPPER FOR LANGCHAIN
+# ===============================================
+class HFLLMWrapper(LLM):
+    def __init__(self, pipeline):
+        self.pipeline = pipeline
+
+    @property
+    def _llm_type(self) -> str:
+        return "hf_pipeline"
+
+    def _call(self, prompt: str, stop: Optional[List[str]] = None) -> str:
+        output = self.pipeline(prompt)
+        if isinstance(output, list) and len(output) > 0:
+            return output[0].get("generated_text", str(output[0]))
+        return str(output)
+
+# Initialize LLM wrapper
+llm = HFLLMWrapper(hf_pipeline) if hf_pipeline else None
 
 # ===============================================
 # MEMORY SYSTEM
@@ -104,12 +126,14 @@ async def root():
 
 @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 not llm:
         raise HTTPException(status_code=500, detail="Model not initialized")
 
+    # --- Process Query ---
     try:
         response = chain.run(query=data.query.strip())
         return {"reply": response.strip()}

requirements.txt

@@ -3,7 +3,6 @@ uvicorn[standard]
 torch
 transformers
 accelerate
-langchain==1.0.0
+langchain>=1.0
 huggingface-hub
 pydantic
-python-multipart