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

agent_api_server/service.py

@@ -0,0 +1,354 @@
+import asyncio
+import multiprocessing
+import os
+import threading
+
+from agent_api_server.config_center.config_center import AsyncConfigCenterClient
+from agent_api_server.configs import global_config
+import logging
+from fastapi import HTTPException
+from fastapi import FastAPI
+from agent_api_server.api.v1.api import api_router
+from pydantic_settings import BaseSettings
+from typing import List, Any
+from fastapi.staticfiles import StaticFiles
+from fastapi.middleware.cors import CORSMiddleware
+from pathlib import Path
+from agent_api_server.listener import ListenerType, create_listener
+from agent_api_server.mcp_convert.mcp_convert import create_mcp_tool_from_agent
+from agent_api_server.service_hub.service_hub import ServiceHubCredentialLoader, CredentialNotFoundError, \
+    InvalidCredentialError, ServiceHubError
+from agent_api_server.sso_service import SSOConfig, SSOService
+from fastapi.responses import Response, JSONResponse
+from fastapi.requests import Request
+from agent_api_server.register.register import AgentRegistry
+from typing import AsyncIterator
+from agent_api_server.cache.redis_cache import AsyncRedisThreadStorage
+from agent_api_server.memeory.postgres import AsyncPostgresCheckpointer
+from contextlib import asynccontextmanager
+from agent_api_server.shared.util_func import parse_agent_config
+
+WELCOME_ART = """
+╦  ┌─┐┌┐┌┌─┐╔═╗┬─┐┌─┐┌─┐┬ ┬
+║  ├─┤││││ ┬║ ╦├┬┘├─┤├─┘├─┤
+╩═╝┴ ┴┘└┘└─┘╚═╝┴└─┴ ┴┴  ┴ ┴
+
+🚀 API: http://{host}:{port}
+📚 API Docs: http://{host}:{port}/docs
+💻 client demo: http://{host}:{port}/site
+"""
+
+logger = logging.getLogger(__name__)
+
+class Settings(BaseSettings):
+    port: int = global_config.SERVER_PORT
+    mcp_port: int = global_config.MCP_SERVER_PORT
+    host: str = "0.0.0.0"
+    workers: int = global_config.SERVER_WORKER_AMOUNT
+    reload: bool = False
+    api_path: str = "/api/v1"
+    cors_origins: List[str] = ["*"]
+
+settings = Settings()
+
+@asynccontextmanager
+async def lifespan(app: FastAPI) -> AsyncIterator[None]:
+    cache = AsyncRedisThreadStorage.get_worker_instance()
+    checkpointer = AsyncPostgresCheckpointer.get_worker_instance()
+
+    try:
+        await checkpointer.initialize()
+        logger.info(f"PostgresSQL connection pool initialized in worker {os.getpid()}")
+
+        await cache.initialize()
+        logger.info(f"Redis connection pool initialized in worker {os.getpid()}")
+        yield
+    except Exception as e:
+        logger.error(f"Failed to initialize connection pool in worker {os.getpid()}")
+        raise
+    finally:
+        try:
+            await AsyncPostgresCheckpointer.close_worker_instance()
+            logger.info(f"PostgresSQL connection pool closed in worker {os.getpid()}")
+
+            await AsyncRedisThreadStorage.close_worker_instance()
+            logger.info(f"Redis connection pool closed in worker {os.getpid()}")
+        except Exception as e:
+            logger.error(f"Error closing connection pool in worker {os.getpid()}, error message is {str(e)}")
+
+
+app = FastAPI(lifespan=lifespan)
+
+CLIENT_DIR = Path(__file__).parent / "client"
+app.mount(
+    "/site",
+    StaticFiles(directory=str(CLIENT_DIR), html=True),
+    name="client"
+)
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=settings.cors_origins,
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+@app.exception_handler(Exception)
+async def universal_handler(request: Request, exc: Exception):
+    if "h11_impl" in str(exc):
+        return Response(status_code=204)
+
+    logger.error(
+        f"Unhandled Exception: {exc}",
+        exc_info=True,
+        extra={"path": request.url.path, "client": request.client}
+    )
+    return JSONResponse(
+        status_code=500,
+        content={"error": "internal_error", "message": f"{exc}"}
+    )
+
+app.include_router(api_router, prefix=settings.api_path)
+
+def print_welcome(host: str, port: int):
+    welcome_msg = WELCOME_ART.format(host=host, port=port)
+    logger.info(welcome_msg)
+
+async def load_agent_config():
+    try:
+        return await parse_agent_config()
+    except HTTPException as e:
+        raise e
+    except Exception as e:
+        raise e
+
+def create_fastapi_app() -> FastAPI:
+    return app
+
+
+async def init_fastapi_server():
+    print_welcome(host=settings.host, port=settings.port)
+
+    agents: list[dict[str, Any]]
+    try:
+        agents = await load_agent_config()
+        logger.info(f"Agent config loaded: {agents}")
+
+        if os.getenv("ENSAAS_SERVICES"):
+            logger.info(f"ENSAAS_SERVICES is not empty, Get POSTGRES_URL and REDIS_URL from ENSAAS_SERVICES")
+
+            try:
+                credentials = ServiceHubCredentialLoader.load_credentials()
+
+                logger.info(f"Get credentials from ENSAAS_SERVICES key, credentials is {credentials}, let's parse POSTGRES_URL, REDIS_URL and Blob information from it!")
+
+                global_config.POSTGRES_URL = credentials.to_postgres_dsn()
+                global_config.REDIS_URL = credentials.to_redis_dsn()
+                blob_credential = credentials.get_blobstore_s3_config()
+
+                if blob_credential != {}:
+                    global_config.STORAGE_ENDPOINT = blob_credential.get('endpoint', '')
+                    global_config.ACCESS_KEY = blob_credential.get('access_key', '')
+                    global_config.SECRET_KEY = blob_credential.get('secret_key', '')
+
+                    os.environ['STORAGE_ENDPOINT'] = global_config.STORAGE_ENDPOINT
+                    os.environ['ACCESS_KEY'] = global_config.ACCESS_KEY
+                    os.environ['SECRET_KEY'] = global_config.SECRET_KEY
+
+                    logger.info(
+                        f"Get POSTGRES_URL {global_config.POSTGRES_URL}, REDIS_URL {global_config.REDIS_URL} "
+                        f"Get STORAGE_ENDPOINT {global_config.STORAGE_ENDPOINT}, ACCESS_KEY {global_config.ACCESS_KEY}, SECRET_KEY {global_config.SECRET_KEY}"
+                        f"from ENSAAS_SERVICES key is ok!")
+                else:
+                    logger.info(
+                        f"Get POSTGRES_URL {global_config.POSTGRES_URL}, REDIS_URL {global_config.REDIS_URL} "
+                        f"from ENSAAS_SERVICES key is ok!")
+
+            except CredentialNotFoundError as e:
+                logger.error(f"Required credential not found: {e}")
+                raise
+            except InvalidCredentialError as e:
+                logger.error(f"Invalid credential format: {e}")
+                raise
+            except ServiceHubError as e:
+                logger.error(f"Failed to load credentials: {e}")
+                raise
+
+        if global_config.ENABLE_CONFIG_CENTER:
+            try:
+                config_client = await AsyncConfigCenterClient.get_instance()
+
+                if os.getenv("ENSAAS_SERVICES"):
+                    with_secret = True
+                else:
+                    with_secret = False
+
+                await config_client.register_to_config_center(agents=agents, with_secret=with_secret)
+                service_config = await config_client.service_config()
+                global_config.MODEL_MANAGER_SERVICE_URL = service_config["MODEL_MANAGER_SERVICE_URL"]
+                os.environ['MODEL_MANAGER_SERVICE_URL'] = global_config.MODEL_MANAGER_SERVICE_URL
+
+                if service_config["SSO_URL"] is not None and service_config["SSO_URL"] != "":
+                    global_config.SSO_URL = service_config["SSO_URL"]
+
+                mcp_gateway_url = await config_client.get_service_internal_url_with_router_path(
+                    service_name='MCPGateway',
+                    router_path='MCPGateway')
+                global_config.MCP_GATEWAY_URL = mcp_gateway_url
+                os.environ['MCP_GATEWAY_URL'] = global_config.MCP_GATEWAY_URL
+
+                logger.info(
+                    f"Got service config from config center: "
+                    f"MODEL_MANAGER_SERVICE_URL={os.environ['MODEL_MANAGER_SERVICE_URL']}, "
+                    f"SSO_URL={global_config.SSO_URL}, "
+                    f"MCP_GATEWAY_URL={global_config.MCP_GATEWAY_URL}")
+
+                if not with_secret:
+                    logger.info("ENABLE_CONFIG_CENTER is True and SECRET_NAME is empty,  get db and public information from config center")
+
+                    public_config = await config_client.public_config()
+                    os.environ['cluster'] = public_config.get('cluster', '')
+                    os.environ['datacenter'] = public_config.get('datacenter', '')
+                    os.environ['workspace'] = public_config.get('workspace', '')
+                    os.environ['namespace'] = public_config.get('namespace', '')
+                    logger.info(f"Got public config from config center: {public_config}")
+
+                    db_credentials = await config_client.database_credential()
+                    global_config.POSTGRES_URL = db_credentials["POSTGRES_URL"]
+                    global_config.REDIS_URL = db_credentials["REDIS_URL"]
+                    logger.info(
+                        f"Got database credentials from config center: POSTGRES_URL={global_config.POSTGRES_URL}, REDIS_URL={global_config.REDIS_URL}")
+
+                config_client.start_periodic_registration(agents, with_secret)
+                logger.info("Started periodic registration to config center successfully")
+
+            except Exception as e:
+                logger.error(f"Failed to get configuration from config center: {e}")
+                raise
+        else:
+            logger.info("ENABLE_CONFIG_CENTER is False, do not get service and db information from config center")
+
+        if global_config.POSTGRES_URL != "":
+            checkpointer = AsyncPostgresCheckpointer.get_worker_instance()
+            try:
+                await checkpointer.initialize()
+                await checkpointer.initialize_schema_with_retry()
+            except Exception as e:
+                logger.error(f"Init schema to postgres database with an un-expected error: {e}")
+                raise
+            finally:
+                logger.info("Init schema to postgres database with successfully and close postgres connection instance")
+                await checkpointer.close_worker_instance()
+
+        try:
+            sso_service = SSOService(sso_config=SSOConfig(
+                sso_address=global_config.SSO_URL,
+                client_id=global_config.CLIENT_ID,
+                client_secret=global_config.CLIENT_SECRET,
+            ))
+            client_id, client_secret = sso_service.get_client_id_and_secret()
+            logger.debug(f"Client ID: {client_id}, Client Secret: {client_secret}")
+
+            global_config.CLIENT_TOKEN = sso_service.generate_client_token(client_id, client_secret)
+            os.environ['CLIENT_TOKEN'] = global_config.CLIENT_TOKEN
+            logger.info(f"Generated Client Token: {os.environ['CLIENT_TOKEN']}")
+        except Exception as e:
+            logger.error(f"Failed during agent registration or listener setup: {e}")
+            raise
+
+        if not global_config.AGENT_AUTO_REGISTRATION:
+            logger.warning(
+                "Agent auto-registration is disabled by configuration (AGENT_AUTO_REGISTRATION=False).\n"
+                "IMPORTANT: Automatic registration to Model Manager is prohibited.\n"
+                "Required action:\n"
+                "1. You must manually register this agent to Model Manager\n"
+                "2. Use the Model Manager SDK to complete registration\n"
+            )
+            return
+
+        try:
+            for agent in agents:
+                agent["agent_url"] = global_config.SERVICE_EXTERNAL_URL
+
+            agent_registry = AgentRegistry(
+                base_url=global_config.MODEL_MANAGER_SERVICE_URL,
+                client_token=global_config.CLIENT_TOKEN
+            )
+            await agent_registry.register_all(agents=agents)
+
+            try:
+                listener_type = determine_listener_type()
+                logger.info(f"Using {listener_type.name} listener based on configuration")
+            except ValueError as e:
+                logger.error(str(e))
+                raise
+
+            for agent in agents:
+                listener = create_listener(
+                    agent["agent_name"],
+                    global_config.CLIENT_TOKEN,
+                    listener_type
+                )
+
+                thread = threading.Thread(
+                    target=listener.run,
+                    daemon=True,
+                    name=f"RedisListener-{agent['agent_name']}"
+                )
+                thread.start()
+                logger.info(f"Started {listener_type.name} listener for agent: {agent['agent_name']}")
+
+        except Exception as e:
+            logger.error(f"Failed during agent registration or listener setup: {e}")
+            raise
+
+    except ValueError as e:
+        logger.error(f"Failed to start agent server with a ValueError: {e}")
+        raise
+    except Exception as e:
+        logger.error(f"Failed to start agent server with an Exception error: {e}")
+        raise
+
+
+def determine_listener_type() -> ListenerType:
+    if hasattr(global_config, 'MODEL_MANAGER_NATS_URL') and global_config.MODEL_MANAGER_NATS_URL:
+        return ListenerType.NATS
+    elif hasattr(global_config, 'MODEL_MANAGER_REDIS_URL') and global_config.MODEL_MANAGER_REDIS_URL:
+        return ListenerType.REDIS
+    else:
+        raise ValueError(
+            "No valid message broker configured. "
+            "Please set either MODEL_MANAGER_NATS_URL or MODEL_MANAGER_REDIS_URL"
+        )
+
+def run_mcp_server(log_cfg):
+    try:
+        mcp = asyncio.run(create_mcp_tool_from_agent())
+        mcp.run(
+            transport="streamable-http",
+            show_banner=False,
+            uvicorn_config={"log_config": log_cfg},
+            host=settings.host,
+            port=settings.mcp_port,
+        )
+    except (KeyboardInterrupt, asyncio.CancelledError):
+        logging.info("MCP server received shutdown signal")
+    except Exception as e:
+        logging.error(f"MCP server failed: {e}")
+        raise
+
+
+def start_mcp_server_process(log_cfg):
+    if not global_config.ENABLE_MCP_SERVER:
+        logger.info("MCP server is disabled by configuration (ENABLE_MCP_SERVER=False)")
+        return None
+
+    process = multiprocessing.Process(
+        target=run_mcp_server,
+        args=(log_cfg,),
+        name="MCP-Server",
+    )
+    process.start()
+    logger.info(f"MCP server process started (PID: {process.pid})")
+    return process

agent_api_server/service_hub/service_hub.py

@@ -0,0 +1,233 @@
+import json
+import os
+import logging
+from typing import Dict, List, Union
+from dataclasses import dataclass, field
+
+logger = logging.getLogger(__name__)
+
+class ServiceHubError(Exception):
+    """Base exception for ServiceHub credential parsing errors"""
+    pass
+
+
+class CredentialNotFoundError(ServiceHubError):
+    """Raised when a specific credential type is not found"""
+    pass
+
+
+class InvalidCredentialError(ServiceHubError):
+    """Raised when credential data is invalid"""
+    pass
+
+
+@dataclass
+class PostgresCredential:
+    async_flag: bool
+    binding_name: str
+    database: str
+    external_hosts: str
+    host: str
+    internal_hosts: str
+    password: str
+    port: int
+    uri: str
+    username: str
+    instance_name: str = ""
+    label: str = ""
+    plan: str = ""
+    service_instance_id: str = ""
+    subscription_id: str = ""
+
+
+@dataclass
+class RedisCredential:
+    binding_name: str
+    external_hosts: str
+    host: str
+    internal_hosts: str
+    password: str
+    port: int
+    uri: str = ""
+    username: str = ""
+    instance_name: str = ""
+    label: str = ""
+    plan: str = ""
+    service_instance_id: str = ""
+    subscription_id: str = ""
+
+@dataclass
+class BlobstoreCredential:
+    """Blobstore (S3-compatible) connection credentials"""
+    binding_name: str
+    access_key: str
+    secret_key: str
+    endpoint: str
+    external_hosts: str
+    internal_hosts: str
+    signature_version: str
+    type: str
+    instance_name: str = ""
+    label: str = ""
+    plan: str = ""
+    service_instance_id: str = ""
+    subscription_id: str = ""
+
+
+@dataclass
+class ServiceHubCredentials:
+    postgres: List[PostgresCredential] = field(default_factory=list)
+    redis: List[RedisCredential] = field(default_factory=list)
+    blobstore: List[BlobstoreCredential] = field(default_factory=list)
+
+    def to_postgres_dsn(self) -> str:
+        if not self.postgres:
+            raise CredentialNotFoundError("PostgreSQL credentials not found")
+
+        cred = self.postgres[0]
+        return f"postgresql://{cred.username}:{cred.password}@{cred.internal_hosts}/{cred.database}"
+
+    def to_redis_dsn(self) -> str:
+        if not self.redis:
+            raise CredentialNotFoundError("Redis credentials not found")
+
+        cred = self.redis[0]
+
+        if not cred.password and not cred.username:
+            return f"redis://{cred.internal_hosts}:{cred.port}"
+
+        if cred.password and not cred.username:
+            return f"redis://:{cred.password}@{cred.internal_hosts}:{cred.port}"
+
+        if cred.username and not cred.password:
+            return f"redis://{cred.username}@{cred.internal_hosts}:{cred.port}"
+
+        if cred.username and cred.password:
+            return f"redis://{cred.username}:{cred.password}@{cred.internal_hosts}:{cred.port}"
+
+    def get_blobstore_s3_config(self) -> Dict[str, str]:
+        """Get Blobstore credentials as S3-compatible configuration dictionary"""
+        if not self.blobstore:
+            logger.warning("Blobstore credentials not found")
+            return {}
+
+        cred = self.blobstore[0]
+        return {
+            "access_key": cred.access_key,
+            "secret_key": cred.secret_key,
+            "endpoint": cred.endpoint,
+            "external_hosts": cred.external_hosts,
+            "internal_hosts": cred.internal_hosts,
+            "signature_version": cred.signature_version,
+            "type": cred.type
+        }
+
+    def get_service_instance_id(self) -> str:
+        """Get the first available service instance ID"""
+        for cred_list in [
+            self.postgres,
+            self.redis,
+            self.blobstore,
+        ]:
+            if cred_list and cred_list[0].service_instance_id:
+                return cred_list[0].service_instance_id
+        raise CredentialNotFoundError("No service instance ID found")
+
+
+class ServiceHubCredentialParser:
+    @staticmethod
+    def _parse_postgres(data: Dict) -> PostgresCredential:
+        return PostgresCredential(
+            async_flag=data.get("async", False),
+            binding_name=data.get("binding_name", ""),
+            database=data.get("credentials", {}).get("database", ""),
+            external_hosts=data.get("credentials", {}).get("externalHosts", ""),
+            host=data.get("credentials", {}).get("host", ""),
+            internal_hosts=data.get("credentials", {}).get("internalHosts", ""),
+            password=data.get("credentials", {}).get("password", ""),
+            port=data.get("credentials", {}).get("port", 0),
+            uri=data.get("credentials", {}).get("uri", ""),
+            username=data.get("credentials", {}).get("username", ""),
+            instance_name=data.get("instance_name", ""),
+            label=data.get("label", ""),
+            plan=data.get("plan", ""),
+            service_instance_id=data.get("serviceInstanceId", ""),
+            subscription_id=data.get("subscriptionId", ""),
+        )
+
+    @staticmethod
+    def _parse_redis(data: Dict) -> RedisCredential:
+        return RedisCredential(
+            binding_name=data.get("binding_name", ""),
+            external_hosts=data.get("credentials", {}).get("externalHosts", ""),
+            host=data.get("credentials", {}).get("host", ""),
+            internal_hosts=data.get("credentials", {}).get("internalHosts", ""),
+            password=data.get("credentials", {}).get("password", ""),
+            port=data.get("credentials", {}).get("port", 0),
+            uri=data.get("credentials", {}).get("uri", ""),
+            username=data.get("credentials", {}).get("username", ""),
+            instance_name=data.get("instance_name", ""),
+            label=data.get("label", ""),
+            plan=data.get("plan", ""),
+            service_instance_id=data.get("serviceInstanceId", ""),
+            subscription_id=data.get("subscriptionId", ""),
+        )
+
+    @staticmethod
+    def _parse_blobstore(data: Dict) -> BlobstoreCredential:
+        return BlobstoreCredential(
+            binding_name=data.get("binding_name", ""),
+            access_key=data.get("credentials", {}).get("accessKey", ""),
+            secret_key=data.get("credentials", {}).get("secretKey", ""),
+            endpoint=data.get("credentials", {}).get("endpoint", ""),
+            external_hosts=data.get("credentials", {}).get("externalHosts", ""),
+            internal_hosts=data.get("credentials", {}).get("internalHosts", ""),
+            signature_version=data.get("credentials", {}).get("signature-version", ""),
+            type=data.get("credentials", {}).get("type", ""),
+            instance_name=data.get("instance_name", ""),
+            label=data.get("label", ""),
+            plan=data.get("plan", ""),
+            service_instance_id=data.get("serviceInstanceId", ""),
+            subscription_id=data.get("subscriptionId", ""),
+        )
+
+    @staticmethod
+    def parse(data: Union[str, bytes, Dict]) -> ServiceHubCredentials:
+        """Parse ServiceHub credentials from JSON data"""
+        if isinstance(data, (str, bytes)):
+            try:
+                if isinstance(data, bytes):
+                    data = data.decode('utf-8')
+                data = json.loads(data)
+            except json.JSONDecodeError as e:
+                raise InvalidCredentialError(f"Invalid JSON data: {e}") from e
+
+        if not isinstance(data, dict):
+            raise InvalidCredentialError("Expected dictionary data")
+
+        credentials = ServiceHubCredentials()
+
+        if "postgresql" in data:
+            credentials.postgres = [
+                ServiceHubCredentialParser._parse_postgres(item)
+                for item in data["postgresql"]
+            ]
+
+        if "redis" in data:
+            credentials.redis = [
+                ServiceHubCredentialParser._parse_redis(item)
+                for item in data["redis"]
+            ]
+
+        if "blobstore" in data:
+            credentials.blobstore = [
+                ServiceHubCredentialParser._parse_blobstore(item)
+                for item in data["blobstore"]
+            ]
+
+        return credentials
+
+class ServiceHubCredentialLoader:
+    @staticmethod
+    def load_credentials() -> ServiceHubCredentials:
+        return ServiceHubCredentialParser.parse(os.getenv("ENSAAS_SERVICES"))