agentic-search/agentic_search.py

98 lines
3.3 KiB
Python

import json
import sys
import argparse
from typing import Optional
from rich.console import Console
from rich.spinner import Spinner
from qwen_agent.agents import Assistant
from transformers import pipeline
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(model: str, server: str, api_key: str, prompt: str) -> None:
"""Run the agent with the given configuration and prompt"""
llm_cfg = {
'model': model,
'model_server': server,
'api_key': api_key,
}
# Define Tools
tools = [
{'mcpServers': { # You can specify the MCP configuration file
'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', # Built-in tools
]
# Define Agent
bot = Assistant(llm=llm_cfg, function_list=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=True, max_tokens=30000):
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)
run_agent(args.model, args.server, args.api_key, prompt_text)
else:
parser.print_help()
if __name__ == '__main__':
main()