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

agent_api_server/shared/util_func.py

@@ -0,0 +1,372 @@
+import os
+from typing import Dict
+from importlib import import_module
+from fastapi import HTTPException
+from pathlib import Path
+import json
+import logging
+import aiofiles
+import asyncio
+import sys
+from typing import List, Any
+from agent_api_server.memeory.postgres import AsyncPostgresCheckpointer
+from langgraph.pregel import Pregel
+
+from agent_api_server.shared.common import ConfigCategory
+
+logger = logging.getLogger(__name__)
+
+BASE_DIR = Path.cwd()
+DATA_DIR = BASE_DIR / 'data'
+DATA_DIR.mkdir(exist_ok=True, mode=0o700)
+LLM_CFG_ENV = DATA_DIR /'.llm_cfg_env'
+
+async def get_all_graph_names() -> List[str]:
+    """
+    获取配置文件中定义的所有graph名称
+
+    Returns:
+        List[str]: 所有可用的graph名称列表
+
+    Raises:
+        HTTPException: 如果配置文件加载失败
+    """
+    try:
+        config = await load_graph_config()
+        if not config.get("graphs"):
+            return []
+        return list(config["graphs"].keys())
+    except Exception as e:
+        raise HTTPException(
+            status_code=500,
+            detail={
+                "error": "config_loading_failed",
+                "message": f"Cannot get graph names: {str(e)}"
+            }
+        )
+
+
+async def load_graph(graph_name: str, config: Dict, with_checker: bool) -> tuple[str, Pregel, None] | tuple[
+    str, Pregel, AsyncPostgresCheckpointer]:
+    if not config.get("graphs") or graph_name not in config["graphs"]:
+        available = list(config["graphs"].keys()) if "graphs" in config else []
+        raise HTTPException(
+            status_code=404,
+            detail={
+                "error": "graph_not_found",
+                "available": available,
+                "message": f"Graph '{graph_name}' not found in config"
+            }
+        )
+
+    graph_path = config["graphs"][graph_name]
+    try:
+        module_path, attr_name = graph_path.rsplit(":", 1)
+        if not module_path or not attr_name:
+            raise ValueError("Empty module path or attribute name")
+    except ValueError as e:
+        raise HTTPException(
+            status_code=400,
+            detail={
+                "error": "invalid_path_format",
+                "expected": "path/to/module.py:attribute_name",
+                "actual": graph_path,
+                "message": str(e)
+            }
+        )
+
+    try:
+        module_path = module_path.replace("./", "").replace("\\", "/")
+        if module_path.endswith(".py"):
+            module_path = module_path[:-3]
+        import_path = module_path.replace("/", ".")
+    except Exception as e:
+        raise HTTPException(
+            status_code=400,
+            detail={
+                "error": "path_processing_failed",
+                "message": f"Cannot process path '{module_path}': {str(e)}"
+            }
+        )
+
+    try:
+        module_dir = str(Path(module_path).parent.resolve())
+        if module_dir not in sys.path:
+            sys.path.insert(0, module_dir)
+
+        module = await asyncio.to_thread(import_module, import_path)
+    except ModuleNotFoundError as e:
+        raise HTTPException(
+            status_code=500,
+            detail={
+                "error": "module_not_found",
+                "path": import_path,
+                "sys_path": str(sys.path),
+                "message": f"Module not found: {str(e)}"
+            }
+        )
+    except Exception as e:
+        raise HTTPException(
+            status_code=500,
+            detail={
+                "error": "module_import_error",
+                "path": import_path,
+                "message": f"Error importing module: {str(e)}"
+            }
+        )
+
+    # 5. 获取graph实例 - 更健壮的实例化
+    try:
+        graph_factory = getattr(module, attr_name)
+
+        # 处理可能是协程或工厂函数的情况
+        if asyncio.iscoroutinefunction(graph_factory):
+            graph = await graph_factory()
+        elif callable(graph_factory):
+            graph = graph_factory()
+        else:
+            graph = graph_factory
+
+        # 确保是 Pregel 实例
+        if not isinstance(graph, Pregel):
+            raise TypeError(f"Expected Pregel instance, got {type(graph)}")
+
+        if with_checker:
+            checkpointer = AsyncPostgresCheckpointer.get_worker_instance()
+            try:
+                graph.checkpointer = await checkpointer.checkpointer()
+            except Exception as e:
+                raise RuntimeError(f"Failed to initialize checkpointer: {str(e)}")
+
+            return graph_name, graph, checkpointer
+
+        return graph_name, graph,None
+    except AttributeError:
+        raise HTTPException(
+            status_code=500,
+            detail={
+                "error": "missing_attribute",
+                "module": import_path,
+                "attribute": attr_name,
+                "available": dir(module)
+            }
+        )
+    except Exception as e:
+        raise HTTPException(
+            status_code=500,
+            detail={
+                "error": "graph_initialization_failed",
+                "message": f"Failed to initialize graph: {str(e)}",
+                "type": type(e).__name__
+            }
+        )
+
+
+def resolve_config_path(file_name: str) -> Path:
+    search_paths = [
+        Path.cwd() / file_name,
+        Path(__file__).parent.parent / file_name,
+    ]
+
+    for path in search_paths:
+        if path.exists():
+            return path
+
+    for parent in Path.cwd().parents:
+        candidate = parent / file_name
+        if candidate.exists():
+            return candidate
+
+    raise FileNotFoundError(f"Could not locate file {file_name}")
+
+
+async def load_graph_config() -> Dict[str, str]:
+    try:
+        json_path = resolve_config_path("langgraph.json")
+        async with aiofiles.open(json_path, mode='r') as f:
+            content = await f.read()
+            config = json.loads(content)
+            logger.info(f"Successfully loaded config from {json_path}")
+            return config
+    except Exception as e:
+        raise HTTPException(
+            status_code=500,
+            detail=f"Config loading failed: {str(e)}"
+        )
+
+
+async def parse_agent_config() -> List[Dict[str, Any]]:
+    """
+    解析langgraph.json配置文件，返回所有Agent的标准化信息数组
+
+    Returns:
+        List[Dict[str, Any]]: 每个Agent的信息字典组成的数组，结构为：
+            [{
+                "agent_name": str,
+                "agent_description": str,
+                "agent_api_version": str,
+                "agent_icon_url": str,      # 预留字段
+                "agent_features": Dict[str, Any]  # 预留字段
+                "agent_label": str, # Agent的标签
+                "has_site": bool # Agent是否要产生chatbot
+                "multilangs": Dict[str, Dict[str, str]]  # 多语言标签,
+                "is_system_agent": bool # Agent是否要注册为系统级Agent，如果注册为系统级Agent则全局租户共用一个，否则分租户使用
+            }]
+
+    Raises:
+        HTTPException: 如果配置文件加载或解析失败，包含详细错误分类
+    """
+    try:
+        config = await load_graph_config()
+
+        if not config.get("graphs") or not config.get("agent_description"):
+            raise ValueError("Missing required 'graphs' or 'agent_description' in config")
+
+        multilangs = config.get("multilangs", {})
+        default_multilang = {
+            "default": {
+                "en_US": "Default",
+                "zh_CN": "默认",
+                "zh_TW": "预设"
+            }
+        }
+
+        if multilangs:
+            multilangs = {**default_multilang, **multilangs}
+        else:
+            multilangs = default_multilang
+
+        agents = []
+        for agent_name in config["graphs"].keys():
+            agents.append({
+                "agent_name": agent_name,
+                "agent_description": config["agent_description"].get(agent_name, ""),
+                "agent_api_version": config.get("agent_api_version", "v0.0.1"),
+                "agent_icon_url": config.get("agent_icon_url", ""),
+                "agent_features": config.get("agent_features", {}),
+                "agent_labels": config.get("agent_labels", []),
+                "has_site": config.get("has_site", False),
+                "is_system_agent": config.get("is_system_agent", False),
+                "multilangs": multilangs
+            })
+
+        return agents
+    except HTTPException as e:
+        raise e
+    except Exception as e:
+        raise e
+
+
+def set_model_config(
+        tool_name: str,
+        model_type: str,
+        model_provider: str,
+        model_name: str,
+        credentials: dict,
+        agent_id: str,
+        ts_id: str = None,
+) -> None:
+    category_map = {
+        'llm': (ConfigCategory.CHAT_PROVIDER.value, ConfigCategory.CHAT_MODEL.value,
+                ConfigCategory.CHAT_CREDENTIALS.value, ConfigCategory.AGENT_ID.value),
+
+        'text-embedding': (ConfigCategory.EMBEDDING_PROVIDER.value, ConfigCategory.EMBEDDING_MODEL.value,
+                           ConfigCategory.EMBEDDING_CREDENTIALS.value, ConfigCategory.AGENT_ID.value),
+
+        'rerank': (ConfigCategory.RERANK_PROVIDER.value, ConfigCategory.RERANK_MODEL.value,
+                   ConfigCategory.RERANK_CREDENTIALS.value, ConfigCategory.AGENT_ID.value)
+    }
+
+    try:
+        provider_suffix, model_suffix, credentials_key_suffix, agent_id_suffix = category_map[model_type]
+    except KeyError:
+        if logger:
+            warning_msg = f"Unsupported model type {model_type}"
+            if tool_name != "default":
+                warning_msg += f" for tool {tool_name}"
+            logger.warning(warning_msg)
+        return
+
+    base_prefix = "" if tool_name == "default" else f"{tool_name.upper()}_"
+
+    provider_key = f"{base_prefix}{provider_suffix}"
+    model_key = f"{base_prefix}{model_suffix}"
+    credentials_key = f"{base_prefix}{credentials_key_suffix}"
+    agent_id_key = f"{base_prefix}{agent_id_suffix}"
+
+    if ts_id:
+        provider_key = f"{provider_key}_{ts_id}"
+        model_key = f"{model_key}_{ts_id}"
+        credentials_key = f"{credentials_key}_{ts_id}"
+        agent_id_key = f"{agent_id_key}_{ts_id}"
+
+    new_vars = {
+        provider_key: model_provider,
+        model_key: model_name,
+        credentials_key: credentials,
+        agent_id_key: agent_id
+    }
+    LLM_INFO_PATH = str(LLM_CFG_ENV) + f"_{ts_id}"
+    update_env_file(LLM_INFO_PATH, new_vars)
+
+    logger.debug(f"\nset model config to {LLM_INFO_PATH} ok\n"
+                 f"provider_key is {provider_key}, provider_value is {model_provider}\n"
+                 f"model_key is {model_key}, model_value is {model_name}\n"
+                 f"agent_id_key is {agent_id_key}, agent id value is {agent_id}\n"
+                 f"credentials_key is {credentials_key}, credentials is {credentials}\n")
+
+def get_env(ts_tenant: str):
+    envs = dict(os.environ)
+    suffix = f"_{ts_tenant}"
+    LLM_INFO_PATH = str(LLM_CFG_ENV) + f"_{ts_tenant}"
+
+    try:
+        env_path = resolve_config_path(LLM_INFO_PATH)
+        if env_path.exists():
+            with open(env_path, 'r') as f:
+                encrypted_lines = f.readlines()
+
+            ll_vars = {}
+            for line in encrypted_lines:
+                line = line.strip()
+                if not line or line.startswith('#'):
+                    continue
+                try:
+                    key, value = line.split('=', 1)
+                    if key.endswith(suffix):
+                        ll_vars[key] = value
+                except ValueError as e:
+                    logger.warning(f"Skipping malformed line in {env_path}: {line}")
+
+            os.environ.update(ll_vars)
+            envs.update(ll_vars)
+    except FileNotFoundError:
+        pass
+    except Exception as e:
+        logger.error(f"Error decrypting environment: {str(e)}")
+    return envs
+
+def update_env_file(env_path, new_vars):
+    existing_vars = {}
+    if os.path.exists(env_path):
+        with open(env_path, 'r') as f:
+            for line in f:
+                line = line.strip()
+                if not line or line.startswith('#'):
+                    continue
+                try:
+                    key, value = line.split('=', 1)
+                    existing_vars[key] = value
+                except ValueError:
+                    continue
+
+    ll_vars = {}
+    for key, value in new_vars.items():
+        ll_vars[key] = value
+
+    existing_vars.update(ll_vars)
+    with open(env_path, 'w') as f:
+        for key, value in existing_vars.items():
+            f.write(f"{key}={value}\n")
+
+    os.chmod(env_path, 0o600)

agent_api_server/sso_service/__init__.py

@@ -0,0 +1 @@
+from .sso_service import SSOConfig, SSOService

agent_api_server/sso_service/sdk/__init__.py

@@ -0,0 +1 @@
+from .client import SSOClient, get_sso_client_id_and_secret

agent_api_server/sso_service/sdk/client.py

@@ -0,0 +1,224 @@
+import os
+import json
+import logging
+from urllib.parse import quote
+from typing import List, Optional
+
+import requests
+
+logger = logging.getLogger(__name__)
+
+CLIENT_CREDENTIALS = "client_credentials"
+CURRENT_SERVICE_NAME = "KB-Insight"
+
+
+class ClientRequest:
+    def __init__(
+            self,
+            datacenter: Optional[str] = None,
+            cluster: Optional[str] = None,
+            workspace: Optional[str] = None,
+            namespace: Optional[str] = None,
+            app_id: Optional[str] = None,
+            app_name: Optional[str] = None,
+            scopes: Optional[List[str]] = None,
+            service_name: Optional[str] = None,
+            redirect_url: Optional[str] = None,
+            **kwargs,
+    ):
+        self.datacenter = datacenter
+        self.cluster = cluster
+        self.workspace = workspace
+        self.namespace = namespace
+        self.app_id = app_id
+        self.app_name = app_name
+        self.scopes = scopes if scopes else ["Admin", "Editor", "Viewer"]
+        self.service_name = service_name
+        self.redirect_url = redirect_url
+
+    def to_dict(self):
+        return {
+            "datacenter": self.datacenter,
+            "cluster": self.cluster,
+            "workspace": self.workspace,
+            "namespace": self.namespace,
+            "appId": self.app_id,
+            "appName": self.app_name,
+            "scopes": self.scopes,
+            "serviceName": self.service_name,
+            "redirectUrl": self.redirect_url,
+        }
+
+
+class ClientInfo(ClientRequest):
+    def __init__(
+            self,
+            client_id: Optional[str] = None,
+            creation_time: Optional[int] = None,
+            last_modified_time: Optional[int] = None,
+            client_secret: Optional[str] = None,
+            client_type: Optional[str] = None,
+            **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.client_id = client_id
+        self.creation_time = creation_time
+        self.last_modified_time = last_modified_time
+        self.client_secret = client_secret
+        self.client_type = client_type
+
+    @classmethod
+    def from_dict(cls, data: dict):
+        # 过滤掉 ClientRequest 不支持的字段
+        client_request_data = {
+            "datacenter": data.get("datacenter"),
+            "cluster": data.get("cluster"),
+            "workspace": data.get("workspace"),
+            "namespace": data.get("namespace"),
+            "app_id": data.get("appId"),
+            "app_name": data.get("appName"),
+            "scopes": data.get("scopes"),
+            "service_name": data.get("serviceName"),
+            "redirect_url": data.get("redirectUrl"),
+        }
+        return cls(
+            client_id=data.get("clientId"),
+            creation_time=data.get("creationTime"),
+            last_modified_time=data.get("lastModifiedTime"),
+            client_secret=data.get("clientSecret"),
+            client_type=data.get("clientType"),
+            **client_request_data,  # 只传递 ClientRequest 支持的字段
+        )
+
+
+class SsoToken:
+    def __init__(
+            self,
+            access_token: Optional[str] = None,
+            token_type: Optional[str] = None,
+            expires_in: Optional[int] = None,
+            refresh_token: Optional[str] = None,
+    ):
+        self.access_token = access_token
+        self.token_type = token_type
+        self.expires_in = expires_in
+        self.refresh_token = refresh_token
+
+    @classmethod
+    def from_dict(cls, data: dict):
+        return cls(
+            access_token=data.get("accessToken"),
+            token_type=data.get("tokenType"),
+            expires_in=data.get("expiresIn"),
+            refresh_token=data.get("refreshToken"),
+        )
+
+
+class SSOClient:
+    def __init__(self, service_name: Optional[str] = None):
+        self.client_info = ClientInfo(
+            datacenter=os.getenv("datacenter"),
+            cluster=os.getenv("cluster"),
+            workspace=os.getenv("workspace"),
+            namespace=os.getenv("namespace"),
+            app_id=os.getenv("appID") or os.getenv("appId"),
+            app_name=service_name or CURRENT_SERVICE_NAME,
+            service_name=service_name or CURRENT_SERVICE_NAME,
+        )
+
+    def query_sso_client(self, address: str, req: ClientRequest) -> ClientInfo:
+        path = f"{address}/clients/{quote(req.app_name)}"
+        path += f"?namespace={req.namespace}&workspace={req.workspace}&cluster={req.cluster}&datacenter={req.datacenter}&serviceName={quote(req.service_name)}&appId={req.app_id}"
+
+        credential = self.new_srp_token_credential(req.service_name)
+        headers = {credential["header_name"]: credential["header_value"]}
+
+        response = requests.get(path, headers=headers)
+        response.raise_for_status()
+        return ClientInfo.from_dict(response.json())
+
+    def register_client(self, address: str, req: ClientRequest) -> ClientInfo:
+        path = f"{address}/clients"
+        body = json.dumps(req.to_dict())
+
+        logger.info({body})
+
+        credential = self.new_srp_token_credential(req.service_name)
+        headers = {
+            credential["header_name"]: credential["header_value"],
+            "Content-Type": "application/json"
+        }
+
+        response = requests.post(path, headers=headers, data=body)
+        response.raise_for_status()
+        return ClientInfo.from_dict(response.json())
+
+    def update_sso_client(self, address: str, client_id: str, req: ClientRequest) -> ClientInfo:
+        path = f"{address}/clients/{client_id}"
+        body = json.dumps(req.to_dict())
+
+        credential = self.new_srp_token_credential(req.service_name)
+        headers = {
+            credential["header_name"]: credential["header_value"],
+            "Content-Type": "application/json"
+        }
+
+        response = requests.put(path, headers=headers, data=body)
+        response.raise_for_status()
+        return ClientInfo.from_dict(response.json())
+
+    def query_sso_client_token(self, address: str) -> SsoToken:
+        path = f"{address}/oauth/token"
+        path += f"?grant_type={CLIENT_CREDENTIALS}&client_id={self.client_info.client_id}&client_secret={self.client_info.client_secret}&duration=eternal"
+
+        response = requests.post(path)
+        response.raise_for_status()
+        return SsoToken.from_dict(response.json())
+
+    def set_sso_client(self, address: str) -> None:
+        if not address.endswith("/v4.0"):
+            address += "/v4.0"
+
+        try:
+            rsp = self.query_sso_client(address, self.client_info)
+        except requests.exceptions.HTTPError as e:
+            if "Not Found" in str(e):
+                rsp = self.register_client(address, self.client_info)
+            else:
+                raise e
+
+        if rsp is None:
+            raise Exception("Failed to set SSO client")
+
+        if rsp.app_id != self.client_info.app_id:
+            try:
+                new_rsp = self.update_sso_client(address, rsp.client_id, self.client_info)
+                rsp = new_rsp
+            except requests.exceptions.HTTPError as e:
+                logger.error(f"Failed to update SSO client: {e}")
+
+        self.client_info = rsp
+        return None
+
+    @staticmethod
+    def new_srp_token_credential(service_name: str) -> dict:
+        from .credential import get_srp_token
+        return {
+            "header_name": "X-Auth-SRPToken",
+            "header_value": get_srp_token(service_name),
+        }
+
+
+def get_sso_client_id_and_secret(address: str, service: str) -> tuple:
+    sso = SSOClient(service)
+    logger.debug(f"{sso.client_info.__dict__}")
+
+    try:
+        sso.set_sso_client(address)
+    except Exception as e:
+        raise ValueError(f"get sso clientId and ClientSecret err: {e}")
+
+    logger.debug(
+        f"ServiceName: {sso.client_info.service_name}, ClientID: {sso.client_info.client_id}, ClientSecret: {sso.client_info.client_secret}")
+
+    return sso.client_info.client_id, sso.client_info.client_secret

agent_api_server/sso_service/sdk/credential.py

@@ -0,0 +1,11 @@
+import time
+from .encoding import AesEncrypt, Base64UrlSafeEncode
+
+SRP_TOKEN_KEY = "ssoisno12345678987654321"
+
+
+def get_srp_token(service_name: str) -> str:
+    cur_time = str(int(time.time()))
+    src = f"{cur_time}-{service_name}"
+    crypted = AesEncrypt(src, SRP_TOKEN_KEY)
+    return Base64UrlSafeEncode(crypted)

agent_api_server/sso_service/sdk/encoding.py

@@ -0,0 +1,22 @@
+import base64
+from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
+from cryptography.hazmat.primitives import padding
+
+
+def AesEncrypt(src: str, key: str) -> bytes:
+    key_bytes = key.encode("utf-8")
+    src_bytes = src.encode("utf-8")
+
+    # Pad the source data
+    padder = padding.PKCS7(algorithms.AES.block_size).padder()
+    padded_data = padder.update(src_bytes) + padder.finalize()
+
+    # Encrypt
+    cipher = Cipher(algorithms.AES(key_bytes), modes.ECB())
+    encryptor = cipher.encryptor()
+    return encryptor.update(padded_data) + encryptor.finalize()
+
+
+def Base64UrlSafeEncode(data: bytes) -> str:
+    encoded = base64.urlsafe_b64encode(data).decode("utf-8")
+    return encoded.rstrip("=")