agent.py

from smolagents import HfApiModel, CodeAgent, DuckDuckGoSearchTool, WikipediaSearchTool, Tool, LiteLLMModel
from langchain_community.tools.tavily_search import TavilySearchResults
from langchain_community.document_loaders import WikipediaLoader, ArxivLoader




class add(Tool):
    name = "add"
    description = """
        This tool adds two integers together and returns an integer.

        Args:
            a: first integer
            b: second integer
        """
    inputs = {
        "a":{
            "type":"integer",
            "description":"first integer"
        },
        "b":{
            "type":"integer",
            "description":"second integer"
        }
    }
    output_type = "integer"

    def forward(self, a: int, b: int):
        return a + b


class subtract(Tool):
    name = "subtract"
    description = """
        This tool subtracts two integers and returns an integer.

        Args:
            a: first integer
            b: second integer
        """
    inputs = {
        "a":{
            "type":"integer",
            "description":"first integer"
        },
        "b":{
            "type":"integer",
            "description":"second integer"
        }
    }
    output_type = "integer"

    def forward(self, a: int, b: int):
        return a - b

class multiply(Tool):
    name = "multiply"
    description = """
        This tool multiplies two integers and returns an integer.

        Args:
            a: first integer
            b: second integer
        """
    inputs = {
        "a":{
            "type":"integer",
            "description":"first integer"
        },
        "b":{
            "type":"integer",
            "description":"second integer"
        }
    }
    output_type = "integer"

    def forward(self, a: int, b: int):
        return a * b

class divide(Tool):
    name = "divide"
    description = """
        This tool divides two integers and returns an integer.

        Args:
            a: first integer
            b: second integer
        """
    inputs = {
        "a":{
            "type":"integer",
            "description":"first integer"
        },
        "b":{
            "type":"integer",
            "description":"second integer"
        }
    }
    output_type = "integer"

    def forward(self, a: int, b: int):
        if b == 0:
            raise ValueError("Cannot divide by zero.")
        return a / b

class modulo(Tool):
    name = "modulo"
    description = """
        This tool returns the modulus of two integers

        Args:
            a: first integer
            b: second integer
        """
    inputs = {
        "a":{
            "type":"integer",
            "description":"first integer"
        },
        "b":{
            "type":"integer",
            "description":"second integer"
        }
    }
    output_type = "integer"

    def forward(self, a: int, b: int):
        return a % b

class WikipediaSearchTool(tool):
    name = "wikipedia_search_tool"
    
    description = """
        Search Wikipedia for a query and return top 2 results.
    
        Args:
            query: the search query.
    """
    
    inputs = {
        "query":{
            "type":"string",
            "description":"the search query"
        }
    }
    output_type = "string"

    def forward(self, query: str) -> str:
        documents = WikipediaLoader(query=query, load_max_docs=2).load()
        condensed_docs = "\n\n---\n\n".join(
        [
            f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content}\n</Document>'
            for doc in documents
        ])
        return {"wikipedia_results": condensed_docs}

class TavilySearchTool(tool):
    name = "tavily_search_tool"
    
    description = """
        Search the internet using the browser Tavily and return the top 3 results.
    
        Args:
            query: the search query.
    """
    
    inputs = {
        "query":{
            "type":"string",
            "description":"the search query"
        }
    }
    output_type = "string"

    def forward(self, query: str) -> str:
        documents = TavilySearchResults(max_results=3).invoke(query=query)
        condensed_docs = "\n\n---\n\n".join(
        [
            f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content}\n</Document>'
            for doc in documents
        ])
        return {"web_search_results": condensed_docs}


@tool
def arvix_search(query: str) -> str:
    """Search Arxiv for a query and return maximum 3 result.
    
    Args:
        query: The search query."""
    search_docs = ArxivLoader(query=query, load_max_docs=3).load()
    formatted_search_docs = "\n\n---\n\n".join(
        [
            f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content[:1000]}\n</Document>'
            for doc in search_docs
        ])
    return {"arvix_results": formatted_search_docs}



tools=[
    add(),
    subtract(),
    multiply(),
    divide(),
    modulo(),
    DuckDuckGoSearchTool(),
    WikipediaSearchTool()
]

model = LiteLLMModel(
  model_id='ollama_chat/qwen2.5:14b'
)

def create_agent():
    agent = CodeAgent(
        model = model,
        tools = tools,
        max_steps=10,
        verbosity_level=2
    )

    return agent