Update main.py

main.py

@@ -1,26 +1,63 @@
-from fastapi import FastAPI
+# main.py
+from fastapi import FastAPI, Query
 from fastapi.middleware.cors import CORSMiddleware
-from pydantic import BaseModel
 from transformers import pipeline
 
-# Load a chat model (you can choose a bigger one)
-generator = pipeline('text-generation', model='gpt2')  # gpt2 is small; for better, use GPT-Neo
+# ---------------------------
+# Initialize FastAPI app
+# ---------------------------
+app = FastAPI(title="Smart AI Backend with Memory")
 
-app = FastAPI()
-
-# Allow CORS for your frontend
+# Allow all origins (so you can call from any frontend)
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
+    allow_credentials=True,
     allow_methods=["*"],
-    allow_headers=["*"]
+    allow_headers=["*"],
 )
 
-class Prompt(BaseModel):
-    prompt: str
+# ---------------------------
+# Load small AI model
+# ---------------------------
+generator = pipeline("text-generation", model="distilgpt2")
+
+# ---------------------------
+# Conversation memory
+# ---------------------------
+conversation_history = []
+
+# ---------------------------
+# API endpoints
+# ---------------------------
+
+@app.get("/")
+def home():
+    return {"message": "Backend is working"}
 
 @app.get("/chat")
-def chat(prompt: str):
-    # Generate response
-    response = generator(prompt, max_length=100, do_sample=True)[0]['generated_text']
-    return {"reply": response}
+def chat(prompt: str = Query(..., description="User input text")):
+    """
+    Returns AI-generated reply based on user prompt
+    and keeps conversation history in memory.
+    """
+    try:
+        # Add user input to conversation history
+        conversation_history.append(f"User: {prompt}")
+
+        # Combine history as input for AI
+        input_text = "\n".join(conversation_history)
+
+        # Generate AI response
+        output = generator(input_text, max_length=150, do_sample=True)
+        response_text = output[0]['generated_text']
+
+        # Optionally, remove repeated prompt from generated text
+        response_clean = response_text.replace(input_text, "").strip()
+
+        # Save AI response to conversation history
+        conversation_history.append(f"AI: {response_clean}")
+
+        return {"reply": response_clean}
+    except Exception as e:
+        return {"error": str(e)}