├── requirements.txt ├── .gitignore ├── pytest.ini ├── .env.example ├── kiro_gateway ├── exceptions.py ├── __init__.py ├── utils.py ├── cache.py ├── debug_logger.py ├── http_client.py ├── models.py ├── tokenizer.py ├── config.py ├── routes.py ├── auth.py └── parsers.py ├── CLA.md ├── manual_api_test.py ├── main.py ├── tests ├── unit │ ├── test_config.py │ ├── test_routes.py │ ├── test_cache.py │ └── test_streaming.py └── integration │ └── test_full_flow.py └── README.md /requirements.txt: -------------------------------------------------------------------------------- 1 | # Prod dependencies 2 | fastapi 3 | uvicorn[standard] 4 | httpx 5 | loguru 6 | requests 7 | python-dotenv 8 | tiktoken 9 | 10 | # Testing dependencies 11 | pytest 12 | pytest-asyncio 13 | hypothesis -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Environment 2 | .env 3 | .env.local 4 | 5 | # IDE 6 | .vscode/ 7 | .idea/ 8 | .shard/ 9 | 10 | # Python 11 | __pycache__/ 12 | *.py[cod] 13 | *$py.class 14 | *.so 15 | .Python 16 | *.egg-info/ 17 | dist/ 18 | build/ 19 | 20 | # Debug logs (generated when DEBUG_LAST_REQUEST=true) 21 | debug_logs*/ 22 | debug*.json 23 | 24 | # Project-specific 25 | _notes/ 26 | requests/ 27 | 28 | # Testing 29 | .pytest_cache/ 30 | .coverage 31 | htmlcov/ -------------------------------------------------------------------------------- /pytest.ini: -------------------------------------------------------------------------------- 1 | [pytest] 2 | # Конфигурация pytest для проекта 3 | testpaths = tests 4 | python_files = test_*.py 5 | python_classes = Test* 6 | python_functions = test_* 7 | 8 | # Добавляем корневую директорию в PYTHONPATH 9 | pythonpath = . 10 | 11 | # Исключаем manual_api_test.py из автоматического запуска 12 | # (это скрипт для ручного тестирования реального API, не unit-тест) 13 | # Чтобы запустить его: python manual_api_test.py 14 | norecursedirs = .git __pycache__ old requests _notes -------------------------------------------------------------------------------- /.env.example: -------------------------------------------------------------------------------- 1 | # Kiro OpenAI Gateway - Environment Configuration 2 | # Copy this file to .env and fill in your values 3 | 4 | # =========================================== 5 | # REQUIRED 6 | # =========================================== 7 | 8 | # Password to protect YOUR proxy server 9 | # This is NOT a token from anywhere - YOU make it up! 10 | # Use this same value as api_key when connecting to your gateway 11 | # Example: "my-super-secret-password-123" or any secure string 12 | PROXY_API_KEY="my-super-secret-password-123" 13 | 14 | # =========================================== 15 | # FIRST WAY: JSON credentials file 16 | # =========================================== 17 | 18 | # Path to JSON credentials file (alternative to REFRESH_TOKEN) 19 | KIRO_CREDS_FILE="~/.aws/sso/cache/kiro-auth-token.json" 20 | 21 | # =========================================== 22 | # SECOND WAY: Kiro refresh token 23 | # =========================================== 24 | 25 | # Your Kiro refresh token obtained from Kiro IDE traffic. 26 | # (Alternative to KIRO_CREDS_FILE) 27 | # REFRESH_TOKEN="your_kiro_refresh_token_here" 28 | 29 | # =========================================== 30 | # OPTIONAL 31 | # =========================================== 32 | 33 | # AWS CodeWhisperer profile ARN (usually auto-detected) 34 | # PROFILE_ARN="arn:aws:codewhisperer:us-east-1:..." 35 | 36 | # AWS region (default: us-east-1) 37 | # KIRO_REGION="us-east-1" 38 | 39 | # =========================================== 40 | # LOGGING 41 | # =========================================== 42 | 43 | # Log level: TRACE, DEBUG, INFO, WARNING, ERROR, CRITICAL 44 | # Default: INFO (recommended for production) 45 | # Set to DEBUG for detailed troubleshooting 46 | # LOG_LEVEL="INFO" 47 | 48 | # =========================================== 49 | # FIRST TOKEN TIMEOUT (Streaming Retry) 50 | # =========================================== 51 | 52 | # Timeout for waiting for the first token from the model (in seconds). 53 | # If the model doesn't respond within this time, the request will be cancelled and retried. 54 | # This helps handle "stuck" requests when the model takes too long to start responding. 55 | # Default: 30 seconds (recommended for production) 56 | # Set a lower value (e.g., 5-10) for more aggressive retry behavior. 57 | # FIRST_TOKEN_TIMEOUT="15" 58 | 59 | # Maximum number of retry attempts when first token timeout occurs. 60 | # After exhausting all attempts, a 504 Gateway Timeout error will be returned. 61 | # Default: 3 attempts 62 | # FIRST_TOKEN_MAX_RETRIES="3" 63 | 64 | # =========================================== 65 | # DEBUG (for development only) 66 | # =========================================== 67 | 68 | # Enable debug logging of requests/responses to files 69 | # DEBUG_LAST_REQUEST=true 70 | 71 | # Directory for debug log files 72 | # DEBUG_DIR="debug_logs" 73 | -------------------------------------------------------------------------------- /kiro_gateway/exceptions.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Kiro OpenAI Gateway 4 | # Copyright (C) 2025 Jwadow 5 | # 6 | # This program is free software: you can redistribute it and/or modify 7 | # it under the terms of the GNU Affero General Public License as published by 8 | # the Free Software Foundation, either version 3 of the License, or 9 | # (at your option) any later version. 10 | # 11 | # This program is distributed in the hope that it will be useful, 12 | # but WITHOUT ANY WARRANTY; without even the implied warranty of 13 | # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 14 | # GNU Affero General Public License for more details. 15 | # 16 | # You should have received a copy of the GNU Affero General Public License 17 | # along with this program. If not, see . 18 | 19 | """ 20 | Обработчики исключений для Kiro Gateway. 21 | 22 | Содержит функции для обработки ошибок валидации и других исключений 23 | в формате, совместимом с JSON-сериализацией. 24 | """ 25 | 26 | from typing import Any, List, Dict 27 | 28 | from fastapi import Request 29 | from fastapi.exceptions import RequestValidationError 30 | from fastapi.responses import JSONResponse 31 | from loguru import logger 32 | 33 | 34 | def sanitize_validation_errors(errors: List[Dict[str, Any]]) -> List[Dict[str, Any]]: 35 | """ 36 | Преобразует ошибки валидации в JSON-сериализуемый формат. 37 | 38 | Pydantic может включать bytes объекты в поле 'input', которые 39 | не сериализуются в JSON. Эта функция конвертирует их в строки. 40 | 41 | Args: 42 | errors: Список ошибок валидации от Pydantic 43 | 44 | Returns: 45 | Список ошибок с bytes преобразованными в строки 46 | """ 47 | sanitized = [] 48 | for error in errors: 49 | sanitized_error = {} 50 | for key, value in error.items(): 51 | if isinstance(value, bytes): 52 | # Конвертируем bytes в строку 53 | sanitized_error[key] = value.decode("utf-8", errors="replace") 54 | elif isinstance(value, (list, tuple)): 55 | # Рекурсивно обрабатываем списки 56 | sanitized_error[key] = [ 57 | v.decode("utf-8", errors="replace") if isinstance(v, bytes) else v 58 | for v in value 59 | ] 60 | else: 61 | sanitized_error[key] = value 62 | sanitized.append(sanitized_error) 63 | return sanitized 64 | 65 | 66 | async def validation_exception_handler(request: Request, exc: RequestValidationError) -> JSONResponse: 67 | """ 68 | Обработчик ошибок валидации Pydantic. 69 | 70 | Логирует детали ошибки и возвращает информативный ответ. 71 | Корректно обрабатывает bytes объекты в ошибках, преобразуя их в строки. 72 | 73 | Args: 74 | request: FastAPI Request объект 75 | exc: Исключение валидации от Pydantic 76 | 77 | Returns: 78 | JSONResponse с деталями ошибки и статусом 422 79 | """ 80 | body = await request.body() 81 | body_str = body.decode("utf-8", errors="replace") 82 | 83 | # Санитизируем ошибки для JSON-сериализации 84 | sanitized_errors = sanitize_validation_errors(exc.errors()) 85 | 86 | logger.error(f"Validation error (422): {sanitized_errors}") 87 | logger.error(f"Request body: {body_str[:500]}...") 88 | 89 | return JSONResponse( 90 | status_code=422, 91 | content={"detail": sanitized_errors, "body": body_str[:500]}, 92 | ) -------------------------------------------------------------------------------- /kiro_gateway/__init__.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Kiro OpenAI Gateway 4 | # Copyright (C) 2025 Jwadow 5 | # 6 | # This program is free software: you can redistribute it and/or modify 7 | # it under the terms of the GNU Affero General Public License as published by 8 | # the Free Software Foundation, either version 3 of the License, or 9 | # (at your option) any later version. 10 | # 11 | # This program is distributed in the hope that it will be useful, 12 | # but WITHOUT ANY WARRANTY; without even the implied warranty of 13 | # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 14 | # GNU Affero General Public License for more details. 15 | # 16 | # You should have received a copy of the GNU Affero General Public License 17 | # along with this program. If not, see . 18 | 19 | """ 20 | Kiro Gateway - OpenAI-совместимый прокси для Kiro API. 21 | 22 | Этот пакет предоставляет модульную архитектуру для проксирования 23 | запросов OpenAI API к Kiro (AWS CodeWhisperer). 24 | 25 | Модули: 26 | - config: Конфигурация и константы 27 | - models: Pydantic модели для OpenAI API 28 | - auth: Менеджер аутентификации Kiro 29 | - cache: Кэш метаданных моделей 30 | - utils: Вспомогательные утилиты 31 | - converters: Конвертация OpenAI <-> Kiro форматов 32 | - parsers: Парсеры AWS SSE потоков 33 | - streaming: Логика стриминга ответов 34 | - http_client: HTTP клиент с retry логикой 35 | - routes: FastAPI роуты 36 | - exceptions: Обработчики исключений 37 | """ 38 | 39 | # Версия импортируется из config.py — единственного источника истины (Single Source of Truth) 40 | # Это позволяет менять версию только в одном месте 41 | from kiro_gateway.config import APP_VERSION as __version__ 42 | 43 | __author__ = "Jwadow" 44 | 45 | # Основные компоненты для удобного импорта 46 | from kiro_gateway.auth import KiroAuthManager 47 | from kiro_gateway.cache import ModelInfoCache 48 | from kiro_gateway.http_client import KiroHttpClient 49 | from kiro_gateway.routes import router 50 | 51 | # Конфигурация 52 | from kiro_gateway.config import ( 53 | PROXY_API_KEY, 54 | REGION, 55 | MODEL_MAPPING, 56 | AVAILABLE_MODELS, 57 | APP_VERSION, 58 | ) 59 | 60 | # Модели 61 | from kiro_gateway.models import ( 62 | ChatCompletionRequest, 63 | ChatMessage, 64 | OpenAIModel, 65 | ModelList, 66 | ) 67 | 68 | # Конвертеры 69 | from kiro_gateway.converters import ( 70 | build_kiro_payload, 71 | extract_text_content, 72 | merge_adjacent_messages, 73 | ) 74 | 75 | # Парсеры 76 | from kiro_gateway.parsers import ( 77 | AwsEventStreamParser, 78 | parse_bracket_tool_calls, 79 | ) 80 | 81 | # Streaming 82 | from kiro_gateway.streaming import ( 83 | stream_kiro_to_openai, 84 | collect_stream_response, 85 | ) 86 | 87 | # Exceptions 88 | from kiro_gateway.exceptions import ( 89 | validation_exception_handler, 90 | sanitize_validation_errors, 91 | ) 92 | 93 | __all__ = [ 94 | # Версия 95 | "__version__", 96 | 97 | # Основные классы 98 | "KiroAuthManager", 99 | "ModelInfoCache", 100 | "KiroHttpClient", 101 | "router", 102 | 103 | # Конфигурация 104 | "PROXY_API_KEY", 105 | "REGION", 106 | "MODEL_MAPPING", 107 | "AVAILABLE_MODELS", 108 | "APP_VERSION", 109 | 110 | # Модели 111 | "ChatCompletionRequest", 112 | "ChatMessage", 113 | "OpenAIModel", 114 | "ModelList", 115 | 116 | # Конвертеры 117 | "build_kiro_payload", 118 | "extract_text_content", 119 | "merge_adjacent_messages", 120 | 121 | # Парсеры 122 | "AwsEventStreamParser", 123 | "parse_bracket_tool_calls", 124 | 125 | # Streaming 126 | "stream_kiro_to_openai", 127 | "collect_stream_response", 128 | 129 | # Exceptions 130 | "validation_exception_handler", 131 | "sanitize_validation_errors", 132 | ] -------------------------------------------------------------------------------- /kiro_gateway/utils.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Kiro OpenAI Gateway 4 | # Copyright (C) 2025 Jwadow 5 | # 6 | # This program is free software: you can redistribute it and/or modify 7 | # it under the terms of the GNU Affero General Public License as published by 8 | # the Free Software Foundation, either version 3 of the License, or 9 | # (at your option) any later version. 10 | # 11 | # This program is distributed in the hope that it will be useful, 12 | # but WITHOUT ANY WARRANTY; without even the implied warranty of 13 | # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 14 | # GNU Affero General Public License for more details. 15 | # 16 | # You should have received a copy of the GNU Affero General Public License 17 | # along with this program. If not, see . 18 | 19 | """ 20 | Вспомогательные утилиты для Kiro Gateway. 21 | 22 | Содержит функции для генерации fingerprint, формирования заголовков 23 | и другие общие утилиты. 24 | """ 25 | 26 | import hashlib 27 | import uuid 28 | from typing import TYPE_CHECKING 29 | 30 | from loguru import logger 31 | 32 | if TYPE_CHECKING: 33 | from kiro_gateway.auth import KiroAuthManager 34 | 35 | 36 | def get_machine_fingerprint() -> str: 37 | """ 38 | Генерирует уникальный fingerprint машины на основе hostname и username. 39 | 40 | Используется для формирования User-Agent, чтобы идентифицировать 41 | конкретную установку gateway. 42 | 43 | Returns: 44 | SHA256 хеш строки "{hostname}-{username}-kiro-gateway" 45 | """ 46 | try: 47 | import socket 48 | import getpass 49 | 50 | hostname = socket.gethostname() 51 | username = getpass.getuser() 52 | unique_string = f"{hostname}-{username}-kiro-gateway" 53 | 54 | return hashlib.sha256(unique_string.encode()).hexdigest() 55 | except Exception as e: 56 | logger.warning(f"Failed to get machine fingerprint: {e}") 57 | return hashlib.sha256(b"default-kiro-gateway").hexdigest() 58 | 59 | 60 | def get_kiro_headers(auth_manager: "KiroAuthManager", token: str) -> dict: 61 | """ 62 | Формирует заголовки для запросов к Kiro API. 63 | 64 | Включает все необходимые заголовки для аутентификации и идентификации: 65 | - Authorization с Bearer токеном 66 | - User-Agent с fingerprint 67 | - Специфичные для AWS CodeWhisperer заголовки 68 | 69 | Args: 70 | auth_manager: Менеджер аутентификации для получения fingerprint 71 | token: Access token для авторизации 72 | 73 | Returns: 74 | Словарь с заголовками для HTTP запроса 75 | """ 76 | fingerprint = auth_manager.fingerprint 77 | 78 | return { 79 | "Authorization": f"Bearer {token}", 80 | "Content-Type": "application/json", 81 | "User-Agent": f"aws-sdk-js/1.0.27 ua/2.1 os/win32#10.0.19044 lang/js md/nodejs#22.21.1 api/codewhispererstreaming#1.0.27 m/E KiroGateway-{fingerprint[:32]}", 82 | "x-amz-user-agent": f"aws-sdk-js/1.0.27 KiroGateway-{fingerprint[:32]}", 83 | "x-amzn-codewhisperer-optout": "true", 84 | "x-amzn-kiro-agent-mode": "vibe", 85 | "amz-sdk-invocation-id": str(uuid.uuid4()), 86 | "amz-sdk-request": "attempt=1; max=3", 87 | } 88 | 89 | 90 | def generate_completion_id() -> str: 91 | """ 92 | Генерирует уникальный ID для chat completion. 93 | 94 | Returns: 95 | ID в формате "chatcmpl-{uuid_hex}" 96 | """ 97 | return f"chatcmpl-{uuid.uuid4().hex}" 98 | 99 | 100 | def generate_conversation_id() -> str: 101 | """ 102 | Генерирует уникальный ID для разговора. 103 | 104 | Returns: 105 | UUID в строковом формате 106 | """ 107 | return str(uuid.uuid4()) 108 | 109 | 110 | def generate_tool_call_id() -> str: 111 | """ 112 | Генерирует уникальный ID для tool call. 113 | 114 | Returns: 115 | ID в формате "call_{uuid_hex[:8]}" 116 | """ 117 | return f"call_{uuid.uuid4().hex[:8]}" -------------------------------------------------------------------------------- /kiro_gateway/cache.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Kiro OpenAI Gateway 4 | # Copyright (C) 2025 Jwadow 5 | # 6 | # This program is free software: you can redistribute it and/or modify 7 | # it under the terms of the GNU Affero General Public License as published by 8 | # the Free Software Foundation, either version 3 of the License, or 9 | # (at your option) any later version. 10 | # 11 | # This program is distributed in the hope that it will be useful, 12 | # but WITHOUT ANY WARRANTY; without even the implied warranty of 13 | # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 14 | # GNU Affero General Public License for more details. 15 | # 16 | # You should have received a copy of the GNU Affero General Public License 17 | # along with this program. If not, see . 18 | 19 | """ 20 | Кэш метаданных моделей для Kiro Gateway. 21 | 22 | Потокобезопасное хранилище информации о доступных моделях 23 | с поддержкой TTL и lazy loading. 24 | """ 25 | 26 | import asyncio 27 | import time 28 | from typing import Any, Dict, List, Optional 29 | 30 | from loguru import logger 31 | 32 | from kiro_gateway.config import MODEL_CACHE_TTL, DEFAULT_MAX_INPUT_TOKENS 33 | 34 | 35 | class ModelInfoCache: 36 | """ 37 | Потокобезопасный кэш для хранения метаданных о моделях. 38 | 39 | Использует Lazy Loading для заполнения - данные загружаются 40 | только при первом обращении или когда кэш устарел. 41 | 42 | Attributes: 43 | cache_ttl: Время жизни кэша в секундах 44 | 45 | Example: 46 | >>> cache = ModelInfoCache() 47 | >>> await cache.update([{"modelId": "claude-sonnet-4", "tokenLimits": {...}}]) 48 | >>> info = cache.get("claude-sonnet-4") 49 | >>> max_tokens = cache.get_max_input_tokens("claude-sonnet-4") 50 | """ 51 | 52 | def __init__(self, cache_ttl: int = MODEL_CACHE_TTL): 53 | """ 54 | Инициализирует кэш моделей. 55 | 56 | Args: 57 | cache_ttl: Время жизни кэша в секундах (по умолчанию из конфига) 58 | """ 59 | self._cache: Dict[str, Dict[str, Any]] = {} 60 | self._lock = asyncio.Lock() 61 | self._last_update: Optional[float] = None 62 | self._cache_ttl = cache_ttl 63 | 64 | async def update(self, models_data: List[Dict[str, Any]]) -> None: 65 | """ 66 | Обновляет кэш моделей. 67 | 68 | Потокобезопасно заменяет содержимое кэша новыми данными. 69 | 70 | Args: 71 | models_data: Список словарей с информацией о моделях. 72 | Каждый словарь должен содержать ключ "modelId". 73 | """ 74 | async with self._lock: 75 | logger.info(f"Updating model cache. Found {len(models_data)} models.") 76 | self._cache = {model["modelId"]: model for model in models_data} 77 | self._last_update = time.time() 78 | 79 | def get(self, model_id: str) -> Optional[Dict[str, Any]]: 80 | """ 81 | Возвращает информацию о модели. 82 | 83 | Args: 84 | model_id: ID модели 85 | 86 | Returns: 87 | Словарь с информацией о модели или None если модель не найдена 88 | """ 89 | return self._cache.get(model_id) 90 | 91 | def get_max_input_tokens(self, model_id: str) -> int: 92 | """ 93 | Возвращает maxInputTokens для модели. 94 | 95 | Args: 96 | model_id: ID модели 97 | 98 | Returns: 99 | Максимальное количество input токенов или DEFAULT_MAX_INPUT_TOKENS 100 | """ 101 | model = self._cache.get(model_id) 102 | if model and model.get("tokenLimits"): 103 | return model["tokenLimits"].get("maxInputTokens") or DEFAULT_MAX_INPUT_TOKENS 104 | return DEFAULT_MAX_INPUT_TOKENS 105 | 106 | def is_empty(self) -> bool: 107 | """ 108 | Проверяет, пуст ли кэш. 109 | 110 | Returns: 111 | True если кэш пуст 112 | """ 113 | return not self._cache 114 | 115 | def is_stale(self) -> bool: 116 | """ 117 | Проверяет, устарел ли кэш. 118 | 119 | Returns: 120 | True если кэш устарел (прошло больше cache_ttl секунд) 121 | или если кэш никогда не обновлялся 122 | """ 123 | if not self._last_update: 124 | return True 125 | return time.time() - self._last_update > self._cache_ttl 126 | 127 | def get_all_model_ids(self) -> List[str]: 128 | """ 129 | Возвращает список всех ID моделей в кэше. 130 | 131 | Returns: 132 | Список ID моделей 133 | """ 134 | return list(self._cache.keys()) 135 | 136 | @property 137 | def size(self) -> int: 138 | """Количество моделей в кэше.""" 139 | return len(self._cache) 140 | 141 | @property 142 | def last_update_time(self) -> Optional[float]: 143 | """Время последнего обновления (timestamp) или None.""" 144 | return self._last_update -------------------------------------------------------------------------------- /kiro_gateway/debug_logger.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Kiro OpenAI Gateway 4 | # Copyright (C) 2025 Jwadow 5 | # 6 | # This program is free software: you can redistribute it and/or modify 7 | # it under the terms of the GNU Affero General Public License as published by 8 | # the Free Software Foundation, either version 3 of the License, or 9 | # (at your option) any later version. 10 | # 11 | # This program is distributed in the hope that it will be useful, 12 | # but WITHOUT ANY WARRANTY; without even the implied warranty of 13 | # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 14 | # GNU Affero General Public License for more details. 15 | # 16 | # You should have received a copy of the GNU Affero General Public License 17 | # along with this program. If not, see . 18 | 19 | """ 20 | Модуль для отладочного логирования последнего запроса. 21 | 22 | Сохраняет данные запроса и потоки ответов в файлы для последующего анализа. 23 | Активен только когда DEBUG_LAST_REQUEST=true в окружении. 24 | """ 25 | 26 | import json 27 | import shutil 28 | from pathlib import Path 29 | from loguru import logger 30 | 31 | from kiro_gateway.config import DEBUG_LAST_REQUEST, DEBUG_DIR 32 | 33 | 34 | class DebugLogger: 35 | """ 36 | Синглтон для управления отладочными логами последнего запроса. 37 | """ 38 | _instance = None 39 | 40 | def __new__(cls): 41 | if cls._instance is None: 42 | cls._instance = super(DebugLogger, cls).__new__(cls) 43 | cls._instance._initialized = False 44 | return cls._instance 45 | 46 | def __init__(self): 47 | if self._initialized: 48 | return 49 | self.debug_dir = Path(DEBUG_DIR) 50 | self._initialized = True 51 | 52 | def prepare_new_request(self): 53 | """ 54 | Очищает папку с логами и создает её заново для нового запроса. 55 | """ 56 | if not DEBUG_LAST_REQUEST: 57 | return 58 | 59 | try: 60 | if self.debug_dir.exists(): 61 | shutil.rmtree(self.debug_dir) 62 | self.debug_dir.mkdir(parents=True, exist_ok=True) 63 | logger.debug(f"[DebugLogger] Directory {self.debug_dir} cleared for new request.") 64 | except Exception as e: 65 | logger.error(f"[DebugLogger] Error preparing directory: {e}") 66 | 67 | def log_request_body(self, body: bytes): 68 | """ 69 | Сохраняет тело запроса (от клиента, OpenAI формат) в JSON файл. 70 | """ 71 | if not DEBUG_LAST_REQUEST: 72 | return 73 | 74 | try: 75 | file_path = self.debug_dir / "request_body.json" 76 | # Пытаемся сохранить как красивый JSON 77 | try: 78 | json_obj = json.loads(body) 79 | with open(file_path, "w", encoding="utf-8") as f: 80 | json.dump(json_obj, f, indent=2, ensure_ascii=False) 81 | except json.JSONDecodeError: 82 | # Если не JSON, пишем как есть (байты -> строка, если получится) 83 | with open(file_path, "wb") as f: 84 | f.write(body) 85 | except Exception as e: 86 | logger.error(f"[DebugLogger] Error writing request_body: {e}") 87 | 88 | def log_kiro_request_body(self, body: bytes): 89 | """ 90 | Сохраняет модифицированное тело запроса (к Kiro API) в JSON файл. 91 | """ 92 | if not DEBUG_LAST_REQUEST: 93 | return 94 | 95 | try: 96 | file_path = self.debug_dir / "kiro_request_body.json" 97 | # Пытаемся сохранить как красивый JSON 98 | try: 99 | json_obj = json.loads(body) 100 | with open(file_path, "w", encoding="utf-8") as f: 101 | json.dump(json_obj, f, indent=2, ensure_ascii=False) 102 | except json.JSONDecodeError: 103 | # Если не JSON, пишем как есть (байты -> строка, если получится) 104 | with open(file_path, "wb") as f: 105 | f.write(body) 106 | except Exception as e: 107 | logger.error(f"[DebugLogger] Error writing kiro_request_body: {e}") 108 | 109 | def log_raw_chunk(self, chunk: bytes): 110 | """ 111 | Дописывает сырой чанк ответа (от провайдера) в файл. 112 | """ 113 | if not DEBUG_LAST_REQUEST: 114 | return 115 | 116 | try: 117 | file_path = self.debug_dir / "response_stream_raw.txt" 118 | with open(file_path, "ab") as f: 119 | f.write(chunk) 120 | except Exception: 121 | # Не логируем ошибку на каждый чанк, чтобы не спамить 122 | pass 123 | 124 | def log_modified_chunk(self, chunk: bytes): 125 | """ 126 | Дописывает модифицированный чанк (клиенту) в файл. 127 | """ 128 | if not DEBUG_LAST_REQUEST: 129 | return 130 | 131 | try: 132 | file_path = self.debug_dir / "response_stream_modified.txt" 133 | with open(file_path, "ab") as f: 134 | f.write(chunk) 135 | except Exception: 136 | pass 137 | 138 | 139 | # Глобальный экземпляр 140 | debug_logger = DebugLogger() -------------------------------------------------------------------------------- /CLA.md: -------------------------------------------------------------------------------- 1 | # Contributor License Agreement (CLA) 2 | 3 | **Kiro OpenAI Gateway** 4 | 5 | Version 1.0 — Effective Date: December 2025 6 | 7 | --- 8 | 9 | ## Introduction 10 | 11 | Thank you for your interest in contributing to **Kiro OpenAI Gateway** (the "Project"), maintained by **Jwadow** (the "Maintainer"). This Contributor License Agreement ("Agreement") documents the rights granted by contributors to the Maintainer. 12 | 13 | By submitting a Contribution to this Project, you accept and agree to the following terms and conditions for your present and future Contributions. 14 | 15 | --- 16 | 17 | ## 1. Definitions 18 | 19 | **"You" (or "Your")** means the copyright owner or legal entity authorized by the copyright owner that is making this Agreement with the Maintainer. 20 | 21 | **"Contribution"** means any original work of authorship, including any modifications or additions to an existing work, that is intentionally submitted by You to the Maintainer for inclusion in the Project. This includes any communication sent to the Project's repositories, issue trackers, mailing lists, or any other communication channel. 22 | 23 | **"Submitted"** means any form of electronic, verbal, or written communication sent to the Maintainer, including but not limited to communication on electronic mailing lists, source code control systems, and issue tracking systems. 24 | 25 | --- 26 | 27 | ## 2. Grant of Copyright License 28 | 29 | Subject to the terms and conditions of this Agreement, You hereby grant to the Maintainer and to recipients of software distributed by the Maintainer a perpetual, worldwide, non-exclusive, no-charge, royalty-free, irrevocable copyright license to: 30 | 31 | - Reproduce, prepare derivative works of, publicly display, publicly perform, sublicense, and distribute Your Contributions and such derivative works 32 | - Relicense the Contribution under any license, including proprietary licenses 33 | 34 | --- 35 | 36 | ## 3. Grant of Patent License 37 | 38 | Subject to the terms and conditions of this Agreement, You hereby grant to the Maintainer and to recipients of software distributed by the Maintainer a perpetual, worldwide, non-exclusive, no-charge, royalty-free, irrevocable patent license to make, have made, use, offer to sell, sell, import, and otherwise transfer the Work, where such license applies only to those patent claims licensable by You that are necessarily infringed by Your Contribution(s) alone or by combination of Your Contribution(s) with the Work to which such Contribution(s) was submitted. 39 | 40 | --- 41 | 42 | ## 4. Representations 43 | 44 | You represent that: 45 | 46 | ### 4.1 Original Work 47 | You are legally entitled to grant the above license. If your employer(s) has rights to intellectual property that you create that includes your Contributions, you represent that: 48 | - You have received permission to make Contributions on behalf of that employer 49 | - Your employer has waived such rights for your Contributions to the Maintainer 50 | - Your employer has executed a separate Corporate CLA with the Maintainer 51 | 52 | ### 4.2 Third-Party Content 53 | If your Contribution includes or is based on any third-party code, you represent that: 54 | - You have identified all such third-party code in your Contribution 55 | - You have provided complete details of any third-party license or other restriction associated with any part of your Contribution 56 | 57 | ### 4.3 No Conflicts 58 | Your Contribution does not violate any agreement or obligation you have with any third party. 59 | 60 | --- 61 | 62 | ## 5. Support and Warranty Disclaimer 63 | 64 | You are not expected to provide support for Your Contributions, except to the extent You desire to provide support. You may provide support for free, for a fee, or not at all. 65 | 66 | **UNLESS REQUIRED BY APPLICABLE LAW OR AGREED TO IN WRITING, YOU PROVIDE YOUR CONTRIBUTIONS ON AN "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING, WITHOUT LIMITATION, ANY WARRANTIES OR CONDITIONS OF TITLE, NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.** 67 | 68 | --- 69 | 70 | ## 6. Notification of Changes 71 | 72 | You agree to notify the Maintainer of any facts or circumstances of which you become aware that would make these representations inaccurate in any respect. 73 | 74 | --- 75 | 76 | ## 7. Moral Rights 77 | 78 | To the fullest extent permitted under applicable law, You hereby waive, and agree not to assert, all of Your "moral rights" in or relating to Your Contributions for the benefit of the Maintainer, its assigns, and their respective direct and indirect sublicensees. 79 | 80 | --- 81 | 82 | ## 8. Governing Law 83 | 84 | This Agreement shall be governed by and construed in accordance with the laws of the jurisdiction in which the Maintainer resides, without regard to its conflict of laws provisions. 85 | 86 | --- 87 | 88 | ## 9. Entire Agreement 89 | 90 | This Agreement constitutes the entire agreement between the parties with respect to the subject matter hereof and supersedes all prior and contemporaneous agreements and understandings, whether written or oral, relating to such subject matter. 91 | 92 | --- 93 | 94 | ## How to Sign This CLA 95 | 96 | By submitting a pull request or other Contribution to this Project, you signify your acceptance of this Agreement. 97 | 98 | For significant contributions, you may be asked to explicitly confirm your acceptance by: 99 | 100 | 1. Adding your name to the [CONTRIBUTORS.md](CONTRIBUTORS.md) file (if it exists) 101 | 2. Commenting "I have read the CLA and I accept its terms" on your pull request 102 | 3. Signing via a CLA bot (if implemented) 103 | 104 | --- 105 | 106 | ## Contact 107 | 108 | If you have questions about this CLA, please open an issue in the repository or contact the Maintainer directly. 109 | 110 | **Maintainer:** Jwadow 111 | **GitHub:** [@jwadow](https://github.com/jwadow) 112 | **Project:** [Kiro OpenAI Gateway](https://github.com/jwadow/kiro-openai-gateway) 113 | 114 | --- 115 | 116 | *This CLA is based on the Apache Individual Contributor License Agreement and has been modified for this project.* -------------------------------------------------------------------------------- /manual_api_test.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Kiro OpenAI Gateway 4 | # Copyright (C) 2025 Jwadow 5 | # 6 | # This program is free software: you can redistribute it and/or modify 7 | # it under the terms of the GNU Affero General Public License as published by 8 | # the Free Software Foundation, either version 3 of the License, or 9 | # (at your option) any later version. 10 | # 11 | # This program is distributed in the hope that it will be useful, 12 | # but WITHOUT ANY WARRANTY; without even the implied warranty of 13 | # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 14 | # GNU Affero General Public License for more details. 15 | # 16 | # You should have received a copy of the GNU Affero General Public License 17 | # along with this program. If not, see . 18 | 19 | import json 20 | import os 21 | import requests 22 | from dotenv import load_dotenv 23 | 24 | # --- Загрузка переменных окружения --- 25 | load_dotenv() 26 | 27 | # --- Конфигурация --- 28 | KIRO_API_HOST = "https://q.us-east-1.amazonaws.com" 29 | TOKEN_URL = "https://prod.us-east-1.auth.desktop.kiro.dev/refreshToken" 30 | REFRESH_TOKEN = os.getenv("REFRESH_TOKEN") 31 | PROFILE_ARN = os.getenv("PROFILE_ARN", "arn:aws:codewhisperer:us-east-1:699475941385:profile/EHGA3GRVQMUK") 32 | 33 | # Глобальные переменные 34 | AUTH_TOKEN = None 35 | HEADERS = { 36 | "Authorization": None, 37 | "Content-Type": "application/json", 38 | "User-Agent": "aws-sdk-js/1.0.27 ua/2.1 os/win32#10.0.19044 lang/js md/nodejs#22.21.1 api/codewhispererstreaming#1.0.27 m/E KiroIDE-0.7.45-31c325a0ff0a9c8dec5d13048f4257462d751fe5b8af4cb1088f1fca45856c64", 39 | "x-amz-user-agent": "aws-sdk-js/1.0.27 KiroIDE-0.7.45-31c325a0ff0a9c8dec5d13048f4257462d751fe5b8af4cb1088f1fca45856c64", 40 | "x-amzn-codewhisperer-optout": "true", 41 | "x-amzn-kiro-agent-mode": "vibe", 42 | } 43 | 44 | 45 | def refresh_auth_token(): 46 | """Refreshes AUTH_TOKEN via Kiro API.""" 47 | global AUTH_TOKEN, HEADERS 48 | print("--- Refreshing Kiro token ---") 49 | 50 | payload = {"refreshToken": REFRESH_TOKEN} 51 | headers = { 52 | "Content-Type": "application/json", 53 | "User-Agent": "KiroIDE-0.7.45-31c325a0ff0a9c8dec5d13048f4257462d751fe5b8af4cb1088f1fca45856c64", 54 | } 55 | 56 | try: 57 | response = requests.post(TOKEN_URL, json=payload, headers=headers) 58 | response.raise_for_status() 59 | data = response.json() 60 | 61 | new_token = data.get("accessToken") 62 | expires_in = data.get("expiresIn") 63 | 64 | if not new_token: 65 | print("ERROR: failed to get accessToken") 66 | return False 67 | 68 | print(f"Token successfully refreshed. Expires in: {expires_in}s") 69 | AUTH_TOKEN = new_token 70 | HEADERS['Authorization'] = f"Bearer {AUTH_TOKEN}" 71 | print("--- Token refresh COMPLETED ---\n") 72 | return True 73 | 74 | except requests.exceptions.RequestException as e: 75 | print(f"ERROR refreshing token: {e}") 76 | if hasattr(e, 'response') and e.response: 77 | print(f"Server response: {e.response.status_code} {e.response.text}") 78 | return False 79 | 80 | 81 | def test_get_models(): 82 | """Tests the ListAvailableModels endpoint.""" 83 | print("--- Testing /ListAvailableModels ---") 84 | url = f"{KIRO_API_HOST}/ListAvailableModels" 85 | params = { 86 | "origin": "AI_EDITOR", 87 | "profileArn": PROFILE_ARN 88 | } 89 | 90 | try: 91 | response = requests.get(url, headers=HEADERS, params=params) 92 | response.raise_for_status() 93 | 94 | print(f"Response status: {response.status_code}") 95 | print("Response (JSON):") 96 | print(json.dumps(response.json(), indent=2, ensure_ascii=False)) 97 | print("--- ListAvailableModels test COMPLETED SUCCESSFULLY ---\n") 98 | return True 99 | except requests.exceptions.RequestException as e: 100 | print(f"ERROR: {e}") 101 | return False 102 | 103 | 104 | def test_generate_content(): 105 | """Tests the generateAssistantResponse endpoint.""" 106 | print("--- Testing /generateAssistantResponse ---") 107 | url = f"{KIRO_API_HOST}/generateAssistantResponse" 108 | 109 | import uuid 110 | payload = { 111 | "conversationState": { 112 | "agentContinuationId": str(uuid.uuid4()), 113 | "agentTaskType": "vibe", 114 | "chatTriggerType": "MANUAL", 115 | "conversationId": str(uuid.uuid4()), 116 | "currentMessage": { 117 | "userInputMessage": { 118 | "content": "Привет! Скажи что-нибудь короткое.", 119 | "modelId": "claude-haiku-4.5", 120 | "origin": "AI_EDITOR", 121 | "userInputMessageContext": { 122 | "tools": [] 123 | } 124 | } 125 | }, 126 | "history": [] 127 | }, 128 | "profileArn": PROFILE_ARN 129 | } 130 | 131 | try: 132 | with requests.post(url, headers=HEADERS, json=payload, stream=True) as response: 133 | response.raise_for_status() 134 | print(f"Response status: {response.status_code}") 135 | print("Streaming response:") 136 | 137 | for chunk in response.iter_content(chunk_size=1024): 138 | if chunk: 139 | # Пытаемся декодировать и найти JSON 140 | chunk_str = chunk.decode('utf-8', errors='ignore') 141 | print(f" Chunk: {chunk_str[:200]}...") 142 | 143 | print("\n--- generateAssistantResponse test COMPLETED ---\n") 144 | return True 145 | except requests.exceptions.RequestException as e: 146 | print(f"ERROR: {e}") 147 | return False 148 | 149 | 150 | if __name__ == "__main__": 151 | print("Starting Kiro API tests...\n") 152 | 153 | token_ok = refresh_auth_token() 154 | 155 | if token_ok: 156 | models_ok = test_get_models() 157 | generate_ok = test_generate_content() 158 | 159 | print("="*40) 160 | if models_ok and generate_ok: 161 | print("All tests passed successfully!") 162 | else: 163 | print("One or more tests failed.") 164 | else: 165 | print("="*40) 166 | print("Failed to refresh token. Tests not started.") 167 | -------------------------------------------------------------------------------- /main.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Kiro OpenAI Gateway 4 | # Copyright (C) 2025 Jwadow 5 | # 6 | # This program is free software: you can redistribute it and/or modify 7 | # it under the terms of the GNU Affero General Public License as published by 8 | # the Free Software Foundation, either version 3 of the License, or 9 | # (at your option) any later version. 10 | # 11 | # This program is distributed in the hope that it will be useful, 12 | # but WITHOUT ANY WARRANTY; without even the implied warranty of 13 | # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 14 | # GNU Affero General Public License for more details. 15 | # 16 | # You should have received a copy of the GNU Affero General Public License 17 | # along with this program. If not, see . 18 | 19 | """ 20 | Kiro API Gateway - OpenAI-compatible interface for Kiro API. 21 | 22 | Application entry point. Creates FastAPI app and connects routes. 23 | 24 | Usage: 25 | uvicorn main:app --host 0.0.0.0 --port 8000 26 | or directly: 27 | python main.py 28 | """ 29 | 30 | import sys 31 | from contextlib import asynccontextmanager 32 | from pathlib import Path 33 | 34 | from fastapi import FastAPI 35 | from fastapi.exceptions import RequestValidationError 36 | from loguru import logger 37 | 38 | from kiro_gateway.config import ( 39 | APP_TITLE, 40 | APP_DESCRIPTION, 41 | APP_VERSION, 42 | REFRESH_TOKEN, 43 | PROFILE_ARN, 44 | REGION, 45 | KIRO_CREDS_FILE, 46 | PROXY_API_KEY, 47 | LOG_LEVEL, 48 | ) 49 | from kiro_gateway.auth import KiroAuthManager 50 | from kiro_gateway.cache import ModelInfoCache 51 | from kiro_gateway.routes import router 52 | from kiro_gateway.exceptions import validation_exception_handler 53 | 54 | 55 | # --- Loguru Configuration --- 56 | logger.remove() 57 | logger.add( 58 | sys.stderr, 59 | level=LOG_LEVEL, 60 | colorize=True, 61 | format="{time:YYYY-MM-DD HH:mm:ss} | {level: <8} | {name}:{function}:{line} - {message}" 62 | ) 63 | 64 | 65 | # --- Configuration Validation --- 66 | def validate_configuration() -> None: 67 | """ 68 | Validates that required configuration is present. 69 | 70 | Checks: 71 | - .env file exists 72 | - Either REFRESH_TOKEN or KIRO_CREDS_FILE is configured 73 | 74 | Raises: 75 | SystemExit: If critical configuration is missing 76 | """ 77 | errors = [] 78 | 79 | # Check if .env file exists 80 | env_file = Path(".env") 81 | env_example = Path(".env.example") 82 | 83 | if not env_file.exists(): 84 | errors.append( 85 | ".env file not found!\n" 86 | "\n" 87 | "To get started:\n" 88 | "1. Create .env or rename from .env.example:\n" 89 | " cp .env.example .env\n" 90 | "\n" 91 | "2. Edit .env and configure your credentials:\n" 92 | " 2.1. Set you super-secret password as PROXY_API_KEY\n" 93 | " 2.2. Set your Kiro credentials:\n" 94 | " - 1 way: KIRO_CREDS_FILE to your Kiro credentials JSON file\n" 95 | " - 2 way: REFRESH_TOKEN from Kiro IDE traffic\n" 96 | "\n" 97 | "See README.md for detailed instructions." 98 | ) 99 | else: 100 | # .env exists, check for credentials 101 | has_refresh_token = bool(REFRESH_TOKEN) 102 | has_creds_file = bool(KIRO_CREDS_FILE) 103 | 104 | # Check if creds file actually exists 105 | if KIRO_CREDS_FILE: 106 | creds_path = Path(KIRO_CREDS_FILE).expanduser() 107 | if not creds_path.exists(): 108 | has_creds_file = False 109 | logger.warning(f"KIRO_CREDS_FILE not found: {KIRO_CREDS_FILE}") 110 | 111 | if not has_refresh_token and not has_creds_file: 112 | errors.append( 113 | "No Kiro credentials configured!\n" 114 | "\n" 115 | " Configure one of the following in your .env file:\n" 116 | "\n" 117 | "Set you super-secret password as PROXY_API_KEY\n" 118 | " PROXY_API_KEY=\"my-super-secret-password-123\"\n" 119 | "\n" 120 | " Option 1 (Recommended): JSON credentials file\n" 121 | " KIRO_CREDS_FILE=\"path/to/your/kiro-credentials.json\"\n" 122 | "\n" 123 | " Option 2: Refresh token\n" 124 | " REFRESH_TOKEN=\"your_refresh_token_here\"\n" 125 | "\n" 126 | " See README.md for how to obtain credentials." 127 | ) 128 | 129 | # Print errors and exit if any 130 | if errors: 131 | logger.error("") 132 | logger.error("=" * 60) 133 | logger.error(" CONFIGURATION ERROR") 134 | logger.error("=" * 60) 135 | for error in errors: 136 | for line in error.split('\n'): 137 | logger.error(f" {line}") 138 | logger.error("=" * 60) 139 | logger.error("") 140 | sys.exit(1) 141 | 142 | # Log successful configuration 143 | if KIRO_CREDS_FILE: 144 | logger.info(f"Using credentials file: {KIRO_CREDS_FILE}") 145 | elif REFRESH_TOKEN: 146 | logger.info("Using refresh token from environment") 147 | 148 | 149 | # Run configuration validation on import 150 | validate_configuration() 151 | 152 | 153 | # --- Lifespan Manager --- 154 | @asynccontextmanager 155 | async def lifespan(app: FastAPI): 156 | """ 157 | Управляет жизненным циклом приложения. 158 | 159 | Создаёт и инициализирует: 160 | - KiroAuthManager для управления токенами 161 | - ModelInfoCache для кэширования моделей 162 | """ 163 | logger.info("Starting application... Creating state managers.") 164 | 165 | # Создаём AuthManager 166 | app.state.auth_manager = KiroAuthManager( 167 | refresh_token=REFRESH_TOKEN, 168 | profile_arn=PROFILE_ARN, 169 | region=REGION, 170 | creds_file=KIRO_CREDS_FILE if KIRO_CREDS_FILE else None 171 | ) 172 | 173 | # Создаём кэш моделей 174 | app.state.model_cache = ModelInfoCache() 175 | 176 | yield 177 | 178 | logger.info("Shutting down application.") 179 | 180 | 181 | # --- FastAPI приложение --- 182 | app = FastAPI( 183 | title=APP_TITLE, 184 | description=APP_DESCRIPTION, 185 | version=APP_VERSION, 186 | lifespan=lifespan 187 | ) 188 | 189 | 190 | # --- Регистрация обработчика ошибок валидации --- 191 | app.add_exception_handler(RequestValidationError, validation_exception_handler) 192 | 193 | 194 | # --- Подключение роутов --- 195 | app.include_router(router) 196 | 197 | 198 | # --- Точка входа --- 199 | if __name__ == "__main__": 200 | import uvicorn 201 | logger.info("Starting Uvicorn server...") 202 | uvicorn.run(app, host="0.0.0.0", port=8000) 203 | -------------------------------------------------------------------------------- /tests/unit/test_config.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | """ 4 | Unit-тесты для модуля конфигурации. 5 | Проверяет загрузку настроек из переменных окружения. 6 | """ 7 | 8 | import pytest 9 | import os 10 | from unittest.mock import patch 11 | 12 | 13 | class TestLogLevelConfig: 14 | """Тесты для настройки LOG_LEVEL.""" 15 | 16 | def test_default_log_level_is_info(self): 17 | """ 18 | Что он делает: Проверяет, что LOG_LEVEL по умолчанию равен INFO. 19 | Цель: Убедиться, что без переменной окружения используется INFO. 20 | 21 | Примечание: Этот тест проверяет логику кода config.py, а не реальное 22 | значение из .env файла. Мы мокируем os.getenv чтобы симулировать 23 | отсутствие переменной окружения. 24 | """ 25 | print("Настройка: Мокируем os.getenv для LOG_LEVEL...") 26 | 27 | # Создаём мок который возвращает None для LOG_LEVEL (симулируя отсутствие переменной) 28 | original_getenv = os.getenv 29 | 30 | def mock_getenv(key, default=None): 31 | if key == "LOG_LEVEL": 32 | print(f"os.getenv('{key}') -> None (мокировано)") 33 | return default # Возвращаем default, симулируя отсутствие переменной 34 | return original_getenv(key, default) 35 | 36 | with patch.object(os, 'getenv', side_effect=mock_getenv): 37 | # Перезагружаем модуль config с мокированным getenv 38 | import importlib 39 | import kiro_gateway.config as config_module 40 | importlib.reload(config_module) 41 | 42 | print(f"LOG_LEVEL: {config_module.LOG_LEVEL}") 43 | print(f"Сравниваем: Ожидалось 'INFO', Получено '{config_module.LOG_LEVEL}'") 44 | assert config_module.LOG_LEVEL == "INFO" 45 | 46 | # Восстанавливаем модуль с реальными значениями 47 | import importlib 48 | import kiro_gateway.config as config_module 49 | importlib.reload(config_module) 50 | 51 | def test_log_level_from_environment(self): 52 | """ 53 | Что он делает: Проверяет загрузку LOG_LEVEL из переменной окружения. 54 | Цель: Убедиться, что значение из окружения используется. 55 | """ 56 | print("Настройка: Устанавливаем LOG_LEVEL=DEBUG...") 57 | 58 | with patch.dict(os.environ, {"LOG_LEVEL": "DEBUG"}): 59 | import importlib 60 | import kiro_gateway.config as config_module 61 | importlib.reload(config_module) 62 | 63 | print(f"LOG_LEVEL: {config_module.LOG_LEVEL}") 64 | print(f"Сравниваем: Ожидалось 'DEBUG', Получено '{config_module.LOG_LEVEL}'") 65 | assert config_module.LOG_LEVEL == "DEBUG" 66 | 67 | def test_log_level_uppercase_conversion(self): 68 | """ 69 | Что он делает: Проверяет преобразование LOG_LEVEL в верхний регистр. 70 | Цель: Убедиться, что lowercase значение преобразуется в uppercase. 71 | """ 72 | print("Настройка: Устанавливаем LOG_LEVEL=warning (lowercase)...") 73 | 74 | with patch.dict(os.environ, {"LOG_LEVEL": "warning"}): 75 | import importlib 76 | import kiro_gateway.config as config_module 77 | importlib.reload(config_module) 78 | 79 | print(f"LOG_LEVEL: {config_module.LOG_LEVEL}") 80 | print(f"Сравниваем: Ожидалось 'WARNING', Получено '{config_module.LOG_LEVEL}'") 81 | assert config_module.LOG_LEVEL == "WARNING" 82 | 83 | def test_log_level_trace(self): 84 | """ 85 | Что он делает: Проверяет установку LOG_LEVEL=TRACE. 86 | Цель: Убедиться, что TRACE уровень поддерживается. 87 | """ 88 | print("Настройка: Устанавливаем LOG_LEVEL=TRACE...") 89 | 90 | with patch.dict(os.environ, {"LOG_LEVEL": "TRACE"}): 91 | import importlib 92 | import kiro_gateway.config as config_module 93 | importlib.reload(config_module) 94 | 95 | print(f"LOG_LEVEL: {config_module.LOG_LEVEL}") 96 | assert config_module.LOG_LEVEL == "TRACE" 97 | 98 | def test_log_level_error(self): 99 | """ 100 | Что он делает: Проверяет установку LOG_LEVEL=ERROR. 101 | Цель: Убедиться, что ERROR уровень поддерживается. 102 | """ 103 | print("Настройка: Устанавливаем LOG_LEVEL=ERROR...") 104 | 105 | with patch.dict(os.environ, {"LOG_LEVEL": "ERROR"}): 106 | import importlib 107 | import kiro_gateway.config as config_module 108 | importlib.reload(config_module) 109 | 110 | print(f"LOG_LEVEL: {config_module.LOG_LEVEL}") 111 | assert config_module.LOG_LEVEL == "ERROR" 112 | 113 | def test_log_level_critical(self): 114 | """ 115 | Что он делает: Проверяет установку LOG_LEVEL=CRITICAL. 116 | Цель: Убедиться, что CRITICAL уровень поддерживается. 117 | """ 118 | print("Настройка: Устанавливаем LOG_LEVEL=CRITICAL...") 119 | 120 | with patch.dict(os.environ, {"LOG_LEVEL": "CRITICAL"}): 121 | import importlib 122 | import kiro_gateway.config as config_module 123 | importlib.reload(config_module) 124 | 125 | print(f"LOG_LEVEL: {config_module.LOG_LEVEL}") 126 | assert config_module.LOG_LEVEL == "CRITICAL" 127 | 128 | 129 | class TestToolDescriptionMaxLengthConfig: 130 | """Тесты для настройки TOOL_DESCRIPTION_MAX_LENGTH.""" 131 | 132 | def test_default_tool_description_max_length(self): 133 | """ 134 | Что он делает: Проверяет значение по умолчанию для TOOL_DESCRIPTION_MAX_LENGTH. 135 | Цель: Убедиться, что по умолчанию используется 10000. 136 | """ 137 | print("Настройка: Удаляем TOOL_DESCRIPTION_MAX_LENGTH из окружения...") 138 | 139 | with patch.dict(os.environ, {}, clear=False): 140 | if "TOOL_DESCRIPTION_MAX_LENGTH" in os.environ: 141 | del os.environ["TOOL_DESCRIPTION_MAX_LENGTH"] 142 | 143 | import importlib 144 | import kiro_gateway.config as config_module 145 | importlib.reload(config_module) 146 | 147 | print(f"TOOL_DESCRIPTION_MAX_LENGTH: {config_module.TOOL_DESCRIPTION_MAX_LENGTH}") 148 | assert config_module.TOOL_DESCRIPTION_MAX_LENGTH == 10000 149 | 150 | def test_tool_description_max_length_from_environment(self): 151 | """ 152 | Что он делает: Проверяет загрузку TOOL_DESCRIPTION_MAX_LENGTH из окружения. 153 | Цель: Убедиться, что значение из окружения используется. 154 | """ 155 | print("Настройка: Устанавливаем TOOL_DESCRIPTION_MAX_LENGTH=5000...") 156 | 157 | with patch.dict(os.environ, {"TOOL_DESCRIPTION_MAX_LENGTH": "5000"}): 158 | import importlib 159 | import kiro_gateway.config as config_module 160 | importlib.reload(config_module) 161 | 162 | print(f"TOOL_DESCRIPTION_MAX_LENGTH: {config_module.TOOL_DESCRIPTION_MAX_LENGTH}") 163 | assert config_module.TOOL_DESCRIPTION_MAX_LENGTH == 5000 164 | 165 | def test_tool_description_max_length_zero_disables(self): 166 | """ 167 | Что он делает: Проверяет, что 0 отключает функцию. 168 | Цель: Убедиться, что TOOL_DESCRIPTION_MAX_LENGTH=0 работает. 169 | """ 170 | print("Настройка: Устанавливаем TOOL_DESCRIPTION_MAX_LENGTH=0...") 171 | 172 | with patch.dict(os.environ, {"TOOL_DESCRIPTION_MAX_LENGTH": "0"}): 173 | import importlib 174 | import kiro_gateway.config as config_module 175 | importlib.reload(config_module) 176 | 177 | print(f"TOOL_DESCRIPTION_MAX_LENGTH: {config_module.TOOL_DESCRIPTION_MAX_LENGTH}") 178 | assert config_module.TOOL_DESCRIPTION_MAX_LENGTH == 0 -------------------------------------------------------------------------------- /kiro_gateway/http_client.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Kiro OpenAI Gateway 4 | # Copyright (C) 2025 Jwadow 5 | # 6 | # This program is free software: you can redistribute it and/or modify 7 | # it under the terms of the GNU Affero General Public License as published by 8 | # the Free Software Foundation, either version 3 of the License, or 9 | # (at your option) any later version. 10 | # 11 | # This program is distributed in the hope that it will be useful, 12 | # but WITHOUT ANY WARRANTY; without even the implied warranty of 13 | # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 14 | # GNU Affero General Public License for more details. 15 | # 16 | # You should have received a copy of the GNU Affero General Public License 17 | # along with this program. If not, see . 18 | 19 | """ 20 | HTTP клиент для Kiro API с поддержкой retry логики. 21 | 22 | Обрабатывает: 23 | - 403: автоматический refresh токена и повтор 24 | - 429: exponential backoff 25 | - 5xx: exponential backoff 26 | - Таймауты: exponential backoff 27 | """ 28 | 29 | import asyncio 30 | from typing import Optional 31 | 32 | import httpx 33 | from fastapi import HTTPException 34 | from loguru import logger 35 | 36 | from kiro_gateway.config import MAX_RETRIES, BASE_RETRY_DELAY, FIRST_TOKEN_TIMEOUT, FIRST_TOKEN_MAX_RETRIES 37 | from kiro_gateway.auth import KiroAuthManager 38 | from kiro_gateway.utils import get_kiro_headers 39 | 40 | 41 | class KiroHttpClient: 42 | """ 43 | HTTP клиент для Kiro API с поддержкой retry логики. 44 | 45 | Автоматически обрабатывает ошибки и повторяет запросы: 46 | - 403: обновляет токен и повторяет 47 | - 429: ждёт с exponential backoff 48 | - 5xx: ждёт с exponential backoff 49 | - Таймауты: ждёт с exponential backoff 50 | Attributes: 51 | auth_manager: Менеджер аутентификации для получения токенов 52 | client: HTTP клиент httpx 53 | 54 | Example: 55 | >>> client = KiroHttpClient(auth_manager) 56 | >>> response = await client.request_with_retry( 57 | ... "POST", 58 | ... "https://api.example.com/endpoint", 59 | ... {"data": "value"}, 60 | ... stream=True 61 | ... ) 62 | """ 63 | 64 | def __init__(self, auth_manager: KiroAuthManager): 65 | """ 66 | Инициализирует HTTP клиент. 67 | 68 | Args: 69 | auth_manager: Менеджер аутентификации 70 | """ 71 | self.auth_manager = auth_manager 72 | self.client: Optional[httpx.AsyncClient] = None 73 | 74 | async def _get_client(self, timeout: float = 300) -> httpx.AsyncClient: 75 | """ 76 | Возвращает или создаёт HTTP клиент. 77 | 78 | Args: 79 | timeout: Таймаут для запросов (секунды) 80 | 81 | Returns: 82 | Активный HTTP клиент 83 | """ 84 | if self.client is None or self.client.is_closed: 85 | self.client = httpx.AsyncClient(timeout=timeout, follow_redirects=True) 86 | return self.client 87 | 88 | async def close(self) -> None: 89 | """Закрывает HTTP клиент.""" 90 | if self.client and not self.client.is_closed: 91 | await self.client.aclose() 92 | 93 | async def request_with_retry( 94 | self, 95 | method: str, 96 | url: str, 97 | json_data: dict, 98 | stream: bool = False, 99 | first_token_timeout: float = None 100 | ) -> httpx.Response: 101 | """ 102 | Выполняет HTTP запрос с retry логикой. 103 | 104 | Автоматически обрабатывает различные типы ошибок: 105 | - 403: обновляет токен через auth_manager.force_refresh() и повторяет 106 | - 429: ждёт с exponential backoff (1s, 2s, 4s) 107 | - 5xx: ждёт с exponential backoff 108 | - Таймауты: ждёт с exponential backoff (для streaming - retry при first token timeout) 109 | 110 | Args: 111 | method: HTTP метод (GET, POST, etc.) 112 | url: URL запроса 113 | json_data: Тело запроса (JSON) 114 | stream: Использовать streaming (по умолчанию False) 115 | first_token_timeout: Таймаут ожидания первого ответа для streaming (секунды). 116 | Если None, используется FIRST_TOKEN_TIMEOUT из config. 117 | 118 | Returns: 119 | httpx.Response с успешным ответом 120 | 121 | Raises: 122 | HTTPException: При неудаче после всех попыток (502) 123 | """ 124 | # Для streaming используем first_token_timeout, для обычных запросов - 300 секунд 125 | if stream: 126 | timeout = first_token_timeout if first_token_timeout is not None else FIRST_TOKEN_TIMEOUT 127 | max_retries = FIRST_TOKEN_MAX_RETRIES 128 | else: 129 | timeout = 300 130 | max_retries = MAX_RETRIES 131 | 132 | client = await self._get_client(timeout) 133 | last_error = None 134 | 135 | for attempt in range(max_retries): 136 | try: 137 | # Получаем актуальный токен 138 | token = await self.auth_manager.get_access_token() 139 | headers = get_kiro_headers(self.auth_manager, token) 140 | 141 | if stream: 142 | req = client.build_request(method, url, json=json_data, headers=headers) 143 | response = await client.send(req, stream=True) 144 | else: 145 | response = await client.request(method, url, json=json_data, headers=headers) 146 | 147 | # Проверяем статус 148 | if response.status_code == 200: 149 | return response 150 | 151 | # 403 - токен истёк, обновляем и повторяем 152 | if response.status_code == 403: 153 | logger.warning(f"Received 403, refreshing token (attempt {attempt + 1}/{MAX_RETRIES})") 154 | await self.auth_manager.force_refresh() 155 | continue 156 | 157 | # 429 - rate limit, ждём и повторяем 158 | if response.status_code == 429: 159 | delay = BASE_RETRY_DELAY * (2 ** attempt) 160 | logger.warning(f"Received 429, waiting {delay}s (attempt {attempt + 1}/{MAX_RETRIES})") 161 | await asyncio.sleep(delay) 162 | continue 163 | 164 | # 5xx - серверная ошибка, ждём и повторяем 165 | if 500 <= response.status_code < 600: 166 | delay = BASE_RETRY_DELAY * (2 ** attempt) 167 | logger.warning(f"Received {response.status_code}, waiting {delay}s (attempt {attempt + 1}/{MAX_RETRIES})") 168 | await asyncio.sleep(delay) 169 | continue 170 | 171 | # Другие ошибки - возвращаем как есть 172 | return response 173 | 174 | except httpx.TimeoutException as e: 175 | last_error = e 176 | if stream: 177 | # Для streaming - это first token timeout, retry без задержки 178 | logger.warning(f"First token timeout after {timeout}s (attempt {attempt + 1}/{max_retries})") 179 | else: 180 | delay = BASE_RETRY_DELAY * (2 ** attempt) 181 | logger.warning(f"Timeout, waiting {delay}s (attempt {attempt + 1}/{max_retries})") 182 | await asyncio.sleep(delay) 183 | 184 | except httpx.RequestError as e: 185 | last_error = e 186 | delay = BASE_RETRY_DELAY * (2 ** attempt) 187 | logger.warning(f"Request error: {e}, waiting {delay}s (attempt {attempt + 1}/{max_retries})") 188 | await asyncio.sleep(delay) 189 | 190 | # Все попытки исчерпаны 191 | if stream: 192 | raise HTTPException( 193 | status_code=504, 194 | detail=f"Model did not respond within {timeout}s after {max_retries} attempts. Please try again." 195 | ) 196 | else: 197 | raise HTTPException( 198 | status_code=502, 199 | detail=f"Failed to complete request after {max_retries} attempts: {last_error}" 200 | ) 201 | 202 | async def __aenter__(self) -> "KiroHttpClient": 203 | """Поддержка async context manager.""" 204 | return self 205 | 206 | async def __aexit__(self, exc_type, exc_val, exc_tb) -> None: 207 | """Закрывает клиент при выходе из контекста.""" 208 | await self.close() -------------------------------------------------------------------------------- /kiro_gateway/models.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Kiro OpenAI Gateway 4 | # Copyright (C) 2025 Jwadow 5 | # 6 | # This program is free software: you can redistribute it and/or modify 7 | # it under the terms of the GNU Affero General Public License as published by 8 | # the Free Software Foundation, either version 3 of the License, or 9 | # (at your option) any later version. 10 | # 11 | # This program is distributed in the hope that it will be useful, 12 | # but WITHOUT ANY WARRANTY; without even the implied warranty of 13 | # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 14 | # GNU Affero General Public License for more details. 15 | # 16 | # You should have received a copy of the GNU Affero General Public License 17 | # along with this program. If not, see . 18 | 19 | """ 20 | Pydantic модели для OpenAI-совместимого API. 21 | 22 | Определяет схемы данных для запросов и ответов, 23 | обеспечивая валидацию и сериализацию. 24 | """ 25 | 26 | import time 27 | from typing import Any, Dict, List, Optional, Union 28 | from typing_extensions import Annotated 29 | from pydantic import BaseModel, Field 30 | 31 | 32 | # ================================================================================================== 33 | # Модели для /v1/models endpoint 34 | # ================================================================================================== 35 | 36 | class OpenAIModel(BaseModel): 37 | """ 38 | Модель данных для описания AI модели в формате OpenAI. 39 | 40 | Используется в ответе эндпоинта /v1/models. 41 | """ 42 | id: str 43 | object: str = "model" 44 | created: int = Field(default_factory=lambda: int(time.time())) 45 | owned_by: str = "anthropic" 46 | description: Optional[str] = None 47 | 48 | 49 | class ModelList(BaseModel): 50 | """ 51 | Список моделей в формате OpenAI. 52 | 53 | Ответ эндпоинта GET /v1/models. 54 | """ 55 | object: str = "list" 56 | data: List[OpenAIModel] 57 | 58 | 59 | # ================================================================================================== 60 | # Модели для /v1/chat/completions endpoint 61 | # ================================================================================================== 62 | 63 | class ChatMessage(BaseModel): 64 | """ 65 | Сообщение в чате в формате OpenAI. 66 | 67 | Поддерживает различные роли (user, assistant, system, tool) 68 | и различные форматы контента (строка, список, объект). 69 | 70 | Attributes: 71 | role: Роль отправителя (user, assistant, system, tool) 72 | content: Содержимое сообщения (может быть строкой, списком или None) 73 | name: Опциональное имя отправителя 74 | tool_calls: Список вызовов инструментов (для assistant) 75 | tool_call_id: ID вызова инструмента (для tool) 76 | """ 77 | role: str 78 | content: Optional[Union[str, List[Any], Any]] = None 79 | name: Optional[str] = None 80 | tool_calls: Optional[List[Any]] = None 81 | tool_call_id: Optional[str] = None 82 | 83 | model_config = {"extra": "allow"} 84 | 85 | 86 | class ToolFunction(BaseModel): 87 | """ 88 | Описание функции инструмента. 89 | 90 | Attributes: 91 | name: Имя функции 92 | description: Описание функции 93 | parameters: JSON Schema параметров функции 94 | """ 95 | name: str 96 | description: Optional[str] = None 97 | parameters: Optional[Dict[str, Any]] = None 98 | 99 | 100 | class Tool(BaseModel): 101 | """ 102 | Инструмент (tool) в формате OpenAI. 103 | 104 | Attributes: 105 | type: Тип инструмента (обычно "function") 106 | function: Описание функции 107 | """ 108 | type: str = "function" 109 | function: ToolFunction 110 | 111 | 112 | class ChatCompletionRequest(BaseModel): 113 | """ 114 | Запрос на генерацию ответа в формате OpenAI Chat Completions API. 115 | 116 | Поддерживает все стандартные поля OpenAI API, включая: 117 | - Базовые параметры (model, messages, stream) 118 | - Параметры генерации (temperature, top_p, max_tokens) 119 | - Tools (function calling) 120 | - Дополнительные параметры (игнорируются, но принимаются для совместимости) 121 | 122 | Attributes: 123 | model: ID модели для генерации 124 | messages: Список сообщений чата 125 | stream: Использовать streaming (по умолчанию False) 126 | temperature: Температура генерации (0-2) 127 | top_p: Top-p sampling 128 | n: Количество вариантов ответа 129 | max_tokens: Максимальное количество токенов в ответе 130 | max_completion_tokens: Альтернативное поле для max_tokens 131 | stop: Стоп-последовательности 132 | presence_penalty: Штраф за повторение тем 133 | frequency_penalty: Штраф за повторение слов 134 | tools: Список доступных инструментов 135 | tool_choice: Стратегия выбора инструмента 136 | """ 137 | model: str 138 | messages: Annotated[List[ChatMessage], Field(min_length=1)] 139 | stream: bool = False 140 | 141 | # Параметры генерации 142 | temperature: Optional[float] = None 143 | top_p: Optional[float] = None 144 | n: Optional[int] = 1 145 | max_tokens: Optional[int] = None 146 | max_completion_tokens: Optional[int] = None 147 | stop: Optional[Union[str, List[str]]] = None 148 | presence_penalty: Optional[float] = None 149 | frequency_penalty: Optional[float] = None 150 | 151 | # Tools (function calling) 152 | tools: Optional[List[Tool]] = None 153 | tool_choice: Optional[Union[str, Dict]] = None 154 | 155 | # Поля для совместимости (игнорируются) 156 | stream_options: Optional[Dict[str, Any]] = None 157 | logit_bias: Optional[Dict[str, float]] = None 158 | logprobs: Optional[bool] = None 159 | top_logprobs: Optional[int] = None 160 | user: Optional[str] = None 161 | seed: Optional[int] = None 162 | parallel_tool_calls: Optional[bool] = None 163 | 164 | model_config = {"extra": "allow"} 165 | 166 | 167 | # ================================================================================================== 168 | # Модели для ответов 169 | # ================================================================================================== 170 | 171 | class ChatCompletionChoice(BaseModel): 172 | """ 173 | Один вариант ответа в Chat Completion. 174 | 175 | Attributes: 176 | index: Индекс варианта 177 | message: Сообщение ответа 178 | finish_reason: Причина завершения (stop, tool_calls, length) 179 | """ 180 | index: int = 0 181 | message: Dict[str, Any] 182 | finish_reason: Optional[str] = None 183 | 184 | 185 | class ChatCompletionUsage(BaseModel): 186 | """ 187 | Информация об использовании токенов. 188 | 189 | Attributes: 190 | prompt_tokens: Количество токенов в запросе 191 | completion_tokens: Количество токенов в ответе 192 | total_tokens: Общее количество токенов 193 | credits_used: Использованные кредиты (специфично для Kiro) 194 | """ 195 | prompt_tokens: int = 0 196 | completion_tokens: int = 0 197 | total_tokens: int = 0 198 | credits_used: Optional[float] = None 199 | 200 | 201 | class ChatCompletionResponse(BaseModel): 202 | """ 203 | Полный ответ Chat Completion (non-streaming). 204 | 205 | Attributes: 206 | id: Уникальный ID ответа 207 | object: Тип объекта ("chat.completion") 208 | created: Timestamp создания 209 | model: Использованная модель 210 | choices: Список вариантов ответа 211 | usage: Информация об использовании токенов 212 | """ 213 | id: str 214 | object: str = "chat.completion" 215 | created: int = Field(default_factory=lambda: int(time.time())) 216 | model: str 217 | choices: List[ChatCompletionChoice] 218 | usage: ChatCompletionUsage 219 | 220 | 221 | class ChatCompletionChunkDelta(BaseModel): 222 | """ 223 | Дельта изменений в streaming chunk. 224 | 225 | Attributes: 226 | role: Роль (только в первом chunk) 227 | content: Новый контент 228 | tool_calls: Новые tool calls 229 | """ 230 | role: Optional[str] = None 231 | content: Optional[str] = None 232 | tool_calls: Optional[List[Dict[str, Any]]] = None 233 | 234 | 235 | class ChatCompletionChunkChoice(BaseModel): 236 | """ 237 | Один вариант в streaming chunk. 238 | 239 | Attributes: 240 | index: Индекс варианта 241 | delta: Дельта изменений 242 | finish_reason: Причина завершения (только в последнем chunk) 243 | """ 244 | index: int = 0 245 | delta: ChatCompletionChunkDelta 246 | finish_reason: Optional[str] = None 247 | 248 | 249 | class ChatCompletionChunk(BaseModel): 250 | """ 251 | Streaming chunk в формате OpenAI. 252 | 253 | Attributes: 254 | id: Уникальный ID ответа 255 | object: Тип объекта ("chat.completion.chunk") 256 | created: Timestamp создания 257 | model: Использованная модель 258 | choices: Список вариантов 259 | usage: Информация об использовании (только в последнем chunk) 260 | """ 261 | id: str 262 | object: str = "chat.completion.chunk" 263 | created: int = Field(default_factory=lambda: int(time.time())) 264 | model: str 265 | choices: List[ChatCompletionChunkChoice] 266 | usage: Optional[ChatCompletionUsage] = None -------------------------------------------------------------------------------- /kiro_gateway/tokenizer.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Kiro OpenAI Gateway 4 | # Copyright (C) 2025 Jwadow 5 | # 6 | # This program is free software: you can redistribute it and/or modify 7 | # it under the terms of the GNU Affero General Public License as published by 8 | # the Free Software Foundation, either version 3 of the License, or 9 | # (at your option) any later version. 10 | # 11 | # This program is distributed in the hope that it will be useful, 12 | # but WITHOUT ANY WARRANTY; without even the implied warranty of 13 | # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 14 | # GNU Affero General Public License for more details. 15 | # 16 | # You should have received a copy of the GNU Affero General Public License 17 | # along with this program. If not, see . 18 | 19 | """ 20 | Модуль для быстрого подсчёта токенов. 21 | 22 | Использует tiktoken (библиотека OpenAI на Rust) для приблизительного 23 | подсчёта токенов. Кодировка cl100k_base близка к токенизации Claude. 24 | 25 | Примечание: Это приблизительный подсчёт, так как точный токенизатор 26 | Claude не является публичным. Anthropic не публикует свой токенизатор, 27 | поэтому используется tiktoken с коэффициентом коррекции. 28 | 29 | Коэффициент коррекции CLAUDE_CORRECTION_FACTOR = 1.15 основан на 30 | эмпирических наблюдениях: Claude токенизирует текст примерно на 15% 31 | больше чем GPT-4 (cl100k_base). Это связано с различиями в BPE словарях. 32 | """ 33 | 34 | from typing import List, Dict, Any, Optional 35 | from loguru import logger 36 | 37 | # Ленивая загрузка tiktoken для ускорения импорта 38 | _encoding = None 39 | 40 | # Коэффициент коррекции для Claude моделей 41 | # Claude токенизирует текст примерно на 15% больше чем GPT-4 (cl100k_base) 42 | # Это эмпирическое значение, основанное на сравнении с context_usage от API 43 | CLAUDE_CORRECTION_FACTOR = 1.15 44 | 45 | 46 | def _get_encoding(): 47 | """ 48 | Ленивая инициализация токенизатора. 49 | 50 | Использует cl100k_base - кодировку для GPT-4/ChatGPT, 51 | которая достаточно близка к токенизации Claude. 52 | 53 | Returns: 54 | tiktoken.Encoding или None если tiktoken недоступен 55 | """ 56 | global _encoding 57 | if _encoding is None: 58 | try: 59 | import tiktoken 60 | _encoding = tiktoken.get_encoding("cl100k_base") 61 | logger.debug("[Tokenizer] Initialized tiktoken with cl100k_base encoding") 62 | except ImportError: 63 | logger.warning( 64 | "[Tokenizer] tiktoken not installed. " 65 | "Token counting will use fallback estimation. " 66 | "Install with: pip install tiktoken" 67 | ) 68 | _encoding = False # Маркер что импорт не удался 69 | except Exception as e: 70 | logger.error(f"[Tokenizer] Failed to initialize tiktoken: {e}") 71 | _encoding = False 72 | return _encoding if _encoding else None 73 | 74 | 75 | def count_tokens(text: str, apply_claude_correction: bool = True) -> int: 76 | """ 77 | Подсчитывает количество токенов в тексте. 78 | 79 | Args: 80 | text: Текст для подсчёта токенов 81 | apply_claude_correction: Применять коэффициент коррекции для Claude (по умолчанию True) 82 | 83 | Returns: 84 | Количество токенов (приблизительное, с коррекцией для Claude) 85 | """ 86 | if not text: 87 | return 0 88 | 89 | encoding = _get_encoding() 90 | if encoding: 91 | try: 92 | base_tokens = len(encoding.encode(text)) 93 | if apply_claude_correction: 94 | return int(base_tokens * CLAUDE_CORRECTION_FACTOR) 95 | return base_tokens 96 | except Exception as e: 97 | logger.warning(f"[Tokenizer] Error encoding text: {e}") 98 | 99 | # Fallback: грубая оценка ~4 символа на токен для английского, 100 | # ~2-3 символа для других языков (берём среднее ~3.5) 101 | # Для Claude добавляем коррекцию 102 | base_estimate = len(text) // 4 + 1 103 | if apply_claude_correction: 104 | return int(base_estimate * CLAUDE_CORRECTION_FACTOR) 105 | return base_estimate 106 | 107 | 108 | def count_message_tokens(messages: List[Dict[str, Any]], apply_claude_correction: bool = True) -> int: 109 | """ 110 | Подсчитывает токены в списке сообщений чата. 111 | 112 | Учитывает структуру сообщений OpenAI/Claude: 113 | - role: ~1 токен 114 | - content: токены текста 115 | - Служебные токены между сообщениями: ~3-4 токена 116 | 117 | Args: 118 | messages: Список сообщений в формате OpenAI 119 | apply_claude_correction: Применять коэффициент коррекции для Claude 120 | 121 | Returns: 122 | Приблизительное количество токенов (с коррекцией для Claude) 123 | """ 124 | if not messages: 125 | return 0 126 | 127 | total_tokens = 0 128 | 129 | for message in messages: 130 | # Базовые токены на сообщение (role, разделители) 131 | total_tokens += 4 # ~4 токена на служебную информацию 132 | 133 | # Токены роли (без коррекции, это короткие строки) 134 | role = message.get("role", "") 135 | total_tokens += count_tokens(role, apply_claude_correction=False) 136 | 137 | # Токены контента 138 | content = message.get("content") 139 | if content: 140 | if isinstance(content, str): 141 | total_tokens += count_tokens(content, apply_claude_correction=False) 142 | elif isinstance(content, list): 143 | # Мультимодальный контент (текст + изображения) 144 | for item in content: 145 | if isinstance(item, dict): 146 | if item.get("type") == "text": 147 | total_tokens += count_tokens(item.get("text", ""), apply_claude_correction=False) 148 | elif item.get("type") == "image_url": 149 | # Изображения занимают ~85-170 токенов в зависимости от размера 150 | total_tokens += 100 # Средняя оценка 151 | 152 | # Токены tool_calls (если есть) 153 | tool_calls = message.get("tool_calls") 154 | if tool_calls: 155 | for tc in tool_calls: 156 | total_tokens += 4 # Служебные токены 157 | func = tc.get("function", {}) 158 | total_tokens += count_tokens(func.get("name", ""), apply_claude_correction=False) 159 | total_tokens += count_tokens(func.get("arguments", ""), apply_claude_correction=False) 160 | 161 | # Токены tool_call_id (для ответов от инструментов) 162 | if message.get("tool_call_id"): 163 | total_tokens += count_tokens(message["tool_call_id"], apply_claude_correction=False) 164 | 165 | # Финальные служебные токены 166 | total_tokens += 3 167 | 168 | # Применяем коррекцию к общему количеству 169 | if apply_claude_correction: 170 | return int(total_tokens * CLAUDE_CORRECTION_FACTOR) 171 | return total_tokens 172 | 173 | 174 | def count_tools_tokens(tools: Optional[List[Dict[str, Any]]], apply_claude_correction: bool = True) -> int: 175 | """ 176 | Подсчитывает токены в определениях инструментов. 177 | 178 | Args: 179 | tools: Список инструментов в формате OpenAI 180 | apply_claude_correction: Применять коэффициент коррекции для Claude 181 | 182 | Returns: 183 | Приблизительное количество токенов (с коррекцией для Claude) 184 | """ 185 | if not tools: 186 | return 0 187 | 188 | total_tokens = 0 189 | 190 | for tool in tools: 191 | total_tokens += 4 # Служебные токены 192 | 193 | if tool.get("type") == "function": 194 | func = tool.get("function", {}) 195 | 196 | # Имя функции 197 | total_tokens += count_tokens(func.get("name", ""), apply_claude_correction=False) 198 | 199 | # Описание функции 200 | total_tokens += count_tokens(func.get("description", ""), apply_claude_correction=False) 201 | 202 | # Параметры (JSON schema) 203 | params = func.get("parameters") 204 | if params: 205 | import json 206 | params_str = json.dumps(params, ensure_ascii=False) 207 | total_tokens += count_tokens(params_str, apply_claude_correction=False) 208 | 209 | # Применяем коррекцию к общему количеству 210 | if apply_claude_correction: 211 | return int(total_tokens * CLAUDE_CORRECTION_FACTOR) 212 | return total_tokens 213 | 214 | 215 | def estimate_request_tokens( 216 | messages: List[Dict[str, Any]], 217 | tools: Optional[List[Dict[str, Any]]] = None, 218 | system_prompt: Optional[str] = None 219 | ) -> Dict[str, int]: 220 | """ 221 | Оценивает общее количество токенов в запросе. 222 | 223 | Args: 224 | messages: Список сообщений 225 | tools: Список инструментов (опционально) 226 | system_prompt: Системный промпт (опционально, если не в messages) 227 | 228 | Returns: 229 | Словарь с детализацией токенов: 230 | - messages_tokens: токены сообщений 231 | - tools_tokens: токены инструментов 232 | - system_tokens: токены системного промпта 233 | - total_tokens: общее количество 234 | """ 235 | messages_tokens = count_message_tokens(messages) 236 | tools_tokens = count_tools_tokens(tools) 237 | system_tokens = count_tokens(system_prompt) if system_prompt else 0 238 | 239 | return { 240 | "messages_tokens": messages_tokens, 241 | "tools_tokens": tools_tokens, 242 | "system_tokens": system_tokens, 243 | "total_tokens": messages_tokens + tools_tokens + system_tokens 244 | } -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 |
2 | 3 | # 🚀 Kiro OpenAI Gateway 4 | 5 | **OpenAI-compatible proxy gateway for Kiro IDE API (AWS CodeWhisperer)** 6 | 7 | [![License: AGPL v3](https://img.shields.io/badge/License-AGPL%20v3-blue.svg)](https://www.gnu.org/licenses/agpl-3.0) 8 | [![Python 3.10+](https://img.shields.io/badge/python-3.10+-blue.svg)](https://www.python.org/downloads/) 9 | [![FastAPI](https://img.shields.io/badge/FastAPI-0.100+-green.svg)](https://fastapi.tiangolo.com/) 10 | 11 | *Use Claude models through any tools that support the OpenAI API* 12 | 13 | [Features](#-features) • [Quick Start](#-quick-start) • [Configuration](#%EF%B8%8F-configuration) • [API Reference](#-api-reference) • [License](#-license) 14 | 15 |
16 | 17 | --- 18 | 19 | ## ✨ Features 20 | 21 | | Feature | Description | 22 | |---------|-------------| 23 | | 🔌 **OpenAI-compatible API** | Works with any OpenAI client out of the box | 24 | | 💬 **Full message history** | Passes complete conversation context | 25 | | 🛠️ **Tool Calling** | Supports function calling in OpenAI format | 26 | | 📡 **Streaming** | Full SSE streaming support | 27 | | 🔄 **Retry Logic** | Automatic retries on errors (403, 429, 5xx) | 28 | | 📋 **Extended model list** | Including versioned models | 29 | | 🔐 **Smart token management** | Automatic refresh before expiration | 30 | | 🧩 **Modular architecture** | Easy to extend with new providers | 31 | 32 | --- 33 | 34 | ## 🚀 Quick Start 35 | 36 | ### Prerequisites 37 | 38 | - Python 3.10+ 39 | - [Kiro IDE](https://kiro.dev/) with logged in account 40 | 41 | ### Installation 42 | 43 | ```bash 44 | # Clone the repository 45 | git clone https://github.com/Jwadow/kiro-openai-gateway.git 46 | cd kiro-openai-gateway 47 | 48 | # Install dependencies 49 | pip install -r requirements.txt 50 | 51 | # Configure (see Configuration section) 52 | cp .env.example .env 53 | # Edit .env with your credentials 54 | 55 | # Start the server 56 | python main.py 57 | ``` 58 | 59 | The server will be available at `http://localhost:8000` 60 | 61 | --- 62 | 63 | ## ⚙️ Configuration 64 | 65 | ### Option 1: JSON Credentials File 66 | 67 | Specify the path to the credentials file: 68 | 69 | ```env 70 | KIRO_CREDS_FILE="~/.aws/sso/cache/kiro-auth-token.json" 71 | 72 | # Password to protect YOUR proxy server (make up any secure string) 73 | # You'll use this as api_key when connecting to your gateway 74 | PROXY_API_KEY="my-super-secret-password-123" 75 | ``` 76 | 77 |
78 | 📄 JSON file format 79 | 80 | ```json 81 | { 82 | "accessToken": "eyJ...", 83 | "refreshToken": "eyJ...", 84 | "expiresAt": "2025-01-12T23:00:00.000Z", 85 | "profileArn": "arn:aws:codewhisperer:us-east-1:...", 86 | "region": "us-east-1" 87 | } 88 | ``` 89 | 90 |
91 | 92 | ### Option 2: Environment Variables (.env file) 93 | 94 | Create a `.env` file in the project root: 95 | 96 | ```env 97 | # Required 98 | REFRESH_TOKEN="your_kiro_refresh_token" 99 | 100 | # Password to protect YOUR proxy server (make up any secure string) 101 | PROXY_API_KEY="my-super-secret-password-123" 102 | 103 | # Optional 104 | PROFILE_ARN="arn:aws:codewhisperer:us-east-1:..." 105 | KIRO_REGION="us-east-1" 106 | ``` 107 | 108 | ### Getting the Refresh Token 109 | 110 | The refresh token can be obtained by intercepting Kiro IDE traffic. Look for requests to: 111 | - `prod.us-east-1.auth.desktop.kiro.dev/refreshToken` 112 | 113 | --- 114 | 115 | ## 📡 API Reference 116 | 117 | ### Endpoints 118 | 119 | | Endpoint | Method | Description | 120 | |----------|--------|-------------| 121 | | `/` | GET | Health check | 122 | | `/health` | GET | Detailed health check | 123 | | `/v1/models` | GET | List available models | 124 | | `/v1/chat/completions` | POST | Chat completions | 125 | 126 | ### Available Models 127 | 128 | | Model | Description | 129 | |-------|-------------| 130 | | `claude-opus-4-5` | Top-tier model | 131 | | `claude-opus-4-5-20251101` | Top-tier model (versioned) | 132 | | `claude-sonnet-4-5` | Enhanced model | 133 | | `claude-sonnet-4-5-20250929` | Enhanced model (versioned) | 134 | | `claude-sonnet-4` | Balanced model | 135 | | `claude-sonnet-4-20250514` | Balanced model (versioned) | 136 | | `claude-haiku-4-5` | Fast model | 137 | | `claude-3-7-sonnet-20250219` | Legacy model | 138 | 139 | --- 140 | 141 | ## 💡 Usage Examples 142 | 143 |
144 | 🔹 Simple cURL Request 145 | 146 | ```bash 147 | curl http://localhost:8000/v1/chat/completions \ 148 | -H "Authorization: Bearer my-super-secret-password-123" \ 149 | -H "Content-Type: application/json" \ 150 | -d '{ 151 | "model": "claude-sonnet-4-5", 152 | "messages": [{"role": "user", "content": "Hello!"}], 153 | "stream": true 154 | }' 155 | ``` 156 | 157 | > **Note:** Replace `my-super-secret-password-123` with the `PROXY_API_KEY` you set in your `.env` file. 158 | 159 |
160 | 161 |
162 | 🔹 Streaming Request 163 | 164 | ```bash 165 | curl http://localhost:8000/v1/chat/completions \ 166 | -H "Authorization: Bearer my-super-secret-password-123" \ 167 | -H "Content-Type: application/json" \ 168 | -d '{ 169 | "model": "claude-sonnet-4-5", 170 | "messages": [ 171 | {"role": "system", "content": "You are a helpful assistant."}, 172 | {"role": "user", "content": "What is 2+2?"} 173 | ], 174 | "stream": true 175 | }' 176 | ``` 177 | 178 |
179 | 180 |
181 | 🔹 With Tool Calling 182 | 183 | ```bash 184 | curl http://localhost:8000/v1/chat/completions \ 185 | -H "Authorization: Bearer my-super-secret-password-123" \ 186 | -H "Content-Type: application/json" \ 187 | -d '{ 188 | "model": "claude-sonnet-4-5", 189 | "messages": [{"role": "user", "content": "What is the weather in London?"}], 190 | "tools": [{ 191 | "type": "function", 192 | "function": { 193 | "name": "get_weather", 194 | "description": "Get weather for a location", 195 | "parameters": { 196 | "type": "object", 197 | "properties": { 198 | "location": {"type": "string", "description": "City name"} 199 | }, 200 | "required": ["location"] 201 | } 202 | } 203 | }] 204 | }' 205 | ``` 206 | 207 |
208 | 209 |
210 | 🐍 Python OpenAI SDK 211 | 212 | ```python 213 | from openai import OpenAI 214 | 215 | client = OpenAI( 216 | base_url="http://localhost:8000/v1", 217 | api_key="my-super-secret-password-123" # Your PROXY_API_KEY from .env 218 | ) 219 | 220 | response = client.chat.completions.create( 221 | model="claude-sonnet-4-5", 222 | messages=[ 223 | {"role": "system", "content": "You are a helpful assistant."}, 224 | {"role": "user", "content": "Hello!"} 225 | ], 226 | stream=True 227 | ) 228 | 229 | for chunk in response: 230 | if chunk.choices[0].delta.content: 231 | print(chunk.choices[0].delta.content, end="") 232 | ``` 233 | 234 |
235 | 236 |
237 | 🦜 LangChain 238 | 239 | ```python 240 | from langchain_openai import ChatOpenAI 241 | 242 | llm = ChatOpenAI( 243 | base_url="http://localhost:8000/v1", 244 | api_key="my-super-secret-password-123", # Your PROXY_API_KEY from .env 245 | model="claude-sonnet-4-5" 246 | ) 247 | 248 | response = llm.invoke("Hello, how are you?") 249 | print(response.content) 250 | ``` 251 | 252 |
253 | 254 | --- 255 | 256 | ## 📁 Project Structure 257 | 258 | ``` 259 | kiro-openai-gateway/ 260 | ├── main.py # Entry point, FastAPI app creation 261 | ├── requirements.txt # Python dependencies 262 | ├── .env.example # Environment configuration example 263 | │ 264 | ├── kiro_gateway/ # Main package 265 | │ ├── __init__.py # Package exports 266 | │ ├── config.py # Configuration and constants 267 | │ ├── models.py # Pydantic models for OpenAI API 268 | │ ├── auth.py # KiroAuthManager - token management 269 | │ ├── cache.py # ModelInfoCache - model caching 270 | │ ├── utils.py # Helper utilities 271 | │ ├── converters.py # OpenAI <-> Kiro conversion 272 | │ ├── parsers.py # AWS SSE stream parsers 273 | │ ├── streaming.py # Response streaming logic 274 | │ ├── http_client.py # HTTP client with retry logic 275 | │ ├── debug_logger.py # Debug logging (optional) 276 | │ └── routes.py # FastAPI routes 277 | │ 278 | ├── tests/ # Tests 279 | │ ├── unit/ # Unit tests 280 | │ └── integration/ # Integration tests 281 | │ 282 | └── debug_logs/ # Debug logs (generated when enabled) 283 | ``` 284 | 285 | --- 286 | 287 | ## 🔧 Debugging 288 | 289 | Debug logging is **disabled by default**. To enable, add to your `.env`: 290 | 291 | ```env 292 | DEBUG_LAST_REQUEST=true 293 | ``` 294 | 295 | When enabled, requests are logged to the `debug_logs/` folder: 296 | 297 | | File | Description | 298 | |------|-------------| 299 | | `request_body.json` | Incoming request from client | 300 | | `kiro_request_body.json` | Request to Kiro API | 301 | | `response_stream_raw.txt` | Raw stream from Kiro | 302 | | `response_stream_modified.txt` | Transformed stream | 303 | 304 | --- 305 | 306 | ## 🧪 Testing 307 | 308 | ```bash 309 | # Run all tests 310 | pytest 311 | 312 | # Run unit tests only 313 | pytest tests/unit/ 314 | 315 | # Run with coverage 316 | pytest --cov=kiro_gateway 317 | ``` 318 | 319 | --- 320 | 321 | ## 🔌 Extending with New Providers 322 | 323 | The modular architecture makes it easy to add support for other providers: 324 | 325 | 1. Create a new module `kiro_gateway/providers/new_provider.py` 326 | 2. Implement the required classes: 327 | - `NewProviderAuthManager` — token management 328 | - `NewProviderConverter` — format conversion 329 | - `NewProviderParser` — response parsing 330 | 3. Add routes to `routes.py` or create a separate router 331 | 332 | --- 333 | 334 | ## 📜 License 335 | 336 | This project is licensed under the **GNU Affero General Public License v3.0 (AGPL-3.0)**. 337 | 338 | This means: 339 | - ✅ You can use, modify, and distribute this software 340 | - ✅ You can use it for commercial purposes 341 | - ⚠️ **You must disclose source code** when you distribute the software 342 | - ⚠️ **Network use is distribution** — if you run a modified version on a server and let others interact with it, you must make the source code available to them 343 | - ⚠️ Modifications must be released under the same license 344 | 345 | See the [LICENSE](LICENSE) file for the full license text. 346 | 347 | ### Why AGPL-3.0? 348 | 349 | AGPL-3.0 ensures that improvements to this software benefit the entire community. If you modify this gateway and deploy it as a service, you must share your improvements with your users. 350 | 351 | ### Contributor License Agreement (CLA) 352 | 353 | By submitting a contribution to this project, you agree to the terms of our [Contributor License Agreement (CLA)](CLA.md). This ensures that: 354 | - You have the right to submit the contribution 355 | - You grant the maintainer rights to use and relicense your contribution 356 | - The project remains legally protected 357 | 358 | --- 359 | 360 | ## 👤 Author 361 | 362 | **Jwadow** — [@Jwadow](https://github.com/jwadow) 363 | 364 | --- 365 | 366 | ## ⚠️ Disclaimer 367 | 368 | This project is not affiliated with, endorsed by, or sponsored by Amazon Web Services (AWS), Anthropic, or Kiro IDE. Use at your own risk and in compliance with the terms of service of the underlying APIs. 369 | 370 | --- 371 | 372 |
373 | 374 | **[⬆ Back to Top](#-kiro-openai-gateway)** 375 | 376 |
377 | -------------------------------------------------------------------------------- /kiro_gateway/config.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Kiro OpenAI Gateway 4 | # Copyright (C) 2025 Jwadow 5 | # 6 | # This program is free software: you can redistribute it and/or modify 7 | # it under the terms of the GNU Affero General Public License as published by 8 | # the Free Software Foundation, either version 3 of the License, or 9 | # (at your option) any later version. 10 | # 11 | # This program is distributed in the hope that it will be useful, 12 | # but WITHOUT ANY WARRANTY; without even the implied warranty of 13 | # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 14 | # GNU Affero General Public License for more details. 15 | # 16 | # You should have received a copy of the GNU Affero General Public License 17 | # along with this program. If not, see . 18 | 19 | """ 20 | Конфигурация Kiro Gateway. 21 | 22 | Централизованное хранение всех настроек, констант и маппингов. 23 | Загружает переменные окружения и предоставляет типизированный доступ к ним. 24 | """ 25 | 26 | import os 27 | import re 28 | from pathlib import Path 29 | from typing import Dict, List, Optional 30 | from dotenv import load_dotenv 31 | 32 | # Загрузка переменных окружения 33 | load_dotenv() 34 | 35 | 36 | def _get_raw_env_value(var_name: str, env_file: str = ".env") -> Optional[str]: 37 | """ 38 | Читает значение переменной из .env файла без обработки escape-последовательностей. 39 | 40 | Это необходимо для корректной работы с путями на Windows, где обратные слэши 41 | (например, D:\\Projects\\file.json) могут быть ошибочно интерпретированы 42 | как escape-последовательности (\\a -> bell, \\n -> newline и т.д.). 43 | 44 | Args: 45 | var_name: Имя переменной окружения 46 | env_file: Путь к .env файлу (по умолчанию ".env") 47 | 48 | Returns: 49 | Сырое значение переменной или None если не найдено 50 | """ 51 | env_path = Path(env_file) 52 | if not env_path.exists(): 53 | return None 54 | 55 | try: 56 | # Читаем файл как есть, без интерпретации 57 | content = env_path.read_text(encoding="utf-8") 58 | 59 | # Ищем переменную с учётом разных форматов: 60 | # VAR="value" или VAR='value' или VAR=value 61 | # Паттерн захватывает значение в кавычках или без них 62 | pattern = rf'^{re.escape(var_name)}=(["\']?)(.+?)\1\s*$' 63 | 64 | for line in content.splitlines(): 65 | line = line.strip() 66 | if line.startswith("#") or not line: 67 | continue 68 | 69 | match = re.match(pattern, line) 70 | if match: 71 | # Возвращаем значение как есть, без обработки escape-последовательностей 72 | return match.group(2) 73 | except Exception: 74 | pass 75 | 76 | return None 77 | 78 | # ================================================================================================== 79 | # Настройки прокси-сервера 80 | # ================================================================================================== 81 | 82 | # API ключ для доступа к прокси (клиенты должны передавать его в Authorization header) 83 | PROXY_API_KEY: str = os.getenv("PROXY_API_KEY", "changeme_proxy_secret") 84 | 85 | # ================================================================================================== 86 | # Kiro API Credentials 87 | # ================================================================================================== 88 | 89 | # Refresh token для обновления access token 90 | REFRESH_TOKEN: str = os.getenv("REFRESH_TOKEN", "") 91 | 92 | # Profile ARN для AWS CodeWhisperer 93 | PROFILE_ARN: str = os.getenv("PROFILE_ARN", "") 94 | 95 | # Регион AWS (по умолчанию us-east-1) 96 | REGION: str = os.getenv("KIRO_REGION", "us-east-1") 97 | 98 | # Путь к файлу с credentials (опционально, альтернатива .env) 99 | # Читаем напрямую из .env чтобы избежать проблем с escape-последовательностями на Windows 100 | # (например, \a в пути D:\Projects\adolf интерпретируется как bell character) 101 | _raw_creds_file = _get_raw_env_value("KIRO_CREDS_FILE") or os.getenv("KIRO_CREDS_FILE", "") 102 | # Нормализуем путь для кроссплатформенной совместимости 103 | KIRO_CREDS_FILE: str = str(Path(_raw_creds_file)) if _raw_creds_file else "" 104 | 105 | # ================================================================================================== 106 | # Kiro API URL Templates 107 | # ================================================================================================== 108 | 109 | # URL для обновления токена 110 | KIRO_REFRESH_URL_TEMPLATE: str = "https://prod.{region}.auth.desktop.kiro.dev/refreshToken" 111 | 112 | # Хост для основного API (generateAssistantResponse) 113 | KIRO_API_HOST_TEMPLATE: str = "https://codewhisperer.{region}.amazonaws.com" 114 | 115 | # Хост для Q API (ListAvailableModels) 116 | KIRO_Q_HOST_TEMPLATE: str = "https://q.{region}.amazonaws.com" 117 | 118 | # ================================================================================================== 119 | # Настройки токенов 120 | # ================================================================================================== 121 | 122 | # Время до истечения токена, когда нужно обновить (в секундах) 123 | # По умолчанию 10 минут - обновляем токен заранее, чтобы избежать ошибок 124 | TOKEN_REFRESH_THRESHOLD: int = 600 125 | 126 | # ================================================================================================== 127 | # Retry конфигурация 128 | # ================================================================================================== 129 | 130 | # Максимальное количество попыток при ошибках 131 | MAX_RETRIES: int = 3 132 | 133 | # Базовая задержка между попытками (секунды) 134 | # Используется exponential backoff: delay * (2 ** attempt) 135 | BASE_RETRY_DELAY: float = 1.0 136 | 137 | # ================================================================================================== 138 | # Маппинг моделей 139 | # ================================================================================================== 140 | 141 | # Внешние имена моделей (OpenAI-совместимые) -> внутренние ID Kiro 142 | # Клиенты используют внешние имена, а мы конвертируем их во внутренние 143 | MODEL_MAPPING: Dict[str, str] = { 144 | # Claude Opus 4.5 - топовая модель 145 | "claude-opus-4-5": "claude-opus-4.5", 146 | "claude-opus-4-5-20251101": "claude-opus-4.5", 147 | 148 | # Claude Haiku 4.5 - быстрая модель 149 | "claude-haiku-4-5": "claude-haiku-4.5", 150 | "claude-haiku-4.5": "claude-haiku-4.5", # Прямой проброс 151 | 152 | # Claude Sonnet 4.5 - улучшенная модель 153 | "claude-sonnet-4-5": "CLAUDE_SONNET_4_5_20250929_V1_0", 154 | "claude-sonnet-4-5-20250929": "CLAUDE_SONNET_4_5_20250929_V1_0", 155 | 156 | # Claude Sonnet 4 - сбалансированная модель 157 | "claude-sonnet-4": "CLAUDE_SONNET_4_20250514_V1_0", 158 | "claude-sonnet-4-20250514": "CLAUDE_SONNET_4_20250514_V1_0", 159 | 160 | # Claude 3.7 Sonnet - legacy модель 161 | "claude-3-7-sonnet-20250219": "CLAUDE_3_7_SONNET_20250219_V1_0", 162 | 163 | # Алиасы для удобства 164 | "auto": "claude-sonnet-4.5", 165 | } 166 | 167 | # Список доступных моделей для эндпоинта /v1/models 168 | # Эти модели будут отображаться клиентам как доступные 169 | AVAILABLE_MODELS: List[str] = [ 170 | "claude-opus-4-5", 171 | "claude-opus-4-5-20251101", 172 | "claude-haiku-4-5", 173 | "claude-sonnet-4-5", 174 | "claude-sonnet-4-5-20250929", 175 | "claude-sonnet-4", 176 | "claude-sonnet-4-20250514", 177 | "claude-3-7-sonnet-20250219", 178 | ] 179 | 180 | # ================================================================================================== 181 | # Настройки кэша моделей 182 | # ================================================================================================== 183 | 184 | # TTL кэша моделей в секундах (1 час) 185 | MODEL_CACHE_TTL: int = 3600 186 | 187 | # Максимальное количество input токенов по умолчанию 188 | DEFAULT_MAX_INPUT_TOKENS: int = 200000 189 | 190 | # ================================================================================================== 191 | # Tool Description Handling (Kiro API Limitations) 192 | # ================================================================================================== 193 | 194 | # Kiro API возвращает ошибку 400 "Improperly formed request" при слишком длинных 195 | # описаниях инструментов в toolSpecification.description. 196 | # 197 | # Решение: Tool Documentation Reference Pattern 198 | # - Если description ≤ лимита → оставляем как есть 199 | # - Если description > лимита: 200 | # * В toolSpecification.description → ссылка на system prompt: 201 | # "[Full documentation in system prompt under '## Tool: {name}']" 202 | # * В system prompt добавляется секция "## Tool: {name}" с полным описанием 203 | # 204 | # Модель видит явную ссылку и точно понимает, где искать полную документацию. 205 | 206 | # Максимальная длина description для tool в символах. 207 | # Описания длиннее этого лимита будут перенесены в system prompt. 208 | # Установите 0 для отключения (не рекомендуется - вызовет ошибки Kiro API). 209 | TOOL_DESCRIPTION_MAX_LENGTH: int = int(os.getenv("TOOL_DESCRIPTION_MAX_LENGTH", "10000")) 210 | 211 | # ================================================================================================== 212 | # Logging Settings 213 | # ================================================================================================== 214 | 215 | # Log level for the application 216 | # Available levels: TRACE, DEBUG, INFO, WARNING, ERROR, CRITICAL 217 | # Default: INFO (recommended for production) 218 | # Set to DEBUG for detailed troubleshooting 219 | LOG_LEVEL: str = os.getenv("LOG_LEVEL", "INFO").upper() 220 | 221 | # ================================================================================================== 222 | # First Token Timeout Settings (Streaming Retry) 223 | # ================================================================================================== 224 | 225 | # Таймаут ожидания первого токена от модели (в секундах). 226 | # Если модель не отвечает в течение этого времени, запрос будет отменён и повторён. 227 | # Это помогает справиться с "зависшими" запросами, когда модель долго думает. 228 | # По умолчанию: 30 секунд (рекомендуется для production) 229 | # Установите меньшее значение (например, 10-15) для более агрессивного retry. 230 | FIRST_TOKEN_TIMEOUT: float = float(os.getenv("FIRST_TOKEN_TIMEOUT", "15")) 231 | 232 | # Максимальное количество попыток при таймауте первого токена. 233 | # После исчерпания всех попыток будет возвращена ошибка. 234 | # По умолчанию: 3 попытки 235 | FIRST_TOKEN_MAX_RETRIES: int = int(os.getenv("FIRST_TOKEN_MAX_RETRIES", "3")) 236 | 237 | # ================================================================================================== 238 | # Debug Settings 239 | # ================================================================================================== 240 | 241 | # If True, the last request will be logged in detail to DEBUG_DIR 242 | # Enable via .env: DEBUG_LAST_REQUEST=true 243 | DEBUG_LAST_REQUEST: bool = os.getenv("DEBUG_LAST_REQUEST", "false").lower() in ("true", "1", "yes") 244 | 245 | # Directory for debug log files 246 | DEBUG_DIR: str = os.getenv("DEBUG_DIR", "debug_logs") 247 | 248 | # ================================================================================================== 249 | # Версия приложения 250 | # ================================================================================================== 251 | 252 | APP_VERSION: str = "1.0.3" 253 | APP_TITLE: str = "Kiro API Gateway" 254 | APP_DESCRIPTION: str = "OpenAI-compatible interface for Kiro API (AWS CodeWhisperer). Made by @jwadow" 255 | 256 | 257 | def get_kiro_refresh_url(region: str) -> str: 258 | """Возвращает URL для обновления токена для указанного региона.""" 259 | return KIRO_REFRESH_URL_TEMPLATE.format(region=region) 260 | 261 | 262 | def get_kiro_api_host(region: str) -> str: 263 | """Возвращает хост API для указанного региона.""" 264 | return KIRO_API_HOST_TEMPLATE.format(region=region) 265 | 266 | 267 | def get_kiro_q_host(region: str) -> str: 268 | """Возвращает хост Q API для указанного региона.""" 269 | return KIRO_Q_HOST_TEMPLATE.format(region=region) 270 | 271 | 272 | def get_internal_model_id(external_model: str) -> str: 273 | """ 274 | Конвертирует внешнее имя модели во внутренний ID Kiro. 275 | 276 | Args: 277 | external_model: Внешнее имя модели (например, "claude-sonnet-4-5") 278 | 279 | Returns: 280 | Внутренний ID модели для Kiro API 281 | """ 282 | return MODEL_MAPPING.get(external_model, external_model) -------------------------------------------------------------------------------- /kiro_gateway/routes.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Kiro OpenAI Gateway 4 | # Copyright (C) 2025 Jwadow 5 | # 6 | # This program is free software: you can redistribute it and/or modify 7 | # it under the terms of the GNU Affero General Public License as published by 8 | # the Free Software Foundation, either version 3 of the License, or 9 | # (at your option) any later version. 10 | # 11 | # This program is distributed in the hope that it will be useful, 12 | # but WITHOUT ANY WARRANTY; without even the implied warranty of 13 | # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 14 | # GNU Affero General Public License for more details. 15 | # 16 | # You should have received a copy of the GNU Affero General Public License 17 | # along with this program. If not, see . 18 | 19 | """ 20 | FastAPI роуты для Kiro Gateway. 21 | 22 | Содержит все эндпоинты API: 23 | - / и /health: Health check 24 | - /v1/models: Список моделей 25 | - /v1/chat/completions: Chat completions 26 | """ 27 | 28 | import json 29 | from datetime import datetime, timezone 30 | 31 | import httpx 32 | from fastapi import APIRouter, Depends, HTTPException, Request, Response, Security 33 | from fastapi.responses import JSONResponse, StreamingResponse 34 | from fastapi.security import APIKeyHeader 35 | from loguru import logger 36 | 37 | from kiro_gateway.config import ( 38 | PROXY_API_KEY, 39 | AVAILABLE_MODELS, 40 | APP_VERSION, 41 | ) 42 | from kiro_gateway.models import ( 43 | OpenAIModel, 44 | ModelList, 45 | ChatCompletionRequest, 46 | ) 47 | from kiro_gateway.auth import KiroAuthManager 48 | from kiro_gateway.cache import ModelInfoCache 49 | from kiro_gateway.converters import build_kiro_payload 50 | from kiro_gateway.streaming import stream_kiro_to_openai, collect_stream_response, stream_with_first_token_retry 51 | from kiro_gateway.http_client import KiroHttpClient 52 | from kiro_gateway.utils import get_kiro_headers, generate_conversation_id 53 | 54 | # Импортируем debug_logger 55 | try: 56 | from kiro_gateway.debug_logger import debug_logger 57 | except ImportError: 58 | debug_logger = None 59 | 60 | 61 | # --- Схема безопасности --- 62 | api_key_header = APIKeyHeader(name="Authorization", auto_error=False) 63 | 64 | 65 | async def verify_api_key(auth_header: str = Security(api_key_header)) -> bool: 66 | """ 67 | Проверяет API ключ в заголовке Authorization. 68 | 69 | Ожидает формат: "Bearer {PROXY_API_KEY}" 70 | 71 | Args: 72 | auth_header: Значение заголовка Authorization 73 | 74 | Returns: 75 | True если ключ валиден 76 | 77 | Raises: 78 | HTTPException: 401 если ключ невалиден или отсутствует 79 | """ 80 | if not auth_header or auth_header != f"Bearer {PROXY_API_KEY}": 81 | logger.warning("Access attempt with invalid API key.") 82 | raise HTTPException(status_code=401, detail="Invalid or missing API Key") 83 | return True 84 | 85 | 86 | # --- Роутер --- 87 | router = APIRouter() 88 | 89 | 90 | @router.get("/") 91 | async def root(): 92 | """ 93 | Health check endpoint. 94 | 95 | Returns: 96 | Статус и версия приложения 97 | """ 98 | return { 99 | "status": "ok", 100 | "message": "Kiro API Gateway is running", 101 | "version": APP_VERSION 102 | } 103 | 104 | 105 | @router.get("/health") 106 | async def health(): 107 | """ 108 | Детальный health check. 109 | 110 | Returns: 111 | Статус, timestamp и версия 112 | """ 113 | return { 114 | "status": "healthy", 115 | "timestamp": datetime.now(timezone.utc).isoformat(), 116 | "version": APP_VERSION 117 | } 118 | 119 | 120 | @router.get("/v1/models", response_model=ModelList, dependencies=[Depends(verify_api_key)]) 121 | async def get_models(request: Request): 122 | """ 123 | Возвращает список доступных моделей. 124 | 125 | Использует статический список моделей с возможностью обновления из API. 126 | Кэширует результаты для уменьшения нагрузки на API. 127 | 128 | Args: 129 | request: FastAPI Request для доступа к app.state 130 | 131 | Returns: 132 | ModelList с доступными моделями 133 | """ 134 | logger.info("Request to /v1/models") 135 | 136 | auth_manager: KiroAuthManager = request.app.state.auth_manager 137 | model_cache: ModelInfoCache = request.app.state.model_cache 138 | 139 | # Пытаемся получить модели из API если кэш пуст или устарел 140 | if model_cache.is_empty() or model_cache.is_stale(): 141 | try: 142 | token = await auth_manager.get_access_token() 143 | headers = get_kiro_headers(auth_manager, token) 144 | 145 | async with httpx.AsyncClient(timeout=30) as client: 146 | response = await client.get( 147 | f"{auth_manager.q_host}/ListAvailableModels", 148 | headers=headers, 149 | params={ 150 | "origin": "AI_EDITOR", 151 | "profileArn": auth_manager.profile_arn or "" 152 | } 153 | ) 154 | 155 | if response.status_code == 200: 156 | data = response.json() 157 | models_list = data.get("models", []) 158 | await model_cache.update(models_list) 159 | logger.info(f"Received {len(models_list)} models from API") 160 | except Exception as e: 161 | logger.warning(f"Failed to fetch models from API: {e}") 162 | 163 | # Возвращаем статический список моделей 164 | openai_models = [ 165 | OpenAIModel( 166 | id=model_id, 167 | owned_by="anthropic", 168 | description="Claude model via Kiro API" 169 | ) 170 | for model_id in AVAILABLE_MODELS 171 | ] 172 | 173 | return ModelList(data=openai_models) 174 | 175 | 176 | @router.post("/v1/chat/completions", dependencies=[Depends(verify_api_key)]) 177 | async def chat_completions(request: Request, request_data: ChatCompletionRequest): 178 | """ 179 | Chat completions endpoint - совместим с OpenAI API. 180 | 181 | Принимает запросы в формате OpenAI и транслирует их в Kiro API. 182 | Поддерживает streaming и non-streaming режимы. 183 | 184 | Args: 185 | request: FastAPI Request для доступа к app.state 186 | request_data: Запрос в формате OpenAI ChatCompletionRequest 187 | 188 | Returns: 189 | StreamingResponse для streaming режима 190 | JSONResponse для non-streaming режима 191 | 192 | Raises: 193 | HTTPException: При ошибках валидации или API 194 | """ 195 | logger.info(f"Request to /v1/chat/completions (model={request_data.model}, stream={request_data.stream})") 196 | 197 | auth_manager: KiroAuthManager = request.app.state.auth_manager 198 | model_cache: ModelInfoCache = request.app.state.model_cache 199 | 200 | # Подготовка отладочных логов 201 | if debug_logger: 202 | debug_logger.prepare_new_request() 203 | 204 | # Логируем входящий запрос 205 | try: 206 | request_body = json.dumps(request_data.model_dump(), ensure_ascii=False, indent=2).encode('utf-8') 207 | if debug_logger: 208 | debug_logger.log_request_body(request_body) 209 | except Exception as e: 210 | logger.warning(f"Failed to log request body: {e}") 211 | 212 | # Ленивое заполнение кэша моделей 213 | if model_cache.is_empty(): 214 | logger.debug("Model cache is empty, skipping forced population") 215 | 216 | # Генерируем ID для разговора 217 | conversation_id = generate_conversation_id() 218 | 219 | # Строим payload для Kiro 220 | try: 221 | kiro_payload = build_kiro_payload( 222 | request_data, 223 | conversation_id, 224 | auth_manager.profile_arn or "" 225 | ) 226 | except ValueError as e: 227 | raise HTTPException(status_code=400, detail=str(e)) 228 | 229 | # Логируем payload для Kiro 230 | try: 231 | kiro_request_body = json.dumps(kiro_payload, ensure_ascii=False, indent=2).encode('utf-8') 232 | if debug_logger: 233 | debug_logger.log_kiro_request_body(kiro_request_body) 234 | except Exception as e: 235 | logger.warning(f"Failed to log Kiro request: {e}") 236 | 237 | # Создаём HTTP клиент с retry логикой 238 | http_client = KiroHttpClient(auth_manager) 239 | url = f"{auth_manager.api_host}/generateAssistantResponse" 240 | try: 241 | # Делаем запрос к Kiro API (для обоих режимов - streaming и non-streaming) 242 | # Это важно: мы ждём ответа от Kiro ПЕРЕД возвратом StreamingResponse, 243 | # чтобы 200 OK означал что Kiro принял запрос и начал отвечать 244 | response = await http_client.request_with_retry( 245 | "POST", 246 | url, 247 | kiro_payload, 248 | stream=True 249 | ) 250 | 251 | if response.status_code != 200: 252 | try: 253 | error_content = await response.aread() 254 | except Exception: 255 | error_content = b"Unknown error" 256 | 257 | await http_client.close() 258 | error_text = error_content.decode('utf-8', errors='replace') 259 | logger.error(f"Error from Kiro API: {response.status_code} - {error_text}") 260 | 261 | # Пытаемся распарсить JSON ответ от Kiro для извлечения сообщения об ошибке 262 | error_message = error_text 263 | try: 264 | error_json = json.loads(error_text) 265 | if "message" in error_json: 266 | error_message = error_json["message"] 267 | if "reason" in error_json: 268 | error_message = f"{error_message} (reason: {error_json['reason']})" 269 | except (json.JSONDecodeError, KeyError): 270 | pass 271 | 272 | # Возвращаем ошибку в формате OpenAI API 273 | return JSONResponse( 274 | status_code=response.status_code, 275 | content={ 276 | "error": { 277 | "message": error_message, 278 | "type": "kiro_api_error", 279 | "code": response.status_code 280 | } 281 | } 282 | ) 283 | 284 | # Подготавливаем данные для fallback подсчёта токенов 285 | # Конвертируем Pydantic модели в словари для токенизатора 286 | messages_for_tokenizer = [msg.model_dump() for msg in request_data.messages] 287 | tools_for_tokenizer = [tool.model_dump() for tool in request_data.tools] if request_data.tools else None 288 | 289 | if request_data.stream: 290 | # Streaming режим 291 | async def stream_wrapper(): 292 | try: 293 | async for chunk in stream_kiro_to_openai( 294 | http_client.client, 295 | response, 296 | request_data.model, 297 | model_cache, 298 | auth_manager, 299 | request_messages=messages_for_tokenizer, 300 | request_tools=tools_for_tokenizer 301 | ): 302 | yield chunk 303 | finally: 304 | await http_client.close() 305 | 306 | return StreamingResponse(stream_wrapper(), media_type="text/event-stream") 307 | 308 | else: 309 | 310 | # Non-streaming режим - собираем весь ответ 311 | openai_response = await collect_stream_response( 312 | http_client.client, 313 | response, 314 | request_data.model, 315 | model_cache, 316 | auth_manager, 317 | request_messages=messages_for_tokenizer, 318 | request_tools=tools_for_tokenizer 319 | ) 320 | 321 | await http_client.close() 322 | return JSONResponse(content=openai_response) 323 | 324 | except HTTPException: 325 | await http_client.close() 326 | raise 327 | except Exception as e: 328 | await http_client.close() 329 | logger.error(f"Internal error: {e}", exc_info=True) 330 | raise HTTPException(status_code=500, detail=f"Internal Server Error: {str(e)}") -------------------------------------------------------------------------------- /kiro_gateway/auth.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Kiro OpenAI Gateway 4 | # Copyright (C) 2025 Jwadow 5 | # 6 | # This program is free software: you can redistribute it and/or modify 7 | # it under the terms of the GNU Affero General Public License as published by 8 | # the Free Software Foundation, either version 3 of the License, or 9 | # (at your option) any later version. 10 | # 11 | # This program is distributed in the hope that it will be useful, 12 | # but WITHOUT ANY WARRANTY; without even the implied warranty of 13 | # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 14 | # GNU Affero General Public License for more details. 15 | # 16 | # You should have received a copy of the GNU Affero General Public License 17 | # along with this program. If not, see . 18 | 19 | """ 20 | Менеджер аутентификации для Kiro API. 21 | 22 | Управляет жизненным циклом токенов доступа: 23 | - Загрузка credentials из .env или JSON файла 24 | - Автоматическое обновление токена при истечении 25 | - Потокобезопасное обновление с использованием asyncio.Lock 26 | """ 27 | 28 | import asyncio 29 | import json 30 | from datetime import datetime, timezone 31 | from pathlib import Path 32 | from typing import Optional 33 | 34 | import httpx 35 | from loguru import logger 36 | 37 | from kiro_gateway.config import ( 38 | TOKEN_REFRESH_THRESHOLD, 39 | get_kiro_refresh_url, 40 | get_kiro_api_host, 41 | get_kiro_q_host, 42 | ) 43 | from kiro_gateway.utils import get_machine_fingerprint 44 | 45 | 46 | class KiroAuthManager: 47 | """ 48 | Управляет жизненным циклом токена для доступа к Kiro API. 49 | 50 | Поддерживает: 51 | - Загрузку credentials из .env или JSON файла 52 | - Автоматическое обновление токена при истечении 53 | - Проверку времени истечения (expiresAt) 54 | - Сохранение обновлённых токенов в файл 55 | 56 | Attributes: 57 | profile_arn: ARN профиля AWS CodeWhisperer 58 | region: Регион AWS 59 | api_host: Хост API для текущего региона 60 | q_host: Хост Q API для текущего региона 61 | fingerprint: Уникальный fingerprint машины 62 | 63 | Example: 64 | >>> auth_manager = KiroAuthManager( 65 | ... refresh_token="your_refresh_token", 66 | ... region="us-east-1" 67 | ... ) 68 | >>> token = await auth_manager.get_access_token() 69 | """ 70 | 71 | def __init__( 72 | self, 73 | refresh_token: Optional[str] = None, 74 | profile_arn: Optional[str] = None, 75 | region: str = "us-east-1", 76 | creds_file: Optional[str] = None 77 | ): 78 | """ 79 | Инициализирует менеджер аутентификации. 80 | 81 | Args: 82 | refresh_token: Refresh token для обновления access token 83 | profile_arn: ARN профиля AWS CodeWhisperer 84 | region: Регион AWS (по умолчанию us-east-1) 85 | creds_file: Путь к JSON файлу с credentials (опционально) 86 | """ 87 | self._refresh_token = refresh_token 88 | self._profile_arn = profile_arn 89 | self._region = region 90 | self._creds_file = creds_file 91 | 92 | self._access_token: Optional[str] = None 93 | self._expires_at: Optional[datetime] = None 94 | self._lock = asyncio.Lock() 95 | 96 | # Динамические URL на основе региона 97 | self._refresh_url = get_kiro_refresh_url(region) 98 | self._api_host = get_kiro_api_host(region) 99 | self._q_host = get_kiro_q_host(region) 100 | 101 | # Fingerprint для User-Agent 102 | self._fingerprint = get_machine_fingerprint() 103 | 104 | # Загружаем credentials из файла если указан 105 | if creds_file: 106 | self._load_credentials_from_file(creds_file) 107 | 108 | def _load_credentials_from_file(self, file_path: str) -> None: 109 | """ 110 | Загружает credentials из JSON файла. 111 | 112 | Поддерживаемые поля в JSON: 113 | - refreshToken: Refresh token 114 | - accessToken: Access token (если уже есть) 115 | - profileArn: ARN профиля 116 | - region: Регион AWS 117 | - expiresAt: Время истечения токена (ISO 8601) 118 | 119 | Args: 120 | file_path: Путь к JSON файлу 121 | """ 122 | try: 123 | path = Path(file_path).expanduser() 124 | if not path.exists(): 125 | logger.warning(f"Credentials file not found: {file_path}") 126 | return 127 | 128 | with open(path, 'r', encoding='utf-8') as f: 129 | data = json.load(f) 130 | 131 | # Загружаем данные из файла 132 | if 'refreshToken' in data: 133 | self._refresh_token = data['refreshToken'] 134 | if 'accessToken' in data: 135 | self._access_token = data['accessToken'] 136 | if 'profileArn' in data: 137 | self._profile_arn = data['profileArn'] 138 | if 'region' in data: 139 | self._region = data['region'] 140 | # Обновляем URL для нового региона 141 | self._refresh_url = get_kiro_refresh_url(self._region) 142 | self._api_host = get_kiro_api_host(self._region) 143 | self._q_host = get_kiro_q_host(self._region) 144 | 145 | # Парсим expiresAt 146 | if 'expiresAt' in data: 147 | try: 148 | expires_str = data['expiresAt'] 149 | # Поддержка разных форматов даты 150 | if expires_str.endswith('Z'): 151 | self._expires_at = datetime.fromisoformat(expires_str.replace('Z', '+00:00')) 152 | else: 153 | self._expires_at = datetime.fromisoformat(expires_str) 154 | except Exception as e: 155 | logger.warning(f"Failed to parse expiresAt: {e}") 156 | 157 | logger.info(f"Credentials loaded from {file_path}") 158 | 159 | except Exception as e: 160 | logger.error(f"Error loading credentials from file: {e}") 161 | 162 | def _save_credentials_to_file(self) -> None: 163 | """ 164 | Сохраняет обновлённые credentials в JSON файл. 165 | 166 | Обновляет существующий файл, сохраняя другие поля. 167 | """ 168 | if not self._creds_file: 169 | return 170 | 171 | try: 172 | path = Path(self._creds_file).expanduser() 173 | 174 | # Читаем существующие данные 175 | existing_data = {} 176 | if path.exists(): 177 | with open(path, 'r', encoding='utf-8') as f: 178 | existing_data = json.load(f) 179 | 180 | # Обновляем данные 181 | existing_data['accessToken'] = self._access_token 182 | existing_data['refreshToken'] = self._refresh_token 183 | if self._expires_at: 184 | existing_data['expiresAt'] = self._expires_at.isoformat() 185 | if self._profile_arn: 186 | existing_data['profileArn'] = self._profile_arn 187 | 188 | # Сохраняем 189 | with open(path, 'w', encoding='utf-8') as f: 190 | json.dump(existing_data, f, indent=2, ensure_ascii=False) 191 | 192 | logger.debug(f"Credentials saved to {self._creds_file}") 193 | 194 | except Exception as e: 195 | logger.error(f"Error saving credentials: {e}") 196 | 197 | def is_token_expiring_soon(self) -> bool: 198 | """ 199 | Проверяет, истекает ли токен в ближайшее время. 200 | 201 | Returns: 202 | True если токен истекает в течение TOKEN_REFRESH_THRESHOLD секунд 203 | или если информация о времени истечения отсутствует 204 | """ 205 | if not self._expires_at: 206 | return True # Если нет информации о времени истечения, считаем что нужно обновить 207 | 208 | now = datetime.now(timezone.utc) 209 | threshold = now.timestamp() + TOKEN_REFRESH_THRESHOLD 210 | 211 | return self._expires_at.timestamp() <= threshold 212 | 213 | async def _refresh_token_request(self) -> None: 214 | """ 215 | Выполняет запрос на обновление токена. 216 | 217 | Отправляет POST запрос к Kiro API для получения нового access token. 218 | Обновляет внутреннее состояние и сохраняет credentials в файл. 219 | 220 | Raises: 221 | ValueError: Если refresh token не установлен или ответ не содержит accessToken 222 | httpx.HTTPError: При ошибке HTTP запроса 223 | """ 224 | if not self._refresh_token: 225 | raise ValueError("Refresh token is not set") 226 | 227 | logger.info("Refreshing Kiro token...") 228 | 229 | payload = {'refreshToken': self._refresh_token} 230 | headers = { 231 | "Content-Type": "application/json", 232 | "User-Agent": f"KiroGateway-{self._fingerprint[:16]}", 233 | } 234 | 235 | async with httpx.AsyncClient(timeout=30) as client: 236 | response = await client.post(self._refresh_url, json=payload, headers=headers) 237 | response.raise_for_status() 238 | data = response.json() 239 | 240 | new_access_token = data.get("accessToken") 241 | new_refresh_token = data.get("refreshToken") 242 | expires_in = data.get("expiresIn", 3600) 243 | new_profile_arn = data.get("profileArn") 244 | 245 | if not new_access_token: 246 | raise ValueError(f"Response does not contain accessToken: {data}") 247 | 248 | # Обновляем данные 249 | self._access_token = new_access_token 250 | if new_refresh_token: 251 | self._refresh_token = new_refresh_token 252 | if new_profile_arn: 253 | self._profile_arn = new_profile_arn 254 | 255 | # Вычисляем время истечения с запасом (минус 60 секунд) 256 | self._expires_at = datetime.now(timezone.utc).replace(microsecond=0) 257 | self._expires_at = datetime.fromtimestamp( 258 | self._expires_at.timestamp() + expires_in - 60, 259 | tz=timezone.utc 260 | ) 261 | 262 | logger.info(f"Token refreshed, expires: {self._expires_at.isoformat()}") 263 | 264 | # Сохраняем в файл 265 | self._save_credentials_to_file() 266 | 267 | async def get_access_token(self) -> str: 268 | """ 269 | Возвращает действительный access_token, обновляя его при необходимости. 270 | 271 | Потокобезопасный метод с использованием asyncio.Lock. 272 | Автоматически обновляет токен если он истёк или скоро истечёт. 273 | 274 | Returns: 275 | Действительный access token 276 | 277 | Raises: 278 | ValueError: Если не удалось получить access token 279 | """ 280 | async with self._lock: 281 | if not self._access_token or self.is_token_expiring_soon(): 282 | await self._refresh_token_request() 283 | 284 | if not self._access_token: 285 | raise ValueError("Failed to obtain access token") 286 | 287 | return self._access_token 288 | 289 | async def force_refresh(self) -> str: 290 | """ 291 | Принудительно обновляет токен. 292 | 293 | Используется при получении 403 ошибки от API. 294 | 295 | Returns: 296 | Новый access token 297 | """ 298 | async with self._lock: 299 | await self._refresh_token_request() 300 | return self._access_token 301 | 302 | @property 303 | def profile_arn(self) -> Optional[str]: 304 | """ARN профиля AWS CodeWhisperer.""" 305 | return self._profile_arn 306 | 307 | @property 308 | def region(self) -> str: 309 | """Регион AWS.""" 310 | return self._region 311 | 312 | @property 313 | def api_host(self) -> str: 314 | """Хост API для текущего региона.""" 315 | return self._api_host 316 | 317 | @property 318 | def q_host(self) -> str: 319 | """Хост Q API для текущего региона.""" 320 | return self._q_host 321 | 322 | @property 323 | def fingerprint(self) -> str: 324 | """Уникальный fingerprint машины.""" 325 | return self._fingerprint -------------------------------------------------------------------------------- /tests/unit/test_routes.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | """ 4 | Unit-тесты для API endpoints (routes.py). 5 | Проверяет работу эндпоинтов /, /health, /v1/models, /v1/chat/completions. 6 | """ 7 | 8 | import pytest 9 | from unittest.mock import AsyncMock, Mock, patch, MagicMock 10 | from datetime import datetime, timezone 11 | 12 | from fastapi import HTTPException 13 | from fastapi.testclient import TestClient 14 | 15 | from kiro_gateway.routes import verify_api_key, router 16 | from kiro_gateway.config import PROXY_API_KEY, APP_VERSION, AVAILABLE_MODELS 17 | 18 | 19 | class TestVerifyApiKey: 20 | """Тесты функции verify_api_key.""" 21 | 22 | @pytest.mark.asyncio 23 | async def test_valid_api_key_returns_true(self): 24 | """ 25 | Что он делает: Проверяет успешную валидацию корректного ключа. 26 | Цель: Убедиться, что валидный ключ проходит проверку. 27 | """ 28 | print("Настройка: Валидный API ключ...") 29 | valid_header = f"Bearer {PROXY_API_KEY}" 30 | 31 | print("Действие: Проверка ключа...") 32 | result = await verify_api_key(valid_header) 33 | 34 | print(f"Сравниваем результат: Ожидалось True, Получено {result}") 35 | assert result is True 36 | 37 | @pytest.mark.asyncio 38 | async def test_invalid_api_key_raises_401(self): 39 | """ 40 | Что он делает: Проверяет отклонение невалидного ключа. 41 | Цель: Убедиться, что невалидный ключ вызывает 401. 42 | """ 43 | print("Настройка: Невалидный API ключ...") 44 | invalid_header = "Bearer wrong_key" 45 | 46 | print("Действие: Проверка ключа...") 47 | with pytest.raises(HTTPException) as exc_info: 48 | await verify_api_key(invalid_header) 49 | 50 | print(f"Проверка: HTTPException с кодом 401...") 51 | assert exc_info.value.status_code == 401 52 | assert "Invalid or missing API Key" in exc_info.value.detail 53 | 54 | @pytest.mark.asyncio 55 | async def test_missing_api_key_raises_401(self): 56 | """ 57 | Что он делает: Проверяет отклонение отсутствующего ключа. 58 | Цель: Убедиться, что отсутствие ключа вызывает 401. 59 | """ 60 | print("Настройка: Отсутствующий API ключ...") 61 | 62 | print("Действие: Проверка ключа...") 63 | with pytest.raises(HTTPException) as exc_info: 64 | await verify_api_key(None) 65 | 66 | print(f"Проверка: HTTPException с кодом 401...") 67 | assert exc_info.value.status_code == 401 68 | 69 | @pytest.mark.asyncio 70 | async def test_empty_api_key_raises_401(self): 71 | """ 72 | Что он делает: Проверяет отклонение пустого ключа. 73 | Цель: Убедиться, что пустая строка вызывает 401. 74 | """ 75 | print("Настройка: Пустой API ключ...") 76 | 77 | print("Действие: Проверка ключа...") 78 | with pytest.raises(HTTPException) as exc_info: 79 | await verify_api_key("") 80 | 81 | print(f"Проверка: HTTPException с кодом 401...") 82 | assert exc_info.value.status_code == 401 83 | 84 | @pytest.mark.asyncio 85 | async def test_wrong_format_raises_401(self): 86 | """ 87 | Что он делает: Проверяет отклонение ключа без Bearer. 88 | Цель: Убедиться, что неправильный формат вызывает 401. 89 | """ 90 | print("Настройка: Ключ без Bearer...") 91 | wrong_format = PROXY_API_KEY # Без "Bearer " 92 | 93 | print("Действие: Проверка ключа...") 94 | with pytest.raises(HTTPException) as exc_info: 95 | await verify_api_key(wrong_format) 96 | 97 | print(f"Проверка: HTTPException с кодом 401...") 98 | assert exc_info.value.status_code == 401 99 | 100 | 101 | class TestRootEndpoint: 102 | """Тесты эндпоинта /.""" 103 | 104 | def test_root_returns_status_ok(self, test_client): 105 | """ 106 | Что он делает: Проверяет ответ корневого эндпоинта. 107 | Цель: Убедиться, что / возвращает статус ok. 108 | """ 109 | print("Действие: GET /...") 110 | response = test_client.get("/") 111 | 112 | print(f"Результат: {response.json()}") 113 | assert response.status_code == 200 114 | assert response.json()["status"] == "ok" 115 | assert "Kiro API Gateway" in response.json()["message"] 116 | 117 | def test_root_returns_version(self, test_client): 118 | """ 119 | Что он делает: Проверяет наличие версии в ответе. 120 | Цель: Убедиться, что версия приложения возвращается. 121 | """ 122 | print("Действие: GET /...") 123 | response = test_client.get("/") 124 | 125 | print(f"Результат: {response.json()}") 126 | assert response.status_code == 200 127 | assert "version" in response.json() 128 | assert response.json()["version"] == APP_VERSION 129 | 130 | 131 | class TestHealthEndpoint: 132 | """Тесты эндпоинта /health.""" 133 | 134 | def test_health_returns_healthy(self, test_client): 135 | """ 136 | Что он делает: Проверяет ответ health эндпоинта. 137 | Цель: Убедиться, что /health возвращает статус healthy. 138 | """ 139 | print("Действие: GET /health...") 140 | response = test_client.get("/health") 141 | 142 | print(f"Результат: {response.json()}") 143 | assert response.status_code == 200 144 | assert response.json()["status"] == "healthy" 145 | 146 | def test_health_returns_timestamp(self, test_client): 147 | """ 148 | Что он делает: Проверяет наличие timestamp в ответе. 149 | Цель: Убедиться, что timestamp возвращается. 150 | """ 151 | print("Действие: GET /health...") 152 | response = test_client.get("/health") 153 | 154 | print(f"Результат: {response.json()}") 155 | assert response.status_code == 200 156 | assert "timestamp" in response.json() 157 | 158 | def test_health_returns_version(self, test_client): 159 | """ 160 | Что он делает: Проверяет наличие версии в ответе. 161 | Цель: Убедиться, что версия приложения возвращается. 162 | """ 163 | print("Действие: GET /health...") 164 | response = test_client.get("/health") 165 | 166 | print(f"Результат: {response.json()}") 167 | assert response.status_code == 200 168 | assert response.json()["version"] == APP_VERSION 169 | 170 | 171 | class TestModelsEndpoint: 172 | """Тесты эндпоинта /v1/models.""" 173 | 174 | def test_models_requires_auth(self, test_client): 175 | """ 176 | Что он делает: Проверяет требование авторизации. 177 | Цель: Убедиться, что без ключа возвращается 401. 178 | """ 179 | print("Действие: GET /v1/models без авторизации...") 180 | response = test_client.get("/v1/models") 181 | 182 | print(f"Статус: {response.status_code}") 183 | assert response.status_code == 401 184 | 185 | def test_models_returns_list(self, test_client, valid_proxy_api_key): 186 | """ 187 | Что он делает: Проверяет возврат списка моделей. 188 | Цель: Убедиться, что /v1/models возвращает список. 189 | """ 190 | print("Действие: GET /v1/models с авторизацией...") 191 | response = test_client.get( 192 | "/v1/models", 193 | headers={"Authorization": f"Bearer {valid_proxy_api_key}"} 194 | ) 195 | 196 | print(f"Результат: {response.json()}") 197 | assert response.status_code == 200 198 | assert response.json()["object"] == "list" 199 | assert "data" in response.json() 200 | 201 | def test_models_returns_available_models(self, test_client, valid_proxy_api_key): 202 | """ 203 | Что он делает: Проверяет наличие доступных моделей. 204 | Цель: Убедиться, что все модели из AVAILABLE_MODELS возвращаются. 205 | """ 206 | print("Действие: GET /v1/models с авторизацией...") 207 | response = test_client.get( 208 | "/v1/models", 209 | headers={"Authorization": f"Bearer {valid_proxy_api_key}"} 210 | ) 211 | 212 | print(f"Результат: {response.json()}") 213 | assert response.status_code == 200 214 | 215 | model_ids = [m["id"] for m in response.json()["data"]] 216 | for expected_model in AVAILABLE_MODELS: 217 | assert expected_model in model_ids, f"Модель {expected_model} не найдена" 218 | 219 | def test_models_format_is_openai_compatible(self, test_client, valid_proxy_api_key): 220 | """ 221 | Что он делает: Проверяет формат ответа на совместимость с OpenAI. 222 | Цель: Убедиться, что формат соответствует OpenAI API. 223 | """ 224 | print("Действие: GET /v1/models с авторизацией...") 225 | response = test_client.get( 226 | "/v1/models", 227 | headers={"Authorization": f"Bearer {valid_proxy_api_key}"} 228 | ) 229 | 230 | print(f"Результат: {response.json()}") 231 | assert response.status_code == 200 232 | 233 | for model in response.json()["data"]: 234 | assert "id" in model 235 | assert "object" in model 236 | assert model["object"] == "model" 237 | assert "owned_by" in model 238 | 239 | 240 | class TestChatCompletionsEndpoint: 241 | """Тесты эндпоинта /v1/chat/completions.""" 242 | 243 | def test_chat_completions_requires_auth(self, test_client): 244 | """ 245 | Что он делает: Проверяет требование авторизации. 246 | Цель: Убедиться, что без ключа возвращается 401. 247 | """ 248 | print("Действие: POST /v1/chat/completions без авторизации...") 249 | response = test_client.post( 250 | "/v1/chat/completions", 251 | json={ 252 | "model": "claude-sonnet-4-5", 253 | "messages": [{"role": "user", "content": "Hello"}] 254 | } 255 | ) 256 | 257 | print(f"Статус: {response.status_code}") 258 | assert response.status_code == 401 259 | 260 | def test_chat_completions_validates_messages(self, test_client, valid_proxy_api_key): 261 | """ 262 | Что он делает: Проверяет валидацию пустых сообщений. 263 | Цель: Убедиться, что пустой список сообщений отклоняется. 264 | """ 265 | print("Действие: POST /v1/chat/completions с пустыми сообщениями...") 266 | response = test_client.post( 267 | "/v1/chat/completions", 268 | headers={"Authorization": f"Bearer {valid_proxy_api_key}"}, 269 | json={ 270 | "model": "claude-sonnet-4-5", 271 | "messages": [] 272 | } 273 | ) 274 | 275 | print(f"Статус: {response.status_code}") 276 | # Pydantic должен отклонить пустой список 277 | assert response.status_code == 422 278 | 279 | def test_chat_completions_validates_model(self, test_client, valid_proxy_api_key): 280 | """ 281 | Что он делает: Проверяет валидацию отсутствующей модели. 282 | Цель: Убедиться, что запрос без модели отклоняется. 283 | """ 284 | print("Действие: POST /v1/chat/completions без модели...") 285 | response = test_client.post( 286 | "/v1/chat/completions", 287 | headers={"Authorization": f"Bearer {valid_proxy_api_key}"}, 288 | json={ 289 | "messages": [{"role": "user", "content": "Hello"}] 290 | } 291 | ) 292 | 293 | print(f"Статус: {response.status_code}") 294 | assert response.status_code == 422 295 | 296 | 297 | class TestChatCompletionsWithMockedKiro: 298 | """Тесты /v1/chat/completions с мокированным Kiro API.""" 299 | 300 | def test_chat_completions_accepts_valid_request_format(self, test_client, valid_proxy_api_key): 301 | """ 302 | Что он делает: Проверяет, что валидный формат запроса принимается. 303 | Цель: Убедиться, что Pydantic валидация проходит для корректного запроса. 304 | """ 305 | print("Настройка: Валидный запрос...") 306 | 307 | # Этот тест проверяет только валидацию запроса 308 | # Реальный вызов к Kiro API будет заблокирован фикстурой block_all_network_calls 309 | # Поэтому мы ожидаем ошибку на этапе HTTP запроса, а не валидации 310 | 311 | print("Действие: POST /v1/chat/completions с валидным запросом...") 312 | response = test_client.post( 313 | "/v1/chat/completions", 314 | headers={"Authorization": f"Bearer {valid_proxy_api_key}"}, 315 | json={ 316 | "model": "claude-sonnet-4-5", 317 | "messages": [{"role": "user", "content": "Hello"}], 318 | "stream": False 319 | } 320 | ) 321 | 322 | print(f"Статус: {response.status_code}") 323 | # Запрос должен пройти валидацию (не 422) 324 | # Но может упасть на этапе HTTP из-за блокировки сети 325 | assert response.status_code != 422 326 | 327 | 328 | class TestChatCompletionsErrorHandling: 329 | """Тесты обработки ошибок в /v1/chat/completions.""" 330 | 331 | def test_invalid_json_returns_422(self, test_client, valid_proxy_api_key): 332 | """ 333 | Что он делает: Проверяет обработку невалидного JSON. 334 | Цель: Убедиться, что невалидный JSON возвращает 422. 335 | """ 336 | print("Действие: POST /v1/chat/completions с невалидным JSON...") 337 | response = test_client.post( 338 | "/v1/chat/completions", 339 | headers={ 340 | "Authorization": f"Bearer {valid_proxy_api_key}", 341 | "Content-Type": "application/json" 342 | }, 343 | content=b"not valid json" 344 | ) 345 | 346 | print(f"Статус: {response.status_code}") 347 | assert response.status_code == 422 348 | 349 | def test_missing_content_in_message_returns_200(self, test_client, valid_proxy_api_key): 350 | """ 351 | Что он делает: Проверяет обработку сообщения без content. 352 | Цель: Убедиться, что сообщение без content допустимо (content опционален). 353 | """ 354 | print("Действие: POST /v1/chat/completions с сообщением без content...") 355 | # Этот тест проверяет валидацию Pydantic 356 | # content может быть None согласно модели 357 | response = test_client.post( 358 | "/v1/chat/completions", 359 | headers={"Authorization": f"Bearer {valid_proxy_api_key}"}, 360 | json={ 361 | "model": "claude-sonnet-4-5", 362 | "messages": [{"role": "user"}] # content отсутствует 363 | } 364 | ) 365 | 366 | print(f"Статус: {response.status_code}") 367 | # Запрос должен пройти валидацию (content опционален) 368 | # Но может упасть на этапе обработки из-за отсутствия мока Kiro API 369 | # Поэтому проверяем, что это не 422 (валидация прошла) 370 | assert response.status_code != 422 or "content" not in str(response.json()) 371 | 372 | 373 | class TestRouterIntegration: 374 | """Тесты интеграции роутера.""" 375 | 376 | def test_router_has_all_endpoints(self): 377 | """ 378 | Что он делает: Проверяет наличие всех эндпоинтов в роутере. 379 | Цель: Убедиться, что все эндпоинты зарегистрированы. 380 | """ 381 | print("Проверка: Эндпоинты в роутере...") 382 | 383 | routes = [route.path for route in router.routes] 384 | 385 | print(f"Найденные роуты: {routes}") 386 | assert "/" in routes 387 | assert "/health" in routes 388 | assert "/v1/models" in routes 389 | assert "/v1/chat/completions" in routes 390 | 391 | def test_router_methods(self): 392 | """ 393 | Что он делает: Проверяет HTTP методы эндпоинтов. 394 | Цель: Убедиться, что методы соответствуют ожиданиям. 395 | """ 396 | print("Проверка: HTTP методы...") 397 | 398 | for route in router.routes: 399 | if route.path == "/": 400 | assert "GET" in route.methods 401 | elif route.path == "/health": 402 | assert "GET" in route.methods 403 | elif route.path == "/v1/models": 404 | assert "GET" in route.methods 405 | elif route.path == "/v1/chat/completions": 406 | assert "POST" in route.methods -------------------------------------------------------------------------------- /tests/integration/test_full_flow.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | """ 4 | Integration-тесты для полного end-to-end flow. 5 | Проверяет взаимодействие всех компонентов системы. 6 | """ 7 | 8 | import pytest 9 | import json 10 | from unittest.mock import AsyncMock, Mock, patch, MagicMock 11 | from datetime import datetime, timezone, timedelta 12 | 13 | from fastapi.testclient import TestClient 14 | import httpx 15 | 16 | from kiro_gateway.config import PROXY_API_KEY, AVAILABLE_MODELS 17 | 18 | 19 | class TestFullChatCompletionFlow: 20 | """Integration-тесты полного flow chat completions.""" 21 | 22 | def test_full_flow_health_to_models_to_chat(self, test_client, valid_proxy_api_key): 23 | """ 24 | Что он делает: Проверяет полный flow от health check до chat completions. 25 | Цель: Убедиться, что все эндпоинты работают вместе. 26 | """ 27 | print("Шаг 1: Health check...") 28 | health_response = test_client.get("/health") 29 | assert health_response.status_code == 200 30 | assert health_response.json()["status"] == "healthy" 31 | print(f"Health: {health_response.json()}") 32 | 33 | print("Шаг 2: Получение списка моделей...") 34 | models_response = test_client.get( 35 | "/v1/models", 36 | headers={"Authorization": f"Bearer {valid_proxy_api_key}"} 37 | ) 38 | assert models_response.status_code == 200 39 | assert len(models_response.json()["data"]) > 0 40 | print(f"Модели: {[m['id'] for m in models_response.json()['data']]}") 41 | 42 | print("Шаг 3: Валидация запроса chat completions...") 43 | # Этот запрос пройдёт валидацию, но упадёт на HTTP из-за блокировки сети 44 | chat_response = test_client.post( 45 | "/v1/chat/completions", 46 | headers={"Authorization": f"Bearer {valid_proxy_api_key}"}, 47 | json={ 48 | "model": "claude-sonnet-4-5", 49 | "messages": [{"role": "user", "content": "Hello"}] 50 | } 51 | ) 52 | # Запрос должен пройти валидацию (не 422) 53 | assert chat_response.status_code != 422 54 | print(f"Chat response status: {chat_response.status_code}") 55 | 56 | def test_authentication_flow(self, test_client, valid_proxy_api_key, invalid_proxy_api_key): 57 | """ 58 | Что он делает: Проверяет flow аутентификации. 59 | Цель: Убедиться, что защищённые эндпоинты требуют авторизации. 60 | """ 61 | print("Шаг 1: Запрос без авторизации...") 62 | no_auth_response = test_client.get("/v1/models") 63 | assert no_auth_response.status_code == 401 64 | print(f"Без авторизации: {no_auth_response.status_code}") 65 | 66 | print("Шаг 2: Запрос с неверным ключом...") 67 | wrong_auth_response = test_client.get( 68 | "/v1/models", 69 | headers={"Authorization": f"Bearer {invalid_proxy_api_key}"} 70 | ) 71 | assert wrong_auth_response.status_code == 401 72 | print(f"Неверный ключ: {wrong_auth_response.status_code}") 73 | 74 | print("Шаг 3: Запрос с верным ключом...") 75 | valid_auth_response = test_client.get( 76 | "/v1/models", 77 | headers={"Authorization": f"Bearer {valid_proxy_api_key}"} 78 | ) 79 | assert valid_auth_response.status_code == 200 80 | print(f"Верный ключ: {valid_auth_response.status_code}") 81 | 82 | def test_openai_compatibility_format(self, test_client, valid_proxy_api_key): 83 | """ 84 | Что он делает: Проверяет совместимость формата ответов с OpenAI API. 85 | Цель: Убедиться, что ответы соответствуют спецификации OpenAI. 86 | """ 87 | print("Проверка формата /v1/models...") 88 | models_response = test_client.get( 89 | "/v1/models", 90 | headers={"Authorization": f"Bearer {valid_proxy_api_key}"} 91 | ) 92 | 93 | assert models_response.status_code == 200 94 | data = models_response.json() 95 | 96 | # Проверяем структуру ответа OpenAI 97 | assert "object" in data 98 | assert data["object"] == "list" 99 | assert "data" in data 100 | assert isinstance(data["data"], list) 101 | 102 | # Проверяем структуру каждой модели 103 | for model in data["data"]: 104 | assert "id" in model 105 | assert "object" in model 106 | assert model["object"] == "model" 107 | assert "owned_by" in model 108 | assert "created" in model 109 | 110 | print(f"Формат соответствует OpenAI API: {len(data['data'])} моделей") 111 | 112 | 113 | class TestRequestValidationFlow: 114 | """Integration-тесты валидации запросов.""" 115 | 116 | def test_chat_completions_request_validation(self, test_client, valid_proxy_api_key): 117 | """ 118 | Что он делает: Проверяет валидацию различных форматов запросов. 119 | Цель: Убедиться, что валидация работает корректно. 120 | """ 121 | print("Тест 1: Пустые сообщения...") 122 | empty_messages = test_client.post( 123 | "/v1/chat/completions", 124 | headers={"Authorization": f"Bearer {valid_proxy_api_key}"}, 125 | json={"model": "claude-sonnet-4-5", "messages": []} 126 | ) 127 | assert empty_messages.status_code == 422 128 | print(f"Пустые сообщения: {empty_messages.status_code}") 129 | 130 | print("Тест 2: Отсутствует model...") 131 | no_model = test_client.post( 132 | "/v1/chat/completions", 133 | headers={"Authorization": f"Bearer {valid_proxy_api_key}"}, 134 | json={"messages": [{"role": "user", "content": "Hello"}]} 135 | ) 136 | assert no_model.status_code == 422 137 | print(f"Без model: {no_model.status_code}") 138 | 139 | print("Тест 3: Отсутствует messages...") 140 | no_messages = test_client.post( 141 | "/v1/chat/completions", 142 | headers={"Authorization": f"Bearer {valid_proxy_api_key}"}, 143 | json={"model": "claude-sonnet-4-5"} 144 | ) 145 | assert no_messages.status_code == 422 146 | print(f"Без messages: {no_messages.status_code}") 147 | 148 | print("Тест 4: Валидный запрос...") 149 | valid_request = test_client.post( 150 | "/v1/chat/completions", 151 | headers={"Authorization": f"Bearer {valid_proxy_api_key}"}, 152 | json={ 153 | "model": "claude-sonnet-4-5", 154 | "messages": [{"role": "user", "content": "Hello"}] 155 | } 156 | ) 157 | # Валидация должна пройти (не 422) 158 | assert valid_request.status_code != 422 159 | print(f"Валидный запрос: {valid_request.status_code}") 160 | 161 | def test_complex_message_formats(self, test_client, valid_proxy_api_key): 162 | """ 163 | Что он делает: Проверяет обработку сложных форматов сообщений. 164 | Цель: Убедиться, что multimodal и tool форматы принимаются. 165 | """ 166 | print("Тест 1: System + User сообщения...") 167 | system_user = test_client.post( 168 | "/v1/chat/completions", 169 | headers={"Authorization": f"Bearer {valid_proxy_api_key}"}, 170 | json={ 171 | "model": "claude-sonnet-4-5", 172 | "messages": [ 173 | {"role": "system", "content": "You are helpful"}, 174 | {"role": "user", "content": "Hello"} 175 | ] 176 | } 177 | ) 178 | assert system_user.status_code != 422 179 | print(f"System + User: {system_user.status_code}") 180 | 181 | print("Тест 2: Multi-turn conversation...") 182 | multi_turn = test_client.post( 183 | "/v1/chat/completions", 184 | headers={"Authorization": f"Bearer {valid_proxy_api_key}"}, 185 | json={ 186 | "model": "claude-sonnet-4-5", 187 | "messages": [ 188 | {"role": "user", "content": "Hello"}, 189 | {"role": "assistant", "content": "Hi there!"}, 190 | {"role": "user", "content": "How are you?"} 191 | ] 192 | } 193 | ) 194 | assert multi_turn.status_code != 422 195 | print(f"Multi-turn: {multi_turn.status_code}") 196 | 197 | print("Тест 3: С tools...") 198 | with_tools = test_client.post( 199 | "/v1/chat/completions", 200 | headers={"Authorization": f"Bearer {valid_proxy_api_key}"}, 201 | json={ 202 | "model": "claude-sonnet-4-5", 203 | "messages": [{"role": "user", "content": "What's the weather?"}], 204 | "tools": [{ 205 | "type": "function", 206 | "function": { 207 | "name": "get_weather", 208 | "description": "Get weather", 209 | "parameters": {"type": "object", "properties": {}} 210 | } 211 | }] 212 | } 213 | ) 214 | assert with_tools.status_code != 422 215 | print(f"С tools: {with_tools.status_code}") 216 | 217 | 218 | class TestErrorHandlingFlow: 219 | """Integration-тесты обработки ошибок.""" 220 | 221 | def test_invalid_json_handling(self, test_client, valid_proxy_api_key): 222 | """ 223 | Что он делает: Проверяет обработку невалидного JSON. 224 | Цель: Убедиться, что невалидный JSON возвращает понятную ошибку. 225 | """ 226 | print("Отправка невалидного JSON...") 227 | response = test_client.post( 228 | "/v1/chat/completions", 229 | headers={ 230 | "Authorization": f"Bearer {valid_proxy_api_key}", 231 | "Content-Type": "application/json" 232 | }, 233 | content=b"not valid json" 234 | ) 235 | 236 | assert response.status_code == 422 237 | print(f"Невалидный JSON: {response.status_code}") 238 | 239 | def test_wrong_content_type_handling(self, test_client, valid_proxy_api_key): 240 | """ 241 | Что он делает: Проверяет обработку неверного Content-Type. 242 | Цель: Убедиться, что неверный Content-Type обрабатывается. 243 | """ 244 | print("Отправка с неверным Content-Type...") 245 | response = test_client.post( 246 | "/v1/chat/completions", 247 | headers={ 248 | "Authorization": f"Bearer {valid_proxy_api_key}", 249 | "Content-Type": "text/plain" 250 | }, 251 | content=b"Hello" 252 | ) 253 | 254 | # Должна быть ошибка валидации 255 | assert response.status_code == 422 256 | print(f"Неверный Content-Type: {response.status_code}") 257 | 258 | 259 | class TestModelsEndpointIntegration: 260 | """Integration-тесты эндпоинта /v1/models.""" 261 | 262 | def test_models_returns_all_available_models(self, test_client, valid_proxy_api_key): 263 | """ 264 | Что он делает: Проверяет, что все модели из конфига возвращаются. 265 | Цель: Убедиться в полноте списка моделей. 266 | """ 267 | print("Получение списка моделей...") 268 | response = test_client.get( 269 | "/v1/models", 270 | headers={"Authorization": f"Bearer {valid_proxy_api_key}"} 271 | ) 272 | 273 | assert response.status_code == 200 274 | 275 | returned_ids = {m["id"] for m in response.json()["data"]} 276 | expected_ids = set(AVAILABLE_MODELS) 277 | 278 | print(f"Возвращённые модели: {returned_ids}") 279 | print(f"Ожидаемые модели: {expected_ids}") 280 | 281 | assert returned_ids == expected_ids 282 | 283 | def test_models_caching_behavior(self, test_client, valid_proxy_api_key): 284 | """ 285 | Что он делает: Проверяет поведение кэширования моделей. 286 | Цель: Убедиться, что повторные запросы работают корректно. 287 | """ 288 | print("Первый запрос моделей...") 289 | response1 = test_client.get( 290 | "/v1/models", 291 | headers={"Authorization": f"Bearer {valid_proxy_api_key}"} 292 | ) 293 | assert response1.status_code == 200 294 | 295 | print("Второй запрос моделей...") 296 | response2 = test_client.get( 297 | "/v1/models", 298 | headers={"Authorization": f"Bearer {valid_proxy_api_key}"} 299 | ) 300 | assert response2.status_code == 200 301 | 302 | # Ответы должны быть идентичны 303 | assert response1.json()["data"] == response2.json()["data"] 304 | print("Кэширование работает корректно") 305 | 306 | 307 | class TestStreamingFlagHandling: 308 | """Integration-тесты обработки флага stream.""" 309 | 310 | def test_stream_true_accepted(self, test_client, valid_proxy_api_key): 311 | """ 312 | Что он делает: Проверяет, что stream=true принимается. 313 | Цель: Убедиться, что streaming режим доступен. 314 | 315 | Примечание: Для streaming режима нужен мок HTTP клиента, 316 | так как запрос выполняется внутри генератора. 317 | """ 318 | print("Запрос с stream=true...") 319 | 320 | # Создаём мок response для streaming 321 | mock_response = AsyncMock() 322 | mock_response.status_code = 200 323 | 324 | async def mock_aiter_bytes(): 325 | yield b'{"content":"Hello"}' 326 | yield b'{"usage":0.5}' 327 | 328 | mock_response.aiter_bytes = mock_aiter_bytes 329 | mock_response.aclose = AsyncMock() 330 | 331 | # Мокируем request_with_retry чтобы вернуть наш мок response 332 | with patch('kiro_gateway.routes.KiroHttpClient') as MockHttpClient: 333 | mock_client_instance = AsyncMock() 334 | mock_client_instance.request_with_retry = AsyncMock(return_value=mock_response) 335 | mock_client_instance.client = AsyncMock() 336 | mock_client_instance.close = AsyncMock() 337 | MockHttpClient.return_value = mock_client_instance 338 | 339 | response = test_client.post( 340 | "/v1/chat/completions", 341 | headers={"Authorization": f"Bearer {valid_proxy_api_key}"}, 342 | json={ 343 | "model": "claude-sonnet-4-5", 344 | "messages": [{"role": "user", "content": "Hello"}], 345 | "stream": True 346 | } 347 | ) 348 | 349 | # Валидация должна пройти и streaming должен работать 350 | assert response.status_code == 200 351 | print(f"stream=true: {response.status_code}") 352 | 353 | def test_stream_false_accepted(self, test_client, valid_proxy_api_key): 354 | """ 355 | Что он делает: Проверяет, что stream=false принимается. 356 | Цель: Убедиться, что non-streaming режим доступен. 357 | """ 358 | print("Запрос с stream=false...") 359 | response = test_client.post( 360 | "/v1/chat/completions", 361 | headers={"Authorization": f"Bearer {valid_proxy_api_key}"}, 362 | json={ 363 | "model": "claude-sonnet-4-5", 364 | "messages": [{"role": "user", "content": "Hello"}], 365 | "stream": False 366 | } 367 | ) 368 | 369 | # Валидация должна пройти 370 | assert response.status_code != 422 371 | print(f"stream=false: {response.status_code}") 372 | 373 | 374 | class TestHealthEndpointIntegration: 375 | """Integration-тесты health endpoints.""" 376 | 377 | def test_root_and_health_consistency(self, test_client): 378 | """ 379 | Что он делает: Проверяет консистентность / и /health. 380 | Цель: Убедиться, что оба эндпоинта возвращают корректный статус. 381 | """ 382 | print("Запрос к /...") 383 | root_response = test_client.get("/") 384 | 385 | print("Запрос к /health...") 386 | health_response = test_client.get("/health") 387 | 388 | assert root_response.status_code == 200 389 | assert health_response.status_code == 200 390 | 391 | # Оба должны показывать "ok" статус 392 | assert root_response.json()["status"] == "ok" 393 | assert health_response.json()["status"] == "healthy" 394 | 395 | # Версии должны совпадать 396 | assert root_response.json()["version"] == health_response.json()["version"] 397 | 398 | print("Health endpoints консистентны") -------------------------------------------------------------------------------- /kiro_gateway/parsers.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # Kiro OpenAI Gateway 4 | # Copyright (C) 2025 Jwadow 5 | # 6 | # This program is free software: you can redistribute it and/or modify 7 | # it under the terms of the GNU Affero General Public License as published by 8 | # the Free Software Foundation, either version 3 of the License, or 9 | # (at your option) any later version. 10 | # 11 | # This program is distributed in the hope that it will be useful, 12 | # but WITHOUT ANY WARRANTY; without even the implied warranty of 13 | # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 14 | # GNU Affero General Public License for more details. 15 | # 16 | # You should have received a copy of the GNU Affero General Public License 17 | # along with this program. If not, see . 18 | 19 | """ 20 | Парсеры для AWS Event Stream формата. 21 | 22 | Содержит классы и функции для: 23 | - Парсинга бинарного AWS SSE потока 24 | - Извлечения JSON событий 25 | - Обработки tool calls 26 | - Дедупликации контента 27 | """ 28 | 29 | import json 30 | import re 31 | from typing import Any, Dict, List, Optional 32 | 33 | from loguru import logger 34 | 35 | from kiro_gateway.utils import generate_tool_call_id 36 | 37 | 38 | def find_matching_brace(text: str, start_pos: int) -> int: 39 | """ 40 | Находит позицию закрывающей скобки с учётом вложенности и строк. 41 | 42 | Использует bracket counting для корректного парсинга вложенных JSON. 43 | Учитывает строки в кавычках и escape-последовательности. 44 | 45 | Args: 46 | text: Текст для поиска 47 | start_pos: Позиция открывающей скобки '{' 48 | 49 | Returns: 50 | Позиция закрывающей скобки или -1 если не найдена 51 | 52 | Example: 53 | >>> find_matching_brace('{"a": {"b": 1}}', 0) 54 | 14 55 | >>> find_matching_brace('{"a": "{}"}', 0) 56 | 10 57 | """ 58 | if start_pos >= len(text) or text[start_pos] != '{': 59 | return -1 60 | 61 | brace_count = 0 62 | in_string = False 63 | escape_next = False 64 | 65 | for i in range(start_pos, len(text)): 66 | char = text[i] 67 | 68 | if escape_next: 69 | escape_next = False 70 | continue 71 | 72 | if char == '\\' and in_string: 73 | escape_next = True 74 | continue 75 | 76 | if char == '"' and not escape_next: 77 | in_string = not in_string 78 | continue 79 | 80 | if not in_string: 81 | if char == '{': 82 | brace_count += 1 83 | elif char == '}': 84 | brace_count -= 1 85 | if brace_count == 0: 86 | return i 87 | 88 | return -1 89 | 90 | 91 | def parse_bracket_tool_calls(response_text: str) -> List[Dict[str, Any]]: 92 | """ 93 | Парсит tool calls в формате [Called func_name with args: {...}]. 94 | 95 | Некоторые модели возвращают tool calls в текстовом формате вместо 96 | структурированного JSON. Эта функция извлекает их. 97 | 98 | Args: 99 | response_text: Текст ответа модели 100 | 101 | Returns: 102 | Список tool calls в формате OpenAI 103 | 104 | Example: 105 | >>> text = "[Called get_weather with args: {\"city\": \"London\"}]" 106 | >>> calls = parse_bracket_tool_calls(text) 107 | >>> calls[0]["function"]["name"] 108 | 'get_weather' 109 | """ 110 | if not response_text or "[Called" not in response_text: 111 | return [] 112 | 113 | tool_calls = [] 114 | pattern = r'\[Called\s+(\w+)\s+with\s+args:\s*' 115 | 116 | for match in re.finditer(pattern, response_text, re.IGNORECASE): 117 | func_name = match.group(1) 118 | args_start = match.end() 119 | 120 | # Ищем начало JSON 121 | json_start = response_text.find('{', args_start) 122 | if json_start == -1: 123 | continue 124 | 125 | # Ищем конец JSON с учётом вложенности 126 | json_end = find_matching_brace(response_text, json_start) 127 | if json_end == -1: 128 | continue 129 | 130 | json_str = response_text[json_start:json_end + 1] 131 | 132 | try: 133 | args = json.loads(json_str) 134 | tool_call_id = generate_tool_call_id() 135 | # index будет добавлен позже при формировании финального ответа 136 | tool_calls.append({ 137 | "id": tool_call_id, 138 | "type": "function", 139 | "function": { 140 | "name": func_name, 141 | "arguments": json.dumps(args) 142 | } 143 | }) 144 | except json.JSONDecodeError: 145 | logger.warning(f"Failed to parse tool call arguments: {json_str[:100]}") 146 | 147 | return tool_calls 148 | 149 | 150 | def deduplicate_tool_calls(tool_calls: List[Dict[str, Any]]) -> List[Dict[str, Any]]: 151 | """ 152 | Удаляет дубликаты tool calls. 153 | 154 | Дедупликация происходит по двум критериям: 155 | 1. По id - если есть несколько tool calls с одинаковым id, оставляем тот у которого 156 | больше аргументов (не пустой "{}") 157 | 2. По name+arguments - удаляем полные дубликаты 158 | 159 | Args: 160 | tool_calls: Список tool calls 161 | 162 | Returns: 163 | Список уникальных tool calls 164 | """ 165 | # Сначала дедупликация по id - оставляем tool call с непустыми аргументами 166 | by_id: Dict[str, Dict[str, Any]] = {} 167 | for tc in tool_calls: 168 | tc_id = tc.get("id", "") 169 | if not tc_id: 170 | # Без id - добавляем как есть (будет дедуплицировано по name+args) 171 | continue 172 | 173 | existing = by_id.get(tc_id) 174 | if existing is None: 175 | by_id[tc_id] = tc 176 | else: 177 | # Есть дубликат по id - оставляем тот у которого больше аргументов 178 | existing_args = existing.get("function", {}).get("arguments", "{}") 179 | current_args = tc.get("function", {}).get("arguments", "{}") 180 | 181 | # Предпочитаем непустые аргументы 182 | if current_args != "{}" and (existing_args == "{}" or len(current_args) > len(existing_args)): 183 | logger.debug(f"Replacing tool call {tc_id} with better arguments: {len(existing_args)} -> {len(current_args)}") 184 | by_id[tc_id] = tc 185 | 186 | # Собираем tool calls: сначала те что с id, потом без id 187 | result_with_id = list(by_id.values()) 188 | result_without_id = [tc for tc in tool_calls if not tc.get("id")] 189 | 190 | # Теперь дедупликация по name+arguments для всех 191 | seen = set() 192 | unique = [] 193 | 194 | for tc in result_with_id + result_without_id: 195 | # Защита от None в function 196 | func = tc.get("function") or {} 197 | func_name = func.get("name") or "" 198 | func_args = func.get("arguments") or "{}" 199 | key = f"{func_name}-{func_args}" 200 | if key not in seen: 201 | seen.add(key) 202 | unique.append(tc) 203 | 204 | if len(tool_calls) != len(unique): 205 | logger.debug(f"Deduplicated tool calls: {len(tool_calls)} -> {len(unique)}") 206 | 207 | return unique 208 | 209 | 210 | class AwsEventStreamParser: 211 | """ 212 | Парсер для AWS Event Stream формата. 213 | 214 | AWS возвращает события в бинарном формате с разделителями :message-type...event. 215 | Этот класс извлекает JSON события из потока и преобразует их в удобный формат. 216 | 217 | Поддерживаемые типы событий: 218 | - content: Текстовый контент ответа 219 | - tool_start: Начало tool call (name, toolUseId) 220 | - tool_input: Продолжение input для tool call 221 | - tool_stop: Завершение tool call 222 | - usage: Информация о потреблении кредитов 223 | - context_usage: Процент использования контекста 224 | 225 | Attributes: 226 | buffer: Буфер для накопления данных 227 | last_content: Последний обработанный контент (для дедупликации) 228 | current_tool_call: Текущий незавершённый tool call 229 | tool_calls: Список завершённых tool calls 230 | 231 | Example: 232 | >>> parser = AwsEventStreamParser() 233 | >>> events = parser.feed(chunk) 234 | >>> for event in events: 235 | ... if event["type"] == "content": 236 | ... print(event["data"]) 237 | """ 238 | 239 | # Паттерны для поиска JSON событий 240 | EVENT_PATTERNS = [ 241 | ('{"content":', 'content'), 242 | ('{"name":', 'tool_start'), 243 | ('{"input":', 'tool_input'), 244 | ('{"stop":', 'tool_stop'), 245 | ('{"followupPrompt":', 'followup'), 246 | ('{"usage":', 'usage'), 247 | ('{"contextUsagePercentage":', 'context_usage'), 248 | ] 249 | 250 | def __init__(self): 251 | """Инициализирует парсер.""" 252 | self.buffer = "" 253 | self.last_content: Optional[str] = None # Для дедупликации повторяющегося контента 254 | self.current_tool_call: Optional[Dict[str, Any]] = None 255 | self.tool_calls: List[Dict[str, Any]] = [] 256 | 257 | def feed(self, chunk: bytes) -> List[Dict[str, Any]]: 258 | """ 259 | Добавляет chunk в буфер и возвращает распарсенные события. 260 | 261 | Args: 262 | chunk: Байты данных из потока 263 | 264 | Returns: 265 | Список событий в формате {"type": str, "data": Any} 266 | """ 267 | try: 268 | self.buffer += chunk.decode('utf-8', errors='ignore') 269 | except Exception: 270 | return [] 271 | 272 | events = [] 273 | 274 | while True: 275 | # Находим ближайший паттерн 276 | earliest_pos = -1 277 | earliest_type = None 278 | 279 | for pattern, event_type in self.EVENT_PATTERNS: 280 | pos = self.buffer.find(pattern) 281 | if pos != -1 and (earliest_pos == -1 or pos < earliest_pos): 282 | earliest_pos = pos 283 | earliest_type = event_type 284 | 285 | if earliest_pos == -1: 286 | break 287 | 288 | # Ищем конец JSON 289 | json_end = find_matching_brace(self.buffer, earliest_pos) 290 | if json_end == -1: 291 | # JSON не полный, ждём больше данных 292 | break 293 | 294 | json_str = self.buffer[earliest_pos:json_end + 1] 295 | self.buffer = self.buffer[json_end + 1:] 296 | 297 | try: 298 | data = json.loads(json_str) 299 | event = self._process_event(data, earliest_type) 300 | if event: 301 | events.append(event) 302 | except json.JSONDecodeError: 303 | logger.warning(f"Failed to parse JSON: {json_str[:100]}") 304 | 305 | return events 306 | 307 | def _process_event(self, data: dict, event_type: str) -> Optional[Dict[str, Any]]: 308 | """ 309 | Обрабатывает распарсенное событие. 310 | 311 | Args: 312 | data: Распарсенный JSON 313 | event_type: Тип события 314 | 315 | Returns: 316 | Обработанное событие или None 317 | """ 318 | if event_type == 'content': 319 | return self._process_content_event(data) 320 | elif event_type == 'tool_start': 321 | return self._process_tool_start_event(data) 322 | elif event_type == 'tool_input': 323 | return self._process_tool_input_event(data) 324 | elif event_type == 'tool_stop': 325 | return self._process_tool_stop_event(data) 326 | elif event_type == 'usage': 327 | return {"type": "usage", "data": data.get('usage', 0)} 328 | elif event_type == 'context_usage': 329 | return {"type": "context_usage", "data": data.get('contextUsagePercentage', 0)} 330 | 331 | return None 332 | 333 | def _process_content_event(self, data: dict) -> Optional[Dict[str, Any]]: 334 | """Обрабатывает событие с контентом.""" 335 | content = data.get('content', '') 336 | 337 | # Пропускаем followupPrompt 338 | if data.get('followupPrompt'): 339 | return None 340 | 341 | # Дедупликация повторяющегося контента 342 | if content == self.last_content: 343 | return None 344 | 345 | self.last_content = content 346 | 347 | return {"type": "content", "data": content} 348 | 349 | def _process_tool_start_event(self, data: dict) -> Optional[Dict[str, Any]]: 350 | """Обрабатывает начало tool call.""" 351 | # Завершаем предыдущий tool call если есть 352 | if self.current_tool_call: 353 | self._finalize_tool_call() 354 | 355 | # input может быть строкой или объектом 356 | input_data = data.get('input', '') 357 | if isinstance(input_data, dict): 358 | input_str = json.dumps(input_data) 359 | else: 360 | input_str = str(input_data) if input_data else '' 361 | 362 | self.current_tool_call = { 363 | "id": data.get('toolUseId', generate_tool_call_id()), 364 | "type": "function", 365 | "function": { 366 | "name": data.get('name', ''), 367 | "arguments": input_str 368 | } 369 | } 370 | 371 | if data.get('stop'): 372 | self._finalize_tool_call() 373 | 374 | return None 375 | 376 | def _process_tool_input_event(self, data: dict) -> Optional[Dict[str, Any]]: 377 | """Обрабатывает продолжение input для tool call.""" 378 | if self.current_tool_call: 379 | # input может быть строкой или объектом 380 | input_data = data.get('input', '') 381 | if isinstance(input_data, dict): 382 | input_str = json.dumps(input_data) 383 | else: 384 | input_str = str(input_data) if input_data else '' 385 | self.current_tool_call['function']['arguments'] += input_str 386 | return None 387 | 388 | def _process_tool_stop_event(self, data: dict) -> Optional[Dict[str, Any]]: 389 | """Обрабатывает завершение tool call.""" 390 | if self.current_tool_call and data.get('stop'): 391 | self._finalize_tool_call() 392 | return None 393 | 394 | def _finalize_tool_call(self) -> None: 395 | """Завершает текущий tool call и добавляет в список.""" 396 | if not self.current_tool_call: 397 | return 398 | 399 | # Пытаемся распарсить и нормализовать arguments как JSON 400 | args = self.current_tool_call['function']['arguments'] 401 | tool_name = self.current_tool_call['function'].get('name', 'unknown') 402 | 403 | logger.debug(f"Finalizing tool call '{tool_name}' with raw arguments: {repr(args)[:200]}") 404 | 405 | if isinstance(args, str): 406 | if args.strip(): 407 | try: 408 | parsed = json.loads(args) 409 | # Убеждаемся что результат - строка JSON 410 | self.current_tool_call['function']['arguments'] = json.dumps(parsed) 411 | logger.debug(f"Tool '{tool_name}' arguments parsed successfully: {list(parsed.keys()) if isinstance(parsed, dict) else type(parsed)}") 412 | except json.JSONDecodeError as e: 413 | # Если не удалось распарсить, оставляем пустой объект 414 | logger.warning(f"Failed to parse tool '{tool_name}' arguments: {e}. Raw: {args[:200]}") 415 | self.current_tool_call['function']['arguments'] = "{}" 416 | else: 417 | # Пустая строка - используем пустой объект 418 | # Это нормальное поведение для дубликатов tool calls от Kiro 419 | logger.debug(f"Tool '{tool_name}' has empty arguments string (will be deduplicated)") 420 | self.current_tool_call['function']['arguments'] = "{}" 421 | elif isinstance(args, dict): 422 | # Если уже объект - сериализуем в строку 423 | self.current_tool_call['function']['arguments'] = json.dumps(args) 424 | logger.debug(f"Tool '{tool_name}' arguments already dict with keys: {list(args.keys())}") 425 | else: 426 | # Неизвестный тип - пустой объект 427 | logger.warning(f"Tool '{tool_name}' has unexpected arguments type: {type(args)}") 428 | self.current_tool_call['function']['arguments'] = "{}" 429 | 430 | self.tool_calls.append(self.current_tool_call) 431 | self.current_tool_call = None 432 | 433 | def get_tool_calls(self) -> List[Dict[str, Any]]: 434 | """ 435 | Возвращает все собранные tool calls. 436 | 437 | Завершает текущий tool call если он не завершён. 438 | Удаляет дубликаты. 439 | 440 | Returns: 441 | Список уникальных tool calls 442 | """ 443 | if self.current_tool_call: 444 | self._finalize_tool_call() 445 | return deduplicate_tool_calls(self.tool_calls) 446 | 447 | def reset(self) -> None: 448 | """Сбрасывает состояние парсера.""" 449 | self.buffer = "" 450 | self.last_content = None 451 | self.current_tool_call = None 452 | self.tool_calls = [] -------------------------------------------------------------------------------- /tests/unit/test_cache.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | """ 4 | Unit-тесты для ModelInfoCache. 5 | Проверяет логику кэширования метаданных моделей. 6 | """ 7 | 8 | import asyncio 9 | import time 10 | import pytest 11 | 12 | from kiro_gateway.cache import ModelInfoCache 13 | from kiro_gateway.config import DEFAULT_MAX_INPUT_TOKENS 14 | 15 | 16 | class TestModelInfoCacheInitialization: 17 | """Тесты инициализации ModelInfoCache.""" 18 | 19 | def test_initialization_creates_empty_cache(self): 20 | """ 21 | Что он делает: Проверяет, что кэш создаётся пустым. 22 | Цель: Убедиться в корректной инициализации. 23 | """ 24 | print("Настройка: Создание ModelInfoCache...") 25 | cache = ModelInfoCache() 26 | 27 | print("Проверка: Кэш пуст при создании...") 28 | print(f"Сравниваем is_empty(): Ожидалось True, Получено {cache.is_empty()}") 29 | assert cache.is_empty() is True 30 | 31 | print(f"Сравниваем size: Ожидалось 0, Получено {cache.size}") 32 | assert cache.size == 0 33 | 34 | def test_initialization_with_custom_ttl(self): 35 | """ 36 | Что он делает: Проверяет создание кэша с кастомным TTL. 37 | Цель: Убедиться, что TTL можно настроить. 38 | """ 39 | print("Настройка: Создание ModelInfoCache с TTL=7200...") 40 | cache = ModelInfoCache(cache_ttl=7200) 41 | 42 | print("Проверка: TTL установлен корректно...") 43 | print(f"Сравниваем _cache_ttl: Ожидалось 7200, Получено {cache._cache_ttl}") 44 | assert cache._cache_ttl == 7200 45 | 46 | def test_initialization_last_update_is_none(self): 47 | """ 48 | Что он делает: Проверяет, что last_update_time изначально None. 49 | Цель: Убедиться, что время обновления не установлено до первого update. 50 | """ 51 | print("Настройка: Создание ModelInfoCache...") 52 | cache = ModelInfoCache() 53 | 54 | print("Проверка: last_update_time изначально None...") 55 | print(f"Сравниваем last_update_time: Ожидалось None, Получено {cache.last_update_time}") 56 | assert cache.last_update_time is None 57 | 58 | 59 | class TestModelInfoCacheUpdate: 60 | """Тесты обновления кэша.""" 61 | 62 | @pytest.mark.asyncio 63 | async def test_update_populates_cache(self, sample_models_data): 64 | """ 65 | Что он делает: Проверяет заполнение кэша данными. 66 | Цель: Убедиться, что update() корректно сохраняет модели. 67 | """ 68 | print("Настройка: Создание ModelInfoCache...") 69 | cache = ModelInfoCache() 70 | 71 | print(f"Действие: Обновление кэша с {len(sample_models_data)} моделями...") 72 | await cache.update(sample_models_data) 73 | 74 | print("Проверка: Кэш заполнен...") 75 | print(f"Сравниваем is_empty(): Ожидалось False, Получено {cache.is_empty()}") 76 | assert cache.is_empty() is False 77 | 78 | print(f"Сравниваем size: Ожидалось {len(sample_models_data)}, Получено {cache.size}") 79 | assert cache.size == len(sample_models_data) 80 | 81 | @pytest.mark.asyncio 82 | async def test_update_sets_last_update_time(self, sample_models_data): 83 | """ 84 | Что он делает: Проверяет установку времени последнего обновления. 85 | Цель: Убедиться, что last_update_time устанавливается после update. 86 | """ 87 | print("Настройка: Создание ModelInfoCache...") 88 | cache = ModelInfoCache() 89 | 90 | before_update = time.time() 91 | print(f"Действие: Обновление кэша (время до: {before_update})...") 92 | await cache.update(sample_models_data) 93 | after_update = time.time() 94 | 95 | print("Проверка: last_update_time установлен в разумных пределах...") 96 | print(f"last_update_time: {cache.last_update_time}") 97 | assert cache.last_update_time is not None 98 | assert before_update <= cache.last_update_time <= after_update 99 | 100 | @pytest.mark.asyncio 101 | async def test_update_replaces_existing_data(self, sample_models_data): 102 | """ 103 | Что он делает: Проверяет замену данных при повторном update. 104 | Цель: Убедиться, что старые данные полностью заменяются. 105 | """ 106 | print("Настройка: Создание ModelInfoCache и первое обновление...") 107 | cache = ModelInfoCache() 108 | await cache.update(sample_models_data) 109 | 110 | print("Действие: Обновление с новыми данными...") 111 | new_data = [{"modelId": "new-model", "tokenLimits": {"maxInputTokens": 50000}}] 112 | await cache.update(new_data) 113 | 114 | print("Проверка: Старые данные заменены...") 115 | print(f"Сравниваем size: Ожидалось 1, Получено {cache.size}") 116 | assert cache.size == 1 117 | 118 | print("Проверка: Старая модель недоступна...") 119 | assert cache.get("claude-sonnet-4") is None 120 | 121 | print("Проверка: Новая модель доступна...") 122 | assert cache.get("new-model") is not None 123 | 124 | @pytest.mark.asyncio 125 | async def test_update_with_empty_list(self): 126 | """ 127 | Что он делает: Проверяет обновление пустым списком. 128 | Цель: Убедиться, что кэш очищается при пустом update. 129 | """ 130 | print("Настройка: Создание ModelInfoCache с данными...") 131 | cache = ModelInfoCache() 132 | await cache.update([{"modelId": "test-model"}]) 133 | 134 | print("Действие: Обновление пустым списком...") 135 | await cache.update([]) 136 | 137 | print("Проверка: Кэш пуст...") 138 | print(f"Сравниваем is_empty(): Ожидалось True, Получено {cache.is_empty()}") 139 | assert cache.is_empty() is True 140 | 141 | 142 | class TestModelInfoCacheGet: 143 | """Тесты получения данных из кэша.""" 144 | 145 | @pytest.mark.asyncio 146 | async def test_get_returns_model_info(self, sample_models_data): 147 | """ 148 | Что он делает: Проверяет получение информации о модели. 149 | Цель: Убедиться, что get() возвращает корректные данные. 150 | """ 151 | print("Настройка: Создание и заполнение кэша...") 152 | cache = ModelInfoCache() 153 | await cache.update(sample_models_data) 154 | 155 | print("Действие: Получение информации о claude-sonnet-4...") 156 | model_info = cache.get("claude-sonnet-4") 157 | 158 | print("Проверка: Информация получена...") 159 | print(f"model_info: {model_info}") 160 | assert model_info is not None 161 | assert model_info["modelId"] == "claude-sonnet-4" 162 | 163 | @pytest.mark.asyncio 164 | async def test_get_returns_none_for_unknown_model(self, sample_models_data): 165 | """ 166 | Что он делает: Проверяет возврат None для неизвестной модели. 167 | Цель: Убедиться, что get() не падает при отсутствии модели. 168 | """ 169 | print("Настройка: Создание и заполнение кэша...") 170 | cache = ModelInfoCache() 171 | await cache.update(sample_models_data) 172 | 173 | print("Действие: Получение информации о несуществующей модели...") 174 | model_info = cache.get("non-existent-model") 175 | 176 | print("Проверка: Возвращён None...") 177 | print(f"Сравниваем model_info: Ожидалось None, Получено {model_info}") 178 | assert model_info is None 179 | 180 | def test_get_from_empty_cache(self): 181 | """ 182 | Что он делает: Проверяет get() из пустого кэша. 183 | Цель: Убедиться, что пустой кэш не вызывает ошибок. 184 | """ 185 | print("Настройка: Создание пустого кэша...") 186 | cache = ModelInfoCache() 187 | 188 | print("Действие: Получение из пустого кэша...") 189 | model_info = cache.get("any-model") 190 | 191 | print("Проверка: Возвращён None...") 192 | print(f"Сравниваем model_info: Ожидалось None, Получено {model_info}") 193 | assert model_info is None 194 | 195 | 196 | class TestModelInfoCacheGetMaxInputTokens: 197 | """Тесты получения maxInputTokens.""" 198 | 199 | @pytest.mark.asyncio 200 | async def test_get_max_input_tokens_returns_value(self, sample_models_data): 201 | """ 202 | Что он делает: Проверяет получение maxInputTokens для модели. 203 | Цель: Убедиться, что значение извлекается из tokenLimits. 204 | """ 205 | print("Настройка: Создание и заполнение кэша...") 206 | cache = ModelInfoCache() 207 | await cache.update(sample_models_data) 208 | 209 | print("Действие: Получение maxInputTokens для claude-sonnet-4...") 210 | max_tokens = cache.get_max_input_tokens("claude-sonnet-4") 211 | 212 | print("Проверка: Значение корректно...") 213 | print(f"Сравниваем max_tokens: Ожидалось 200000, Получено {max_tokens}") 214 | assert max_tokens == 200000 215 | 216 | @pytest.mark.asyncio 217 | async def test_get_max_input_tokens_returns_default_for_unknown(self, sample_models_data): 218 | """ 219 | Что он делает: Проверяет возврат дефолта для неизвестной модели. 220 | Цель: Убедиться, что возвращается DEFAULT_MAX_INPUT_TOKENS. 221 | """ 222 | print("Настройка: Создание и заполнение кэша...") 223 | cache = ModelInfoCache() 224 | await cache.update(sample_models_data) 225 | 226 | print("Действие: Получение maxInputTokens для неизвестной модели...") 227 | max_tokens = cache.get_max_input_tokens("unknown-model") 228 | 229 | print("Проверка: Возвращён дефолт...") 230 | print(f"Сравниваем max_tokens: Ожидалось {DEFAULT_MAX_INPUT_TOKENS}, Получено {max_tokens}") 231 | assert max_tokens == DEFAULT_MAX_INPUT_TOKENS 232 | 233 | @pytest.mark.asyncio 234 | async def test_get_max_input_tokens_returns_default_when_no_token_limits(self): 235 | """ 236 | Что он делает: Проверяет возврат дефолта при отсутствии tokenLimits. 237 | Цель: Убедиться, что модель без tokenLimits не ломает логику. 238 | """ 239 | print("Настройка: Создание кэша с моделью без tokenLimits...") 240 | cache = ModelInfoCache() 241 | await cache.update([{"modelId": "model-without-limits"}]) 242 | 243 | print("Действие: Получение maxInputTokens...") 244 | max_tokens = cache.get_max_input_tokens("model-without-limits") 245 | 246 | print("Проверка: Возвращён дефолт...") 247 | print(f"Сравниваем max_tokens: Ожидалось {DEFAULT_MAX_INPUT_TOKENS}, Получено {max_tokens}") 248 | assert max_tokens == DEFAULT_MAX_INPUT_TOKENS 249 | 250 | @pytest.mark.asyncio 251 | async def test_get_max_input_tokens_returns_default_when_max_input_is_none(self): 252 | """ 253 | Что он делает: Проверяет возврат дефолта при maxInputTokens=None. 254 | Цель: Убедиться, что None в tokenLimits обрабатывается корректно. 255 | """ 256 | print("Настройка: Создание кэша с моделью с maxInputTokens=None...") 257 | cache = ModelInfoCache() 258 | await cache.update([{ 259 | "modelId": "model-with-null", 260 | "tokenLimits": {"maxInputTokens": None} 261 | }]) 262 | 263 | print("Действие: Получение maxInputTokens...") 264 | max_tokens = cache.get_max_input_tokens("model-with-null") 265 | 266 | print("Проверка: Возвращён дефолт...") 267 | print(f"Сравниваем max_tokens: Ожидалось {DEFAULT_MAX_INPUT_TOKENS}, Получено {max_tokens}") 268 | assert max_tokens == DEFAULT_MAX_INPUT_TOKENS 269 | 270 | 271 | class TestModelInfoCacheIsEmpty: 272 | """Тесты проверки пустоты кэша.""" 273 | 274 | def test_is_empty_returns_true_for_new_cache(self): 275 | """ 276 | Что он делает: Проверяет is_empty() для нового кэша. 277 | Цель: Убедиться, что новый кэш считается пустым. 278 | """ 279 | print("Настройка: Создание нового кэша...") 280 | cache = ModelInfoCache() 281 | 282 | print("Проверка: is_empty() возвращает True...") 283 | print(f"Сравниваем is_empty(): Ожидалось True, Получено {cache.is_empty()}") 284 | assert cache.is_empty() is True 285 | 286 | @pytest.mark.asyncio 287 | async def test_is_empty_returns_false_after_update(self, sample_models_data): 288 | """ 289 | Что он делает: Проверяет is_empty() после заполнения. 290 | Цель: Убедиться, что заполненный кэш не считается пустым. 291 | """ 292 | print("Настройка: Создание и заполнение кэша...") 293 | cache = ModelInfoCache() 294 | await cache.update(sample_models_data) 295 | 296 | print("Проверка: is_empty() возвращает False...") 297 | print(f"Сравниваем is_empty(): Ожидалось False, Получено {cache.is_empty()}") 298 | assert cache.is_empty() is False 299 | 300 | 301 | class TestModelInfoCacheIsStale: 302 | """Тесты проверки устаревания кэша.""" 303 | 304 | def test_is_stale_returns_true_for_new_cache(self): 305 | """ 306 | Что он делает: Проверяет is_stale() для нового кэша. 307 | Цель: Убедиться, что кэш без обновлений считается устаревшим. 308 | """ 309 | print("Настройка: Создание нового кэша...") 310 | cache = ModelInfoCache() 311 | 312 | print("Проверка: is_stale() возвращает True...") 313 | print(f"Сравниваем is_stale(): Ожидалось True, Получено {cache.is_stale()}") 314 | assert cache.is_stale() is True 315 | 316 | @pytest.mark.asyncio 317 | async def test_is_stale_returns_false_after_recent_update(self, sample_models_data): 318 | """ 319 | Что он делает: Проверяет is_stale() сразу после обновления. 320 | Цель: Убедиться, что свежий кэш не считается устаревшим. 321 | """ 322 | print("Настройка: Создание и заполнение кэша...") 323 | cache = ModelInfoCache() 324 | await cache.update(sample_models_data) 325 | 326 | print("Проверка: is_stale() возвращает False...") 327 | print(f"Сравниваем is_stale(): Ожидалось False, Получено {cache.is_stale()}") 328 | assert cache.is_stale() is False 329 | 330 | @pytest.mark.asyncio 331 | async def test_is_stale_returns_true_after_ttl_expires(self, sample_models_data): 332 | """ 333 | Что он делает: Проверяет is_stale() после истечения TTL. 334 | Цель: Убедиться, что кэш считается устаревшим после TTL. 335 | """ 336 | print("Настройка: Создание кэша с TTL=0.1 секунды...") 337 | cache = ModelInfoCache(cache_ttl=0.1) 338 | await cache.update(sample_models_data) 339 | 340 | print("Действие: Ожидание истечения TTL...") 341 | await asyncio.sleep(0.2) 342 | 343 | print("Проверка: is_stale() возвращает True...") 344 | print(f"Сравниваем is_stale(): Ожидалось True, Получено {cache.is_stale()}") 345 | assert cache.is_stale() is True 346 | 347 | 348 | class TestModelInfoCacheGetAllModelIds: 349 | """Тесты получения списка ID моделей.""" 350 | 351 | def test_get_all_model_ids_returns_empty_for_new_cache(self): 352 | """ 353 | Что он делает: Проверяет get_all_model_ids() для пустого кэша. 354 | Цель: Убедиться, что возвращается пустой список. 355 | """ 356 | print("Настройка: Создание пустого кэша...") 357 | cache = ModelInfoCache() 358 | 359 | print("Действие: Получение списка ID моделей...") 360 | model_ids = cache.get_all_model_ids() 361 | 362 | print("Проверка: Список пуст...") 363 | print(f"Сравниваем model_ids: Ожидалось [], Получено {model_ids}") 364 | assert model_ids == [] 365 | 366 | @pytest.mark.asyncio 367 | async def test_get_all_model_ids_returns_all_ids(self, sample_models_data): 368 | """ 369 | Что он делает: Проверяет get_all_model_ids() для заполненного кэша. 370 | Цель: Убедиться, что возвращаются все ID моделей. 371 | """ 372 | print("Настройка: Создание и заполнение кэша...") 373 | cache = ModelInfoCache() 374 | await cache.update(sample_models_data) 375 | 376 | print("Действие: Получение списка ID моделей...") 377 | model_ids = cache.get_all_model_ids() 378 | 379 | print("Проверка: Все ID присутствуют...") 380 | expected_ids = [m["modelId"] for m in sample_models_data] 381 | print(f"Сравниваем model_ids: Ожидалось {expected_ids}, Получено {model_ids}") 382 | assert set(model_ids) == set(expected_ids) 383 | 384 | 385 | class TestModelInfoCacheThreadSafety: 386 | """Тесты потокобезопасности кэша.""" 387 | 388 | @pytest.mark.asyncio 389 | async def test_concurrent_updates_dont_corrupt_cache(self, sample_models_data): 390 | """ 391 | Что он делает: Проверяет потокобезопасность при параллельных update. 392 | Цель: Убедиться, что asyncio.Lock защищает от race conditions. 393 | """ 394 | print("Настройка: Создание кэша...") 395 | cache = ModelInfoCache() 396 | 397 | async def update_with_data(data): 398 | await cache.update(data) 399 | 400 | print("Действие: 10 параллельных обновлений...") 401 | tasks = [] 402 | for i in range(10): 403 | data = [{"modelId": f"model-{i}", "tokenLimits": {"maxInputTokens": 100000 + i}}] 404 | tasks.append(update_with_data(data)) 405 | 406 | await asyncio.gather(*tasks) 407 | 408 | print("Проверка: Кэш содержит данные последнего обновления...") 409 | # Из-за race condition, мы не знаем какое обновление было последним, 410 | # но кэш должен содержать ровно одну модель 411 | print(f"Сравниваем size: Ожидалось 1, Получено {cache.size}") 412 | assert cache.size == 1 413 | 414 | print("Проверка: Кэш не повреждён...") 415 | model_ids = cache.get_all_model_ids() 416 | assert len(model_ids) == 1 417 | assert model_ids[0].startswith("model-") 418 | 419 | @pytest.mark.asyncio 420 | async def test_concurrent_reads_are_safe(self, sample_models_data): 421 | """ 422 | Что он делает: Проверяет безопасность параллельных чтений. 423 | Цель: Убедиться, что множественные get() не вызывают проблем. 424 | """ 425 | print("Настройка: Создание и заполнение кэша...") 426 | cache = ModelInfoCache() 427 | await cache.update(sample_models_data) 428 | 429 | print("Действие: 100 параллельных чтений...") 430 | async def read_model(): 431 | return cache.get("claude-sonnet-4") 432 | 433 | results = await asyncio.gather(*[read_model() for _ in range(100)]) 434 | 435 | print("Проверка: Все чтения вернули одинаковый результат...") 436 | assert all(r is not None for r in results) 437 | assert all(r["modelId"] == "claude-sonnet-4" for r in results) -------------------------------------------------------------------------------- /tests/unit/test_streaming.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | """ 4 | Unit-тесты для streaming модуля. 5 | Проверяет логику добавления index к tool_calls и защиту от None значений. 6 | """ 7 | 8 | import pytest 9 | import json 10 | from unittest.mock import AsyncMock, MagicMock, patch 11 | 12 | from kiro_gateway.streaming import ( 13 | stream_kiro_to_openai, 14 | collect_stream_response 15 | ) 16 | 17 | 18 | @pytest.fixture 19 | def mock_model_cache(): 20 | """Мок для ModelInfoCache.""" 21 | cache = MagicMock() 22 | cache.get_max_input_tokens.return_value = 200000 23 | return cache 24 | 25 | 26 | @pytest.fixture 27 | def mock_auth_manager(): 28 | """Мок для KiroAuthManager.""" 29 | manager = MagicMock() 30 | return manager 31 | 32 | 33 | @pytest.fixture 34 | def mock_http_client(): 35 | """Мок для httpx.AsyncClient.""" 36 | client = AsyncMock() 37 | return client 38 | 39 | 40 | class TestStreamingToolCallsIndex: 41 | """Тесты для добавления index к tool_calls в streaming ответах.""" 42 | 43 | @pytest.mark.asyncio 44 | async def test_tool_calls_have_index_field(self, mock_http_client, mock_model_cache, mock_auth_manager): 45 | """ 46 | Что он делает: Проверяет, что tool_calls в streaming ответе содержат поле index. 47 | Цель: Убедиться, что OpenAI API spec соблюдается для streaming tool calls. 48 | """ 49 | print("Настройка: Мок tool calls без index...") 50 | tool_calls = [ 51 | { 52 | "id": "call_123", 53 | "type": "function", 54 | "function": { 55 | "name": "get_weather", 56 | "arguments": '{"location": "Moscow"}' 57 | } 58 | } 59 | ] 60 | 61 | print("Настройка: Мок парсера...") 62 | mock_parser = MagicMock() 63 | mock_parser.feed.return_value = [] 64 | mock_parser.get_tool_calls.return_value = tool_calls 65 | 66 | print("Настройка: Мок response...") 67 | mock_response = AsyncMock() 68 | mock_response.status_code = 200 69 | 70 | async def mock_aiter_bytes(): 71 | yield b'{"content":"test"}' 72 | 73 | mock_response.aiter_bytes = mock_aiter_bytes 74 | mock_response.aclose = AsyncMock() 75 | 76 | print("Действие: Сбор streaming chunks...") 77 | chunks = [] 78 | 79 | with patch('kiro_gateway.streaming.AwsEventStreamParser', return_value=mock_parser): 80 | with patch('kiro_gateway.streaming.parse_bracket_tool_calls', return_value=[]): 81 | async for chunk in stream_kiro_to_openai( 82 | mock_http_client, mock_response, "test-model", 83 | mock_model_cache, mock_auth_manager 84 | ): 85 | chunks.append(chunk) 86 | 87 | print(f"Получено chunks: {len(chunks)}") 88 | 89 | # Ищем chunk с tool_calls 90 | tool_calls_found = False 91 | for chunk in chunks: 92 | if isinstance(chunk, str) and "tool_calls" in chunk: 93 | if chunk.startswith("data: "): 94 | json_str = chunk[6:].strip() 95 | if json_str != "[DONE]": 96 | data = json.loads(json_str) 97 | if "choices" in data and data["choices"]: 98 | delta = data["choices"][0].get("delta", {}) 99 | if "tool_calls" in delta: 100 | print(f"Tool calls в delta: {delta['tool_calls']}") 101 | for tc in delta["tool_calls"]: 102 | print(f"Проверяем index в tool call: {tc}") 103 | assert "index" in tc, "Tool call должен содержать поле index" 104 | tool_calls_found = True 105 | 106 | assert tool_calls_found, "Tool calls chunk не найден" 107 | 108 | @pytest.mark.asyncio 109 | async def test_multiple_tool_calls_have_sequential_indices(self, mock_http_client, mock_model_cache, mock_auth_manager): 110 | """ 111 | Что он делает: Проверяет, что несколько tool_calls имеют последовательные индексы. 112 | Цель: Убедиться, что индексы начинаются с 0 и идут последовательно. 113 | """ 114 | print("Настройка: Несколько tool calls...") 115 | tool_calls = [ 116 | {"id": "call_1", "type": "function", "function": {"name": "func1", "arguments": "{}"}}, 117 | {"id": "call_2", "type": "function", "function": {"name": "func2", "arguments": "{}"}}, 118 | {"id": "call_3", "type": "function", "function": {"name": "func3", "arguments": "{}"}} 119 | ] 120 | 121 | print("Настройка: Мок парсера...") 122 | mock_parser = MagicMock() 123 | mock_parser.feed.return_value = [] 124 | mock_parser.get_tool_calls.return_value = tool_calls 125 | 126 | print("Настройка: Мок response...") 127 | mock_response = AsyncMock() 128 | mock_response.status_code = 200 129 | 130 | async def mock_aiter_bytes(): 131 | yield b'{"content":"test"}' 132 | 133 | mock_response.aiter_bytes = mock_aiter_bytes 134 | mock_response.aclose = AsyncMock() 135 | 136 | print("Действие: Сбор streaming chunks...") 137 | chunks = [] 138 | 139 | with patch('kiro_gateway.streaming.AwsEventStreamParser', return_value=mock_parser): 140 | with patch('kiro_gateway.streaming.parse_bracket_tool_calls', return_value=[]): 141 | async for chunk in stream_kiro_to_openai( 142 | mock_http_client, mock_response, "test-model", 143 | mock_model_cache, mock_auth_manager 144 | ): 145 | chunks.append(chunk) 146 | 147 | # Ищем chunk с tool_calls 148 | for chunk in chunks: 149 | if isinstance(chunk, str) and "tool_calls" in chunk: 150 | if chunk.startswith("data: "): 151 | json_str = chunk[6:].strip() 152 | if json_str != "[DONE]": 153 | data = json.loads(json_str) 154 | if "choices" in data and data["choices"]: 155 | delta = data["choices"][0].get("delta", {}) 156 | if "tool_calls" in delta: 157 | indices = [tc["index"] for tc in delta["tool_calls"]] 158 | print(f"Индексы: {indices}") 159 | assert indices == [0, 1, 2], f"Индексы должны быть [0, 1, 2], получено {indices}" 160 | 161 | 162 | class TestStreamingToolCallsNoneProtection: 163 | """Тесты для защиты от None значений в tool_calls.""" 164 | 165 | @pytest.mark.asyncio 166 | async def test_handles_none_function_name(self, mock_http_client, mock_model_cache, mock_auth_manager): 167 | """ 168 | Что он делает: Проверяет обработку None в function.name. 169 | Цель: Убедиться, что None заменяется на пустую строку. 170 | """ 171 | print("Настройка: Tool call с None name...") 172 | tool_calls = [ 173 | { 174 | "id": "call_1", 175 | "type": "function", 176 | "function": { 177 | "name": None, 178 | "arguments": '{"a": 1}' 179 | } 180 | } 181 | ] 182 | 183 | mock_parser = MagicMock() 184 | mock_parser.feed.return_value = [] 185 | mock_parser.get_tool_calls.return_value = tool_calls 186 | 187 | mock_response = AsyncMock() 188 | mock_response.status_code = 200 189 | 190 | async def mock_aiter_bytes(): 191 | yield b'{"content":"test"}' 192 | 193 | mock_response.aiter_bytes = mock_aiter_bytes 194 | mock_response.aclose = AsyncMock() 195 | 196 | print("Действие: Сбор streaming chunks...") 197 | chunks = [] 198 | 199 | with patch('kiro_gateway.streaming.AwsEventStreamParser', return_value=mock_parser): 200 | with patch('kiro_gateway.streaming.parse_bracket_tool_calls', return_value=[]): 201 | async for chunk in stream_kiro_to_openai( 202 | mock_http_client, mock_response, "test-model", 203 | mock_model_cache, mock_auth_manager 204 | ): 205 | chunks.append(chunk) 206 | 207 | # Проверяем, что не было исключений и chunks собраны 208 | print(f"Получено chunks: {len(chunks)}") 209 | assert len(chunks) > 0 210 | 211 | # Проверяем, что name заменён на пустую строку 212 | for chunk in chunks: 213 | if isinstance(chunk, str) and "tool_calls" in chunk: 214 | if chunk.startswith("data: "): 215 | json_str = chunk[6:].strip() 216 | if json_str != "[DONE]": 217 | data = json.loads(json_str) 218 | if "choices" in data and data["choices"]: 219 | delta = data["choices"][0].get("delta", {}) 220 | if "tool_calls" in delta: 221 | for tc in delta["tool_calls"]: 222 | assert tc["function"]["name"] == "", "None name должен быть заменён на пустую строку" 223 | 224 | @pytest.mark.asyncio 225 | async def test_handles_none_function_arguments(self, mock_http_client, mock_model_cache, mock_auth_manager): 226 | """ 227 | Что он делает: Проверяет обработку None в function.arguments. 228 | Цель: Убедиться, что None заменяется на "{}". 229 | """ 230 | print("Настройка: Tool call с None arguments...") 231 | tool_calls = [ 232 | { 233 | "id": "call_1", 234 | "type": "function", 235 | "function": { 236 | "name": "test_func", 237 | "arguments": None 238 | } 239 | } 240 | ] 241 | 242 | mock_parser = MagicMock() 243 | mock_parser.feed.return_value = [] 244 | mock_parser.get_tool_calls.return_value = tool_calls 245 | 246 | mock_response = AsyncMock() 247 | mock_response.status_code = 200 248 | 249 | async def mock_aiter_bytes(): 250 | yield b'{"content":"test"}' 251 | 252 | mock_response.aiter_bytes = mock_aiter_bytes 253 | mock_response.aclose = AsyncMock() 254 | 255 | print("Действие: Сбор streaming chunks...") 256 | chunks = [] 257 | 258 | with patch('kiro_gateway.streaming.AwsEventStreamParser', return_value=mock_parser): 259 | with patch('kiro_gateway.streaming.parse_bracket_tool_calls', return_value=[]): 260 | async for chunk in stream_kiro_to_openai( 261 | mock_http_client, mock_response, "test-model", 262 | mock_model_cache, mock_auth_manager 263 | ): 264 | chunks.append(chunk) 265 | 266 | print(f"Получено chunks: {len(chunks)}") 267 | assert len(chunks) > 0 268 | 269 | # Проверяем, что arguments заменён на "{}" 270 | for chunk in chunks: 271 | if isinstance(chunk, str) and "tool_calls" in chunk: 272 | if chunk.startswith("data: "): 273 | json_str = chunk[6:].strip() 274 | if json_str != "[DONE]": 275 | data = json.loads(json_str) 276 | if "choices" in data and data["choices"]: 277 | delta = data["choices"][0].get("delta", {}) 278 | if "tool_calls" in delta: 279 | for tc in delta["tool_calls"]: 280 | # None должен быть заменён на "{}" или пустую строку 281 | assert tc["function"]["arguments"] is not None 282 | 283 | @pytest.mark.asyncio 284 | async def test_handles_none_function_object(self, mock_http_client, mock_model_cache, mock_auth_manager): 285 | """ 286 | Что он делает: Проверяет обработку None вместо function объекта. 287 | Цель: Убедиться, что None function обрабатывается без ошибок. 288 | """ 289 | print("Настройка: Tool call с None function...") 290 | tool_calls = [ 291 | { 292 | "id": "call_1", 293 | "type": "function", 294 | "function": None 295 | } 296 | ] 297 | 298 | mock_parser = MagicMock() 299 | mock_parser.feed.return_value = [] 300 | mock_parser.get_tool_calls.return_value = tool_calls 301 | 302 | mock_response = AsyncMock() 303 | mock_response.status_code = 200 304 | 305 | async def mock_aiter_bytes(): 306 | yield b'{"content":"test"}' 307 | 308 | mock_response.aiter_bytes = mock_aiter_bytes 309 | mock_response.aclose = AsyncMock() 310 | 311 | print("Действие: Сбор streaming chunks...") 312 | chunks = [] 313 | 314 | with patch('kiro_gateway.streaming.AwsEventStreamParser', return_value=mock_parser): 315 | with patch('kiro_gateway.streaming.parse_bracket_tool_calls', return_value=[]): 316 | async for chunk in stream_kiro_to_openai( 317 | mock_http_client, mock_response, "test-model", 318 | mock_model_cache, mock_auth_manager 319 | ): 320 | chunks.append(chunk) 321 | 322 | print(f"Получено chunks: {len(chunks)}") 323 | assert len(chunks) > 0 324 | 325 | 326 | class TestCollectStreamResponseToolCalls: 327 | """Тесты для collect_stream_response с tool_calls.""" 328 | 329 | @pytest.mark.asyncio 330 | async def test_collected_tool_calls_have_no_index(self, mock_http_client, mock_model_cache, mock_auth_manager): 331 | """ 332 | Что он делает: Проверяет, что собранные tool_calls не содержат поле index. 333 | Цель: Убедиться, что для non-streaming ответа index удаляется. 334 | """ 335 | print("Настройка: Tool calls...") 336 | tool_calls = [ 337 | { 338 | "id": "call_1", 339 | "type": "function", 340 | "function": {"name": "func1", "arguments": '{"a": 1}'} 341 | } 342 | ] 343 | 344 | mock_parser = MagicMock() 345 | mock_parser.feed.return_value = [] 346 | mock_parser.get_tool_calls.return_value = tool_calls 347 | 348 | mock_response = AsyncMock() 349 | mock_response.status_code = 200 350 | 351 | async def mock_aiter_bytes(): 352 | yield b'{"content":"Hello"}' 353 | 354 | mock_response.aiter_bytes = mock_aiter_bytes 355 | mock_response.aclose = AsyncMock() 356 | 357 | print("Действие: Сбор полного ответа...") 358 | 359 | with patch('kiro_gateway.streaming.AwsEventStreamParser', return_value=mock_parser): 360 | with patch('kiro_gateway.streaming.parse_bracket_tool_calls', return_value=[]): 361 | result = await collect_stream_response( 362 | mock_http_client, mock_response, "test-model", 363 | mock_model_cache, mock_auth_manager 364 | ) 365 | 366 | print(f"Результат: {result}") 367 | 368 | if "choices" in result and result["choices"]: 369 | message = result["choices"][0].get("message", {}) 370 | if "tool_calls" in message: 371 | for tc in message["tool_calls"]: 372 | print(f"Tool call: {tc}") 373 | assert "index" not in tc, "Non-streaming tool_calls не должны содержать index" 374 | 375 | @pytest.mark.asyncio 376 | async def test_collected_tool_calls_have_required_fields(self, mock_http_client, mock_model_cache, mock_auth_manager): 377 | """ 378 | Что он делает: Проверяет, что собранные tool_calls содержат все обязательные поля. 379 | Цель: Убедиться, что id, type, function присутствуют. 380 | """ 381 | print("Настройка: Tool calls...") 382 | tool_calls = [ 383 | { 384 | "id": "call_abc", 385 | "type": "function", 386 | "function": {"name": "get_weather", "arguments": '{"city": "Moscow"}'} 387 | } 388 | ] 389 | 390 | mock_parser = MagicMock() 391 | mock_parser.feed.return_value = [] 392 | mock_parser.get_tool_calls.return_value = tool_calls 393 | 394 | mock_response = AsyncMock() 395 | mock_response.status_code = 200 396 | 397 | async def mock_aiter_bytes(): 398 | yield b'{"content":""}' 399 | 400 | mock_response.aiter_bytes = mock_aiter_bytes 401 | mock_response.aclose = AsyncMock() 402 | 403 | print("Действие: Сбор полного ответа...") 404 | 405 | with patch('kiro_gateway.streaming.AwsEventStreamParser', return_value=mock_parser): 406 | with patch('kiro_gateway.streaming.parse_bracket_tool_calls', return_value=[]): 407 | result = await collect_stream_response( 408 | mock_http_client, mock_response, "test-model", 409 | mock_model_cache, mock_auth_manager 410 | ) 411 | 412 | print(f"Результат: {result}") 413 | 414 | if "choices" in result and result["choices"]: 415 | message = result["choices"][0].get("message", {}) 416 | if "tool_calls" in message: 417 | for tc in message["tool_calls"]: 418 | print(f"Проверяем tool call: {tc}") 419 | assert "id" in tc, "Tool call должен содержать id" 420 | assert "type" in tc, "Tool call должен содержать type" 421 | assert "function" in tc, "Tool call должен содержать function" 422 | assert "name" in tc["function"], "Function должен содержать name" 423 | assert "arguments" in tc["function"], "Function должен содержать arguments" 424 | 425 | @pytest.mark.asyncio 426 | async def test_handles_none_in_collected_tool_calls(self, mock_http_client, mock_model_cache, mock_auth_manager): 427 | """ 428 | Что он делает: Проверяет обработку None значений в собранных tool_calls. 429 | Цель: Убедиться, что None заменяются на дефолтные значения. 430 | """ 431 | print("Настройка: Tool calls с None значениями...") 432 | tool_calls = [ 433 | { 434 | "id": "call_1", 435 | "type": "function", 436 | "function": None 437 | } 438 | ] 439 | 440 | mock_parser = MagicMock() 441 | mock_parser.feed.return_value = [] 442 | mock_parser.get_tool_calls.return_value = tool_calls 443 | 444 | mock_response = AsyncMock() 445 | mock_response.status_code = 200 446 | 447 | async def mock_aiter_bytes(): 448 | yield b'{"content":""}' 449 | 450 | mock_response.aiter_bytes = mock_aiter_bytes 451 | mock_response.aclose = AsyncMock() 452 | 453 | print("Действие: Сбор полного ответа...") 454 | 455 | with patch('kiro_gateway.streaming.AwsEventStreamParser', return_value=mock_parser): 456 | with patch('kiro_gateway.streaming.parse_bracket_tool_calls', return_value=[]): 457 | result = await collect_stream_response( 458 | mock_http_client, mock_response, "test-model", 459 | mock_model_cache, mock_auth_manager 460 | ) 461 | 462 | print(f"Результат: {result}") 463 | 464 | # Проверяем, что не было исключений 465 | assert "choices" in result --------------------------------------------------------------------------------