From b2eba19fb1df5f8f658faf27472e980bbd5fcc5a Mon Sep 17 00:00:00 2001
From: ipu <uwu@ipu.moe>
Date: Fri, 19 Sep 2025 12:09:26 +0300
Subject: [PATCH] add user_session

---
 ...9_19_1203-31359fcda8a7_add_user_session.py | 45 +++++++++++++++++++
 src/database.py                               |  8 ++++
 src/services/chat_service.py                  | 37 ++++++++++++---
 3 files changed, 83 insertions(+), 7 deletions(-)
 create mode 100644 alembic/versions/2025_09_19_1203-31359fcda8a7_add_user_session.py

diff --git a/alembic/versions/2025_09_19_1203-31359fcda8a7_add_user_session.py b/alembic/versions/2025_09_19_1203-31359fcda8a7_add_user_session.py
new file mode 100644
index 0000000..5ea0d00
--- /dev/null
+++ b/alembic/versions/2025_09_19_1203-31359fcda8a7_add_user_session.py
@@ -0,0 +1,45 @@
+"""add user_session
+
+Revision ID: 31359fcda8a7
+Revises: 57f67bce2bec
+Create Date: 2025-09-19 12:03:40.032535
+
+"""
+from typing import Sequence, Union
+
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision: str = '31359fcda8a7'
+down_revision: Union[str, Sequence[str], None] = '57f67bce2bec'
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    """Upgrade schema."""
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.create_table('user_session',
+    sa.Column('id', sa.BigInteger(), nullable=False),
+    sa.Column('user_id', sa.String(), nullable=True),
+    sa.Column('session_id', sa.String(), nullable=True),
+    sa.PrimaryKeyConstraint('id')
+    )
+    with op.batch_alter_table('user_session', schema=None) as batch_op:
+        batch_op.create_index(batch_op.f('ix_user_session_id'), ['id'], unique=False)
+        batch_op.create_index(batch_op.f('ix_user_session_user_id'), ['user_id'], unique=False)
+
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    """Downgrade schema."""
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table('user_session', schema=None) as batch_op:
+        batch_op.drop_index(batch_op.f('ix_user_session_user_id'))
+        batch_op.drop_index(batch_op.f('ix_user_session_id'))
+
+    op.drop_table('user_session')
+    # ### end Alembic commands ###
diff --git a/src/database.py b/src/database.py
index 19c461c..8cbb5f1 100644
--- a/src/database.py
+++ b/src/database.py
@@ -34,5 +34,13 @@ class Webpage(Base):
     description = Column(String)
 
 
+class UserSession(Base):
+    __tablename__ = "user_session"
+
+    id = Column(BigInteger, primary_key=True, index=True)
+    user_id = Column(String, index=True)
+    session_id = Column(String)
+
+
 engine = create_engine(settings.DATABASE_URL)
 Session = sessionmaker(autocommit=False, autoflush=False, bind=engine)
diff --git a/src/services/chat_service.py b/src/services/chat_service.py
index 0aefdf7..705a0a1 100644
--- a/src/services/chat_service.py
+++ b/src/services/chat_service.py
@@ -7,6 +7,7 @@ from src.models import ApplicantParam, ChatHook, PlansParam, InsuranceChatContex
 from .session_service import session_service
 from ..api.v1.models import Source, HistoryItem
 from ..config import settings
+from ..database import Session, UserSession
 
 
 class ChatService:
@@ -134,15 +135,37 @@ class ChatService:
         
         # For now, return empty list - this would be populated when RAG is implemented
         return []
-    
-    async def process_insurance_chat(self, message: str, session_id: Optional[str] = None, uid: Optional[int] = None, current_page: Optional[str] = None, application: Optional[dict] = None) -> Dict[str, Any]:
+
+
+    async def get_user_session(self, uid: str) -> str | None:
+        with Session() as session:
+            session = session.query(UserSession).filter(UserSession.user_id == uid).first()
+            if not session:
+                return None
+            return session.session_id
+
+    async def create_user_session(self, uid: str, session_id: str):
+        with Session() as session:
+            user_session = UserSession(
+                user_id=uid,
+                session_id=session_id,
+            )
+            session.add(user_session)
+            session.commit()
+
+
+    async def process_insurance_chat(self, message: str, session_id: Optional[str] = None, uid: Optional[str] = None, current_page: Optional[str] = None, application: Optional[dict] = None) -> Dict[str, Any]:
         """Process an insurance chat request"""
         try:
-            if not session_id:
-                session_id = await session_service.create_session(agent_id=settings.TALESTORM_AGENT_ID)
-            
-            elif not await session_service.validate_session(session_id):
-                session_id = await session_service.create_session(agent_id=settings.TALESTORM_AGENT_ID)
+            if not session_id or not await session_service.validate_session(session_id):
+                if uid:
+                    session_id = await self.get_user_session(uid)
+                    if not session_id or not await session_service.validate_session(session_id):
+                        session_id = await session_service.create_session(agent_id=settings.TALESTORM_AGENT_ID)
+                    await self.create_user_session(uid, session_id)
+                else:
+                    session_id = await session_service.create_session(agent_id=settings.TALESTORM_AGENT_ID)
+
 
             instructions = ""
             if uid: