kailash 0.3.1__py3-none-any.whl → 0.4.0__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- kailash/__init__.py +33 -1
- kailash/access_control/__init__.py +129 -0
- kailash/access_control/managers.py +461 -0
- kailash/access_control/rule_evaluators.py +467 -0
- kailash/access_control_abac.py +825 -0
- kailash/config/__init__.py +27 -0
- kailash/config/database_config.py +359 -0
- kailash/database/__init__.py +28 -0
- kailash/database/execution_pipeline.py +499 -0
- kailash/middleware/__init__.py +306 -0
- kailash/middleware/auth/__init__.py +33 -0
- kailash/middleware/auth/access_control.py +436 -0
- kailash/middleware/auth/auth_manager.py +422 -0
- kailash/middleware/auth/jwt_auth.py +477 -0
- kailash/middleware/auth/kailash_jwt_auth.py +616 -0
- kailash/middleware/communication/__init__.py +37 -0
- kailash/middleware/communication/ai_chat.py +989 -0
- kailash/middleware/communication/api_gateway.py +802 -0
- kailash/middleware/communication/events.py +470 -0
- kailash/middleware/communication/realtime.py +710 -0
- kailash/middleware/core/__init__.py +21 -0
- kailash/middleware/core/agent_ui.py +890 -0
- kailash/middleware/core/schema.py +643 -0
- kailash/middleware/core/workflows.py +396 -0
- kailash/middleware/database/__init__.py +63 -0
- kailash/middleware/database/base.py +113 -0
- kailash/middleware/database/base_models.py +525 -0
- kailash/middleware/database/enums.py +106 -0
- kailash/middleware/database/migrations.py +12 -0
- kailash/{api/database.py → middleware/database/models.py} +183 -291
- kailash/middleware/database/repositories.py +685 -0
- kailash/middleware/database/session_manager.py +19 -0
- kailash/middleware/mcp/__init__.py +38 -0
- kailash/middleware/mcp/client_integration.py +585 -0
- kailash/middleware/mcp/enhanced_server.py +576 -0
- kailash/nodes/__init__.py +25 -3
- kailash/nodes/admin/__init__.py +35 -0
- kailash/nodes/admin/audit_log.py +794 -0
- kailash/nodes/admin/permission_check.py +864 -0
- kailash/nodes/admin/role_management.py +823 -0
- kailash/nodes/admin/security_event.py +1519 -0
- kailash/nodes/admin/user_management.py +944 -0
- kailash/nodes/ai/a2a.py +24 -7
- kailash/nodes/ai/ai_providers.py +1 -0
- kailash/nodes/ai/embedding_generator.py +11 -11
- kailash/nodes/ai/intelligent_agent_orchestrator.py +99 -11
- kailash/nodes/ai/llm_agent.py +407 -2
- kailash/nodes/ai/self_organizing.py +85 -10
- kailash/nodes/api/auth.py +287 -6
- kailash/nodes/api/rest.py +151 -0
- kailash/nodes/auth/__init__.py +17 -0
- kailash/nodes/auth/directory_integration.py +1228 -0
- kailash/nodes/auth/enterprise_auth_provider.py +1328 -0
- kailash/nodes/auth/mfa.py +2338 -0
- kailash/nodes/auth/risk_assessment.py +872 -0
- kailash/nodes/auth/session_management.py +1093 -0
- kailash/nodes/auth/sso.py +1040 -0
- kailash/nodes/base.py +344 -13
- kailash/nodes/base_cycle_aware.py +4 -2
- kailash/nodes/base_with_acl.py +1 -1
- kailash/nodes/code/python.py +293 -12
- kailash/nodes/compliance/__init__.py +9 -0
- kailash/nodes/compliance/data_retention.py +1888 -0
- kailash/nodes/compliance/gdpr.py +2004 -0
- kailash/nodes/data/__init__.py +22 -2
- kailash/nodes/data/async_connection.py +469 -0
- kailash/nodes/data/async_sql.py +757 -0
- kailash/nodes/data/async_vector.py +598 -0
- kailash/nodes/data/readers.py +767 -0
- kailash/nodes/data/retrieval.py +360 -1
- kailash/nodes/data/sharepoint_graph.py +397 -21
- kailash/nodes/data/sql.py +94 -5
- kailash/nodes/data/streaming.py +68 -8
- kailash/nodes/data/vector_db.py +54 -4
- kailash/nodes/enterprise/__init__.py +13 -0
- kailash/nodes/enterprise/batch_processor.py +741 -0
- kailash/nodes/enterprise/data_lineage.py +497 -0
- kailash/nodes/logic/convergence.py +31 -9
- kailash/nodes/logic/operations.py +14 -3
- kailash/nodes/mixins/__init__.py +8 -0
- kailash/nodes/mixins/event_emitter.py +201 -0
- kailash/nodes/mixins/mcp.py +9 -4
- kailash/nodes/mixins/security.py +165 -0
- kailash/nodes/monitoring/__init__.py +7 -0
- kailash/nodes/monitoring/performance_benchmark.py +2497 -0
- kailash/nodes/rag/__init__.py +284 -0
- kailash/nodes/rag/advanced.py +1615 -0
- kailash/nodes/rag/agentic.py +773 -0
- kailash/nodes/rag/conversational.py +999 -0
- kailash/nodes/rag/evaluation.py +875 -0
- kailash/nodes/rag/federated.py +1188 -0
- kailash/nodes/rag/graph.py +721 -0
- kailash/nodes/rag/multimodal.py +671 -0
- kailash/nodes/rag/optimized.py +933 -0
- kailash/nodes/rag/privacy.py +1059 -0
- kailash/nodes/rag/query_processing.py +1335 -0
- kailash/nodes/rag/realtime.py +764 -0
- kailash/nodes/rag/registry.py +547 -0
- kailash/nodes/rag/router.py +837 -0
- kailash/nodes/rag/similarity.py +1854 -0
- kailash/nodes/rag/strategies.py +566 -0
- kailash/nodes/rag/workflows.py +575 -0
- kailash/nodes/security/__init__.py +19 -0
- kailash/nodes/security/abac_evaluator.py +1411 -0
- kailash/nodes/security/audit_log.py +91 -0
- kailash/nodes/security/behavior_analysis.py +1893 -0
- kailash/nodes/security/credential_manager.py +401 -0
- kailash/nodes/security/rotating_credentials.py +760 -0
- kailash/nodes/security/security_event.py +132 -0
- kailash/nodes/security/threat_detection.py +1103 -0
- kailash/nodes/testing/__init__.py +9 -0
- kailash/nodes/testing/credential_testing.py +499 -0
- kailash/nodes/transform/__init__.py +10 -2
- kailash/nodes/transform/chunkers.py +592 -1
- kailash/nodes/transform/processors.py +484 -14
- kailash/nodes/validation.py +321 -0
- kailash/runtime/access_controlled.py +1 -1
- kailash/runtime/async_local.py +41 -7
- kailash/runtime/docker.py +1 -1
- kailash/runtime/local.py +474 -55
- kailash/runtime/parallel.py +1 -1
- kailash/runtime/parallel_cyclic.py +1 -1
- kailash/runtime/testing.py +210 -2
- kailash/utils/migrations/__init__.py +25 -0
- kailash/utils/migrations/generator.py +433 -0
- kailash/utils/migrations/models.py +231 -0
- kailash/utils/migrations/runner.py +489 -0
- kailash/utils/secure_logging.py +342 -0
- kailash/workflow/__init__.py +16 -0
- kailash/workflow/cyclic_runner.py +3 -4
- kailash/workflow/graph.py +70 -2
- kailash/workflow/resilience.py +249 -0
- kailash/workflow/templates.py +726 -0
- {kailash-0.3.1.dist-info → kailash-0.4.0.dist-info}/METADATA +253 -20
- kailash-0.4.0.dist-info/RECORD +223 -0
- kailash/api/__init__.py +0 -17
- kailash/api/__main__.py +0 -6
- kailash/api/studio_secure.py +0 -893
- kailash/mcp/__main__.py +0 -13
- kailash/mcp/server_new.py +0 -336
- kailash/mcp/servers/__init__.py +0 -12
- kailash-0.3.1.dist-info/RECORD +0 -136
- {kailash-0.3.1.dist-info → kailash-0.4.0.dist-info}/WHEEL +0 -0
- {kailash-0.3.1.dist-info → kailash-0.4.0.dist-info}/entry_points.txt +0 -0
- {kailash-0.3.1.dist-info → kailash-0.4.0.dist-info}/licenses/LICENSE +0 -0
- {kailash-0.3.1.dist-info → kailash-0.4.0.dist-info}/top_level.txt +0 -0
@@ -0,0 +1,38 @@
|
|
1
|
+
"""
|
2
|
+
Enterprise MCP (Model Context Protocol) for Kailash Middleware
|
3
|
+
|
4
|
+
Consolidates existing Kailash MCP implementations into the middleware layer
|
5
|
+
with enhanced features for agent-frontend communication.
|
6
|
+
|
7
|
+
Features:
|
8
|
+
- Enhanced MCP server with production capabilities
|
9
|
+
- MCP client integration with middleware
|
10
|
+
- Tool discovery and registration
|
11
|
+
- Real-time MCP event streaming
|
12
|
+
- AI agent integration patterns
|
13
|
+
"""
|
14
|
+
|
15
|
+
from .client_integration import (
|
16
|
+
MCPClientConfig,
|
17
|
+
MCPServerConnection,
|
18
|
+
MiddlewareMCPClient,
|
19
|
+
)
|
20
|
+
from .enhanced_server import (
|
21
|
+
MCPResourceNode,
|
22
|
+
MCPServerConfig,
|
23
|
+
MCPToolNode,
|
24
|
+
MiddlewareMCPServer,
|
25
|
+
)
|
26
|
+
|
27
|
+
# Legacy MCP imports removed - all MCP functionality is now in middleware
|
28
|
+
|
29
|
+
__all__ = [
|
30
|
+
# Middleware MCP components
|
31
|
+
"MiddlewareMCPServer",
|
32
|
+
"MCPServerConfig",
|
33
|
+
"MCPToolNode",
|
34
|
+
"MCPResourceNode",
|
35
|
+
"MiddlewareMCPClient",
|
36
|
+
"MCPClientConfig",
|
37
|
+
"MCPServerConnection",
|
38
|
+
]
|
@@ -0,0 +1,585 @@
|
|
1
|
+
"""
|
2
|
+
Enhanced MCP Client for Kailash Middleware
|
3
|
+
|
4
|
+
Integrates existing Kailash MCP client implementations with middleware-specific
|
5
|
+
features for agent-frontend communication and real-time updates.
|
6
|
+
"""
|
7
|
+
|
8
|
+
import asyncio
|
9
|
+
import json
|
10
|
+
import logging
|
11
|
+
import uuid
|
12
|
+
from datetime import datetime, timezone
|
13
|
+
from typing import Any, Callable, Dict, List, Optional
|
14
|
+
|
15
|
+
# Import Kailash SDK components
|
16
|
+
from kailash.nodes.base import Node, NodeParameter
|
17
|
+
from kailash.nodes.code import PythonCodeNode
|
18
|
+
from kailash.runtime.local import LocalRuntime
|
19
|
+
from kailash.workflow.builder import WorkflowBuilder
|
20
|
+
|
21
|
+
# Import existing Kailash MCP components
|
22
|
+
try:
|
23
|
+
from kailash.mcp import MCPClient
|
24
|
+
|
25
|
+
_KAILASH_MCP_AVAILABLE = True
|
26
|
+
except ImportError:
|
27
|
+
_KAILASH_MCP_AVAILABLE = False
|
28
|
+
|
29
|
+
# Import middleware components
|
30
|
+
from ..communication.events import EventStream, EventType
|
31
|
+
from ..core.agent_ui import AgentUIMiddleware
|
32
|
+
|
33
|
+
logger = logging.getLogger(__name__)
|
34
|
+
|
35
|
+
|
36
|
+
class MCPClientConfig:
|
37
|
+
"""Configuration for Middleware MCP Client using Kailash patterns."""
|
38
|
+
|
39
|
+
def __init__(self):
|
40
|
+
self.name = "kailash-middleware-mcp-client"
|
41
|
+
self.version = "1.0.0"
|
42
|
+
self.description = "Enhanced MCP client built with Kailash SDK"
|
43
|
+
|
44
|
+
# Connection settings
|
45
|
+
self.connection_timeout = 30
|
46
|
+
self.request_timeout = 60
|
47
|
+
self.max_retries = 3
|
48
|
+
self.retry_delay = 1.0
|
49
|
+
|
50
|
+
# Middleware features
|
51
|
+
self.enable_events = True
|
52
|
+
self.enable_caching = True
|
53
|
+
self.cache_ttl = 300
|
54
|
+
self.enable_streaming = True
|
55
|
+
|
56
|
+
|
57
|
+
class MCPServerConnection:
|
58
|
+
"""Represents a connection to an MCP server using Kailash patterns."""
|
59
|
+
|
60
|
+
def __init__(
|
61
|
+
self,
|
62
|
+
server_name: str,
|
63
|
+
connection_config: Dict[str, Any],
|
64
|
+
client_instance: "MiddlewareMCPClient",
|
65
|
+
):
|
66
|
+
self.server_name = server_name
|
67
|
+
self.connection_config = connection_config
|
68
|
+
self.client = client_instance
|
69
|
+
self.connection_id = str(uuid.uuid4())
|
70
|
+
|
71
|
+
# Connection state
|
72
|
+
self.connected = False
|
73
|
+
self.last_connection = None
|
74
|
+
self.connection_attempts = 0
|
75
|
+
|
76
|
+
# Capabilities cache
|
77
|
+
self.server_capabilities = {}
|
78
|
+
self.available_tools = {}
|
79
|
+
self.available_resources = {}
|
80
|
+
|
81
|
+
# Kailash MCP client if available
|
82
|
+
self.mcp_client = None
|
83
|
+
if _KAILASH_MCP_AVAILABLE:
|
84
|
+
try:
|
85
|
+
self.mcp_client = MCPClient()
|
86
|
+
except Exception as e:
|
87
|
+
logger.warning(f"Could not initialize MCP client: {e}")
|
88
|
+
|
89
|
+
async def connect(self) -> bool:
|
90
|
+
"""Connect to MCP server using Kailash patterns."""
|
91
|
+
try:
|
92
|
+
self.connection_attempts += 1
|
93
|
+
|
94
|
+
# Use existing Kailash MCP client if available
|
95
|
+
if self.mcp_client:
|
96
|
+
# This would use the actual MCP client connection
|
97
|
+
# For now, simulate connection
|
98
|
+
pass
|
99
|
+
|
100
|
+
# Simulate successful connection
|
101
|
+
self.connected = True
|
102
|
+
self.last_connection = datetime.now(timezone.utc)
|
103
|
+
|
104
|
+
# Discover capabilities
|
105
|
+
await self._discover_capabilities()
|
106
|
+
|
107
|
+
# Emit middleware event
|
108
|
+
if self.client.event_stream:
|
109
|
+
await self.client._emit_client_event(
|
110
|
+
"server_connected",
|
111
|
+
{
|
112
|
+
"server_name": self.server_name,
|
113
|
+
"connection_id": self.connection_id,
|
114
|
+
"capabilities": self.server_capabilities,
|
115
|
+
},
|
116
|
+
)
|
117
|
+
|
118
|
+
logger.info(f"Connected to MCP server: {self.server_name}")
|
119
|
+
return True
|
120
|
+
|
121
|
+
except Exception as e:
|
122
|
+
logger.error(f"Failed to connect to MCP server {self.server_name}: {e}")
|
123
|
+
return False
|
124
|
+
|
125
|
+
async def _discover_capabilities(self):
|
126
|
+
"""Discover server capabilities using Kailash workflows."""
|
127
|
+
|
128
|
+
# Create capability discovery workflow
|
129
|
+
discovery_workflow = WorkflowBuilder()
|
130
|
+
|
131
|
+
discoverer = PythonCodeNode(
|
132
|
+
name="discover_capabilities",
|
133
|
+
code="""
|
134
|
+
# Discover MCP server capabilities using Kailash patterns
|
135
|
+
server_name = input_data.get('server_name')
|
136
|
+
|
137
|
+
# Simulate capability discovery
|
138
|
+
capabilities = {
|
139
|
+
'server_info': {
|
140
|
+
'name': server_name,
|
141
|
+
'version': '1.0.0',
|
142
|
+
'implementation': 'Kailash MCP Server'
|
143
|
+
},
|
144
|
+
'features': {
|
145
|
+
'tools': True,
|
146
|
+
'resources': True,
|
147
|
+
'prompts': True,
|
148
|
+
'streaming': True
|
149
|
+
}
|
150
|
+
}
|
151
|
+
|
152
|
+
# Simulate available tools
|
153
|
+
tools = {
|
154
|
+
f'{server_name}_search': {
|
155
|
+
'description': f'Search tool for {server_name}',
|
156
|
+
'parameters': {
|
157
|
+
'query': {'type': 'string', 'required': True}
|
158
|
+
}
|
159
|
+
},
|
160
|
+
f'{server_name}_process': {
|
161
|
+
'description': f'Process data with {server_name}',
|
162
|
+
'parameters': {
|
163
|
+
'data': {'type': 'object', 'required': True}
|
164
|
+
}
|
165
|
+
}
|
166
|
+
}
|
167
|
+
|
168
|
+
# Simulate available resources
|
169
|
+
resources = {
|
170
|
+
f'{server_name}://data': {
|
171
|
+
'description': f'Data resources from {server_name}',
|
172
|
+
'type': 'application/json'
|
173
|
+
}
|
174
|
+
}
|
175
|
+
|
176
|
+
result = {
|
177
|
+
'capabilities': capabilities,
|
178
|
+
'tools': tools,
|
179
|
+
'resources': resources,
|
180
|
+
'discovery_time': datetime.now().isoformat()
|
181
|
+
}
|
182
|
+
""",
|
183
|
+
)
|
184
|
+
|
185
|
+
discovery_workflow.add_node(discoverer)
|
186
|
+
workflow = discovery_workflow.build()
|
187
|
+
|
188
|
+
# Execute discovery
|
189
|
+
runtime = LocalRuntime()
|
190
|
+
results, _ = runtime.execute(
|
191
|
+
workflow, parameters={"server_name": self.server_name}
|
192
|
+
)
|
193
|
+
|
194
|
+
discovery_result = results.get("discover_capabilities", {})
|
195
|
+
|
196
|
+
if discovery_result:
|
197
|
+
self.server_capabilities = discovery_result.get("capabilities", {})
|
198
|
+
self.available_tools = discovery_result.get("tools", {})
|
199
|
+
self.available_resources = discovery_result.get("resources", {})
|
200
|
+
|
201
|
+
async def call_tool(
|
202
|
+
self, tool_name: str, arguments: Dict[str, Any], session_id: str = None
|
203
|
+
) -> Dict[str, Any]:
|
204
|
+
"""Call MCP tool using Kailash patterns."""
|
205
|
+
|
206
|
+
if not self.connected:
|
207
|
+
return {
|
208
|
+
"success": False,
|
209
|
+
"error": f"Not connected to server {self.server_name}",
|
210
|
+
}
|
211
|
+
|
212
|
+
if tool_name not in self.available_tools:
|
213
|
+
return {
|
214
|
+
"success": False,
|
215
|
+
"error": f"Tool {tool_name} not available on server {self.server_name}",
|
216
|
+
"available_tools": list(self.available_tools.keys()),
|
217
|
+
}
|
218
|
+
|
219
|
+
# Create tool execution workflow
|
220
|
+
execution_workflow = WorkflowBuilder()
|
221
|
+
|
222
|
+
executor = PythonCodeNode(
|
223
|
+
name="execute_tool",
|
224
|
+
code="""
|
225
|
+
# Execute MCP tool using Kailash patterns
|
226
|
+
tool_name = input_data.get('tool_name')
|
227
|
+
arguments = input_data.get('arguments', {})
|
228
|
+
server_name = input_data.get('server_name')
|
229
|
+
|
230
|
+
# Simulate tool execution
|
231
|
+
execution_result = {
|
232
|
+
'tool_name': tool_name,
|
233
|
+
'server_name': server_name,
|
234
|
+
'arguments': arguments,
|
235
|
+
'result': f'Executed {tool_name} on {server_name} with args: {arguments}',
|
236
|
+
'execution_time': datetime.now().isoformat(),
|
237
|
+
'success': True
|
238
|
+
}
|
239
|
+
|
240
|
+
result = {'tool_result': execution_result}
|
241
|
+
""",
|
242
|
+
)
|
243
|
+
|
244
|
+
execution_workflow.add_node(executor)
|
245
|
+
workflow = execution_workflow.build()
|
246
|
+
|
247
|
+
# Execute tool call
|
248
|
+
runtime = LocalRuntime()
|
249
|
+
results, _ = runtime.execute(
|
250
|
+
workflow,
|
251
|
+
parameters={
|
252
|
+
"tool_name": tool_name,
|
253
|
+
"arguments": arguments,
|
254
|
+
"server_name": self.server_name,
|
255
|
+
},
|
256
|
+
)
|
257
|
+
|
258
|
+
tool_result = results.get("execute_tool", {}).get("tool_result", {})
|
259
|
+
|
260
|
+
# Emit middleware event
|
261
|
+
if self.client.event_stream:
|
262
|
+
await self.client._emit_client_event(
|
263
|
+
"tool_executed",
|
264
|
+
{
|
265
|
+
"server_name": self.server_name,
|
266
|
+
"tool_name": tool_name,
|
267
|
+
"arguments": arguments,
|
268
|
+
"session_id": session_id,
|
269
|
+
"success": tool_result.get("success", False),
|
270
|
+
},
|
271
|
+
)
|
272
|
+
|
273
|
+
return {
|
274
|
+
"success": True,
|
275
|
+
"server_name": self.server_name,
|
276
|
+
"tool_result": tool_result,
|
277
|
+
}
|
278
|
+
|
279
|
+
async def get_resource(
|
280
|
+
self, resource_uri: str, session_id: str = None
|
281
|
+
) -> Dict[str, Any]:
|
282
|
+
"""Get MCP resource using Kailash patterns."""
|
283
|
+
|
284
|
+
if not self.connected:
|
285
|
+
return {
|
286
|
+
"success": False,
|
287
|
+
"error": f"Not connected to server {self.server_name}",
|
288
|
+
}
|
289
|
+
|
290
|
+
# Create resource access workflow
|
291
|
+
access_workflow = WorkflowBuilder()
|
292
|
+
|
293
|
+
accessor = PythonCodeNode(
|
294
|
+
name="access_resource",
|
295
|
+
code="""
|
296
|
+
# Access MCP resource using Kailash patterns
|
297
|
+
resource_uri = input_data.get('resource_uri')
|
298
|
+
server_name = input_data.get('server_name')
|
299
|
+
|
300
|
+
# Simulate resource access
|
301
|
+
resource_data = {
|
302
|
+
'uri': resource_uri,
|
303
|
+
'server_name': server_name,
|
304
|
+
'content': f'Resource content from {resource_uri} on {server_name}',
|
305
|
+
'content_type': 'application/json',
|
306
|
+
'access_time': datetime.now().isoformat(),
|
307
|
+
'success': True
|
308
|
+
}
|
309
|
+
|
310
|
+
result = {'resource_data': resource_data}
|
311
|
+
""",
|
312
|
+
)
|
313
|
+
|
314
|
+
access_workflow.add_node(accessor)
|
315
|
+
workflow = access_workflow.build()
|
316
|
+
|
317
|
+
# Execute resource access
|
318
|
+
runtime = LocalRuntime()
|
319
|
+
results, _ = runtime.execute(
|
320
|
+
workflow,
|
321
|
+
parameters={"resource_uri": resource_uri, "server_name": self.server_name},
|
322
|
+
)
|
323
|
+
|
324
|
+
resource_data = results.get("access_resource", {}).get("resource_data", {})
|
325
|
+
|
326
|
+
# Emit middleware event
|
327
|
+
if self.client.event_stream:
|
328
|
+
await self.client._emit_client_event(
|
329
|
+
"resource_accessed",
|
330
|
+
{
|
331
|
+
"server_name": self.server_name,
|
332
|
+
"resource_uri": resource_uri,
|
333
|
+
"session_id": session_id,
|
334
|
+
"success": resource_data.get("success", False),
|
335
|
+
},
|
336
|
+
)
|
337
|
+
|
338
|
+
return {
|
339
|
+
"success": True,
|
340
|
+
"server_name": self.server_name,
|
341
|
+
"resource_data": resource_data,
|
342
|
+
}
|
343
|
+
|
344
|
+
async def disconnect(self):
|
345
|
+
"""Disconnect from MCP server."""
|
346
|
+
if self.connected:
|
347
|
+
self.connected = False
|
348
|
+
|
349
|
+
# Emit middleware event
|
350
|
+
if self.client.event_stream:
|
351
|
+
await self.client._emit_client_event(
|
352
|
+
"server_disconnected",
|
353
|
+
{
|
354
|
+
"server_name": self.server_name,
|
355
|
+
"connection_id": self.connection_id,
|
356
|
+
},
|
357
|
+
)
|
358
|
+
|
359
|
+
logger.info(f"Disconnected from MCP server: {self.server_name}")
|
360
|
+
|
361
|
+
|
362
|
+
class MiddlewareMCPClient:
|
363
|
+
"""
|
364
|
+
Enhanced MCP Client built with Kailash SDK components.
|
365
|
+
|
366
|
+
Integrates with the middleware layer for real-time events,
|
367
|
+
session management, and agent-UI communication.
|
368
|
+
"""
|
369
|
+
|
370
|
+
def __init__(
|
371
|
+
self,
|
372
|
+
config: MCPClientConfig = None,
|
373
|
+
event_stream: EventStream = None,
|
374
|
+
agent_ui: AgentUIMiddleware = None,
|
375
|
+
):
|
376
|
+
self.config = config or MCPClientConfig()
|
377
|
+
self.event_stream = event_stream
|
378
|
+
self.agent_ui = agent_ui
|
379
|
+
|
380
|
+
# Client state
|
381
|
+
self.client_id = str(uuid.uuid4())
|
382
|
+
self.server_connections: Dict[str, MCPServerConnection] = {}
|
383
|
+
|
384
|
+
# Kailash runtime for workflows
|
385
|
+
self.runtime = LocalRuntime()
|
386
|
+
|
387
|
+
# Cache for tool/resource discovery
|
388
|
+
self._capability_cache = {}
|
389
|
+
self._cache_timestamps = {}
|
390
|
+
|
391
|
+
async def add_server(
|
392
|
+
self, server_name: str, connection_config: Dict[str, Any]
|
393
|
+
) -> bool:
|
394
|
+
"""Add MCP server connection."""
|
395
|
+
|
396
|
+
if server_name in self.server_connections:
|
397
|
+
logger.warning(f"Server {server_name} already exists")
|
398
|
+
return False
|
399
|
+
|
400
|
+
# Create server connection
|
401
|
+
connection = MCPServerConnection(server_name, connection_config, self)
|
402
|
+
|
403
|
+
# Attempt to connect
|
404
|
+
success = await connection.connect()
|
405
|
+
|
406
|
+
if success:
|
407
|
+
self.server_connections[server_name] = connection
|
408
|
+
|
409
|
+
# Emit middleware event
|
410
|
+
if self.event_stream:
|
411
|
+
await self._emit_client_event(
|
412
|
+
"server_added",
|
413
|
+
{
|
414
|
+
"server_name": server_name,
|
415
|
+
"connection_config": connection_config,
|
416
|
+
},
|
417
|
+
)
|
418
|
+
|
419
|
+
return success
|
420
|
+
|
421
|
+
async def remove_server(self, server_name: str) -> bool:
|
422
|
+
"""Remove MCP server connection."""
|
423
|
+
|
424
|
+
if server_name not in self.server_connections:
|
425
|
+
return False
|
426
|
+
|
427
|
+
connection = self.server_connections[server_name]
|
428
|
+
await connection.disconnect()
|
429
|
+
|
430
|
+
del self.server_connections[server_name]
|
431
|
+
|
432
|
+
# Emit middleware event
|
433
|
+
if self.event_stream:
|
434
|
+
await self._emit_client_event(
|
435
|
+
"server_removed", {"server_name": server_name}
|
436
|
+
)
|
437
|
+
|
438
|
+
return True
|
439
|
+
|
440
|
+
async def discover_all_capabilities(self) -> Dict[str, Any]:
|
441
|
+
"""Discover capabilities from all connected servers."""
|
442
|
+
|
443
|
+
capabilities = {}
|
444
|
+
|
445
|
+
for server_name, connection in self.server_connections.items():
|
446
|
+
if connection.connected:
|
447
|
+
capabilities[server_name] = {
|
448
|
+
"server_capabilities": connection.server_capabilities,
|
449
|
+
"available_tools": connection.available_tools,
|
450
|
+
"available_resources": connection.available_resources,
|
451
|
+
}
|
452
|
+
|
453
|
+
return capabilities
|
454
|
+
|
455
|
+
async def call_tool(
|
456
|
+
self,
|
457
|
+
server_name: str,
|
458
|
+
tool_name: str,
|
459
|
+
arguments: Dict[str, Any],
|
460
|
+
session_id: str = None,
|
461
|
+
) -> Dict[str, Any]:
|
462
|
+
"""Call tool on specific MCP server."""
|
463
|
+
|
464
|
+
if server_name not in self.server_connections:
|
465
|
+
return {
|
466
|
+
"success": False,
|
467
|
+
"error": f"Server {server_name} not found",
|
468
|
+
"available_servers": list(self.server_connections.keys()),
|
469
|
+
}
|
470
|
+
|
471
|
+
connection = self.server_connections[server_name]
|
472
|
+
return await connection.call_tool(tool_name, arguments, session_id)
|
473
|
+
|
474
|
+
async def get_resource(
|
475
|
+
self, server_name: str, resource_uri: str, session_id: str = None
|
476
|
+
) -> Dict[str, Any]:
|
477
|
+
"""Get resource from specific MCP server."""
|
478
|
+
|
479
|
+
if server_name not in self.server_connections:
|
480
|
+
return {
|
481
|
+
"success": False,
|
482
|
+
"error": f"Server {server_name} not found",
|
483
|
+
"available_servers": list(self.server_connections.keys()),
|
484
|
+
}
|
485
|
+
|
486
|
+
connection = self.server_connections[server_name]
|
487
|
+
return await connection.get_resource(resource_uri, session_id)
|
488
|
+
|
489
|
+
async def broadcast_tool_call(
|
490
|
+
self, tool_name: str, arguments: Dict[str, Any], session_id: str = None
|
491
|
+
) -> Dict[str, Dict[str, Any]]:
|
492
|
+
"""Call tool on all servers that support it."""
|
493
|
+
|
494
|
+
results = {}
|
495
|
+
|
496
|
+
for server_name, connection in self.server_connections.items():
|
497
|
+
if connection.connected and tool_name in connection.available_tools:
|
498
|
+
result = await connection.call_tool(tool_name, arguments, session_id)
|
499
|
+
results[server_name] = result
|
500
|
+
|
501
|
+
return results
|
502
|
+
|
503
|
+
async def get_client_stats(self) -> Dict[str, Any]:
|
504
|
+
"""Get MCP client statistics."""
|
505
|
+
|
506
|
+
connected_servers = sum(
|
507
|
+
1 for conn in self.server_connections.values() if conn.connected
|
508
|
+
)
|
509
|
+
|
510
|
+
total_tools = sum(
|
511
|
+
len(conn.available_tools)
|
512
|
+
for conn in self.server_connections.values()
|
513
|
+
if conn.connected
|
514
|
+
)
|
515
|
+
|
516
|
+
total_resources = sum(
|
517
|
+
len(conn.available_resources)
|
518
|
+
for conn in self.server_connections.values()
|
519
|
+
if conn.connected
|
520
|
+
)
|
521
|
+
|
522
|
+
return {
|
523
|
+
"client_info": {
|
524
|
+
"client_id": self.client_id,
|
525
|
+
"name": self.config.name,
|
526
|
+
"version": self.config.version,
|
527
|
+
},
|
528
|
+
"connections": {
|
529
|
+
"total_servers": len(self.server_connections),
|
530
|
+
"connected_servers": connected_servers,
|
531
|
+
"disconnected_servers": len(self.server_connections)
|
532
|
+
- connected_servers,
|
533
|
+
},
|
534
|
+
"capabilities": {
|
535
|
+
"total_tools": total_tools,
|
536
|
+
"total_resources": total_resources,
|
537
|
+
},
|
538
|
+
"server_details": {
|
539
|
+
server_name: {
|
540
|
+
"connected": conn.connected,
|
541
|
+
"tools_count": len(conn.available_tools),
|
542
|
+
"resources_count": len(conn.available_resources),
|
543
|
+
"last_connection": (
|
544
|
+
conn.last_connection.isoformat()
|
545
|
+
if conn.last_connection
|
546
|
+
else None
|
547
|
+
),
|
548
|
+
}
|
549
|
+
for server_name, conn in self.server_connections.items()
|
550
|
+
},
|
551
|
+
"implementation": "Kailash SDK Middleware MCP Client",
|
552
|
+
}
|
553
|
+
|
554
|
+
async def _emit_client_event(self, event_type: str, data: Dict[str, Any]):
|
555
|
+
"""Emit MCP client event to middleware event stream."""
|
556
|
+
|
557
|
+
from ..events import WorkflowEvent
|
558
|
+
|
559
|
+
event = WorkflowEvent(
|
560
|
+
type=EventType.SYSTEM_STATUS,
|
561
|
+
workflow_id="mcp_client",
|
562
|
+
data={
|
563
|
+
"mcp_client_event": event_type,
|
564
|
+
"client_id": self.client_id,
|
565
|
+
"client_name": self.config.name,
|
566
|
+
**data,
|
567
|
+
},
|
568
|
+
)
|
569
|
+
|
570
|
+
await self.event_stream.emit(event)
|
571
|
+
|
572
|
+
async def disconnect_all(self):
|
573
|
+
"""Disconnect from all MCP servers."""
|
574
|
+
|
575
|
+
for connection in self.server_connections.values():
|
576
|
+
await connection.disconnect()
|
577
|
+
|
578
|
+
# Emit shutdown event
|
579
|
+
if self.event_stream:
|
580
|
+
await self._emit_client_event(
|
581
|
+
"client_shutdown",
|
582
|
+
{"disconnected_servers": list(self.server_connections.keys())},
|
583
|
+
)
|
584
|
+
|
585
|
+
logger.info(f"MCP client {self.client_id} disconnected from all servers")
|