Spaces:
Running on Zero
Running on Zero
| import os | |
| import yaml | |
| from dotenv import load_dotenv | |
| from smolagents import CodeAgent, InferenceClientModel, load_tool | |
| from Gradio_UI import GradioUI | |
| from tools.final_answer import FinalAnswerTool | |
| from tools.get_current_time_in_timezone import get_current_time_in_timezone | |
| from tools.get_weather import get_weather | |
| from tools.search_wikipedia import search_wikipedia | |
| from tools.visit_webpage import VisitWebpageTool | |
| from tools.web_search import DuckDuckGoSearchTool | |
| load_dotenv() | |
| # --- Core Tools & Agent Setup --- | |
| final_answer = FinalAnswerTool() | |
| web_search = DuckDuckGoSearchTool() | |
| visit_webpage = VisitWebpageTool() | |
| hf_token = os.getenv("HF_TOKEN") | |
| if not hf_token: | |
| raise RuntimeError( | |
| "HF_TOKEN is required. Local model execution is disabled in this project." | |
| ) | |
| # Image generation tool from Hugging Face Hub | |
| image_generation_tool = None | |
| try: | |
| image_generation_tool = load_tool( | |
| "agents-course/text-to-image", | |
| trust_remote_code=True, | |
| token=hf_token, | |
| ) | |
| except Exception as e: | |
| print(f"Notice: text-to-image tool could not be loaded: {e}") | |
| # Model configuration (cloud by default; Ollama supported for local learning/testing) | |
| model_id = os.getenv("MODEL_ID", "Qwen/Qwen2.5-Coder-32B-Instruct") | |
| if "ollama" in model_id.lower(): | |
| from smolagents import LiteLLMModel | |
| model = LiteLLMModel( | |
| model_id=model_id, | |
| api_base=os.getenv("OLLAMA_API_BASE", "http://127.0.0.1:11434"), | |
| num_ctx=int(os.getenv("OLLAMA_NUM_CTX", "8192")), | |
| temperature=0.5, | |
| ) | |
| else: | |
| model = InferenceClientModel( | |
| model_id=model_id, | |
| token=hf_token, | |
| max_tokens=2096, | |
| temperature=0.5, | |
| custom_role_conversions=None, | |
| ) | |
| # Load prompt templates | |
| prompts_file = os.path.join(os.path.dirname(__file__), "prompts.yaml") | |
| with open(prompts_file, "r") as stream: | |
| prompt_templates = yaml.safe_load(stream) | |
| # --- Specialist Agents (Unit 2: Hierarchical Multi-Agent System) --- | |
| # Specialist 1: Web & Wikipedia Researcher | |
| research_agent = CodeAgent( | |
| model=model, | |
| tools=[web_search, search_wikipedia, visit_webpage], | |
| max_steps=5, | |
| verbosity_level=1, | |
| name="web_researcher", | |
| description="A specialist researcher that searches DuckDuckGo, queries Wikipedia, and reads webpage content to gather accurate facts and summaries. ALWAYS call this agent with a 'task' argument, e.g. web_researcher(task='Find the capital of Japan').", | |
| prompt_templates=prompt_templates | |
| ) | |
| # Specialist 2: Weather & Timezone Specialist | |
| local_info_agent = CodeAgent( | |
| model=model, | |
| tools=[get_current_time_in_timezone, get_weather], | |
| max_steps=4, | |
| verbosity_level=1, | |
| name="local_info_specialist", | |
| description="A specialist agent that fetches real-time live weather reports and calculates current time in timezones worldwide. ALWAYS call this agent with a 'task' argument, e.g. local_info_specialist(task='Get current weather in Tokyo'). Never call weather() or get_weather() directly from the manager.", | |
| prompt_templates=prompt_templates | |
| ) | |
| # Make managed agents tolerant of common argument names (query/question/etc.) | |
| def _adapt_managed_agent(agent): | |
| original_call = agent.__call__ | |
| def _call(task=None, **kwargs): | |
| if task is None: | |
| for key in ("task", "query", "question", "request", "prompt"): | |
| if key in kwargs: | |
| task = kwargs.pop(key) | |
| break | |
| if task is None and kwargs: | |
| task = list(kwargs.values())[0] | |
| kwargs = {} | |
| return original_call(task, **kwargs) | |
| agent.__call__ = _call | |
| return agent | |
| research_agent = _adapt_managed_agent(research_agent) | |
| local_info_agent = _adapt_managed_agent(local_info_agent) | |
| # Manager Agent: Coordinates specialists and synthesizes the final response | |
| manager_tools = [final_answer] | |
| if image_generation_tool is not None: | |
| manager_tools.append(image_generation_tool) | |
| manager_agent = CodeAgent( | |
| model=model, | |
| tools=manager_tools, | |
| managed_agents=[research_agent, local_info_agent], | |
| max_steps=8, | |
| verbosity_level=1, | |
| planning_interval=None, | |
| name="ArkAssistant_Manager", | |
| description="A senior manager agent that breaks down user tasks and delegates them to specialized team members (web researcher and local info specialist).", | |
| prompt_templates=prompt_templates | |
| ) | |
| # Main agent interface for Gradio | |
| agent = manager_agent | |
| if __name__ == "__main__": | |
| GradioUI(agent).launch() | |