|
|
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())
|
|
|
|
