diff --git a/oai.py b/oai.py index ad0c358..c9411f0 100644 --- a/oai.py +++ b/oai.py @@ -70,6 +70,35 @@ except Exception as e: models_data = [] text_models = [] +# **Function to fetch credit information from OpenRouter API** +def get_credits(api_key: str, base_url: str = OPENROUTER_BASE_URL) -> Optional[dict]: + """ + Fetch credit information from OpenRouter API. + + Returns a dict with 'total_credits', 'used_credits', 'credits_left' or None on error. + Based on OpenRouter's /credits endpoint ([openrouter.ai/docs/limits](https://openrouter.ai/docs/limits)). + """ + if not api_key: + return None + url = f"{base_url}/credits" + headers = {"Authorization": f"Bearer {api_key}"} + try: + response = requests.get(url, headers=headers) + response.raise_for_status() + data = response.json().get('data', {}) + total_credits = float(data.get('total_credits', 0)) + total_usage = float(data.get('total_usage', 0)) + credits_left = total_credits - total_usage + return { + 'total_credits': f"${total_credits:.2f}", + 'used_credits': f"${total_usage:.2f}", + 'credits_left': f"${credits_left:.2f}" + } + except Exception as e: + # Gracefully handle errors (e.g., invalid API key, network issues) + console.print(f"[bold red]Error fetching credits: {e}[/]") + return None + # Function to clear the screen using ANSI escape sequences (cross-platform) def clear_screen(): """Clear the terminal screen using ANSI escape codes, with fallback to newlines.""" @@ -81,6 +110,7 @@ def clear_screen(): # Fallback: Fill with newlines (simpler, per [medium.com](https://medium.com/@ryan_forrester_/c-screen-clearing-how-to-guide-cff5bf764ccd)) print("\n" * 100) + @app.command() def chat(): """Start the oAI chat app with OpenRouter models.""" @@ -153,7 +183,7 @@ def chat(): console.print("[bold red]Invalid input. Enter a number.[/]") continue - # Handle /config command (updated to include stream toggle) + # Handle /config command (**UPDATED:** Now includes credit info) if user_input.startswith("/config"): args = user_input[8:].strip().lower() # Get args after "/config" if args == "api": @@ -195,16 +225,37 @@ def chat(): table.add_row("Streaming", "Enabled" if STREAM_ENABLED == "on" else "Disabled") table.add_row("Database", str(database) or "[Not set]") table.add_row("Current Model", "[Not set]" if selected_model is None else str(selected_model["name"])) + + # Fetch and display credit info + credits = get_credits(API_KEY, OPENROUTER_BASE_URL) + if credits: + table.add_row("Total Credits", credits['total_credits']) + table.add_row("Used Credits", credits['used_credits']) + table.add_row("Credits Left", credits['credits_left']) + else: + table.add_row("Total Credits", "[Unavailable - Check API key]") + table.add_row("Used Credits", "[Unavailable - Check API key]") + table.add_row("Credits Left", "[Unavailable - Check API key]") + console.print(Panel(table, title="[bold green]Current Configurations[/]", title_align="left")) continue + # **Handle /credits command to display credits left** + if user_input.lower() == "/credits": + credits = get_credits(API_KEY, OPENROUTER_BASE_URL) + if credits: + console.print(f"[bold green]Credits left: {credits['credits_left']}[/]") + else: + console.print("[bold red]Unable to fetch credits. Check your API key or network.[/]") + continue + # Handle /clear command to clear the screen if user_input.lower() == "/clear": clear_screen() console.print("[bold cyan]Screen cleared. Ready for your next input![/]") continue - # Handle /help command (updated: now includes /clear) + # Handle /help command (**UPDATED:** Now includes /credits) if user_input.lower() == "/help": help_table = Table("Command", "Description", "Example", show_header=True, header_style="bold cyan") help_table.add_row( @@ -229,8 +280,13 @@ def chat(): ) help_table.add_row( "/config", - "View all current configurations.", - "/config\n(Displays table of API Key, Base URL, etc.)" + "View all current configurations, including credits.", + "/config\n(Displays table with credits info)" + ) + help_table.add_row( + "/credits", + "Display credits left on your OpenRouter account.", + "/credits\n[bold green]Credits left: $4.23[/bold green]" ) help_table.add_row( "/clear",