Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/client/cli/commands/chat.py: 0%

1import json 

2import sys 

3from collections.abc import Mapping, Sequence 

4from typing import Final 

5 

6import click 

7import requests 

8from rich.console import Console 

9from rich.panel import Panel 

10from rich.prompt import Prompt 

11from rich.table import Table 

12from typing_extensions import NotRequired, ReadOnly, TypedDict 

13 

14from ... import Client 

15from ...chat import ChatClient 

16from ._cli_context import cli_context_values 

17 

18 

19class _MessagesView(TypedDict): 

20 messages: ReadOnly[list[dict[str, str]]] 

21 

22 

23class _StreamDelta(TypedDict): 

24 content: ReadOnly[NotRequired[str]] 

25 

26 

27class _StreamChoice(TypedDict): 

28 delta: ReadOnly[NotRequired[_StreamDelta]] 

29 

30 

31class _StreamChunkView(TypedDict): 

32 choices: ReadOnly[Sequence[_StreamChoice]] 

33 

34 

35class _StreamErrorBody(TypedDict): 

36 error: ReadOnly[NotRequired[Mapping[str, object]]] 

37 

38 

39class _ErrorBodyView(TypedDict): 

40 body: ReadOnly[_StreamErrorBody] 

41 

42 

43def _get_available_models(ctx: click.Context) -> Sequence[Mapping[str, object]]: 

44 """Get list of available models from the proxy server""" 

45 try: 

46 context: Final = cli_context_values(ctx) 

47 client: Final = Client(base_url=context["base_url"], api_key=context["api_key"]) 

48 models_list: Final = client.models.list() 

49 # Ensure we return a list of dictionaries 

50 if isinstance(models_list, list): 

51 # Filter to ensure all items are dictionaries 

52 return [model for model in models_list if isinstance(model, dict)] 

53 return [] 

54 except Exception as e: 

55 click.echo(f"Warning: Could not fetch models list: {e}", err=True) 

56 return [] 

57 

58 

59def _select_model(console: Console, available_models: Sequence[Mapping[str, object]]) -> str | None: 

60 """Interactive model selection""" 

61 if not available_models: 

62 console.print("[yellow]No models available or could not fetch models list.[/yellow]") 

63 model_name: Final = Prompt.ask("Please enter a model name") 

64 return model_name if model_name.strip() else None 

65 

66 # Display available models in a table 

67 table: Final = Table(title="Available Models") 

68 table.add_column("Index", style="cyan", no_wrap=True) 

69 table.add_column("Model ID", style="green") 

70 table.add_column("Owned By", style="yellow") 

71 MAX_MODELS_TO_DISPLAY: Final = 200 

72 

73 models_to_display: Final = available_models[:MAX_MODELS_TO_DISPLAY] 

74 for i, model in enumerate(models_to_display): # Limit to first 200 models 

75 table.add_row(str(i + 1), str(model.get("id", "")), str(model.get("owned_by", ""))) 

76 

77 if len(available_models) > MAX_MODELS_TO_DISPLAY: 

78 console.print(f"\n[dim]... and {len(available_models) - MAX_MODELS_TO_DISPLAY} more models[/dim]") 

79 

80 console.print(table) 

81 

82 while True: 

83 try: 

84 choice = Prompt.ask( 

85 "\nSelect a model by entering the index number (or type a model name directly)", 

86 default="1", 

87 ).strip() 

88 

89 # Try to parse as index 

90 try: 

91 index = int(choice) - 1 

92 if 0 <= index < len(available_models): 

93 return str(available_models[index]["id"]) 

94 else: 

95 console.print( 

96 f"[red]Invalid index. Please enter a number between 1 and {len(available_models)}[/red]" 

97 ) 

98 continue 

99 except ValueError: 

100 # Not a number, treat as model name 

101 if choice: 

102 return choice 

103 else: 

104 console.print("[red]Please enter a valid model name or index[/red]") 

105 continue 

106 

107 except KeyboardInterrupt: 

108 console.print("\n[yellow]Model selection cancelled.[/yellow]") 

109 return None 

110 

111 

112@click.command() 

113@click.argument("model", required=False) 

114@click.option( 

115 "--temperature", 

116 "-t", 

117 type=float, 

118 default=0.7, 

119 help="Sampling temperature between 0 and 2 (default: 0.7)", 

120) 

121@click.option( 

122 "--max-tokens", 

123 type=int, 

124 help="Maximum number of tokens to generate", 

125) 

126@click.option( 

127 "--system", 

128 "-s", 

129 type=str, 

130 help="System message to set the behavior of the assistant", 

131) 

132@click.pass_context 

133def chat( 

134 ctx: click.Context, 

135 model: str | None, 

136 temperature: float, 

137 max_tokens: int | None = None, 

138 system: str | None = None, 

139): 

140 """Interactive chat with streaming responses 

141 

142 Examples: 

143 

144 # Chat with a specific model 

145 lite chat gpt-4 

146 

147 # Chat without specifying model (will show model selection) 

148 lite chat 

149 

150 # Chat with custom settings 

151 lite chat gpt-4 --temperature 0.9 --system "You are a helpful coding assistant" 

152 """ 

153 console: Final = Console() 

154 

155 # If no model specified, show model selection 

156 if not model: 

157 available_models: Final = _get_available_models(ctx) 

158 model = _select_model(console, available_models) 

159 if not model: 

160 console.print("[red]No model selected. Exiting.[/red]") 

161 return 

162 

163 context: Final = cli_context_values(ctx) 

164 client: Final = ChatClient(context["base_url"], context["api_key"]) 

165 

166 # Initialize conversation history 

167 messages: list[dict[str, str]] = [] 

168 

169 # Add system message if provided 

170 if system: 

171 messages.append({"role": "system", "content": system}) 

172 

173 # Display welcome message 

174 console.print( 

175 Panel.fit( 

176 f"[bold blue]LiteLLM Interactive Chat[/bold blue]\n" 

177 f"Model: [green]{model}[/green]\n" 

178 f"Temperature: [yellow]{temperature}[/yellow]\n" 

179 f"Max Tokens: [yellow]{max_tokens or 'unlimited'}[/yellow]\n\n" 

180 f"Type your messages and press Enter. Type '/quit' or '/exit' to end the session.\n" 

181 f"Type '/help' for more commands.", 

182 title="Chat Session", 

183 ) 

184 ) 

185 

186 try: 

187 while True: 

188 # Get user input 

189 try: 

190 user_input = console.input("\n[bold cyan]You:[/bold cyan] ").strip() 

191 except (EOFError, KeyboardInterrupt): 

192 console.print("\n[yellow]Chat session ended.[/yellow]") 

193 break 

194 

195 # Handle special commands 

196 should_exit, messages, new_model = _handle_special_commands(console, user_input, messages, system, ctx) 

197 

198 if should_exit: 

199 break 

200 if new_model: 

201 model = new_model 

202 

203 # Check if this was a special command that was handled (not a normal message) 

204 if ( 

205 user_input.lower().startswith( 

206 ( 

207 "/quit", 

208 "/exit", 

209 "/q", 

210 "/help", 

211 "/clear", 

212 "/history", 

213 "/save", 

214 "/load", 

215 "/model", 

216 ) 

217 ) 

218 or not user_input 

219 ): 

220 continue 

221 

222 # Add user message to conversation 

223 messages.append({"role": "user", "content": user_input}) 

224 

225 # Display assistant label 

226 console.print("\n[bold green]Assistant:[/bold green]") 

227 

228 # Stream the response 

229 assistant_content = _stream_response( 

230 console=console, 

231 client=client, 

232 model=model, 

233 messages=messages, 

234 temperature=temperature, 

235 max_tokens=max_tokens, 

236 ) 

237 

238 # Add assistant message to conversation history 

239 if assistant_content: 

240 messages.append({"role": "assistant", "content": assistant_content}) 

241 else: 

242 console.print("[red]Error: No content received from the model[/red]") 

243 

244 except KeyboardInterrupt: 

245 console.print("\n[yellow]Chat session interrupted.[/yellow]") 

246 

247 

248def _show_help(console: Console): 

249 """Show help for interactive chat commands""" 

250 help_text: Final = """ 

251[bold]Interactive Chat Commands:[/bold] 

252 

253[cyan]/help[/cyan] - Show this help message 

254[cyan]/quit[/cyan] - Exit the chat session (also /exit, /q) 

255[cyan]/clear[/cyan] - Clear conversation history 

256[cyan]/history[/cyan] - Show conversation history 

257[cyan]/model[/cyan] - Switch to a different model 

258[cyan]/save <name>[/cyan] - Save conversation to file 

259[cyan]/load <name>[/cyan] - Load conversation from file 

260 

261[bold]Tips:[/bold] 

262- Your conversation history is maintained during the session 

263- Use Ctrl+C to interrupt at any time 

264- Responses are streamed in real-time 

265- You can switch models mid-conversation with /model 

266 """ 

267 console.print(Panel(help_text, title="Help")) 

268 

269 

270def _show_history(console: Console, messages: list[dict[str, str]]): 

271 """Show conversation history""" 

272 if not messages: 

273 console.print("[yellow]No conversation history.[/yellow]") 

274 return 

275 

276 console.print(Panel.fit("[bold]Conversation History[/bold]", title="History")) 

277 

278 for i, message in enumerate(messages, 1): 

279 role = message["role"] 

280 content = message["content"] 

281 

282 if role == "system": 

283 console.print(f"[dim]{i}. [bold magenta]System:[/bold magenta] {content}[/dim]") 

284 elif role == "user": 

285 console.print(f"{i}. [bold cyan]You:[/bold cyan] {content}") 

286 elif role == "assistant": 

287 console.print( 

288 f"{i}. [bold green]Assistant:[/bold green] {content[:100]}{'...' if len(content) > 100 else ''}" 

289 ) 

290 

291 

292def _save_conversation(console: Console, messages: list[dict[str, str]], command: str): 

293 """Save conversation to a file""" 

294 parts: Final = command.split() 

295 if len(parts) < 2: 

296 console.print("[red]Usage: /save <filename>[/red]") 

297 return 

298 

299 filename = parts[1] 

300 if not filename.endswith(".json"): 

301 filename += ".json" 

302 

303 try: 

304 with open(filename, "w") as f: 

305 json.dump(messages, f, indent=2) 

306 console.print(f"[green]Conversation saved to {filename}[/green]") 

307 except Exception as e: 

308 console.print(f"[red]Error saving conversation: {e}[/red]") 

309 

310 

311def _load_conversation(console: Console, command: str, system: str | None) -> list[dict[str, str]]: 

312 """Load conversation from a file""" 

313 parts: Final = command.split() 

314 if len(parts) < 2: 

315 console.print("[red]Usage: /load <filename>[/red]") 

316 return [] 

317 

318 filename = parts[1] 

319 if not filename.endswith(".json"): 

320 filename += ".json" 

321 

322 try: 

323 with open(filename, "r") as f: 

324 loaded: Final[_MessagesView] = {"messages": json.load(f)} 

325 console.print(f"[green]Conversation loaded from {filename}[/green]") 

326 return loaded["messages"] 

327 except FileNotFoundError: 

328 console.print(f"[red]File not found: {filename}[/red]") 

329 except Exception as e: 

330 console.print(f"[red]Error loading conversation: {e}[/red]") 

331 

332 # Return empty list or just system message if load failed 

333 if system: 

334 return [{"role": "system", "content": system}] 

335 return [] 

336 

337 

338def _handle_special_commands( 

339 console: Console, 

340 user_input: str, 

341 messages: list[dict[str, str]], 

342 system: str | None, 

343 ctx: click.Context, 

344) -> tuple[bool, list[dict[str, str]], str | None]: 

345 """Handle special chat commands. Returns (should_exit, updated_messages, updated_model)""" 

346 if user_input.lower() in ["/quit", "/exit", "/q"]: 

347 console.print("[yellow]Chat session ended.[/yellow]") 

348 return True, messages, None 

349 elif user_input.lower() == "/help": 

350 _show_help(console) 

351 return False, messages, None 

352 elif user_input.lower() == "/clear": 

353 cleared_messages: Final[list[dict[str, str]]] = [{"role": "system", "content": system}] if system else [] 

354 console.print("[green]Conversation history cleared.[/green]") 

355 return False, cleared_messages, None 

356 elif user_input.lower() == "/history": 

357 _show_history(console, messages) 

358 return False, messages, None 

359 elif user_input.lower().startswith("/save"): 

360 _save_conversation(console, messages, user_input) 

361 return False, messages, None 

362 elif user_input.lower().startswith("/load"): 

363 new_messages = _load_conversation(console, user_input, system) 

364 return False, new_messages, None 

365 elif user_input.lower() == "/model": 

366 available_models: Final = _get_available_models(ctx) 

367 new_model: Final = _select_model(console, available_models) 

368 if new_model: 

369 console.print(f"[green]Switched to model: {new_model}[/green]") 

370 return False, messages, new_model 

371 return False, messages, None 

372 elif not user_input: 

373 return False, messages, None 

374 

375 # Not a special command 

376 return False, messages, None 

377 

378 

379def _stream_response( 

380 console: Console, 

381 client: ChatClient, 

382 model: str, 

383 messages: list[dict[str, str]], 

384 temperature: float, 

385 max_tokens: int | None, 

386) -> str | None: 

387 """Stream the model response and return the complete content""" 

388 try: 

389 assistant_content = "" 

390 for chunk in client.completions_stream( 

391 model=model, 

392 messages=messages, 

393 temperature=temperature, 

394 max_tokens=max_tokens, 

395 ): 

396 streamed: _StreamChunkView = {"choices": chunk.get("choices", ())} 

397 if len(streamed["choices"]) > 0: 

398 delta = streamed["choices"][0].get("delta", {}) 

399 content = delta.get("content", "") 

400 if content: 

401 assistant_content += content 

402 console.print(content, end="") 

403 sys.stdout.flush() 

404 

405 console.print() # Add newline after streaming 

406 return assistant_content if assistant_content else None 

407 

408 except requests.exceptions.HTTPError as e: 

409 console.print(f"\n[red]Error: HTTP {e.response.status_code}[/red]") 

410 try: 

411 error_body: Final[_ErrorBodyView] = {"body": e.response.json()} 

412 console.print(f"[red]{error_body['body'].get('error', {}).get('message', 'Unknown error')}[/red]") 

413 except json.JSONDecodeError: 

414 console.print(f"[red]{e.response.text}[/red]") 

415 return None 

416 except Exception as e: 

417 console.print(f"\n[red]Error: {e}[/red]") 

418 return None