115 lines
3.8 KiB
Python
115 lines
3.8 KiB
Python
import json
|
|
import sys
|
|
import argparse
|
|
from typing import Optional, List, Dict, Any
|
|
from dataclasses import dataclass
|
|
from rich.console import Console
|
|
from rich.spinner import Spinner
|
|
from qwen_agent.agents import Assistant
|
|
from transformers import pipeline
|
|
|
|
@dataclass
|
|
class AgentConfig:
|
|
model: str
|
|
server: str
|
|
api_key: str
|
|
max_tokens: int = 30000
|
|
enable_thinking: bool = True
|
|
tools: Optional[List[Dict[str, Any]]] = None
|
|
|
|
def __post_init__(self):
|
|
if self.tools is None:
|
|
self.tools = [
|
|
{'mcpServers': {
|
|
'time': {
|
|
'command': 'uvx',
|
|
'args': ['mcp-server-time', '--local-timezone=Europe/London']
|
|
},
|
|
"fetch": {
|
|
"command": "uvx",
|
|
"args": ["mcp-server-fetch"]
|
|
},
|
|
"ddg-search": {
|
|
"command": "npx",
|
|
"args": ["-y", "duckduckgo-mcp-server"]
|
|
},
|
|
}},
|
|
'code_interpreter',
|
|
]
|
|
|
|
def setup_argparse():
|
|
parser = argparse.ArgumentParser(description='Qwen3 Agent CLI')
|
|
parser.add_argument('--model', default='qwen3:32b',
|
|
help='Model identifier (default: qwen3:32b)')
|
|
parser.add_argument('--server', default='http://localhost:11434/v1',
|
|
help='Model server URL (default: http://localhost:11434/v1)')
|
|
parser.add_argument('--api-key', default='EMPTY',
|
|
help='API key for the model server (default: EMPTY)')
|
|
|
|
subparsers = parser.add_subparsers(dest='command', help='Available commands')
|
|
|
|
# Prompt command
|
|
prompt_parser = subparsers.add_parser('prompt', help='Run agent with a prompt')
|
|
prompt_parser.add_argument('text', nargs='?', default='-',
|
|
help='Prompt text or "-" for stdin (default: -)')
|
|
|
|
return parser
|
|
|
|
def read_prompt(text: str) -> str:
|
|
"""Read prompt from argument or stdin if text is '-'"""
|
|
if text == '-':
|
|
return sys.stdin.read().strip()
|
|
return text
|
|
|
|
def run_agent(config: AgentConfig, prompt: str) -> None:
|
|
"""Run the agent with the given configuration and prompt"""
|
|
llm_cfg = {
|
|
'model': config.model,
|
|
'model_server': config.server,
|
|
'api_key': config.api_key,
|
|
}
|
|
|
|
# Define Agent
|
|
bot = Assistant(llm=llm_cfg, function_list=config.tools)
|
|
console = Console()
|
|
|
|
# Streaming generation
|
|
messages = [{'role': 'user', 'content': prompt}]
|
|
|
|
final_responses = None
|
|
try:
|
|
with console.status("[bold blue]Thinking...", spinner="dots") as status:
|
|
for responses in bot.run(messages=messages,
|
|
enable_thinking=config.enable_thinking,
|
|
max_tokens=config.max_tokens):
|
|
final_responses = responses.pop()
|
|
except Exception as e:
|
|
console.print(f"[bold red]An error occurred during agent execution:[/] {e}")
|
|
|
|
# Pretty-print the final response object
|
|
if final_responses:
|
|
console.print("\n[bold green]--- Full Response Object ---[/]")
|
|
console.print(json.dumps(final_responses, indent=2))
|
|
console.print("\n[bold green]--- Extracted Content ---[/]")
|
|
console.print(final_responses.get('content', 'No content found in response.'))
|
|
else:
|
|
console.print("[bold red]No final response received from the agent.[/]")
|
|
|
|
def main():
|
|
parser = setup_argparse()
|
|
args = parser.parse_args()
|
|
|
|
if args.command == 'prompt':
|
|
prompt_text = read_prompt(args.text)
|
|
config = AgentConfig(
|
|
model=args.model,
|
|
server=args.server,
|
|
api_key=args.api_key
|
|
)
|
|
run_agent(config, prompt_text)
|
|
else:
|
|
parser.print_help()
|
|
|
|
if __name__ == '__main__':
|
|
main()
|