agent-api-server 2.1.7__py3-none-any.whl

agent_api_server/listener.py

@@ -0,0 +1,530 @@
+import json
+import threading
+import time
+import logging
+import asyncio
+from typing import Optional, Dict, Any, List
+import redis
+import nats
+from nats.js.api import StreamConfig, RetentionPolicy, StorageType, ConsumerConfig
+from model_manage_client import ModelManageClient
+from agent_api_server.callback_handler import notify_callbacks
+from agent_api_server.configs import global_config
+from enum import Enum, auto
+from dataclasses import dataclass
+from agent_api_server.shared.util_func import set_model_config
+
+logger = logging.getLogger(__name__)
+
+
+class ListenerType(Enum):
+    REDIS = auto()
+    NATS = auto()
+
+
+@dataclass
+class ModelUpdateEvent:
+    event_type: str
+    tenant_id: str
+    tool_name: Optional[str] = None
+    provider: Optional[str] = None
+    model: Optional[str] = None
+    credentials: Optional[Dict[str, Any]] = None
+    model_type: Optional[str] = None
+
+
+class MessageListener:
+    def __init__(self, agent_name: str, client_token: str):
+        self.agent_name = agent_name
+        self.client_token = client_token
+        self.m_client = ModelManageClient(global_config.MODEL_MANAGER_SERVICE_URL, client_token)
+        self.running = False
+
+    @staticmethod
+    def parse_nats_message(message_data: Dict[str, Any]) -> List[ModelUpdateEvent]:
+        events = []
+        try:
+            event_type = message_data.get('event_type', '')
+            message = message_data.get('message', {})
+
+            if isinstance(message, dict):
+                logger.debug("Message is a dictionary, processing...")
+                for tool_name, tool_data in message.items():
+                    if isinstance(tool_data, dict):
+                        for model_type, model_data in tool_data.items():
+                            events.append(ModelUpdateEvent(
+                                event_type=event_type,
+                                tenant_id=message_data.get('tenant_id', ''),
+                                tool_name=tool_name,
+                                provider=model_data.get('provider'),
+                                model=model_data.get('model'),
+                                credentials=model_data.get('credentials', {}),
+                                model_type=model_type
+                            ))
+
+        except Exception as e:
+            logger.error(f"Failed to parse message: {e}")
+
+        return events
+
+    @staticmethod
+    def parse_redis_message(message_data: Dict[str, Any]) -> List[ModelUpdateEvent]:
+        events = []
+        try:
+            event_type = message_data.get('event', '')
+            message = message_data.get('message', {})
+
+            if isinstance(message, dict):
+                logger.debug("Message is a dictionary, processing...")
+                for tool_name, tool_data in message.items():
+                    for model_type, model_data in tool_data.items():
+                        events.append(ModelUpdateEvent(
+                            event_type=event_type,
+                            tenant_id='',
+                            tool_name=tool_name,
+                            provider=model_data.get('provider'),
+                            model=model_data.get('model'),
+                            credentials=model_data.get('credentials', {}),
+                            model_type=model_type
+                        ))
+
+        except Exception as e:
+            logger.error(f"Failed to parse message: {e}")
+
+        return events
+
+    async def handle_events(self, events: List[ModelUpdateEvent], ts_id: str):
+        try:
+            agent_info = self.m_client.get_agent(agent_name=self.agent_name)
+            agent_id = agent_info[0].get('agent_id')
+            logger.info(f"Processing events for agent: {self.agent_name}, agent_id: {agent_id}")
+
+            for event in events:
+                if event.event_type in ('delete model', 'delete'):
+                    tool_name = event.tool_name or 'default'
+                    logger.info(f"""
+                        Processing {event.event_type} event:
+                        - Tenant ID: {ts_id}
+                        - Tool: {tool_name}
+                        - Provider: {event.provider}
+                        - Model: {event.model}
+                        - Model Type: {event.model_type}
+                        """)
+
+                    set_model_config(
+                        tool_name=tool_name,
+                        model_type=event.model_type,
+                        model_name=event.model,
+                        model_provider=event.provider,
+                        credentials="",
+                        agent_id=agent_id,
+                        ts_id=ts_id
+                    )
+
+                    notify_callbacks(
+                        ts_id=ts_id,
+                        client_token=self.client_token,
+                        model_type=event.model_type
+                    )
+                elif event.event_type in ('update model', 'update'):
+                    tool_name = event.tool_name or 'default'
+                    logger.info(f"""
+                        Processing {event.event_type} event:
+                        - Tenant ID: {ts_id}
+                        - Tool: {tool_name}
+                        - Provider: {event.provider}
+                        - Model: {event.model}
+                        - Model Type: {event.model_type}
+                        """)
+
+                    set_model_config(
+                        tool_name=tool_name,
+                        model_type=event.model_type,
+                        model_name=event.model,
+                        model_provider=event.provider,
+                        credentials=json.dumps(event.credentials) if event.credentials else "",
+                        agent_id=agent_id,
+                        ts_id=ts_id
+                    )
+
+                    notify_callbacks(
+                        ts_id=ts_id,
+                        client_token=self.client_token,
+                        model_type=event.model_type
+                    )
+
+        except Exception as e:
+            logger.error(f"Error processing events: {e}", exc_info=True)
+
+
+class RedisListener(MessageListener):
+    def __init__(self, agent_name: str, client_token: str):
+        super().__init__(agent_name, client_token)
+        self.redis_client: Optional[redis.Redis] = None
+        self.pubsub: Optional[redis.client.PubSub] = None
+        self.base_retry_delay = 5
+        self.max_retry_delay = 60
+
+    def run(self):
+        self.running = True
+        retry_delay = self.base_retry_delay
+
+        while self.running:
+            try:
+                self._connect()
+                retry_delay = self.base_retry_delay
+                self._listen()
+            except redis.ConnectionError as e:
+                logger.error(f"Redis connection error: {e}")
+                time.sleep(retry_delay)
+                retry_delay = min(retry_delay * 2, self.max_retry_delay)
+            except Exception as e:
+                logger.error(f"Unexpected error: {e}")
+                time.sleep(retry_delay)
+            finally:
+                self._disconnect()
+
+    def _connect(self):
+        try:
+            self.redis_client = redis.Redis.from_url(
+                url=global_config.MODEL_MANAGER_REDIS_URL,
+                socket_timeout=10,
+                socket_keepalive=True,
+                health_check_interval=15,
+                retry_on_timeout=True,
+                socket_connect_timeout=5
+            )
+
+            if not self.redis_client.ping():
+                raise ConnectionError("Redis ping failed")
+
+            self.pubsub = self.redis_client.pubsub()
+            self.pubsub.psubscribe(f'/model/update/{self.agent_name}/*')
+            logger.info(f"Redis listener connected for agent: {self.agent_name}")
+
+        except Exception as e:
+            self._disconnect()
+            raise ConnectionError(f"Redis connection failed: {e}")
+
+    def _listen(self):
+        while self.running:
+            try:
+                message = self.pubsub.get_message(
+                    ignore_subscribe_messages=True,
+                    timeout=30
+                )
+
+                if message is None:
+                    continue
+
+                logger.debug(f"Received raw message from Redis: {message}")
+
+                if message['type'] == 'pmessage':
+                    try:
+                        channel = message['channel'].decode()
+                        data = message['data'].decode()
+                        ts_id = channel.split('/')[-1]
+
+                        message_data = json.loads(data)
+                        events = self.parse_redis_message(message_data)
+
+                        asyncio.run(self.handle_events(events, ts_id))
+                    except json.JSONDecodeError as e:
+                        logger.error(f"Failed to decode message JSON: {e}")
+                    except Exception as e:
+                        logger.error(f"Error processing message: {e}")
+
+            except redis.ConnectionError:
+                logger.warning("Redis connection lost during listening")
+                raise
+            except Exception as e:
+                logger.error(f"Unexpected error in listen loop: {e}")
+                time.sleep(1)
+
+    def _disconnect(self):
+        try:
+            if self.pubsub:
+                self.pubsub.close()
+                self.pubsub = None
+        except Exception as e:
+            logger.warning(f"Error closing pubsub: {e}")
+
+        try:
+            if self.redis_client:
+                self.redis_client.close()
+                self.redis_client = None
+        except Exception as e:
+            logger.warning(f"Error closing redis client: {e}")
+
+    def stop(self):
+        self.running = False
+        self._disconnect()
+        logger.info("Redis listener stopped")
+
+
+class NATSListener(MessageListener):
+    def __init__(self, agent_name: str, client_token: str):
+        super().__init__(agent_name, client_token)
+        self.nc = None
+        self.js = None
+        self.subscription = None
+        self.loop = asyncio.new_event_loop()
+        self.reconnect_attempts = 0
+        self.max_reconnect_attempts = 10
+        self.base_reconnect_delay = 5
+        self.thread = threading.Thread(target=self._start_loop, daemon=True)
+        self.thread.start()
+
+    def _start_loop(self):
+        asyncio.set_event_loop(self.loop)
+        self.loop.run_forever()
+
+    def run(self):
+        self.running = True
+        asyncio.run_coroutine_threadsafe(self._connect_and_listen(), self.loop)
+
+    async def _connect_and_listen(self):
+        while self.running and self.reconnect_attempts < self.max_reconnect_attempts:
+            try:
+                await self._connect()
+                await self._ensure_stream_exists()
+                await self._subscribe_to_existing_consumer()
+                self.reconnect_attempts = 0
+                await self._monitor_connection()
+            except Exception as e:
+                logger.error(f"NATS connection error: {e}")
+                self.reconnect_attempts += 1
+                delay = min(self.base_reconnect_delay * (2 ** (self.reconnect_attempts - 1)), 60)
+                logger.info(f"Attempting to reconnect in {delay} seconds... (attempt {self.reconnect_attempts})")
+                await asyncio.sleep(delay)
+            finally:
+                await self._close()
+
+    async def _connect(self):
+        try:
+            self.nc = await nats.connect(
+                servers=[global_config.MODEL_MANAGER_NATS_URL],
+                reconnect_time_wait=5,
+                max_reconnect_attempts=-1,
+                disconnected_cb=self._on_disconnected,
+                reconnected_cb=self._on_reconnected,
+                closed_cb=self._on_closed,
+                error_cb=self._on_error
+            )
+            self.js = self.nc.jetstream()
+            logger.info(f"NATS connected for agent: {self.agent_name}")
+        except Exception as e:
+            logger.error(f"NATS connection failed: {e}")
+            raise
+
+    async def _ensure_stream_exists(self):
+        stream_name = "model-management-service"
+        subject = f"model-management-service.model.>"
+
+        try:
+            await self.js.stream_info(stream_name)
+            logger.debug(f"Stream {stream_name} already exists")
+        except Exception as e:
+            if "stream not found" in str(e).lower():
+                stream_config = StreamConfig(
+                    name=stream_name,
+                    subjects=[subject],
+                    retention=RetentionPolicy.LIMITS,
+                    max_msgs=10000,
+                    max_bytes=1024 * 1024 * 10,
+                    max_age=3600 * 24,
+                    storage=StorageType.FILE
+                )
+
+                try:
+                    await self.js.add_stream(config=stream_config)
+                    logger.info(f"Created stream {stream_name} with subject {subject}")
+                except Exception as create_error:
+                    logger.error(f"Failed to create stream {stream_name}: {create_error}")
+                    raise
+            else:
+                logger.error(f"Error checking stream existence: {e}")
+                raise
+
+    async def _subscribe_to_existing_consumer(self):
+        stream_name = "model-management-service"
+        subject = f"model-management-service.model.agent.{self.agent_name}.>"
+        durable_name = f"{self.agent_name}-consumer"
+
+        async def message_handler(msg):
+            try:
+                data = msg.data.decode()
+                message_data = json.loads(data)
+
+                prefix = f"model-management-service.model.agent.{self.agent_name}."
+                ts_id = msg.subject[len(prefix):]
+
+                logger.debug(f"Received raw message from NATS: {message_data}")
+
+                events = self.parse_nats_message(message_data)
+                await self.handle_events(events, ts_id)
+                await msg.ack()
+            except Exception as e:
+                logger.error(f"Message handling failed: {e}")
+
+        try:
+            if self.subscription:
+                try:
+                    await self.subscription.unsubscribe()
+                    logger.info(f"Unsubscribed from previous subscription")
+                    await asyncio.sleep(1)
+                except Exception as e:
+                    logger.warning(f"Error unsubscribing: {e}")
+                finally:
+                    self.subscription = None
+
+            consumer_info = await self.js.consumer_info(stream_name, durable_name)
+            logger.info(f"Found existing consumer: {durable_name}")
+
+            desired_config = ConsumerConfig(
+                durable_name=durable_name,
+                deliver_subject=f"deliver.{durable_name}",
+                ack_wait=30 * 1000,
+                deliver_policy="all",
+                ack_policy="explicit",
+                max_deliver=5,
+                filter_subject=subject
+            )
+
+            current_config = consumer_info.config
+            need_recreate = (
+                    current_config.deliver_subject != desired_config.deliver_subject or
+                    current_config.ack_wait != desired_config.ack_wait or
+                    current_config.deliver_policy != desired_config.deliver_policy or
+                    current_config.ack_policy != desired_config.ack_policy or
+                    current_config.max_deliver != desired_config.max_deliver or
+                    current_config.filter_subject != desired_config.filter_subject
+            )
+
+            if need_recreate:
+                logger.info(f"Consumer config mismatch, recreating consumer: {durable_name}")
+                try:
+                    await self.js.delete_consumer(stream_name, durable_name)
+                    logger.info(f"Deleted old consumer: {durable_name}")
+
+                    await self.js.add_consumer(stream_name, config=desired_config)
+                    logger.info(f"Created new consumer with updated config: {durable_name}")
+                except Exception as delete_error:
+                    logger.error(f"Failed to recreate consumer: {delete_error}")
+                    raise
+
+            self.subscription = await self.js.subscribe(
+                subject=subject,
+                durable=durable_name,
+                cb=message_handler,
+                stream=stream_name
+            )
+            logger.info(f"Subscribed to NATS consumer: {durable_name}")
+
+        except Exception as e:
+            if "consumer not found" in str(e).lower():
+                logger.info(f"Creating new consumer: {durable_name}")
+                consumer_config = ConsumerConfig(
+                    durable_name=durable_name,
+                    deliver_subject=f"deliver.{durable_name}",
+                    ack_wait=30 * 1000,
+                    deliver_policy="all",
+                    ack_policy="explicit",
+                    max_deliver=5,
+                    filter_subject=subject
+                )
+
+                await self.js.add_consumer(stream_name, config=consumer_config)
+
+                self.subscription = await self.js.subscribe(
+                    subject=subject,
+                    durable=durable_name,
+                    cb=message_handler,
+                    stream=stream_name
+                )
+                logger.info(f"Created and subscribed to new NATS consumer: {durable_name}")
+            else:
+                logger.error(f"Failed to subscribe: {e}")
+                raise
+
+    async def _monitor_connection(self):
+        while self.running:
+            try:
+                if not self.nc or self.nc.is_closed:
+                    logger.warning("NATS connection not available, attempting to reconnect...")
+                    await self._connect()
+                    continue
+
+                if not await self._check_consumer_status():
+                    logger.warning("Consumer not active, waiting before resubscribe...")
+                    await asyncio.sleep(10)
+                    logger.warning("Attempting to resubscribe...")
+                    await self._subscribe_to_existing_consumer()
+
+                await asyncio.sleep(5)
+            except Exception as e:
+                logger.error(f"Connection monitoring error: {e}")
+                await asyncio.sleep(10)
+
+    async def _check_consumer_status(self):
+        if not self.js:
+            return False
+
+        try:
+            durable_name = f"{self.agent_name}-consumer"
+            await self.js.consumer_info("model-management-service", durable_name)
+            return True
+        except Exception as e:
+            logger.warning(f"Error checking consumer status: {e}")
+            return False
+
+    async def _close(self):
+        if self.subscription:
+            try:
+                await self.subscription.unsubscribe()
+                self.subscription = None
+                logger.info("NATS subscription unsubscribed")
+            except Exception as e:
+                logger.warning(f"Error unsubscribing: {e}")
+
+        if self.nc and not self.nc.is_closed:
+            try:
+                await self.nc.close()
+                logger.info("NATS connection closed")
+            except Exception as e:
+                logger.warning(f"Error closing NATS connection: {e}")
+
+    @staticmethod
+    async def _on_disconnected():
+        logger.warning("NATS disconnected")
+
+    async def _on_reconnected(self):
+        logger.info("NATS reconnected")
+        if self.running:
+            await asyncio.sleep(2)
+            try:
+                await self._subscribe_to_existing_consumer()
+            except Exception as e:
+                logger.error(f"Failed to resubscribe: {e}")
+
+    @staticmethod
+    async def _on_closed():
+        logger.info("NATS connection closed")
+
+    @staticmethod
+    async def _on_error(e):
+        logger.error(f"NATS error: {e}")
+
+    def stop(self):
+        self.running = False
+        asyncio.run_coroutine_threadsafe(self._close(), self.loop)
+        logger.info("NATS listener stopped")
+
+
+def create_listener(agent_name: str, client_token: str, listener_type: ListenerType):
+    if listener_type == ListenerType.REDIS:
+        return RedisListener(agent_name, client_token)
+    elif listener_type == ListenerType.NATS:
+        return NATSListener(agent_name, client_token)
+    else:
+        raise ValueError(f"Unsupported listener type: {listener_type}")

agent_api_server/log/formatters.py

@@ -0,0 +1,122 @@
+import logging
+import logging.config
+import json
+from pathlib import Path
+from typing import Dict, Any
+from datetime import datetime, timezone
+
+_logging_config: Dict[str, Any] = {}
+_current_log_level = logging.INFO
+
+
+def setup_logging() -> Dict[str, Any]:
+    global _logging_config, _current_log_level
+
+    current_dir = Path(__file__).parent.absolute()
+    config_path = current_dir / "logging.json"
+
+    if not config_path.exists():
+        raise FileNotFoundError(f"Logging config not found at: {config_path}")
+
+    with open(config_path, 'r', encoding='utf-8') as f:
+        _logging_config = json.load(f)
+
+    _current_log_level = _logging_config.get('root', {}).get('level', logging.INFO)
+    if isinstance(_current_log_level, str):
+        _current_log_level = logging.getLevelName(_current_log_level.upper())
+
+    _apply_logging_config()
+    return _logging_config
+
+
+def _apply_logging_config():
+    logging.config.dictConfig(_logging_config)
+
+    for logger_name in logging.root.manager.loggerDict:
+        logger = logging.getLogger(logger_name)
+        if not logger.handlers:
+            logger.setLevel(_current_log_level)
+
+
+def set_log_level(level: str | int):
+    global _current_log_level, _logging_config
+
+    if isinstance(level, str):
+        level = level.upper()
+        try:
+            level = logging.getLevelName(level) if level.isdigit() else getattr(logging, level)
+        except AttributeError:
+            raise ValueError(f"Invalid log level: {level}")
+
+    if level not in (logging.DEBUG, logging.INFO, logging.WARNING, logging.ERROR, logging.CRITICAL):
+        raise ValueError(f"Invalid log level value: {level}")
+
+    _current_log_level = level
+    if 'root' in _logging_config:
+        _logging_config['root']['level'] = logging.getLevelName(level)
+
+    _apply_logging_config()
+    logging.info(f"Log level changed to {logging.getLevelName(level)}")
+
+
+def get_log_level() -> str:
+    return logging.getLevelName(_current_log_level)
+
+class ColorFormatter(logging.Formatter):
+    COLORS = {
+        'timestamp': '\x1b[38;5;244m',  #浅灰色时间戳
+        'debug': '\x1b[36m',  # Cyan
+        'info': '\x1b[32m',  # Green
+        'warning': '\x1b[33m',  # Yellow
+        'error': '\x1b[31m',  # Red
+        'critical': '\x1b[41;1m',  # 红色背景
+        'module': '\x1b[38;5;33m',  # 亮蓝色
+        'thread': '\x1b[38;5;33m',  # 亮蓝色
+        'reset': '\x1b[0m'
+    }
+
+    def format(self, record):
+        colors = {
+            'timestamp': self.COLORS['timestamp'],
+            'level': self.COLORS.get(record.levelname.lower(), ''),
+            'module': self.COLORS['module'],
+            'thread': self.COLORS['thread'],
+            'reset': self.COLORS['reset']
+        }
+
+        timestamp = datetime.fromtimestamp(record.created, tz=timezone.utc).strftime('%Y-%m-%dT%H:%M:%S.%f')[:23] + 'Z'
+
+        message = self._format_message(record)
+
+        log_template = (
+            "{timestamp_color}{timestamp}{reset} "
+            "[{level_color}{levelname:8}{reset}] "
+            "{message} "
+            "[{module_color}{module}{reset}] "
+            "thread_name={thread_color}{thread}{reset}"
+        )
+
+        return log_template.format(
+            timestamp_color=colors['timestamp'],
+            timestamp=timestamp,
+            reset=colors['reset'],
+            level_color=colors['level'],
+            levelname=record.levelname.lower(),
+            message=message,
+            module_color=colors['module'],
+            module=record.module,
+            thread_color=colors['thread'],
+            thread=record.threadName
+        )
+
+    @staticmethod
+    def _format_message(record):
+        try:
+            if record.args:
+                if isinstance(record.args, (tuple, dict)):
+                    return record.msg % record.args
+                elif isinstance(record.args, dict):
+                    return record.msg.format(**record.args)
+        except Exception:
+            pass
+        return record.msg

agent_api_server/log/logging.json

@@ -0,0 +1,50 @@
+{
+  "version": 1,
+  "incremental": false,
+  "disable_existing_loggers": false,
+  "formatters": {
+    "standard": {
+      "()": "agent_api_server.log.formatters.ColorFormatter",
+      "format": "%(asctime)s.%(msecs)03dZ [%(levelname)-8s] %(message)s [%(name)s] thread=%(threadName)s",
+      "datefmt": "%Y-%m-%d %H:%M:%S"
+    }
+  },
+  "handlers": {
+    "console": {
+      "class": "logging.StreamHandler",
+      "formatter": "standard",
+      "stream": "ext://sys.stdout"
+    }
+  },
+  "root": {
+    "handlers": ["console"],
+    "level": "INFO"
+  },
+  "loggers": {
+    "fastmcp": {
+      "handlers": ["console"],
+      "level": "INFO",
+      "propagate": false
+    },
+    "FastMCP": {
+      "handlers": ["console"],
+      "level": "INFO",
+      "propagate": false
+    },
+    "FastMCP.server": {
+      "handlers": ["console"],
+      "level": "INFO",
+      "propagate": false
+    },
+    "uvicorn": {
+      "handlers": ["console"],
+      "level": "INFO",
+      "propagate": false
+    },
+    "uvicorn.error": {
+      "handlers": ["console"],
+      "level": "INFO",
+      "propagate": false
+    }
+  }
+}