dtSpark 1.0.4__py3-none-any.whl

dtSpark/database/tool_permissions.py

@@ -0,0 +1,186 @@
+"""
+Tool permissions operations module.
+
+This module handles:
+- Checking tool permissions for conversations
+- Setting tool permission states
+- Managing first-time tool usage prompts
+"""
+
+import sqlite3
+import logging
+from datetime import datetime
+from typing import Optional, Dict, List
+
+
+# Permission states
+PERMISSION_ALLOWED = 'allowed'  # Run this time and all future times
+PERMISSION_DENIED = 'denied'    # Never run this tool
+
+
+def check_tool_permission(conn: sqlite3.Connection, conversation_id: int,
+                          tool_name: str, user_guid: Optional[str] = None) -> Optional[str]:
+    """
+    Check the permission state for a tool in a conversation.
+    Returns None if no record exists (first-time usage, should prompt).
+
+    Args:
+        conn: Database connection
+        conversation_id: ID of the conversation
+        tool_name: Name of the tool
+        user_guid: User GUID for multi-user support
+
+    Returns:
+        Permission state ('allowed', 'denied') or None if no record exists
+    """
+    cursor = conn.cursor()
+    cursor.execute('''
+        SELECT permission_state
+        FROM conversation_tool_permissions
+        WHERE conversation_id = ? AND tool_name = ? AND user_guid = ?
+    ''', (conversation_id, tool_name, user_guid))
+
+    row = cursor.fetchone()
+    if row is None:
+        # No record exists - first time usage, should prompt user
+        return None
+    return row['permission_state']
+
+
+def set_tool_permission(conn: sqlite3.Connection, conversation_id: int,
+                        tool_name: str, permission_state: str,
+                        user_guid: Optional[str] = None) -> bool:
+    """
+    Set the permission state for a tool in a conversation.
+
+    Args:
+        conn: Database connection
+        conversation_id: ID of the conversation
+        tool_name: Name of the tool
+        permission_state: Permission state ('allowed' or 'denied')
+        user_guid: User GUID for multi-user support
+
+    Returns:
+        True if successful, False otherwise
+    """
+    if permission_state not in [PERMISSION_ALLOWED, PERMISSION_DENIED]:
+        logging.error(f"Invalid permission state: {permission_state}")
+        return False
+
+    try:
+        cursor = conn.cursor()
+        now = datetime.now().isoformat()
+
+        cursor.execute('''
+            INSERT INTO conversation_tool_permissions
+                (conversation_id, tool_name, permission_state, granted_at, updated_at, user_guid)
+            VALUES (?, ?, ?, ?, ?, ?)
+            ON CONFLICT(conversation_id, tool_name) DO UPDATE SET
+                permission_state = excluded.permission_state,
+                updated_at = excluded.updated_at
+        ''', (conversation_id, tool_name, permission_state, now, now, user_guid))
+
+        conn.commit()
+        logging.info(f"Tool permission '{tool_name}' set to '{permission_state}' for conversation {conversation_id}")
+        return True
+
+    except Exception as e:
+        logging.error(f"Failed to set tool permission: {e}")
+        conn.rollback()
+        return False
+
+
+def get_all_tool_permissions(conn: sqlite3.Connection, conversation_id: int,
+                             user_guid: Optional[str] = None) -> List[Dict]:
+    """
+    Get all tool permissions for a conversation.
+
+    Args:
+        conn: Database connection
+        conversation_id: ID of the conversation
+        user_guid: User GUID for multi-user support
+
+    Returns:
+        List of dicts with 'tool_name', 'permission_state', 'granted_at', 'updated_at'
+    """
+    cursor = conn.cursor()
+    cursor.execute('''
+        SELECT tool_name, permission_state, granted_at, updated_at
+        FROM conversation_tool_permissions
+        WHERE conversation_id = ? AND user_guid = ?
+        ORDER BY updated_at DESC
+    ''', (conversation_id, user_guid))
+
+    permissions = []
+    for row in cursor.fetchall():
+        permissions.append({
+            'tool_name': row['tool_name'],
+            'permission_state': row['permission_state'],
+            'granted_at': row['granted_at'],
+            'updated_at': row['updated_at']
+        })
+
+    return permissions
+
+
+def delete_tool_permission(conn: sqlite3.Connection, conversation_id: int,
+                          tool_name: str, user_guid: Optional[str] = None) -> bool:
+    """
+    Delete a tool permission record (reset to first-time usage behavior).
+
+    Args:
+        conn: Database connection
+        conversation_id: ID of the conversation
+        tool_name: Name of the tool
+        user_guid: User GUID for multi-user support
+
+    Returns:
+        True if successful, False otherwise
+    """
+    try:
+        cursor = conn.cursor()
+        cursor.execute('''
+            DELETE FROM conversation_tool_permissions
+            WHERE conversation_id = ? AND tool_name = ? AND user_guid = ?
+        ''', (conversation_id, tool_name, user_guid))
+
+        conn.commit()
+        deleted_count = cursor.rowcount
+        if deleted_count > 0:
+            logging.info(f"Deleted tool permission for '{tool_name}' in conversation {conversation_id}")
+        return True
+
+    except Exception as e:
+        logging.error(f"Failed to delete tool permission: {e}")
+        conn.rollback()
+        return False
+
+
+def is_tool_allowed(conn: sqlite3.Connection, conversation_id: int,
+                   tool_name: str, user_guid: Optional[str] = None) -> Optional[bool]:
+    """
+    Check if a tool is allowed to run.
+    Returns None if permission should be requested from user (first-time usage).
+    Returns True if allowed, False if denied.
+
+    Args:
+        conn: Database connection
+        conversation_id: ID of the conversation
+        tool_name: Name of the tool
+        user_guid: User GUID for multi-user support
+
+    Returns:
+        True if allowed, False if denied, None if should prompt user
+    """
+    permission_state = check_tool_permission(conn, conversation_id, tool_name, user_guid)
+
+    if permission_state is None:
+        # No record exists - should prompt user
+        return None
+    elif permission_state == PERMISSION_ALLOWED:
+        return True
+    elif permission_state == PERMISSION_DENIED:
+        return False
+    else:
+        logging.warning(f"Unknown permission state: {permission_state}")
+        return None

dtSpark/database/usage.py

@@ -0,0 +1,167 @@
+"""
+Usage tracking module.
+
+This module handles:
+- Recording token and cost usage
+- Retrieving usage within time windows
+- Usage summary and reporting
+- Cleanup of old usage data
+"""
+
+import sqlite3
+import logging
+from datetime import datetime
+from typing import List, Dict, Optional, Tuple
+
+
+def record_usage(conn: sqlite3.Connection, conversation_id: int, model_id: str,
+                region: str, input_tokens: int, output_tokens: int, cost: float,
+                timestamp: datetime, user_guid: str = None):
+    """
+    Record usage for token management and billing.
+
+    Args:
+        conn: Database connection
+        conversation_id: ID of the conversation
+        model_id: Bedrock model ID
+        region: AWS region
+        input_tokens: Number of input tokens
+        output_tokens: Number of output tokens
+        cost: Cost in USD
+        timestamp: Timestamp of usage
+        user_guid: User GUID for multi-user support
+    """
+    cursor = conn.cursor()
+    cursor.execute('''
+        INSERT INTO usage_tracking
+        (conversation_id, model_id, region, input_tokens, output_tokens, cost, timestamp, user_guid)
+        VALUES (?, ?, ?, ?, ?, ?, ?, ?)
+    ''', (conversation_id, model_id, region, input_tokens, output_tokens, cost, timestamp, user_guid))
+    conn.commit()
+
+
+def get_usage_in_window(conn: sqlite3.Connection, window_start: datetime,
+                        user_guid: str = None) -> float:
+    """
+    Get total cost for usage since window_start.
+
+    Args:
+        conn: Database connection
+        window_start: Start of rolling window
+        user_guid: User GUID for multi-user support
+
+    Returns:
+        Total cost in USD
+    """
+    cursor = conn.cursor()
+    cursor.execute('''
+        SELECT SUM(cost) as total_cost
+        FROM usage_tracking
+        WHERE timestamp >= ? AND user_guid = ?
+    ''', (window_start, user_guid))
+
+    result = cursor.fetchone()
+    return result['total_cost'] if result['total_cost'] is not None else 0.0
+
+
+def get_oldest_usage_in_window(conn: sqlite3.Connection,
+                               window_start: datetime, user_guid: str = None) -> Optional[datetime]:
+    """
+    Get timestamp of oldest usage in the rolling window.
+
+    Args:
+        conn: Database connection
+        window_start: Start of rolling window
+        user_guid: User GUID for multi-user support
+
+    Returns:
+        Timestamp of oldest usage or None
+    """
+    cursor = conn.cursor()
+    cursor.execute('''
+        SELECT MIN(timestamp) as oldest_timestamp
+        FROM usage_tracking
+        WHERE timestamp >= ? AND user_guid = ?
+    ''', (window_start, user_guid))
+
+    result = cursor.fetchone()
+    if result and result['oldest_timestamp']:
+        return datetime.fromisoformat(result['oldest_timestamp'])
+    return None
+
+
+def get_token_usage_in_window(conn: sqlite3.Connection,
+                              window_start: datetime, user_guid: str = None) -> Tuple[int, int]:
+    """
+    Get total token usage (input and output separately) since window_start.
+
+    Args:
+        conn: Database connection
+        window_start: Start of rolling window
+        user_guid: User GUID for multi-user support
+
+    Returns:
+        Tuple of (total_input_tokens, total_output_tokens)
+    """
+    cursor = conn.cursor()
+    cursor.execute('''
+        SELECT
+            SUM(input_tokens) as total_input,
+            SUM(output_tokens) as total_output
+        FROM usage_tracking
+        WHERE timestamp >= ? AND user_guid = ?
+    ''', (window_start, user_guid))
+
+    result = cursor.fetchone()
+    total_input = result['total_input'] if result['total_input'] is not None else 0
+    total_output = result['total_output'] if result['total_output'] is not None else 0
+    return int(total_input), int(total_output)
+
+
+def get_usage_summary(conn: sqlite3.Connection, window_start: datetime,
+                      user_guid: str = None) -> List[Dict]:
+    """
+    Get detailed usage summary for the rolling window.
+
+    Args:
+        conn: Database connection
+        window_start: Start of rolling window
+        user_guid: User GUID for multi-user support
+
+    Returns:
+        List of usage records with model, tokens, and costs
+    """
+    cursor = conn.cursor()
+    cursor.execute('''
+        SELECT model_id, region,
+               SUM(input_tokens) as total_input_tokens,
+               SUM(output_tokens) as total_output_tokens,
+               SUM(cost) as total_cost,
+               COUNT(*) as request_count
+        FROM usage_tracking
+        WHERE timestamp >= ? AND user_guid = ?
+        GROUP BY model_id, region
+        ORDER BY total_cost DESC
+    ''', (window_start, user_guid))
+
+    return [dict(row) for row in cursor.fetchall()]
+
+
+def cleanup_old_usage(conn: sqlite3.Connection, cutoff_date: datetime,
+                      user_guid: str = None):
+    """
+    Clean up usage records older than cutoff_date.
+
+    Args:
+        conn: Database connection
+        cutoff_date: Delete records older than this date
+        user_guid: User GUID for multi-user support (for safety filtering)
+    """
+    cursor = conn.cursor()
+    cursor.execute('''
+        DELETE FROM usage_tracking
+        WHERE timestamp < ? AND user_guid = ?
+    ''', (cutoff_date, user_guid))
+    deleted_count = cursor.rowcount
+    conn.commit()
+    logging.info(f"Cleaned up {deleted_count} old usage records for user {user_guid}")

dtSpark/files/__init__.py

@@ -0,0 +1,4 @@
+"""File management module."""
+from .manager import FileManager
+
+__all__ = ['FileManager']

dtSpark/files/manager.py

@@ -0,0 +1,322 @@
+"""
+File Manager module for processing various file types for conversation context.
+
+This module provides functionality for:
+- Extracting text from documents (PDF, DOCX, TXT, MD, CSV)
+- Processing images as base64-encoded data
+- Token counting for file content
+"""
+
+import base64
+import logging
+import mimetypes
+from pathlib import Path
+from typing import Dict, Optional, Tuple
+
+
+class FileManager:
+    """Manages file processing for conversation context."""
+
+    # Supported file extensions
+    SUPPORTED_TEXT_FILES = {'.txt', '.md', '.csv', '.json', '.yaml', '.yml', '.xml', '.log'}
+    SUPPORTED_CODE_FILES = {
+        '.py', '.js', '.ts', '.jsx', '.tsx',  # Python, JavaScript, TypeScript
+        '.java', '.kt', '.scala',              # JVM languages
+        '.c', '.cpp', '.cc', '.h', '.hpp',     # C/C++
+        '.cs',                                  # C#
+        '.go', '.rs', '.rb',                   # Go, Rust, Ruby
+        '.php', '.swift', '.m',                # PHP, Swift, Objective-C
+        '.sh', '.bash', '.zsh', '.ps1',        # Shell scripts
+        '.sql', '.r',                          # SQL, R
+        '.html', '.css', '.scss', '.sass',     # Web
+        '.vue', '.svelte',                     # Web frameworks
+        '.toml', '.ini', '.cfg', '.conf',      # Config files
+        '.dockerfile', '.tf', '.hcl',          # Infrastructure
+    }
+    SUPPORTED_DOCUMENT_FILES = {'.pdf', '.docx', '.doc'}
+    SUPPORTED_IMAGE_FILES = {'.png', '.jpg', '.jpeg', '.gif', '.bmp', '.webp'}
+
+    def __init__(self, bedrock_service=None):
+        """
+        Initialise the file manager.
+
+        Args:
+            bedrock_service: Optional BedrockService instance for token counting
+        """
+        self.bedrock_service = bedrock_service
+
+    @classmethod
+    def is_supported(cls, file_path: str) -> bool:
+        """
+        Check if a file type is supported.
+
+        Args:
+            file_path: Path to the file
+
+        Returns:
+            True if the file type is supported
+        """
+        ext = Path(file_path).suffix.lower()
+        return ext in (cls.SUPPORTED_TEXT_FILES | cls.SUPPORTED_CODE_FILES | cls.SUPPORTED_DOCUMENT_FILES | cls.SUPPORTED_IMAGE_FILES)
+
+    def process_file(self, file_path: str) -> Dict:
+        """
+        Process a file and extract its content.
+
+        Args:
+            file_path: Path to the file
+
+        Returns:
+            Dictionary containing:
+            - filename: Original filename
+            - file_type: File extension
+            - file_size: Size in bytes
+            - content_text: Extracted text content (for documents)
+            - content_base64: Base64 encoded content (for images)
+            - mime_type: MIME type (for images)
+            - token_count: Token count of extracted content
+            - error: Error message if processing failed
+        """
+        path = Path(file_path)
+
+        if not path.exists():
+            return {'error': f"File not found: {file_path}"}
+
+        if not self.is_supported(file_path):
+            return {'error': f"Unsupported file type: {path.suffix}"}
+
+        result = {
+            'filename': path.name,
+            'file_type': path.suffix.lower(),
+            'file_size': path.stat().st_size,
+            'content_text': None,
+            'content_base64': None,
+            'mime_type': None,
+            'token_count': 0
+        }
+
+        ext = path.suffix.lower()
+
+        try:
+            # Process text and code files
+            if ext in (self.SUPPORTED_TEXT_FILES | self.SUPPORTED_CODE_FILES):
+                result['content_text'] = self._extract_text_file(file_path)
+                if result['content_text']:
+                    result['token_count'] = self._count_tokens(result['content_text'])
+
+            # Process documents
+            elif ext in self.SUPPORTED_DOCUMENT_FILES:
+                if ext == '.pdf':
+                    result['content_text'] = self._extract_pdf(file_path)
+                elif ext in {'.docx', '.doc'}:
+                    result['content_text'] = self._extract_docx(file_path)
+
+                if result['content_text']:
+                    result['token_count'] = self._count_tokens(result['content_text'])
+
+            # Process images
+            elif ext in self.SUPPORTED_IMAGE_FILES:
+                result['content_base64'], result['mime_type'] = self._encode_image(file_path)
+                # Images don't have text token count, but we could estimate based on size
+                # For now, we'll use a simple heuristic: ~1 token per 750 bytes
+                result['token_count'] = max(1, result['file_size'] // 750)
+
+            logging.info(f"Processed file: {path.name} ({result['file_type']}, {result['token_count']} tokens)")
+
+        except Exception as e:
+            logging.error(f"Error processing file {file_path}: {e}")
+            result['error'] = str(e)
+
+        return result
+
+    def _extract_text_file(self, file_path: str) -> str:
+        """
+        Extract text from a plain text file.
+
+        Args:
+            file_path: Path to the file
+
+        Returns:
+            Extracted text content
+        """
+        try:
+            # Try UTF-8 first
+            with open(file_path, 'r', encoding='utf-8') as f:
+                return f.read()
+        except UnicodeDecodeError:
+            # Fall back to latin-1 if UTF-8 fails
+            try:
+                with open(file_path, 'r', encoding='latin-1') as f:
+                    return f.read()
+            except Exception as e:
+                logging.error(f"Failed to read text file {file_path}: {e}")
+                raise
+
+    def _extract_pdf(self, file_path: str) -> str:
+        """
+        Extract text from a PDF file.
+
+        Args:
+            file_path: Path to the PDF file
+
+        Returns:
+            Extracted text content
+        """
+        try:
+            import pypdf
+
+            text_parts = []
+            with open(file_path, 'rb') as f:
+                pdf_reader = pypdf.PdfReader(f)
+                for page_num, page in enumerate(pdf_reader.pages):
+                    text = page.extract_text()
+                    if text.strip():
+                        text_parts.append(f"--- Page {page_num + 1} ---\n{text}")
+
+            return '\n\n'.join(text_parts)
+
+        except ImportError:
+            raise ImportError("pypdf library is required for PDF processing. Install with: pip install pypdf")
+        except Exception as e:
+            logging.error(f"Failed to extract PDF {file_path}: {e}")
+            raise
+
+    def _extract_docx(self, file_path: str) -> str:
+        """
+        Extract text from a DOCX file.
+
+        Args:
+            file_path: Path to the DOCX file
+
+        Returns:
+            Extracted text content
+        """
+        try:
+            import docx
+
+            doc = docx.Document(file_path)
+            text_parts = []
+
+            # Extract paragraphs
+            for para in doc.paragraphs:
+                if para.text.strip():
+                    text_parts.append(para.text)
+
+            # Extract tables
+            for table in doc.tables:
+                for row in table.rows:
+                    row_text = ' | '.join(cell.text.strip() for cell in row.cells)
+                    if row_text.strip():
+                        text_parts.append(row_text)
+
+            return '\n\n'.join(text_parts)
+
+        except ImportError:
+            raise ImportError("python-docx library is required for DOCX processing. Install with: pip install python-docx")
+        except Exception as e:
+            logging.error(f"Failed to extract DOCX {file_path}: {e}")
+            raise
+
+    def _encode_image(self, file_path: str) -> Tuple[str, str]:
+        """
+        Encode an image as base64.
+
+        Args:
+            file_path: Path to the image file
+
+        Returns:
+            Tuple of (base64_string, mime_type)
+        """
+        try:
+            # Read and encode image
+            with open(file_path, 'rb') as f:
+                image_data = f.read()
+
+            base64_string = base64.b64encode(image_data).decode('utf-8')
+
+            # Determine MIME type
+            mime_type, _ = mimetypes.guess_type(file_path)
+            if not mime_type:
+                # Default to common image types
+                ext = Path(file_path).suffix.lower()
+                mime_map = {
+                    '.jpg': 'image/jpeg',
+                    '.jpeg': 'image/jpeg',
+                    '.png': 'image/png',
+                    '.gif': 'image/gif',
+                    '.bmp': 'image/bmp',
+                    '.webp': 'image/webp'
+                }
+                mime_type = mime_map.get(ext, 'application/octet-stream')
+
+            return base64_string, mime_type
+
+        except Exception as e:
+            logging.error(f"Failed to encode image {file_path}: {e}")
+            raise
+
+    def _count_tokens(self, text: str) -> int:
+        """
+        Count tokens in text.
+
+        Args:
+            text: Text to count tokens for
+
+        Returns:
+            Token count
+        """
+        if self.bedrock_service:
+            return self.bedrock_service.count_tokens(text)
+        else:
+            # Rough estimate: ~4 characters per token
+            return max(1, len(text) // 4)
+
+    @classmethod
+    def get_supported_extensions(cls) -> str:
+        """
+        Get a formatted string of supported file extensions.
+
+        Returns:
+            Comma-separated list of supported extensions
+        """
+        all_extensions = sorted(
+            cls.SUPPORTED_TEXT_FILES | cls.SUPPORTED_DOCUMENT_FILES | cls.SUPPORTED_IMAGE_FILES
+        )
+        return ', '.join(all_extensions)
+
+    @classmethod
+    def scan_directory(cls, directory_path: str, recursive: bool = False) -> list:
+        """
+        Scan a directory for supported files.
+
+        Args:
+            directory_path: Path to the directory to scan
+            recursive: If True, scan subdirectories as well
+
+        Returns:
+            List of absolute file paths for supported files
+        """
+        from pathlib import Path
+
+        dir_path = Path(directory_path)
+
+        if not dir_path.exists():
+            raise FileNotFoundError(f"Directory not found: {directory_path}")
+
+        if not dir_path.is_dir():
+            raise NotADirectoryError(f"Not a directory: {directory_path}")
+
+        supported_files = []
+
+        if recursive:
+            # Recursively find all files
+            for file_path in dir_path.rglob('*'):
+                if file_path.is_file() and cls.is_supported(str(file_path)):
+                    supported_files.append(str(file_path.absolute()))
+        else:
+            # Only scan immediate directory
+            for file_path in dir_path.iterdir():
+                if file_path.is_file() and cls.is_supported(str(file_path)):
+                    supported_files.append(str(file_path.absolute()))
+
+        return sorted(supported_files)