130 lines
4.2 KiB
Python
130 lines
4.2 KiB
Python
"""Statistics screen for oAI TUI."""
|
|
|
|
from textual.app import ComposeResult
|
|
from textual.containers import Container, Vertical
|
|
from textual.screen import ModalScreen
|
|
from textual.widgets import Button, Static
|
|
|
|
from oai.core.session import ChatSession
|
|
|
|
|
|
class StatsScreen(ModalScreen[None]):
|
|
"""Modal screen displaying session statistics."""
|
|
|
|
DEFAULT_CSS = """
|
|
StatsScreen {
|
|
align: center middle;
|
|
}
|
|
|
|
StatsScreen > Container {
|
|
width: 70;
|
|
height: auto;
|
|
background: #1e1e1e;
|
|
border: solid #555555;
|
|
}
|
|
|
|
StatsScreen .header {
|
|
dock: top;
|
|
width: 100%;
|
|
height: auto;
|
|
background: #2d2d2d;
|
|
color: #cccccc;
|
|
padding: 0 2;
|
|
}
|
|
|
|
StatsScreen .content {
|
|
width: 100%;
|
|
height: auto;
|
|
background: #1e1e1e;
|
|
padding: 2;
|
|
color: #cccccc;
|
|
}
|
|
|
|
StatsScreen .footer {
|
|
dock: bottom;
|
|
width: 100%;
|
|
height: auto;
|
|
background: #2d2d2d;
|
|
padding: 1 2;
|
|
align: center middle;
|
|
}
|
|
"""
|
|
|
|
def __init__(self, session: ChatSession):
|
|
super().__init__()
|
|
self.session = session
|
|
|
|
def compose(self) -> ComposeResult:
|
|
"""Compose the screen."""
|
|
with Container():
|
|
yield Static("[bold]Session Statistics[/]", classes="header")
|
|
with Vertical(classes="content"):
|
|
yield Static(self._get_stats_text(), markup=True)
|
|
with Vertical(classes="footer"):
|
|
yield Button("Close", id="close", variant="primary")
|
|
|
|
def _get_stats_text(self) -> str:
|
|
"""Generate the statistics text."""
|
|
stats = self.session.stats
|
|
|
|
# Calculate averages
|
|
avg_input = stats.total_input_tokens // stats.message_count if stats.message_count > 0 else 0
|
|
avg_output = stats.total_output_tokens // stats.message_count if stats.message_count > 0 else 0
|
|
avg_cost = stats.total_cost / stats.message_count if stats.message_count > 0 else 0
|
|
|
|
# Get model info
|
|
model_name = "None"
|
|
model_context = "N/A"
|
|
if self.session.selected_model:
|
|
model_name = self.session.selected_model.get("name", "Unknown")
|
|
model_context = str(self.session.selected_model.get("context_length", "N/A"))
|
|
|
|
# MCP status
|
|
mcp_status = "Disabled"
|
|
if self.session.mcp_manager and self.session.mcp_manager.enabled:
|
|
mode = self.session.mcp_manager.mode
|
|
if mode == "files":
|
|
write = " (Write)" if self.session.mcp_manager.write_enabled else ""
|
|
mcp_status = f"Enabled - Files{write}"
|
|
elif mode == "database":
|
|
db_idx = self.session.mcp_manager.selected_db_index
|
|
if db_idx is not None:
|
|
db_name = self.session.mcp_manager.databases[db_idx]["name"]
|
|
mcp_status = f"Enabled - Database ({db_name})"
|
|
|
|
return f"""
|
|
[bold cyan]═══ SESSION INFO ═══[/]
|
|
[bold]Messages:[/] {stats.message_count}
|
|
[bold]Current Model:[/] {model_name}
|
|
[bold]Context Length:[/] {model_context}
|
|
[bold]Memory:[/] {"Enabled" if self.session.memory_enabled else "Disabled"}
|
|
[bold]Online Mode:[/] {"Enabled" if self.session.online_enabled else "Disabled"}
|
|
[bold]MCP:[/] {mcp_status}
|
|
|
|
[bold cyan]═══ TOKEN USAGE ═══[/]
|
|
[bold]Input Tokens:[/] {stats.total_input_tokens:,}
|
|
[bold]Output Tokens:[/] {stats.total_output_tokens:,}
|
|
[bold]Total Tokens:[/] {stats.total_tokens:,}
|
|
|
|
[bold]Avg Input/Msg:[/] {avg_input:,}
|
|
[bold]Avg Output/Msg:[/] {avg_output:,}
|
|
|
|
[bold cyan]═══ COSTS ═══[/]
|
|
[bold]Total Cost:[/] ${stats.total_cost:.6f}
|
|
[bold]Avg Cost/Msg:[/] ${avg_cost:.6f}
|
|
|
|
[bold cyan]═══ HISTORY ═══[/]
|
|
[bold]History Size:[/] {len(self.session.history)} entries
|
|
[bold]Current Index:[/] {self.session.current_index + 1 if self.session.history else 0}
|
|
[bold]Memory Start:[/] {self.session.memory_start_index + 1}
|
|
"""
|
|
|
|
def on_button_pressed(self, event: Button.Pressed) -> None:
|
|
"""Handle button press."""
|
|
self.dismiss()
|
|
|
|
def on_key(self, event) -> None:
|
|
"""Handle keyboard shortcuts."""
|
|
if event.key in ("escape", "enter"):
|
|
self.dismiss()
|