Update src/core/services/database_service.py

src/core/services/database_service.py

@@ -1,4 +1,4 @@
-"""Enhanced database service with comprehensive data model and synthetic data generation"""
+"""Enhanced database service with persistent storage and account management"""
 import sqlite3
 from contextlib import contextmanager
 from pathlib import Path
@@ -8,15 +8,28 @@ from datetime import datetime, timedelta
 import random
 from faker import Faker
 import uuid
-
+import os
 
 class DatabaseService:
-    def __init__(self, db_path: str = "data/robata.db"):
-        """Initialize database service"""
+    def __init__(self, db_path: str = "/data/robata.db"):
+        """Initialize database service with persistent storage"""
+        # Use /data directory for persistence in Hugging Face Spaces
         self.db_path = db_path
         self.fake = Faker()
         Path(db_path).parent.mkdir(parents=True, exist_ok=True)
         self.setup_database()
+        
+        # Only generate data if database is empty
+        if not self._has_data():
+            self.generate_synthetic_data()
+
+    def _has_data(self) -> bool:
+        """Check if database has existing data"""
+        with self.get_db() as conn:
+            cursor = conn.cursor()
+            cursor.execute("SELECT COUNT(*) FROM users")
+            user_count = cursor.fetchone()[0]
+            return user_count > 0
 
     @contextmanager
     def get_db(self):
@@ -29,6 +42,7 @@ class DatabaseService:
             conn.close()
 
     def setup_database(self):
+        """Set up database schema"""
         with self.get_db() as conn:
             c = conn.cursor()
             c.executescript('''
@@ -64,133 +78,148 @@ class DatabaseService:
                     FOREIGN KEY (parent_account_id) REFERENCES accounts (id),
                     FOREIGN KEY (account_owner_id) REFERENCES users (id)
                 );
-                
-                CREATE TABLE IF NOT EXISTS contacts (
-                    id TEXT PRIMARY KEY,
-                    account_id TEXT,
-                    first_name TEXT,
-                    last_name TEXT,
-                    email TEXT,
-                    phone TEXT,
-                    title TEXT,
-                    department TEXT,
-                    reports_to_id TEXT,
-                    influence_level TEXT,
-                    engagement_score REAL,
-                    preferences TEXT,
-                    created_at TEXT,
-                    last_contacted TEXT,
-                    FOREIGN KEY (account_id) REFERENCES accounts (id),
-                    FOREIGN KEY (reports_to_id) REFERENCES contacts (id)
-                );
-                CREATE TABLE IF NOT EXISTS interactions (
-                    id TEXT PRIMARY KEY,
-                    type TEXT,
-                    account_id TEXT,
-                    owner_id TEXT,
-                    transcript TEXT,
-                    summary TEXT,
-                    sentiment_score REAL,
-                    metadata TEXT,
-                    created_at TEXT,
-                    FOREIGN KEY (account_id) REFERENCES accounts (id),
-                    FOREIGN KEY (owner_id) REFERENCES users (id)
-                );
             ''')
 
     def generate_synthetic_data(self):
-        """Generate comprehensive synthetic data"""
+        """Generate synthetic users and accounts with proper linkage"""
         with self.get_db() as conn:
             c = conn.cursor()
-
-            # Configuration
-            NUM_USERS = 50
-            NUM_ACCOUNTS = 50
-            NUM_CONTACTS = 300
-
-            # Helper function to generate timestamp within last year
-            def random_date():
-                days = random.randint(0, 365)
-                return (datetime.now() - timedelta(days=days)).isoformat()
-
+            
             # Generate Users
-            roles = ['sales_rep'] * 35 + ['regional_lead'] * 10 + ['head_of_sales'] * 5
-            departments = ['Sales', 'Consulting', 'Technology', 'Operations']
-            regions = ['North', 'South', 'East', 'West', 'Central']
-
             users = []
-            for _ in range(NUM_USERS):
-                role = random.choice(roles)
-                dept = random.choice(departments)
+            user_ids = []  # Keep track of user IDs for account assignment
+            
+            # Predefined test user
+            test_user_id = str(uuid.uuid4())
+            users.append({
+                'id': test_user_id,
+                'email': 'test@example.com',  # Default test login
+                'name': 'Test User',
+                'role': 'sales_rep',
+                'department': 'Sales',
+                'title': 'Senior Sales Representative',
+                'region': 'North',
+                'quota': 1000000.0,
+                'created_at': datetime.now().isoformat(),
+                'last_login': datetime.now().isoformat()
+            })
+            user_ids.append(test_user_id)
+            
+            # Generate additional users
+            for _ in range(10):
+                user_id = str(uuid.uuid4())
+                user_ids.append(user_id)
                 users.append({
-                    'id': str(uuid.uuid4()),
+                    'id': user_id,
                     'email': self.fake.company_email(),
                     'name': self.fake.name(),
-                    'role': role,
-                    'department': dept,
-                    'title': f"Senior {dept} {role.title()}",
-                    'region': random.choice(regions),
+                    'role': random.choice(['sales_rep', 'regional_lead', 'head_of_sales']),
+                    'department': random.choice(['Sales', 'Consulting', 'Technology']),
+                    'title': 'Senior Sales Representative',
+                    'region': random.choice(['North', 'South', 'East', 'West']),
                     'quota': random.uniform(500000, 2000000),
-                    'created_at': random_date(),
-                    'last_login': random_date()
+                    'created_at': datetime.now().isoformat(),
+                    'last_login': datetime.now().isoformat()
+                })
+            
+            # Insert users
+            c.executemany('''
+                INSERT OR REPLACE INTO users VALUES (
+                    :id, :email, :name, :role, :department, :title, :region,
+                    :quota, :created_at, :last_login
+                )
+            ''', users)
+            
+            # Generate Accounts
+            accounts = []
+            industries = ['Technology', 'Healthcare', 'Financial Services', 'Manufacturing', 'Retail']
+            
+            # Ensure test user has accounts
+            for _ in range(3):
+                accounts.append({
+                    'id': str(uuid.uuid4()),
+                    'name': self.fake.company(),
+                    'parent_account_id': None,
+                    'industry': random.choice(industries),
+                    'status': 'active',
+                    'website': self.fake.url(),
+                    'annual_revenue': random.uniform(1000000, 100000000),
+                    'employee_count': random.randint(50, 10000),
+                    'technology_stack': json.dumps(['Python', 'React', 'AWS']),
+                    'region': 'North',
+                    'address': self.fake.address(),
+                    'account_owner_id': test_user_id,  # Assign to test user
+                    'engagement_score': random.uniform(0, 100),
+                    'created_at': datetime.now().isoformat(),
+                    'last_activity_at': datetime.now().isoformat()
                 })
+            
+            # Generate additional accounts
+            for user_id in user_ids:
+                for _ in range(random.randint(2, 5)):
+                    accounts.append({
+                        'id': str(uuid.uuid4()),
+                        'name': self.fake.company(),
+                        'parent_account_id': None,
+                        'industry': random.choice(industries),
+                        'status': 'active',
+                        'website': self.fake.url(),
+                        'annual_revenue': random.uniform(1000000, 100000000),
+                        'employee_count': random.randint(50, 10000),
+                        'technology_stack': json.dumps(['Python', 'React', 'AWS']),
+                        'region': random.choice(['North', 'South', 'East', 'West']),
+                        'address': self.fake.address(),
+                        'account_owner_id': user_id,
+                        'engagement_score': random.uniform(0, 100),
+                        'created_at': datetime.now().isoformat(),
+                        'last_activity_at': datetime.now().isoformat()
+                    })
+            
+            # Insert accounts
+            c.executemany('''
+                INSERT OR REPLACE INTO accounts VALUES (
+                    :id, :name, :parent_account_id, :industry, :status, :website,
+                    :annual_revenue, :employee_count, :technology_stack, :region,
+                    :address, :account_owner_id, :engagement_score, :created_at,
+                    :last_activity_at
+                )
+            ''', accounts)
+            
+            conn.commit()
 
-            c.executemany('''INSERT INTO users VALUES (
-                :id, :email, :name, :role, :department, :title, :region,
-                :quota, :created_at, :last_login
-            )''', users)
+    def add_account(self, account_data: Dict[str, Any]) -> str:
+        """Add a new account to the database"""
+        account_id = str(uuid.uuid4())
+        account_data['id'] = account_id
+        account_data['created_at'] = datetime.now().isoformat()
+        account_data['last_activity_at'] = datetime.now().isoformat()
+        
+        with self.get_db() as conn:
+            c = conn.cursor()
+            placeholders = ', '.join(['?' for _ in account_data])
+            columns = ', '.join(account_data.keys())
+            sql = f'INSERT INTO accounts ({columns}) VALUES ({placeholders})'
+            c.execute(sql, list(account_data.values()))
+            conn.commit()
+        
+        return account_id
 
     def get_user_by_email(self, email: str) -> Optional[Dict]:
         """Get user by email address"""
         with self.get_db() as conn:
-            # Print all users for debugging
-            all_users = conn.execute("SELECT email FROM users").fetchall()
-            print(f"Available users: {[u[0] for u in all_users]}")
-            
             cursor = conn.execute(
                 "SELECT * FROM users WHERE email = ?",
                 (email,)
             )
             row = cursor.fetchone()
-            if row:
-                return dict(row)
-            print(f"No user found for email: {email}")
-            return None
-            
+            return dict(row) if row else None
+
     def get_user_accounts(self, user_id: str) -> List[Dict]:
         """Get accounts associated with user"""
         with self.get_db() as conn:
             cursor = conn.execute("""
-                SELECT DISTINCT a.*
-                FROM accounts a
-                WHERE a.account_owner_id = ?
-                ORDER BY a.name
+                SELECT * FROM accounts 
+                WHERE account_owner_id = ?
+                ORDER BY name
             """, (user_id,))
-            return [dict(row) for row in cursor.fetchall()]
-
-    def get_account_metrics(self, account_id: str) -> Dict:
-        """Get metrics for a specific account"""
-        with self.get_db() as conn:
-            contacts_count = conn.execute(
-                "SELECT COUNT(*) FROM contacts WHERE account_id = ?",
-                (account_id,)
-            ).fetchone()[0]
-            
-            return {
-                'contact_count': contacts_count,
-                'interaction_count': 0,  # Default for now
-                'avg_sentiment': 0  # Default for now
-            }
-
-    def get_recent_interactions(self, user_id: str, limit: int = 10) -> List[Dict]:
-        """Get recent interactions for a user"""
-        with self.get_db() as conn:
-            cursor = conn.execute("""
-                SELECT i.*, a.name as account_name
-                FROM interactions i
-                JOIN accounts a ON i.account_id = a.id
-                WHERE i.owner_id = ?
-                ORDER BY i.created_at DESC
-                LIMIT ?
-            """, (user_id, limit))
-            return [dict(row) for row in cursor.fetchall()]
+            return [dict(row) for row in cursor.fetchall()]