YO / db.py
NEXAS's picture
Upload 12 files
b3cb0b5 verified
Raw
History Blame Contribute Delete
14.6 kB
import os
from datetime import datetime, timedelta
from typing import List, Optional, Dict, Any
from dotenv import load_dotenv
from supabase import create_client, Client
from cache import cache # Import Redis cache
import uuid
load_dotenv()
SUPABASE_URL = os.getenv("SUPABASE_URL")
SUPABASE_KEY = os.getenv("SUPABASE_KEY")
class Database:
def __init__(self):
self.client = None
if SUPABASE_URL and SUPABASE_KEY:
try:
self.client: Client = create_client(SUPABASE_URL, SUPABASE_KEY)
except Exception as e:
print(f"Failed to initialize Supabase client: {e}")
# In-memory mock storage
self.mock_users = [
{"contact_number": "5550101", "name": "Alice Test", "created_at": datetime.now().isoformat()},
{"contact_number": "9730102", "name": "Naresh", "created_at": datetime.now().isoformat()}
]
self.mock_appointments = [
{
"id": "mock_apt_1",
"contact_number": "555-0101",
"appointment_time": "2026-01-22T10:00:00",
"status": "confirmed",
"purpose": "Checkup",
"created_at": datetime.now().isoformat()
}
]
self.mock_summaries = []
self.mock_chat_messages = []
self.cache = cache
# Initialize Mock Slots (Next 10 days)
self.mock_slots = []
base_time = datetime.now().replace(minute=0, second=0, microsecond=0)
for d in range(1, 11):
day = base_time + timedelta(days=d)
for h in [9, 10, 14, 16]:
slot_time = day.replace(hour=h).isoformat()
self.mock_slots.append({"slot_time": slot_time, "is_booked": False})
def get_available_slots(self) -> List[str]:
"""Get list of available slot times."""
if self.client:
try:
response = self.client.table("appointment_slots")\
.select("slot_time")\
.eq("is_booked", False)\
.gt("slot_time", datetime.now().isoformat())\
.order("slot_time")\
.execute()
return [row["slot_time"] for row in response.data]
except Exception as e:
print(f"Error fetching slots from DB: {e}")
# Mock fallback
# Filter mock slots that are in future and not booked
now_str = datetime.now().isoformat()
return [s["slot_time"] for s in self.mock_slots if not s["is_booked"] and s["slot_time"] > now_str]
def get_user(self, contact_number: str) -> Optional[Dict[str, Any]]:
"""Check if a user exists by contact number (with caching)."""
# Normalize input: remove non-digit characters
contact_number = "".join(filter(str.isdigit, str(contact_number)))
# Try cache first
cache_key = f"user:{contact_number}"
cached_user = self.cache.get(cache_key)
if cached_user:
return cached_user
if self.client:
try:
response = self.client.table("users").select("*").eq("contact_number", contact_number).execute()
if response.data:
user = response.data[0]
# Cache for 1 hour
self.cache.set(cache_key, user, ttl=3600)
return user
except Exception as e:
print(f"Error fetching user from DB (falling back to mock): {e}")
# Mock implementation fallback
for user in self.mock_users:
if user["contact_number"] == contact_number:
return user
return None
def create_user(self, contact_number: str, name: str = "Unknown") -> Optional[Dict[str, Any]]:
"""Create a new user."""
if self.client:
try:
data = {"contact_number": contact_number, "name": name}
# Use upsert to handle potential race conditions or existing users
response = self.client.table("users").upsert(data).execute()
if response.data:
return response.data[0]
except Exception as e:
print(f"Error creating user in DB (falling back to mock): {e}")
# Mock implementation fallback
new_user = {"contact_number": contact_number, "name": name, "created_at": datetime.now().isoformat()}
self.mock_users.append(new_user)
return new_user
def get_user_appointments(self, contact_number: str) -> List[Dict[str, Any]]:
"""Fetch past and upcoming appointments for a user."""
if self.client:
try:
response = self.client.table("appointments")\
.select("*")\
.eq("contact_number", contact_number)\
.order("appointment_time", desc=True)\
.execute()
return response.data
except Exception as e:
print(f"Error fetching appointments from DB (falling back to mock): {e}")
# Mock implementation fallback
return [
apt for apt in self.mock_appointments
if apt["contact_number"] == contact_number and apt["status"] != "cancelled"
]
def check_slot_availability(self, appointment_time: datetime) -> bool:
"""Check if a slot is valid and available."""
time_str = appointment_time.isoformat()
if self.client:
try:
# Check appointment_slots table for validity and availability
response = self.client.table("appointment_slots")\
.select("*")\
.eq("slot_time", time_str)\
.eq("is_booked", False)\
.execute()
return len(response.data) > 0
except Exception as e:
print(f"Error checking availability in DB (falling back to mock): {e}")
# Mock fallback
for slot in self.mock_slots:
if slot["slot_time"] == time_str:
return not slot["is_booked"]
return False
def book_appointment(self, contact_number: str, appointment_time: str, purpose: str = "General") -> Optional[Dict[str, Any]]:
"""Book an appointment and mark slot as booked."""
if self.client:
try:
# 1. Insert into appointments
data = {
"contact_number": contact_number,
"appointment_time": appointment_time,
"status": "confirmed",
"purpose": purpose
}
response = self.client.table("appointments").insert(data).execute()
# 2. Mark slot as booked
self.client.table("appointment_slots")\
.update({"is_booked": True})\
.eq("slot_time", appointment_time)\
.execute()
if response.data:
return response.data[0]
except Exception as e:
print(f"Error booking appointment in DB (falling back to mock): {e}")
# Mock implementation fallback
import random
apt_id = f"APT-{random.randint(1000, 9999)}"
new_apt = {
"id": apt_id,
"contact_number": contact_number,
"appointment_time": appointment_time,
"status": "confirmed",
"purpose": purpose,
"created_at": datetime.now().isoformat()
}
self.mock_appointments.append(new_apt)
# Mark mock slot as booked
for slot in self.mock_slots:
if slot["slot_time"] == appointment_time:
slot["is_booked"] = True
return new_apt
def cancel_appointment(self, appointment_id: str) -> bool:
"""Cancel an appointment."""
if self.client:
try:
response = self.client.table("appointments")\
.update({"status": "cancelled"})\
.eq("id", appointment_id)\
.execute()
return True
except Exception as e:
print(f"Error cancelling appointment in DB (falling back to mock): {e}")
# Mock implementation fallback
for apt in self.mock_appointments:
if apt["id"] == appointment_id:
apt["status"] = "cancelled"
return True
return False
def modify_appointment(self, appointment_id: str, new_time: str) -> bool:
"""Modify appointment time."""
if self.client:
try:
response = self.client.table("appointments")\
.update({"appointment_time": new_time})\
.eq("id", appointment_id)\
.execute()
return True
except Exception as e:
print(f"Error modifying appointment in DB (falling back to mock): {e}")
# Mock implementation fallback
for apt in self.mock_appointments:
if apt["id"] == appointment_id:
apt["appointment_time"] = new_time
return True
return False
def save_summary(self, contact_number: str, summary: str) -> bool:
"""Save the conversation summary."""
if self.client:
try:
data = {
"contact_number": contact_number,
"summary": summary,
"created_at": datetime.now().isoformat()
}
# Assuming a 'conversations' table exists
self.client.table("conversations").insert(data).execute()
return True
except Exception as e:
print(f"Error saving summary in DB (falling back to mock): {e}")
# Mock implementation fallback
print(f"Mock saving summary for {contact_number}: {summary}")
self.mock_summaries.append({
"contact_number": contact_number,
"summary": summary,
"created_at": datetime.now().isoformat()
})
return True
def save_chat_message(self, session_id: str, contact_number: str, role: str, content: str, tool_name: str = None, tool_args: dict = None) -> bool:
"""Save a single chat message to the database"""
if self.client:
try:
data = {
"session_id": session_id,
"contact_number": contact_number,
"role": role,
"content": content,
"tool_name": tool_name,
"tool_args": tool_args,
"created_at": datetime.now().isoformat()
}
self.client.table("chat_messages").insert(data).execute()
return True
except Exception as e:
print(f"Error saving chat message to DB (falling back to mock): {e}")
# Mock fallback
self.mock_chat_messages.append({
"session_id": session_id,
"contact_number": contact_number,
"role": role,
"content": content,
"tool_name": tool_name,
"tool_args": tool_args,
"created_at": datetime.now().isoformat()
})
return True
def save_chat_transcript(self, session_id: str, contact_number: str, messages: list) -> bool:
"""Save entire chat transcript (batch insert)"""
if not messages:
return False
if self.client:
try:
# Prepare batch data
data = []
for msg in messages:
data.append({
"session_id": session_id,
"contact_number": contact_number,
"role": msg.get("role"),
"content": msg.get("content"),
"tool_name": msg.get("tool_name"),
"tool_args": msg.get("tool_args"),
"created_at": datetime.now().isoformat()
})
# Batch insert
self.client.table("chat_messages").insert(data).execute()
print(f"✅ Saved {len(data)} chat messages to database")
return True
except Exception as e:
print(f"Error saving chat transcript to DB (falling back to mock): {e}")
# Mock fallback
for msg in messages:
self.mock_chat_messages.append({
"session_id": session_id,
"contact_number": contact_number,
"role": msg.get("role"),
"content": msg.get("content"),
"tool_name": msg.get("tool_name"),
"tool_args": msg.get("tool_args"),
"created_at": datetime.now().isoformat()
})
print(f"Mock saved {len(messages)} chat messages")
return True
def get_chat_history(self, contact_number: str, limit: int = 100) -> list:
"""Get chat history for a user"""
if self.client:
try:
response = self.client.table("chat_messages")\
.select("*")\
.eq("contact_number", contact_number)\
.order("created_at", desc=True)\
.limit(limit)\
.execute()
return response.data if response.data else []
except Exception as e:
print(f"Error fetching chat history: {e}")
# Mock fallback
return [msg for msg in self.mock_chat_messages if msg["contact_number"] == contact_number][-limit:]
# Hardcoded slots for the 'fetch_slots' requirement
AVAILABLE_SLOTS = [
"2026-01-22T09:00:00",
"2026-01-22T10:00:00",
"2026-01-22T14:00:00",
"2026-01-23T11:00:00",
"2026-01-23T15:00:00"
]