File size: 4,070 Bytes
1f1b4db |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 |
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 no API key selected or available, do registration
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())
|