Creating Your First AIPerf Plugin

1
"""My custom AIPerf plugins."""

1
"""Custom endpoint implementations."""
2
3
from my_plugins.endpoints.echo_endpoint import EchoEndpoint
4
5
__all__ = ["EchoEndpoint"]

1
"""Echo endpoint for demonstration purposes."""
2
from __future__ import annotations
3
from typing import Any
4
5
from aiperf.common.models import ParsedResponse, RequestInfo, TextResponseData, InferenceServerResponse
6
from aiperf.endpoints.base_endpoint import BaseEndpoint
7
8
9
class EchoEndpoint(BaseEndpoint):
10
    """Echo endpoint that sends text and receives it back."""
11
12
    # ─────────────────────────────────────────────────────────────────────────
13
    # REQUIRED: Format outgoing request
14
    # ─────────────────────────────────────────────────────────────────────────
15
    def format_payload(self, request_info: RequestInfo) -> dict[str, Any]:
16
        turn = request_info.turns[-1]
17
        model_endpoint = request_info.model_endpoint
18
        texts = [content for text in turn.texts for content in text.contents if content]
19
        return {
20
            "text": texts[0] if texts else "",
21
            "model": turn.model or model_endpoint.primary_model_name,
22
            "max_tokens": turn.max_tokens,
23
            "stream": model_endpoint.endpoint.streaming,
24
        }
25
26
    # ─────────────────────────────────────────────────────────────────────────
27
    # REQUIRED: Parse incoming response
28
    # ─────────────────────────────────────────────────────────────────────────
29
    def parse_response(self, response: InferenceServerResponse) -> ParsedResponse | None:
30
        if json_obj := response.get_json():
31
            if text := json_obj.get("echo") or json_obj.get("text"):
32
                return ParsedResponse(perf_ns=response.perf_ns, data=TextResponseData(text=text))
33
            # Fallback: auto-detect common response formats
34
            if data := self.auto_detect_and_extract(json_obj):
35
                return ParsedResponse(perf_ns=response.perf_ns, data=data)
36
        if text := response.get_text():
37
            return ParsedResponse(perf_ns=response.perf_ns, data=TextResponseData(text=text))
38
        return None

1
"""Minimal Echo API server for testing the EchoEndpoint plugin."""
2
from __future__ import annotations
3
4
import asyncio
5
6
import cyclopts
7
import orjson
8
import uvicorn
9
from fastapi import FastAPI
10
from fastapi.responses import ORJSONResponse, StreamingResponse
11
12
app = FastAPI()
13
cli = cyclopts.App()
14
15
@app.post("/echo", response_model=None)
16
async def echo(body: dict) -> ORJSONResponse | StreamingResponse:
17
    echo_text = f"[echo] {body.get('text', '')}"
18
    model = body.get("model", "echo-model")
19
20
    if not body.get("stream"):
21
        return ORJSONResponse({"echo": echo_text, "model": model})
22
23
    async def sse():
24
        for i, word in enumerate(echo_text.split()):
25
            chunk = orjson.dumps({"echo": word if i == 0 else f" {word}", "model": model})
26
            yield b"data: " + chunk + b"\n\n"
27
            await asyncio.sleep(0.02)
28
        yield b"data: [DONE]\n\n"
29
30
    return StreamingResponse(sse(), media_type="text/event-stream")
31
32
33
@cli.default
34
def main(host: str = "127.0.0.1", port: int = 8000) -> None:
35
    uvicorn.run(app, host=host, port=port)
36
37
38
if __name__ == "__main__":
39
    cli()

1
"""Tests for the Echo endpoint."""
2
import pytest
3
from my_plugins.endpoints.echo_endpoint import EchoEndpoint
4
5
6
class TestEchoEndpoint:
7
    def test_format_payload(self, mock_model_endpoint, mock_request_info):
8
        endpoint = EchoEndpoint(model_endpoint=mock_model_endpoint)
9
        payload = endpoint.format_payload(mock_request_info)
10
        assert "text" in payload and "model" in payload
11
12
    def test_parse_response(self, mock_model_endpoint, mock_response):
13
        endpoint = EchoEndpoint(model_endpoint=mock_model_endpoint)
14
        result = endpoint.parse_response(mock_response)
15
        assert result is not None and result.data.text