Skip to main content

Custom Code

aiXplain allows developers to onboard custom code and utility models as tools within the platform. This guide walks through different techniques for integrating external functionalities, including:

  • Creating and onboarding utility models.
  • Using custom Python code directly within an agent.
  • Searching and managing onboarded models.

Creating and Onboarding a Utility Model​

Utility models allow developers to integrate external services as AI tools. These models can be created from Python functions and onboarded onto the aiXplain platform.

Example: Google Maps Directions Utility​

The following utility model fetches driving directions between two locations using the Google Maps API.

def main(start_location: str, end_location: str):
    """
    Get driving directions from start_location to end_location
    """
    import os
    try:
        import googlemaps
    except:
        os.system("pip install googlemaps")
        import googlemaps
    import json
    from datetime import datetime

    # Set up your Google Maps API key
    API_KEY = 'YOUR_GOOGLE_MAPS_API_KEY'  # Replace with your actual API key

    # Initialize the Google Maps client
    gmaps = googlemaps.Client(key=API_KEY)

    # Request directions
    now = datetime.now()
    directions_result = gmaps.directions(
        start_location,
        end_location,
        mode="driving",
        departure_time=now
    )

    try:
        directions = []
        for step in directions_result[0]['legs'][0]['steps']:
            directions.append({
                "instructions": step['html_instructions'],
                "duration": step['duration']['text'],
                "distance": step['distance']['text']
            })
        directions = json.dumps(directions)
    except:
        directions = "No directions found. Please check your input."
    
    return directions

Onboarding the Utility Model​

Once the function is defined, you can onboard it as a utility model in aiXplain.

from aixplain.enums import DataType
from aixplain.factories import ModelFactory
from aixplain.modules.model.utility_model import UtilityModelInput

utility_model = ModelFactory.create_utility_model(
    name="Google Maps V5",
    code=main
)

Now, this model can be executed like any other AI model:

response = utility_model.run({
    "start_location": "Riyadh Marriott Hotel",
    "end_location": "Sushi Yoshi Riyadh",
})
response
Show output

Using Custom Code Directly​

Alternatively, you can embed Python code directly into an agent instead of onboarding it as a separate utility model.

from aixplain.factories import AgentFactory
from aixplain.enums import Function, Supplier

ROLE = """You are a chatbot designed to assist tourists in the MENA region.

Your primary goal is to answer questions related to hotels, attractions, transportation, and any other tourism-related inquiries in an ethical and assertive manner.

If you are unable to provide an answer or require more information, please ask the user for further details or let them know that you are unable to provide the information they seek."""

agent = AgentFactory.create(
    name="MENA Tourist Assistant",
    description=ROLE,
    tools=[
        AgentFactory.create_model_tool(model="6736411cf127849667606689"), # Tavily search
        AgentFactory.create_model_tool(model="66f6ac496eb563510e3503d1"), # Google Places API
        AgentFactory.create_model_tool(model=utility_model.id), # Google Maps
        AgentFactory.create_model_tool(function=Function.TRANSLATION, supplier=Supplier.MICROSOFT), # Translation
    ]
)

This approach eliminates the need to onboard the function as a separate utility model and allows direct execution within the agent.

response = agent.run("What are good restaurants in Riyadh to eat sushi?")
print(response["data"]["output"])
Show output

Creating and Managing Utility Assets​

If aiXplain does not have a tool you need, you can create one as a Utility Model.

Example: Wikipedia Search Utility​

The following function queries Wikipedia and returns a summary.

def main(search_term: str, sentences: int):
    """
    Searches Wikipedia for a given term and returns a summary.

    Args:
        search_term (str): The term to search for on Wikipedia.
        sentences (int): Number of sentences to include in the summary.

    Returns:
        str: A summary of the Wikipedia article, or an error message if no results are found.
    """
    import os
    try:
        import wikipedia
    except:
        os.system("pip install wikipedia")
        import wikipedia
    
    try:
        search_results = wikipedia.search(search_term)
        search_results = wikipedia.summary(search_results[0], sentences=sentences)
    except:
        search_results = "No search results found. Please try another search term."
    
    return search_results

Onboarding the Wikipedia Utility​

from aixplain.factories import ModelFactory

wiki_model = ModelFactory.create_utility_model(
    name="Wikipedia Search",
    code=main,
)

To execute the model:

response = wiki_model.run({
    "search_term": "Friendship paradox",
    "sentences": 2
})
print(response)

Once a utility model is onboarded, it can be used as a tool in aiXplain agents, allowing seamless integration into AI-driven workflows.