Model Upload Examples
Learn how to upload and customize your own models on the Clarifai platform
This section provides examples that guide you through uploading custom models to Clarifai. You’ll learn how to adapt pre-built examples, configure model settings, and prepare your project for deployment on the platform.
tip
All the examples here are available in the Clarifai Runners Examples GitHub repository. You can use the Clarifai CLI to download a model template. For example:
clarifai model init --github-url https://github.com/Clarifai/runners-examples/tree/main/local-runners/ollama-model-upload
To use these examples, create a project directory and organize your files as shown below. This structure is required for successfully uploading models to the Clarifai platform:
your_model_directory/
├── 1/
│ └── model.py
├── requirements.txt
└── config.yaml
your_model_directory/– Root directory containing all files related to your custom model.1/– A subdirectory (named1) that contains the main model file.
model.py– Defines your model logic, including loading the model and handling inference.requirements.txt– Lists all Python dependencies required to run your model.config.yaml– Includes metadata and configuration details for building the model, such as compute resources and environment settings.
Hello World
1/model.py
from clarifai.runners.models.model_class import ModelClass
from clarifai.runners.utils.data_utils import Param
from typing import Iterator
import random
import string
class MyModel(ModelClass):
"""This is a model that does some string manipulation."""
def load_model(self):
"""Nothing to load for this model."""
@ModelClass.method
def predict(self, prompt: str, number_of_letters: int = Param(default=3, description="number of letters to add")) -> str:
"""Function to append some string information"""
return new_str(prompt, number_of_letters)
@ModelClass.method
def generate(self, prompt: str = "", number_of_letters: int = Param(default=3, description="number of letters to add")) -> Iterator[str]:
"""Example yielding a whole batch of streamed stuff back."""
for i in range(10): # fake something iterating generating 10 times.
yield new_str(str(i) + "-" + prompt, number_of_letters)
@ModelClass.method
def s(self, input_iterator: Iterator[str], number_of_letters: int = Param(default=3, description="number of letters to add")) -> Iterator[str]:
"""Example yielding getting an iterator and yielding back results."""
for i, inp in enumerate(input_iterator):
yield new_str(inp, number_of_letters)
def new_str(input_str: str, number_of_letters: int = 3) -> str:
"""Append a dash and random letters to the input string."""
random_letters = ''.join(random.choices(string.ascii_letters, k=number_of_letters))
return f"{input_str}-{random_letters}\n"
def test_predict() -> None:
"""Test the predict method of MyModel by printing its output."""
model = MyModel()
model.load_model()
print("Testing predict method:")
output = model.predict("TestPredict", number_of_letters=5)
print(output, end="\n")
def test_generate() -> None:
"""Test the generate method of MyModel by printing its outputs."""
model = MyModel()
model.load_model()
print("Testing generate method:")
for output in model.generate("Test", number_of_letters=5):
print(output, end="\n")
if __name__ == "__main__":
test_predict()
test_generate()
requirements.txt
clarifai
config.yaml
model:
id: "my-model-id"
user_id: "my-user-id"
app_id: "my-app-id"
model_type_id: "text-to-text"
build_info:
python_version: "3.12"
inference_compute_info:
cpu_limit: 50m
cpu_memory: 250Mi
num_accelerators: 0
NSFW Image Classifier
1/model.py
import os
from typing import List, Iterator
# Third-party imports
import cv2
import torch
from PIL import Image as PILImage
from transformers import AutoModelForImageClassification, ViTImageProcessor
# Clarifai imports
from clarifai.runners.models.visual_classifier_class import VisualClassifierClass
from clarifai.runners.utils.data_types import Concept, Image, Video
from clarifai.runners.models.model_builder import ModelBuilder
from clarifai.utils.logging import logger
class ImageClassifierModel(VisualClassifierClass):
"""A custom runner that classifies images and outputs concepts."""
def load_model(self):
"""Load the model and processor."""
model_path = os.path.dirname(os.path.dirname(__file__))
builder = ModelBuilder(model_path, download_validation_only=True)
checkpoints = builder.download_checkpoints(stage="runtime")
self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
logger.info(f"Running on device: {self.device}")
self.model = AutoModelForImageClassification.from_pretrained(checkpoints,).to(self.device)
self.model_labels = self.model.config.id2label
self.processor = ViTImageProcessor.from_pretrained(checkpoints)
logger.info("Done loading!")
@VisualClassifierClass.method
def predict(self, image: Image) -> List[Concept]:
"""Predict concepts for a list of images."""
pil_image = VisualClassifierClass.preprocess_image(image.bytes)
inputs = self.processor(images=pil_image, return_tensors="pt")
inputs = {name: tensor.to(self.device) for name, tensor in inputs.items()}
with torch.no_grad():
logits = self.model(**inputs).logits
outputs = VisualClassifierClass.process_concepts(logits, self.model_labels)
return outputs[0]
@VisualClassifierClass.method
def generate(self, video: Video) -> Iterator[List[Concept]]:
"""Generate concepts for frames extracted from a video."""
video_bytes = video.bytes
frame_generator = VisualClassifierClass.video_to_frames(video_bytes)
for frame in frame_generator:
image = VisualClassifierClass.preprocess_image(frame.image.bytes)
inputs = self.processor(images=image, return_tensors="pt")
inputs = {name: tensor.to(self.device) for name, tensor in inputs.items()}
with torch.no_grad():
logits = self.model(**inputs).logits
outputs = VisualClassifierClass.process_concepts(logits, self.model_labels) # Yield concepts for each frame
yield outputs[0]
@VisualClassifierClass.method
def stream_image(self, image_stream: Iterator[Image]) -> Iterator[List[Concept]]:
"""Stream process image inputs."""
for image in image_stream:
result = self.predict(image)
yield result
@VisualClassifierClass.method
def stream_video(self, video_stream: Iterator[Video]) -> Iterator[List[Concept]]:
"""Stream process video inputs."""
for video in video_stream:
for frame_result in self.generate(video):
yield frame_result
def test(self):
"""Test the model functionality."""
import requests # Import moved here as it's only used for testing
# Test configuration
TEST_URLS = {
"images": [
"https://samples.clarifai.com/metro-north.jpg",
"https://samples.clarifai.com/dog.tiff"
],
"video": "https://samples.clarifai.com/beer.mp4"
}
def get_test_data(url):
return Image(bytes=requests.get(url).content)
def get_test_video():
return Video(bytes=requests.get(TEST_URLS["video"]).content)
def run_test(name, test_fn):
logger.info(f"\nTesting {name}...")
try:
test_fn()
logger.info(f"{name} test completed successfully")
except Exception as e:
logger.error(f"Error in {name} test: {e}")
# Test predict
def test_predict():
result = self.predict(get_test_data(TEST_URLS["images"][0]))
logger.info(f"Predict result: {result}")
# Test generate
def test_generate():
for classifications in self.generate(get_test_video()):
logger.info(f"First frame classifications: {classifications}")
break
# Test stream
def test_stream():
# Split into two separate test functions for clarity
def test_stream_image():
images = [get_test_data(url) for url in TEST_URLS["images"]]
for result in self.stream_image(iter(images)):
logger.info(f"Image stream result: {result}")
def test_stream_video():
for result in self.stream_video(iter([get_test_video()])):
logger.info(f"Video stream result: {result}")
break # Just test first frame
logger.info("\nTesting image streaming...")
test_stream_image()
logger.info("\nTesting video streaming...")
test_stream_video()
# Run all tests
for test_name, test_fn in [
("predict", test_predict),
("generate", test_generate),
("stream", test_stream)
]:
run_test(test_name, test_fn)
requirements.txt
torch==2.6.0
transformers>=4.51.1
pillow==10.4.0
requests==2.32.3
timm==1.0.12
opencv-python-headless==4.10.0.84
clarifai>=11.5.0,<12.0.0
config.yaml
# This is the sample config file for the image-classification model.
model:
id: "nsfw_image_detection"
user_id: "user_id"
app_id: "app_id"
model_type_id: "visual-detector"
build_info:
python_version: "3.11"
inference_compute_info:
cpu_limit: "2"
cpu_memory: "2Gi"
num_accelerators: 1
accelerator_type: ["NVIDIA-*"]
accelerator_memory: "3Gi"
checkpoints:
type: "huggingface"
repo_id: "Falconsai/nsfw_image_detection"
DETR ResNet Image Detector
1/model.py
# Standard library imports
import os
from typing import List, Dict, Any, Iterator
# Third-party imports
import cv2
import torch
from PIL import Image as PILImage
from transformers import DetrForObjectDetection, DetrImageProcessor
# Clarifai imports
from clarifai.runners.models.model_builder import ModelBuilder
from clarifai.runners.models.visual_detector_class import VisualDetectorClass
from clarifai.runners.utils.data_types import Image, Video, Region, Frame
from clarifai.utils.logging import logger
def detect_objects(
images: List[PILImage],
model: DetrForObjectDetection,
processor: DetrImageProcessor,
device: str
) -> Dict[str, Any]:
"""Process images through the DETR model to detect objects.
Args:
images: List of preprocessed images
model: DETR model instance
processor: Image processor for DETR
device: Computation device (CPU/GPU)
Returns:
Detection results from the model
"""
model_inputs = processor(images=images, return_tensors="pt").to(device)
model_inputs = {name: tensor.to(device) for name, tensor in model_inputs.items()}
model_output = model(**model_inputs)
results = processor.post_process_object_detection(model_output)
return results
class MyRunner(VisualDetectorClass):
"""A custom runner for DETR object detection model that processes images and videos"""
def load_model(self):
"""Load the model here."""
model_path = os.path.dirname(os.path.dirname(__file__))
builder = ModelBuilder(model_path, download_validation_only=True)
checkpoint_path = builder.download_checkpoints(stage="runtime")
self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
logger.info(f"Running on device: {self.device}")
self.model = DetrForObjectDetection.from_pretrained(checkpoint_path).to(self.device)
self.processor = DetrImageProcessor.from_pretrained(checkpoint_path)
self.model.eval()
self.threshold = 0.9
self.model_labels = self.model.config.id2label
logger.info("Done loading!")
@VisualDetectorClass.method
def predict(self, image: Image) -> List[Region]:
"""Process a single image and return detected objects."""
image_bytes = image.bytes
image = VisualDetectorClass.preprocess_image(image_bytes)
with torch.no_grad():
results = detect_objects([image], self.model, self.processor, self.device)
outputs = VisualDetectorClass.process_detections(results, self.threshold, self.model_labels)
return outputs[0] # Return detections for single image
@VisualDetectorClass.method
def generate(self, video: Video) -> Iterator[Frame]:
"""Process video frames and yield detected objects for each frame."""
frame_generator = VisualDetectorClass.video_to_frames(video.bytes)
for frame in frame_generator:
with torch.no_grad():
image = VisualDetectorClass.preprocess_image(frame.image.bytes)
results = detect_objects([image], self.model, self.processor, self.device)
outputs = VisualDetectorClass.process_detections(results, self.threshold, self.model_labels)
frame.regions = outputs[0] # Assign detections to the frame
yield frame # Yield the frame with detections
@VisualDetectorClass.method
def stream_image(self, image_stream: Iterator[Image]) -> Iterator[List[Region]]:
"""Stream process image inputs."""
for image in image_stream:
result = self.predict(image)
yield result
@VisualDetectorClass.method
def stream_video(self, video_stream: Iterator[Video]) -> Iterator[Frame]:
"""Stream process video inputs."""
for video in video_stream:
for frame_result in self.generate(video):
yield frame_result
def test(self):
"""Test the model functionality."""
import requests # Import moved here as it's only used for testing
# Test configuration
TEST_URLS = {
"images": [
"https://samples.clarifai.com/metro-north.jpg",
"https://samples.clarifai.com/dog.tiff"
],
"video": "https://samples.clarifai.com/beer.mp4"
}
def get_test_data(url):
return Image(bytes=requests.get(url).content)
def get_test_video():
return Video(bytes=requests.get(TEST_URLS["video"]).content)
def run_test(name, test_fn):
logger.info(f"\nTesting {name}...")
try:
test_fn()
logger.info(f"{name} test completed successfully")
except Exception as e:
logger.error(f"Error in {name} test: {e}")
# Test predict
def test_predict():
result = self.predict(get_test_data(TEST_URLS["images"][0]))
logger.info(f"Predict result: {result}")
# Test generate
def test_generate():
for detections in self.generate(get_test_video()):
logger.info(f"First frame detections: {detections}")
break
# Test stream
def test_stream():
# Split into two separate test functions for clarity
def test_stream_image():
images = [get_test_data(url) for url in TEST_URLS["images"]]
for result in self.stream_image(iter(images)):
logger.info(f"Image stream result: {result}")
def test_stream_video():
for result in self.stream_video(iter([get_test_video()])):
logger.info(f"Video stream result: {result}")
break # Just test first frame
logger.info("\nTesting image streaming...")
test_stream_image()
logger.info("\nTesting video streaming...")
test_stream_video()
# Run all tests
for test_name, test_fn in [
("predict", test_predict),
("generate", test_generate),
("stream", test_stream)
]:
run_test(test_name, test_fn)
requirements.txt
torch==2.6.0
transformers>=4.51.1
pillow==10.4.0
requests==2.32.3
timm==1.0.12
opencv-python-headless==4.10.0.84
clarifai>=11.4.10,<12.0.0
config.yaml
# This is the sample config file for the image-detection model.
model:
id: "detr-resnet-50"
user_id: "user_id"
app_id: "app_id"
model_type_id: "visual-detector"
build_info:
python_version: "3.11"
inference_compute_info:
cpu_limit: "4"
cpu_memory: "2Gi"
num_accelerators: 1
accelerator_type: ["NVIDIA-*"]
accelerator_memory: "5Gi"
checkpoints:
type: "huggingface"
repo_id: "facebook/detr-resnet-50"
hf_token: "hf_token"
Image Segmenter
Mask2Former ADE
1/model.py
import os
from typing import Iterator, List, Tuple
from clarifai.runners.models.visual_detector_class import VisualDetectorClass
from clarifai.runners.utils import data_types as dt
from clarifai.utils.logging import logger
import yaml
import torch
from PIL import Image as PILImage
from transformers import AutoImageProcessor, Mask2FormerForUniversalSegmentation
ROOT = os.path.dirname(__file__)
class MyRunner(VisualDetectorClass):
def _load_concepts(self, config_path, name, model_path):
with open(config_path, "r") as f:
data = yaml.safe_load(f)
# Map Clarifai concept name to id and reverse
self.conceptid2name = {each["id"] : each["name"] for each in data.get("concepts", [])}
self.conceptname2id = {each["name"] : each["id"] for each in data.get("concepts", [])}
def load_model(self):
"""Load the model here."""
checkpoint_path = "facebook/mask2former-swin-tiny-ade-semantic"
self.device = 'cuda' #if torch.cuda.is_available() else 'cpu'
logger.info(f"Running on device: {self.device}")
self.model = Mask2FormerForUniversalSegmentation.from_pretrained(
checkpoint_path, trust_remote_code=True).to(self.device)
self.processor = AutoImageProcessor.from_pretrained(checkpoint_path, trust_remote_code=True)
self.model.eval()
# Load clarifai concept
config_path = os.path.join(ROOT, "../config.yaml")
self._load_concepts(config_path, "mask2former-ade", checkpoint_path)
logger.info("Done loading!")
@VisualDetectorClass.method
def model_predict(
self,
images: List[PILImage.Image]
) -> List[List[dt.Region]]:
inputs = self.processor(images=images, return_tensors="pt").to(self.device)
with torch.no_grad():
outputs = self.model(**inputs)
target_sizes = [image.size[::-1] for image in images]
results = self.processor.post_process_semantic_segmentation(outputs, target_sizes=target_sizes)
outputs = []
for i, all_masks_tensor in enumerate(results):
masks = []
for clss_id in all_masks_tensor.unique().tolist():
label = self.model.config.id2label[clss_id]
mask = torch.zeros_like(all_masks_tensor)
mask[all_masks_tensor == clss_id] = 255
mask = mask.cpu().numpy() if self.device == "cuda" else mask.numpy()
mask = PILImage.fromarray(mask.astype("uint8"))
region = dt.Region(
mask=mask,
concepts=[dt.Concept(id=self.conceptname2id[label], name=label)]
)
masks.append(region)
outputs.append(masks)
return outputs
@VisualDetectorClass.method
def predict(self, image: dt.Image) -> List[dt.Region]:
return self.model_predict([image.to_pil()])[0]
@VisualDetectorClass.method
def stream(self, images: Iterator[dt.Image]) -> Iterator[dt.Region]:
for each in images:
yield self.predict(image=each)
@VisualDetectorClass.method
def generate(self, video: dt.Video) -> Iterator[dt.Region]:
for frame in self.video_to_frames(video.bytes):
yield self.predict(image=frame.image)
def test(self):
import requests
image = dt.Image(bytes=requests.get("https://samples.clarifai.com/metro-north.jpg").content)
video = dt.Video(bytes=requests.get("https://samples.clarifai.com/beer.mp4").content)
logger.info("# -------- Test predict/detect -------------")
logger.info(f"{self.predict(image=image)}")
logger.info("# -------- Test generate -------------")
n=5
for i, each in enumerate(self.generate(video=video)):
print(each)
if i > n:
break
logger.info("# -------- Test stream -------------")
def iteration():
for each in [image]*10:
yield each
for i, each in enumerate(self.stream(images=iteration())):
print(each)
if i > n:
break
requirements.txt
torch==2.6.0
tokenizers>=0.19.1
transformers>=4.44.2
pillow>=10.4.0
requests>=2.32.3
timm
opencv-python-headless==4.10.0.84
numpy
aiohttp
scipy>=1.10.1
config.yaml
model:
id: facebook_mask2former-swin-tiny-ade-semantic
user_id: a
app_id: b
model_type_id: visual-segmenter
build_info:
python_version: '3.12'
inference_compute_info:
cpu_limit: '2'
cpu_memory: 8Gi
num_accelerators: 1
accelerator_type:
- NVIDIA-*
accelerator_memory: 21Gi
concepts:
- id: id-mask2former-ade-0
name: wall
- id: id-mask2former-ade-1
name: building
- id: id-mask2former-ade-2
name: sky
- id: id-mask2former-ade-3
name: floor
- id: id-mask2former-ade-4
name: tree
- id: id-mask2former-ade-5
name: ceiling
- id: id-mask2former-ade-6
name: road
- id: id-mask2former-ade-7
name: 'bed '
- id: id-mask2former-ade-8
name: windowpane
- id: id-mask2former-ade-9
name: grass
- id: id-mask2former-ade-10
name: cabinet
- id: id-mask2former-ade-11
name: sidewalk
- id: id-mask2former-ade-12
name: person
- id: id-mask2former-ade-13
name: earth
- id: id-mask2former-ade-14
name: door
- id: id-mask2former-ade-15
name: table
- id: id-mask2former-ade-16
name: mountain
- id: id-mask2former-ade-17
name: plant
- id: id-mask2former-ade-18
name: curtain
- id: id-mask2former-ade-19
name: chair
- id: id-mask2former-ade-20
name: car
- id: id-mask2former-ade-21
name: water
- id: id-mask2former-ade-22
name: painting
- id: id-mask2former-ade-23
name: sofa
- id: id-mask2former-ade-24
name: shelf
- id: id-mask2former-ade-25
name: house
- id: id-mask2former-ade-26
name: sea
- id: id-mask2former-ade-27
name: mirror
- id: id-mask2former-ade-28
name: rug
- id: id-mask2former-ade-29
name: field
- id: id-mask2former-ade-30
name: armchair
- id: id-mask2former-ade-31
name: seat
- id: id-mask2former-ade-32
name: fence
- id: id-mask2former-ade-33
name: desk
- id: id-mask2former-ade-34
name: rock
- id: id-mask2former-ade-35
name: wardrobe
- id: id-mask2former-ade-36
name: lamp
- id: id-mask2former-ade-37
name: bathtub
- id: id-mask2former-ade-38
name: railing
- id: id-mask2former-ade-39
name: cushion
- id: id-mask2former-ade-40
name: base
- id: id-mask2former-ade-41
name: box
- id: id-mask2former-ade-42
name: column
- id: id-mask2former-ade-43
name: signboard
- id: id-mask2former-ade-44
name: chest of drawers
- id: id-mask2former-ade-45
name: counter
- id: id-mask2former-ade-46
name: sand
- id: id-mask2former-ade-47
name: sink
- id: id-mask2former-ade-48
name: skyscraper
- id: id-mask2former-ade-49
name: fireplace
- id: id-mask2former-ade-50
name: refrigerator
- id: id-mask2former-ade-51
name: grandstand
- id: id-mask2former-ade-52
name: path
- id: id-mask2former-ade-53
name: stairs
- id: id-mask2former-ade-54
name: runway
- id: id-mask2former-ade-55
name: case
- id: id-mask2former-ade-56
name: pool table
- id: id-mask2former-ade-57
name: pillow
- id: id-mask2former-ade-58
name: screen door
- id: id-mask2former-ade-59
name: stairway
- id: id-mask2former-ade-60
name: river
- id: id-mask2former-ade-61
name: bridge
- id: id-mask2former-ade-62
name: bookcase
- id: id-mask2former-ade-63
name: blind
- id: id-mask2former-ade-64
name: coffee table
- id: id-mask2former-ade-65
name: toilet
- id: id-mask2former-ade-66
name: flower
- id: id-mask2former-ade-67
name: book
- id: id-mask2former-ade-68
name: hill
- id: id-mask2former-ade-69
name: bench
- id: id-mask2former-ade-70
name: countertop
- id: id-mask2former-ade-71
name: stove
- id: id-mask2former-ade-72
name: palm
- id: id-mask2former-ade-73
name: kitchen island
- id: id-mask2former-ade-74
name: computer
- id: id-mask2former-ade-75
name: swivel chair
- id: id-mask2former-ade-76
name: boat
- id: id-mask2former-ade-77
name: bar
- id: id-mask2former-ade-78
name: arcade machine
- id: id-mask2former-ade-79
name: hovel
- id: id-mask2former-ade-80
name: bus
- id: id-mask2former-ade-81
name: towel
- id: id-mask2former-ade-82
name: light
- id: id-mask2former-ade-83
name: truck
- id: id-mask2former-ade-84
name: tower
- id: id-mask2former-ade-85
name: chandelier
- id: id-mask2former-ade-86
name: awning
- id: id-mask2former-ade-87
name: streetlight
- id: id-mask2former-ade-88
name: booth
- id: id-mask2former-ade-89
name: television receiver
- id: id-mask2former-ade-90
name: airplane
- id: id-mask2former-ade-91
name: dirt track
- id: id-mask2former-ade-92
name: apparel
- id: id-mask2former-ade-93
name: pole
- id: id-mask2former-ade-94
name: land
- id: id-mask2former-ade-95
name: bannister
- id: id-mask2former-ade-96
name: escalator
- id: id-mask2former-ade-97
name: ottoman
- id: id-mask2former-ade-98
name: bottle
- id: id-mask2former-ade-99
name: buffet
- id: id-mask2former-ade-100
name: poster
- id: id-mask2former-ade-101
name: stage
- id: id-mask2former-ade-102
name: van
- id: id-mask2former-ade-103
name: ship
- id: id-mask2former-ade-104
name: fountain
- id: id-mask2former-ade-105
name: conveyer belt
- id: id-mask2former-ade-106
name: canopy
- id: id-mask2former-ade-107
name: washer
- id: id-mask2former-ade-108
name: plaything
- id: id-mask2former-ade-109
name: swimming pool
- id: id-mask2former-ade-110
name: stool
- id: id-mask2former-ade-111
name: barrel
- id: id-mask2former-ade-112
name: basket
- id: id-mask2former-ade-113
name: waterfall
- id: id-mask2former-ade-114
name: tent
- id: id-mask2former-ade-115
name: bag
- id: id-mask2former-ade-116
name: minibike
- id: id-mask2former-ade-117
name: cradle
- id: id-mask2former-ade-118
name: oven
- id: id-mask2former-ade-119
name: ball
- id: id-mask2former-ade-120
name: food
- id: id-mask2former-ade-121
name: step
- id: id-mask2former-ade-122
name: tank
- id: id-mask2former-ade-123
name: trade name
- id: id-mask2former-ade-124
name: microwave
- id: id-mask2former-ade-125
name: pot
- id: id-mask2former-ade-126
name: animal
- id: id-mask2former-ade-127
name: bicycle
- id: id-mask2former-ade-128
name: lake
- id: id-mask2former-ade-129
name: dishwasher
- id: id-mask2former-ade-130
name: screen
- id: id-mask2former-ade-131
name: blanket
- id: id-mask2former-ade-132
name: sculpture
- id: id-mask2former-ade-133
name: hood
- id: id-mask2former-ade-134
name: sconce
- id: id-mask2former-ade-135
name: vase
- id: id-mask2former-ade-136
name: traffic light
- id: id-mask2former-ade-137
name: tray
- id: id-mask2former-ade-138
name: ashcan
- id: id-mask2former-ade-139
name: fan
- id: id-mask2former-ade-140
name: pier
- id: id-mask2former-ade-141
name: crt screen
- id: id-mask2former-ade-142
name: plate
- id: id-mask2former-ade-143
name: monitor
- id: id-mask2former-ade-144
name: bulletin board
- id: id-mask2former-ade-145
name: shower
- id: id-mask2former-ade-146
name: radiator
- id: id-mask2former-ade-147
name: glass
- id: id-mask2former-ade-148
name: clock
- id: id-mask2former-ade-149
name: flag
Image-Text-to-Image
Stable Diffusion 2 Depth
1/model.py
from typing import List
import os
from clarifai.runners.models.model_builder import ModelBuilder
from clarifai.runners.models.openai_class import OpenAIModelClass
from clarifai.runners.utils.data_types import Image
from clarifai.runners.utils.data_utils import Param
from clarifai.utils.logging import logger
from diffusers import StableDiffusionDepth2ImgPipeline
from diffusers.utils import load_image
import torch
class StableDiffusion(OpenAIModelClass):
"""
A Model that integrates with the Clarifai platform and uses the FluxFillPipeline for image inpainting.
"""
client = True # This will be set in load_model method
model = True # This will be set in load_model method
def load_model(self):
"""Load the model here."""
model_path = os.path.dirname(os.path.dirname(__file__))
builder = ModelBuilder(model_path, download_validation_only=True)
hf_token = builder.config["checkpoints"]["hf_token"]
self.client = StableDiffusionDepth2ImgPipeline.from_pretrained("stabilityai/stable-diffusion-2-depth",
torch_dtype=torch.float16,
use_safetensors=True,
token=hf_token)
self.client.to("cuda" if torch.cuda.is_available() else "cpu")
logger.info("stable-diffusion model loaded successfully.")
@OpenAIModelClass.method
def predict(self,
prompt: str,
image: Image = None,
negative_prompt: str = None,
mask: Image = None,
strength: float = 0.8
) -> Image:
"""
Predict method that uses the FluxFillPipeline to inpaint images based on the provided prompt.
"""
if image:
if image.url:
image = load_image(image.url)
elif image.bytes:
image=image.to_pil()
response = self.client(prompt=prompt,
image=image,
negative_prompt=negative_prompt,
strength=strength).images[0]
return Image.from_pil(pil_image = response)
requirements.txt
tokenizers==0.21.0
transformers>=4.48
diffusers==0.32.2
accelerate==1.2.0
optimum==1.23.3
xformers
einops==0.8.0
requests==2.32.3
numpy>2.0
torch==2.5.1
clarifai
clarifai-protocol
config.yaml
model:
id: "stable-diffusion-2-depth"
user_id: "user_id"
app_id: "app_id"
model_type_id: "multimodal-to-text"
build_info:
python_version: '3.11'
inference_compute_info:
cpu_limit: '3'
cpu_memory: 15Gi
num_accelerators: 1
accelerator_type: ["NVIDIA-*"]
accelerator_memory: 6Gi
checkpoints:
type: huggingface
repo_id: "stabilityai/stable-diffusion-2-depth"
hf_token: "hf_token"
when: runtime
LLM
SmolLM2 1.7B Instruct (SGLang)
1/model.py
import os
import sys
sys.path.append(os.path.dirname(__file__))
from typing import Iterator, List
from clarifai.runners.models.model_builder import ModelBuilder
from clarifai.runners.models.openai_class import OpenAIModelClass
from clarifai.runners.utils.data_utils import Param
from clarifai.runners.utils.openai_convertor import build_openai_messages
from clarifai.utils.logging import logger
from openai import OpenAI
from openai_server_starter import OpenAI_APIServer
##################
class SglangModel(OpenAIModelClass):
"""
A custom runner that integrates with the Clarifai platform and uses Server inference
to process inputs, including text.
"""
client = True # This will be set in load_model method
model = True # This will be set in load_model method
def load_model(self):
"""Load the model here and start the server."""
os.path.join(os.path.dirname(__file__))
# Use downloaded checkpoints.
# Or if you intend to download checkpoint at runtime, set hf id instead. For example:
# checkpoints = "Qwen/Qwen2-7B-Instruct"
# server args were generated by `upload` module
server_args = {
'dtype': 'auto',
'kv_cache_dtype': 'auto',
'tp_size': 1,
'load_format': 'auto',
'context_length': None,
'device': 'cuda',
'port': 23333,
'host': '0.0.0.0',
'mem_fraction_static': 0.9,
'max_total_tokens': '8192',
'max_prefill_tokens': None,
'schedule_policy': 'fcfs',
'schedule_conservativeness': 1.0,
'checkpoints': 'runtime'}
# if checkpoints == "checkpoints" => assign to checkpoints var aka local checkpoints path
stage = server_args.get("checkpoints")
if stage in ["build", "runtime"]:
#checkpoints = os.path.join(os.path.dirname(__file__), "checkpoints")
config_path = os.path.dirname(os.path.dirname(__file__))
builder = ModelBuilder(config_path, download_validation_only=True)
checkpoints = builder.download_checkpoints(stage=stage)
server_args.update({"checkpoints": checkpoints})
if server_args.get("additional_list_args") == ['']:
server_args.pop("additional_list_args")
# Start server
# This line were generated by `upload` module
self.server = OpenAI_APIServer.from_sglang_backend(**server_args)
# Create client
self.client = OpenAI(
api_key="notset",
base_url=SglangModel.make_api_url(self.server.host, self.server.port))
self.model = self._get_model()
logger.info(f"OpenAI {self.model} model loaded successfully!")
def _get_model(self):
try:
return self.client.models.list().data[0].id
except Exception as e:
raise ConnectionError("Failed to retrieve model ID from API") from e
@staticmethod
def make_api_url(host: str, port: int, version: str = "v1") -> str:
return f"http://{host}:{port}/{version}"
@OpenAIModelClass.method
def predict(self,
prompt: str,
chat_history: List[dict] = None,
max_tokens: int = Param(default=512, description="The maximum number of tokens to generate. Shorter token lengths will provide faster performance.", ),
temperature: float = Param(default=0.7, description="A decimal number that determines the degree of randomness in the response", ),
top_p: float = Param(default=0.8, description="An alternative to sampling with temperature, where the model considers the results of the tokens with top_p probability mass.", )
) -> str:
"""This is the method that will be called when the runner is run. It takes in an input and
returns an output.
"""
openai_messages = build_openai_messages(prompt=prompt, messages=chat_history)
response = self.client.chat.completions.create(
model=self.model,
messages=openai_messages,
max_completion_tokens=max_tokens,
temperature=temperature,
top_p=top_p)
if response.usage and response.usage.prompt_tokens and response.usage.completion_tokens:
self.set_output_context(prompt_tokens=response.usage.prompt_tokens, completion_tokens=response.usage.completion_tokens)
return response.choices[0].message.content
@OpenAIModelClass.method
def generate(self,
prompt: str,
chat_history: List[dict] = None,
max_tokens: int = Param(default=512, description="The maximum number of tokens to generate. Shorter token lengths will provide faster performance.", ),
temperature: float = Param(default=0.7, description="A decimal number that determines the degree of randomness in the response", ),
top_p: float = Param(default=0.8, description="An alternative to sampling with temperature, where the model considers the results of the tokens with top_p probability mass.", )
) -> Iterator[str]:
"""Example yielding a whole batch of streamed stuff back."""
openai_messages = build_openai_messages(prompt=prompt, messages=chat_history)
for chunk in self.client.chat.completions.create(
model=self.model,
messages=openai_messages,
max_completion_tokens=max_tokens,
temperature=temperature,
top_p=top_p,
stream=True):
if chunk.choices:
text = (chunk.choices[0].delta.content
if (chunk and chunk.choices[0].delta.content) is not None else '')
yield text
# This method is needed to test the model with the test-locally CLI command.
def test(self):
"""Test the model here."""
try:
print("Testing predict...")
# Test predict
print(self.predict(prompt="Hello, how are you?",))
except Exception as e:
print("Error in predict", e)
try:
print("Testing generate...")
# Test generate
for each in self.generate(prompt="Hello, how are you?",):
print(each, end=" ")
except Exception as e:
print("Error in generate", e)
1/openai_server_starter.py
import os
import signal
import subprocess
import sys
import threading
from typing import List
import psutil
from clarifai.utils.logging import logger
PYTHON_EXEC = sys.executable
def kill_process_tree(parent_pid, include_parent: bool = True, skip_pid: int = None):
"""Kill the process and all its child processes."""
if parent_pid is None:
parent_pid = os.getpid()
include_parent = False
try:
itself = psutil.Process(parent_pid)
except psutil.NoSuchProcess:
return
children = itself.children(recursive=True)
for child in children:
if child.pid == skip_pid:
continue
try:
child.kill()
except psutil.NoSuchProcess:
pass
if include_parent:
try:
itself.kill()
# Sometime processes cannot be killed with SIGKILL (e.g, PID=1 launched by kubernetes),
# so we send an additional signal to kill them.
itself.send_signal(signal.SIGQUIT)
except psutil.NoSuchProcess:
pass
class OpenAI_APIServer:
def __init__(self, **kwargs):
self.server_started_event = threading.Event()
self.process = None
self.backend = None
self.server_thread = None
def __del__(self, *exc):
# This is important
# close the server when exit the program
self.close()
def close(self):
if self.process:
try:
kill_process_tree(self.process.pid)
except:
self.process.terminate()
if self.server_thread:
self.server_thread.join()
def wait_for_startup(self):
self.server_started_event.wait()
def validate_if_server_start(self, line: str):
line_lower = line.lower()
if self.backend in ["vllm", "sglang", "lmdeploy"]:
if self.backend == "vllm":
return "application startup complete" in line_lower or "vllm api server on" in line_lower
else:
return f" running on http://{self.host}:" in line.strip()
elif self.backend == "llamacpp":
return "waiting for new tasks" in line_lower
elif self.backend == "tgi":
return "Connected" in line.strip()
def _start_server(self, cmds):
try:
env = os.environ.copy()
env["VLLM_USAGE_SOURCE"] = "production-docker-image"
self.process = subprocess.Popen(
cmds,
stdout=subprocess.PIPE,
stderr=subprocess.STDOUT,
text=True,
)
for line in self.process.stdout:
logger.info("Server Log: " + line.strip())
if self.validate_if_server_start(line):
self.server_started_event.set()
# break
except Exception as e:
if self.process:
self.process.terminate()
raise RuntimeError(f"Failed to start Server server: {e}")
def start_server_thread(self, cmds: str):
try:
# Start the server in a separate thread
self.server_thread = threading.Thread(target=self._start_server, args=(cmds,), daemon=None)
self.server_thread.start()
# Wait for the server to start
self.wait_for_startup()
except Exception as e:
raise Exception(e)
@classmethod
def from_sglang_backend(
cls,
checkpoints,
dtype: str = "auto",
kv_cache_dtype: str = "auto",
tp_size: int = 1,
quantization: str = None,
load_format: str = "auto",
context_length: str = None,
device: str = "cuda",
port=23333,
host="0.0.0.0",
chat_template: str = None,
mem_fraction_static: float = 0.8,
max_running_requests: int = None,
max_total_tokens: int = None,
max_prefill_tokens: int = None,
schedule_policy: str = "fcfs",
schedule_conservativeness: float = 1.0,
cpu_offload_gb: int = 0,
additional_list_args: List[str] = [],
):
"""Start SGlang OpenAI compatible server.
Args:
checkpoints (str): model id or path.
dtype (str, optional): Dtype used for the model {"auto", "half", "float16", "bfloat16", "float", "float32"}. Defaults to "auto".
kv_cache_dtype (str, optional): Dtype of the kv cache, defaults to the dtype. Defaults to "auto".
tp_size (int, optional): The number of GPUs the model weights get sharded over. Mainly for saving memory rather than for high throughput. Defaults to 1.
quantization (str, optional): Quantization format {"awq","fp8","gptq","marlin","gptq_marlin","awq_marlin","bitsandbytes","gguf","modelopt","w8a8_int8"}. Defaults to None.
load_format (str, optional): The format of the model weights to load:\n* `auto`: will try to load the weights in the safetensors format and fall back to the pytorch bin format if safetensors format is not available.\n* `pt`: will load the weights in the pytorch bin format. \n* `safetensors`: will load the weights in the safetensors format. \n* `npcache`: will load the weights in pytorch format and store a numpy cache to speed up the loading. \n* `dummy`: will initialize the weights with random values, which is mainly for profiling.\n* `gguf`: will load the weights in the gguf format. \n* `bitsandbytes`: will load the weights using bitsandbytes quantization."\n* `layered`: loads weights layer by layer so that one can quantize a layer before loading another to make the peak memory envelope smaller.\n. Defaults to "auto".\n
context_length (str, optional): The model's maximum context length. Defaults to None (will use the value from the model's config.json instead). Defaults to None.
device (str, optional): The device type {"cuda", "xpu", "hpu", "cpu"}. Defaults to "cuda".
port (int, optional): Port number. Defaults to 23333.
host (str, optional): Host name. Defaults to "0.0.0.0".
chat_template (str, optional): The buliltin chat template name or the path of the chat template file. This is only used for OpenAI-compatible API server.. Defaults to None.
mem_fraction_static (float, optional): The fraction of the memory used for static allocation (model weights and KV cache memory pool). Use a smaller value if you see out-of-memory errors. Defaults to 0.8.
max_running_requests (int, optional): The maximum number of running requests.. Defaults to None.
max_total_tokens (int, optional): The maximum number of tokens in the memory pool. If not specified, it will be automatically calculated based on the memory usage fraction. This option is typically used for development and debugging purposes.. Defaults to None.
max_prefill_tokens (int, optional): The maximum number of tokens in a prefill batch. The real bound will be the maximum of this value and the model's maximum context length. Defaults to None.
schedule_policy (str, optional): The scheduling policy of the requests {"lpm", "random", "fcfs", "dfs-weight"}. Defaults to "fcfs".
schedule_conservativeness (float, optional): How conservative the schedule policy is. A larger value means more conservative scheduling. Use a larger value if you see requests being retracted frequently. Defaults to 1.0.
cpu_offload_gb (int, optional): How many GBs of RAM to reserve for CPU offloading. Defaults to 0.
additional_list_args (List[str], optional): additional args to run subprocess cmd e.g. ["--arg-name", "arg value"]. See more at [github](https://github.com/sgl-project/sglang/blob/1baa9e6cf90b30aaa7dae51c01baa25229e8f7d5/python/sglang/srt/server_args.py#L298). Defaults to [].
Returns:
_type_: _description_
"""
from sglang.utils import execute_shell_command, wait_for_server
cmds = [
PYTHON_EXEC, '-m', 'sglang.launch_server', '--model-path', checkpoints, '--dtype',
str(dtype), '--device',
str(device), '--kv-cache-dtype',
str(kv_cache_dtype), '--tp-size',
str(tp_size), '--load-format',
str(load_format), '--mem-fraction-static',
str(mem_fraction_static), '--schedule-policy',
str(schedule_policy), '--schedule-conservativeness',
str(schedule_conservativeness), '--port',
str(port), '--host',
host, "--trust-remote-code"
]
if chat_template:
cmds += ["--chat-template", chat_template]
if quantization:
cmds += [
'--quantization',
quantization,
]
if context_length:
cmds += [
'--context-length',
context_length,
]
if max_running_requests:
cmds += [
'--max-running-requests',
max_running_requests,
]
if max_total_tokens:
cmds += [
'--max-total-tokens',
max_total_tokens,
]
if max_prefill_tokens:
cmds += [
'--max-prefill-tokens',
max_prefill_tokens,
]
if additional_list_args:
cmds += additional_list_args
print("CMDS to run `sglang` server: ", " ".join(cmds), "\n")
_self = cls()
_self.host = host
_self.port = port
_self.backend = "sglang"
# _self.start_server_thread(cmds)
# new_path = os.environ["PATH"] + ":/sbin"
# _self.process = subprocess.Popen(cmds, text=True, stderr=subprocess.STDOUT, env={**os.environ, "PATH": new_path})
_self.process = execute_shell_command(" ".join(cmds))
logger.info("Waiting for " + f"http://{_self.host}:{_self.port}")
wait_for_server(f"http://{_self.host}:{_self.port}")
logger.info("Done")
return _self
requirements.txt
torch==2.6.0
tokenizers==0.21.1
accelerate==1.2.0
optimum==1.23.3
xformers
einops==0.8.0
packaging
ninja
qwen-vl-utils==0.0.8
timm==1.0.12
openai
clarifai>=11.5.0,<12.0.0
psutil
--extra-index-url https://flashinfer.ai/whl/cu124/torch2.4/
flashinfer
sglang[all]==0.4.6
transformers==4.51.1
config.yaml
# Config file for the Sglang runner
model:
id: "SmolLM2-1_7B-Instruct"
user_id: "USER_ID"
app_id: "APP_ID"
model_type_id: "text-to-text"
build_info:
python_version: "3.11"
inference_compute_info:
cpu_limit: "1"
cpu_memory: "6Gi"
num_accelerators: 1
accelerator_type: ["NVIDIA-*"]
accelerator_memory: "44Gi"
checkpoints:
type: "huggingface"
repo_id: "HuggingFaceTB/SmolLM2-1.7B-Instruct"
hf_token: "hf_token"
LLaMA 3.2 1B Instruct (Hugging Face)
1/model.py
from typing import List, Iterator
from threading import Thread
import os
import torch
from clarifai.runners.models.model_class import ModelClass
from clarifai.utils.logging import logger
from clarifai.runners.models.model_builder import ModelBuilder
from clarifai.runners.utils.openai_convertor import openai_response
from clarifai.runners.utils.data_utils import Param
from transformers import (AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer)
class MyModel(ModelClass):
"""A custom runner for llama-3.2-1b-instruct llm that integrates with the Clarifai platform"""
def load_model(self):
"""Load the model here."""
self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
logger.info(f"Running on device: {self.device}")
# Load checkpoints
model_path = os.path.dirname(os.path.dirname(__file__))
builder = ModelBuilder(model_path, download_validation_only=True)
self.checkpoints = builder.download_checkpoints(stage="runtime")
# Load model and tokenizer
self.tokenizer = AutoTokenizer.from_pretrained(self.checkpoints,)
self.tokenizer.pad_token = self.tokenizer.eos_token # Set pad token to eos token
self.model = AutoModelForCausalLM.from_pretrained(
self.checkpoints,
low_cpu_mem_usage=True,
device_map=self.device,
torch_dtype=torch.bfloat16,
)
self.streamer = TextIteratorStreamer(tokenizer=self.tokenizer, skip_prompt=True, skip_special_tokens=True)
self.chat_template = None
logger.info("Done loading!")
@ModelClass.method
def predict(self,
prompt: str ="",
chat_history: List[dict] = None,
max_tokens: int = Param(default=512, description="The maximum number of tokens to generate. Shorter token lengths will provide faster performance.", ),
temperature: float = Param(default=0.7, description="A decimal number that determines the degree of randomness in the response", ),
top_p: float = Param(default=0.8, description="An alternative to sampling with temperature, where the model considers the results of the tokens with top_p probability mass.", )) -> str:
"""
Predict the response for the given prompt and chat history using the model.
"""
# Construct chat-style messages
messages = chat_history if chat_history else []
if prompt:
messages.append({
"role": "user",
"content": [{"type": "text", "text": prompt}]
})
inputs = self.tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_tensors="pt").to(self.model.device)
generation_kwargs = {
"input_ids": inputs,
"do_sample": True,
"max_new_tokens": max_tokens,
"temperature": temperature,
"top_p": top_p,
"eos_token_id": self.tokenizer.eos_token_id,
}
output = self.model.generate(**generation_kwargs)
generated_tokens = output[0][inputs.shape[-1]:]
return self.tokenizer.decode(generated_tokens, skip_special_tokens=True)
@ModelClass.method
def generate(self,
prompt: str="",
chat_history: List[dict] = None,
max_tokens: int = Param(default=512, description="The maximum number of tokens to generate. Shorter token lengths will provide faster performance.", ),
temperature: float = Param(default=0.7, description="A decimal number that determines the degree of randomness in the response", ),
top_p: float = Param(default=0.8, description="An alternative to sampling with temperature, where the model considers the results of the tokens with top_p probability mass.", )) -> Iterator[str]:
"""Stream generated text tokens from a prompt + optional chat history."""
# Construct chat-style messages
messages = chat_history if chat_history else []
if prompt:
messages.append({
"role": "user",
"content": [{"type": "text", "text": prompt}]
})
response = self.chat(
messages=messages,
max_tokens=max_tokens,
temperature=temperature,
top_p=top_p
)
for each in response:
if 'choices' in each and 'delta' in each['choices'][0] and 'content' in each['choices'][0]['delta']:
yield each['choices'][0]['delta']['content']
@ModelClass.method
def chat(self,
messages: List[dict],
max_tokens: int = Param(default=512, description="The maximum number of tokens to generate. Shorter token lengths will provide faster performance.", ),
temperature: float = Param(default=0.7, description="A decimal number that determines the degree of randomness in the response", ),
top_p: float = Param(default=0.8, description="An alternative to sampling with temperature, where the model considers the results of the tokens with top_p probability mass.", )
) -> Iterator[dict]:
"""
Stream back JSON dicts for assistant messages.
Example return format:
{"role": "assistant", "content": [{"type": "text", "text": "response here"}]}
"""
# Tokenize using chat template
inputs = self.tokenizer.apply_chat_template(
messages,
tokenize=True,
add_generation_prompt=True,
return_tensors="pt"
).to(self.model.device)
generation_kwargs = {
"input_ids": inputs,
"do_sample": True,
"max_new_tokens": max_tokens,
"temperature": temperature,
"top_p": top_p,
"eos_token_id": self.tokenizer.eos_token_id,
"streamer": self.streamer
}
thread = Thread(target=self.model.generate, kwargs=generation_kwargs)
thread.start()
# Accumulate response text
for chunk in openai_response(self.streamer):
yield chunk
thread.join()
def test(self):
"""Test the model here."""
try:
print("Testing predict...")
# Test predict
print(self.predict(prompt="What is the capital of India?",))
except Exception as e:
print("Error in predict", e)
try:
print("Testing generate...")
# Test generate
for each in self.generate(prompt="What is the capital of India?",):
print(each, end="")
print()
except Exception as e:
print("Error in generate", e)
try:
print("Testing chat...")
messages = [
{"role": "system", "content": "You are an helpful assistant."},
{"role": "user", "content": "What is the capital of India?"},
]
for each in self.chat(messages=messages,):
print(each, end="")
print()
except Exception as e:
print("Error in generate", e)
requirements.txt
torch==2.5.1
tokenizers>=0.21.0
transformers>=4.47.0
accelerate>=1.2.0
scipy==1.10.1
optimum>=1.23.3
protobuf==5.27.3
einops>=0.8.0
requests==2.32.3
clarifai>=11.4.1
config.yaml
model:
id: "Llama-3_2-1B-Instruct"
user_id: "user_id"
app_id: "app_id"
model_type_id: "text-to-text"
build_info:
python_version: "3.11"
inference_compute_info:
cpu_limit: "1"
cpu_memory: "13Gi"
num_accelerators: 1
accelerator_type: ["NVIDIA-*"]
accelerator_memory: "44Gi"
checkpoints:
type: "huggingface"
repo_id: "unsloth/Llama-3.2-1B-Instruct"
hf_token: "hf_token"
when: "runtime"
LLaMA 3.2 3B Instruct (LMDeploy)
1/model.py
import os
import sys
from typing import Iterator, List
import json
from clarifai.runners.models.model_builder import ModelBuilder
from clarifai.runners.models.openai_class import OpenAIModelClass
from clarifai.utils.logging import logger
from clarifai.runners.utils.data_utils import Param
from clarifai.runners.utils.openai_convertor import build_openai_messages
from openai import OpenAI
PYTHON_EXEC = sys.executable
def lmdeploy_openai_server(checkpoints, **kwargs):
"""Start lmdeploy OpenAI compatible server."""
from clarifai.runners.utils.model_utils import execute_shell_command, wait_for_server, terminate_process
# Start building the command
cmds = [
PYTHON_EXEC, '-m', 'lmdeploy', 'serve', 'api_server', checkpoints,
]
# Add all parameters from kwargs to the command
for key, value in kwargs.items():
if value is None: # Skip None values
continue
param_name = key.replace('_', '-')
if isinstance(value, bool):
if value: # Only add the flag if True
cmds.append(f'--{param_name}')
else:
cmds.extend([f'--{param_name}', str(value)])
# Create server instance
server = type('Server', (), {
'host': kwargs.get('server_name', '0.0.0.0'),
'port': kwargs.get('server_port', 23333),
'backend': "lmdeploy",
'process': None
})()
try:
server.process = execute_shell_command(" ".join(cmds))
logger.info("Waiting for " + f"http://{server.host}:{server.port}")
wait_for_server(f"http://{server.host}:{server.port}")
logger.info("Server started successfully at " + f"http://{server.host}:{server.port}")
except Exception as e:
logger.error(f"Failed to start lmdeploy server: {str(e)}")
if server.process:
terminate_process(server.process)
raise RuntimeError(f"Failed to start lmdeploy server: {str(e)}")
return server
class LMDeployModel(OpenAIModelClass):
"""
A custom runner that integrates with the Clarifai platform and uses lmdeploy framework for inference
"""
client = True
model = True
def load_model(self):
"""Load the model here and start the server."""
# server args were generated by `upload` module
server_args = {'backend': 'turbomind', 'cache_max_entry_count': 0.95,
'tp': 1, 'max_prefill_token_num': 8192, 'dtype': 'auto',
'model_format': None, 'quant_policy': 0, 'chat_template': 'llama3_2', 'max_batch_size': 16,
'device': 'cuda', 'server_name': '0.0.0.0', 'server_port': 23333,
'tool_call_parser': 'llama3',
}
model_path = os.path.dirname(os.path.dirname(__file__))
builder = ModelBuilder(model_path, download_validation_only=True)
model_config = builder.config
stage = model_config["checkpoints"]['when']
checkpoints = builder.config["checkpoints"]['repo_id']
if stage in ["build", "runtime"]:
checkpoints = builder.download_checkpoints(stage=stage)
# Start server
self.server = lmdeploy_openai_server(checkpoints, **server_args)
# Create client
self.client = OpenAI(
api_key="notset",
base_url= f"http://{self.server.host}:{self.server.port}/v1")
self.model = self.client.models.list().data[0].id
@OpenAIModelClass.method
def predict(self,
prompt: str,
chat_history: List[dict] = None,
tools: List[dict] = None,
tool_choice: str = None,
max_tokens: int = Param(default=512, description="The maximum number of tokens to generate. Shorter token lengths will provide faster performance.", ),
temperature: float = Param(default=0.7, description="A decimal number that determines the degree of randomness in the response", ),
top_p: float = Param(default=0.95, description="An alternative to sampling with temperature, where the model considers the results of the tokens with top_p probability mass."),
) -> str:
"""
This method is used to predict the response for the given prompt and chat history using the model and tools.
"""
if tools is not None and tool_choice is None:
tool_choice = "auto"
messages = build_openai_messages(prompt=prompt, messages=chat_history)
response = self.client.chat.completions.create(
model=self.model,
messages=messages,
tools=tools,
tool_choice=tool_choice,
max_completion_tokens=max_tokens,
temperature=temperature,
top_p=top_p)
if response.choices[0] and response.choices[0].message.tool_calls:
# If the response contains tool calls, return as a string
tool_calls = response.choices[0].message.tool_calls
tool_calls_json = json.dumps([tc.to_dict() for tc in tool_calls], indent=2)
return tool_calls_json
else:
# Otherwise, return the content of the first choice
return response.choices[0].message.content
@OpenAIModelClass.method
def generate(self,
prompt: str,
chat_history: List[dict] = None,
tools: List[dict] = None,
tool_choice: str = None,
max_tokens: int = Param(default=512, description="The maximum number of tokens to generate. Shorter token lengths will provide faster performance.", ),
temperature: float = Param(default=0.7, description="A decimal number that determines the degree of randomness in the response", ),
top_p: float = Param(default=0.95, description="An alternative to sampling with temperature, where the model considers the results of the tokens with top_p probability mass.")) -> Iterator[str]:
"""
This method is used to stream generated text tokens from a prompt + optional chat history and tools.
"""
messages = build_openai_messages(prompt=prompt, messages=chat_history)
response = self.client.chat.completions.create(
model=self.model,
messages=messages,
tools=tools,
tool_choice=tool_choice,
max_completion_tokens=max_tokens,
temperature=temperature,
top_p=top_p,
stream=True)
for chunk in response:
if chunk.choices:
if chunk.choices[0].delta.tool_calls:
# If the response contains tool calls, return the first one as a string
tool_calls = chunk.choices[0].delta.tool_calls
tool_calls_json = [tc.to_dict() for tc in tool_calls]
# Convert to JSON string
json_string = json.dumps(tool_calls_json, indent=2)
# Yield the JSON string
yield json_string
else:
# Otherwise, return the content of the first choice
text = (chunk.choices[0].delta.content
if (chunk and chunk.choices[0].delta.content) is not None else '')
yield text
# This method is needed to test the model with the test-locally CLI command.
def test(self):
"""Test the model here."""
try:
print("Testing predict...")
# Test predict
print(self.predict(prompt="Hello, how are you?",))
except Exception as e:
print("Error in predict", e)
try:
print("Testing generate...")
# Test generate
for each in self.generate(prompt="Hello, how are you?",):
print(each, end=" ")
except Exception as e:
print("Error in generate", e)
requirements.txt
tokenizers>=0.21.0
accelerate>=1.2.0
optimum>=1.23.3
einops>=0.8.0
packaging
ninja
timm
openai
clarifai
clarifai-protocol
psutil
torch==2.6.0
lmdeploy==0.8.0
transformers>=4.51.2
partial-json-parser
config.yaml
model:
id: Llama-3_2-3B-Instruct
user_id:
app_id:
model_type_id: text-to-text
build_info:
python_version: '3.11'
inference_compute_info:
cpu_limit: '1'
cpu_memory: 6Gi
num_accelerators: 1
accelerator_type:
- NVIDIA-*
accelerator_memory: 20Gi
num_threads: 64
checkpoints:
type: huggingface
repo_id: meta-llama/Llama-3.2-3B-Instruct
hf_token:
when: runtime
Gemma 3 1B Instruct (vLLM)
1/model.py
import os
import sys
from typing import List, Iterator
from clarifai.runners.models.model_builder import ModelBuilder
from clarifai.runners.models.openai_class import OpenAIModelClass
from openai import OpenAI
from clarifai.runners.utils.openai_convertor import build_openai_messages
from clarifai.runners.utils.data_utils import Param
from clarifai.utils.logging import logger
PYTHON_EXEC = sys.executable
def vllm_openai_server(checkpoints, **kwargs):
"""Start vLLM OpenAI compatible server."""
from clarifai.runners.utils.model_utils import execute_shell_command, wait_for_server, terminate_process
# Start building the command
cmds = [
PYTHON_EXEC, '-m', 'vllm.entrypoints.openai.api_server', '--model', checkpoints,
]
# Add all parameters from kwargs to the command
for key, value in kwargs.items():
if value is None: # Skip None values
continue
param_name = key.replace('_', '-')
if isinstance(value, bool):
if value: # Only add the flag if True
cmds.append(f'--{param_name}')
else:
cmds.extend([f'--{param_name}', str(value)])
# Create server instance
server = type('Server', (), {
'host': kwargs.get('host', '0.0.0.0'),
'port': kwargs.get('port', 23333),
'backend': "vllm",
'process': None
})()
try:
server.process = execute_shell_command(" ".join(cmds))
logger.info("Waiting for " + f"http://{server.host}:{server.port}")
wait_for_server(f"http://{server.host}:{server.port}")
logger.info("Server started successfully at " + f"http://{server.host}:{server.port}")
except Exception as e:
logger.error(f"Failed to start vllm server: {str(e)}")
if server.process:
terminate_process(server.process)
raise RuntimeError(f"Failed to start vllm server: {str(e)}")
return server
class VLLMLlamaModel(OpenAIModelClass):
"""
A Model that integrates with the Clarifai platform and uses vLLM framework for inference to run the Llama 3.1 8B model with tool calling capabilities.
"""
client = True # This will be set in load_model method
model = True # This will be set in load_model method
def load_model(self):
"""Load the model here and start the server."""
os.path.join(os.path.dirname(__file__))
# This is the path to the chat template file and you can get this chat template from vLLM repo(https://github.com/vllm-project/vllm/blob/main/examples/tool_chat_template_llama3.1_json.jinja)
server_args = {
'max_model_len': 2048,
#'gpu_memory_utilization': 0.8,
'dtype': 'auto',
'task': 'auto',
'kv_cache_dtype': 'auto',
'tensor_parallel_size': 1,
'quantization': None,
'cpu_offload_gb': 5.0,
'chat_template': None,
'port': 23333,
'host': 'localhost',
}
model_path = os.path.dirname(os.path.dirname(__file__))
builder = ModelBuilder(model_path, download_validation_only=True)
model_config = builder.config
stage = model_config["checkpoints"]['when']
checkpoints = builder.config["checkpoints"]['repo_id']
if stage in ["build", "runtime"]:
checkpoints = builder.download_checkpoints(stage=stage)
# Start server
self.server = vllm_openai_server(checkpoints, **server_args)
# CLIent initialization
self.client = OpenAI(
api_key="notset",
base_url=f'http://{self.server.host}:{self.server.port}/v1')
self.model = self.client.models.list().data[0].id
@OpenAIModelClass.method
def predict(self,
prompt: str,
chat_history: List[dict] = None,
tools: List[dict] = None,
tool_choice: str = None,
max_tokens: int = Param(default=512, description="The maximum number of tokens to generate. Shorter token lengths will provide faster performance.", ),
temperature: float = Param(default=0.7, description="A decimal number that determines the degree of randomness in the response", ),
top_p: float = Param(default=0.95, description="An alternative to sampling with temperature, where the model considers the results of the tokens with top_p probability mass."),
) -> str:
"""
This method is used to predict the response for the given prompt and chat history using the model and tools.
"""
if tools is not None and tool_choice is None:
tool_choice = "auto"
messages = build_openai_messages(prompt=prompt, messages=chat_history)
response = self.client.chat.completions.create(
model=self.model,
messages=messages,
tools=tools,
tool_choice=tool_choice,
max_completion_tokens=max_tokens,
temperature=temperature,
top_p=top_p)
if response.choices[0] and response.choices[0].message.tool_calls:
import json
# If the response contains tool calls, return as a string
tool_calls = response.choices[0].message.tool_calls
tool_calls_json = json.dumps([tc.to_dict() for tc in tool_calls], indent=2)
return tool_calls_json
else:
# Otherwise, return the content of the first choice
return response.choices[0].message.content
@OpenAIModelClass.method
def generate(self,
prompt: str,
chat_history: List[dict] = None,
tools: List[dict] = None,
tool_choice: str = None,
max_tokens: int = Param(default=512, description="The maximum number of tokens to generate. Shorter token lengths will provide faster performance.", ),
temperature: float = Param(default=0.7, description="A decimal number that determines the degree of randomness in the response", ),
top_p: float = Param(default=0.95, description="An alternative to sampling with temperature, where the model considers the results of the tokens with top_p probability mass.")) -> Iterator[str]:
"""
This method is used to stream generated text tokens from a prompt + optional chat history and tools.
"""
messages = build_openai_messages(prompt=prompt, messages=chat_history)
response = self.client.chat.completions.create(
model=self.model,
messages=messages,
tools=tools,
tool_choice=tool_choice,
max_completion_tokens=max_tokens,
temperature=temperature,
top_p=top_p,
stream=True)
for chunk in response:
if chunk.choices:
if chunk.choices[0].delta.tool_calls:
# If the response contains tool calls, return the first one as a string
import json
tool_calls = chunk.choices[0].delta.tool_calls
tool_calls_json = [tc.to_dict() for tc in tool_calls]
# Convert to JSON string
json_string = json.dumps(tool_calls_json, indent=2)
# Yield the JSON string
yield json_string
else:
# Otherwise, return the content of the first choice
text = (chunk.choices[0].delta.content
if (chunk and chunk.choices[0].delta.content) is not None else '')
yield text
requirements.txt
tokenizers==0.21.0
accelerate==1.2.0
optimum==1.23.3
# xformers
einops==0.8.0
packaging
ninja
qwen-vl-utils==0.0.8
timm
openai
clarifai
psutil
torch==2.6.0
vllm==0.8.0
transformers==4.50.1
backoff==2.2.1
peft>=0.13.2
soundfile>=0.13.1
scipy==1.15.2
librosa
decord
config.yaml
build_info:
python_version: '3.12'
checkpoints:
hf_token: "<your_hf_token>"
repo_id: google/gemma-3-1b-it
type: huggingface
when: runtime
inference_compute_info:
accelerator_memory: 5Gi
accelerator_type:
- NVIDIA-*
cpu_limit: '1'
cpu_memory: 5Gi
num_accelerators: 1
model:
app_id: APP_ID
id: MODEL_ID
model_type_id: text-to-text
user_id: USER_ID
LLaMA 3.1 8B Tool Calling (vLLM)
1/model.py
import os
import sys
from typing import List, Iterator
from clarifai.runners.models.model_builder import ModelBuilder
from clarifai.runners.models.openai_class import OpenAIModelClass
from openai import OpenAI
from clarifai.runners.utils.openai_convertor import build_openai_messages
from clarifai.runners.utils.data_utils import Param
from clarifai.utils.logging import logger
PYTHON_EXEC = sys.executable
def vllm_openai_server(checkpoints, **kwargs):
"""Start vLLM OpenAI compatible server."""
from clarifai.runners.utils.model_utils import execute_shell_command, wait_for_server, terminate_process
# Start building the command
cmds = [
PYTHON_EXEC, '-m', 'vllm.entrypoints.openai.api_server', '--model', checkpoints,
]
# Add all parameters from kwargs to the command
for key, value in kwargs.items():
if value is None: # Skip None values
continue
param_name = key.replace('_', '-')
if isinstance(value, bool):
if value: # Only add the flag if True
cmds.append(f'--{param_name}')
else:
cmds.extend([f'--{param_name}', str(value)])
# Create server instance
server = type('Server', (), {
'host': kwargs.get('host', '0.0.0.0'),
'port': kwargs.get('port', 23333),
'backend': "vllm",
'process': None
})()
try:
server.process = execute_shell_command(" ".join(cmds))
logger.info("Waiting for " + f"http://{server.host}:{server.port}")
wait_for_server(f"http://{server.host}:{server.port}")
logger.info("Server started successfully at " + f"http://{server.host}:{server.port}")
except Exception as e:
logger.error(f"Failed to start vllm server: {str(e)}")
if server.process:
terminate_process(server.process)
raise RuntimeError(f"Failed to start vllm server: {str(e)}")
return server
class VLLMLlamaModel(OpenAIModelClass):
"""
A Model that integrates with the Clarifai platform and uses vLLM framework for inference to run the Llama 3.1 8B model with tool calling capabilities.
"""
client = True # This will be set in load_model method
model = True # This will be set in load_model method
def load_model(self):
"""Load the model here and start the server."""
os.path.join(os.path.dirname(__file__))
# This is the path to the chat template file and you can get this chat template from vLLM repo(https://github.com/vllm-project/vllm/blob/main/examples/tool_chat_template_llama3.1_json.jinja)
chat_template = 'examples/tool_chat_template_llama3.1_json.jinja'
server_args = {
'max_model_len': 2048,
'gpu_memory_utilization': 0.8,
'dtype': 'auto',
'task': 'auto',
'kv_cache_dtype': 'auto',
'tensor_parallel_size': 1,
'quantization': None,
'chat_template': chat_template,
'cpu_offload_gb': 0.0,
'port': 23333,
'host': 'localhost',
"enable_auto_tool_choice": True,
'tool_call_parser': "llama3_json",
}
model_path = os.path.dirname(os.path.dirname(__file__))
builder = ModelBuilder(model_path, download_validation_only=True)
model_config = builder.config
stage = model_config["checkpoints"]['when']
checkpoints = builder.config["checkpoints"]['repo_id']
if stage in ["build", "runtime"]:
checkpoints = builder.download_checkpoints(stage=stage)
# Start server
self.server = vllm_openai_server(checkpoints, **server_args)
# CLIent initialization
self.client = OpenAI(
api_key="notset",
base_url=f'http://{self.server.host}:{self.server.port}/v1')
self.model = self.client.models.list().data[0].id
@OpenAIModelClass.method
def predict(self,
prompt: str,
chat_history: List[dict] = None,
tools: List[dict] = None,
tool_choice: str = None,
max_tokens: int = Param(default=512, description="The maximum number of tokens to generate. Shorter token lengths will provide faster performance.", ),
temperature: float = Param(default=0.7, description="A decimal number that determines the degree of randomness in the response", ),
top_p: float = Param(default=0.95, description="An alternative to sampling with temperature, where the model considers the results of the tokens with top_p probability mass."),
) -> str:
"""
This method is used to predict the response for the given prompt and chat history using the model and tools.
"""
if tools is not None and tool_choice is None:
tool_choice = "auto"
messages = build_openai_messages(prompt=prompt, messages=chat_history)
response = self.client.chat.completions.create(
model=self.model,
messages=messages,
tools=tools,
tool_choice=tool_choice,
max_completion_tokens=max_tokens,
temperature=temperature,
top_p=top_p)
if response.choices[0] and response.choices[0].message.tool_calls:
import json
# If the response contains tool calls, return as a string
tool_calls = response.choices[0].message.tool_calls
tool_calls_json = json.dumps([tc.to_dict() for tc in tool_calls], indent=2)
return tool_calls_json
else:
# Otherwise, return the content of the first choice
return response.choices[0].message.content
@OpenAIModelClass.method
def generate(self,
prompt: str,
chat_history: List[dict] = None,
tools: List[dict] = None,
tool_choice: str = None,
max_tokens: int = Param(default=512, description="The maximum number of tokens to generate. Shorter token lengths will provide faster performance.", ),
temperature: float = Param(default=0.7, description="A decimal number that determines the degree of randomness in the response", ),
top_p: float = Param(default=0.95, description="An alternative to sampling with temperature, where the model considers the results of the tokens with top_p probability mass.")) -> Iterator[str]:
"""
This method is used to stream generated text tokens from a prompt + optional chat history and tools.
"""
messages = build_openai_messages(prompt=prompt, messages=chat_history)
response = self.client.chat.completions.create(
model=self.model,
messages=messages,
tools=tools,
tool_choice=tool_choice,
max_completion_tokens=max_tokens,
temperature=temperature,
top_p=top_p,
stream=True)
for chunk in response:
if chunk.choices:
if chunk.choices[0].delta.tool_calls:
# If the response contains tool calls, return the first one as a string
import json
tool_calls = chunk.choices[0].delta.tool_calls
tool_calls_json = [tc.to_dict() for tc in tool_calls]
# Convert to JSON string
json_string = json.dumps(tool_calls_json, indent=2)
# Yield the JSON string
yield json_string
else:
# Otherwise, return the content of the first choice
text = (chunk.choices[0].delta.content
if (chunk and chunk.choices[0].delta.content) is not None else '')
yield text
requirements.txt
tokenizers==0.21.0
accelerate==1.2.0
optimum==1.23.3
# xformers
einops==0.8.0
packaging
ninja
qwen-vl-utils==0.0.8
timm
openai
clarifai
psutil
torch==2.6.0
vllm==0.8.0
transformers==4.50.1
backoff==2.2.1
peft>=0.13.2
soundfile>=0.13.1
scipy==1.15.2
librosa
decord
config.yaml
model:
id: "llama-3_1-8B-instruct-tool-calling"
user_id: "user_id"
app_id: "app_id"
model_type_id: "text-to-text"
build_info:
python_version: '3.12'
inference_compute_info:
cpu_limit: '1'
cpu_memory: 12Gi
num_accelerators: 1
accelerator_type: ["NVIDIA-*"]
accelerator_memory: 44Gi
checkpoints:
type: huggingface
repo_id: meta-llama/Llama-3.1-8B-Instruct
hf_token: "hf_token"
when: runtime
Local Runners
Ollama
1/model.py
import json
import os
from typing import Iterator, List
from clarifai.runners.models.openai_class import OpenAIModelClass
from clarifai.runners.utils.data_types import Image
from clarifai.runners.utils.data_utils import Param
from clarifai.runners.utils.openai_convertor import build_openai_messages
from clarifai.utils.logging import logger
from openai import OpenAI
# Set default host
if not os.environ.get('OLLAMA_HOST'):
os.environ["OLLAMA_HOST"] = '127.0.0.1:23333'
OLLAMA_HOST = os.environ.get('OLLAMA_HOST')
if not os.environ.get('OLLAMA_CONTEXT_LENGTH'):
# Set default context length if not set
os.environ["OLLAMA_CONTEXT_LENGTH"] = '8192' # Default context length for Llama 3.2
OLLAMA_CONTEXT_LENGTH = os.environ.get('OLLAMA_CONTEXT_LENGTH')
def run_ollama_server(model_name: str = 'llama3.2'):
"""
start the Ollama server.
"""
from clarifai.runners.utils.model_utils import execute_shell_command, terminate_process
try:
logger.info(f"Starting Ollama server in the host: {OLLAMA_HOST}")
start_process = execute_shell_command("ollama serve")
if start_process:
pull_model = execute_shell_command(f"ollama pull {model_name}")
logger.info(f"Model {model_name} pulled successfully.")
logger.info(f"Ollama server started successfully on {OLLAMA_HOST}")
except Exception as e:
logger.error(f"Error starting Ollama server: {e}")
if 'start_process' in locals():
terminate_process(start_process)
raise RuntimeError(f"Failed to start Ollama server: {e}")
# Check if Image has content before building messages
def has_image_content(image: Image) -> bool:
"""Check if Image object has either bytes or URL."""
return bool(getattr(image, 'url', None) or getattr(image, 'bytes', None))
class OllamaModelClass(OpenAIModelClass):
client = True
model = True
def load_model(self):
"""
Load the Ollama model.
"""
# set the model name here or via OLLAMA_MODEL_NAME
self.model = os.environ.get("OLLAMA_MODEL_NAME", 'llama3.2') #'devstral:latest')
# start ollama server
run_ollama_server(model_name=self.model)
self.client = OpenAI(api_key="notset", base_url=f"http://{OLLAMA_HOST}/v1")
logger.info(f"Ollama model loaded successfully: {self.model}")
@OpenAIModelClass.method
def predict(
self,
prompt: str,
image: Image = None,
images: List[Image] = None,
chat_history: List[dict] = None,
tools: List[dict] = None,
tool_choice: str = None,
max_tokens: int = Param(
default=2048,
description="The maximum number of tokens to generate. Shorter token lengths will provide faster performance.",
),
temperature: float = Param(
default=0.7,
description="A decimal number that determines the degree of randomness in the response",
),
top_p: float = Param(
default=0.95,
description="An alternative to sampling with temperature, where the model considers the results of the tokens with top_p probability mass.",
),
) -> str:
"""
This method is used to predict the response for the given prompt and chat history using the model and tools.
"""
if tools is not None and tool_choice is None:
tool_choice = "auto"
img_content = image if has_image_content(image) else None
messages = build_openai_messages(
prompt=prompt, image=img_content, images=images, messages=chat_history
)
response = self.client.chat.completions.create(
model=self.model,
messages=messages,
tools=tools,
tool_choice=tool_choice,
max_completion_tokens=max_tokens,
temperature=temperature,
top_p=top_p,
)
if response.usage is not None:
self.set_output_context(
prompt_tokens=response.usage.prompt_tokens,
completion_tokens=response.usage.completion_tokens,
)
if len(response.choices) == 0:
# still need to send the usage back.
return ""
if response.choices[0] and response.choices[0].message.tool_calls:
# If the response contains tool calls, return as a string
tool_calls = response.choices[0].message.tool_calls
tool_calls_json = json.dumps([tc.to_dict() for tc in tool_calls], indent=2)
return tool_calls_json
else:
# Otherwise, return the content of the first choice
return response.choices[0].message.content
@OpenAIModelClass.method
def generate(
self,
prompt: str,
image: Image = None,
images: List[Image] = None,
chat_history: List[dict] = None,
tools: List[dict] = None,
tool_choice: str = None,
max_tokens: int = Param(
default=2048,
description="The maximum number of tokens to generate. Shorter token lengths will provide faster performance.",
),
temperature: float = Param(
default=0.7,
description="A decimal number that determines the degree of randomness in the response",
),
top_p: float = Param(
default=0.95,
description="An alternative to sampling with temperature, where the model considers the results of the tokens with top_p probability mass.",
),
) -> Iterator[str]:
"""
This method is used to stream generated text tokens from a prompt + optional chat history and tools.
"""
if tools is not None and tool_choice is None:
tool_choice = "auto"
img_content = image if has_image_content(image) else None
messages = build_openai_messages(
prompt=prompt, image=img_content, images=images, messages=chat_history
)
for chunk in self.client.chat.completions.create(
model=self.model,
messages=messages,
tools=tools,
tool_choice=tool_choice,
max_completion_tokens=max_tokens,
temperature=temperature,
top_p=top_p,
stream=True,
stream_options={"include_usage": True},
):
if chunk.usage is not None:
if chunk.usage.prompt_tokens or chunk.usage.completion_tokens:
self.set_output_context(
prompt_tokens=chunk.usage.prompt_tokens,
completion_tokens=chunk.usage.completion_tokens,
)
if len(chunk.choices) == 0: # still need to send the usage back.
yield ""
if chunk.choices:
if chunk.choices[0].delta.tool_calls:
# If the response contains tool calls, return the first one as a string
import json
tool_calls = chunk.choices[0].delta.tool_calls
tool_calls_json = [tc.to_dict() for tc in tool_calls]
# Convert to JSON string
json_string = json.dumps(tool_calls_json, indent=2)
# Yield the JSON string
yield json_string
else:
# Otherwise, return the content of the first choice
text = (
chunk.choices[0].delta.content
if (chunk and chunk.choices[0].delta.content) is not None
else ''
)
yield text
requirements.txt
ollama
clarifai
openai
config.yaml
model:
app_id: local-dev-runner-app
id: local-dev-model
model_type_id: text-to-text
user_id: clarifai-user-id
build_info:
python_version: '3.12'
inference_compute_info:
cpu_limit: '3'
cpu_memory: 14Gi
num_accelerators: 0
MCP
Browser Tools
1/model.py
from typing import Annotated
from urllib.parse import urljoin, urlparse
from fastmcp import FastMCP # use fastmcp v2 not the built in mcp
from pydantic import Field
server = FastMCP(
"browser-tools-mcp-server",
instructions="Web browsing, scraping, and search tools for gathering information from the internet",
stateless_http=True,
)
def make_http_request(
url: str, method: str = "GET", headers: dict = None, timeout: int = 30
) -> tuple[bool, dict]:
"""Make an HTTP request with proper error handling."""
import requests
from requests.adapters import HTTPAdapter
from urllib3.util.retry import Retry
try:
# Setup session with retries
session = requests.Session()
retry_strategy = Retry(
total=3,
backoff_factor=1,
status_forcelist=[429, 500, 502, 503, 504],
)
adapter = HTTPAdapter(max_retries=retry_strategy)
session.mount("http://", adapter)
session.mount("https://", adapter)
# Default headers
default_headers = {
"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36"
}
if headers:
default_headers.update(headers)
response = session.request(method, url, headers=default_headers, timeout=timeout)
return True, {
"status_code": response.status_code,
"headers": dict(response.headers),
"text": response.text,
"url": response.url,
}
except Exception as e:
return False, {"error": str(e)}
@server.tool("fetch_webpage", description="Fetch and return the content of a webpage")
def fetch_webpage(
url: Annotated[str, Field(description="URL of the webpage to fetch")],
include_headers: Annotated[
bool, Field(description="Include HTTP headers in response")
] = False,
max_length: Annotated[
int, Field(description="Maximum content length to return", ge=100, le=50000)
] = 10000,
) -> str:
"""Fetch the content of a webpage and return the text."""
success, response = make_http_request(url)
if not success:
return f"Error fetching webpage: {response.get('error', 'Unknown error')}"
content = response.get('text', '')
# Truncate if too long
if len(content) > max_length:
content = (
content[:max_length]
+ f"\n\n... (content truncated, showing first {max_length} characters)"
)
result = f"URL: {response.get('url', url)}\n"
result += f"Status Code: {response.get('status_code')}\n\n"
if include_headers:
headers = response.get('headers', {})
result += "Headers:\n"
for key, value in headers.items():
result += f" {key}: {value}\n"
result += "\n"
result += "Content:\n"
result += content
return result
@server.tool("extract_text", description="Extract clean text from HTML content")
def extract_text(
url: Annotated[str, Field(description="URL of the webpage")],
remove_scripts: Annotated[bool, Field(description="Remove script and style tags")] = True,
max_length: Annotated[
int, Field(description="Maximum text length to return", ge=100, le=50000)
] = 5000,
) -> str:
"""Extract clean text from a webpage by removing HTML tags."""
success, response = make_http_request(url)
if not success:
return f"Error fetching webpage: {response.get('error', 'Unknown error')}"
try:
from bs4 import BeautifulSoup
soup = BeautifulSoup(response.get('text', ''), 'html.parser')
# Remove script and style elements if requested
if remove_scripts:
for script in soup(["script", "style"]):
script.decompose()
# Extract text
text = soup.get_text()
# Clean up whitespace
lines = (line.strip() for line in text.splitlines())
chunks = (phrase.strip() for line in lines for phrase in line.split(" "))
text = ' '.join(chunk for chunk in chunks if chunk)
# Truncate if too long
if len(text) > max_length:
text = (
text[:max_length]
+ f"\n\n... (text truncated, showing first {max_length} characters)"
)
return f"URL: {response.get('url', url)}\nExtracted Text:\n\n{text}"
except ImportError:
return "Error: BeautifulSoup4 not available for HTML parsing"
except Exception as e:
return f"Error extracting text: {str(e)}"
@server.tool("search_google", description="Search Google and return results")
def search_google(
query: Annotated[str, Field(description="Search query")],
num_results: Annotated[int, Field(description="Number of results to return", ge=1, le=20)] = 5,
safe_search: Annotated[bool, Field(description="Enable safe search")] = True,
) -> str:
"""Search Google and return search results. Note: This is a mock implementation."""
# This is a simplified mock implementation
# In production, you'd use Google Custom Search API or similar service
# Mock search results based on common queries
mock_results = {
"python": [
{
"title": "Welcome to Python.org",
"url": "https://www.python.org/",
"snippet": "The official home of the Python Programming Language",
},
{
"title": "Python Tutorial",
"url": "https://docs.python.org/3/tutorial/",
"snippet": "Python is an easy to learn, powerful programming language",
},
{
"title": "Learn Python",
"url": "https://www.learnpython.org/",
"snippet": "Learn Python programming with interactive tutorials",
},
],
"machine learning": [
{
"title": "Machine Learning | Coursera",
"url": "https://www.coursera.org/learn/machine-learning",
"snippet": "Learn Machine Learning online with courses from top universities",
},
{
"title": "Machine Learning - Wikipedia",
"url": "https://en.wikipedia.org/wiki/Machine_learning",
"snippet": "Machine learning is a method of data analysis that automates analytical model building",
},
{
"title": "Introduction to Machine Learning",
"url": "https://scikit-learn.org/",
"snippet": "Simple and efficient tools for predictive data analysis",
},
],
}
# Find relevant results
results = []
query_lower = query.lower()
for key, search_results in mock_results.items():
if any(word in query_lower for word in key.split()):
results.extend(search_results)
if not results:
# Generic results for unknown queries
results = [
{
"title": f"Search results for: {query}",
"url": "https://www.google.com/search?q=" + query.replace(" ", "+"),
"snippet": f"Various results related to {query}",
},
{
"title": f"Wikipedia: {query}",
"url": f"https://en.wikipedia.org/wiki/{query.replace(' ', '_')}",
"snippet": f"Wikipedia article about {query}",
},
]
# Limit results
results = results[:num_results]
# Format output
output = f"Google Search Results for '{query}':\n\n"
for i, result in enumerate(results, 1):
output += f"{i}. {result['title']}\n"
output += f" URL: {result['url']}\n"
output += f" {result['snippet']}\n\n"
return output
@server.tool("extract_links", description="Extract all links from a webpage")
def extract_links(
url: Annotated[str, Field(description="URL of the webpage")],
filter_domain: Annotated[
bool, Field(description="Only return links from the same domain")
] = False,
max_links: Annotated[
int, Field(description="Maximum number of links to return", ge=1, le=100)
] = 20,
) -> str:
"""Extract all links from a webpage."""
success, response = make_http_request(url)
if not success:
return f"Error fetching webpage: {response.get('error', 'Unknown error')}"
try:
from bs4 import BeautifulSoup
soup = BeautifulSoup(response.get('text', ''), 'html.parser')
# Find all links
links = []
base_domain = urlparse(url).netloc
for link in soup.find_all('a', href=True):
href = link['href']
text = link.get_text(strip=True)
# Convert relative URLs to absolute
if href.startswith('/'):
href = urljoin(url, href)
elif not href.startswith(('http://', 'https://')):
href = urljoin(url, href)
# Filter by domain if requested
if filter_domain:
link_domain = urlparse(href).netloc
if link_domain != base_domain:
continue
links.append({"url": href, "text": text or "No text"})
# Remove duplicates and limit
unique_links = []
seen_urls = set()
for link in links:
if link["url"] not in seen_urls:
unique_links.append(link)
seen_urls.add(link["url"])
if len(unique_links) >= max_links:
break
# Format output
output = f"Links extracted from {url} ({len(unique_links)} unique links):\n\n"
for i, link in enumerate(unique_links, 1):
output += f"{i}. {link['text']}\n"
output += f" URL: {link['url']}\n\n"
return output
except ImportError:
return "Error: BeautifulSoup4 not available for HTML parsing"
except Exception as e:
return f"Error extracting links: {str(e)}"
@server.tool("take_screenshot", description="Take a screenshot of a webpage")
def take_screenshot(
url: Annotated[str, Field(description="URL of the webpage")],
width: Annotated[int, Field(description="Screenshot width", ge=100, le=3840)] = 1280,
height: Annotated[int, Field(description="Screenshot height", ge=100, le=2160)] = 720,
full_page: Annotated[bool, Field(description="Capture full page")] = False,
) -> str:
"""Take a screenshot of a webpage using headless browser (mock implementation)."""
# This is a mock implementation
# In production, you'd use Selenium, Playwright, or similar
try:
# Mock screenshot functionality
screenshot_path = f"/tmp/screenshot_{url.replace('://', '_').replace('/', '_')}.png"
# Simulate screenshot creation
return (
f"Screenshot would be saved to: {screenshot_path}\n"
f"URL: {url}\n"
f"Dimensions: {width}x{height}\n"
f"Full page: {full_page}\n\n"
f"Note: This is a mock implementation. In production, this would use:\n"
f"- Selenium WebDriver\n"
f"- Playwright\n"
f"- Puppeteer\n"
f"- Or similar headless browser automation tools"
)
except Exception as e:
return f"Error taking screenshot: {str(e)}"
@server.tool(
"check_website_status",
description="Check if a website is accessible and get status information",
)
def check_website_status(
url: Annotated[str, Field(description="URL to check")],
check_ssl: Annotated[bool, Field(description="Check SSL certificate validity")] = True,
) -> str:
"""Check website accessibility and status."""
success, response = make_http_request(url, timeout=10)
if not success:
return f"Website is not accessible: {response.get('error', 'Unknown error')}"
status_code = response.get('status_code')
headers = response.get('headers', {})
result = f"Website Status for {url}:\n\n"
result += f"Status Code: {status_code}\n"
result += f"Status: {'✓ Accessible' if 200 <= status_code < 400 else '✗ Error'}\n"
result += f"Server: {headers.get('server', 'Unknown')}\n"
result += f"Content-Type: {headers.get('content-type', 'Unknown')}\n"
result += f"Content-Length: {headers.get('content-length', 'Unknown')}\n"
result += f"Last-Modified: {headers.get('last-modified', 'Unknown')}\n"
if check_ssl and url.startswith('https://'):
result += "\nSSL Status: ✓ HTTPS enabled"
return result
@server.tool("search_webpage_content", description="Search for specific content within a webpage")
def search_webpage_content(
url: Annotated[str, Field(description="URL of the webpage")],
search_term: Annotated[str, Field(description="Term to search for")],
case_sensitive: Annotated[bool, Field(description="Case sensitive search")] = False,
max_matches: Annotated[
int, Field(description="Maximum number of matches to return", ge=1, le=50)
] = 10,
) -> str:
"""Search for specific content within a webpage."""
success, response = make_http_request(url)
if not success:
return f"Error fetching webpage: {response.get('error', 'Unknown error')}"
try:
from bs4 import BeautifulSoup
soup = BeautifulSoup(response.get('text', ''), 'html.parser')
text = soup.get_text()
# Perform search
if not case_sensitive:
text = text.lower()
search_term = search_term.lower()
matches = []
lines = text.split('\n')
for line_num, line in enumerate(lines, 1):
if search_term in line:
# Get context around the match
context_start = max(0, line.find(search_term) - 50)
context_end = min(len(line), line.find(search_term) + len(search_term) + 50)
context = line[context_start:context_end]
matches.append(
{"line": line_num, "context": context.strip(), "full_line": line.strip()}
)
if len(matches) >= max_matches:
break
if not matches:
return f"No matches found for '{search_term}' in {url}"
result = f"Found {len(matches)} matches for '{search_term}' in {url}:\n\n"
for i, match in enumerate(matches, 1):
result += f"{i}. Line {match['line']}: ...{match['context']}...\n\n"
return result
except ImportError:
return "Error: BeautifulSoup4 not available for HTML parsing"
except Exception as e:
return f"Error searching content: {str(e)}"
# Static resource
@server.resource("config://browser_settings")
def get_browser_settings():
return {
"default_timeout": 30,
"max_content_length": 10000,
"supported_formats": ["html", "text", "json"],
"user_agent": "MCP-Browser-Tools/1.0",
"screenshot_formats": ["png", "jpg"],
}
# Dynamic resource template
@server.resource("site://{domain}/info")
def get_site_info(domain: str):
return {
"domain": domain,
"note": "Use check_website_status tool to get actual site information",
}
@server.prompt()
def web_research_prompt(research_type: str) -> str:
"""Generate prompts for web research tasks."""
prompts = {
"content": "To research web content:\n1. Use search_google to find relevant pages\n2. Use fetch_webpage to get page content\n3. Use extract_text for clean text extraction\n4. Use search_webpage_content to find specific information",
"links": "To analyze website links:\n1. Use extract_links to get all links\n2. Filter by domain if needed\n3. Check link status with check_website_status",
"monitoring": "To monitor websites:\n1. Use check_website_status for availability\n2. Use fetch_webpage to check content changes\n3. Use take_screenshot for visual monitoring",
}
return prompts.get(research_type, f"Web research guidance for: {research_type}")
from clarifai.runners.models.mcp_class import MCPModelClass
class MyModelClass(MCPModelClass):
def get_server(self) -> FastMCP:
return server
client.py
import asyncio
import os
from clarifai.urls.helper import ClarifaiUrlHelper
from fastmcp import Client
from fastmcp.client.transports import StreamableHttpTransport
PAT = os.environ['CLARIFAI_PAT']
url = ClarifaiUrlHelper().mcp_api_url() # get url from the current clarifai config
print(url)
transport = StreamableHttpTransport(url=url, headers={"Authorization": "Bearer " + PAT})
async def main():
print("=== Browser Tools MCP Server Examples ===\n")
async with Client(transport) as client:
# List available tools first
print("Available tools:")
tools = await client.list_tools()
for tool in tools:
print(f"- {tool.name}: {tool.description}")
print("\n" + "=" * 50 + "\n")
# Example 1: Fetch webpage content
print("1. Fetching webpage content:")
try:
result = await client.call_tool(
"fetch_webpage", {"url": "https://httpbin.org/get", "max_length": 2000}
)
print(result[0].text)
except Exception as e:
print(f"Error: {e}")
print("\n" + "=" * 50 + "\n")
# Example 2: Extract clean text from webpage
print("2. Extracting clean text:")
try:
result = await client.call_tool(
"extract_text", {"url": "https://httpbin.org/html", "max_length": 1000}
)
print(result[0].text)
except Exception as e:
print(f"Error: {e}")
print("\n" + "=" * 50 + "\n")
# Example 3: Search Google (mock results)
print("3. Searching Google (mock implementation):")
try:
result = await client.call_tool(
"search_google", {"query": "python programming", "num_results": 3}
)
print(result[0].text)
except Exception as e:
print(f"Error: {e}")
print("\n" + "=" * 50 + "\n")
# Example 4: Check website status
print("4. Checking website status:")
try:
result = await client.call_tool(
"check_website_status", {"url": "https://httpbin.org", "check_ssl": True}
)
print(result[0].text)
except Exception as e:
print(f"Error: {e}")
print("\n" + "=" * 50 + "\n")
# Example 5: Extract links from webpage
print("5. Extracting links from webpage:")
try:
result = await client.call_tool(
"extract_links", {"url": "https://httpbin.org", "max_links": 5}
)
print(result[0].text)
except Exception as e:
print(f"Error: {e}")
print("\n" + "=" * 50 + "\n")
if __name__ == "__main__":
asyncio.run(main())
requirements.txt
clarifai
anyio==4.9.0
mcp==1.9.0
fastmcp==2.3.4
requests>=2.31.0
beautifulsoup4==4.12.2
lxml==4.9.3
config.yaml
build_info:
python_version: '3.11'
inference_compute_info:
cpu_limit: 1000m
cpu_memory: 1Gi
num_accelerators: 0
model:
app_id: mcp-examples-app
id: browser-tools-mcp-server
model_type_id: text-to-text
user_id: mcp-examples-user
Code Execution
1/model.py
import io
import tarfile
from typing import Annotated, Any, Dict
import docker
from clarifai.runners.models.mcp_class import MCPModelClass
from fastmcp import FastMCP
from pydantic import Field
server = FastMCP(
"python-execution-server",
instructions="Execute Python code securely in Docker containers",
stateless_http=True,
)
_docker_client = None
def get_docker_client():
"""Get or create Docker client."""
global _docker_client
if _docker_client is None:
try:
_docker_client = docker.from_env()
_docker_client.ping()
except Exception as e:
# Try alternative connection methods for different Docker setups
try:
# TODO: replace base_url with your local machine's docker socket
_docker_client = docker.DockerClient(
base_url='unix:///Users/YOUR_USER_NAME/.rd/docker.sock'
)
_docker_client.ping()
except:
try:
_docker_client = docker.DockerClient(base_url='unix://var/run/docker.sock')
_docker_client.ping()
except:
raise Exception(f"Cannot connect to Docker daemon. Original error: {e}")
return _docker_client
def execute_python_code_fresh_container(code: str) -> Dict[str, Any]:
"""
Execute Python code in a fresh Docker container (OpenAI approach).
Each execution gets a completely clean environment.
"""
try:
client = get_docker_client()
# Pull Python image if not present
try:
client.images.get("python:3.11")
except docker.errors.ImageNotFound:
client.images.pull("python:3.11")
# Create a temporary tar archive containing the script (like OpenAI)
script_name = "script.py"
tarstream = io.BytesIO()
with tarfile.open(fileobj=tarstream, mode="w") as tar:
script_bytes = code.encode("utf-8")
tarinfo = tarfile.TarInfo(name=script_name)
tarinfo.size = len(script_bytes)
tar.addfile(tarinfo, io.BytesIO(script_bytes))
tarstream.seek(0)
# Start fresh container
container = client.containers.create("python:3.11", command="sleep infinity", detach=True)
try:
container.start()
# Put the script into the container
container.put_archive(path="/tmp", data=tarstream.read())
# Execute the script
exec_result = container.exec_run(f"python /tmp/{script_name}")
return {
"stdout": exec_result.output.decode("utf-8", errors='replace'),
"stderr": "",
"status": exec_result.exit_code,
}
finally:
# Always clean up container
container.remove(force=True)
except docker.errors.ContainerError as e:
return {"stdout": "", "stderr": str(e), "status": 1}
except Exception as e:
return {"stdout": "", "stderr": f"Execution error: {str(e)}", "status": 1}
def execute_with_packages(code: str, packages: list = None) -> Dict[str, Any]:
"""
Execute Python code with pre-installed packages.
This is the key enhancement over #1 - allows package installation.
"""
if packages:
# Prepend package installation to the code
install_code = "\n".join(
[
f"import subprocess; subprocess.run(['pip', 'install', '{pkg}'], check=True)"
for pkg in packages
]
)
full_code = f"{install_code}\n\n{code}"
else:
full_code = code
return execute_python_code_fresh_container(full_code)
@server.tool(
"execute_with_packages", description="Execute Python code with packages pre-installed"
)
def execute_with_packages_tool(
code: Annotated[str, Field(description="Python code to execute")],
packages: Annotated[
list[str], Field(description="List of packages to install before execution")
] = None,
) -> str:
"""
Execute Python code with specified packages installed on top of the base Python image.
This enables users to work with the full Python ecosystem.
Example: execute_with_packages("import requests; print(requests.get('https://httpbin.org/json').json())", ["requests"])
"""
if not code.strip():
return "Error: No code provided"
result = execute_with_packages(code, packages or [])
if result["status"] == 0:
output = "--- Execution Successful ---\n"
if packages:
output += f"Packages installed: {', '.join(packages)}\n\n"
if result["stdout"].strip():
output += result["stdout"]
else:
output += "(No output - use print() to see results)"
else:
output = f"--- Execution Error (status: {result['status']}) ---\n"
if result["stderr"].strip():
output += result["stderr"]
if result["stdout"].strip():
output += "\n--- Output ---\n" + result["stdout"]
return output
class MyModel(MCPModelClass):
def get_server(self) -> FastMCP:
"""Return the FastMCP server instance."""
return server
client.py
import asyncio
import os
from clarifai.urls.helper import ClarifaiUrlHelper
from fastmcp import Client
from fastmcp.client.transports import StreamableHttpTransport
PAT = os.environ['CLARIFAI_PAT']
url = ClarifaiUrlHelper().mcp_api_url()
transport = StreamableHttpTransport(url=url, headers={"Authorization": "Bearer " + PAT})
async def main():
async with Client(transport) as client:
# List available tools
print("=== Available Tools ===")
tools = await client.list_tools()
for tool in tools:
print(f"- {tool.name}: {tool.description}")
print()
# Test 1: execute_with_packages
print("=== Test 1: execute_with_packages ===")
code_with_pkg = 'import requests; print(f"Requests version: {requests.__version__}"); print("Package imported successfully!")'
result = await client.call_tool(
"execute_with_packages", {"code": code_with_pkg, "packages": ["requests"]}
)
print(result.content[0].text)
if __name__ == "__main__":
asyncio.run(main())
requirements.txt
# Clarifai SDK - required
clarifai
fastmcp
pydantic
docker
config.yaml
build_info:
python_version: '3.11'
inference_compute_info:
cpu_limit: 1
cpu_memory: 500Mi
num_accelerators: 0
model:
app_id: code-execution-app
id: code-execution-model
model_type_id: text-to-text
user_id: your_user_id
Code Execution Without Docker Version
1/model.py
import subprocess
import tempfile
import os
from typing import Annotated, Any, Dict
from clarifai.runners.models.mcp_class import MCPModelClass
from fastmcp import FastMCP
from pydantic import Field
server = FastMCP(
"python-execution-server",
instructions="Execute Python code using local Python environment",
stateless_http=True,
)
def execute_python_code_fresh_container(code: str) -> Dict[str, Any]:
"""
Execute Python code using local Python environment.
Each execution gets a clean temporary file.
"""
try:
# Create a temporary file for the code
with tempfile.NamedTemporaryFile(mode='w', suffix='.py', delete=False) as temp_file:
temp_file.write(code)
temp_path = temp_file.name
try:
# Execute the script using subprocess
result = subprocess.run(
["python", temp_path],
capture_output=True,
text=True,
check=False
)
return {
"stdout": result.stdout,
"stderr": result.stderr,
"status": result.returncode,
}
finally:
# Clean up temp file
try:
os.unlink(temp_path)
except Exception:
pass
except Exception as e:
return {"stdout": "", "stderr": f"Execution error: {str(e)}", "status": 1}
def execute_with_packages(code: str, packages: list = None) -> Dict[str, Any]:
"""
Execute Python code with pre-installed packages.
This is the key enhancement over #1 - allows package installation.
"""
if packages:
# Prepend package installation to the code
install_code = "\n".join(
[
f"import subprocess; subprocess.run(['pip', 'install', '{pkg}'], check=True)"
for pkg in packages
]
)
full_code = f"{install_code}\n\n{code}"
else:
full_code = code
return execute_python_code_fresh_container(full_code)
@server.tool(
"execute_with_packages", description="Execute Python code with packages pre-installed"
)
def execute_with_packages_tool(
code: Annotated[str, Field(description="Python code to execute")],
packages: Annotated[
list[str], Field(description="List of packages to install before execution")
] = None,
) -> str:
"""
Execute Python code with specified packages installed on top of the base Python image.
This enables users to work with the full Python ecosystem.
Example: execute_with_packages("import requests; print(requests.get('https://httpbin.org/json').json())", ["requests"])
"""
if not code.strip():
return "Error: No code provided"
result = execute_with_packages(code, packages or [])
if result["status"] == 0:
output = "--- Execution Successful ---\n"
if packages:
output += f"Packages installed: {', '.join(packages)}\n\n"
if result["stdout"].strip():
output += result["stdout"]
else:
output += "(No output - use print() to see results)"
else:
output = f"--- Execution Error (status: {result['status']}) ---\n"
if result["stderr"].strip():
output += result["stderr"]
if result["stdout"].strip():
output += "\n--- Output ---\n" + result["stdout"]
return output
class MyModel(MCPModelClass):
def get_server(self) -> FastMCP:
"""Return the FastMCP server instance."""
return server
client.py
import asyncio
import os
from clarifai.urls.helper import ClarifaiUrlHelper
from fastmcp import Client
from fastmcp.client.transports import StreamableHttpTransport
PAT = os.environ['CLARIFAI_PAT']
url = ClarifaiUrlHelper().mcp_api_url()
transport = StreamableHttpTransport(url=url, headers={"Authorization": "Bearer " + PAT})
async def main():
async with Client(transport) as client:
# List available tools
print("=== Available Tools ===")
tools = await client.list_tools()
for tool in tools:
print(f"- {tool.name}: {tool.description}")
print()
# Test 1: execute_with_packages
print("=== Test 1: execute_with_packages ===")
code_with_pkg = 'import requests; print(f"Requests version: {requests.__version__}"); print("Package imported successfully!")'
result = await client.call_tool(
"execute_with_packages", {"code": code_with_pkg, "packages": ["requests"]}
)
print(result.content[0].text)
if __name__ == "__main__":
asyncio.run(main())
requirements.txt
# Clarifai SDK - required
clarifai==11.7.5
fastmcp
pydantic
config.yaml
build_info:
python_version: '3.11'
inference_compute_info:
cpu_limit: 1
cpu_memory: 500Mi
num_accelerators: 0
model:
app_id: code-execution-app-no-docker
id: code-execution-model-no-docker
model_type_id: text-to-text
user_id: your_user_id
Google Drive
1/model.py
import os
from typing import Annotated
from clarifai.utils.logging import logger
from fastmcp import FastMCP # use fastmcp v2 not the built in mcp
from pydantic import Field
server = FastMCP(
"google-drive-mcp-server",
instructions="Google Drive operations for file storage, sharing, and collaboration",
stateless_http=True,
)
def get_drive_service():
"""Create and return a Google Drive service object."""
try:
from google.auth.transport.requests import Request
from google.oauth2.credentials import Credentials
from google_auth_oauthlib.flow import InstalledAppFlow
from googleapiclient.discovery import build
# Define the scope
SCOPES = ['https://www.googleapis.com/auth/drive']
creds = None
# Check for existing token
if os.path.exists('token.json'):
creds = Credentials.from_authorized_user_file('token.json', SCOPES)
# If there are no valid credentials, use mock data
if not creds or not creds.valid:
logger.warning("Google Drive credentials not available, using mock data")
return None
service = build('drive', 'v3', credentials=creds)
return service
except ImportError:
logger.error("Google API client libraries not available")
return None
except Exception as e:
logger.error(f"Failed to create Drive service: {str(e)}")
return None
def handle_drive_operation(operation_func):
"""Decorator to handle Drive operations with proper error handling."""
def wrapper(*args, **kwargs):
try:
service = get_drive_service()
if not service:
# Return mock data if service not available
return operation_func(None, *args, **kwargs)
return operation_func(service, *args, **kwargs)
except Exception as e:
return f"Google Drive operation failed: {str(e)}"
return wrapper
@server.tool("drive_list_files", description="List files in Google Drive")
def drive_list_files(
folder_id: Annotated[str, Field(description="Folder ID to list files from (optional)")] = "",
file_type: Annotated[
str,
Field(
description="Filter by file type (document, spreadsheet, presentation, folder, etc.)"
),
] = "",
max_results: Annotated[
int, Field(description="Maximum number of files to return", ge=1, le=100)
] = 20,
order_by: Annotated[
str, Field(description="Order by: name, modifiedTime, createdTime")
] = "modifiedTime desc",
) -> str:
"""List files in Google Drive."""
@handle_drive_operation
def _list_files(service, folder_id, file_type, max_results, order_by):
if not service:
# Mock response
mock_files = [
{
"id": "1abc123",
"name": "Project Document.docx",
"mimeType": "application/vnd.google-apps.document",
"modifiedTime": "2024-01-15T10:30:00Z",
"size": "2048576",
"owners": [{"displayName": "John Doe"}],
},
{
"id": "2def456",
"name": "Budget Spreadsheet.xlsx",
"mimeType": "application/vnd.google-apps.spreadsheet",
"modifiedTime": "2024-01-14T14:22:00Z",
"size": "1048576",
"owners": [{"displayName": "Jane Smith"}],
},
{
"id": "3ghi789",
"name": "Presentation.pptx",
"mimeType": "application/vnd.google-apps.presentation",
"modifiedTime": "2024-01-13T09:15:00Z",
"size": "5242880",
"owners": [{"displayName": "Bob Johnson"}],
},
]
output = f"Google Drive Files (showing {len(mock_files)} files):\n\n"
for file in mock_files:
file_size_mb = int(file.get('size', 0)) / (1024 * 1024)
output += f"• {file['name']}\n"
output += f" ID: {file['id']}\n"
output += f" Type: {file['mimeType'].split('.')[-1] if '.' in file['mimeType'] else 'Google App'}\n"
output += f" Size: {file_size_mb:.2f} MB\n"
output += f" Modified: {file['modifiedTime']}\n"
output += f" Owner: {file['owners'][0]['displayName']}\n\n"
return output
# Real implementation would use service.files().list()
query_parts = []
if folder_id:
query_parts.append(f"'{folder_id}' in parents")
if file_type:
mime_type_map = {
"document": "application/vnd.google-apps.document",
"spreadsheet": "application/vnd.google-apps.spreadsheet",
"presentation": "application/vnd.google-apps.presentation",
"folder": "application/vnd.google-apps.folder",
"pdf": "application/pdf",
}
if file_type in mime_type_map:
query_parts.append(f"mimeType='{mime_type_map[file_type]}'")
query = " and ".join(query_parts) if query_parts else None
results = (
service.files()
.list(
q=query,
pageSize=max_results,
orderBy=order_by,
fields="files(id,name,mimeType,size,modifiedTime,owners)",
)
.execute()
)
files = results.get('files', [])
if not files:
return "No files found in Google Drive"
output = f"Google Drive Files ({len(files)} files):\n\n"
for file in files:
file_size = int(file.get('size', 0))
file_size_mb = file_size / (1024 * 1024) if file_size > 0 else 0
output += f"• {file['name']}\n"
output += f" ID: {file['id']}\n"
output += f" Type: {file['mimeType'].split('.')[-1] if '.' in file['mimeType'] else 'Google App'}\n"
output += f" Size: {file_size_mb:.2f} MB\n"
output += f" Modified: {file['modifiedTime']}\n"
if file.get('owners'):
output += f" Owner: {file['owners'][0]['displayName']}\n"
output += "\n"
return output
return _list_files(folder_id, file_type, max_results, order_by)
@server.tool("drive_upload_file", description="Upload a file to Google Drive")
def drive_upload_file(
file_path: Annotated[str, Field(description="Local file path to upload")],
file_name: Annotated[str, Field(description="Name for the file in Drive (optional)")] = "",
folder_id: Annotated[str, Field(description="Folder ID to upload to (optional)")] = "",
description: Annotated[str, Field(description="File description")] = "",
) -> str:
"""Upload a file to Google Drive."""
@handle_drive_operation
def _upload_file(service, file_path, file_name, folder_id, description):
if not os.path.exists(file_path):
return f"Error: File '{file_path}' does not exist"
if not service:
# Mock response
file_name = file_name or os.path.basename(file_path)
file_size = os.path.getsize(file_path)
file_size_mb = file_size / (1024 * 1024)
return (
f"Mock upload successful!\n"
f"File: {file_name}\n"
f"Size: {file_size_mb:.2f} MB\n"
f"Drive ID: mock-file-id-{hash(file_path) % 10000}\n"
f"Note: This is a mock response. Install Google API libraries for actual uploads."
)
from googleapiclient.http import MediaFileUpload
file_name = file_name or os.path.basename(file_path)
file_metadata = {'name': file_name, 'description': description}
if folder_id:
file_metadata['parents'] = [folder_id]
media = MediaFileUpload(file_path, resumable=True)
file = (
service.files()
.create(body=file_metadata, media_body=media, fields='id,name,size,webViewLink')
.execute()
)
file_size = int(file.get('size', 0))
file_size_mb = file_size / (1024 * 1024) if file_size > 0 else 0
return (
f"Successfully uploaded '{file_name}' to Google Drive\n"
f"File ID: {file['id']}\n"
f"Size: {file_size_mb:.2f} MB\n"
f"View Link: {file.get('webViewLink', 'N/A')}"
)
return _upload_file(file_path, file_name, folder_id, description)
@server.tool("drive_download_file", description="Download a file from Google Drive")
def drive_download_file(
file_id: Annotated[str, Field(description="Google Drive file ID")],
local_path: Annotated[str, Field(description="Local path to save the file")] = "",
export_format: Annotated[
str, Field(description="Export format for Google Docs (pdf, docx, etc.)")
] = "",
) -> str:
"""Download a file from Google Drive."""
@handle_drive_operation
def _download_file(service, file_id, local_path, export_format):
if not service:
# Mock response
local_path = local_path or f"downloaded_file_{file_id}"
return (
f"Mock download successful!\n"
f"File ID: {file_id}\n"
f"Saved to: {local_path}\n"
f"Note: This is a mock response. Install Google API libraries for actual downloads."
)
# Get file metadata
file_metadata = service.files().get(fileId=file_id).execute()
file_name = file_metadata['name']
if not local_path:
local_path = file_name
# Check if it's a Google Workspace file that needs to be exported
google_mime_types = [
'application/vnd.google-apps.document',
'application/vnd.google-apps.spreadsheet',
'application/vnd.google-apps.presentation',
]
if file_metadata['mimeType'] in google_mime_types:
if not export_format:
export_format = 'pdf' # Default export format
export_mime_types = {
'pdf': 'application/pdf',
'docx': 'application/vnd.openxmlformats-officedocument.wordprocessingml.document',
'xlsx': 'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet',
'pptx': 'application/vnd.openxmlformats-officedocument.presentationml.presentation',
}
request = service.files().export_media(
fileId=file_id,
mimeType=export_mime_types.get(export_format, export_mime_types['pdf']),
)
else:
request = service.files().get_media(fileId=file_id)
# Download the file
import io
fh = io.BytesIO()
downloader = MediaIoBaseDownload(fh, request)
done = False
while done is False:
status, done = downloader.next_chunk()
# Save to local file
with open(local_path, 'wb') as f:
f.write(fh.getvalue())
file_size = os.path.getsize(local_path)
file_size_mb = file_size / (1024 * 1024)
return (
f"Successfully downloaded '{file_name}' from Google Drive\n"
f"Saved to: {local_path}\n"
f"Size: {file_size_mb:.2f} MB"
)
return _download_file(file_id, local_path, export_format)
@server.tool("drive_share_file", description="Share a file or folder in Google Drive")
def drive_share_file(
file_id: Annotated[str, Field(description="Google Drive file or folder ID")],
email: Annotated[str, Field(description="Email address to share with (optional)")] = "",
role: Annotated[
str, Field(description="Permission role: reader, writer, commenter")
] = "reader",
anyone_can_view: Annotated[
bool, Field(description="Make file viewable by anyone with the link")
] = False,
) -> str:
"""Share a file or folder in Google Drive."""
@handle_drive_operation
def _share_file(service, file_id, email, role, anyone_can_view):
if not service:
# Mock response
share_url = f"https://drive.google.com/file/d/{file_id}/view"
return (
f"Mock sharing successful!\n"
f"File ID: {file_id}\n"
f"Share URL: {share_url}\n"
f"Permissions: {role}\n"
f"Note: This is a mock response. Install Google API libraries for actual sharing."
)
permissions_created = []
# Share with specific email if provided
if email:
permission = {'type': 'user', 'role': role, 'emailAddress': email}
result = (
service.permissions()
.create(fileId=file_id, body=permission, sendNotificationEmail=True)
.execute()
)
permissions_created.append(f"Shared with {email} as {role}")
# Make viewable by anyone with link if requested
if anyone_can_view:
permission = {'type': 'anyone', 'role': 'reader'}
service.permissions().create(fileId=file_id, body=permission).execute()
permissions_created.append("Made viewable by anyone with the link")
# Get the file's web view link
file_metadata = service.files().get(fileId=file_id, fields='name,webViewLink').execute()
result = f"Successfully shared '{file_metadata['name']}'\n"
result += f"File ID: {file_id}\n"
result += f"Share URL: {file_metadata.get('webViewLink', 'N/A')}\n"
if permissions_created:
result += "\nPermissions:\n"
for perm in permissions_created:
result += f" • {perm}\n"
return result
return _share_file(file_id, email, role, anyone_can_view)
@server.tool("drive_create_folder", description="Create a new folder in Google Drive")
def drive_create_folder(
folder_name: Annotated[str, Field(description="Name for the new folder")],
parent_folder_id: Annotated[str, Field(description="Parent folder ID (optional)")] = "",
description: Annotated[str, Field(description="Folder description")] = "",
) -> str:
"""Create a new folder in Google Drive."""
@handle_drive_operation
def _create_folder(service, folder_name, parent_folder_id, description):
if not service:
# Mock response
return (
f"Mock folder creation successful!\n"
f"Folder name: {folder_name}\n"
f"Folder ID: mock-folder-id-{hash(folder_name) % 10000}\n"
f"Note: This is a mock response. Install Google API libraries for actual folder creation."
)
folder_metadata = {
'name': folder_name,
'mimeType': 'application/vnd.google-apps.folder',
'description': description,
}
if parent_folder_id:
folder_metadata['parents'] = [parent_folder_id]
folder = (
service.files().create(body=folder_metadata, fields='id,name,webViewLink').execute()
)
return (
f"Successfully created folder '{folder_name}'\n"
f"Folder ID: {folder['id']}\n"
f"View Link: {folder.get('webViewLink', 'N/A')}"
)
return _create_folder(folder_name, parent_folder_id, description)
@server.tool("drive_delete_file", description="Delete a file or folder from Google Drive")
def drive_delete_file(
file_id: Annotated[str, Field(description="Google Drive file or folder ID to delete")],
permanent: Annotated[bool, Field(description="Permanently delete (bypass trash)")] = False,
) -> str:
"""Delete a file or folder from Google Drive."""
@handle_drive_operation
def _delete_file(service, file_id, permanent):
if not service:
# Mock response
return (
f"Mock deletion successful!\n"
f"File ID: {file_id}\n"
f"Permanent: {permanent}\n"
f"Note: This is a mock response. Install Google API libraries for actual deletion."
)
# Get file metadata before deletion
try:
file_metadata = service.files().get(fileId=file_id, fields='name').execute()
file_name = file_metadata['name']
except:
return f"Error: File with ID '{file_id}' not found"
if permanent:
service.files().delete(fileId=file_id).execute()
return f"Permanently deleted '{file_name}' (ID: {file_id})"
else:
# Move to trash
service.files().update(fileId=file_id, body={'trashed': True}).execute()
return f"Moved '{file_name}' to trash (ID: {file_id})"
return _delete_file(file_id, permanent)
@server.tool("drive_search_files", description="Search for files in Google Drive")
def drive_search_files(
query: Annotated[str, Field(description="Search query")],
file_type: Annotated[str, Field(description="File type filter")] = "",
max_results: Annotated[int, Field(description="Maximum number of results", ge=1, le=100)] = 10,
) -> str:
"""Search for files in Google Drive."""
@handle_drive_operation
def _search_files(service, query, file_type, max_results):
if not service:
# Mock response
mock_results = [
{
"id": "search1",
"name": f"Document about {query}.docx",
"mimeType": "application/vnd.google-apps.document",
},
{
"id": "search2",
"name": f"{query} Spreadsheet.xlsx",
"mimeType": "application/vnd.google-apps.spreadsheet",
},
]
output = f"Search results for '{query}' (mock data):\n\n"
for i, file in enumerate(mock_results, 1):
output += f"{i}. {file['name']}\n"
output += f" ID: {file['id']}\n"
output += f" Type: {file['mimeType'].split('.')[-1]}\n\n"
return output
search_query = f"name contains '{query}'"
if file_type:
mime_type_map = {
"document": "application/vnd.google-apps.document",
"spreadsheet": "application/vnd.google-apps.spreadsheet",
"presentation": "application/vnd.google-apps.presentation",
}
if file_type in mime_type_map:
search_query += f" and mimeType='{mime_type_map[file_type]}'"
results = (
service.files()
.list(
q=search_query, pageSize=max_results, fields="files(id,name,mimeType,modifiedTime)"
)
.execute()
)
files = results.get('files', [])
if not files:
return f"No files found matching '{query}'"
output = f"Search results for '{query}' ({len(files)} files):\n\n"
for i, file in enumerate(files, 1):
output += f"{i}. {file['name']}\n"
output += f" ID: {file['id']}\n"
output += f" Type: {file['mimeType'].split('.')[-1] if '.' in file['mimeType'] else 'Google App'}\n"
output += f" Modified: {file['modifiedTime']}\n\n"
return output
return _search_files(query, file_type, max_results)
# Static resource
@server.resource("config://drive_settings")
def get_drive_settings():
return {
"api_version": "v3",
"scopes": ["https://www.googleapis.com/auth/drive"],
"supported_formats": ["pdf", "docx", "xlsx", "pptx"],
"max_upload_size": "5TB",
"permission_roles": ["reader", "writer", "commenter", "owner"],
}
# Dynamic resource template
@server.resource("drive://file/{file_id}/info")
def get_file_info(file_id: str):
return {
"file_id": file_id,
"note": "Use drive_list_files or other tools to get actual file information",
}
@server.prompt()
def drive_workflow_prompt(workflow_type: str) -> str:
"""Generate prompts for Google Drive workflows."""
prompts = {
"collaboration": "For collaboration workflow:\n1. Use drive_upload_file to share documents\n2. Use drive_share_file to set permissions\n3. Use drive_create_folder for organization\n4. Grant appropriate roles (reader, writer, commenter)",
"backup": "For backup workflow:\n1. Use drive_create_folder for organized storage\n2. Use drive_upload_file to backup important files\n3. Consider using drive_share_file for team access\n4. Regular sync using automated tools",
"document_management": "For document management:\n1. Create folder structure with drive_create_folder\n2. Upload files with drive_upload_file\n3. Use drive_search_files to find documents\n4. Share selectively with drive_share_file",
}
return prompts.get(workflow_type, f"Google Drive workflow guidance for: {workflow_type}")
from clarifai.runners.models.mcp_class import MCPModelClass
class MyModelClass(MCPModelClass):
def get_server(self) -> FastMCP:
return server
client.py
import asyncio
import os
import tempfile
from clarifai.urls.helper import ClarifaiUrlHelper
from fastmcp import Client
from fastmcp.client.transports import StreamableHttpTransport
PAT = os.environ['CLARIFAI_PAT']
url = ClarifaiUrlHelper().mcp_api_url() # get url from the current clarifai config
transport = StreamableHttpTransport(url=url, headers={"Authorization": "Bearer " + PAT})
async def main():
print("=== Google Drive MCP Server Examples ===\n")
async with Client(transport) as client:
# List available tools first
print("Available tools:")
tools = await client.list_tools()
for tool in tools:
print(f"- {tool.name}: {tool.description}")
print("\n" + "=" * 50 + "\n")
# Example 1: List files in Google Drive
print("1. Listing files in Google Drive:")
try:
result = await client.call_tool(
"drive_list_files", {"max_results": 5, "order_by": "modifiedTime desc"}
)
print(result[0].text)
except Exception as e:
print(f"Error: {e}")
print("\n" + "=" * 50 + "\n")
# Example 2: Search for files
print("2. Searching for files:")
try:
result = await client.call_tool(
"drive_search_files", {"query": "project", "max_results": 3}
)
print(result[0].text)
except Exception as e:
print(f"Error: {e}")
print("\n" + "=" * 50 + "\n")
# Example 3: Create a folder
print("3. Creating a folder:")
try:
result = await client.call_tool(
"drive_create_folder",
{"folder_name": "MCP Demo Folder", "description": "Folder created by MCP example"},
)
print(result[0].text)
except Exception as e:
print(f"Error: {e}")
print("\n" + "=" * 50 + "\n")
# Example 4: Upload a file (creates a temporary file for demo)
print("4. Uploading a file:")
# Create a temporary file for demonstration
with tempfile.NamedTemporaryFile(mode='w', delete=False, suffix='.txt') as f:
f.write("This is a test file for Google Drive upload demonstration.")
temp_file_path = f.name
try:
result = await client.call_tool(
"drive_upload_file",
{
"file_path": temp_file_path,
"file_name": "MCP Demo File.txt",
"description": "File uploaded by MCP example",
},
)
print(result[0].text)
except Exception as e:
print(f"Error: {e}")
finally:
# Clean up temporary file
os.unlink(temp_file_path)
print("\n" + "=" * 50 + "\n")
# Example 5: Share a file
print("5. Sharing a file:")
try:
result = await client.call_tool(
"drive_share_file", {"file_id": "mock-file-id-123", "anyone_can_view": True}
)
print(result[0].text)
except Exception as e:
print(f"Error: {e}")
print("\n" + "=" * 50 + "\n")
# Example 6: Download a file
print("6. Downloading a file:")
try:
result = await client.call_tool(
"drive_download_file",
{
"file_id": "mock-file-id-123",
"local_path": "downloaded_document.pdf",
"export_format": "pdf",
},
)
print(result[0].text)
except Exception as e:
print(f"Error: {e}")
print("\n" + "=" * 50 + "\n")
print("\n" + "=" * 50)
print("Note: Google Drive authentication required for actual operations:")
print("- Set up OAuth 2.0 credentials in Google Cloud Console")
print("- Download credentials.json file")
print("- Run authentication flow to generate token.json")
print("- This example uses mock data for demonstration")
print("=" * 50)
if __name__ == "__main__":
asyncio.run(main())
requirements.txt
clarifai
anyio==4.9.0
mcp==1.9.0
fastmcp==2.3.4
google-api-python-client==2.111.0
google-auth-httplib2==0.2.0
google-auth-oauthlib==1.1.0
config.yaml
build_info:
python_version: '3.11'
inference_compute_info:
cpu_limit: 1000m
cpu_memory: 1Gi
num_accelerators: 0
model:
app_id: mcp-examples-app
id: google-drive-mcp-server
model_type_id: text-to-text
user_id: mcp-examples-user
Math
1/model.py
from fastmcp import FastMCP
from pydantic import Field
from typing import Annotated
from clarifai.runners.models.mcp_class import MCPModelClass
server = FastMCP("my-mcp-server", instructions="", stateless_http=True)
@server.tool("addition_tool", description="Add two numbers")
def addition_tool(
a: Annotated[float, Field(description="First number")],
b: Annotated[float, Field(description="Second number")],
) -> float:
"""Add two numbers"""
return a + b
@server.tool("subtraction_tool", description="Subtract two numbers")
def subtraction_tool(
a: Annotated[float, Field(description="First number")],
b: Annotated[float, Field(description="Second number")],
) -> float:
"""Subtract two numbers"""
return a - b
@server.tool("multiplication_tool", description="Multiply two numbers")
def multiplication_tool(
a: Annotated[float, Field(description="First number")],
b: Annotated[float, Field(description="Second number")],
) -> float:
"""Multiply two numbers"""
return a * b
@server.tool("division_tool", description="Divide two numbers")
def division_tool(
a: Annotated[float, Field(description="First number")],
b: Annotated[float, Field(description="Second number")],
) -> float:
"""Divide two numbers"""
return a / b
class MyModel(MCPModelClass):
def get_server(self) -> FastMCP:
"""Return the FastMCP server instance."""
return server
client.py
import asyncio
import os
from clarifai.urls.helper import ClarifaiUrlHelper
from fastmcp import Client
from fastmcp.client.transports import StreamableHttpTransport
PAT = os.environ['CLARIFAI_PAT']
url = ClarifaiUrlHelper().mcp_api_url() # get url from the current clarifai config
transport = StreamableHttpTransport(url=url, headers={"Authorization": "Bearer " + PAT})
async def main():
async with Client(transport) as client:
tools = await client.list_tools()
for tool in tools:
print(f"- {tool.name}: {tool.description}")
result = await client.call_tool("addition_tool", {"a": 10.1, "b": 5.2})
print(f"10.1 + 5.2 = {result.content[0].text}")
result = await client.call_tool("subtraction_tool", {"a": 10.1, "b": 3.2})
print(f"10.1 - 3.2 = {result.content[0].text}")
result = await client.call_tool("multiplication_tool", {"a": 4.1, "b": 7.2})
print(f"4.1 * 7.2 = {result.content[0].text}")
result = await client.call_tool("division_tool", {"a": 20.1, "b": 4.2})
print(f"20.1 / 4.2 = {result.content[0].text}")
if __name__ == "__main__":
asyncio.run(main())
requirements.txt
# Clarifai SDK - required
clarifai
fastmcp
pydantic
config.yaml
build_info:
python_version: '3.11'
inference_compute_info:
cpu_limit: 1
cpu_memory: 500Mi
num_accelerators: 0
model:
app_id: mcp-math-app
id: mcp-math-model
model_type_id: text-to-text
user_id: mcp-math-user
Postgres
1/model.py
import os
from typing import Annotated, Any
from clarifai.utils.logging import logger
from fastmcp import FastMCP # use fastmcp v2 not the built in mcp
from pydantic import Field
server = FastMCP(
"postgres-mcp-server",
instructions="PostgreSQL database operations and management",
stateless_http=True,
)
def get_postgres_connection(host: str, port: int, user: str, password: str, database: str):
"""Create a PostgreSQL connection. Returns connection object or None if failed."""
try:
import psycopg2
connection = psycopg2.connect(
host=host, port=port, user=user, password=password, database=database
)
connection.autocommit = True
return connection
except Exception as e:
logger.error(f"PostgreSQL connection failed: {str(e)}")
return None
def execute_postgres_query(connection, query: str, params: tuple = None) -> tuple[bool, Any]:
"""Execute a PostgreSQL query and return results."""
try:
from psycopg2.extras import RealDictCursor
cursor = connection.cursor(cursor_factory=RealDictCursor)
cursor.execute(query, params or ())
# Handle different query types
if query.strip().upper().startswith(('SELECT', 'WITH', 'SHOW', 'EXPLAIN')):
results = cursor.fetchall()
# Convert to list of dicts for JSON serialization
results = [dict(row) for row in results]
cursor.close()
return True, results
else:
# For INSERT, UPDATE, DELETE, etc.
affected_rows = cursor.rowcount
cursor.close()
return True, {"affected_rows": affected_rows, "message": "Query executed successfully"}
except Exception as e:
return False, {"error": str(e)}
@server.tool("postgres_connect", description="Test PostgreSQL database connection")
def postgres_connect(
host: Annotated[str, Field(description="PostgreSQL host")] = "localhost",
port: Annotated[int, Field(description="PostgreSQL port", ge=1, le=65535)] = 5432,
user: Annotated[str, Field(description="PostgreSQL username")] = "postgres",
password: Annotated[str, Field(description="PostgreSQL password")] = "",
database: Annotated[str, Field(description="Database name")] = "postgres",
) -> str:
"""Test connection to PostgreSQL database."""
connection = get_postgres_connection(host, port, user, password, database)
if connection:
try:
cursor = connection.cursor()
cursor.execute("SELECT version();")
version = cursor.fetchone()[0]
cursor.close()
connection.close()
return f"Successfully connected to PostgreSQL database '{database}' on {host}:{port}\nPostgreSQL Version: {version}"
except Exception as e:
connection.close()
return f"Connected but failed to get version: {str(e)}"
else:
return f"Failed to connect to PostgreSQL database '{database}' on {host}:{port}"
@server.tool("postgres_execute_query", description="Execute a PostgreSQL query")
def postgres_execute_query(
query: Annotated[str, Field(description="SQL query to execute")],
host: Annotated[str, Field(description="PostgreSQL host")] = "localhost",
port: Annotated[int, Field(description="PostgreSQL port", ge=1, le=65535)] = 5432,
user: Annotated[str, Field(description="PostgreSQL username")] = "postgres",
password: Annotated[str, Field(description="PostgreSQL password")] = "",
database: Annotated[str, Field(description="Database name")] = "postgres",
limit: Annotated[
int, Field(description="Limit results for SELECT queries", ge=1, le=1000)
] = 100,
) -> str:
"""Execute a SQL query on PostgreSQL database."""
connection = get_postgres_connection(host, port, user, password, database)
if not connection:
return f"Failed to connect to PostgreSQL database '{database}' on {host}:{port}"
try:
# Add LIMIT to SELECT queries if not already present
print("QQQQQQQQQQQQQQQQQQQQQQQQQQQQQQQ")
print(query)
if query.strip().upper().startswith('SELECT') and 'LIMIT' not in query.upper():
l = limit
query = f"{query.strip().rstrip(';')} LIMIT {l};"
print(query)
success, result = execute_postgres_query(connection, query)
connection.close()
if not success:
return f"Query failed: {result.get('error', 'Unknown error')}"
if isinstance(result, list):
# Format SELECT results
if not result:
return "Query executed successfully. No rows returned."
# Create a formatted table
output = f"Query results ({len(result)} rows):\n\n"
if result:
# Get column names
columns = list(result[0].keys())
# Calculate column widths
col_widths = {}
for col in columns:
col_widths[col] = max(
len(col), max(len(str(row.get(col, ''))) for row in result)
)
# Create header
header = " | ".join(col.ljust(col_widths[col]) for col in columns)
separator = " | ".join("-" * col_widths[col] for col in columns)
output += header + "\n"
output += separator + "\n"
# Add data rows
for row in result:
row_str = " | ".join(
str(row.get(col, '')).ljust(col_widths[col]) for col in columns
)
output += row_str + "\n"
return output
else:
# Non-SELECT query result
return f"Query executed successfully. {result.get('message', '')} Affected rows: {result.get('affected_rows', 0)}"
except Exception as e:
if connection:
connection.close()
return f"Error executing query: {str(e)}"
@server.tool("postgres_list_tables", description="List all tables in the database")
def postgres_list_tables(
host: Annotated[str, Field(description="PostgreSQL host")] = "localhost",
port: Annotated[int, Field(description="PostgreSQL port", ge=1, le=65535)] = 5432,
user: Annotated[str, Field(description="PostgreSQL username")] = "postgres",
password: Annotated[str, Field(description="PostgreSQL password")] = "",
database: Annotated[str, Field(description="Database name")] = "postgres",
schema: Annotated[str, Field(description="Schema name")] = "public",
) -> str:
"""List all tables in the PostgreSQL database."""
query = f"""
SELECT table_name, table_type
FROM information_schema.tables
WHERE table_schema = '{schema}'
ORDER BY table_name;
"""
return postgres_execute_query(query, host, port, user, password, database)
@server.tool("postgres_describe_table", description="Describe the structure of a table")
def postgres_describe_table(
table_name: Annotated[str, Field(description="Table name to describe")],
host: Annotated[str, Field(description="PostgreSQL host")] = "localhost",
port: Annotated[int, Field(description="PostgreSQL port", ge=1, le=65535)] = 5432,
user: Annotated[str, Field(description="PostgreSQL username")] = "postgres",
password: Annotated[str, Field(description="PostgreSQL password")] = "",
database: Annotated[str, Field(description="Database name")] = "postgres",
schema: Annotated[str, Field(description="Schema name")] = "public",
) -> str:
"""Describe the structure of a PostgreSQL table."""
query = f"""
SELECT
column_name,
data_type,
character_maximum_length,
is_nullable,
column_default
FROM information_schema.columns
WHERE table_schema = '{schema}' AND table_name = '{table_name}'
ORDER BY ordinal_position;
"""
return postgres_execute_query(query, host, port, user, password, database)
@server.tool("postgres_list_databases", description="List all databases")
def postgres_list_databases(
host: Annotated[str, Field(description="PostgreSQL host")] = "localhost",
port: Annotated[int, Field(description="PostgreSQL port", ge=1, le=65535)] = 5432,
user: Annotated[str, Field(description="PostgreSQL username")] = "postgres",
password: Annotated[str, Field(description="PostgreSQL password")] = "",
) -> str:
"""List all databases on the PostgreSQL server."""
# Use postgres database as default for this query
query = "SELECT datname as database_name FROM pg_database WHERE datistemplate = false ORDER BY datname;"
return postgres_execute_query(query, host, port, user, password, "postgres")
@server.tool("postgres_table_stats", description="Get statistics for a table")
def postgres_table_stats(
table_name: Annotated[str, Field(description="Table name")],
host: Annotated[str, Field(description="PostgreSQL host")] = "localhost",
port: Annotated[int, Field(description="PostgreSQL port", ge=1, le=65535)] = 5432,
user: Annotated[str, Field(description="PostgreSQL username")] = "postgres",
password: Annotated[str, Field(description="PostgreSQL password")] = "",
database: Annotated[str, Field(description="Database name")] = "postgres",
schema: Annotated[str, Field(description="Schema name")] = "public",
) -> str:
"""Get statistics for a PostgreSQL table."""
query = f"""
SELECT
schemaname,
tablename,
attname as column_name,
n_distinct,
correlation
FROM pg_stats
WHERE schemaname = '{schema}' AND tablename = '{table_name}'
ORDER BY attname;
"""
return postgres_execute_query(query, host, port, user, password, database)
@server.tool("postgres_table_size", description="Get size information for tables")
def postgres_table_size(
host: Annotated[str, Field(description="PostgreSQL host")] = "localhost",
port: Annotated[int, Field(description="PostgreSQL port", ge=1, le=65535)] = 5432,
user: Annotated[str, Field(description="PostgreSQL username")] = "postgres",
password: Annotated[str, Field(description="PostgreSQL password")] = "",
database: Annotated[str, Field(description="Database name")] = "postgres",
schema: Annotated[str, Field(description="Schema name")] = "public",
) -> str:
"""Get size information for all tables in the schema."""
query = f"""
SELECT
table_name,
pg_size_pretty(pg_total_relation_size(schemaname||'.'||tablename)) AS size,
pg_size_pretty(pg_relation_size(schemaname||'.'||tablename)) AS table_size,
pg_size_pretty(pg_total_relation_size(schemaname||'.'||tablename) - pg_relation_size(schemaname||'.'||tablename)) AS index_size
FROM pg_tables
WHERE schemaname = '{schema}'
ORDER BY pg_total_relation_size(schemaname||'.'||tablename) DESC;
"""
return postgres_execute_query(query, host, port, user, password, database)
@server.tool("postgres_active_connections", description="Show active database connections")
def postgres_active_connections(
host: Annotated[str, Field(description="PostgreSQL host")] = "localhost",
port: Annotated[int, Field(description="PostgreSQL port", ge=1, le=65535)] = 5432,
user: Annotated[str, Field(description="PostgreSQL username")] = "postgres",
password: Annotated[str, Field(description="PostgreSQL password")] = "",
database: Annotated[str, Field(description="Database name")] = "postgres",
) -> str:
"""Show active connections to the PostgreSQL database."""
query = """
SELECT
pid,
usename,
datname,
client_addr,
application_name,
state,
query_start,
state_change
FROM pg_stat_activity
WHERE state = 'active'
ORDER BY query_start DESC;
"""
return postgres_execute_query(query, host, port, user, password, database)
@server.tool("postgres_create_backup", description="Create a backup using pg_dump")
def postgres_create_backup(
backup_type: Annotated[str, Field(description="Backup type: 'database' or 'table'")],
target_name: Annotated[str, Field(description="Database name or table name")],
host: Annotated[str, Field(description="PostgreSQL host")] = "localhost",
port: Annotated[int, Field(description="PostgreSQL port", ge=1, le=65535)] = 5432,
user: Annotated[str, Field(description="PostgreSQL username")] = "postgres",
password: Annotated[str, Field(description="PostgreSQL password")] = "",
database: Annotated[str, Field(description="Database name (for table backup)")] = "postgres",
) -> str:
"""Create a backup using pg_dump."""
import subprocess
try:
env = os.environ.copy()
env['PGPASSWORD'] = password
if backup_type == "database":
cmd = [
"pg_dump",
f"--host={host}",
f"--port={port}",
f"--username={user}",
"--format=custom",
"--no-password",
target_name,
]
backup_file = f"{target_name}_backup.dump"
elif backup_type == "table":
cmd = [
"pg_dump",
f"--host={host}",
f"--port={port}",
f"--username={user}",
"--format=custom",
"--no-password",
f"--table={target_name}",
database,
]
backup_file = f"{database}_{target_name}_backup.dump"
else:
return "Error: backup_type must be 'database' or 'table'"
# Execute pg_dump
with open(backup_file, 'wb') as f:
result = subprocess.run(
cmd, stdout=f, stderr=subprocess.PIPE, env=env, timeout=300, check=False
)
if result.returncode == 0:
return f"Backup created successfully: {backup_file}"
else:
return f"Backup failed: {result.stderr.decode()}"
except Exception as e:
return f"Backup failed: {str(e)}"
@server.tool("postgres_analyze_table", description="Analyze a table to update statistics")
def postgres_analyze_table(
table_name: Annotated[str, Field(description="Table name to analyze")],
host: Annotated[str, Field(description="PostgreSQL host")] = "localhost",
port: Annotated[int, Field(description="PostgreSQL port", ge=1, le=65535)] = 5432,
user: Annotated[str, Field(description="PostgreSQL username")] = "postgres",
password: Annotated[str, Field(description="PostgreSQL password")] = "",
database: Annotated[str, Field(description="Database name")] = "postgres",
) -> str:
"""Analyze a PostgreSQL table to update statistics."""
return postgres_execute_query(f"ANALYZE {table_name};", host, port, user, password, database)
# Static resource
@server.resource("config://postgres_settings")
def get_postgres_settings():
return {
"default_port": 5432,
"default_host": "localhost",
"default_schema": "public",
"max_query_limit": 1000,
"supported_operations": [
"SELECT",
"INSERT",
"UPDATE",
"DELETE",
"CREATE",
"DROP",
"ALTER",
],
"backup_formats": ["custom", "plain", "tar"],
}
# Dynamic resource template
@server.resource("database://{database_name}/schema_info")
def get_database_schema_info(database_name: str):
return {
"database": database_name,
"note": "Use postgres_list_tables and postgres_describe_table tools to get actual schema information",
}
@server.prompt()
def postgres_query_prompt(query_type: str) -> str:
"""Generate prompts for PostgreSQL query construction."""
prompts = {
"select": "To write a SELECT query:\nSELECT column1, column2 FROM schema.table_name WHERE condition ORDER BY column LIMIT n;",
"insert": "To write an INSERT query:\nINSERT INTO schema.table_name (column1, column2) VALUES (value1, value2);",
"update": "To write an UPDATE query:\nUPDATE schema.table_name SET column1 = value1 WHERE condition;",
"delete": "To write a DELETE query:\nDELETE FROM schema.table_name WHERE condition;",
"create": "To create a table:\nCREATE TABLE schema.table_name (column1 datatype constraints, column2 datatype constraints);",
"index": "To create an index:\nCREATE INDEX index_name ON schema.table_name (column1, column2);",
"jsonb": "For JSONB operations:\nSELECT data->>'key' FROM table WHERE data @> '{\"key\": \"value\"}';",
}
return prompts.get(query_type, f"PostgreSQL query guidance for: {query_type}")
from clarifai.runners.models.mcp_class import MCPModelClass
class MyModelClass(MCPModelClass):
def get_server(self) -> FastMCP:
return server
client.py
import asyncio
import os
from clarifai.urls.helper import ClarifaiUrlHelper
from fastmcp import Client
from fastmcp.client.transports import StreamableHttpTransport
PAT = os.environ['CLARIFAI_PAT']
url = ClarifaiUrlHelper().mcp_api_url() # get url from the current clarifai config
transport = StreamableHttpTransport(url=url, headers={"Authorization": "Bearer " + PAT})
async def main():
print("=== PostgreSQL MCP Server Examples ===\n")
# Note: These examples assume you have PostgreSQL credentials
# Set environment variables: POSTGRES_USER, POSTGRES_PASSWORD, POSTGRES_HOST, POSTGRES_DATABASE
postgres_config = {
"host": os.environ.get("POSTGRES_HOST", "localhost"),
"port": os.environ.get("POSTGRES_PORT", "5432"),
"user": os.environ.get("POSTGRES_USER", "postgres"),
"password": os.environ.get("POSTGRES_PASSWORD", ""),
"database": os.environ.get("POSTGRES_DATABASE", "postgres"),
}
async with Client(transport) as client:
# List available tools first
print("Available tools:")
tools = await client.list_tools()
for tool in tools:
print(f"- {tool.name}: {tool.description}")
print("\n" + "=" * 50 + "\n")
# Example 1: Test PostgreSQL connection
print("1. Testing PostgreSQL connection:")
try:
result = await client.call_tool("postgres_connect", postgres_config)
print(result[0].text)
except Exception as e:
print(f"Error: {e}")
print("\n" + "=" * 50 + "\n")
# Example 2: List databases
print("2. Listing databases:")
try:
result = await client.call_tool(
"postgres_list_databases",
{
"host": postgres_config["host"],
"port": postgres_config["port"],
"user": postgres_config["user"],
"password": postgres_config["password"],
},
)
print(result[0].text)
except Exception as e:
print(f"Error: {e}")
print("\n" + "=" * 50 + "\n")
# Example 3: List tables in database
print("3. Listing tables:")
try:
result = await client.call_tool("postgres_list_tables", postgres_config)
print(result[0].text)
except Exception as e:
print(f"Error: {e}")
print("\n" + "=" * 50 + "\n")
# Example 4: Execute a sample query
print("4. Executing a sample query (SELECT version):")
try:
result = await client.call_tool(
"postgres_execute_query",
{
**postgres_config,
"query": "SELECT version() as postgres_version, current_timestamp as current_time;",
},
)
print(result[0].text)
except Exception as e:
print(f"Error: {e}")
print("\n" + "=" * 50 + "\n")
# Example 5: Get table sizes
print("5. Getting table sizes:")
try:
result = await client.call_tool("postgres_table_size", postgres_config)
print(result[0].text)
except Exception as e:
print(f"Error: {e}")
print("\n" + "=" * 50 + "\n")
print("\n" + "=" * 50)
print("Note: Set these environment variables for actual PostgreSQL connections:")
print("- POSTGRES_HOST (default: localhost)")
print("- POSTGRES_USER (default: postgres)")
print("- POSTGRES_PASSWORD")
print("- POSTGRES_DATABASE (default: postgres)")
print("=" * 50)
if __name__ == "__main__":
asyncio.run(main())
requirements.txt
clarifai
anyio==4.9.0
mcp==1.9.0
fastmcp==2.3.4
psycopg2-binary==2.9.9
config.yaml
build_info:
python_version: '3.11'
inference_compute_info:
cpu_limit: 1000m
cpu_memory: 1Gi
num_accelerators: 0
model:
app_id: mcp-examples-app
id: postgres-mcp-server
model_type_id: text-to-text
user_id: mcp-examples-user
Web Search
1/model.py
from __future__ import annotations
import os, re
from dataclasses import dataclass, field
from typing import Dict, List, Optional, Union
from urllib.parse import urlparse, parse_qs, unquote
import httpx # type: ignore
from bs4 import BeautifulSoup # type: ignore
from mcp.server.fastmcp import FastMCP, Context # type: ignore
USER_AGENT = (
"Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) "
"AppleWebKit/537.36 (KHTML, like Gecko) Chrome/116.0 Safari/537.36"
)
try:
# Optional SDK for DuckDuckGo (pip install duckduckgo_search)
from ddgs import DDGS # type: ignore
except ImportError:
DDGS = None
# ---------------- Data Models ----------------
@dataclass
class SearchResult:
id: int
title: str
url: str
snippet: str
@dataclass
class PageCache:
url: str
text_lines: List[str]
@dataclass
class SessionState:
last_query: Optional[str] = None
results: List[SearchResult] = field(default_factory=list)
opened_pages: Dict[str, PageCache] = field(default_factory=dict)
last_open_url: Optional[str] = None
async def ddg_sdk_search(query: str, topn: int, region: str, safesearch: str) -> List[SearchResult]:
"""
Use duckduckgo_search AsyncDDGS instead of manual HTML scraping.
Falls back to manual method if SDK not available.
"""
if DDGS is None:
raise RuntimeError("duckduckgo_search not installed")
# The library expects: region like 'wt-wt', safesearch one of: 'moderate','off','strict'
safe = safesearch.lower()
if safe not in {"moderate", "off", "strict"}:
safe = "moderate"
results: List[SearchResult] = []
try:
results_gen = DDGS().text(query, region=region, safesearch=safe, max_results=topn)
for r in results_gen:
# r keys commonly: title, href, body
title = (r.get("title") or "").strip()
url = (r.get("href") or "").strip()
snippet = (r.get("body") or "").strip()
if title and url.startswith("http"):
results.append(SearchResult(id=len(results) + 1, title=title, url=url, snippet=snippet))
if len(results) >= topn:
break
except Exception as e:
print(f"[ddg_sdk_search] Exception: {e}")
raise
return results
def _clean_duckduckgo_href(raw: str) -> str:
print(f"[_clean_duckduckgo_href] raw={raw}")
if not raw:
return ""
if raw.startswith("//"):
raw = "https:" + raw
p = urlparse(raw)
if p.netloc.endswith("duckduckgo.com"):
if p.path.startswith("/l/"):
qs = parse_qs(p.query)
if qs.get("uddg"):
target = unquote(qs["uddg"][0])
if "duckduckgo.com" in urlparse(target).netloc:
return ""
return target
if p.path.endswith(".js") or p.path.startswith("/y.js"):
return ""
if not p.scheme:
return ""
return raw
def parse_duckduckgo(html: str, topn: int) -> List[SearchResult]:
print(f"[parse_duckduckgo] Parsing HTML, topn={topn}")
soup = BeautifulSoup(html, "html.parser")
results: List[SearchResult] = []
for block in soup.select("div.result"):
a = block.select_one("a.result__a")
if not a:
continue
title = a.get_text(" ", strip=True)
url = _clean_duckduckgo_href(a.get("href") or "")
if not url or "duckduckgo.com" in urlparse(url).netloc:
continue
snip_el = block.select_one(".result__snippet, .snippet")
snippet = snip_el.get_text(" ", strip=True) if snip_el else ""
if not title or not url.startswith("http"):
continue
results.append(SearchResult(id=len(results) + 1, title=title, url=url, snippet=snippet))
if len(results) >= topn:
break
print(f"[parse_duckduckgo] Found {len(results)} results")
return results
def format_search_results(results: List[SearchResult]) -> str:
print(f"[format_search_results] Formatting {len(results)} results")
if not results:
return "No results."
return "\n\n".join(
f"[{r.id}] {r.title}\nURL: {r.url}\nSnippet: {r.snippet or '(no snippet)'}"
for r in results
)
async def fetch_page(url: str, timeout: int = 20) -> PageCache:
print(f"[fetch_page] Fetching URL: {url} with timeout={timeout}")
headers = {"User-Agent": USER_AGENT, "Accept-Language": "en-US,en;q=0.9"}
try:
async with httpx.AsyncClient(timeout=timeout, follow_redirects=True) as client:
r = await client.get(url, headers=headers)
print(f"[fetch_page] HTTP status: {r.status_code}")
r.raise_for_status()
html = r.text
except Exception as e:
print(f"[fetch_page] Exception: {e}")
raise
soup = BeautifulSoup(html, "html.parser")
for tag in soup(["script", "style", "noscript"]):
tag.decompose()
text = soup.get_text("\n")
lines = [ln.rstrip() for ln in text.splitlines() if ln.strip()]
print(f"[fetch_page] Extracted {len(lines)} lines")
return PageCache(url=url, text_lines=lines)
# ---------------- FastMCP Server (in-process) ----------------
mcp = FastMCP(
name="duckduckgo-browser",
instructions=r"""
Tool for browsing.
The `cursor` appears in brackets before each browsing display: `[{cursor}]`.
Cite information from the tool using the following format:
`【{cursor}†L{line_start}(-L{line_end})?】`, for example: `【6†L9-L11】` or `【8†L3】`.
Do not quote more than 10 words directly from the tool output.
sources=web
""".strip(),
port=8001,
)
@mcp.tool(
name="search",
title="Search for information",
description=
"Searches for information related to `query` and displays `topn` results.",
)
async def search(query: str, topn: int = 10,
region: str = "wt-wt", safesearch: str = "moderate") -> str:
print(f"[search] Query: {query}, topn={topn}, region={region}, safesearch={safesearch}")
try:
print("[search] Using DuckDuckGo SDK")
results = await ddg_sdk_search(query, topn, region, safesearch)
except Exception as e: # noqa: BLE001
print(f"[search] Exception: {e}")
return f"Search error: {e}"
return format_search_results(results)
@mcp.tool(
name="open",
title="Open a link or page",
description="""
Opens the link `id` from the page indicated by `cursor` starting at line number `loc`, showing `num_lines` lines.
Valid link ids are displayed with the formatting: `【{id}†.*】`.
If `cursor` is not provided, the most recent page is implied.
If `id` is a string, it is treated as a fully qualified URL associated with `source`.
If `loc` is not provided, the viewport will be positioned at the beginning of the document or centered on the most relevant passage, if available.
Use this function without `id` to scroll to a new location of an opened page.
""".strip(),
)
async def open(id: Union[int, str] = -1,
loc: int = 0, num_lines: int = 60) -> str:
print(f"[open] id={id}, loc={loc}, num_lines={num_lines}")
if isinstance(id, int) and id != -1:
print("[open] Invalid id (int but not -1)")
return "Result id not found. Run search."
else:
if not isinstance(id, str) or not id.startswith("http"):
print("[open] Invalid id (not a valid URL)")
return "Provide result id or full http(s) URL."
url = id
try:
page = await fetch_page(url)
except Exception as e: # noqa: BLE001
print(f"[open] Exception: {e}")
return f"Open failed: {e}"
total = len(page.text_lines)
if loc < 0: loc = 0
if num_lines <= 0: num_lines = 60
end = min(total, loc + num_lines)
body = "\n".join(f"L{loc+i+1}: {line}" for i, line in enumerate(page.text_lines[loc:end]))
print(f"[open] Returning lines {loc+1}-{end} of {total}")
return f"URL: {url}\nLines {loc+1}-{end} of {total}\n" + "-"*60 + "\n" + body
@mcp.tool(
name="find",
title="Find pattern in page",
description=
"Finds exact matches of `pattern` in the current page, or the page given by `cursor`.",
)
async def find(pattern: str, url: Optional[str] = None, max_matches: int = 50) -> str:
print(f"[find] pattern={pattern}, url={url}, max_matches={max_matches}")
if not url:
print("[find] No URL provided")
return "No page open."
try:
page = await fetch_page(url)
except Exception as e:
print(f"[find] Exception: {e}")
return f"Fetch failed: {e}"
rx = re.compile(re.escape(pattern), re.IGNORECASE)
hits: List[str] = []
for i, line in enumerate(page.text_lines, start=1):
if rx.search(line):
hits.append(f"L{i}: {rx.sub(lambda m: '**'+m.group(0)+'**', line)}")
if len(hits) >= max_matches:
break
print(f"[find] Found {len(hits)} matches")
return "\n".join(hits) if hits else f"No matches for '{pattern}'."
@mcp.resource("config://browser_search_settings")
def get_browser_search_settings() -> Dict[str, str]:
return {
"available_tools": ["search", "open", "find"],
}
@mcp.prompt()
def web_browsing_prompt(task_type: str) -> str:
"""
Generate a usage prompt for web browsing and extraction tasks, tailored to the available tools in this file.
"""
prompts = {
"content_research": (
"To research web content:\n"
"1. Use the 'search' tool to find relevant pages.\n"
"2. Use 'open' to read the content of a result or URL.\n"
"3. Use 'find' to locate specific information within an opened page."
),
"data_extraction": (
"To extract structured data:\n"
"1. Use 'search' to find pages with the data you need.\n"
"2. Use 'open' to load the page content.\n"
"3. Use 'find' to extract or locate specific patterns or fields."
),
"website_analysis": (
"To analyze a website:\n"
"1. Use 'search' to discover relevant pages.\n"
"2. Use 'open' to inspect the content of those pages.\n"
"3. Use 'find' to search for keywords, features, or patterns."
),
"competitive_research": (
"To research competitors:\n"
"1. Use 'search' to find competitor websites or pages.\n"
"2. Use 'open' to review their content.\n"
"3. Use 'find' to compare features, pricing, or other details."
),
"market_research": (
"To conduct market research:\n"
"1. Use 'search' for industry trends and news.\n"
"2. Use 'open' to read relevant articles or reports.\n"
"3. Use 'find' to extract market insights or statistics."
),
"content_monitoring": (
"To monitor web content:\n"
"1. Use 'search' to discover new or updated content.\n"
"2. Use 'open' to review the latest pages.\n"
"3. Use 'find' to detect changes or specific updates."
),
}
# Default fallback if task_type is not recognized
return prompts.get(
task_type,
(
"Web browsing guidance:\n"
"• Use 'search' to find information.\n"
"• Use 'open' to read a page.\n"
"• Use 'find' to locate details within a page."
),
)
from clarifai.runners.models.mcp_class import MCPModelClass
class MyBrowserSearchToolClass(MCPModelClass):
def get_server(self) -> FastMCP:
return mcp
# Main function to run the MCP server
if __name__ == "__main__":
import asyncio
import sys
# Simple approach - just run the server
try:
asyncio.run(mcp.run())
except KeyboardInterrupt:
print("Server stopped by user", file=sys.stderr)
sys.exit(0)
except Exception as e:
print(f"Server error: {e}", file=sys.stderr)
sys.exit(1)
client.py
import asyncio
import os
import json
from openai import AsyncOpenAI
from fastmcp import Client
from fastmcp.client.transports import StreamableHttpTransport
from clarifai.urls.helper import ClarifaiUrlHelper
API_KEY = os.getenv("CLARIFAI_PAT", "None") # Set env var; avoid hardcoding secrets.
transport = StreamableHttpTransport(
url=ClarifaiUrlHelper().mcp_api_url(),
headers={"Authorization": f"Bearer {API_KEY}"},
)
openai_client = AsyncOpenAI(
api_key=API_KEY,
base_url="https://api.clarifai.com/v2/ext/openai/v1"
)
MODEL_ID = "https://clarifai.com/openai/chat-completion/models/gpt-oss-120b"
def format_tools_to_openai_function(tools):
return [
{
"type": "function",
"function": {
"name": tool.name,
"description": f"[{tool.name}] {tool.description}",
"parameters": tool.inputSchema,
},
}
for tool in tools
]
async def first_model_call(messages, tools):
return await openai_client.chat.completions.create(
model=MODEL_ID,
messages=messages,
temperature=0.4,
tools=tools,
tool_choice="auto",
stream=False,
)
async def final_model_call(messages):
# Force answer without further tool use
return await openai_client.chat.completions.create(
model=MODEL_ID,
messages=messages,
temperature=0.4,
tool_choice="none",
stream=False,
)
async def run_two_step_answer(user_prompt: str):
# 1. Base messages
messages = [
{
"role": "system",
"content": (
"You can call tools only in the FIRST response if needed to gather info. "
"After tool results are provided, you MUST produce the final answer without more tool calls."
),
},
{"role": "user", "content": user_prompt},
]
# Load tool definitions
async with Client(transport) as client:
tools_raw = await client.list_tools()
tools = format_tools_to_openai_function(tools_raw)
# 2. First model pass (may request tool calls)
first_resp = await first_model_call(messages, tools)
first_msg = first_resp.choices[0].message
messages.append({
"role": "assistant",
"content": first_msg.content,
"tool_calls": getattr(first_msg, "tool_calls", None)
})
# If no tool calls, just answer now
if not first_msg.tool_calls:
print("Assistant (no tools needed):", first_msg.content)
return first_msg.content
# 3. Execute ALL requested tool calls once
async with Client(transport) as client:
for tool_call in first_msg.tool_calls:
tool_name = tool_call.function.name
raw_args = tool_call.function.arguments or "{}"
try:
args = json.loads(raw_args)
except json.JSONDecodeError:
args = {}
print(f"\n== Executing tool: {tool_name} | args: {raw_args}")
try:
result = await client.call_tool(tool_name, arguments=args)
# Collect text parts
if hasattr(result, "content"):
parts = [getattr(seg, "text", "") for seg in result.content if getattr(seg, "text", "")]
result_text = "\n".join(parts) if parts else str(result)
else:
result_text = str(result)
except Exception as e:
result_text = f"Tool {tool_name} failed: {e}"
print(result_text)
if len(result_text) > 4000:
result_text = result_text[:3500] + "\n...[truncated]..."
messages.append({
"role": "tool",
"tool_call_id": tool_call.id,
"content": result_text
})
# 4. Final model call (NO more tools)
messages.append({
"role": "system",
"content": "Use the tool outputs above to craft the final answer. Do not call tools again."
})
final_resp = await final_model_call(messages)
final_msg = final_resp.choices[0].message.content
print("\nFinal Answer:\n", final_msg)
return final_msg
async def main():
await run_two_step_answer("Who won the 2025 Ballon d'Or?")
if __name__ == "__main__":
asyncio.run(main())
requirements.txt
clarifai==11.7.5
anyio==4.9.0
mcp==1.9.0
fastmcp==2.3.4
requests>=2.31.0
beautifulsoup4==4.12.2
lxml>=4.9.3
ddgs
config.yaml
build_info:
python_version: '3.11'
inference_compute_info:
cpu_limit: 1000m
cpu_memory: 1Gi
num_accelerators: 0
model:
app_id: test-mcp
id: browser-search-mcp-model
model_type_id: text-to-text
user_id: clarifai-user-id
Multimodal Models
Qwen2.5 VL 3B Instruct
1/model.py
import os
import sys
sys.path.append(os.path.dirname(__file__))
from typing import Iterator, List
from clarifai.runners.models.model_builder import ModelBuilder
from clarifai.runners.models.openai_class import OpenAIModelClass
from clarifai.runners.utils.data_types import Image
from clarifai.runners.utils.data_utils import Param
from clarifai.runners.utils.openai_convertor import build_openai_messages
from clarifai.utils.logging import logger
from openai import OpenAI
from openai_server_starter import OpenAI_APIServer
class VLLMModel(OpenAIModelClass):
"""
A custom runner that integrates with the Clarifai platform and uses Server inference
to process inputs, including text and images.
"""
client = True # This will be set in load_model method
model = True # This will be set in load_model method
def load_model(self):
"""Load the model here and start the server."""
os.path.join(os.path.dirname(__file__))
server_args = {
'max_model_len': '8192',
'gpu_memory_utilization': 0.95,
'dtype': 'auto',
'task': 'auto',
'kv_cache_dtype': 'auto',
'tensor_parallel_size': 1,
'chat_template': None,
'cpu_offload_gb': 0.0,
'quantization': None,
'port': 23333,
'host': 'localhost',
'checkpoints': 'runtime'
}
stage = server_args.get("checkpoints")
if stage in ["build", "runtime"]:
config_path = os.path.dirname(os.path.dirname(__file__))
builder = ModelBuilder(config_path, download_validation_only=True)
checkpoints = builder.download_checkpoints(stage=stage)
server_args.update({"checkpoints": checkpoints})
if server_args.get("additional_list_args") == ['']:
server_args.pop("additional_list_args")
# Start server
# This line were generated by `upload` module
self.server = OpenAI_APIServer.from_vllm_backend(**server_args)
self.client = OpenAI(
api_key="notset",
base_url=VLLMModel.make_api_url(self.server.host, self.server.port))
self.model = self._get_model()
logger.info(f"OpenAI {self.model} model loaded successfully!")
def _get_model(self):
try:
return self.client.models.list().data[0].id
except Exception as e:
raise ConnectionError("Failed to retrieve model ID from API") from e
@staticmethod
def make_api_url(host: str, port: int, version: str = "v1") -> str:
return f"http://{host}:{port}/{version}"
@OpenAIModelClass.method
def predict(self,
prompt: str,
image: Image = None,
images: List[Image] = None,
chat_history: List[dict] = None,
max_tokens: int = Param(default=512, description="The maximum number of tokens to generate. Shorter token lengths will provide faster performance.", ),
temperature: float = Param(default=0.7, description="A decimal number that determines the degree of randomness in the response", ),
top_p: float = Param(default=0.95, description="An alternative to sampling with temperature, where the model considers the results of the tokens with top_p probability mass.", )
) -> str:
"""This is the method that will be called when the runner is run. It takes in an input and
returns an output.
"""
openai_messages = build_openai_messages(prompt=prompt, image=image, images=images, messages=chat_history)
response = self.client.chat.completions.create(
model=self.model,
messages=openai_messages,
max_completion_tokens=max_tokens,
temperature=temperature,
top_p=top_p)
if response.usage and response.usage.prompt_tokens and response.usage.completion_tokens:
self.set_output_context(prompt_tokens=response.usage.prompt_tokens, completion_tokens=response.usage.completion_tokens)
return response.choices[0].message.content
@OpenAIModelClass.method
def generate(self,
prompt: str,
image: Image = None,
images: List[Image] = None,
chat_history: List[dict] = None,
max_tokens: int = Param(default=512, description="The maximum number of tokens to generate. Shorter token lengths will provide faster performance.", ),
temperature: float = Param(default=0.7, description="A decimal number that determines the degree of randomness in the response", ),
top_p: float = Param(default=0.95, description="An alternative to sampling with temperature, where the model considers the results of the tokens with top_p probability mass.", )
) -> Iterator[str]:
"""Example yielding a whole batch of streamed stuff back."""
openai_messages = build_openai_messages(prompt=prompt, image=image, images=images, messages=chat_history)
for chunk in self.client.chat.completions.create(
model=self.model,
messages=openai_messages,
max_completion_tokens=max_tokens,
temperature=temperature,
top_p=top_p,
stream=True):
if chunk.choices:
text = (chunk.choices[0].delta.content
if (chunk and chunk.choices[0].delta.content) is not None else '')
yield text
def test(self):
try:
print("Testing predict...")
print(self.predict(prompt="Explain why cat can't fly"))
print(self.predict(prompt="Describe this image", image=Image.from_url("https://samples.clarifai.com/metro-north.jpg")))
except Exception as e:
print(f"Error in predict {e}")
try:
print("Testing generate...")
for each in self.generate(prompt="Explain why cat can't fly"):
print(each, end=" ")
print()
for each in self.generate(prompt="Describe this image", image=Image.from_url("https://samples.clarifai.com/metro-north.jpg")):
print(each, end=" ")
print()
except Exception as e:
print(f"Error in generate {e}")
1/openai_server_starter.py
import os
import signal
import subprocess
import sys
import threading
from typing import List
import psutil
from clarifai.utils.logging import logger
PYTHON_EXEC = sys.executable
def kill_process_tree(parent_pid, include_parent: bool = True, skip_pid: int = None):
"""Kill the process and all its child processes."""
if parent_pid is None:
parent_pid = os.getpid()
include_parent = False
try:
itself = psutil.Process(parent_pid)
except psutil.NoSuchProcess:
return
children = itself.children(recursive=True)
for child in children:
if child.pid == skip_pid:
continue
try:
child.kill()
except psutil.NoSuchProcess:
pass
if include_parent:
try:
itself.kill()
# Sometime processes cannot be killed with SIGKILL (e.g, PID=1 launched by kubernetes),
# so we send an additional signal to kill them.
itself.send_signal(signal.SIGQUIT)
except psutil.NoSuchProcess:
pass
class OpenAI_APIServer:
def __init__(self, **kwargs):
self.server_started_event = threading.Event()
self.process = None
self.backend = None
self.server_thread = None
def __del__(self, *exc):
# This is important
# close the server when exit the program
self.close()
def close(self):
if self.process:
try:
kill_process_tree(self.process.pid)
except:
self.process.terminate()
if self.server_thread:
self.server_thread.join()
def wait_for_startup(self):
self.server_started_event.wait()
def validate_if_server_start(self, line: str):
line_lower = line.lower()
if self.backend in ["vllm", "sglang", "lmdeploy"]:
if self.backend == "vllm":
return "application startup complete" in line_lower or "vllm api server on" in line_lower
else:
return f" running on http://{self.host}:" in line.strip()
elif self.backend == "llamacpp":
return "waiting for new tasks" in line_lower
elif self.backend == "tgi":
return "Connected" in line.strip()
def _start_server(self, cmds):
try:
env = os.environ.copy()
env["VLLM_USAGE_SOURCE"] = "production-docker-image"
self.process = subprocess.Popen(
cmds,
stdout=subprocess.PIPE,
stderr=subprocess.STDOUT,
text=True,
)
for line in self.process.stdout:
logger.info("Server Log: " + line.strip())
if self.validate_if_server_start(line):
self.server_started_event.set()
# break
except Exception as e:
if self.process:
self.process.terminate()
raise RuntimeError(f"Failed to start Server server: {e}")
def start_server_thread(self, cmds: str):
try:
# Start the server in a separate thread
self.server_thread = threading.Thread(target=self._start_server, args=(cmds,), daemon=None)
self.server_thread.start()
# Wait for the server to start
self.wait_for_startup()
except Exception as e:
raise Exception(e)
@classmethod
def from_vllm_backend(cls,
checkpoints,
limit_mm_per_prompt: str = '',
max_model_len: float = None,
gpu_memory_utilization: float = 0.9,
dtype="auto",
task="auto",
kv_cache_dtype: str = "auto",
tensor_parallel_size=1,
chat_template: str = None,
cpu_offload_gb: float = 0.,
quantization: str = None,
port=23333,
host="localhost",
additional_list_args: List[str] = []):
"""Run VLLM OpenAI compatible server
Args:
checkpoints (str): model id or path
limit_mm_per_prompt (str, optional): For each multimodal plugin, limit how many input instances to allow for each prompt. Expects a comma-separated list of items, e.g.: image=16,video=2 allows a maximum of 16 images and 2 videos per prompt. Defaults to 1 for each modality.
max_model_len (float, optional):Model context length. If unspecified, will be automatically derived from the model config. Defaults to None.
gpu_memory_utilization (float, optional): The fraction of GPU memory to be used for the model executor, which can range from 0 to 1. For example, a value of 0.5 would imply 50% GPU memory utilization. If unspecified, will use the default value of 0.9. This is a per-instance limit, and only applies to the current vLLM instance.It does not matter if you have another vLLM instance running on the same GPU. For example, if you have two vLLM instances running on the same GPU, you can set the GPU memory utilization to 0.5 for each instance. Defaults to 0.9.
dtype (str, optional): dtype. Defaults to "float16".
task (str, optional): The task to use the model for. Each vLLM instance only supports one task, even if the same model can be used for multiple tasks. When the model only supports one task, "auto" can be used to select it; otherwise, you must specify explicitly which task to use. Choices {auto, generate, embedding, embed, classify, score, reward, transcription}. Defaults to "auto".
kv_cache_dtype (str, optional): Data type for kv cache storage. If “auto”, will use model data type. CUDA 11.8+ supports fp8 (=fp8_e4m3) and fp8_e5m2. ROCm (AMD GPU) supports fp8 (=fp8_e4m3). Defaults to "auto".
tensor_parallel_size (int, optional): n gpus. Defaults to 1.
chat_template (str, optional): The file path to the chat template, or the template in single-line form for the specified model. Defaults to None.
cpu_offload_gb (float, optional): The space in GiB to offload to CPU, per GPU. Default is 0, which means no offloading. Intuitively, this argument can be seen as a virtual way to increase the GPU memory size. For example, if you have one 24 GB GPU and set this to 10, virtually you can think of it as a 34 GB GPU. Then you can load a 13B model with BF16 weight, which requires at least 26GB GPU memory. Note that this requires fast CPU-GPU interconnect, as part of the model is loaded from CPU memory to GPU memory on the fly in each model forward pass. Defaults to 0.
quantization (str, optional): quantization format {aqlm,awq,deepspeedfp,tpu_int8,fp8,fbgemm_fp8,modelopt,marlin,gguf,gptq_marlin_24,gptq_marlin,awq_marlin,gptq,compressed-tensors,bitsandbytes,qqq,hqq,experts_int8,neuron_quant,ipex,quark,moe_wna16,None}. Defaults to None.
port (int, optional): port. Defaults to 23333.
host (str, optional): host name. Defaults to "localhost".
additional_list_args (List[str], optional): additional args to run subprocess cmd e.g. ["--arg-name", "arg value"]. See more at [this document](https://docs.vllm.ai/en/latest/serving/openai_compatible_server.html#vllm-serve). Defaults to [].
"""
cmds = [
PYTHON_EXEC, '-m', 'vllm.entrypoints.openai.api_server', '--model', checkpoints, '--dtype',
str(dtype), '--task',
str(task), '--kv-cache-dtype',
str(kv_cache_dtype), '--tensor-parallel-size',
str(tensor_parallel_size), '--gpu-memory-utilization',
str(gpu_memory_utilization), '--cpu-offload-gb',
str(cpu_offload_gb), '--port',
str(port), '--host',
str(host), "--trust-remote-code"
]
if quantization:
cmds += [
'--quantization',
str(quantization),
]
if chat_template:
cmds += [
'--chat-template',
str(chat_template),
]
if max_model_len:
cmds += [
'--max-model-len',
str(max_model_len),
]
if limit_mm_per_prompt:
cmds += [
'--limit-mm-per-prompt',
str(limit_mm_per_prompt),
]
if additional_list_args != []:
cmds += additional_list_args
print("CMDS to run vllm server: ", cmds)
_self = cls()
_self.host = host
_self.port = port
_self.backend = "vllm"
_self.start_server_thread(cmds)
import time
time.sleep(5)
return _self
requirements.txt
torch==2.6.0
tokenizers==0.21.1
accelerate==1.2.0
optimum==1.23.3
xformers
einops==0.8.0
packaging
ninja
qwen-vl-utils==0.0.8
timm==1.0.12
openai
clarifai>=11.5.0,<12.0.0
psutil
vllm==0.8.5
transformers==4.51.1
config.yaml
# Config file for the vLLM runner
model:
id: "Qwen2_5-VL-3B-Instruct"
user_id: "user_id"
app_id: "app_id"
model_type_id: "multimodal-to-text"
build_info:
python_version: "3.11"
inference_compute_info:
cpu_limit: "3"
cpu_memory: "14Gi"
num_accelerators: 1
accelerator_type: ["NVIDIA-*"]
accelerator_memory: "44Gi"
checkpoints:
type: "huggingface"
repo_id: "Qwen/Qwen2.5-VL-3B-Instruct"
hf_token: "hf_token"
OCR
NanoNets OCR Small
1/model.py
import os
from io import BytesIO
# Third-party imports
import torch
from PIL import Image as PILImage
from transformers import AutoTokenizer, AutoProcessor, AutoModelForImageTextToText
# Clarifai imports
from clarifai.runners.models.model_builder import ModelBuilder
from clarifai.runners.models.model_class import ModelClass
from clarifai.runners.utils.data_types import Image
from clarifai.runners.utils.data_utils import Param
from clarifai.utils.logging import logger
DEFAULT_PROMPT = """Extract the text from the above document as if you were reading it naturally. Return the tables in html format. Return the equations in LaTeX representation. If there is an image in the document and image caption is not present, add a small description of the image inside the <img></img> tag; otherwise, add the image caption inside <img></img>. Watermarks should be wrapped in brackets. Ex: <watermark>OFFICIAL COPY</watermark>. Page numbers should be wrapped in brackets. Ex: <page_number>14</page_number> or <page_number>9/22</page_number>. Prefer using ☐ and ☑ for check boxes."""
def preprocess_image(image_bytes: bytes) -> PILImage:
"""Convert image bytes to PIL Image."""
return PILImage.open(BytesIO(image_bytes)).convert("RGB")
class MyRunner(ModelClass):
"""A custom runner that loads the OCR model and runs it on the input image."""
def load_model(self):
"""Load the model here."""
model_path = os.path.dirname(os.path.dirname(__file__))
builder = ModelBuilder(model_path, download_validation_only=True)
checkpoint_path = builder.download_checkpoints(stage="runtime")
self.device = "cuda" if torch.cuda.is_available() else "cpu"
logger.info(f"Running on device: {self.device}")
self.model = AutoModelForImageTextToText.from_pretrained(checkpoint_path).to(
self.device
)
self.model.eval()
self.tokenizer = AutoTokenizer.from_pretrained(checkpoint_path)
self.processor = AutoProcessor.from_pretrained(checkpoint_path)
logger.info("Done loading!")
@ModelClass.method
def predict(
self,
image: Image,
prompt: str = Param(
default=DEFAULT_PROMPT, description="The prompt to use for the OCR model."
),
max_new_tokens: int = Param(
default=512,
description="The maximum number of tokens to generate. Shorter token lengths will provide faster performance.",
),
) -> str:
"""This is the method that will be called when the runner is run. It takes in an input and returns an output."""
messages = [
{"role": "system", "content": "You are a helpful assistant."},
{
"role": "user",
"content": [
{
"type": "image_url",
"image_url": f"data:image/png;base64,{image.bytes}",
},
{"type": "text", "text": prompt},
],
},
]
image = preprocess_image(image.bytes)
text = self.processor.apply_chat_template(
messages, tokenize=False, add_generation_prompt=True
)
inputs = self.processor(
text=[text], images=[image], padding=True, return_tensors="pt"
)
inputs = inputs.to(self.model.device)
output_ids = self.model.generate(
**inputs, max_new_tokens=max_new_tokens, do_sample=False
)
generated_ids = [
output_ids[len(input_ids) :]
for input_ids, output_ids in zip(inputs.input_ids, output_ids)
]
output_text = self.processor.batch_decode(
generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True
)
return output_text[0]
def test(self):
"""Test the model with a sample image."""
import requests # Import moved here as it's only used for testing
# Load a sample image from the IAM database
url = "https://dl.a9t9.com/ocr/solarcell.jpg"
image = Image(bytes=requests.get(url).content)
# image = Image.from_url(url)
generated_text = self.predict(image)
# Log the detected text
logger.info(f"Detected text:\n{generated_text}")
requirements.txt
clarifai>=11.5.2,<12.0.0
pillow==10.4.0
torch==2.6.0
torchvision==0.21.0
transformers>=4.51.1
config.yaml
model:
id: "nanonets-ocr-s"
user_id: "user_id"
app_id: "app_id"
model_type_id: "multimodal-to-text"
build_info:
python_version: "3.11"
inference_compute_info:
cpu_limit: "2"
cpu_memory: "18Gi"
num_accelerators: 1
accelerator_type: ["NVIDIA-*"]
accelerator_memory: "18Gi"
checkpoints:
type: "huggingface"
repo_id: "nanonets/Nanonets-OCR-s"
Text Embedder
Jina Embeddings v3
1/model.py
from typing import List, Iterator
import os
from clarifai.runners.models.openai_class import OpenAIModelClass
from clarifai.runners.models.model_builder import ModelBuilder
import onnxruntime
import numpy as np
from transformers import AutoTokenizer, PretrainedConfig
from clarifai.utils.logging import logger
#Helper function Mean pool function
def mean_pooling(model_output: np.ndarray, attention_mask: np.ndarray):
token_embeddings = model_output
input_mask_expanded = np.expand_dims(attention_mask, axis=-1)
input_mask_expanded = np.broadcast_to(input_mask_expanded, token_embeddings.shape)
sum_embeddings = np.sum(token_embeddings * input_mask_expanded, axis=1)
sum_mask = np.clip(np.sum(input_mask_expanded, axis=1), a_min=1e-9, a_max=None)
return sum_embeddings / sum_mask
class Jinaai_embedding_v2(OpenAIModelClass):
"""
A custom runner that integrates with the Clarifai platform and uses Jinaai embedding v2 model to process inputs, including text and images.
"""
client = True
model = "jinaai-embedding-v3"
def load_model(self):
"""Load the model here and start the server."""
# Load checkpoints
model_path = os.path.dirname(os.path.dirname(__file__))
builder = ModelBuilder(model_path, download_validation_only=True)
logger.info(f"\nDownloading Jinaai {self.model} model checkpoints...\n")
self.checkpoints = builder.download_checkpoints(stage='runtime')
logger.info(f"Checkpoints downloaded to {self.checkpoints}")
#logger.info("Loading Jinaai embedding v3 model...")
self.tokenizer = AutoTokenizer.from_pretrained(self.checkpoints)
self.config = PretrainedConfig.from_pretrained(self.checkpoints)
#logger.info(f"Tokenizer and config loaded from 'jinaai/jinaai-embeddings-v3'")
self.session = onnxruntime.InferenceSession(self.checkpoints + '/onnx/model.onnx')
# log that system is ready
logger.info(f"Jinaai {self.model} model loaded successfully!")
def tokenize_and_embed(self, input: str):
"""
Tokenize the input text and return the embedding vector.
Args:
input (str): The input text to be tokenized and embedded.
Returns:
np.ndarray: The embedding vector for the input text.
"""
# Tokenize input
input_text = self.tokenizer(input, return_tensors='np')
# Prepare inputs for ONNX model
task_type = 'text-matching'
task_id = np.array(self.config.lora_adaptations.index(task_type), dtype=np.int64)
inputs = {
'input_ids': input_text['input_ids'],
'attention_mask': input_text['attention_mask'],
'task_id': task_id
}
# Run model
outputs = self.session.run(None, inputs)[0]
return outputs, input_text
@OpenAIModelClass.method
def predict(self,
input: str,) -> List[float]:
"""
Predict method to process the input text and return the embedding vector.
Args:
input (str): The input text to be processed.
Returns:
List[float]: The embedding vector for the input text.
"""
# Tokenize and embed the input text
outputs, input_text = self.tokenize_and_embed(input)
# Apply mean pooling and normalization to the model outputs
embeddings = mean_pooling(outputs, input_text["attention_mask"])
embeddings = embeddings / np.linalg.norm(embeddings, ord=2, axis=1, keepdims=True)
return embeddings[0].tolist()
requirements.txt
transformers #4.52.4
clarifai #11.4.10
onnxruntime
numpy
config.yaml
model:
id: "jinaai-embeddings-v3"
user_id: "user_id"
app_id: "app_id"
model_type_id: "text-embedder"
build_info:
python_version: "3.11"
inference_compute_info:
cpu_limit: "1"
cpu_memory: "5Gi"
num_accelerators: 0
checkpoints:
type: "huggingface"
repo_id: "jinaai/jina-embeddings-v3"
Text-to-Image
FLUX Schnell
1/model.py
import os
from typing import List
from clarifai.runners.models.model_class import ModelClass
from clarifai.runners.utils.data_utils import Param
from clarifai.runners.utils.data_types import Image
from clarifai.runners.models.model_builder import ModelBuilder
from diffusers import FluxPipeline
import torch
class TextToImageModel(ModelClass):
"""
A custom runner for the FLUX model that integrates with the Clarifai platform.
"""
def load_model(self):
"""Load the model here."""
# "black-forest-labs/FLUX.1-schnell"
self.device = "cuda"
model_path = os.path.dirname(os.path.dirname(__file__))
builder = ModelBuilder(model_path, download_validation_only=True)
checkpoints = builder.download_checkpoints(stage="runtime")
# load model and scheduler
self.pipeline = FluxPipeline.from_pretrained(
checkpoints,
torch_dtype=torch.bfloat16
)
self.pipeline = self.pipeline.to(self.device)
@ModelClass.method
def predict(
self,
prompt: str,
num_inference_steps: int = Param(default=28, description="The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference."),
guidance_scale: float = Param(default=3.5, description="The `guidance_scale` controls how strongly the model follows the conditioning input during generation."),
negative_prompt: str = Param(default="", description="The prompt to guide what to not include in image generation. Ignored when not using guidance (guidance_scale < 1)"),
true_cfg_scale: float = Param(default=1.0, description="When > 1.0 and a provided negative_prompt, enables true classifier-free guidance"),
height: int = Param(default=1024, description="The height in pixels of the generated image. This is set to 1024 by default for the best results."),
width: int = Param(default=1024, description="The width in pixels of the generated image. This is set to 1024 by default for the best results."),
max_sequence_length: int = Param(default=256, description="Maximum sequence length to use with the prompt"),
seed: int = Param(default=None, description="Seed value to make generation deterministic."),
# No need
sigmas: List[float] = None,
) -> Image:
image = self.pipeline(
prompt=prompt,
negative_prompt=negative_prompt,
guidance_scale=guidance_scale,
num_inference_steps=num_inference_steps,
max_sequence_length=max_sequence_length,
width=width,
height=height,
true_cfg_scale=true_cfg_scale,
generator=torch.Generator("cpu").manual_seed(seed) if seed else None,
sigmas=sigmas,
).images[0]
# this is important, delete all model cache to avoid OOM
torch.cuda.empty_cache()
return Image.from_pil(image)
@ModelClass.method
def create(
self,
prompt: List[str],
prompt_2: List[str] = None,
negative_prompt: List[str] = None,
negative_prompt_2: List[str] = None,
true_cfg_scale: float = 1.0,
height: int = None,
width: int = None,
max_sequence_length: int = 256,
num_inference_steps: int = 28,
guidance_scale: float = 3.5,
seed: int = None,
sigmas: List[float] = None,
) -> List[Image]:
"""
Generate an image from the given prompt using the FLUX model.
Args:
* prompt (`List[str]`): The prompt or prompts to guide the image generation.
* prompt_2 (`List[str]`, *optional*): The prompt to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is will be used instead.
* negative_prompt (`List[str]`, *optional*): The prompt to guide what to not include in image generation. Ignored when not using guidance (guidance_scale < 1).
* negative_prompt_2 (`List[str]`, *optional*): The negative_prompt to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `negative_prompt` is will be used instead.
* height (`int`, *optional*, defaults to model.unet.config.sample_size * model.vae_scale_factor): The height in pixels of the generated image. This is set to 1024 by default for the best results.
* width (`int`, *optional*, defaults to model.unet.config.sample_size * model.vae_scale_factor): The width in pixels of the generated image. This is set to 1024 by default for the best results.
* num_inference_steps (`int`, *optional*, defaults to 28): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference.
* guidance_scale (`float`, *optional*, defaults to 3.5):
Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
`guidance_scale` is defined as `w` of equation 2. of [Imagen
Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >
1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,
usually at the expense of lower image quality.
* seed (`int`, *optional*, defaults to None):
Seed value passed to `torch.Generator("cpu").manual_seed(seed)` (see more [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)) to make generation deterministic.
* sigmas (`List[float]`, *optional*): Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used.
* max_sequence_length (`int` defaults to `256`): Maximum sequence length to use with the prompt.
(see more at this [doc](https://huggingface.co/docs/diffusers/v0.32.2/en/api/pipelines/flux#diffusers.FluxPipeline.__call__))
"""
assert isinstance(prompt, list), ValueError("prompt must be a list of string")
assert len(prompt) <= 4, ValueError(
f"The provided prompt length ({len(prompt)}) exceeds the maximum limit (4). Please reduce the number of prompts in `prompt`.")
images = self.pipeline(
prompt=prompt,
prompt_2=prompt_2,
negative_prompt=negative_prompt,
negative_prompt_2=negative_prompt_2,
guidance_scale=guidance_scale,
num_inference_steps=num_inference_steps,
max_sequence_length=max_sequence_length,
width=width,
height=height,
true_cfg_scale=true_cfg_scale,
generator=torch.Generator("cpu").manual_seed(seed) if seed else None,
sigmas=sigmas,
).images
# this is important, delete all model cache to avoid OOM
torch.cuda.empty_cache()
return [Image.from_pil(image) for image in images]
def test(self):
"""
Test cases only executed when running `clarifai model test-locally`
"""
image = self.predict(
prompt="A Ghibli animated orange cat, panicked about a deadline, sits in front of a Banana-brand laptop.",
negative_prompt="Ugly, cute", guidance_scale=7)
print(image)
images = self.create(
prompt=["A Ghibli animated orange cat, panicked about a deadline, sits in front of a Banana-brand laptop."]*3,
negative_prompt=["Ugly, cute"]*2, guidance_scale=7)
print(images)
requirements.txt
tokenizers==0.21.0
transformers>=4.48
diffusers==0.32.2
accelerate==1.2.0
optimum==1.23.3
xformers
einops==0.8.0
requests==2.32.3
sentencepiece==0.2.0
numpy>2.0
ninja
aiohttp
packaging
torch==2.5.1
clarifai
clarifai-protocol
config.yaml
model:
id: "flux_1-schnell"
user_id:
app_id:
model_type_id: "text-to-image"
build_info:
python_version: "3.11"
inference_compute_info:
cpu_limit: "3"
cpu_memory: "18Gi"
num_accelerators: 1
accelerator_type: ["NVIDIA-L40S", "NVIDIA-A100", "NVIDIA-H100"]
accelerator_memory: "44Gi"
checkpoints:
type: huggingface
repo_id: black-forest-labs/FLUX.1-schnell
hf_token: