| | from managers.chat_manager import ChatManager
|
| | from api.telemetry import TelemetryAPI
|
| | from config.api_keys import APIKeyManager
|
| | from config.models import ModelID
|
| | from auth.register import Registration
|
| | from auth.login import open_auth_url
|
| | from api.chat.chat_api import ChatAPI, ChatConfig
|
| |
|
| | import urllib3
|
| | import asyncio
|
| |
|
| | urllib3.disable_warnings()
|
| |
|
| | from rich.console import Console
|
| | from rich.progress import Progress, SpinnerColumn, TextColumn
|
| |
|
| | console = Console()
|
| |
|
| |
|
| | def select_api_key(key_manager: APIKeyManager) -> str:
|
| | """Let user select an API key from available keys"""
|
| | keys = key_manager.list_keys()
|
| | if not keys:
|
| | return None
|
| |
|
| | console.print("\n[bold cyan]Available API keys:[/]")
|
| | for i, key in enumerate(keys):
|
| | console.print(f"[green]{i + 1}[/]. {key}")
|
| |
|
| | while True:
|
| | try:
|
| | choice = int(input("\nSelect API key (number) or 0 for new login: "))
|
| | if choice == 0:
|
| | return None
|
| | if 1 <= choice <= len(keys):
|
| | return keys[choice - 1]
|
| | except ValueError:
|
| | return keys[0]
|
| | print("Invalid selection. Please try again.")
|
| |
|
| |
|
| | async def main():
|
| | key_manager = APIKeyManager()
|
| | chat_manager = ChatManager()
|
| | api_key = select_api_key(key_manager)
|
| |
|
| |
|
| | if api_key is None:
|
| | _ = open_auth_url()
|
| | firebase_token = input("Please paste your firebase token: ")
|
| | api_key = await Registration.register_user(firebase_token.strip())
|
| | key_manager.add_key(api_key)
|
| | print(f"Registered successfully. API key: {api_key}")
|
| |
|
| | with Progress(
|
| | SpinnerColumn(),
|
| | TextColumn("[progress.description]{task.description}"),
|
| | console=console,
|
| | ) as progress:
|
| | progress.add_task("Initializing Chat API...", total=None)
|
| | tele = TelemetryAPI(api_key=api_key)
|
| | await tele.do_telemetry()
|
| |
|
| | await chat_manager.add_chat(api_key)
|
| |
|
| | progress.stop()
|
| |
|
| | with Progress(
|
| | SpinnerColumn(),
|
| | TextColumn("[progress.description]{task.description}"),
|
| | console=console,
|
| | ) as progress:
|
| | _ = progress.add_task("Sending message...", total=None)
|
| |
|
| | config = ChatConfig(
|
| | model_id=ModelID.MODEL_CLAUDE_3_OPUS_20240229,
|
| | temperature=0.7,
|
| | max_tokens=8192,
|
| | )
|
| |
|
| | message = [
|
| | {"role": "user", "content": "1234567890"},
|
| | {"role": "assistant", "content": "1234567890"},
|
| | {"role": "user", "content": "1234567890"},
|
| | {"role": "assistant", "content": "1234567890"},
|
| | {"role": "user", "content": "1234567890"},
|
| | {"role": "assistant", "content": "1234567890"},
|
| | {"role": "user", "content": "1234567890"},
|
| | {"role": "assistant", "content": "1234567890"},
|
| | {"role": "user", "content": "1234567890"},
|
| | {"role": "assistant", "content": "1234567890"},
|
| | {"role": "user", "content": "tell me a joke"},
|
| | ]
|
| | first_chunk = True
|
| | chat = await chat_manager.get_chat()
|
| |
|
| | tokens = 0
|
| |
|
| | async for chunk, token_count in chat.send_message(
|
| | messages=message,
|
| | config=config,
|
| | system_prompt="You are a helpful assistant.",
|
| | stream=True,
|
| | ):
|
| | if first_chunk:
|
| | progress.stop()
|
| | first_chunk = False
|
| | console.print("=" * 20)
|
| | if chunk:
|
| | console.print(
|
| | chunk, end="", markup=True, highlight=True, emoji=True
|
| | )
|
| | tokens += token_count
|
| |
|
| | console.print(f"\n[bold green]Tokens: {tokens}[/]")
|
| |
|
| |
|
| | if __name__ == "__main__":
|
| | asyncio.run(main())
|
| |
|
