#%%writefile HFHub.py

#utility file

from smolagents.default_tools import __all__ as dft_tools
from huggingface_hub import login, InferenceClient
import logging
from smolagents import ToolCollection
from smolagents import HfApiModel, CodeAgent, Tool, LiteLLMModel
from smolagents import ToolCollection
from smolagents import HfApiModel, CodeAgent, Tool

publishtoken = None;

tool_Collections: dict[str, ToolCollection] = {}


def all_tools() -> list[Tool]:
    all_tools = []
    for collection in tool_Collections.values():
      if isinstance(collection, ToolCollection):
        for tool in collection.tools:
            all_tools.append(tool)
      else:
        for tool in collection:
            all_tools.append(tool)

    return all_tools

def filter_tools(tools):
    base_tools_names = ['web search']
# dft_tools
    # logging.warning(f"Filtering tools, excluding: {exclude_tool}")  # Log the exclusion
    filtered_tools = [tool for tool in tools if tool.name not in  base_tools_names]
    logging.warning(f"Number of tools after filtering: {len(filtered_tools)}") # Log the result
    return filtered_tools

additional_authorized_imports=["smolagents","subprocess","typing","os", "inspect","open", "requests"]


litellm_api_keys = {
    'xai': '',
    'HF': '',
    'grok': '',
    'anthropic': '',
    'openAI': '',
    # ...
}

def get_litellm_api_key(key_name):
    return litellm_api_keys.get(key_name, '')  

def login_for_publication() :
    login(publishtoken)


def load_collection_from_space(agent : CodeAgent, collection_slug: str = "Mightypeacock/agent-tools-6777c9699c231b7a1e87fa31" ) -> ToolCollection:

    if collection_slug not in tool_Collections:
        tool_collection = ToolCollection(
            collection_slug=collection_slug,
            # token=publishtoken,
            trust_remote_code=True
        )
        tool_Collections[collection_slug] = tool_collection

        for tool in tool_collection.tools:
            if agent.toolbox.tools.get(tool.name) is None:
              agent.toolbox.add_tool(tool)
            else:
              agent.toolbox.update_tool(tool)

    return all_tools()


def createAgent():

    agent = CodeAgent(
        tools=filter_tools(all_tools()),
        # all_tools(),#[tool for col in tool_Collections.values() for tool in col.tools],  # need to flatmap those probably
        # tools=[],
        # model=HfApiModel("microsoft/phi-4"),
        # model=HfApiModel("Leonuraht/Phi-4"), #space

        # model = LiteLLMModel(model_id=
        #             # "gpt-4o"
        #             "anthropic/claude-3-5-sonnet-20240620"
        #             # xai/<any-model-on-xai> : grok-beta
        # ),

        model=HfApiModel(
        #     # "lmstudio-community/phi-4-GGUF" #working
            ),

        additional_authorized_imports=additional_authorized_imports,
        add_base_tools=True,
        planning_interval=None,
        # use_e2b_executor=True,
        verbose=True

        )

    for tool in agent.toolbox.tools.values():
        if tool not in all_tools():
            if "base tools" not in tool_Collections:
                tool_Collections["base tools"] = []

            tool_Collections["base tools"].append(tool)

    return agent