h2ogpte 1.6.41rc5__py3-none-any.whl → 1.6.43rc1__py3-none-any.whl

h2ogpte/cli/commands/command_handlers/status.py

@@ -0,0 +1,33 @@
+from rich.table import Table
+
+from ...core.app import get_app_state
+
+
+async def handle_status(args: str) -> bool:
+    """Show session status."""
+    app = get_app_state()
+
+    # Show session status
+    app.session.display_status()
+
+    # Show connection status
+    status_table = Table(title="Connection Status", show_header=False)
+    status_table.add_column("Service", style="cyan")
+    status_table.add_column("Status", style="white")
+
+    rag_status = (
+        "[green]Connected[/green]"
+        if app.rag_manager.connected
+        else "[red]Disconnected[/red]"
+    )
+    agent_status = (
+        "[green]Connected[/green]"
+        if app.agent_manager.connected
+        else "[red]Disconnected[/red]"
+    )
+
+    status_table.add_row("RAG System", rag_status)
+    status_table.add_row("Agent System", agent_status)
+
+    app.console.print(status_table)
+    return True

h2ogpte/cli/commands/dispatcher.py

@@ -0,0 +1,79 @@
+from typing import Dict, Callable, Awaitable
+
+from .command_handlers.help import handle_help
+from .command_handlers.status import handle_status
+from .command_handlers.clear import handle_clear
+from .command_handlers.history import handle_history
+from .command_handlers.exit import handle_exit
+from .command_handlers.config import handle_config
+from .command_handlers.rag import handle_register, handle_upload, handle_analyze
+from .command_handlers.agent import handle_agent
+from .command_handlers.chat import handle_chat
+from .command_handlers.research_agent import handle_research_agent
+from .command_handlers.session import handle_save, handle_load, handle_session
+from .command_handlers.collection import handle_create_collection
+from .command_handlers.disconnect import handle_disconnect
+
+from ..core.app import get_app_state
+
+
+def _sanitize_command_for_history(command: str, cmd: str) -> str:
+    if cmd == "/register":
+        return "/register [hidden credentials]"
+    elif cmd in ["/config"] and "api" in command.lower():
+        return f"{cmd} [hidden sensitive data]"
+    else:
+        return command
+
+
+COMMAND_HANDLERS: Dict[str, Callable[[str], Awaitable[bool]]] = {
+    "/help": handle_help,
+    "/status": handle_status,
+    "/clear": handle_clear,
+    "/history": handle_history,
+    "/exit": handle_exit,
+    "/quit": handle_exit,
+    "/config": handle_config,
+    "/register": handle_register,
+    "/upload": handle_upload,
+    "/analyze": handle_analyze,
+    "/agent": handle_agent,
+    "/research": handle_research_agent,
+    "/save": handle_save,
+    "/load": handle_load,
+    "/session": handle_session,
+    "/collection": handle_create_collection,
+    "/disconnect": handle_disconnect,
+}
+
+
+async def dispatch_command(command: str) -> bool:
+    if not command:
+        return True
+
+    app = get_app_state()
+    parts = command.split(maxsplit=1)
+    cmd = parts[0].lower()
+    args = parts[1] if len(parts) > 1 else ""
+
+    history_command = _sanitize_command_for_history(command, cmd)
+    app.session.add_to_history(history_command, None)
+
+    if cmd in COMMAND_HANDLERS:
+        try:
+            result = await COMMAND_HANDLERS[cmd](args)
+            return result if result is not None else True
+        except Exception as e:
+            app.ui.show_error(f"Command failed: {e}")
+            app.console.print(f"[dim]Debug: {type(e).__name__}: {e}[/dim]")
+            return True
+    else:
+        if command.startswith("/"):
+            app.ui.show_error(f"Unknown command: {cmd}")
+            app.ui.show_info("Type /help for available commands")
+        else:
+            try:
+                await handle_chat(command)
+            except Exception as e:
+                app.ui.show_error(f"Chat failed: {e}")
+        return True

h2ogpte/cli/core/app.py

@@ -0,0 +1,105 @@
+from typing import Optional
+from rich.console import Console
+
+from .config import settings
+from .session import Session
+from ..integrations.rag import RAGManager
+from ..integrations.agent import AgentManager
+from ..utils.file_manager import FileManager, FileUploader, DirectoryAnalyzer
+from ..ui.prompts import UIManager
+
+
+class AppState:
+    """Global application state container."""
+
+    def __init__(self):
+        self.console = Console()
+        self.session = Session()
+        self.ui = UIManager()
+        self.settings = settings
+        self.rag_manager = RAGManager()
+        self.agent_manager = AgentManager()
+        self.file_manager = FileManager()
+        self.file_uploader = FileUploader(self.file_manager)
+        self.dir_analyzer = DirectoryAnalyzer(self.file_manager)
+        self._cleanup_done = False
+
+        # Register interrupt handler
+        self.session.register_interrupt_handler(self._handle_interrupt)
+
+    async def try_auto_reconnect(self):
+        """Try to automatically reconnect using saved credentials."""
+        try:
+            if await self.rag_manager.auto_reconnect(self.settings):
+                await self.update_status_bar()
+                self.ui.show_success(
+                    f"Reconnected as: {await self.rag_manager.get_username()}"
+                )
+                return True
+        except Exception as e:
+            # Silent failure - just continue without connection
+            pass
+        return False
+
+    def _handle_interrupt(self):
+        """Handle interrupt signal."""
+        if self.agent_manager.session:
+            self.agent_manager.session.interrupt()
+
+    async def update_status_bar(self):
+        """Update the status bar with current connection info."""
+        # Check if we have connection info
+        connected = self.rag_manager.connected or self.agent_manager.connected
+
+        # Get actual username from RAG manager
+        username = None
+        collection = None
+        chat_session = None
+        if self.rag_manager.connected:
+            username = await self.rag_manager.get_username()
+            collection = await self.rag_manager.get_collection_name()
+            chat_session = await self.rag_manager.get_chat_session_name()
+
+        # Use RAG chat session if available, otherwise use manual session
+        session = chat_session or getattr(self, "_current_session", None)
+
+        self.ui.update_status(
+            connected=connected,
+            username=username,
+            collection=collection,
+            session=session,
+        )
+
+    async def set_session(self, session_name: str):
+        """Set the current session name."""
+        self._current_session = session_name
+        await self.update_status_bar()
+
+    async def cleanup(self):
+        """Clean up all resources."""
+        if self._cleanup_done:
+            return
+
+        await self.rag_manager.close()
+        await self.agent_manager.close()
+        self.session.save_session()
+        self._cleanup_done = True
+
+
+# Global app instance - initialized once
+app_state: Optional[AppState] = None
+
+
+def get_app_state() -> AppState:
+    """Get the global app state instance."""
+    global app_state
+    if app_state is None:
+        app_state = AppState()
+    return app_state
+
+
+def initialize_app() -> AppState:
+    """Initialize the global app state."""
+    global app_state
+    app_state = AppState()
+    return app_state

h2ogpte/cli/core/config.py

@@ -0,0 +1,199 @@
+from pathlib import Path
+from typing import Optional
+
+import toml
+from pydantic import BaseModel, Field
+from pydantic_settings import BaseSettings
+from pydantic_settings import SettingsConfigDict
+from rich.console import Console
+
+from .encryption import SecureStorage
+
+console = Console()
+
+
+class RAGConfig(BaseModel):
+    """RAG system configuration."""
+
+    endpoint: str = Field(default="", description="RAG system endpoint URL")
+    api_key: str = Field(default="", description="API key for RAG system")
+    collection_name: str = Field(
+        default="default", description="Collection name in RAG"
+    )
+    chunk_size: int = Field(
+        default=1000, description="Chunk size for document processing"
+    )
+    chunk_overlap: int = Field(default=200, description="Overlap between chunks")
+
+
+class AgentConfig(BaseModel):
+    """Agent system configuration."""
+
+    endpoint: str = Field(default="", description="Agent system endpoint URL")
+    api_key: str = Field(default="", description="API key for agent system")
+    model: str = Field(default="gpt-4", description="Model to use for agent")
+    temperature: float = Field(default=0.7, description="Temperature for generation")
+    max_tokens: int = Field(default=2000, description="Maximum tokens for response")
+    timeout: int = Field(default=300, description="Timeout in seconds")
+
+
+class UIConfig(BaseModel):
+    """UI configuration."""
+
+    theme: str = Field(default="monokai", description="Syntax highlighting theme")
+    show_progress: bool = Field(default=True, description="Show progress bars")
+    auto_complete: bool = Field(default=True, description="Enable autocomplete")
+    history_size: int = Field(default=1000, description="Command history size")
+    animation: bool = Field(default=True, description="Enable animations")
+
+
+class Settings(BaseSettings):
+    """Main application settings."""
+
+    app_name: str = "h2oGPTe-CLI"
+    debug: bool = False
+
+    config_dir: Path = Field(default_factory=lambda: Path.home() / ".h2ogpte-cli")
+    data_dir: Path = Field(
+        default_factory=lambda: Path.home() / ".h2ogpte-cli" / "data"
+    )
+    cache_dir: Path = Field(
+        default_factory=lambda: Path.home() / ".h2ogpte-cli" / "cache"
+    )
+    logs_dir: Path = Field(
+        default_factory=lambda: Path.home() / ".h2ogpte-cli" / "logs"
+    )
+
+    rag: RAGConfig = Field(default_factory=RAGConfig)
+    agent: AgentConfig = Field(default_factory=AgentConfig)
+    ui: UIConfig = Field(default_factory=UIConfig)
+
+    _secure_storage: Optional[SecureStorage] = None
+
+    model_config = SettingsConfigDict(
+        env_file=(
+            ".env",
+            ".env.production",
+            ".env.test",
+            ".env.development",
+            ".env.local",
+        ),
+        env_prefix="H2OGPTE_CLI_",
+        env_file_encoding="utf-8",
+        env_nested_delimiter="__",
+        extra="ignore",
+    )
+
+    def _get_secure_storage(self) -> SecureStorage:
+        """Get or create secure storage instance."""
+        if self._secure_storage is None:
+            self._secure_storage = SecureStorage(self.config_dir)
+        return self._secure_storage
+
+    def get_rag_api_key(self) -> str:
+        """Get decrypted RAG API key."""
+        storage = self._get_secure_storage()
+        if storage.is_encrypted(self.rag.api_key):
+            return storage.decrypt(self.rag.api_key)
+        return self.rag.api_key
+
+    def get_agent_api_key(self) -> str:
+        """Get decrypted Agent API key."""
+        storage = self._get_secure_storage()
+        if storage.is_encrypted(self.agent.api_key):
+            return storage.decrypt(self.agent.api_key)
+        return self.agent.api_key
+
+    def set_rag_api_key(self, api_key: str):
+        """Set encrypted RAG API key."""
+        storage = self._get_secure_storage()
+        self.rag.api_key = storage.encrypt(api_key)
+
+    def set_agent_api_key(self, api_key: str):
+        """Set encrypted Agent API key."""
+        storage = self._get_secure_storage()
+        self.agent.api_key = storage.encrypt(api_key)
+
+    def save(self, path: Optional[Path] = None):
+        """Save configuration to file."""
+        if path is None:
+            path = self.config_dir / "config.toml"
+
+        path.parent.mkdir(parents=True, exist_ok=True)
+
+        config_dict = {
+            "app": {
+                "name": self.app_name,
+                "debug": self.debug,
+            },
+            "directories": {
+                "config": str(self.config_dir),
+                "data": str(self.data_dir),
+                "cache": str(self.cache_dir),
+                "logs": str(self.logs_dir),
+            },
+            "rag": self.rag.model_dump(),
+            "agent": self.agent.model_dump(),
+            "ui": self.ui.model_dump(),
+        }
+
+        with open(path, "w") as f:
+            toml.dump(config_dict, f)
+
+        console.print(f"[green]✓[/green] Configuration saved to {path}")
+
+    @classmethod
+    def load(cls, path: Optional[Path] = None) -> "Settings":
+        """Load configuration from file."""
+        if path is None:
+            path = Path.home() / ".h2ogpte-cli" / "config.toml"
+
+        if not path.exists():
+            return cls()
+
+        try:
+            with open(path, "r") as f:
+                config_dict = toml.load(f)
+
+            # Flatten the configuration
+            flat_config = {}
+
+            if "app" in config_dict:
+                app_config = config_dict["app"]
+                # Map config fields to Settings fields
+                if "name" in app_config:
+                    flat_config["app_name"] = app_config["name"]
+                if "version" in app_config:
+                    flat_config["version"] = app_config["version"]
+                if "debug" in app_config:
+                    flat_config["debug"] = app_config["debug"]
+
+            if "directories" in config_dict:
+                dirs = config_dict["directories"]
+                for key in ["config_dir", "data_dir", "cache_dir", "logs_dir"]:
+                    if key.replace("_dir", "") in dirs:
+                        flat_config[key] = Path(dirs[key.replace("_dir", "")])
+
+            if "rag" in config_dict:
+                flat_config["rag"] = RAGConfig(**config_dict["rag"])
+
+            if "agent" in config_dict:
+                flat_config["agent"] = AgentConfig(**config_dict["agent"])
+
+            if "ui" in config_dict:
+                flat_config["ui"] = UIConfig(**config_dict["ui"])
+
+            return cls(**flat_config)
+
+        except Exception as e:
+            console.print(f"[yellow]⚠[/yellow] Error loading config: {e}")
+            return cls()
+
+    def ensure_directories(self):
+        """Ensure all required directories exist."""
+        for dir_path in [self.config_dir, self.data_dir, self.cache_dir, self.logs_dir]:
+            dir_path.mkdir(parents=True, exist_ok=True)
+
+
+# Global settings instance
+settings = Settings.load()

h2ogpte/cli/core/encryption.py

@@ -0,0 +1,104 @@
+import os
+import base64
+import hashlib
+from pathlib import Path
+from typing import Optional
+from cryptography.fernet import Fernet
+from cryptography.hazmat.primitives import hashes
+from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
+
+
+class SecureStorage:
+    """Handles encryption/decryption of sensitive data."""
+
+    def __init__(self, config_dir: Path):
+        self.config_dir = config_dir
+        self.key_file = config_dir / ".keyfile"
+        self._cipher_suite = None
+
+    def _get_machine_id(self) -> bytes:
+        """Generate a machine-specific identifier."""
+        # Use machine-specific data for key generation
+        import platform
+        import socket
+
+        machine_data = (
+            platform.node()
+            + platform.machine()
+            + platform.processor()
+            + str(os.getuid() if hasattr(os, "getuid") else "windows")
+        ).encode()
+
+        return hashlib.sha256(machine_data).digest()
+
+    def _get_or_create_key(self) -> bytes:
+        """Get or create encryption key."""
+        if self.key_file.exists():
+            with open(self.key_file, "rb") as f:
+                return f.read()
+
+        # Create new key based on machine ID
+        machine_id = self._get_machine_id()
+
+        # Use PBKDF2 to derive a key from machine ID
+        kdf = PBKDF2HMAC(
+            algorithm=hashes.SHA256(),
+            length=32,
+            salt=b"hbot_salt_2024",  # Static salt for consistency
+            iterations=100000,
+        )
+        key = base64.urlsafe_b64encode(kdf.derive(machine_id))
+
+        # Save key to file
+        self.config_dir.mkdir(parents=True, exist_ok=True)
+        with open(self.key_file, "wb") as f:
+            f.write(key)
+
+        # Set restrictive permissions
+        os.chmod(self.key_file, 0o600)
+
+        return key
+
+    def _get_cipher_suite(self) -> Fernet:
+        """Get cipher suite for encryption/decryption."""
+        if self._cipher_suite is None:
+            key = self._get_or_create_key()
+            self._cipher_suite = Fernet(key)
+        return self._cipher_suite
+
+    def encrypt(self, data: str) -> str:
+        """Encrypt sensitive data."""
+        if not data:
+            return ""
+
+        cipher_suite = self._get_cipher_suite()
+        encrypted_data = cipher_suite.encrypt(data.encode())
+        return base64.urlsafe_b64encode(encrypted_data).decode()
+
+    def decrypt(self, encrypted_data: str) -> str:
+        """Decrypt sensitive data."""
+        if not encrypted_data:
+            return ""
+
+        try:
+            cipher_suite = self._get_cipher_suite()
+            decoded_data = base64.urlsafe_b64decode(encrypted_data.encode())
+            decrypted_data = cipher_suite.decrypt(decoded_data)
+            return decrypted_data.decode()
+        except Exception:
+            # Return empty string if decryption fails (corrupted or wrong key)
+            return ""
+
+    def is_encrypted(self, data: str) -> bool:
+        """Check if data appears to be encrypted."""
+        if not data:
+            return False
+
+        # Simple heuristic: encrypted data should be base64-like and long
+        try:
+            # Check if it's valid base64
+            base64.urlsafe_b64decode(data.encode())
+            # Encrypted data should be long (>50 chars) and not look like a normal API key
+            return len(data) > 50 and not data.startswith("sk-")
+        except:
+            return False

h2ogpte/cli/core/session.py

@@ -0,0 +1,171 @@
+import asyncio
+import signal
+import sys
+from typing import Optional, Dict, Any, List, Callable
+from datetime import datetime
+from pathlib import Path
+import json
+from rich.console import Console
+from rich.live import Live
+from rich.table import Table
+from rich.panel import Panel
+from rich.layout import Layout
+from rich.text import Text
+from rich.progress import Progress, SpinnerColumn, TextColumn, BarColumn
+
+console = Console()
+
+
+class Session:
+    """Manages a CLI session with state, history, and interruption handling."""
+
+    def __init__(self):
+        self.start_time = datetime.now()
+        self.history: List[Dict[str, Any]] = []
+        self.context: Dict[str, Any] = {}
+        self.is_running = False
+        self.interrupted = False
+        self.current_task: Optional[asyncio.Task] = None
+        self.interrupt_handlers: List[Callable] = []
+        self.progress = Progress(
+            SpinnerColumn(),
+            TextColumn("[progress.description]{task.description}"),
+            BarColumn(),
+            TextColumn("[progress.percentage]{task.percentage:>3.0f}%"),
+        )
+
+        # Setup interrupt handling
+        signal.signal(signal.SIGINT, self._handle_interrupt)
+
+    def _handle_interrupt(self, signum, frame):
+        """Handle interrupt signal (Ctrl+C)."""
+        self.interrupted = True
+        console.print("\n[yellow]⚠[/yellow] Interrupt received. Processing...")
+
+        # Call registered interrupt handlers
+        for handler in self.interrupt_handlers:
+            try:
+                handler()
+            except Exception as e:
+                console.print(f"[red]Error in interrupt handler: {e}[/red]")
+
+        # Cancel current task if running
+        if self.current_task and not self.current_task.done():
+            self.current_task.cancel()
+
+    def register_interrupt_handler(self, handler: Callable):
+        """Register a handler to be called on interrupt."""
+        self.interrupt_handlers.append(handler)
+
+    def add_to_history(self, command: str, result: Any, success: bool = True):
+        """Add a command and its result to history."""
+        entry = {
+            "timestamp": datetime.now().isoformat(),
+            "command": command,
+            "result": str(result) if result else None,
+            "success": success,
+        }
+        self.history.append(entry)
+
+    def get_context(self, key: str, default: Any = None) -> Any:
+        """Get a value from session context."""
+        return self.context.get(key, default)
+
+    def set_context(self, key: str, value: Any):
+        """Set a value in session context."""
+        self.context[key] = value
+
+    def clear_context(self):
+        """Clear session context."""
+        self.context.clear()
+
+    async def run_with_interrupt(self, coro, description: str = "Processing..."):
+        """Run a coroutine with interrupt handling."""
+        self.is_running = True
+        self.interrupted = False
+
+        task_id = self.progress.add_task(description, total=None)
+
+        try:
+            with self.progress:
+                self.current_task = asyncio.create_task(coro)
+                result = await self.current_task
+                self.progress.update(task_id, completed=100)
+                return result
+
+        except asyncio.CancelledError:
+            console.print("[yellow]✗[/yellow] Task cancelled")
+            return None
+
+        except Exception as e:
+            console.print(f"[red]✗[/red] Error: {e}")
+            return None
+
+        finally:
+            self.is_running = False
+            self.current_task = None
+            self.progress.remove_task(task_id)
+
+    def save_session(self, path: Optional[Path] = None):
+        """Save session to file."""
+        if path is None:
+            from .config import settings
+
+            path = (
+                settings.data_dir
+                / f"session_{self.start_time.strftime('%Y%m%d_%H%M%S')}.json"
+            )
+
+        path.parent.mkdir(parents=True, exist_ok=True)
+
+        session_data = {
+            "start_time": self.start_time.isoformat(),
+            "history": self.history,
+            "context": self.context,
+        }
+
+        with open(path, "w") as f:
+            json.dump(session_data, f, indent=2, default=str)
+
+        # console.print(f"[green]✓[/green] Session saved to {path}")
+
+    @classmethod
+    def load_session(cls, path: Path) -> "Session":
+        """Load session from file."""
+        with open(path, "r") as f:
+            session_data = json.load(f)
+
+        session = cls()
+        session.start_time = datetime.fromisoformat(session_data["start_time"])
+        session.history = session_data["history"]
+        session.context = session_data["context"]
+
+        return session
+
+    def display_status(self):
+        """Display current session status."""
+        layout = Layout()
+
+        # Session info
+        info_table = Table(title="Session Information", show_header=False)
+        info_table.add_row("Started", self.start_time.strftime("%Y-%m-%d %H:%M:%S"))
+        info_table.add_row(
+            "Duration", str(datetime.now() - self.start_time).split(".")[0]
+        )
+        info_table.add_row("Commands", str(len(self.history)))
+        info_table.add_row("Context Items", str(len(self.context)))
+
+        # Recent history
+        history_table = Table(title="Recent History")
+        history_table.add_column("Time", style="dim")
+        history_table.add_column("Command")
+        history_table.add_column("Status")
+
+        for entry in self.history[-5:]:
+            time_str = datetime.fromisoformat(entry["timestamp"]).strftime("%H:%M:%S")
+            status = "[green]✓[/green]" if entry["success"] else "[red]✗[/red]"
+            history_table.add_row(time_str, entry["command"][:50], status)
+
+        layout.split_column(Layout(Panel(info_table)), Layout(Panel(history_table)))
+
+        console.print(layout)