tamar-model-client 0.1.18__py3-none-any.whl → 0.1.20__py3-none-any.whl

tamar_model_client/core/__init__.py

@@ -0,0 +1,34 @@
+"""
+Core components for Tamar Model Client
+
+This package contains shared components used by both sync and async clients.
+"""
+
+from .utils import (
+    is_effective_value,
+    serialize_value,
+    remove_none_from_dict,
+    generate_request_id,
+    set_request_id,
+    get_request_id
+)
+
+from .logging_setup import (
+    setup_logger,
+    RequestIdFilter,
+    MAX_MESSAGE_LENGTH
+)
+
+__all__ = [
+    # Utils
+    'is_effective_value',
+    'serialize_value',
+    'remove_none_from_dict',
+    'generate_request_id',
+    'set_request_id',
+    'get_request_id',
+    # Logging
+    'setup_logger',
+    'RequestIdFilter',
+    'MAX_MESSAGE_LENGTH',
+]

tamar_model_client/core/base_client.py

@@ -0,0 +1,168 @@
+"""
+Base client class for Tamar Model Client
+
+This module provides the base client class with shared initialization logic
+and configuration management for both sync and async clients.
+"""
+
+import os
+import logging
+from typing import Optional, Dict, Any
+from abc import ABC, abstractmethod
+
+from ..auth import JWTAuthHandler
+from ..error_handler import GrpcErrorHandler, ErrorRecoveryStrategy
+from .logging_setup import MAX_MESSAGE_LENGTH, setup_logger
+
+
+class BaseClient(ABC):
+    """
+    基础客户端抽象类
+    
+    提供同步和异步客户端的共享功能：
+    - 配置管理
+    - 认证设置
+    - 连接选项构建
+    - 错误处理器初始化
+    """
+    
+    def __init__(
+            self,
+            server_address: Optional[str] = None,
+            jwt_secret_key: Optional[str] = None,
+            jwt_token: Optional[str] = None,
+            default_payload: Optional[dict] = None,
+            token_expires_in: int = 3600,
+            max_retries: Optional[int] = None,
+            retry_delay: Optional[float] = None,
+            logger_name: str = None,
+    ):
+        """
+        初始化基础客户端
+        
+        Args:
+            server_address: gRPC 服务器地址，格式为 "host:port"
+            jwt_secret_key: JWT 签名密钥，用于生成认证令牌
+            jwt_token: 预生成的 JWT 令牌（可选）
+            default_payload: JWT 令牌的默认载荷
+            token_expires_in: JWT 令牌过期时间（秒）
+            max_retries: 最大重试次数（默认从环境变量读取）
+            retry_delay: 初始重试延迟（秒，默认从环境变量读取）
+            logger_name: 日志记录器名称
+            
+        Raises:
+            ValueError: 当服务器地址未提供时
+        """
+        # === 服务端地址配置 ===
+        self.server_address = server_address or os.getenv("MODEL_MANAGER_SERVER_ADDRESS")
+        if not self.server_address:
+            raise ValueError("Server address must be provided via argument or environment variable.")
+        
+        # 默认调用超时时间
+        self.default_invoke_timeout = float(os.getenv("MODEL_MANAGER_SERVER_INVOKE_TIMEOUT", 30.0))
+        
+        # === JWT 认证配置 ===
+        self.jwt_secret_key = jwt_secret_key or os.getenv("MODEL_MANAGER_SERVER_JWT_SECRET_KEY")
+        self.jwt_handler = JWTAuthHandler(self.jwt_secret_key) if self.jwt_secret_key else None
+        self.jwt_token = jwt_token  # 用户传入的预生成 Token（可选）
+        self.default_payload = default_payload
+        self.token_expires_in = token_expires_in
+        
+        # === TLS/Authority 配置 ===
+        self.use_tls = os.getenv("MODEL_MANAGER_SERVER_GRPC_USE_TLS", "true").lower() == "true"
+        self.default_authority = os.getenv("MODEL_MANAGER_SERVER_GRPC_DEFAULT_AUTHORITY")
+        
+        # === 重试配置 ===
+        self.max_retries = max_retries if max_retries is not None else int(
+            os.getenv("MODEL_MANAGER_SERVER_GRPC_MAX_RETRIES", 3))
+        self.retry_delay = retry_delay if retry_delay is not None else float(
+            os.getenv("MODEL_MANAGER_SERVER_GRPC_RETRY_DELAY", 1.0))
+        
+        # === 日志配置 ===
+        self.logger = setup_logger(logger_name or __name__)
+        
+        # === 错误处理器 ===
+        self.error_handler = GrpcErrorHandler(self.logger)
+        self.recovery_strategy = ErrorRecoveryStrategy(self)
+        
+        # === 连接状态 ===
+        self._closed = False
+    
+    def build_channel_options(self) -> list:
+        """
+        构建 gRPC 通道选项
+        
+        Returns:
+            list: gRPC 通道配置选项列表
+            
+        包含的配置：
+        - 消息大小限制
+        - Keepalive 设置（30秒ping间隔，10秒超时）
+        - 连接生命周期管理（1小时最大连接时间）
+        - 性能优化选项（带宽探测、内置重试）
+        """
+        options = [
+            # 消息大小限制
+            ('grpc.max_send_message_length', MAX_MESSAGE_LENGTH),
+            ('grpc.max_receive_message_length', MAX_MESSAGE_LENGTH),
+            
+            # Keepalive 核心配置
+            ('grpc.keepalive_time_ms', 30000),  # 30秒发送一次 keepalive ping
+            ('grpc.keepalive_timeout_ms', 10000),  # ping 响应超时时间 10秒
+            ('grpc.keepalive_permit_without_calls', True),  # 空闲时也发送 keepalive
+            ('grpc.http2.max_pings_without_data', 2),  # 无数据时最大 ping 次数
+            
+            # 连接管理增强配置
+            ('grpc.http2.min_time_between_pings_ms', 10000),  # ping 最小间隔 10秒
+            ('grpc.http2.max_connection_idle_ms', 300000),  # 最大空闲时间 5分钟
+            ('grpc.http2.max_connection_age_ms', 3600000),  # 连接最大生存时间 1小时
+            ('grpc.http2.max_connection_age_grace_ms', 5000),  # 优雅关闭时间 5秒
+            
+            # 性能相关配置
+            ('grpc.http2.bdp_probe', 1),  # 启用带宽延迟探测
+            ('grpc.enable_retries', 1),  # 启用内置重试
+        ]
+        
+        if self.default_authority:
+            options.append(("grpc.default_authority", self.default_authority))
+            
+        return options
+    
+    def _build_auth_metadata(self, request_id: str) -> list:
+        """
+        构建认证元数据
+        
+        为每个请求构建包含认证信息和请求ID的gRPC元数据。
+        JWT令牌会在每次请求时重新生成以确保有效性。
+        
+        Args:
+            request_id: 当前请求的唯一标识符
+            
+        Returns:
+            list: gRPC元数据列表，包含请求ID和认证令牌
+        """
+        metadata = [("x-request-id", request_id)]  # 将 request_id 添加到 headers
+        
+        if self.jwt_handler:
+            self.jwt_token = self.jwt_handler.encode_token(
+                self.default_payload, 
+                expires_in=self.token_expires_in
+            )
+            metadata.append(("authorization", f"Bearer {self.jwt_token}"))
+            
+        return metadata
+    
+    @abstractmethod
+    def close(self):
+        """关闭客户端连接（由子类实现）"""
+        pass
+    
+    @abstractmethod
+    def __enter__(self):
+        """进入上下文管理器（由子类实现）"""
+        pass
+    
+    @abstractmethod
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """退出上下文管理器（由子类实现）"""
+        pass

tamar_model_client/core/logging_setup.py

@@ -0,0 +1,84 @@
+"""
+Logging configuration for Tamar Model Client
+
+This module provides centralized logging setup for both sync and async clients.
+It includes request ID tracking, JSON formatting, and consistent log configuration.
+"""
+
+import logging
+from typing import Optional
+
+from ..json_formatter import JSONFormatter
+from .utils import get_request_id
+
+# gRPC 消息长度限制（32位系统兼容）
+MAX_MESSAGE_LENGTH = 2 ** 31 - 1
+
+
+class RequestIdFilter(logging.Filter):
+    """
+    自定义日志过滤器，向日志记录中添加 request_id
+    
+    这个过滤器从 ContextVar 中获取当前请求的 ID，
+    并将其添加到日志记录中，便于追踪和调试。
+    """
+
+    def filter(self, record):
+        """
+        过滤日志记录，添加 request_id 字段
+        
+        Args:
+            record: 日志记录对象
+            
+        Returns:
+            bool: 总是返回 True，表示记录应被处理
+        """
+        # 从 ContextVar 中获取当前的 request_id
+        record.request_id = get_request_id()
+        return True
+
+
+def setup_logger(logger_name: str, level: int = logging.INFO) -> logging.Logger:
+    """
+    设置并配置logger
+    
+    为指定的logger配置处理器、格式化器和过滤器。
+    如果logger已经有处理器，则不会重复配置。
+    
+    Args:
+        logger_name: logger的名称
+        level: 日志级别，默认为 INFO
+        
+    Returns:
+        logging.Logger: 配置好的logger实例
+        
+    特性：
+    - 使用 JSON 格式化器提供结构化日志输出
+    - 添加请求ID过滤器用于请求追踪
+    - 避免重复配置
+    """
+    logger = logging.getLogger(logger_name)
+    
+    # 仅在没有处理器时配置，避免重复配置
+    if not logger.hasHandlers():
+        # 创建控制台日志处理器
+        console_handler = logging.StreamHandler()
+        
+        # 使用自定义的 JSON 格式化器，提供结构化日志输出
+        formatter = JSONFormatter()
+        console_handler.setFormatter(formatter)
+        
+        # 为logger添加处理器
+        logger.addHandler(console_handler)
+        
+        # 设置日志级别
+        logger.setLevel(level)
+        
+        # 添加自定义的请求ID过滤器，用于请求追踪
+        logger.addFilter(RequestIdFilter())
+        
+        # 关键：设置 propagate = False，防止日志传播到父logger
+        # 这样可以避免测试脚本的日志格式影响客户端日志
+        logger.propagate = False
+    
+    return logger

tamar_model_client/core/request_builder.py

@@ -0,0 +1,221 @@
+"""
+Request building logic for Tamar Model Client
+
+This module handles the construction of gRPC request objects from
+model request objects, including provider-specific field validation.
+"""
+
+import json
+from typing import Dict, Any, Set
+
+from ..enums import ProviderType, InvokeType
+from ..generated import model_service_pb2
+from ..schemas.inputs import (
+    ModelRequest,
+    BatchModelRequest,
+    BatchModelRequestItem,
+    UserContext,
+    GoogleGenAiInput,
+    GoogleVertexAIImagesInput,
+    OpenAIResponsesInput,
+    OpenAIChatCompletionsInput,
+    OpenAIImagesInput,
+    OpenAIImagesEditInput
+)
+from .utils import is_effective_value, serialize_value, remove_none_from_dict
+
+
+class RequestBuilder:
+    """
+    请求构建器
+    
+    负责将高级的 ModelRequest 对象转换为 gRPC 协议所需的请求对象，
+    包括参数验证、序列化和提供商特定的字段处理。
+    """
+    
+    @staticmethod
+    def get_allowed_fields(provider: ProviderType, invoke_type: InvokeType) -> Set[str]:
+        """
+        获取特定提供商和调用类型组合所允许的字段
+        
+        Args:
+            provider: 提供商类型
+            invoke_type: 调用类型
+            
+        Returns:
+            Set[str]: 允许的字段名集合
+            
+        Raises:
+            ValueError: 当提供商和调用类型组合不受支持时
+        """
+        match (provider, invoke_type):
+            case (ProviderType.GOOGLE, InvokeType.GENERATION):
+                return set(GoogleGenAiInput.model_fields.keys())
+            case (ProviderType.GOOGLE, InvokeType.IMAGE_GENERATION):
+                return set(GoogleVertexAIImagesInput.model_fields.keys())
+            case ((ProviderType.OPENAI | ProviderType.AZURE), InvokeType.RESPONSES | InvokeType.GENERATION):
+                return set(OpenAIResponsesInput.model_fields.keys())
+            case ((ProviderType.OPENAI | ProviderType.AZURE), InvokeType.CHAT_COMPLETIONS):
+                return set(OpenAIChatCompletionsInput.model_fields.keys())
+            case ((ProviderType.OPENAI | ProviderType.AZURE), InvokeType.IMAGE_GENERATION):
+                return set(OpenAIImagesInput.model_fields.keys())
+            case ((ProviderType.OPENAI | ProviderType.AZURE), InvokeType.IMAGE_EDIT_GENERATION):
+                return set(OpenAIImagesEditInput.model_fields.keys())
+            case _:
+                raise ValueError(
+                    f"Unsupported provider/invoke_type combination: {provider} + {invoke_type}"
+                )
+    
+    @staticmethod
+    def build_grpc_extra_fields(model_request: ModelRequest) -> Dict[str, Any]:
+        """
+        构建 gRPC 请求的额外字段
+        
+        根据提供商和调用类型，过滤并序列化请求中的参数。
+        
+        Args:
+            model_request: 模型请求对象
+            
+        Returns:
+            Dict[str, Any]: 序列化后的额外字段字典
+            
+        Raises:
+            ValueError: 当构建请求失败时
+        """
+        try:
+            # 获取允许的字段集合
+            allowed_fields = RequestBuilder.get_allowed_fields(
+                model_request.provider,
+                model_request.invoke_type
+            )
+            
+            # 将 ModelRequest 转换为字典，只包含已设置的字段
+            model_request_dict = model_request.model_dump(exclude_unset=True)
+            
+            # 构建 gRPC 请求参数
+            grpc_request_kwargs = {}
+            for field in allowed_fields:
+                if field in model_request_dict:
+                    value = model_request_dict[field]
+                    
+                    # 跳过无效的值
+                    if not is_effective_value(value):
+                        continue
+                    
+                    # 序列化不支持的类型
+                    grpc_request_kwargs[field] = serialize_value(value)
+            
+            # 清理序列化后的参数中的 None 值
+            grpc_request_kwargs = remove_none_from_dict(grpc_request_kwargs)
+            
+            return grpc_request_kwargs
+            
+        except Exception as e:
+            raise ValueError(f"构建请求失败: {str(e)}") from e
+    
+    @staticmethod
+    def build_single_request(model_request: ModelRequest) -> model_service_pb2.ModelRequestItem:
+        """
+        构建单个模型请求的 gRPC 对象
+        
+        Args:
+            model_request: 模型请求对象
+            
+        Returns:
+            model_service_pb2.ModelRequestItem: gRPC 请求对象
+            
+        Raises:
+            ValueError: 当构建请求失败时
+        """
+        # 构建额外字段
+        extra_fields = RequestBuilder.build_grpc_extra_fields(model_request)
+        
+        # 创建 gRPC 请求对象
+        return model_service_pb2.ModelRequestItem(
+            provider=model_request.provider.value,
+            channel=model_request.channel.value if model_request.channel else "",
+            invoke_type=model_request.invoke_type.value,
+            stream=model_request.stream or False,
+            org_id=model_request.user_context.org_id or "",
+            user_id=model_request.user_context.user_id or "",
+            client_type=model_request.user_context.client_type or "",
+            extra=extra_fields
+        )
+    
+    @staticmethod
+    def build_batch_request_item(
+        batch_item: "BatchModelRequestItem", 
+        user_context: "UserContext"
+    ) -> model_service_pb2.ModelRequestItem:
+        """
+        构建批量请求中的单个项目
+        
+        Args:
+            batch_item: 批量请求项
+            user_context: 用户上下文（来自父BatchModelRequest）
+            
+        Returns:
+            model_service_pb2.ModelRequestItem: gRPC 请求对象
+        """
+        # 构建额外字段
+        extra_fields = RequestBuilder.build_grpc_extra_fields(batch_item)
+        
+        # 添加 custom_id 如果存在
+        if hasattr(batch_item, 'custom_id') and batch_item.custom_id:
+            request_item = model_service_pb2.ModelRequestItem(
+                provider=batch_item.provider.value,
+                channel=batch_item.channel.value if batch_item.channel else "",
+                invoke_type=batch_item.invoke_type.value,
+                stream=batch_item.stream or False,
+                org_id=user_context.org_id or "",
+                user_id=user_context.user_id or "",
+                client_type=user_context.client_type or "",
+                custom_id=batch_item.custom_id,
+                extra=extra_fields
+            )
+        else:
+            request_item = model_service_pb2.ModelRequestItem(
+                provider=batch_item.provider.value,
+                channel=batch_item.channel.value if batch_item.channel else "",
+                invoke_type=batch_item.invoke_type.value,
+                stream=batch_item.stream or False,
+                org_id=user_context.org_id or "",
+                user_id=user_context.user_id or "",
+                client_type=user_context.client_type or "",
+                extra=extra_fields
+            )
+            
+        # 添加 priority 如果存在
+        if hasattr(batch_item, 'priority') and batch_item.priority is not None:
+            request_item.priority = batch_item.priority
+            
+        return request_item
+    
+    @staticmethod
+    def build_batch_request(batch_request: BatchModelRequest) -> model_service_pb2.ModelRequest:
+        """
+        构建批量请求的 gRPC 对象
+        
+        Args:
+            batch_request: 批量请求对象
+            
+        Returns:
+            model_service_pb2.ModelRequest: gRPC 批量请求对象
+            
+        Raises:
+            ValueError: 当构建请求失败时
+        """
+        items = []
+        
+        for batch_item in batch_request.items:
+            # 为每个请求项构建 gRPC 对象，传入 user_context
+            request_item = RequestBuilder.build_batch_request_item(
+                batch_item, 
+                batch_request.user_context
+            )
+            items.append(request_item)
+        
+        # 创建批量请求对象
+        return model_service_pb2.ModelRequest(
+            items=items
+        )

tamar_model_client/core/response_handler.py

@@ -0,0 +1,136 @@
+"""
+Response handling logic for Tamar Model Client
+
+This module provides utilities for processing gRPC responses and
+converting them to client response objects.
+"""
+
+import json
+from typing import Optional, Dict, Any
+
+from ..schemas import ModelResponse, BatchModelResponse
+
+
+class ResponseHandler:
+    """
+    响应处理器
+    
+    负责将 gRPC 响应转换为客户端响应对象，
+    包括 JSON 解析、错误处理和数据结构转换。
+    """
+    
+    @staticmethod
+    def build_model_response(grpc_response) -> ModelResponse:
+        """
+        从 gRPC 响应构建 ModelResponse 对象
+        
+        Args:
+            grpc_response: gRPC 服务返回的响应对象
+            
+        Returns:
+            ModelResponse: 客户端响应对象
+        """
+        return ModelResponse(
+            content=grpc_response.content,
+            usage=ResponseHandler._parse_json_field(grpc_response.usage),
+            error=grpc_response.error or None,
+            raw_response=ResponseHandler._parse_json_field(grpc_response.raw_response),
+            request_id=grpc_response.request_id if grpc_response.request_id else None,
+        )
+    
+    @staticmethod
+    def build_batch_response(grpc_response) -> BatchModelResponse:
+        """
+        从 gRPC 批量响应构建 BatchModelResponse 对象
+        
+        Args:
+            grpc_response: gRPC 服务返回的批量响应对象
+            
+        Returns:
+            BatchModelResponse: 客户端批量响应对象
+        """
+        responses = []
+        for response_item in grpc_response.items:
+            model_response = ResponseHandler.build_model_response(response_item)
+            responses.append(model_response)
+        
+        return BatchModelResponse(
+            responses=responses,
+            request_id=grpc_response.request_id if grpc_response.request_id else None
+        )
+    
+    @staticmethod
+    def _parse_json_field(json_str: Optional[str]) -> Optional[Dict[str, Any]]:
+        """
+        安全地解析 JSON 字符串
+        
+        Args:
+            json_str: 待解析的 JSON 字符串
+            
+        Returns:
+            Optional[Dict[str, Any]]: 解析后的字典，或 None（如果输入为空）
+        """
+        if not json_str:
+            return None
+        
+        try:
+            return json.loads(json_str)
+        except json.JSONDecodeError:
+            # 如果解析失败，返回原始字符串作为错误信息
+            return {"error": "JSON parse error", "raw": json_str}
+    
+    @staticmethod
+    def build_log_data(
+        model_request,
+        response: Optional[ModelResponse] = None,
+        duration: Optional[float] = None,
+        error: Optional[Exception] = None,
+        stream_stats: Optional[Dict[str, Any]] = None
+    ) -> Dict[str, Any]:
+        """
+        构建日志数据
+        
+        为请求和响应日志构建结构化的数据字典。
+        
+        Args:
+            model_request: 原始请求对象
+            response: 响应对象（可选）
+            duration: 请求持续时间（秒）
+            error: 错误对象（可选）
+            stream_stats: 流式响应统计信息（可选）
+            
+        Returns:
+            Dict[str, Any]: 日志数据字典
+        """
+        data = {
+            "provider": model_request.provider.value,
+            "invoke_type": model_request.invoke_type.value,
+            "model": getattr(model_request, 'model', None),
+            "stream": getattr(model_request, 'stream', False),
+        }
+        
+        # 添加用户上下文信息（如果有）
+        if hasattr(model_request, 'user_context'):
+            data.update({
+                "org_id": model_request.user_context.org_id,
+                "user_id": model_request.user_context.user_id,
+                "client_type": model_request.user_context.client_type
+            })
+        
+        # 添加响应信息
+        if response:
+            if hasattr(response, 'content') and response.content:
+                data["content_length"] = len(response.content)
+            if hasattr(response, 'usage'):
+                data["usage"] = response.usage
+        
+        # 添加流式响应统计
+        if stream_stats:
+            data.update(stream_stats)
+        
+        # 添加错误信息
+        if error:
+            data["error_type"] = type(error).__name__
+            data["error_message"] = str(error)
+        
+        return data