tamar-model-client 0.1.16__tar.gz → 0.1.18__tar.gz

{tamar_model_client-0.1.16 → tamar_model_client-0.1.18}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tamar-model-client
-Version: 0.1.16
+Version: 0.1.18
 Summary: A Python SDK for interacting with the Model Manager gRPC service
 Home-page: http://gitlab.tamaredge.top/project-tap/AgentOS/model-manager-client
 Author: Oscar Ou

{tamar_model_client-0.1.16 → tamar_model_client-0.1.18}/setup.py

@@ -2,7 +2,7 @@ from setuptools import setup, find_packages
 
 setup(
     name="tamar-model-client",
-    version="0.1.16",
+    version="0.1.18",
     description="A Python SDK for interacting with the Model Manager gRPC service",
     author="Oscar Ou",
     author_email="oscar.ou@tamaredge.ai",

{tamar_model_client-0.1.16 → tamar_model_client-0.1.18}/tamar_model_client/__init__.py

@@ -1,6 +1,7 @@
 from .sync_client import TamarModelClient
 from .async_client import AsyncTamarModelClient
 from .exceptions import ModelManagerClientError, ConnectionError, ValidationError
+from .json_formatter import JSONFormatter
 
 __all__ = [
     "TamarModelClient",
@@ -8,4 +9,5 @@ __all__ = [
     "ModelManagerClientError",
     "ConnectionError",
     "ValidationError",
+    "JSONFormatter",
 ]

{tamar_model_client-0.1.16 → tamar_model_client-0.1.18}/tamar_model_client/async_client.py

@@ -4,6 +4,7 @@ import base64
 import json
 import logging
 import os
+import time
 import uuid
 from contextvars import ContextVar
 
@@ -19,7 +20,8 @@ from .exceptions import ConnectionError
 from .schemas import ModelRequest, ModelResponse, BatchModelRequest, BatchModelResponse
 from .generated import model_service_pb2, model_service_pb2_grpc
 from .schemas.inputs import GoogleGenAiInput, OpenAIResponsesInput, OpenAIChatCompletionsInput, \
-    GoogleVertexAIImagesInput, OpenAIImagesInput
+    GoogleVertexAIImagesInput, OpenAIImagesInput, OpenAIImagesEditInput
+from .json_formatter import JSONFormatter
 
 logger = logging.getLogger(__name__)
 
@@ -40,8 +42,8 @@ if not logger.hasHandlers():
     # 创建日志处理器，输出到控制台
     console_handler = logging.StreamHandler()
 
-    # 设置日志格式
-    formatter = logging.Formatter('%(asctime)s [%(levelname)s] [%(request_id)s] %(message)s')
+    # 使用 JSON 格式化器
+    formatter = JSONFormatter()
     console_handler.setFormatter(formatter)
 
     # 为当前记录器添加处理器
@@ -181,26 +183,31 @@ class AsyncTamarModelClient:
                 # 对于取消的情况进行指数退避重试
                 if isinstance(e, grpc.aio.AioRpcError) and e.code() == grpc.StatusCode.CANCELLED:
                     retry_count += 1
-                    logger.warning(f"❌ RPC cancelled, retrying {retry_count}/{self.max_retries}...")
+                    logger.info(f"❌ RPC cancelled, retrying {retry_count}/{self.max_retries}...",
+                                 extra={"log_type": "info", "data": {"retry_count": retry_count, "max_retries": self.max_retries, "error_code": "CANCELLED"}})
                     if retry_count < self.max_retries:
                         delay = self.retry_delay * (2 ** (retry_count - 1))
                         await asyncio.sleep(delay)
                     else:
-                        logger.error("❌ Max retry reached for CANCELLED")
+                        logger.error("❌ Max retry reached for CANCELLED",
+                                   extra={"log_type": "info", "data": {"error_code": "CANCELLED", "max_retries_reached": True}})
                         raise
                 # 针对其他 RPC 错误类型，如暂时的连接问题、服务器超时等
                 elif isinstance(e, grpc.RpcError) and e.code() in {grpc.StatusCode.UNAVAILABLE,
                                                                    grpc.StatusCode.DEADLINE_EXCEEDED}:
                     retry_count += 1
-                    logger.warning(f"❌ gRPC error {e.code()}, retrying {retry_count}/{self.max_retries}...")
+                    logger.info(f"❌ gRPC error {e.code()}, retrying {retry_count}/{self.max_retries}...",
+                                 extra={"log_type": "info", "data": {"retry_count": retry_count, "max_retries": self.max_retries, "error_code": str(e.code())}})
                     if retry_count < self.max_retries:
                         delay = self.retry_delay * (2 ** (retry_count - 1))
                         await asyncio.sleep(delay)
                     else:
-                        logger.error(f"❌ Max retry reached for {e.code()}")
+                        logger.error(f"❌ Max retry reached for {e.code()}",
+                                   extra={"log_type": "info", "data": {"error_code": str(e.code()), "max_retries_reached": True}})
                         raise
                 else:
-                    logger.error(f"❌ Non-retryable gRPC error: {e}", exc_info=True)
+                    logger.error(f"❌ Non-retryable gRPC error: {e}", exc_info=True,
+                               extra={"log_type": "info", "data": {"error_code": str(e.code()) if hasattr(e, 'code') else None, "retryable": False}})
                     raise
 
     async def _retry_request_stream(self, func, *args, **kwargs):
@@ -212,26 +219,31 @@ class AsyncTamarModelClient:
                 # 对于取消的情况进行指数退避重试
                 if isinstance(e, grpc.aio.AioRpcError) and e.code() == grpc.StatusCode.CANCELLED:
                     retry_count += 1
-                    logger.warning(f"❌ RPC cancelled, retrying {retry_count}/{self.max_retries}...")
+                    logger.info(f"❌ RPC cancelled, retrying {retry_count}/{self.max_retries}...",
+                                 extra={"log_type": "info", "data": {"retry_count": retry_count, "max_retries": self.max_retries, "error_code": "CANCELLED"}})
                     if retry_count < self.max_retries:
                         delay = self.retry_delay * (2 ** (retry_count - 1))
                         await asyncio.sleep(delay)
                     else:
-                        logger.error("❌ Max retry reached for CANCELLED")
+                        logger.error("❌ Max retry reached for CANCELLED",
+                                   extra={"log_type": "info", "data": {"error_code": "CANCELLED", "max_retries_reached": True}})
                         raise
                 # 针对其他 RPC 错误类型，如暂时的连接问题、服务器超时等
                 elif isinstance(e, grpc.RpcError) and e.code() in {grpc.StatusCode.UNAVAILABLE,
                                                                    grpc.StatusCode.DEADLINE_EXCEEDED}:
                     retry_count += 1
-                    logger.warning(f"❌ gRPC error {e.code()}, retrying {retry_count}/{self.max_retries}...")
+                    logger.info(f"❌ gRPC error {e.code()}, retrying {retry_count}/{self.max_retries}...",
+                                 extra={"log_type": "info", "data": {"retry_count": retry_count, "max_retries": self.max_retries, "error_code": str(e.code())}})
                     if retry_count < self.max_retries:
                         delay = self.retry_delay * (2 ** (retry_count - 1))
                         await asyncio.sleep(delay)
                     else:
-                        logger.error(f"❌ Max retry reached for {e.code()}")
+                        logger.error(f"❌ Max retry reached for {e.code()}",
+                                   extra={"log_type": "info", "data": {"error_code": str(e.code()), "max_retries_reached": True}})
                         raise
                 else:
-                    logger.error(f"❌ Non-retryable gRPC error: {e}", exc_info=True)
+                    logger.error(f"❌ Non-retryable gRPC error: {e}", exc_info=True,
+                               extra={"log_type": "info", "data": {"error_code": str(e.code()) if hasattr(e, 'code') else None, "retryable": False}})
                     raise
 
     def _build_auth_metadata(self, request_id: str) -> list:
@@ -266,32 +278,40 @@ class AsyncTamarModelClient:
                         credentials,
                         options=options
                     )
-                    logger.info("🔐 Using secure gRPC channel (TLS enabled)")
+                    logger.info("🔐 Using secure gRPC channel (TLS enabled)",
+                              extra={"log_type": "info", "data": {"tls_enabled": True, "server_address": self.server_address}})
                 else:
                     self.channel = grpc.aio.insecure_channel(
                         self.server_address,
                         options=options
                     )
-                    logger.info("🔓 Using insecure gRPC channel (TLS disabled)")
+                    logger.info("🔓 Using insecure gRPC channel (TLS disabled)",
+                              extra={"log_type": "info", "data": {"tls_enabled": False, "server_address": self.server_address}})
                 await self.channel.channel_ready()
                 self.stub = model_service_pb2_grpc.ModelServiceStub(self.channel)
-                logger.info(f"✅ gRPC channel initialized to {self.server_address}")
+                logger.info(f"✅ gRPC channel initialized to {self.server_address}",
+                          extra={"log_type": "info", "data": {"status": "success", "server_address": self.server_address}})
                 return
             except grpc.FutureTimeoutError as e:
-                logger.error(f"❌ gRPC channel initialization timed out: {str(e)}", exc_info=True)
+                logger.error(f"❌ gRPC channel initialization timed out: {str(e)}", exc_info=True,
+                           extra={"log_type": "info", "data": {"error_type": "timeout", "server_address": self.server_address}})
             except grpc.RpcError as e:
-                logger.error(f"❌ gRPC channel initialization failed: {str(e)}", exc_info=True)
+                logger.error(f"❌ gRPC channel initialization failed: {str(e)}", exc_info=True,
+                           extra={"log_type": "info", "data": {"error_type": "rpc_error", "server_address": self.server_address}})
             except Exception as e:
-                logger.error(f"❌ Unexpected error during channel initialization: {str(e)}", exc_info=True)
+                logger.error(f"❌ Unexpected error during channel initialization: {str(e)}", exc_info=True,
+                           extra={"log_type": "info", "data": {"error_type": "unexpected", "server_address": self.server_address}})
 
             retry_count += 1
             if retry_count > self.max_retries:
-                logger.error(f"❌ Failed to initialize gRPC channel after {self.max_retries} retries.", exc_info=True)
+                logger.error(f"❌ Failed to initialize gRPC channel after {self.max_retries} retries.", exc_info=True,
+                           extra={"log_type": "info", "data": {"max_retries_reached": True, "server_address": self.server_address}})
                 raise ConnectionError(f"❌ Failed to initialize gRPC channel after {self.max_retries} retries.")
 
             # 指数退避：延迟时间 = retry_delay * (2 ^ (retry_count - 1))
             delay = self.retry_delay * (2 ** (retry_count - 1))
-            logger.info(f"🚀 Retrying connection (attempt {retry_count}/{self.max_retries}) after {delay:.2f}s delay...")
+            logger.info(f"🚀 Retrying connection (attempt {retry_count}/{self.max_retries}) after {delay:.2f}s delay...",
+                      extra={"log_type": "info", "data": {"retry_count": retry_count, "max_retries": self.max_retries, "delay": delay}})
             await asyncio.sleep(delay)
 
     async def _stream(self, request, metadata, invoke_timeout) -> AsyncIterator[ModelResponse]:
@@ -303,6 +323,66 @@ class AsyncTamarModelClient:
                 raw_response=json.loads(response.raw_response) if response.raw_response else None,
                 request_id=response.request_id if response.request_id else None,
             )
+    
+    async def _stream_with_logging(self, request, metadata, invoke_timeout, start_time, model_request) -> AsyncIterator[ModelResponse]:
+        """流式响应的包装器，用于记录完整的响应日志"""
+        total_content = ""
+        final_usage = None
+        error_occurred = None
+        chunk_count = 0
+        
+        try:
+            async for response in self._stream(request, metadata, invoke_timeout):
+                chunk_count += 1
+                if response.content:
+                    total_content += response.content
+                if response.usage:
+                    final_usage = response.usage
+                if response.error:
+                    error_occurred = response.error
+                yield response
+            
+            # 流式响应完成，记录成功日志
+            duration = time.time() - start_time
+            logger.info(
+                f"✅ Stream completed successfully | chunks: {chunk_count}",
+                extra={
+                    "log_type": "response",
+                    "uri": f"/invoke/{model_request.provider.value}/{model_request.invoke_type.value}",
+                    "duration": duration,
+                    "data": {
+                        "provider": model_request.provider.value,
+                        "invoke_type": model_request.invoke_type.value,
+                        "model": model_request.model,
+                        "stream": True,
+                        "chunks_count": chunk_count,
+                        "total_length": len(total_content),
+                        "usage": final_usage
+                    }
+                }
+            )
+        except Exception as e:
+            # 流式响应出错，记录错误日志
+            duration = time.time() - start_time
+            logger.error(
+                f"❌ Stream failed after {chunk_count} chunks: {str(e)}",
+                exc_info=True,
+                extra={
+                    "log_type": "response",
+                    "uri": f"/invoke/{model_request.provider.value}/{model_request.invoke_type.value}",
+                    "duration": duration,
+                    "data": {
+                        "provider": model_request.provider.value,
+                        "invoke_type": model_request.invoke_type.value,
+                        "model": model_request.model,
+                        "stream": True,
+                        "chunks_count": chunk_count,
+                        "error_type": type(e).__name__,
+                        "partial_content_length": len(total_content)
+                    }
+                }
+            )
+            raise
 
     async def _invoke_request(self, request, metadata, invoke_timeout):
         async for response in self.stub.Invoke(request, metadata=metadata, timeout=invoke_timeout):
@@ -345,8 +425,22 @@ class AsyncTamarModelClient:
         metadata = self._build_auth_metadata(request_id)  # 将 request_id 加入到请求头
 
         # 记录开始日志
+        start_time = time.time()
         logger.info(
-            f"🔵 Request Start | request_id: {request_id} | provider: {model_request.provider} | invoke_type: {model_request.invoke_type} | model_request: {model_request}")
+            f"🔵 Request Start | request_id: {request_id} | provider: {model_request.provider} | invoke_type: {model_request.invoke_type}",
+            extra={
+                "log_type": "request",
+                "uri": f"/invoke/{model_request.provider.value}/{model_request.invoke_type.value}",
+                "data": {
+                    "provider": model_request.provider.value,
+                    "invoke_type": model_request.invoke_type.value,
+                    "model": model_request.model,
+                    "stream": model_request.stream,
+                    "org_id": model_request.user_context.org_id,
+                    "user_id": model_request.user_context.user_id,
+                    "client_type": model_request.user_context.client_type
+                }
+            })
 
         # 动态根据 provider/invoke_type 决定使用哪个 input 字段
         try:
@@ -363,6 +457,8 @@ class AsyncTamarModelClient:
                     allowed_fields = OpenAIChatCompletionsInput.model_fields.keys()
                 case ((ProviderType.OPENAI | ProviderType.AZURE), InvokeType.IMAGE_GENERATION):
                     allowed_fields = OpenAIImagesInput.model_fields.keys()
+                case ((ProviderType.OPENAI | ProviderType.AZURE), InvokeType.IMAGE_EDIT_GENERATION):
+                    allowed_fields = OpenAIImagesEditInput.model_fields.keys()
                 case _:
                     raise ValueError(
                         f"Unsupported provider/invoke_type combination: {model_request.provider} + {model_request.invoke_type}")
@@ -402,16 +498,63 @@ class AsyncTamarModelClient:
         try:
             invoke_timeout = timeout or self.default_invoke_timeout
             if model_request.stream:
-                return await self._retry_request_stream(self._stream, request, metadata, invoke_timeout)
+                # 对于流式响应，使用带日志记录的包装器
+                stream_generator = await self._retry_request_stream(self._stream, request, metadata, invoke_timeout)
+                return self._stream_with_logging(request, metadata, invoke_timeout, start_time, model_request)
             else:
-                return await self._retry_request(self._invoke_request, request, metadata, invoke_timeout)
+                result = await self._retry_request(self._invoke_request, request, metadata, invoke_timeout)
+                
+                # 记录非流式响应的成功日志
+                duration = time.time() - start_time
+                logger.info(
+                    f"✅ Request completed successfully",
+                    extra={
+                        "log_type": "response",
+                        "uri": f"/invoke/{model_request.provider.value}/{model_request.invoke_type.value}",
+                        "duration": duration,
+                        "data": {
+                            "provider": model_request.provider.value,
+                            "invoke_type": model_request.invoke_type.value,
+                            "model": model_request.model,
+                            "stream": False,
+                            "content_length": len(result.content) if result.content else 0,
+                            "usage": result.usage
+                        }
+                    }
+                )
+                return result
         except grpc.RpcError as e:
+            duration = time.time() - start_time
             error_message = f"❌ Invoke gRPC failed: {str(e)}"
-            logger.error(error_message, exc_info=True)
+            logger.error(error_message, exc_info=True,
+                       extra={
+                           "log_type": "response",
+                           "uri": f"/invoke/{model_request.provider.value}/{model_request.invoke_type.value}",
+                           "duration": duration,
+                           "data": {
+                               "error_type": "grpc_error",
+                               "error_code": str(e.code()) if hasattr(e, 'code') else None,
+                               "provider": model_request.provider.value,
+                               "invoke_type": model_request.invoke_type.value,
+                               "model": model_request.model
+                           }
+                       })
             raise e
         except Exception as e:
+            duration = time.time() - start_time
             error_message = f"❌ Invoke other error: {str(e)}"
-            logger.error(error_message, exc_info=True)
+            logger.error(error_message, exc_info=True,
+                       extra={
+                           "log_type": "response",
+                           "uri": f"/invoke/{model_request.provider.value}/{model_request.invoke_type.value}",
+                           "duration": duration,
+                           "data": {
+                               "error_type": "other_error",
+                               "provider": model_request.provider.value,
+                               "invoke_type": model_request.invoke_type.value,
+                               "model": model_request.model
+                           }
+                       })
             raise e
 
     async def invoke_batch(self, batch_request_model: BatchModelRequest, timeout: Optional[float] = None,
@@ -442,8 +585,19 @@ class AsyncTamarModelClient:
         metadata = self._build_auth_metadata(request_id)  # 将 request_id 加入到请求头
 
         # 记录开始日志
+        start_time = time.time()
         logger.info(
-            f"🔵 Batch Request Start | request_id: {request_id} | batch_size: {len(batch_request_model.items)} | batch_request_model: {batch_request_model}")
+            f"🔵 Batch Request Start | request_id: {request_id} | batch_size: {len(batch_request_model.items)}",
+            extra={
+                "log_type": "request",
+                "uri": "/batch_invoke",
+                "data": {
+                    "batch_size": len(batch_request_model.items),
+                    "org_id": batch_request_model.user_context.org_id,
+                    "user_id": batch_request_model.user_context.user_id,
+                    "client_type": batch_request_model.user_context.client_type
+                }
+            })
 
         # 构造批量请求
         items = []
@@ -461,6 +615,8 @@ class AsyncTamarModelClient:
                         allowed_fields = OpenAIChatCompletionsInput.model_fields.keys()
                     case ((ProviderType.OPENAI | ProviderType.AZURE), InvokeType.IMAGE_GENERATION):
                         allowed_fields = OpenAIImagesInput.model_fields.keys()
+                    case ((ProviderType.OPENAI | ProviderType.AZURE), InvokeType.IMAGE_EDIT_GENERATION):
+                        allowed_fields = OpenAIImagesEditInput.model_fields.keys()
                     case _:
                         raise ValueError(
                             f"Unsupported provider/invoke_type combination: {model_request_item.provider} + {model_request_item.invoke_type}")
@@ -516,17 +672,54 @@ class AsyncTamarModelClient:
                     error=res_item.error or None,
                     custom_id=res_item.custom_id if res_item.custom_id else None
                 ))
-            return BatchModelResponse(
+            batch_response = BatchModelResponse(
                 request_id=response.request_id if response.request_id else None,
                 responses=result
             )
+            
+            # 记录成功日志
+            duration = time.time() - start_time
+            logger.info(
+                f"✅ Batch request completed successfully",
+                extra={
+                    "log_type": "response",
+                    "uri": "/batch_invoke",
+                    "duration": duration,
+                    "data": {
+                        "batch_size": len(batch_request_model.items),
+                        "responses_count": len(result)
+                    }
+                }
+            )
+            return batch_response
         except grpc.RpcError as e:
+            duration = time.time() - start_time
             error_message = f"❌ BatchInvoke gRPC failed: {str(e)}"
-            logger.error(error_message, exc_info=True)
+            logger.error(error_message, exc_info=True,
+                       extra={
+                           "log_type": "response",
+                           "uri": "/batch_invoke",
+                           "duration": duration,
+                           "data": {
+                               "error_type": "grpc_error",
+                               "error_code": str(e.code()) if hasattr(e, 'code') else None,
+                               "batch_size": len(batch_request_model.items)
+                           }
+                       })
             raise e
         except Exception as e:
+            duration = time.time() - start_time
             error_message = f"❌ BatchInvoke other error: {str(e)}"
-            logger.error(error_message, exc_info=True)
+            logger.error(error_message, exc_info=True,
+                       extra={
+                           "log_type": "response",
+                           "uri": "/batch_invoke",
+                           "duration": duration,
+                           "data": {
+                               "error_type": "other_error",
+                               "batch_size": len(batch_request_model.items)
+                           }
+                       })
             raise e
 
     async def close(self):
@@ -534,7 +727,8 @@ class AsyncTamarModelClient:
         if self.channel and not self._closed:
             await self.channel.close()
             self._closed = True
-            logger.info("✅ gRPC channel closed")
+            logger.info("✅ gRPC channel closed",
+                      extra={"log_type": "info", "data": {"status": "success"}})
 
     def _safe_sync_close(self):
         """进程退出时自动关闭 channel（事件循环处理兼容）"""
@@ -546,7 +740,8 @@ class AsyncTamarModelClient:
                 else:
                     loop.run_until_complete(self.close())
             except Exception as e:
-                logger.warning(f"❌ gRPC channel close failed at exit: {e}")
+                logger.info(f"❌ gRPC channel close failed at exit: {e}",
+                             extra={"log_type": "info", "data": {"status": "failed", "error": str(e)}})
 
     async def __aenter__(self):
         """支持 async with 自动初始化连接"""

{tamar_model_client-0.1.16 → tamar_model_client-0.1.18}/tamar_model_client/enums/invoke.py

@@ -7,4 +7,5 @@ class InvokeType(str, Enum):
     CHAT_COMPLETIONS = "chat-completions"
 
     GENERATION = "generation"  # 生成类，默认的值
-    IMAGE_GENERATION = "image-generation"
+    IMAGE_GENERATION = "image-generation"
+    IMAGE_EDIT_GENERATION = "image-edit-generation"

tamar_model_client-0.1.18/tamar_model_client/json_formatter.py

@@ -0,0 +1,26 @@
+import json
+import logging
+from datetime import datetime
+
+
+class JSONFormatter(logging.Formatter):
+    def format(self, record):
+        # log_type 只能是 request、response 或 info
+        log_type = getattr(record, "log_type", "info")
+        if log_type not in ["request", "response", "info"]:
+            log_type = "info"
+            
+        log_data = {
+            "timestamp": datetime.fromtimestamp(record.created).isoformat(),
+            "level": record.levelname,
+            "type": log_type,
+            "uri": getattr(record, "uri", None),
+            "request_id": getattr(record, "request_id", None),
+            "data": getattr(record, "data", None),
+            "message": record.getMessage(),
+            "duration": getattr(record, "duration", None),
+        }
+        # 增加 trace 支持
+        if hasattr(record, "trace"):
+            log_data["trace"] = getattr(record, "trace")
+        return json.dumps(log_data, ensure_ascii=False)