chore: initial snapshot for gitea/github upload
This commit is contained in:
@@ -0,0 +1,406 @@
|
||||
import json
|
||||
import sys
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
import click
|
||||
import requests
|
||||
from rich.console import Console
|
||||
from rich.panel import Panel
|
||||
from rich.prompt import Prompt
|
||||
from rich.table import Table
|
||||
|
||||
from ... import Client
|
||||
from ...chat import ChatClient
|
||||
|
||||
|
||||
def _get_available_models(ctx: click.Context) -> List[Dict[str, Any]]:
|
||||
"""Get list of available models from the proxy server"""
|
||||
try:
|
||||
client = Client(base_url=ctx.obj["base_url"], api_key=ctx.obj["api_key"])
|
||||
models_list = client.models.list()
|
||||
# Ensure we return a list of dictionaries
|
||||
if isinstance(models_list, list):
|
||||
# Filter to ensure all items are dictionaries
|
||||
return [model for model in models_list if isinstance(model, dict)]
|
||||
return []
|
||||
except Exception as e:
|
||||
click.echo(f"Warning: Could not fetch models list: {e}", err=True)
|
||||
return []
|
||||
|
||||
|
||||
def _select_model(
|
||||
console: Console, available_models: List[Dict[str, Any]]
|
||||
) -> Optional[str]:
|
||||
"""Interactive model selection"""
|
||||
if not available_models:
|
||||
console.print(
|
||||
"[yellow]No models available or could not fetch models list.[/yellow]"
|
||||
)
|
||||
model_name = Prompt.ask("Please enter a model name")
|
||||
return model_name if model_name.strip() else None
|
||||
|
||||
# Display available models in a table
|
||||
table = Table(title="Available Models")
|
||||
table.add_column("Index", style="cyan", no_wrap=True)
|
||||
table.add_column("Model ID", style="green")
|
||||
table.add_column("Owned By", style="yellow")
|
||||
MAX_MODELS_TO_DISPLAY = 200
|
||||
|
||||
models_to_display: List[Dict[str, Any]] = available_models[:MAX_MODELS_TO_DISPLAY]
|
||||
for i, model in enumerate(models_to_display): # Limit to first 200 models
|
||||
table.add_row(
|
||||
str(i + 1), str(model.get("id", "")), str(model.get("owned_by", ""))
|
||||
)
|
||||
|
||||
if len(available_models) > MAX_MODELS_TO_DISPLAY:
|
||||
console.print(
|
||||
f"\n[dim]... and {len(available_models) - MAX_MODELS_TO_DISPLAY} more models[/dim]"
|
||||
)
|
||||
|
||||
console.print(table)
|
||||
|
||||
while True:
|
||||
try:
|
||||
choice = Prompt.ask(
|
||||
"\nSelect a model by entering the index number (or type a model name directly)",
|
||||
default="1",
|
||||
).strip()
|
||||
|
||||
# Try to parse as index
|
||||
try:
|
||||
index = int(choice) - 1
|
||||
if 0 <= index < len(available_models):
|
||||
return available_models[index]["id"]
|
||||
else:
|
||||
console.print(
|
||||
f"[red]Invalid index. Please enter a number between 1 and {len(available_models)}[/red]"
|
||||
)
|
||||
continue
|
||||
except ValueError:
|
||||
# Not a number, treat as model name
|
||||
if choice:
|
||||
return choice
|
||||
else:
|
||||
console.print("[red]Please enter a valid model name or index[/red]")
|
||||
continue
|
||||
|
||||
except KeyboardInterrupt:
|
||||
console.print("\n[yellow]Model selection cancelled.[/yellow]")
|
||||
return None
|
||||
|
||||
|
||||
@click.command()
|
||||
@click.argument("model", required=False)
|
||||
@click.option(
|
||||
"--temperature",
|
||||
"-t",
|
||||
type=float,
|
||||
default=0.7,
|
||||
help="Sampling temperature between 0 and 2 (default: 0.7)",
|
||||
)
|
||||
@click.option(
|
||||
"--max-tokens",
|
||||
type=int,
|
||||
help="Maximum number of tokens to generate",
|
||||
)
|
||||
@click.option(
|
||||
"--system",
|
||||
"-s",
|
||||
type=str,
|
||||
help="System message to set the behavior of the assistant",
|
||||
)
|
||||
@click.pass_context
|
||||
def chat(
|
||||
ctx: click.Context,
|
||||
model: Optional[str],
|
||||
temperature: float,
|
||||
max_tokens: Optional[int] = None,
|
||||
system: Optional[str] = None,
|
||||
):
|
||||
"""Interactive chat with streaming responses
|
||||
|
||||
Examples:
|
||||
|
||||
# Chat with a specific model
|
||||
litellm-proxy chat gpt-4
|
||||
|
||||
# Chat without specifying model (will show model selection)
|
||||
litellm-proxy chat
|
||||
|
||||
# Chat with custom settings
|
||||
litellm-proxy chat gpt-4 --temperature 0.9 --system "You are a helpful coding assistant"
|
||||
"""
|
||||
console = Console()
|
||||
|
||||
# If no model specified, show model selection
|
||||
if not model:
|
||||
available_models = _get_available_models(ctx)
|
||||
model = _select_model(console, available_models)
|
||||
if not model:
|
||||
console.print("[red]No model selected. Exiting.[/red]")
|
||||
return
|
||||
|
||||
client = ChatClient(ctx.obj["base_url"], ctx.obj["api_key"])
|
||||
|
||||
# Initialize conversation history
|
||||
messages: List[Dict[str, Any]] = []
|
||||
|
||||
# Add system message if provided
|
||||
if system:
|
||||
messages.append({"role": "system", "content": system})
|
||||
|
||||
# Display welcome message
|
||||
console.print(
|
||||
Panel.fit(
|
||||
f"[bold blue]LiteLLM Interactive Chat[/bold blue]\n"
|
||||
f"Model: [green]{model}[/green]\n"
|
||||
f"Temperature: [yellow]{temperature}[/yellow]\n"
|
||||
f"Max Tokens: [yellow]{max_tokens or 'unlimited'}[/yellow]\n\n"
|
||||
f"Type your messages and press Enter. Type '/quit' or '/exit' to end the session.\n"
|
||||
f"Type '/help' for more commands.",
|
||||
title="🤖 Chat Session",
|
||||
)
|
||||
)
|
||||
|
||||
try:
|
||||
while True:
|
||||
# Get user input
|
||||
try:
|
||||
user_input = console.input("\n[bold cyan]You:[/bold cyan] ").strip()
|
||||
except (EOFError, KeyboardInterrupt):
|
||||
console.print("\n[yellow]Chat session ended.[/yellow]")
|
||||
break
|
||||
|
||||
# Handle special commands
|
||||
should_exit, messages, new_model = _handle_special_commands(
|
||||
console, user_input, messages, system, ctx
|
||||
)
|
||||
|
||||
if should_exit:
|
||||
break
|
||||
if new_model:
|
||||
model = new_model
|
||||
|
||||
# Check if this was a special command that was handled (not a normal message)
|
||||
if (
|
||||
user_input.lower().startswith(
|
||||
(
|
||||
"/quit",
|
||||
"/exit",
|
||||
"/q",
|
||||
"/help",
|
||||
"/clear",
|
||||
"/history",
|
||||
"/save",
|
||||
"/load",
|
||||
"/model",
|
||||
)
|
||||
)
|
||||
or not user_input
|
||||
):
|
||||
continue
|
||||
|
||||
# Add user message to conversation
|
||||
messages.append({"role": "user", "content": user_input})
|
||||
|
||||
# Display assistant label
|
||||
console.print("\n[bold green]Assistant:[/bold green]")
|
||||
|
||||
# Stream the response
|
||||
assistant_content = _stream_response(
|
||||
console=console,
|
||||
client=client,
|
||||
model=model,
|
||||
messages=messages,
|
||||
temperature=temperature,
|
||||
max_tokens=max_tokens,
|
||||
)
|
||||
|
||||
# Add assistant message to conversation history
|
||||
if assistant_content:
|
||||
messages.append({"role": "assistant", "content": assistant_content})
|
||||
else:
|
||||
console.print("[red]Error: No content received from the model[/red]")
|
||||
|
||||
except KeyboardInterrupt:
|
||||
console.print("\n[yellow]Chat session interrupted.[/yellow]")
|
||||
|
||||
|
||||
def _show_help(console: Console):
|
||||
"""Show help for interactive chat commands"""
|
||||
help_text = """
|
||||
[bold]Interactive Chat Commands:[/bold]
|
||||
|
||||
[cyan]/help[/cyan] - Show this help message
|
||||
[cyan]/quit[/cyan] - Exit the chat session (also /exit, /q)
|
||||
[cyan]/clear[/cyan] - Clear conversation history
|
||||
[cyan]/history[/cyan] - Show conversation history
|
||||
[cyan]/model[/cyan] - Switch to a different model
|
||||
[cyan]/save <name>[/cyan] - Save conversation to file
|
||||
[cyan]/load <name>[/cyan] - Load conversation from file
|
||||
|
||||
[bold]Tips:[/bold]
|
||||
- Your conversation history is maintained during the session
|
||||
- Use Ctrl+C to interrupt at any time
|
||||
- Responses are streamed in real-time
|
||||
- You can switch models mid-conversation with /model
|
||||
"""
|
||||
console.print(Panel(help_text, title="Help"))
|
||||
|
||||
|
||||
def _show_history(console: Console, messages: List[Dict[str, Any]]):
|
||||
"""Show conversation history"""
|
||||
if not messages:
|
||||
console.print("[yellow]No conversation history.[/yellow]")
|
||||
return
|
||||
|
||||
console.print(Panel.fit("[bold]Conversation History[/bold]", title="History"))
|
||||
|
||||
for i, message in enumerate(messages, 1):
|
||||
role = message["role"]
|
||||
content = message["content"]
|
||||
|
||||
if role == "system":
|
||||
console.print(
|
||||
f"[dim]{i}. [bold magenta]System:[/bold magenta] {content}[/dim]"
|
||||
)
|
||||
elif role == "user":
|
||||
console.print(f"{i}. [bold cyan]You:[/bold cyan] {content}")
|
||||
elif role == "assistant":
|
||||
console.print(
|
||||
f"{i}. [bold green]Assistant:[/bold green] {content[:100]}{'...' if len(content) > 100 else ''}"
|
||||
)
|
||||
|
||||
|
||||
def _save_conversation(console: Console, messages: List[Dict[str, Any]], command: str):
|
||||
"""Save conversation to a file"""
|
||||
parts = command.split()
|
||||
if len(parts) < 2:
|
||||
console.print("[red]Usage: /save <filename>[/red]")
|
||||
return
|
||||
|
||||
filename = parts[1]
|
||||
if not filename.endswith(".json"):
|
||||
filename += ".json"
|
||||
|
||||
try:
|
||||
with open(filename, "w") as f:
|
||||
json.dump(messages, f, indent=2)
|
||||
console.print(f"[green]Conversation saved to {filename}[/green]")
|
||||
except Exception as e:
|
||||
console.print(f"[red]Error saving conversation: {e}[/red]")
|
||||
|
||||
|
||||
def _load_conversation(
|
||||
console: Console, command: str, system: Optional[str]
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""Load conversation from a file"""
|
||||
parts = command.split()
|
||||
if len(parts) < 2:
|
||||
console.print("[red]Usage: /load <filename>[/red]")
|
||||
return []
|
||||
|
||||
filename = parts[1]
|
||||
if not filename.endswith(".json"):
|
||||
filename += ".json"
|
||||
|
||||
try:
|
||||
with open(filename, "r") as f:
|
||||
messages = json.load(f)
|
||||
console.print(f"[green]Conversation loaded from {filename}[/green]")
|
||||
return messages
|
||||
except FileNotFoundError:
|
||||
console.print(f"[red]File not found: {filename}[/red]")
|
||||
except Exception as e:
|
||||
console.print(f"[red]Error loading conversation: {e}[/red]")
|
||||
|
||||
# Return empty list or just system message if load failed
|
||||
if system:
|
||||
return [{"role": "system", "content": system}]
|
||||
return []
|
||||
|
||||
|
||||
def _handle_special_commands(
|
||||
console: Console,
|
||||
user_input: str,
|
||||
messages: List[Dict[str, Any]],
|
||||
system: Optional[str],
|
||||
ctx: click.Context,
|
||||
) -> tuple[bool, List[Dict[str, Any]], Optional[str]]:
|
||||
"""Handle special chat commands. Returns (should_exit, updated_messages, updated_model)"""
|
||||
if user_input.lower() in ["/quit", "/exit", "/q"]:
|
||||
console.print("[yellow]Chat session ended.[/yellow]")
|
||||
return True, messages, None
|
||||
elif user_input.lower() == "/help":
|
||||
_show_help(console)
|
||||
return False, messages, None
|
||||
elif user_input.lower() == "/clear":
|
||||
new_messages = []
|
||||
if system:
|
||||
new_messages.append({"role": "system", "content": system})
|
||||
console.print("[green]Conversation history cleared.[/green]")
|
||||
return False, new_messages, None
|
||||
elif user_input.lower() == "/history":
|
||||
_show_history(console, messages)
|
||||
return False, messages, None
|
||||
elif user_input.lower().startswith("/save"):
|
||||
_save_conversation(console, messages, user_input)
|
||||
return False, messages, None
|
||||
elif user_input.lower().startswith("/load"):
|
||||
new_messages = _load_conversation(console, user_input, system)
|
||||
return False, new_messages, None
|
||||
elif user_input.lower() == "/model":
|
||||
available_models = _get_available_models(ctx)
|
||||
new_model = _select_model(console, available_models)
|
||||
if new_model:
|
||||
console.print(f"[green]Switched to model: {new_model}[/green]")
|
||||
return False, messages, new_model
|
||||
return False, messages, None
|
||||
elif not user_input:
|
||||
return False, messages, None
|
||||
|
||||
# Not a special command
|
||||
return False, messages, None
|
||||
|
||||
|
||||
def _stream_response(
|
||||
console: Console,
|
||||
client: ChatClient,
|
||||
model: str,
|
||||
messages: List[Dict[str, Any]],
|
||||
temperature: float,
|
||||
max_tokens: Optional[int],
|
||||
) -> Optional[str]:
|
||||
"""Stream the model response and return the complete content"""
|
||||
try:
|
||||
assistant_content = ""
|
||||
for chunk in client.completions_stream(
|
||||
model=model,
|
||||
messages=messages,
|
||||
temperature=temperature,
|
||||
max_tokens=max_tokens,
|
||||
):
|
||||
if "choices" in chunk and len(chunk["choices"]) > 0:
|
||||
delta = chunk["choices"][0].get("delta", {})
|
||||
content = delta.get("content", "")
|
||||
if content:
|
||||
assistant_content += content
|
||||
console.print(content, end="")
|
||||
sys.stdout.flush()
|
||||
|
||||
console.print() # Add newline after streaming
|
||||
return assistant_content if assistant_content else None
|
||||
|
||||
except requests.exceptions.HTTPError as e:
|
||||
console.print(f"\n[red]Error: HTTP {e.response.status_code}[/red]")
|
||||
try:
|
||||
error_body = e.response.json()
|
||||
console.print(
|
||||
f"[red]{error_body.get('error', {}).get('message', 'Unknown error')}[/red]"
|
||||
)
|
||||
except json.JSONDecodeError:
|
||||
console.print(f"[red]{e.response.text}[/red]")
|
||||
return None
|
||||
except Exception as e:
|
||||
console.print(f"\n[red]Error: {str(e)}[/red]")
|
||||
return None
|
||||
Reference in New Issue
Block a user