Skip to content

Extending VT.ai

This guide explains how to extend VT.ai with new capabilities. VT.ai is designed to be modular and extensible, allowing developers to add new features, models, and integrations.

Adding New Model Providers

VT.ai uses LiteLLM to abstract model providers, making it relatively straightforward to add support for new AI models.

Steps to Add a New Provider

  1. Update Provider Configuration

Edit vtai/utils/llm_providers_config.py to add your new provider:

# Add to the PROVIDERS dictionary
PROVIDERS = {
    # ...existing providers...
    "new_provider": {
        "name": "New Provider",
        "models": ["new-model-1", "new-model-2"],
        "env_var": "NEW_PROVIDER_API_KEY",
        "icon": "new_provider.png",  # Add icon to vtai/resources/
    }
}
  1. Add Provider Icon

Place the provider's icon in vtai/resources/ directory.

  1. Update Model Mappings

Update the model mappings in the same file to include your new models:

# Add to MODEL_PROVIDER_MAP
MODEL_PROVIDER_MAP.update({
    "new-model-1": "new_provider",
    "new-model-2": "new_provider",
})
  1. Implement Special Handling (if needed)

If your provider requires special handling, add custom logic to the appropriate utility files:

  • For conversation handling: vtai/utils/conversation_handlers.py
  • For media processing: vtai/utils/media_processors.py

Extending the Semantic Router

The semantic router classifies user queries to determine the best handler. You can add new intents to the router for specialized handling.

Steps to Add a New Intent

  1. Update Routing Layers

Edit vtai/router/layers.json to add your new intent:

{
  "intents": [
    // ...existing intents...
    {
      "name": "my_new_intent",
      "description": "Description of what this intent handles",
      "examples": [
        "Example query 1",
        "Example query 2",
        "Example query 3"
      ]
    }
  ]
}
  1. Train the Router

Run the router trainer to update the embeddings:

python -m vtai.router.trainer
  1. Add Handler Function

Create a handler function in vtai/utils/conversation_handlers.py:

async def handle_my_new_intent(message, messages, client, **kwargs):
    """
    Handle queries matching my_new_intent.

    Args:
        message: The user message
        messages: Message history
        client: LLM client
        **kwargs: Additional arguments
    """
    # Your handler implementation here
    # ...

    # Send response
    await cl.Message(content=response).send()
  1. Update Router Configuration

Add your handler to the routing configuration in the app initialization:

# In vtai/utils/config.py
route_layer = SemanticRouter(
    # ...existing configuration...
    routes=[
        # ...existing routes...
        Route(
            name="my_new_intent",
            description="Handles my new intent",
            handler=handle_my_new_intent,
        ),
    ]
)

Adding New Assistant Tools

You can extend VT.ai with new assistant tools for specialized capabilities.

Steps to Add a New Tool

  1. Define Tool Interface

Create a new tool definition in vtai/tools/ or extend an existing one:

# In vtai/tools/my_new_tool.py
from typing import Dict, Any

async def my_new_tool_function(args: Dict[str, Any]) -> Dict[str, Any]:
    """
    Implement a new tool capability.

    Args:
        args: Tool arguments

    Returns:
        Tool results
    """
    # Tool implementation
    # ...

    return {"result": "Output from the tool"}
  1. Register the Tool

Add your tool to the assistant configuration in vtai/utils/assistant_tools.py:

# Add to the tools list
ASSISTANT_TOOLS = [
    # ...existing tools...
    {
        "type": "function",
        "function": {
            "name": "my_new_tool",
            "description": "Description of what this tool does",
            "parameters": {
                "type": "object",
                "properties": {
                    "param1": {
                        "type": "string",
                        "description": "Description of param1"
                    },
                    # Add more parameters as needed
                },
                "required": ["param1"]
            }
        }
    }
]
  1. Implement Tool Processing

Add a tool processor in vtai/app.py to handle the tool execution:

# Add to the process_function_tool logic to handle your tool
if function_name == "my_new_tool":
    from tools.my_new_tool import my_new_tool_function
    result = await my_new_tool_function(function_args)
    return {
        "output": result,
        "tool_call_id": tool_call.id
    }

Extending Image Generation

VT.ai now uses GPT-Image-1 for image generation with enhanced capabilities and configuration options. Here's how to extend or customize the image generation functionality:

Customizing Image Generation Settings

You can modify the default settings for GPT-Image-1 by updating the configuration in vtai/utils/llm_providers_config.py:

# Example: Changing default image generation settings
DEFAULT_IMAGE_GEN_IMAGE_SIZE = "1536x1024"  # Default to landscape format
DEFAULT_IMAGE_GEN_BACKGROUND = "transparent"  # Default to transparent backgrounds
DEFAULT_IMAGE_GEN_OUTPUT_FORMAT = "png"  # Default to PNG format
DEFAULT_IMAGE_GEN_MODERATION = "auto"  # Default moderation level
DEFAULT_IMAGE_GEN_OUTPUT_COMPRESSION = 90  # Higher quality default

Extending Image Processing Features

The image generation process is handled in vtai/utils/media_processors.py. To extend this functionality:

  1. Add New Image Processing Features

You can add post-processing features for generated images:

async def process_generated_image(image_data, format="jpeg"):
    """
    Apply custom processing to generated images.

    Args:
        image_data: Raw image data
        format: Image format (jpeg, png, webp)

    Returns:
        Processed image data
    """
    from PIL import Image, ImageFilter
    import io

    # Convert bytes to PIL Image
    image = Image.open(io.BytesIO(image_data))

    # Apply custom processing
    image = image.filter(ImageFilter.SHARPEN)

    # Convert back to bytes
    buffer = io.BytesIO()
    image.save(buffer, format=format.upper())

    return buffer.getvalue()
  1. Modify Image Generation UI

To add custom UI elements for your new image settings, update the settings in vtai/utils/settings_builder.py:

# Add a new setting for image generation
settings_components.append(
    cl.Select(
        id="my_custom_image_setting",
        label="🖼️ My Custom Setting",
        values=["option1", "option2", "option3"],
        initial_value="option1"
    )
)
  1. Add Custom Image Metadata

You can modify how image metadata is displayed by updating the handle_trigger_async_image_gen function:

# Add custom metadata to the image display
metadata = {
    # ...existing metadata...
    "Custom Info": "Your custom information",
    "Processing": f"Applied {your_custom_process}"
}

Saving and Managing Generated Images

Generated images are now saved in an imgs directory with timestamped filenames. To extend this functionality:

  1. Custom Storage Locations

You can modify where images are stored:

# Example: Change image storage location or naming convention
timestamp = int(time.time())
img_dir = Path("custom_images_folder")
img_dir.mkdir(exist_ok=True)
img_path = img_dir / f"custom_prefix_{timestamp}.{format}"
  1. Image Organization Features

You could add features to organize images by prompt or category:

# Organize images by category derived from prompt
def get_category_from_prompt(prompt):
    # Simple category extraction
    if "landscape" in prompt.lower():
        return "landscapes"
    elif "portrait" in prompt.lower():
        return "portraits"
    return "miscellaneous"

category = get_category_from_prompt(query)
img_dir = Path(f"imgs/{category}")
img_dir.mkdir(exist_ok=True, parents=True)

Creating Custom UI Components

VT.ai uses Chainlit for the web interface. You can extend the UI with custom components.

Steps to Add Custom UI Elements

  1. Define Custom Component

Create a custom component in a new file in the vtai directory:

# In vtai/ui/custom_component.py
import chainlit as cl

async def create_custom_component(data):
    """Create a custom UI component."""
    component = cl.Component(
        name="custom_component",
        props={"data": data},
        url="custom_component.js"  # Will need to create this
    )
    await component.send()
  1. Add Client-Side Implementation

If needed, create JavaScript for client-side functionality in .chainlit/public/custom_component.js.

  1. Integrate Component

Use your component in the appropriate part of the application, such as in conversation handlers.

Testing Extensions

To test your extensions:

  1. Run Unit Tests

Create tests for your new functionality in the tests directory:

# Run tests for your extension
pytest tests/unit/test_my_extension.py
  1. Manual Testing

Run VT.ai in development mode to test your changes interactively:

chainlit run vtai/app.py -w

Best Practices

When extending VT.ai, follow these best practices:

  1. Maintain Backward Compatibility: Ensure your changes don't break existing functionality
  2. Follow the Existing Pattern: Maintain the project's coding style and architecture
  3. Add Tests: Include tests for your new functionality
  4. Update Documentation: Document your extensions in the codebase and in the docs directory
  5. Handle Errors Gracefully: Implement proper error handling for your new features
  6. Optimize Performance: Consider the performance impact of your changes