ai-robot-core/ai-service/app/models/mid/tool_registry.py

223 lines
7.6 KiB
Python
Raw Normal View History

"""
Tool Registry models for Mid Platform.
[AC-IDMP-19] Tool Registry 治理模型
Reference: spec/intent-driven-mid-platform/openapi.provider.yaml
"""
import uuid
from dataclasses import dataclass, field
from datetime import datetime
from enum import Enum
from typing import Any
from sqlalchemy import JSON, Column
from sqlmodel import Field, Index, SQLModel
class ToolStatus(str, Enum):
"""工具状态"""
ENABLED = "enabled"
DISABLED = "disabled"
DEPRECATED = "deprecated"
class ToolAuthType(str, Enum):
"""工具鉴权类型"""
NONE = "none"
API_KEY = "api_key"
OAUTH = "oauth"
CUSTOM = "custom"
@dataclass
class ToolAuthConfig:
"""
[AC-IDMP-19] 工具鉴权配置
"""
auth_type: ToolAuthType = ToolAuthType.NONE
required_scopes: list[str] = field(default_factory=list)
api_key_header: str | None = None
oauth_url: str | None = None
custom_validator: str | None = None
def to_dict(self) -> dict[str, Any]:
result = {"auth_type": self.auth_type.value}
if self.required_scopes:
result["required_scopes"] = self.required_scopes
if self.api_key_header:
result["api_key_header"] = self.api_key_header
if self.oauth_url:
result["oauth_url"] = self.oauth_url
if self.custom_validator:
result["custom_validator"] = self.custom_validator
return result
@classmethod
def from_dict(cls, data: dict[str, Any]) -> "ToolAuthConfig":
return cls(
auth_type=ToolAuthType(data.get("auth_type", "none")),
required_scopes=data.get("required_scopes", []),
api_key_header=data.get("api_key_header"),
oauth_url=data.get("oauth_url"),
custom_validator=data.get("custom_validator"),
)
@dataclass
class ToolTimeoutPolicy:
"""
[AC-IDMP-19] 工具超时策略
Reference: openapi.provider.yaml - TimeoutProfile
"""
per_tool_timeout_ms: int = 30000
end_to_end_timeout_ms: int = 120000
retry_count: int = 0
retry_delay_ms: int = 100
def to_dict(self) -> dict[str, Any]:
return {
"per_tool_timeout_ms": self.per_tool_timeout_ms,
"end_to_end_timeout_ms": self.end_to_end_timeout_ms,
"retry_count": self.retry_count,
"retry_delay_ms": self.retry_delay_ms,
}
@classmethod
def from_dict(cls, data: dict[str, Any]) -> "ToolTimeoutPolicy":
return cls(
per_tool_timeout_ms=data.get("per_tool_timeout_ms", 30000),
end_to_end_timeout_ms=data.get("end_to_end_timeout_ms", 120000),
retry_count=data.get("retry_count", 0),
retry_delay_ms=data.get("retry_delay_ms", 100),
)
@dataclass
class ToolDefinition:
"""
[AC-IDMP-19] 工具定义内存模型
包含
- name: 工具名称
- type: 工具类型 (internal | mcp)
- version: 版本号
- timeout_policy: 超时策略
- auth_config: 鉴权配置
- is_enabled: 启停状态
"""
name: str
type: str = "internal"
version: str = "1.0.0"
description: str | None = None
timeout_policy: ToolTimeoutPolicy = field(default_factory=ToolTimeoutPolicy)
auth_config: ToolAuthConfig = field(default_factory=ToolAuthConfig)
is_enabled: bool = True
metadata: dict[str, Any] = field(default_factory=dict)
created_at: datetime = field(default_factory=datetime.utcnow)
updated_at: datetime = field(default_factory=datetime.utcnow)
def to_dict(self) -> dict[str, Any]:
return {
"name": self.name,
"type": self.type,
"version": self.version,
"description": self.description,
"timeout_policy": self.timeout_policy.to_dict(),
"auth_config": self.auth_config.to_dict(),
"is_enabled": self.is_enabled,
"metadata": self.metadata,
"created_at": self.created_at.isoformat(),
"updated_at": self.updated_at.isoformat(),
}
@classmethod
def from_dict(cls, data: dict[str, Any]) -> "ToolDefinition":
return cls(
name=data["name"],
type=data.get("type", "internal"),
version=data.get("version", "1.0.0"),
description=data.get("description"),
timeout_policy=ToolTimeoutPolicy.from_dict(data.get("timeout_policy", {})),
auth_config=ToolAuthConfig.from_dict(data.get("auth_config", {})),
is_enabled=data.get("is_enabled", True),
metadata=data.get("metadata", {}),
created_at=datetime.fromisoformat(data["created_at"]) if data.get("created_at") else datetime.utcnow(),
updated_at=datetime.fromisoformat(data["updated_at"]) if data.get("updated_at") else datetime.utcnow(),
)
class ToolRegistryEntity(SQLModel, table=True):
"""
[AC-IDMP-19] 工具注册表数据库实体
支持动态配置更新
"""
__tablename__ = "tool_registry"
__table_args__ = (
Index("ix_tool_registry_tenant_name", "tenant_id", "name", unique=True),
Index("ix_tool_registry_tenant_enabled", "tenant_id", "is_enabled"),
)
id: uuid.UUID = Field(default_factory=uuid.uuid4, primary_key=True)
tenant_id: str = Field(..., description="租户ID", index=True)
name: str = Field(..., description="工具名称", max_length=128)
type: str = Field(default="internal", description="工具类型: internal | mcp")
version: str = Field(default="1.0.0", description="版本号", max_length=32)
description: str | None = Field(default=None, description="工具描述")
timeout_policy: dict[str, Any] | None = Field(
default=None,
sa_column=Column("timeout_policy", JSON, nullable=True),
description="超时策略配置"
)
auth_config: dict[str, Any] | None = Field(
default=None,
sa_column=Column("auth_config", JSON, nullable=True),
description="鉴权配置"
)
is_enabled: bool = Field(default=True, description="是否启用")
metadata_: dict[str, Any] | None = Field(
default=None,
sa_column=Column("metadata", JSON, nullable=True),
description="扩展元数据"
)
created_at: datetime = Field(default_factory=datetime.utcnow, description="创建时间")
updated_at: datetime = Field(default_factory=datetime.utcnow, description="更新时间")
def to_definition(self) -> ToolDefinition:
"""转换为内存模型"""
return ToolDefinition(
name=self.name,
type=self.type,
version=self.version,
description=self.description,
timeout_policy=ToolTimeoutPolicy.from_dict(self.timeout_policy or {}),
auth_config=ToolAuthConfig.from_dict(self.auth_config or {}),
is_enabled=self.is_enabled,
metadata=self.metadata_ or {},
created_at=self.created_at,
updated_at=self.updated_at,
)
class ToolRegistryCreate(SQLModel):
"""创建工具注册请求"""
name: str = Field(..., max_length=128)
type: str = "internal"
version: str = "1.0.0"
description: str | None = None
timeout_policy: dict[str, Any] | None = None
auth_config: dict[str, Any] | None = None
is_enabled: bool = True
metadata_: dict[str, Any] | None = None
class ToolRegistryUpdate(SQLModel):
"""更新工具注册请求"""
type: str | None = None
version: str | None = None
description: str | None = None
timeout_policy: dict[str, Any] | None = None
auth_config: dict[str, Any] | None = None
is_enabled: bool | None = None
metadata_: dict[str, Any] | None = None