Chat Messages send and receive and also load the history

app/db/models.py

@@ -20,13 +20,33 @@ class UserDB(Base):
     password = Column(String(255), nullable=False)
 
     role = Column(String(20), nullable=False)  # "head" or "member"
-    family_code = Column(String(10), nullable=True)  # Only for members
+    family_code = Column(String(10), nullable=True)
+
     otp_code = Column(String(6), nullable=True)
     otp_expiry = Column(DateTime, nullable=True)
-    created_at = Column(DateTime(timezone=True), server_default=func.now())
-    family = relationship("Family", back_populates="head", uselist=False)
-    usage = relationship("UserUsage", back_populates="user", uselist=False)
 
+    created_at = Column(
+        DateTime(timezone=True),
+        server_default=func.now()
+    )
+
+    family = relationship(
+        "Family",
+        back_populates="head",
+        uselist=False
+    )
+
+    usage = relationship(
+        "UserUsage",
+        back_populates="user",
+        uselist=False
+    )
+
+    chats = relationship(
+        "ChatMessage",
+        back_populates="user",
+        cascade="all, delete-orphan"
+    )
 class UserUsage(Base):
     __tablename__ = "user_usage"
     id = Column(Integer, primary_key=True, index=True)
@@ -41,35 +61,3 @@ class UserUsage(Base):
     user = relationship("UserDB", back_populates="usage")
 
 
-# class PostDB(Base):
-#     __tablename__ = "posts"
-#     id = Column(Integer, primary_key=True, index=True)
-#     user_id = Column(Integer, ForeignKey("user_accounts.id", ondelete="CASCADE"), nullable=False)
-#     content = Column(Text, nullable=True)
-#     image_url = Column(String(300), nullable=True)
-#     created_at = Column(DateTime, default=datetime.utcnow, nullable=False)
-
-#     user = relationship("UserDB")
-#     comments = relationship("CommentDB", backref="post", cascade="all, delete-orphan")
-#     likes = relationship("LikeDB", back_populates="post", cascade="all, delete-orphan")
-
-# class CommentDB(Base):
-#     __tablename__ = "comments"
-#     id = Column(Integer, primary_key=True, index=True)
-#     content = Column(Text, nullable=False)
-#     created_at = Column(DateTime, default=datetime.utcnow, nullable=False)
-#     post_id = Column(Integer, ForeignKey("posts.id", ondelete="CASCADE"), nullable=False)
-#     user_id = Column(Integer, ForeignKey("user_accounts.id", ondelete="CASCADE"), nullable=False)
-#     user = relationship("UserDB")
-
-# class LikeDB(Base):
-#     __tablename__ = "likes"
-#     id = Column(Integer, primary_key=True, index=True)
-#     user_id = Column(Integer, ForeignKey("user_accounts.id", ondelete="CASCADE"), nullable=False)
-#     post_id = Column(Integer, ForeignKey("posts.id", ondelete="CASCADE"), nullable=False)
-#     created_at = Column(DateTime, default=datetime.utcnow, nullable=False)
-
-#     __table_args__ = (UniqueConstraint("user_id", "post_id", name="unique_like"),)
-
-#     user = relationship("UserDB")
-#     post = relationship("PostDB", back_populates="likes")

app/db/models_chat.py

@@ -0,0 +1,30 @@
+from sqlalchemy import Column, Integer, String, Text, DateTime, ForeignKey
+from sqlalchemy.orm import relationship
+from sqlalchemy.sql import func
+
+from app.db.database import Base
+
+
+class ChatMessage(Base):
+    __tablename__ = "chat_messages"
+
+    id = Column(Integer, primary_key=True, index=True)
+
+    user_id = Column(
+        Integer,
+        ForeignKey("user_accounts.id", ondelete="CASCADE"),
+        nullable=False
+    )
+
+    # user OR assistant
+    role = Column(String(20), nullable=False)
+
+    message = Column(Text, nullable=False)
+
+    created_at = Column(
+        DateTime(timezone=True),
+        server_default=func.now()
+    )
+
+    # relationship
+    user = relationship("UserDB", back_populates="chats")

app/routers/chat_ws.py

@@ -11,7 +11,7 @@ from app.services.date_extractor import DateExtractor
 from datetime import datetime
 import calendar
 from app.db.models import UserUsage
-
+from app.db.models_chat import ChatMessage
 router = APIRouter()
 ai = OpenAIService()
 date_extractor = DateExtractor()
@@ -298,8 +298,10 @@ async def chat_socket(websocket: WebSocket):
                 continue
 
             user_message = str(data["content"]).strip()
-            usage = get_or_create_usage(db, user_id)
+            #  Save user message
 
+            save_chat(db, user_id, "user", user_message)
+            usage = get_or_create_usage(db, user_id)
             print(f"[Chat] user_id={user_id} | count={usage.request_count} | paid={usage.is_paid}")
 
             # ── Usage gate ──
@@ -319,26 +321,31 @@ async def chat_socket(websocket: WebSocket):
 
             # ── Greeting — no usage count ──
             if intent == "greeting":
+
+                greeting_message = (
+                    "Hello! I'm FinanceAI, your personal finance assistant 💰\n\n"
+                    "I can help you with:\n"
+                    "• Your monthly budget breakdown\n"
+                    "• Spending predictions for next month\n"
+                    "• Tips to control and reduce your expenses\n\n"
+                    "What would you like to know?"
+                )
+                save_chat(db, user_id, "assistant", greeting_message)
                 await websocket.send_json({
                     "type": "assistant_message",
                     "allowed": True,
-                    "content": (
-                        "Hello! I'm FinanceAI, your personal finance assistant 💰\n\n"
-                        "I can help you with:\n"
-                        "• Your monthly budget breakdown\n"
-                        "• Spending predictions for next month\n"
-                        "• Tips to control and reduce your expenses\n\n"
-                        "What would you like to know?"
-                    ),
+                    "content": greeting_message,
                     "used_requests": usage.request_count,
                     "remaining_requests": max(MAX_FREE_REQUESTS - usage.request_count, 0),
                     "is_paid": usage.is_paid
                 })
-                continue
 
+                continue
             # ── General non-finance chat ──
             if intent == "general":
                 ai_reply = ai.chat_with_context(user_message=user_message, finance_data=None)
+                save_chat(db, user_id, "assistant", ai_reply)
+                
                 await send_and_count(websocket, db, usage, ai_reply)
                 continue
 
@@ -395,6 +402,8 @@ async def chat_socket(websocket: WebSocket):
                 }
 
                 ai_reply = ai.chat_with_context(user_message=user_message, finance_data=finance_context)
+                save_chat(db, user_id, "assistant", ai_reply)
+
                 await send_and_count(websocket, db, usage, ai_reply)
                 continue
 
@@ -429,6 +438,8 @@ async def chat_socket(websocket: WebSocket):
                 }
 
                 ai_reply = ai.chat_with_context(user_message=user_message, finance_data=finance_context)
+                save_chat(db, user_id, "assistant", ai_reply)
+
                 await send_and_count(websocket, db, usage, ai_reply)
                 continue
 
@@ -471,10 +482,50 @@ async def chat_socket(websocket: WebSocket):
             }
 
             ai_reply = ai.chat_with_context(user_message=user_message, finance_data=finance_context)
+            save_chat(db, user_id, "assistant", ai_reply)
+
             await send_and_count(websocket, db, usage, ai_reply)
 
     except WebSocketDisconnect:
         pass
 
     finally:
-        db.close()
+        db.close()
+    
+
+def save_chat(db: Session, user_id: int, role: str, message: str):
+        chat = ChatMessage(
+            user_id=user_id,
+            role=role,
+            message=message
+        )
+
+        db.add(chat)
+        db.commit()
+
+@router.get("/history")
+def get_chat_history(
+    current_user: UserDB = Depends(get_current_user),
+    db: Session = Depends(get_db)
+):
+    chats = (
+        db.query(ChatMessage)
+        .filter(ChatMessage.user_id == current_user.id)
+        .order_by(ChatMessage.created_at.asc())
+        .all()
+    )
+
+    response = []
+
+    for chat in chats:
+        response.append({
+            "id": chat.id,
+            "role": chat.role,
+            "message": chat.message,
+            "created_at": chat.created_at
+        })
+
+    return {
+        "status": True,
+        "messages": response
+    }