mem-llm 2.0.0__py3-none-any.whl

mem_llm/data_export_import.py

@@ -0,0 +1,640 @@
+"""
+Data Export/Import System
+Supports multiple formats and databases: JSON, CSV, SQLite, PostgreSQL, MongoDB
+"""
+
+import json
+import csv
+import sqlite3
+from datetime import datetime
+from typing import Dict, List, Optional, Any, Union
+from pathlib import Path
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+class DataExporter:
+    """Export memory data to various formats and databases"""
+    
+    def __init__(self, memory_manager):
+        """
+        Args:
+            memory_manager: MemoryManager or SQLMemoryManager instance
+        """
+        self.memory = memory_manager
+        
+    def export_to_json(self, user_id: str, output_file: str) -> Dict[str, Any]:
+        """
+        Export user data to JSON file
+        
+        Args:
+            user_id: User ID to export
+            output_file: Output JSON file path
+            
+        Returns:
+            Export statistics
+        """
+        try:
+            # Get all user data
+            conversations = self.memory.get_recent_conversations(user_id, limit=1000)
+            profile = getattr(self.memory, 'user_profiles', {}).get(user_id, {})
+            
+            data = {
+                'user_id': user_id,
+                'export_date': datetime.now().isoformat(),
+                'conversations': conversations,
+                'profile': profile,
+                'metadata': {
+                    'total_conversations': len(conversations),
+                    'format': 'json',
+                    'version': '1.0'
+                }
+            }
+            
+            # Write to file
+            output_path = Path(output_file)
+            output_path.parent.mkdir(parents=True, exist_ok=True)
+            
+            with open(output_path, 'w', encoding='utf-8') as f:
+                json.dump(data, f, ensure_ascii=False, indent=2)
+            
+            logger.info(f"Exported {len(conversations)} conversations to {output_file}")
+            
+            return {
+                'success': True,
+                'file': str(output_path),
+                'conversations': len(conversations),
+                'size_bytes': output_path.stat().st_size
+            }
+            
+        except Exception as e:
+            logger.error(f"JSON export failed: {e}")
+            return {'success': False, 'error': str(e)}
+    
+    def export_to_csv(self, user_id: str, output_file: str) -> Dict[str, Any]:
+        """
+        Export conversations to CSV file
+        
+        Args:
+            user_id: User ID to export
+            output_file: Output CSV file path
+            
+        Returns:
+            Export statistics
+        """
+        try:
+            conversations = self.memory.get_recent_conversations(user_id, limit=1000)
+            
+            output_path = Path(output_file)
+            output_path.parent.mkdir(parents=True, exist_ok=True)
+            
+            with open(output_path, 'w', newline='', encoding='utf-8') as f:
+                writer = csv.DictWriter(f, fieldnames=[
+                    'timestamp', 'user_message', 'bot_response', 'metadata'
+                ])
+                writer.writeheader()
+                
+                for conv in conversations:
+                    writer.writerow({
+                        'timestamp': conv.get('timestamp', ''),
+                        'user_message': conv.get('user_message', ''),
+                        'bot_response': conv.get('bot_response', ''),
+                        'metadata': json.dumps(conv.get('metadata', {}))
+                    })
+            
+            logger.info(f"Exported {len(conversations)} conversations to CSV: {output_file}")
+            
+            return {
+                'success': True,
+                'file': str(output_path),
+                'conversations': len(conversations),
+                'size_bytes': output_path.stat().st_size
+            }
+            
+        except Exception as e:
+            logger.error(f"CSV export failed: {e}")
+            return {'success': False, 'error': str(e)}
+    
+    def export_to_sqlite(self, user_id: str, db_file: str) -> Dict[str, Any]:
+        """
+        Export to SQLite database
+        
+        Args:
+            user_id: User ID to export
+            db_file: SQLite database file path
+            
+        Returns:
+            Export statistics
+        """
+        try:
+            conversations = self.memory.get_recent_conversations(user_id, limit=1000)
+            
+            db_path = Path(db_file)
+            db_path.parent.mkdir(parents=True, exist_ok=True)
+            
+            conn = sqlite3.connect(str(db_path))
+            cursor = conn.cursor()
+            
+            # Create table
+            cursor.execute('''
+                CREATE TABLE IF NOT EXISTS conversations (
+                    id INTEGER PRIMARY KEY AUTOINCREMENT,
+                    user_id TEXT NOT NULL,
+                    timestamp TEXT NOT NULL,
+                    user_message TEXT NOT NULL,
+                    bot_response TEXT NOT NULL,
+                    metadata TEXT
+                )
+            ''')
+            
+            # Insert conversations
+            for conv in conversations:
+                cursor.execute('''
+                    INSERT INTO conversations 
+                    (user_id, timestamp, user_message, bot_response, metadata)
+                    VALUES (?, ?, ?, ?, ?)
+                ''', (
+                    user_id,
+                    conv.get('timestamp', datetime.now().isoformat()),
+                    conv.get('user_message', ''),
+                    conv.get('bot_response', ''),
+                    json.dumps(conv.get('metadata', {}))
+                ))
+            
+            conn.commit()
+            conn.close()
+            
+            logger.info(f"Exported {len(conversations)} conversations to SQLite: {db_file}")
+            
+            return {
+                'success': True,
+                'file': str(db_path),
+                'conversations': len(conversations),
+                'size_bytes': db_path.stat().st_size
+            }
+            
+        except Exception as e:
+            logger.error(f"SQLite export failed: {e}")
+            return {'success': False, 'error': str(e)}
+    
+    def export_to_postgresql(self, user_id: str, connection_string: str) -> Dict[str, Any]:
+        """
+        Export to PostgreSQL database
+        
+        Args:
+            user_id: User ID to export
+            connection_string: PostgreSQL connection string
+                             (e.g., "postgresql://user:pass@localhost/dbname")
+            
+        Returns:
+            Export statistics
+        """
+        try:
+            import psycopg2
+            from psycopg2.extensions import ISOLATION_LEVEL_AUTOCOMMIT
+        except ImportError:
+            return {
+                'success': False,
+                'error': 'psycopg2 not installed. Install: pip install psycopg2-binary'
+            }
+        
+        try:
+            conversations = self.memory.get_recent_conversations(user_id, limit=1000)
+            
+            # Parse connection string to get database name
+            import re
+            match = re.search(r'/([^/]+)(?:\?|$)', connection_string)
+            db_name = match.group(1) if match else None
+            
+            # Try to connect, if database doesn't exist, create it
+            try:
+                conn = psycopg2.connect(connection_string)
+            except psycopg2.OperationalError as e:
+                if "does not exist" in str(e) and db_name:
+                    # Connect to default 'postgres' database to create new one
+                    base_conn_string = connection_string.rsplit('/', 1)[0] + '/postgres'
+                    temp_conn = psycopg2.connect(base_conn_string)
+                    temp_conn.set_isolation_level(ISOLATION_LEVEL_AUTOCOMMIT)
+                    temp_cursor = temp_conn.cursor()
+                    temp_cursor.execute(f'CREATE DATABASE {db_name}')
+                    temp_cursor.close()
+                    temp_conn.close()
+                    logger.info(f"Created PostgreSQL database: {db_name}")
+                    
+                    # Now connect to the new database
+                    conn = psycopg2.connect(connection_string)
+                else:
+                    raise
+            
+            cursor = conn.cursor()
+            
+            # Create table if not exists
+            cursor.execute('''
+                CREATE TABLE IF NOT EXISTS conversations (
+                    id SERIAL PRIMARY KEY,
+                    user_id VARCHAR(255) NOT NULL,
+                    timestamp TIMESTAMP NOT NULL,
+                    user_message TEXT NOT NULL,
+                    bot_response TEXT NOT NULL,
+                    metadata JSONB
+                )
+            ''')
+            
+            # Insert conversations
+            for conv in conversations:
+                cursor.execute('''
+                    INSERT INTO conversations 
+                    (user_id, timestamp, user_message, bot_response, metadata)
+                    VALUES (%s, %s, %s, %s, %s)
+                ''', (
+                    user_id,
+                    conv.get('timestamp', datetime.now().isoformat()),
+                    conv.get('user_message', ''),
+                    conv.get('bot_response', ''),
+                    json.dumps(conv.get('metadata', {}))
+                ))
+            
+            conn.commit()
+            cursor.close()
+            conn.close()
+            
+            logger.info(f"Exported {len(conversations)} conversations to PostgreSQL")
+            
+            return {
+                'success': True,
+                'database': 'postgresql',
+                'conversations': len(conversations),
+                'database_created': db_name is not None
+            }
+            
+        except Exception as e:
+            logger.error(f"PostgreSQL export failed: {e}")
+            return {'success': False, 'error': str(e)}
+    
+    def export_to_mongodb(self, user_id: str, connection_string: str, 
+                         database: str = 'mem_llm', collection: str = 'conversations') -> Dict[str, Any]:
+        """
+        Export to MongoDB database
+        
+        Args:
+            user_id: User ID to export
+            connection_string: MongoDB connection string
+                             (e.g., "mongodb://localhost:27017/")
+            database: Database name
+            collection: Collection name
+            
+        Returns:
+            Export statistics
+        """
+        try:
+            from pymongo import MongoClient
+        except ImportError:
+            return {
+                'success': False,
+                'error': 'pymongo not installed. Install: pip install pymongo'
+            }
+        
+        try:
+            conversations = self.memory.get_recent_conversations(user_id, limit=1000)
+            
+            client = MongoClient(connection_string)
+            
+            # MongoDB automatically creates database and collection if they don't exist
+            db = client[database]
+            coll = db[collection]
+            
+            # Check if this is a new database/collection
+            is_new_db = database not in client.list_database_names()
+            is_new_collection = collection not in db.list_collection_names()
+            
+            if is_new_db:
+                logger.info(f"Creating MongoDB database: {database}")
+            if is_new_collection:
+                logger.info(f"Creating MongoDB collection: {collection}")
+            
+            # Prepare documents
+            documents = []
+            for conv in conversations:
+                doc = {
+                    'user_id': user_id,
+                    'timestamp': conv.get('timestamp', datetime.now().isoformat()),
+                    'user_message': conv.get('user_message', ''),
+                    'bot_response': conv.get('bot_response', ''),
+                    'metadata': conv.get('metadata', {}),
+                    'export_date': datetime.now()
+                }
+                documents.append(doc)
+            
+            # Insert documents
+            result = coll.insert_many(documents)
+            
+            client.close()
+            
+            logger.info(f"Exported {len(conversations)} conversations to MongoDB")
+            
+            return {
+                'success': True,
+                'database': 'mongodb',
+                'conversations': len(conversations),
+                'inserted_ids': len(result.inserted_ids),
+                'database_created': is_new_db,
+                'collection_created': is_new_collection
+            }
+            
+        except Exception as e:
+            logger.error(f"MongoDB export failed: {e}")
+            return {'success': False, 'error': str(e)}
+
+
+class DataImporter:
+    """Import memory data from various formats and databases"""
+    
+    def __init__(self, memory_manager):
+        """
+        Args:
+            memory_manager: MemoryManager or SQLMemoryManager instance
+        """
+        self.memory = memory_manager
+    
+    def import_from_json(self, input_file: str, user_id: Optional[str] = None) -> Dict[str, Any]:
+        """
+        Import user data from JSON file
+        
+        Args:
+            input_file: Input JSON file path
+            user_id: Override user ID (use file's user_id if None)
+            
+        Returns:
+            Import statistics
+        """
+        try:
+            input_path = Path(input_file)
+            
+            if not input_path.exists():
+                return {'success': False, 'error': f'File not found: {input_file}'}
+            
+            with open(input_path, 'r', encoding='utf-8') as f:
+                data = json.load(f)
+            
+            # Get user ID
+            target_user_id = user_id or data.get('user_id')
+            if not target_user_id:
+                return {'success': False, 'error': 'No user_id specified'}
+            
+            # Import conversations
+            conversations = data.get('conversations', [])
+            imported = 0
+            
+            for conv in conversations:
+                self.memory.add_conversation(
+                    target_user_id,
+                    conv.get('user_message', ''),
+                    conv.get('bot_response', ''),
+                    conv.get('metadata', {})
+                )
+                imported += 1
+            
+            # Import profile if available
+            if 'profile' in data and hasattr(self.memory, 'update_profile'):
+                self.memory.update_profile(target_user_id, data['profile'])
+            
+            logger.info(f"Imported {imported} conversations from {input_file}")
+            
+            return {
+                'success': True,
+                'file': str(input_path),
+                'user_id': target_user_id,
+                'conversations': imported
+            }
+            
+        except Exception as e:
+            logger.error(f"JSON import failed: {e}")
+            return {'success': False, 'error': str(e)}
+    
+    def import_from_csv(self, input_file: str, user_id: str) -> Dict[str, Any]:
+        """
+        Import conversations from CSV file
+        
+        Args:
+            input_file: Input CSV file path
+            user_id: User ID for imported conversations
+            
+        Returns:
+            Import statistics
+        """
+        try:
+            input_path = Path(input_file)
+            
+            if not input_path.exists():
+                return {'success': False, 'error': f'File not found: {input_file}'}
+            
+            imported = 0
+            
+            with open(input_path, 'r', encoding='utf-8') as f:
+                reader = csv.DictReader(f)
+                
+                for row in reader:
+                    metadata = {}
+                    if row.get('metadata'):
+                        try:
+                            metadata = json.loads(row['metadata'])
+                        except:
+                            pass
+                    
+                    self.memory.add_conversation(
+                        user_id,
+                        row.get('user_message', ''),
+                        row.get('bot_response', ''),
+                        metadata
+                    )
+                    imported += 1
+            
+            logger.info(f"Imported {imported} conversations from CSV: {input_file}")
+            
+            return {
+                'success': True,
+                'file': str(input_path),
+                'user_id': user_id,
+                'conversations': imported
+            }
+            
+        except Exception as e:
+            logger.error(f"CSV import failed: {e}")
+            return {'success': False, 'error': str(e)}
+    
+    def import_from_sqlite(self, db_file: str, user_id: str) -> Dict[str, Any]:
+        """
+        Import from SQLite database
+        
+        Args:
+            db_file: SQLite database file path
+            user_id: User ID to import data for
+            
+        Returns:
+            Import statistics
+        """
+        try:
+            db_path = Path(db_file)
+            
+            if not db_path.exists():
+                return {'success': False, 'error': f'Database not found: {db_file}'}
+            
+            conn = sqlite3.connect(str(db_path))
+            conn.row_factory = sqlite3.Row
+            cursor = conn.cursor()
+            
+            # Query conversations
+            cursor.execute('''
+                SELECT timestamp, user_message, bot_response, metadata
+                FROM conversations
+                WHERE user_id = ?
+                ORDER BY timestamp
+            ''', (user_id,))
+            
+            imported = 0
+            for row in cursor.fetchall():
+                metadata = {}
+                if row['metadata']:
+                    try:
+                        metadata = json.loads(row['metadata'])
+                    except:
+                        pass
+                
+                self.memory.add_conversation(
+                    user_id,
+                    row['user_message'],
+                    row['bot_response'],
+                    metadata
+                )
+                imported += 1
+            
+            conn.close()
+            
+            logger.info(f"Imported {imported} conversations from SQLite: {db_file}")
+            
+            return {
+                'success': True,
+                'file': str(db_path),
+                'user_id': user_id,
+                'conversations': imported
+            }
+            
+        except Exception as e:
+            logger.error(f"SQLite import failed: {e}")
+            return {'success': False, 'error': str(e)}
+    
+    def import_from_postgresql(self, connection_string: str, user_id: str) -> Dict[str, Any]:
+        """
+        Import from PostgreSQL database
+        
+        Args:
+            connection_string: PostgreSQL connection string
+            user_id: User ID to import data for
+            
+        Returns:
+            Import statistics
+        """
+        try:
+            import psycopg2
+            from psycopg2.extras import RealDictCursor
+        except ImportError:
+            return {
+                'success': False,
+                'error': 'psycopg2 not installed. Install: pip install psycopg2-binary'
+            }
+        
+        try:
+            conn = psycopg2.connect(connection_string)
+            cursor = conn.cursor(cursor_factory=RealDictCursor)
+            
+            # Query conversations
+            cursor.execute('''
+                SELECT timestamp, user_message, bot_response, metadata
+                FROM conversations
+                WHERE user_id = %s
+                ORDER BY timestamp
+            ''', (user_id,))
+            
+            imported = 0
+            for row in cursor.fetchall():
+                metadata = row['metadata'] if isinstance(row['metadata'], dict) else {}
+                
+                self.memory.add_conversation(
+                    user_id,
+                    row['user_message'],
+                    row['bot_response'],
+                    metadata
+                )
+                imported += 1
+            
+            cursor.close()
+            conn.close()
+            
+            logger.info(f"Imported {imported} conversations from PostgreSQL")
+            
+            return {
+                'success': True,
+                'database': 'postgresql',
+                'user_id': user_id,
+                'conversations': imported
+            }
+            
+        except Exception as e:
+            logger.error(f"PostgreSQL import failed: {e}")
+            return {'success': False, 'error': str(e)}
+    
+    def import_from_mongodb(self, connection_string: str, user_id: str,
+                           database: str = 'mem_llm', collection: str = 'conversations') -> Dict[str, Any]:
+        """
+        Import from MongoDB database
+        
+        Args:
+            connection_string: MongoDB connection string
+            user_id: User ID to import data for
+            database: Database name
+            collection: Collection name
+            
+        Returns:
+            Import statistics
+        """
+        try:
+            from pymongo import MongoClient
+        except ImportError:
+            return {
+                'success': False,
+                'error': 'pymongo not installed. Install: pip install pymongo'
+            }
+        
+        try:
+            client = MongoClient(connection_string)
+            db = client[database]
+            coll = db[collection]
+            
+            # Query conversations
+            documents = coll.find({'user_id': user_id}).sort('timestamp', 1)
+            
+            imported = 0
+            for doc in documents:
+                self.memory.add_conversation(
+                    user_id,
+                    doc.get('user_message', ''),
+                    doc.get('bot_response', ''),
+                    doc.get('metadata', {})
+                )
+                imported += 1
+            
+            client.close()
+            
+            logger.info(f"Imported {imported} conversations from MongoDB")
+            
+            return {
+                'success': True,
+                'database': 'mongodb',
+                'user_id': user_id,
+                'conversations': imported
+            }
+            
+        except Exception as e:
+            logger.error(f"MongoDB import failed: {e}")
+            return {'success': False, 'error': str(e)}