#%%writefile HFHub.py

#utility file

from typing import List, Set
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] = {}
loaded_tools: Set[Tool] = set()  

#kueep separate list, maybe should only refer to agent tools merged list ? 
# def all_tools_names2() -> list[Tool]:
#     all_tools = []
#     for collection in tool_Collections.values():
#       if isinstance(collection, ToolCollection):
#         for tool in collection.tools:
#             all_tools.append(tool.name) #new we return key, not the tool
#       else:
#         for tool in collection: # base tools : only keys str
#             all_tools.append(tool)

#     return all_tools

def filter_tools(tools):
    base_tools_names = ['web search']
    filtered_tools = [tool for tool in tools if tool 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: #or overwite ?
        tool_collection = ToolCollection(
            collection_slug=collection_slug,
            # token=publishtoken,
            trust_remote_code=True
        )

        # tool_Collections[collection_slug] = tool_collection

        #replace this by slug > tools names:
        tool_Collections[collection_slug] = [tool.name for tool in tool_collection.tools]

        for tool in tool_collection.tools:
            if agent.tools.get(tool.name) is None: #or overwrite ?

                agent.tools[tool.name]=tool
                loaded_tools.add(tool)

            else:
                agent.tools[tool.name]=tool



    return all_tools_names2()

def all_tools() -> List[Tool]:
    return list(loaded_tools)

def createAgent():
    agent = CodeAgent(
        tools=filter_tools(all_tools()),  
        model=HfApiModel(
            # Model configuration can be added here
        ),
        additional_authorized_imports=additional_authorized_imports,
        add_base_tools=True,
        planning_interval=None,
    )

    # Update the base tools collection
    for tool in agent.tools:  # Changed from agent.toolbox.tools.values() to agent.tools
        if tool not in all_tools():
            if "base tools" not in tool_Collections:
                tool_Collections["base tools"] = []

            tool_Collections["base tools"].append(tool)
            logging.debug(f"Base tool added: {tool}")  # Debug line


    return agent