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

agent_api_server/config_center/config_center.py

@@ -0,0 +1,239 @@
+import asyncio
+import aiohttp
+import logging
+import os
+import threading
+import time
+from typing import List, Dict
+from urllib.parse import urljoin
+
+from agent_api_server.configs import global_config
+
+logger = logging.getLogger(__name__)
+
+
+class AsyncConfigCenterClient:
+    _instance = None
+
+    def __init__(self):
+        self.url = os.getenv("CONFIG_CENTER_ADDRESS") or "http://configcenter:10086"
+        if self.url.endswith("/"):
+            self.url = self.url.rstrip("/")
+        self._agents = []
+        self._session = None
+        logger.info(f"ConfigCenter Client registered URL is: {self.url}")
+
+    async def initialize(self):
+        if self._session is None or self._session.closed:
+            self._session = aiohttp.ClientSession()
+            logger.info("ConfigCenter Client initialized")
+
+    async def close(self):
+        if self._session:
+            await self._session.close()
+            self._session = None
+            logger.info("ConfigCenter Client closed")
+
+    async def ensure_initialized(self):
+        if self._session is None or self._session.closed:
+            await self.initialize()
+
+    async def register_to_config_center(self, agents: List[Dict], with_secret: bool) -> None:
+        if not agents:
+            raise ValueError("Agents list cannot be empty")
+
+        await self.ensure_initialized()
+        self._agents = agents
+
+        for agent in agents:
+            if "agent_name" not in agent:
+                raise ValueError(f"Agent is missing 'agent_name' field: {agent}")
+            if not agent.get("agent_labels"):
+                raise ValueError(f"Agent is missing 'agent_labels' field: {agent}")
+
+        await self._register_agents(agents, global_config.ENABLE_MCP_SERVER, with_secret)
+
+    async def _register_agents(self, agents: List[Dict], is_mcp: bool, with_secret: bool) -> None:
+        config_endpoint = f"{self.url}/v1/services"
+
+        main_agent = agents[0]
+        if not main_agent.get("agent_labels"):
+            raise ValueError("agent_labels is required but empty")
+
+        service_name = main_agent["agent_labels"][0]
+        payload = {
+            "name": service_name,
+            "serviceInstances": [] if with_secret else [
+                {"schema": "postgres"},
+                {"schema": "redis"},
+            ],
+            "routers": {}
+        }
+
+        for agent in agents:
+            agent_name = agent["agent_name"]
+            payload["routers"][agent_name] = {
+                "name": agent_name,
+                "address": os.getenv("SERVICE_EXTERNAL_URL")
+            }
+
+        if is_mcp:
+            mcp_router_name = f"{service_name}_mcp"
+            payload["routers"][mcp_router_name] = {
+                "name": mcp_router_name,
+                "address": urljoin(os.getenv("MCP_SERVER_EXTERNAL_URL"), "mcp"),
+                "isMcpServer": True,
+                "mcpServerParams": {
+                    "isAuth": "true",
+                }
+            }
+
+        max_retries = 3
+        for attempt in range(max_retries):
+            try:
+                async with self._session.post(config_endpoint, json=payload, timeout=30) as response:
+                    response.raise_for_status()
+                    router_names = list(payload['routers'].keys())
+                    logger.info(
+                        f"Service '{service_name}' registered successfully with {len(router_names)} routers: {router_names}")
+                    return
+            except Exception as e:
+                logger.error(f"Attempt {attempt + 1}: Failed to register service '{service_name}', "
+                             f"request url: {config_endpoint}; error {e}")
+                if attempt < max_retries - 1:
+                    await asyncio.sleep(2 ** attempt)
+                else:
+                    logger.error(f"Failed to register service '{service_name}' after 3 attempts")
+                    raise e
+
+    def start_periodic_registration(self, agents: List[Dict], with_secret: bool, interval: int = 5):
+        if not agents:
+            raise ValueError("Agents list cannot be empty")
+
+        self._agents = agents
+
+        def registration_loop():
+            logger.info(f"🚀 Starting periodic registration thread with interval {interval}s")
+
+            loop = asyncio.new_event_loop()
+            asyncio.set_event_loop(loop)
+
+            session = aiohttp.ClientSession(loop=loop)
+            original_session = self._session
+            self._session = session
+
+            try:
+                while True:
+                    try:
+                        logger.info(f"🔄 Registration started")
+                        loop.run_until_complete(self._run_registration_iteration(with_secret))
+                        logger.info(f"✅ Registration completed successfully")
+                        time.sleep(interval)
+                    except Exception as e:
+                        logger.error(f"❌ Error in registration: {e}")
+                        time.sleep(5)
+            finally:
+                loop.run_until_complete(session.close())
+                loop.close()
+                self._session = original_session
+                logger.info("⏹️ Periodic registration thread stopped")
+
+        thread = threading.Thread(
+            target=registration_loop,
+            name="config_center_registration",
+            daemon=True
+        )
+        thread.start()
+        logger.info("✅ Periodic registration thread started")
+
+    async def _run_registration_iteration(self, with_secret: bool):
+        await self.ensure_initialized()
+        await self._register_agents(self._agents, global_config.ENABLE_MCP_SERVER, with_secret)
+
+    async def database_credential(self) -> dict:
+        if not self._agents:
+            raise ValueError("No agents registered yet")
+
+        await self.ensure_initialized()
+
+        main_agent = self._agents[0]
+        if not main_agent.get("agent_labels"):
+            raise ValueError("agent_labels is required but empty")
+
+        service_name = main_agent["agent_labels"][0]
+        config_endpoint = urljoin(self.url, f"/v1/{service_name}/serviceintance")
+
+        async with self._session.get(config_endpoint) as response:
+            response.raise_for_status()
+            data = await response.json()
+
+        postgres_url = data.get("postgres", "")
+        if postgres_url:
+            postgres_url = postgres_url.replace("postgres://", "postgresql://", 1)
+
+        redis_url = data.get("redis", "")
+        return {
+            "POSTGRES_URL": postgres_url,
+            "REDIS_URL": redis_url
+        }
+
+    async def public_config(self) -> dict:
+        await self.ensure_initialized()
+        config_endpoint = urljoin(self.url, "/v1/public/configs")
+
+        async with self._session.get(config_endpoint) as response:
+            response.raise_for_status()
+            config_data = await response.json()
+
+        return {
+            "cluster": config_data.get("cluster"),
+            "datacenter": config_data.get("datacenter"),
+            "workspace": config_data.get("workspace"),
+            "namespace": config_data.get("namespace")
+        }
+
+    async def service_config(self) -> dict:
+        await self.ensure_initialized()
+        config_endpoint = urljoin(self.url, "/v1/services")
+
+        async with self._session.get(config_endpoint) as response:
+            response.raise_for_status()
+            config_data = await response.json()
+
+        sso_api_url = config_data.get("SSO", {}).get("routers", {}).get("API", {}).get("address")
+        model_management_url = config_data.get("model-management", {}).get("routers", {}).get("API", {}).get("address")
+
+        return {
+            "SSO_URL": sso_api_url,
+            "MODEL_MANAGER_SERVICE_URL": model_management_url,
+        }
+
+    async def get_service_internal_url_with_router_path(self, service_name: str, router_path: str) -> str:
+        await self.ensure_initialized()
+        config_endpoint = urljoin(self.url, "/v1/services")
+
+        async with self._session.get(config_endpoint) as response:
+            response.raise_for_status()
+            config_data = await response.json()
+
+        url = config_data.get(service_name, {}).get("routers", {}).get(router_path, {}).get("address")
+
+        return url
+
+    async def get_service_public_url_with_router_path(self, service_name: str, router_path: str) -> str:
+        await self.ensure_initialized()
+        config_endpoint = urljoin(self.url, "/v1/services")
+
+        async with self._session.get(config_endpoint) as response:
+            response.raise_for_status()
+            config_data = await response.json()
+
+        url = config_data.get(service_name, {}).get("routers", {}).get(router_path, {}).get("publicAddress")
+
+        return url
+
+    @staticmethod
+    async def get_instance() -> "AsyncConfigCenterClient":
+        if AsyncConfigCenterClient._instance is None:
+            AsyncConfigCenterClient._instance = AsyncConfigCenterClient()
+        return AsyncConfigCenterClient._instance

agent_api_server/configs/__init__.py

@@ -0,0 +1,3 @@
+from .config import Config
+
+global_config = Config()

agent_api_server/configs/config.py

@@ -0,0 +1,163 @@
+from pydantic_settings import BaseSettings, SettingsConfigDict
+from pydantic import Field
+from pathlib import Path
+from typing import Optional
+from dotenv import load_dotenv
+
+
+class SSOConfig(BaseSettings):
+    SSO_URL: str = Field(
+        default="http://sso/v4.0", description="The URL of SSO service"
+    )
+    CLIENT_ID: str = Field(default="", description="The client-id of server")
+    CLIENT_SECRET: str = Field(default="", description="The client secret of server")
+    CLIENT_TOKEN: str = Field(default="", description="The client token of server")
+
+
+class CacheConfig(BaseSettings):
+    REDIS_URL: Optional[str] = Field(
+        default="redis://localhost:6379/0",
+        description="Redis server URL in format redis://host:port/db"
+    )
+    MAX_THREADS: int = Field(
+        default=1000,
+        description="Maximum number of threads for concurrent operations"
+    )
+
+class PostgresConfig(BaseSettings):
+    POSTGRES_URL: Optional[str] = Field(
+        default="postgresql://user:password@host:port/dbname",
+        description="complete PostgresSQL URL，formater is：postgresql://user:password@host:port/dbname"
+    )
+
+    # 连接池参数
+    POSTGRES_POOL_MIN_SIZE: int = Field(
+        default=1,
+        description="mini connection count for postgres"
+    )
+
+    POSTGRES_POOL_MAX_SIZE: int = Field(
+        default=10,
+        description="max connection count for postgres"
+    )
+
+    POSTGRES_POOL_TIMEOUT: float = Field(
+        default=30.0,
+        description="time out for postgres connection"
+    )
+
+    POSTGRES_POOL_RECYCLE: int = Field(
+        default=3600,
+        description="recycle time for postgres connection"
+    )
+
+    POSTGRES_POOL_MAX_INACTIVE: int = Field(
+        default=600,
+        description="inactive time for postgres"
+    )
+
+class LoggingConfig(BaseSettings):
+    LOG_LEVEL: str = Field(
+        default="INFO", description="Logging level for the application"
+    )
+
+class ModelManagerServiceConfig(BaseSettings):
+    MODEL_MANAGER_SERVICE_URL: str = Field(
+        default="https://api-am-ensaas.axa.wise-paas.com.cn", description="The URL of model manager service url"
+    )
+
+    SERVICE_EXTERNAL_URL: str = Field(
+        default="http://172.21.84.86:8080", description="The URL of agent"
+    )
+
+    MODEL_MANAGER_REDIS_URL: str = Field(
+        default="redis://localhost:6379/0", description="Redis server URL which used by model_manager, format is: redis://host:port/db"
+    )
+
+    MODEL_MANAGER_NATS_URL: str =  Field(
+        default="", description="NATS server URL which used by model_manager, format is: nats://localhost:4222"
+    )
+
+class WebServerConfig(BaseSettings):
+    SERVER_PORT: int = Field(
+        default=8080, description="The port of web server"
+    )
+
+    SERVER_WORKER_AMOUNT: int = Field(
+        default=1, description="The amount of web server."
+    )
+
+    AGENT_AUTO_REGISTRATION: bool = Field(
+        default=False,
+        description="Whether to enable automatic agent registration through API server. "
+                    "If True, agents will be automatically registered via API. "
+                    "If False, manual registration is required."
+    )
+
+    ENABLE_MCP_SERVER: bool = Field(
+        default=False,
+        description="Whether to enable mcp server"
+    )
+
+    MCP_SERVER_PORT: int = Field(
+        default=8081, description="The port of mcp server"
+    )
+
+    MCP_SERVER_EXTERNAL_URL: str = Field(
+        default="http://172.21.84.86:8081", description="The URL of agent"
+    )
+
+    MCP_GATEWAY_URL: str = Field(
+        default="http://172.21.92.186:10105", description="The URL of MCP gateway"
+    )
+
+
+class ConfigCenterConfig(BaseSettings):
+    ENABLE_CONFIG_CENTER: bool =Field(
+        default=False, description="Whether to enable config center to get service information."
+    )
+
+    CONFIG_CENTER_ADDRESS: str = Field(
+        default='', description="The URL of config center."
+    )
+
+class StorageConfig(BaseSettings):
+    STORAGE_ENDPOINT: str =  Field(
+        default='', description="The URL of storage(like: s3)."
+    )
+
+    ACCESS_KEY: str =  Field(
+        default='', description="The access key of storage(like: s3)."
+    )
+
+    SECRET_KEY: str =  Field(
+        default='', description="The secret key of storage(like: s3)."
+    )
+
+
+def find_and_load_dotenv(filename: str = '.env') -> Optional[str]:
+    search_paths = [
+        Path.cwd() / filename,
+        Path(__file__).parent.parent / filename,
+    ]
+
+    for path in search_paths:
+        if path.exists():
+            load_dotenv(path, override=True)
+            return str(path)
+
+    for parent in Path.cwd().parents:
+        candidate = parent / filename
+        if candidate.exists():
+            load_dotenv(candidate, override=True)
+            return str(candidate)
+
+    return None
+
+class Config(PostgresConfig, CacheConfig, LoggingConfig, SSOConfig, WebServerConfig, ModelManagerServiceConfig, ConfigCenterConfig, StorageConfig):
+    model_config = SettingsConfigDict(
+        env_file=find_and_load_dotenv(),
+        env_file_encoding="utf-8",
+        extra="ignore",
+    )
+