Coverage for agentos/protocols/__init__.py: 74%

1""" 

2AgentOS Protocols — Standardized agent communication interfaces. 

3 

4Modules: 

5- registry.py: Agent Registry with service discovery, heartbeat, load balancing 

6- grpc.py: gRPC-based A2A protocol with streaming and TLS/mTLS 

7""" 

8 

9from agentos.protocols.grpc import ( 

10 SERVICE_NAME, 

11 DefaultAgentService, 

12 GrpcAgentService, 

13 GrpcClient, 

14 GrpcClientConfig, 

15 GrpcFrameCodec, 

16 GrpcHeartbeat, 

17 GrpcServer, 

18 GrpcServerConfig, 

19 GrpcStatusCode, 

20 GrpcStreamChunk, 

21 GrpcTaskRequest, 

22 GrpcTaskResponse, 

23 TaskStatus, 

24 create_self_signed_cert, 

25) 

26from agentos.protocols.registry import AgentRegistry 

27 

28__all__ = [ 

29 # Registry 

30 "AgentRegistry", 

31 # gRPC 

32 "GrpcTaskRequest", 

33 "GrpcTaskResponse", 

34 "GrpcHeartbeat", 

35 "GrpcStreamChunk", 

36 "TaskStatus", 

37 "GrpcStatusCode", 

38 "GrpcAgentService", 

39 "DefaultAgentService", 

40 "GrpcServer", 

41 "GrpcServerConfig", 

42 "GrpcClient", 

43 "GrpcClientConfig", 

44 "GrpcFrameCodec", 

45 "SERVICE_NAME", 

46 "create_self_signed_cert", 

47] 

48 

49 

50# ── Compatibility exports (required by agentos/__init__.py) ── 

51 

52from dataclasses import dataclass, field 

53from enum import StrEnum 

54 

55 

56class CapabilityDomain(StrEnum): 

57 TEXT = "text" 

58 CODE = "code" 

59 MULTIMODAL = "multimodal" 

60 TOOL_USE = "tool_use" 

61 SAFETY = "safety" 

62 

63 

64class QoSLevel(StrEnum): 

65 BEST_EFFORT = "best_effort" 

66 AT_LEAST_ONCE = "at_least_once" 

67 EXACTLY_ONCE = "exactly_once" 

68 

69 

70@dataclass 

71class AgentCapability: 

72 name: str = "" 

73 domain: CapabilityDomain = CapabilityDomain.TEXT 

74 version: str = "1.0" 

75 description: str = "" 

76 parameters: dict = field(default_factory=dict) 

77 

78 

79@dataclass 

80class AgentContract: 

81 name: str = "" 

82 version: str = "1.0" 

83 capabilities: list = field(default_factory=list) 

84 qos: QoSLevel = QoSLevel.BEST_EFFORT 

85 endpoint: str = "" 

86 

87 def model_dump(self) -> dict: 

88 return { 

89 "name": self.name, 

90 "version": self.version, 

91 "capabilities": [c.name if hasattr(c, "name") else str(c) for c in self.capabilities], 

92 "qos": self.qos.value, 

93 "endpoint": self.endpoint, 

94 } 

95 

96 

97@dataclass 

98class MatchScore: 

99 score: float = 0.0 

100 matches: list = field(default_factory=list) 

101 

102 

103class CapabilityMatcher: 

104 """匹配 Agent 能力。""" 

105 

106 def __init__(self): 

107 pass 

108 

109 def match(self, required: list, available: list) -> MatchScore: 

110 resolved = set() 

111 for cap in required: 

112 for avail in available: 

113 cap_name = cap if isinstance(cap, str) else cap.name 

114 avail_name = avail if isinstance(avail, str) else avail.name 

115 if cap_name == avail_name: 

116 resolved.add(cap_name) 

117 score = len(resolved) / max(len(required), 1) 

118 return MatchScore(score=score, matches=list(resolved)) 

119 

120 

121class ContractRegistry: 

122 """Agent 合约注册中心。""" 

123 

124 def __init__(self): 

125 self._contracts: dict[str, AgentContract] = {} 

126 

127 def register(self, contract: AgentContract) -> None: 

128 self._contracts[contract.name] = contract 

129 

130 def get(self, name: str) -> AgentContract: 

131 return self._contracts.get(name) 

132 

133 def list_all(self) -> list[str]: 

134 return list(self._contracts.keys()) 

135 

136 

137__all__ += [ 

138 "AgentContract", 

139 "AgentCapability", 

140 "CapabilityDomain", 

141 "QoSLevel", 

142 "CapabilityMatcher", 

143 "ContractRegistry", 

144 "MatchScore", 

145]