vecboost 0.3.0-rc.1

High-performance embedding vector service written in Rust
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125

# Copyright (c) 2025-2026 Kirky.X🌠
# SPDX-License-Identifier: Apache-2.0

"""API 客户端和模拟器

提供统一的 API 客户端和模拟器基类，用于 API 接口测试。
"""

from datetime import datetime
from typing import Any, Dict, List, Tuple, Type, TypeVar, Callable

from perf.config import MAX_TEXT_LENGTH

T = TypeVar("T")


class APIClient:
    """统一的 API 客户端"""

    def __init__(self, simulator: Any):
        self.simulator = simulator

    def get(self, endpoint: str) -> Tuple[int, Dict[str, Any]]:
        return self.simulator.handle_request(endpoint, "GET", {})

    def post(self, endpoint: str, data: Dict[str, Any]) -> Tuple[int, Dict[str, Any]]:
        return self.simulator.handle_request(endpoint, "POST", data)


class BaseAPISimulator:
    """API 模拟器基类"""

    def __init__(self, embedding_service: Type[T]):
        self.service = embedding_service

    def health_check(self) -> Tuple[int, Dict[str, Any]]:
        return 200, {
            "status": "healthy",
            "service": "vecboost",
            "version": "0.1.0",
            "timestamp": datetime.now().isoformat(),
        }

    def embed(self, request: Dict[str, Any]) -> Tuple[int, Dict[str, Any]]:
        if request is None:
            request = {}
        text = request.get("text", "")
        if not text:
            return 400, {
                "error": "Invalid input",
                "message": "text field is required and cannot be empty",
            }
        if len(text) > MAX_TEXT_LENGTH:
            return 400, {
                "error": "Invalid input",
                "message": f"text length exceeds maximum allowed length ({MAX_TEXT_LENGTH} characters)",
            }
        result = self.service.embed(text)
        return 200, result

    def similarity(self, request: Dict[str, Any]) -> Tuple[int, Dict[str, Any]]:
        source = request.get("source", "")
        target = request.get("target", "")
        if not source:
            return 400, {
                "error": "Invalid input",
                "message": "source field is required and cannot be empty",
            }
        if not target:
            return 400, {
                "error": "Invalid input",
                "message": "target field is required and cannot be empty",
            }
        if len(source) > MAX_TEXT_LENGTH or len(target) > MAX_TEXT_LENGTH:
            return 400, {
                "error": "Invalid input",
                "message": f"text length exceeds maximum allowed length ({MAX_TEXT_LENGTH} characters)",
            }
        result = self.service.similarity(source, target)
        return 200, result

    def embed_batch(self, request: Dict[str, Any]) -> Tuple[int, Dict[str, Any]]:
        texts = request.get("texts", [])
        if not texts:
            return 400, {
                "error": "Invalid input",
                "message": "texts field is required and cannot be empty",
            }
        if len(texts) > 100:
            return 400, {
                "error": "Invalid input",
                "message": "batch size exceeds maximum allowed (100 texts)",
            }
        if any(len(text) > MAX_TEXT_LENGTH for text in texts):
            return 400, {
                "error": "Invalid input",
                "message": f"text length exceeds maximum allowed length ({MAX_TEXT_LENGTH} characters)",
            }
        results = self.service.embed_batch(texts)
        return 200, {"embeddings": results}

    def handle_request(
        self, endpoint: str, method: str, request: Dict[str, Any]
    ) -> Tuple[int, Dict[str, Any]]:
        handler_map: Dict[Tuple[str, str], Callable[[], Tuple[int, Dict[str, Any]]]] = {
            ("/health", "GET"): self.health_check,
            ("/api/1/embed", "POST"): lambda: self.embed(request),
            ("/api/1/similarity", "POST"): lambda: self.similarity(request),
            ("/api/1/embed/batch", "POST"): lambda: self.embed_batch(request),
        }
        key = (endpoint, method)
        if key in handler_map:
            return handler_map[key]()
        return 404, {"error": "Not Found", "message": f"Endpoint {endpoint} not found"}


class TestAPISimulator(BaseAPISimulator):
    """测试 API 模拟器"""

    pass


def create_api_client(simulator: BaseAPISimulator) -> APIClient:
    """工厂函数：创建 API 客户端"""
    return APIClient(simulator)