diff --git a/agent.py b/agent.py new file mode 100644 index 0000000..8ba834a --- /dev/null +++ b/agent.py @@ -0,0 +1,68 @@ +import json +from dataclasses import dataclass +from typing import Optional, List, Dict, Any +from rich.console import Console +from qwen_agent.agents import Assistant + +@dataclass +class Agent: + model: str + server: str + api_key: str + max_tokens: int = 30000 + enable_thinking: bool = True + tools: Optional[List[Dict[str, Any]]] = None + console: Console = Console() + + 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 run(self, prompt: str) -> None: + """Run the agent with the given prompt""" + llm_cfg = { + 'model': self.model, + 'model_server': self.server, + 'api_key': self.api_key, + } + + # Define Agent + bot = Assistant(llm=llm_cfg, function_list=self.tools) + + # Streaming generation + messages = [{'role': 'user', 'content': prompt}] + + final_responses = None + try: + with self.console.status("[bold blue]Thinking...", spinner="dots") as status: + for responses in bot.run(messages=messages, + enable_thinking=self.enable_thinking, + max_tokens=self.max_tokens): + final_responses = responses.pop() + except Exception as e: + self.console.print(f"[bold red]An error occurred during agent execution:[/] {e}") + + # Pretty-print the final response object + if final_responses: + self.console.print("\n[bold green]--- Full Response Object ---[/]") + self.console.print(json.dumps(final_responses, indent=2)) + self.console.print("\n[bold green]--- Extracted Content ---[/]") + self.console.print(final_responses.get('content', 'No content found in response.')) + else: + self.console.print("[bold red]No final response received from the agent.[/]") diff --git a/agentic_search.py b/agentic_search.py index 1ee99c4..250da5c 100644 --- a/agentic_search.py +++ b/agentic_search.py @@ -1,41 +1,6 @@ -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', - ] +from agent import Agent def setup_argparse(): parser = argparse.ArgumentParser(description='Qwen3 Agent CLI') @@ -61,39 +26,6 @@ def read_prompt(text: str) -> str: 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() @@ -101,12 +33,12 @@ def main(): if args.command == 'prompt': prompt_text = read_prompt(args.text) - config = AgentConfig( + agent = Agent( model=args.model, server=args.server, api_key=args.api_key ) - run_agent(config, prompt_text) + agent.run(prompt_text) else: parser.print_help()