mcp-use 1.3.12__py3-none-any.whl → 1.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.
Potentially problematic release.
This version of mcp-use might be problematic. Click here for more details.
- mcp_use/__init__.py +1 -1
- mcp_use/adapters/.deprecated +0 -0
- mcp_use/adapters/__init__.py +18 -7
- mcp_use/adapters/base.py +12 -185
- mcp_use/adapters/langchain_adapter.py +12 -219
- mcp_use/agents/adapters/__init__.py +17 -0
- mcp_use/agents/adapters/anthropic.py +93 -0
- mcp_use/agents/adapters/base.py +316 -0
- mcp_use/agents/adapters/google.py +103 -0
- mcp_use/agents/adapters/langchain_adapter.py +212 -0
- mcp_use/agents/adapters/openai.py +111 -0
- mcp_use/agents/base.py +1 -1
- mcp_use/agents/managers/__init__.py +19 -0
- mcp_use/agents/managers/base.py +36 -0
- mcp_use/agents/managers/server_manager.py +131 -0
- mcp_use/agents/managers/tools/__init__.py +15 -0
- mcp_use/agents/managers/tools/base_tool.py +19 -0
- mcp_use/agents/managers/tools/connect_server.py +69 -0
- mcp_use/agents/managers/tools/disconnect_server.py +43 -0
- mcp_use/agents/managers/tools/get_active_server.py +29 -0
- mcp_use/agents/managers/tools/list_servers_tool.py +53 -0
- mcp_use/agents/managers/tools/search_tools.py +328 -0
- mcp_use/agents/mcpagent.py +386 -485
- mcp_use/agents/prompts/system_prompt_builder.py +1 -1
- mcp_use/agents/remote.py +15 -2
- mcp_use/auth/.deprecated +0 -0
- mcp_use/auth/__init__.py +19 -4
- mcp_use/auth/bearer.py +11 -12
- mcp_use/auth/oauth.py +11 -620
- mcp_use/auth/oauth_callback.py +16 -207
- mcp_use/client/__init__.py +1 -0
- mcp_use/client/auth/__init__.py +6 -0
- mcp_use/client/auth/bearer.py +23 -0
- mcp_use/client/auth/oauth.py +629 -0
- mcp_use/client/auth/oauth_callback.py +215 -0
- mcp_use/client/client.py +356 -0
- mcp_use/client/config.py +106 -0
- mcp_use/client/connectors/__init__.py +20 -0
- mcp_use/client/connectors/base.py +470 -0
- mcp_use/client/connectors/http.py +304 -0
- mcp_use/client/connectors/sandbox.py +332 -0
- mcp_use/client/connectors/stdio.py +109 -0
- mcp_use/client/connectors/utils.py +13 -0
- mcp_use/client/connectors/websocket.py +257 -0
- mcp_use/client/exceptions.py +31 -0
- mcp_use/client/middleware/__init__.py +50 -0
- mcp_use/client/middleware/logging.py +31 -0
- mcp_use/client/middleware/metrics.py +314 -0
- mcp_use/client/middleware/middleware.py +266 -0
- mcp_use/client/session.py +162 -0
- mcp_use/client/task_managers/__init__.py +20 -0
- mcp_use/client/task_managers/base.py +145 -0
- mcp_use/client/task_managers/sse.py +84 -0
- mcp_use/client/task_managers/stdio.py +69 -0
- mcp_use/client/task_managers/streamable_http.py +86 -0
- mcp_use/client/task_managers/websocket.py +68 -0
- mcp_use/client.py +12 -344
- mcp_use/config.py +20 -97
- mcp_use/connectors/.deprecated +0 -0
- mcp_use/connectors/__init__.py +46 -20
- mcp_use/connectors/base.py +12 -455
- mcp_use/connectors/http.py +13 -300
- mcp_use/connectors/sandbox.py +13 -306
- mcp_use/connectors/stdio.py +13 -104
- mcp_use/connectors/utils.py +15 -8
- mcp_use/connectors/websocket.py +13 -252
- mcp_use/exceptions.py +33 -18
- mcp_use/logging.py +1 -1
- mcp_use/managers/.deprecated +0 -0
- mcp_use/managers/__init__.py +56 -17
- mcp_use/managers/base.py +13 -31
- mcp_use/managers/server_manager.py +13 -119
- mcp_use/managers/tools/__init__.py +45 -15
- mcp_use/managers/tools/base_tool.py +5 -16
- mcp_use/managers/tools/connect_server.py +5 -67
- mcp_use/managers/tools/disconnect_server.py +5 -41
- mcp_use/managers/tools/get_active_server.py +5 -26
- mcp_use/managers/tools/list_servers_tool.py +5 -51
- mcp_use/managers/tools/search_tools.py +17 -321
- mcp_use/middleware/.deprecated +0 -0
- mcp_use/middleware/__init__.py +89 -50
- mcp_use/middleware/logging.py +14 -26
- mcp_use/middleware/metrics.py +30 -303
- mcp_use/middleware/middleware.py +39 -246
- mcp_use/session.py +13 -149
- mcp_use/task_managers/.deprecated +0 -0
- mcp_use/task_managers/__init__.py +48 -20
- mcp_use/task_managers/base.py +13 -140
- mcp_use/task_managers/sse.py +13 -79
- mcp_use/task_managers/stdio.py +13 -64
- mcp_use/task_managers/streamable_http.py +15 -81
- mcp_use/task_managers/websocket.py +13 -63
- mcp_use/telemetry/events.py +58 -0
- mcp_use/telemetry/telemetry.py +71 -1
- mcp_use/telemetry/utils.py +1 -1
- mcp_use/types/.deprecated +0 -0
- mcp_use/types/sandbox.py +13 -18
- {mcp_use-1.3.12.dist-info → mcp_use-1.4.0.dist-info}/METADATA +68 -43
- mcp_use-1.4.0.dist-info/RECORD +111 -0
- mcp_use/cli.py +0 -581
- mcp_use-1.3.12.dist-info/RECORD +0 -64
- mcp_use-1.3.12.dist-info/licenses/LICENSE +0 -21
- /mcp_use/{observability → agents/observability}/__init__.py +0 -0
- /mcp_use/{observability → agents/observability}/callbacks_manager.py +0 -0
- /mcp_use/{observability → agents/observability}/laminar.py +0 -0
- /mcp_use/{observability → agents/observability}/langfuse.py +0 -0
- {mcp_use-1.3.12.dist-info → mcp_use-1.4.0.dist-info}/WHEEL +0 -0
- {mcp_use-1.3.12.dist-info → mcp_use-1.4.0.dist-info}/entry_points.txt +0 -0
mcp_use/agents/mcpagent.py
CHANGED
|
@@ -3,6 +3,12 @@ MCP: Main integration module with customizable system prompt.
|
|
|
3
3
|
|
|
4
4
|
This module provides the main MCPAgent class that integrates all components
|
|
5
5
|
to provide a simple interface for using MCP tools with different LLMs.
|
|
6
|
+
|
|
7
|
+
LangChain 1.0.0 Migration:
|
|
8
|
+
- The agent uses create_agent() from langchain.agents which returns a CompiledStateGraph
|
|
9
|
+
- New methods: astream_simplified() and run_v2() leverage the built-in astream() from
|
|
10
|
+
CompiledStateGraph which handles the agent loop internally
|
|
11
|
+
- Legacy methods: stream() and run() use manual step-by-step execution for backward compatibility
|
|
6
12
|
"""
|
|
7
13
|
|
|
8
14
|
import logging
|
|
@@ -10,34 +16,30 @@ import time
|
|
|
10
16
|
from collections.abc import AsyncGenerator, AsyncIterator
|
|
11
17
|
from typing import TypeVar
|
|
12
18
|
|
|
13
|
-
from langchain.agents import
|
|
14
|
-
from langchain.agents.
|
|
15
|
-
from
|
|
16
|
-
from
|
|
17
|
-
from
|
|
18
|
-
from
|
|
19
|
-
from langchain_core.agents import AgentAction, AgentFinish
|
|
20
|
-
from langchain_core.exceptions import OutputParserException
|
|
19
|
+
from langchain.agents import create_agent
|
|
20
|
+
from langchain.agents.middleware import ModelCallLimitMiddleware
|
|
21
|
+
from langchain_core.agents import AgentAction
|
|
22
|
+
from langchain_core.globals import set_debug
|
|
23
|
+
from langchain_core.language_models import BaseLanguageModel
|
|
24
|
+
from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, SystemMessage, ToolMessage
|
|
21
25
|
from langchain_core.runnables.schema import StreamEvent
|
|
22
26
|
from langchain_core.tools import BaseTool
|
|
23
|
-
from langchain_core.utils.input import get_color_mapping
|
|
24
27
|
from pydantic import BaseModel
|
|
25
28
|
|
|
26
|
-
from mcp_use.
|
|
27
|
-
from mcp_use.
|
|
28
|
-
from mcp_use.
|
|
29
|
-
from mcp_use.telemetry.utils import extract_model_info
|
|
30
|
-
|
|
31
|
-
from ..adapters.langchain_adapter import LangChainAdapter
|
|
32
|
-
from ..logging import logger
|
|
33
|
-
from ..managers.base import BaseServerManager
|
|
34
|
-
from ..managers.server_manager import ServerManager
|
|
29
|
+
from mcp_use.agents.adapters.langchain_adapter import LangChainAdapter
|
|
30
|
+
from mcp_use.agents.managers.base import BaseServerManager
|
|
31
|
+
from mcp_use.agents.managers.server_manager import ServerManager
|
|
35
32
|
|
|
36
33
|
# Import observability manager
|
|
37
|
-
from
|
|
38
|
-
from .prompts.system_prompt_builder import create_system_message
|
|
39
|
-
from .prompts.templates import DEFAULT_SYSTEM_PROMPT_TEMPLATE, SERVER_MANAGER_SYSTEM_PROMPT_TEMPLATE
|
|
40
|
-
from .remote import RemoteAgent
|
|
34
|
+
from mcp_use.agents.observability import ObservabilityManager
|
|
35
|
+
from mcp_use.agents.prompts.system_prompt_builder import create_system_message
|
|
36
|
+
from mcp_use.agents.prompts.templates import DEFAULT_SYSTEM_PROMPT_TEMPLATE, SERVER_MANAGER_SYSTEM_PROMPT_TEMPLATE
|
|
37
|
+
from mcp_use.agents.remote import RemoteAgent
|
|
38
|
+
from mcp_use.client import MCPClient
|
|
39
|
+
from mcp_use.client.connectors.base import BaseConnector
|
|
40
|
+
from mcp_use.logging import logger
|
|
41
|
+
from mcp_use.telemetry.telemetry import Telemetry, telemetry
|
|
42
|
+
from mcp_use.telemetry.utils import extract_model_info
|
|
41
43
|
|
|
42
44
|
set_debug(logger.level == logging.DEBUG)
|
|
43
45
|
|
|
@@ -151,7 +153,7 @@ class MCPAgent:
|
|
|
151
153
|
self.server_manager = ServerManager(self.client, self.adapter)
|
|
152
154
|
|
|
153
155
|
# State tracking - initialize _tools as empty list
|
|
154
|
-
self._agent_executor
|
|
156
|
+
self._agent_executor = None
|
|
155
157
|
self._system_message: SystemMessage | None = None
|
|
156
158
|
self._tools: list[BaseTool] = []
|
|
157
159
|
|
|
@@ -186,7 +188,8 @@ class MCPAgent:
|
|
|
186
188
|
logger.info(f"✅ Created {len(self._sessions)} new sessions")
|
|
187
189
|
|
|
188
190
|
# Create LangChain tools directly from the client using the adapter
|
|
189
|
-
|
|
191
|
+
await self.adapter.create_all(self.client)
|
|
192
|
+
self._tools = self.adapter.tools + self.adapter.resources + self.adapter.prompts
|
|
190
193
|
logger.info(f"🛠️ Created {len(self._tools)} LangChain tools from client")
|
|
191
194
|
else:
|
|
192
195
|
# Using direct connector - only establish connection
|
|
@@ -198,7 +201,10 @@ class MCPAgent:
|
|
|
198
201
|
await connector.connect()
|
|
199
202
|
|
|
200
203
|
# Create LangChain tools using the adapter with connectors
|
|
201
|
-
|
|
204
|
+
await self.adapter._create_tools_from_connectors(connectors_to_use)
|
|
205
|
+
await self.adapter._create_resources_from_connectors(connectors_to_use)
|
|
206
|
+
await self.adapter._create_prompts_from_connectors(connectors_to_use)
|
|
207
|
+
self._tools = self.adapter.tools + self.adapter.resources + self.adapter.prompts
|
|
202
208
|
logger.info(f"🛠️ Created {len(self._tools)} LangChain tools from connectors")
|
|
203
209
|
|
|
204
210
|
# Get all tools for system message generation
|
|
@@ -275,7 +281,7 @@ class MCPAgent:
|
|
|
275
281
|
msg for msg in self._conversation_history if not isinstance(msg, SystemMessage)
|
|
276
282
|
]
|
|
277
283
|
|
|
278
|
-
def _create_agent(self)
|
|
284
|
+
def _create_agent(self):
|
|
279
285
|
"""Create the LangChain agent with the configured system message.
|
|
280
286
|
|
|
281
287
|
Returns:
|
|
@@ -287,42 +293,23 @@ class MCPAgent:
|
|
|
287
293
|
if self._system_message:
|
|
288
294
|
system_content = self._system_message.content
|
|
289
295
|
|
|
290
|
-
if self.memory_enabled:
|
|
291
|
-
# Query already in chat_history — don't re-inject it
|
|
292
|
-
prompt = ChatPromptTemplate.from_messages(
|
|
293
|
-
[
|
|
294
|
-
("system", system_content),
|
|
295
|
-
MessagesPlaceholder(variable_name="chat_history"),
|
|
296
|
-
("human", "{input}"),
|
|
297
|
-
MessagesPlaceholder(variable_name="agent_scratchpad"),
|
|
298
|
-
]
|
|
299
|
-
)
|
|
300
|
-
else:
|
|
301
|
-
# No memory — inject input directly
|
|
302
|
-
prompt = ChatPromptTemplate.from_messages(
|
|
303
|
-
[
|
|
304
|
-
("system", system_content),
|
|
305
|
-
("human", "{input}"),
|
|
306
|
-
MessagesPlaceholder(variable_name="agent_scratchpad"),
|
|
307
|
-
]
|
|
308
|
-
)
|
|
309
|
-
|
|
310
296
|
tool_names = [tool.name for tool in self._tools]
|
|
311
297
|
logger.info(f"🧠 Agent ready with tools: {', '.join(tool_names)}")
|
|
312
298
|
|
|
313
|
-
#
|
|
314
|
-
|
|
299
|
+
# Create middleware to enforce max_steps
|
|
300
|
+
# ModelCallLimitMiddleware limits the number of model calls, which corresponds to agent steps
|
|
301
|
+
middleware = [ModelCallLimitMiddleware(run_limit=self.max_steps)]
|
|
302
|
+
|
|
303
|
+
# Use the standard create_agent with middleware
|
|
304
|
+
agent = create_agent(
|
|
305
|
+
model=self.llm, tools=self._tools, system_prompt=system_content, middleware=middleware, debug=self.verbose
|
|
306
|
+
)
|
|
315
307
|
|
|
316
|
-
|
|
317
|
-
|
|
318
|
-
|
|
319
|
-
tools=self._tools,
|
|
320
|
-
max_iterations=self.max_steps,
|
|
321
|
-
verbose=self.verbose,
|
|
322
|
-
callbacks=self.callbacks,
|
|
308
|
+
logger.debug(
|
|
309
|
+
f"Created agent with max_steps={self.max_steps} (via ModelCallLimitMiddleware) "
|
|
310
|
+
f"and {len(self.callbacks)} callbacks"
|
|
323
311
|
)
|
|
324
|
-
|
|
325
|
-
return executor
|
|
312
|
+
return agent
|
|
326
313
|
|
|
327
314
|
def get_conversation_history(self) -> list[BaseMessage]:
|
|
328
315
|
"""Get the current conversation history.
|
|
@@ -394,16 +381,14 @@ class MCPAgent:
|
|
|
394
381
|
|
|
395
382
|
async def _consume_and_return(
|
|
396
383
|
self,
|
|
397
|
-
generator: AsyncGenerator[
|
|
384
|
+
generator: AsyncGenerator[str | T, None],
|
|
398
385
|
) -> tuple[str | T, int]:
|
|
399
|
-
"""Consume the generator and return the final result.
|
|
386
|
+
"""Consume the stream generator and return the final result.
|
|
400
387
|
|
|
401
|
-
This
|
|
402
|
-
In Python, async generators cannot return values directly, so we expect
|
|
403
|
-
the final result to be yielded as a special marker.
|
|
388
|
+
This is used by the run() method with the astream implementation.
|
|
404
389
|
|
|
405
390
|
Args:
|
|
406
|
-
generator: The async generator
|
|
391
|
+
generator: The async generator from astream.
|
|
407
392
|
|
|
408
393
|
Returns:
|
|
409
394
|
A tuple of (final_result, steps_taken). final_result can be a string
|
|
@@ -412,415 +397,13 @@ class MCPAgent:
|
|
|
412
397
|
final_result = ""
|
|
413
398
|
steps_taken = 0
|
|
414
399
|
async for item in generator:
|
|
415
|
-
#
|
|
416
|
-
|
|
417
|
-
|
|
418
|
-
|
|
419
|
-
# If it's not a tuple, it might be structured output (Pydantic model)
|
|
420
|
-
elif not isinstance(item, tuple):
|
|
421
|
-
final_result = item
|
|
422
|
-
break
|
|
423
|
-
# Otherwise it's a step tuple, just consume it
|
|
424
|
-
else:
|
|
425
|
-
steps_taken += 1
|
|
400
|
+
# The last item yielded is always the final result
|
|
401
|
+
final_result = item
|
|
402
|
+
# Count steps as the number of tools used during execution
|
|
403
|
+
steps_taken = len(self.tools_used_names)
|
|
426
404
|
return final_result, steps_taken
|
|
427
405
|
|
|
428
|
-
|
|
429
|
-
self,
|
|
430
|
-
query: str,
|
|
431
|
-
max_steps: int | None = None,
|
|
432
|
-
manage_connector: bool = True,
|
|
433
|
-
external_history: list[BaseMessage] | None = None,
|
|
434
|
-
track_execution: bool = True,
|
|
435
|
-
output_schema: type[T] | None = None,
|
|
436
|
-
) -> AsyncGenerator[tuple[AgentAction, str] | str | T, None]:
|
|
437
|
-
"""Run the agent and yield intermediate steps as an async generator.
|
|
438
|
-
|
|
439
|
-
Args:
|
|
440
|
-
query: The query to run.
|
|
441
|
-
max_steps: Optional maximum number of steps to take.
|
|
442
|
-
manage_connector: Whether to handle the connector lifecycle internally.
|
|
443
|
-
external_history: Optional external history to use instead of the
|
|
444
|
-
internal conversation history.
|
|
445
|
-
track_execution: Whether to track execution for telemetry.
|
|
446
|
-
output_schema: Optional Pydantic BaseModel class for structured output.
|
|
447
|
-
If provided, the agent will attempt structured output at finish points
|
|
448
|
-
and continue execution if required information is missing.
|
|
449
|
-
|
|
450
|
-
Yields:
|
|
451
|
-
Intermediate steps as (AgentAction, str) tuples, followed by the final result.
|
|
452
|
-
If output_schema is provided, yields structured output as instance of the schema.
|
|
453
|
-
"""
|
|
454
|
-
# Delegate to remote agent if in remote mode
|
|
455
|
-
if self._is_remote and self._remote_agent:
|
|
456
|
-
async for item in self._remote_agent.stream(
|
|
457
|
-
query, max_steps, manage_connector, external_history, track_execution, output_schema
|
|
458
|
-
):
|
|
459
|
-
yield item
|
|
460
|
-
return
|
|
461
|
-
|
|
462
|
-
result = ""
|
|
463
|
-
initialized_here = False
|
|
464
|
-
start_time = time.time()
|
|
465
|
-
steps_taken = 0
|
|
466
|
-
success = False
|
|
467
|
-
|
|
468
|
-
# Schema-aware setup for structured output
|
|
469
|
-
structured_llm = None
|
|
470
|
-
schema_description = ""
|
|
471
|
-
if output_schema:
|
|
472
|
-
query = self._enhance_query_with_schema(query, output_schema)
|
|
473
|
-
structured_llm = self.llm.with_structured_output(output_schema)
|
|
474
|
-
# Get schema description for feedback
|
|
475
|
-
schema_fields = []
|
|
476
|
-
try:
|
|
477
|
-
for field_name, field_info in output_schema.model_fields.items():
|
|
478
|
-
description = getattr(field_info, "description", "") or field_name
|
|
479
|
-
required = not hasattr(field_info, "default") or field_info.default is None
|
|
480
|
-
schema_fields.append(f"- {field_name}: {description} {'(required)' if required else '(optional)'}")
|
|
481
|
-
|
|
482
|
-
schema_description = "\n".join(schema_fields)
|
|
483
|
-
except Exception as e:
|
|
484
|
-
logger.warning(f"Could not extract schema details: {e}")
|
|
485
|
-
schema_description = f"Schema: {output_schema.__name__}"
|
|
486
|
-
|
|
487
|
-
try:
|
|
488
|
-
# Initialize if needed
|
|
489
|
-
if manage_connector and not self._initialized:
|
|
490
|
-
await self.initialize()
|
|
491
|
-
initialized_here = True
|
|
492
|
-
elif not self._initialized and self.auto_initialize:
|
|
493
|
-
await self.initialize()
|
|
494
|
-
initialized_here = True
|
|
495
|
-
|
|
496
|
-
# Check if initialization succeeded
|
|
497
|
-
if not self._agent_executor:
|
|
498
|
-
raise RuntimeError("MCP agent failed to initialize")
|
|
499
|
-
|
|
500
|
-
steps = max_steps or self.max_steps
|
|
501
|
-
if self._agent_executor:
|
|
502
|
-
self._agent_executor.max_iterations = steps
|
|
503
|
-
|
|
504
|
-
display_query = query[:50].replace("\n", " ") + "..." if len(query) > 50 else query.replace("\n", " ")
|
|
505
|
-
logger.info(f"💬 Received query: '{display_query}'")
|
|
506
|
-
|
|
507
|
-
# Use the provided history or the internal history
|
|
508
|
-
history_to_use = external_history if external_history is not None else self._conversation_history
|
|
509
|
-
|
|
510
|
-
# Convert messages to format expected by LangChain agent input
|
|
511
|
-
# Exclude the main system message as it's part of the agent's prompt
|
|
512
|
-
langchain_history = []
|
|
513
|
-
for msg in history_to_use:
|
|
514
|
-
if isinstance(msg, HumanMessage):
|
|
515
|
-
langchain_history.append(msg)
|
|
516
|
-
elif isinstance(msg, AIMessage):
|
|
517
|
-
langchain_history.append(msg)
|
|
518
|
-
|
|
519
|
-
intermediate_steps: list[tuple[AgentAction, str]] = []
|
|
520
|
-
inputs = {"input": query, "chat_history": langchain_history}
|
|
521
|
-
|
|
522
|
-
# Construct a mapping of tool name to tool for easy lookup
|
|
523
|
-
name_to_tool_map = {tool.name: tool for tool in self._tools}
|
|
524
|
-
color_mapping = get_color_mapping([tool.name for tool in self._tools], excluded_colors=["green", "red"])
|
|
525
|
-
|
|
526
|
-
logger.info(f"🏁 Starting agent execution with max_steps={steps}")
|
|
527
|
-
|
|
528
|
-
# Track whether agent finished successfully vs reached max iterations
|
|
529
|
-
agent_finished_successfully = False
|
|
530
|
-
result = None
|
|
531
|
-
|
|
532
|
-
# Create a run manager with our callbacks if we have any - ONCE for the entire execution
|
|
533
|
-
run_manager = None
|
|
534
|
-
if self.callbacks:
|
|
535
|
-
# Create an async callback manager with our callbacks
|
|
536
|
-
from langchain_core.callbacks.manager import AsyncCallbackManager
|
|
537
|
-
|
|
538
|
-
callback_manager = AsyncCallbackManager.configure(
|
|
539
|
-
inheritable_callbacks=self.callbacks,
|
|
540
|
-
local_callbacks=self.callbacks,
|
|
541
|
-
)
|
|
542
|
-
# Create a run manager for this chain execution
|
|
543
|
-
run_manager = await callback_manager.on_chain_start(
|
|
544
|
-
{"name": "MCPAgent (mcp-use)"},
|
|
545
|
-
inputs,
|
|
546
|
-
)
|
|
547
|
-
|
|
548
|
-
for step_num in range(steps):
|
|
549
|
-
steps_taken = step_num + 1
|
|
550
|
-
# --- Check for tool updates if using server manager ---
|
|
551
|
-
if self.use_server_manager and self.server_manager:
|
|
552
|
-
current_tools = self.server_manager.tools
|
|
553
|
-
current_tool_names = {tool.name for tool in current_tools}
|
|
554
|
-
existing_tool_names = {tool.name for tool in self._tools}
|
|
555
|
-
|
|
556
|
-
if current_tool_names != existing_tool_names:
|
|
557
|
-
logger.info(
|
|
558
|
-
f"🔄 Tools changed before step {step_num + 1}, updating agent."
|
|
559
|
-
f"New tools: {', '.join(current_tool_names)}"
|
|
560
|
-
)
|
|
561
|
-
self._tools = current_tools
|
|
562
|
-
# Regenerate system message with ALL current tools
|
|
563
|
-
await self._create_system_message_from_tools(self._tools)
|
|
564
|
-
# Recreate the agent executor with the new tools and system message
|
|
565
|
-
self._agent_executor = self._create_agent()
|
|
566
|
-
self._agent_executor.max_iterations = steps
|
|
567
|
-
# Update maps for this iteration
|
|
568
|
-
name_to_tool_map = {tool.name: tool for tool in self._tools}
|
|
569
|
-
color_mapping = get_color_mapping(
|
|
570
|
-
[tool.name for tool in self._tools], excluded_colors=["green", "red"]
|
|
571
|
-
)
|
|
572
|
-
|
|
573
|
-
logger.info(f"👣 Step {step_num + 1}/{steps}")
|
|
574
|
-
|
|
575
|
-
# --- Plan and execute the next step ---
|
|
576
|
-
try:
|
|
577
|
-
retry_count = 0
|
|
578
|
-
next_step_output = None
|
|
579
|
-
|
|
580
|
-
while retry_count <= self.max_retries_per_step:
|
|
581
|
-
try:
|
|
582
|
-
# Use the internal _atake_next_step which handles planning and execution
|
|
583
|
-
# This requires providing the necessary context like maps and intermediate steps
|
|
584
|
-
next_step_output = await self._agent_executor._atake_next_step(
|
|
585
|
-
name_to_tool_map=name_to_tool_map,
|
|
586
|
-
color_mapping=color_mapping,
|
|
587
|
-
inputs=inputs,
|
|
588
|
-
intermediate_steps=intermediate_steps,
|
|
589
|
-
run_manager=run_manager,
|
|
590
|
-
)
|
|
591
|
-
|
|
592
|
-
# If we get here, the step succeeded, break out of retry loop
|
|
593
|
-
break
|
|
594
|
-
|
|
595
|
-
except Exception as e:
|
|
596
|
-
if not self.retry_on_error or retry_count >= self.max_retries_per_step:
|
|
597
|
-
logger.error(f"❌ Validation error during step {step_num + 1}: {e}")
|
|
598
|
-
result = f"Agent stopped due to a validation error: {str(e)}"
|
|
599
|
-
success = False
|
|
600
|
-
yield result
|
|
601
|
-
return
|
|
602
|
-
|
|
603
|
-
retry_count += 1
|
|
604
|
-
logger.warning(
|
|
605
|
-
f"⚠️ Validation error, retrying ({retry_count}/{self.max_retries_per_step}): {e}"
|
|
606
|
-
)
|
|
607
|
-
|
|
608
|
-
# Create concise feedback for the LLM about the validation error
|
|
609
|
-
error_message = f"Error: {str(e)}"
|
|
610
|
-
inputs["input"] = error_message
|
|
611
|
-
|
|
612
|
-
# Continue to next iteration of retry loop
|
|
613
|
-
continue
|
|
614
|
-
|
|
615
|
-
# Process the output
|
|
616
|
-
if isinstance(next_step_output, AgentFinish):
|
|
617
|
-
logger.info(f"✅ Agent finished at step {step_num + 1}")
|
|
618
|
-
agent_finished_successfully = True
|
|
619
|
-
output_value = next_step_output.return_values.get("output", "No output generated")
|
|
620
|
-
result = self._normalize_output(output_value)
|
|
621
|
-
# End the chain if we have a run manager
|
|
622
|
-
if run_manager:
|
|
623
|
-
await run_manager.on_chain_end({"output": result})
|
|
624
|
-
|
|
625
|
-
# If structured output is requested, attempt to create it
|
|
626
|
-
if output_schema and structured_llm:
|
|
627
|
-
try:
|
|
628
|
-
logger.info("🔧 Attempting structured output...")
|
|
629
|
-
structured_result = await self._attempt_structured_output(
|
|
630
|
-
result, structured_llm, output_schema, schema_description
|
|
631
|
-
)
|
|
632
|
-
|
|
633
|
-
# Add the final response to conversation history if memory is enabled
|
|
634
|
-
if self.memory_enabled:
|
|
635
|
-
self.add_to_history(AIMessage(content=f"Structured result: {structured_result}"))
|
|
636
|
-
|
|
637
|
-
logger.info("✅ Structured output successful")
|
|
638
|
-
success = True
|
|
639
|
-
yield structured_result
|
|
640
|
-
return
|
|
641
|
-
|
|
642
|
-
except Exception as e:
|
|
643
|
-
logger.warning(f"⚠️ Structured output failed: {e}")
|
|
644
|
-
# Continue execution to gather missing information
|
|
645
|
-
missing_info_prompt = f"""
|
|
646
|
-
The current result cannot be formatted into the required structure.
|
|
647
|
-
Error: {str(e)}
|
|
648
|
-
|
|
649
|
-
Current information: {result}
|
|
650
|
-
|
|
651
|
-
Please continue working to gather the missing information needed for:
|
|
652
|
-
{schema_description}
|
|
653
|
-
|
|
654
|
-
Focus on finding the specific missing details.
|
|
655
|
-
"""
|
|
656
|
-
|
|
657
|
-
# Add this as feedback and continue the loop
|
|
658
|
-
inputs["input"] = missing_info_prompt
|
|
659
|
-
if self.memory_enabled:
|
|
660
|
-
self.add_to_history(HumanMessage(content=missing_info_prompt))
|
|
661
|
-
|
|
662
|
-
logger.info("🔄 Continuing execution to gather missing information...")
|
|
663
|
-
continue
|
|
664
|
-
else:
|
|
665
|
-
# Regular execution without structured output
|
|
666
|
-
break
|
|
667
|
-
|
|
668
|
-
# If it's actions/steps, add to intermediate steps and yield them
|
|
669
|
-
intermediate_steps.extend(next_step_output)
|
|
670
|
-
|
|
671
|
-
# Yield each step and track tool usage
|
|
672
|
-
for agent_step in next_step_output:
|
|
673
|
-
yield agent_step
|
|
674
|
-
action, observation = agent_step
|
|
675
|
-
reasoning = getattr(action, "log", "")
|
|
676
|
-
if reasoning:
|
|
677
|
-
reasoning_str = reasoning.replace("\n", " ")
|
|
678
|
-
if len(reasoning_str) > 300:
|
|
679
|
-
reasoning_str = reasoning_str[:297] + "..."
|
|
680
|
-
logger.info(f"💭 Reasoning: {reasoning_str}")
|
|
681
|
-
tool_name = action.tool
|
|
682
|
-
self.tools_used_names.append(tool_name)
|
|
683
|
-
tool_input_str = str(action.tool_input)
|
|
684
|
-
# Truncate long inputs for readability
|
|
685
|
-
if len(tool_input_str) > 100:
|
|
686
|
-
tool_input_str = tool_input_str[:97] + "..."
|
|
687
|
-
logger.info(f"🔧 Tool call: {tool_name} with input: {tool_input_str}")
|
|
688
|
-
# Truncate long outputs for readability
|
|
689
|
-
observation_str = str(observation)
|
|
690
|
-
if len(observation_str) > 100:
|
|
691
|
-
observation_str = observation_str[:97] + "..."
|
|
692
|
-
observation_str = observation_str.replace("\n", " ")
|
|
693
|
-
logger.info(f"📄 Tool result: {observation_str}")
|
|
694
|
-
|
|
695
|
-
# Check for return_direct on the last action taken
|
|
696
|
-
if len(next_step_output) > 0:
|
|
697
|
-
last_step: tuple[AgentAction, str] = next_step_output[-1]
|
|
698
|
-
tool_return = self._agent_executor._get_tool_return(last_step)
|
|
699
|
-
if tool_return is not None:
|
|
700
|
-
logger.info(f"🏆 Tool returned directly at step {step_num + 1}")
|
|
701
|
-
agent_finished_successfully = True
|
|
702
|
-
result = tool_return.return_values.get("output", "No output generated")
|
|
703
|
-
result = self._normalize_output(result)
|
|
704
|
-
break
|
|
705
|
-
|
|
706
|
-
except OutputParserException as e:
|
|
707
|
-
logger.error(f"❌ Output parsing error during step {step_num + 1}: {e}")
|
|
708
|
-
result = f"Agent stopped due to a parsing error: {str(e)}"
|
|
709
|
-
if run_manager:
|
|
710
|
-
await run_manager.on_chain_error(e)
|
|
711
|
-
break
|
|
712
|
-
except Exception as e:
|
|
713
|
-
logger.error(f"❌ Error during agent execution step {step_num + 1}: {e}")
|
|
714
|
-
import traceback
|
|
715
|
-
|
|
716
|
-
traceback.print_exc()
|
|
717
|
-
# End the chain with error if we have a run manager
|
|
718
|
-
if run_manager:
|
|
719
|
-
await run_manager.on_chain_error(e)
|
|
720
|
-
result = f"Agent stopped due to an error: {str(e)}"
|
|
721
|
-
break
|
|
722
|
-
|
|
723
|
-
# --- Loop finished ---
|
|
724
|
-
if not result:
|
|
725
|
-
if agent_finished_successfully:
|
|
726
|
-
# Agent finished successfully but returned empty output
|
|
727
|
-
result = "Agent completed the task successfully."
|
|
728
|
-
logger.info("✅ Agent finished successfully with empty output")
|
|
729
|
-
else:
|
|
730
|
-
# Agent actually reached max iterations
|
|
731
|
-
logger.warning(f"⚠️ Agent stopped after reaching max iterations ({steps})")
|
|
732
|
-
result = f"Agent stopped after reaching the maximum number of steps ({steps})."
|
|
733
|
-
if run_manager:
|
|
734
|
-
await run_manager.on_chain_end({"output": result})
|
|
735
|
-
|
|
736
|
-
# If structured output was requested but not achieved, attempt one final time
|
|
737
|
-
if output_schema and structured_llm and not success:
|
|
738
|
-
try:
|
|
739
|
-
logger.info("🔧 Final attempt at structured output...")
|
|
740
|
-
structured_result = await self._attempt_structured_output(
|
|
741
|
-
result, structured_llm, output_schema, schema_description
|
|
742
|
-
)
|
|
743
|
-
|
|
744
|
-
# Add the final response to conversation history if memory is enabled
|
|
745
|
-
if self.memory_enabled:
|
|
746
|
-
self.add_to_history(AIMessage(content=f"Structured result: {structured_result}"))
|
|
747
|
-
|
|
748
|
-
logger.info("✅ Final structured output successful")
|
|
749
|
-
success = True
|
|
750
|
-
yield structured_result
|
|
751
|
-
return
|
|
752
|
-
|
|
753
|
-
except Exception as e:
|
|
754
|
-
logger.error(f"❌ Final structured output attempt failed: {e}")
|
|
755
|
-
raise RuntimeError(f"Failed to generate structured output after {steps} steps: {str(e)}") from e
|
|
756
|
-
|
|
757
|
-
if self.memory_enabled:
|
|
758
|
-
self.add_to_history(HumanMessage(content=query))
|
|
759
|
-
|
|
760
|
-
if self.memory_enabled and not output_schema:
|
|
761
|
-
self.add_to_history(AIMessage(content=self._normalize_output(result)))
|
|
762
|
-
|
|
763
|
-
logger.info(f"🎉 Agent execution complete in {time.time() - start_time} seconds")
|
|
764
|
-
if not success:
|
|
765
|
-
success = True
|
|
766
|
-
|
|
767
|
-
# Yield the final result (only for non-structured output)
|
|
768
|
-
if not output_schema:
|
|
769
|
-
yield result
|
|
770
|
-
|
|
771
|
-
except Exception as e:
|
|
772
|
-
logger.error(f"❌ Error running query: {e}")
|
|
773
|
-
if initialized_here and manage_connector:
|
|
774
|
-
logger.info("🧹 Cleaning up resources after initialization error in stream")
|
|
775
|
-
await self.close()
|
|
776
|
-
raise
|
|
777
|
-
|
|
778
|
-
finally:
|
|
779
|
-
# Track comprehensive execution data
|
|
780
|
-
execution_time_ms = int((time.time() - start_time) * 1000)
|
|
781
|
-
|
|
782
|
-
server_count = 0
|
|
783
|
-
if self.client:
|
|
784
|
-
server_count = len(self.client.get_all_active_sessions())
|
|
785
|
-
elif self.connectors:
|
|
786
|
-
server_count = len(self.connectors)
|
|
787
|
-
|
|
788
|
-
conversation_history_length = len(self._conversation_history) if self.memory_enabled else 0
|
|
789
|
-
|
|
790
|
-
# Safely access _tools in case initialization failed
|
|
791
|
-
tools_available = getattr(self, "_tools", [])
|
|
792
|
-
|
|
793
|
-
if track_execution:
|
|
794
|
-
self.telemetry.track_agent_execution(
|
|
795
|
-
execution_method="stream",
|
|
796
|
-
query=query,
|
|
797
|
-
success=success,
|
|
798
|
-
model_provider=self._model_provider,
|
|
799
|
-
model_name=self._model_name,
|
|
800
|
-
server_count=server_count,
|
|
801
|
-
server_identifiers=[connector.public_identifier for connector in self.connectors],
|
|
802
|
-
total_tools_available=len(tools_available),
|
|
803
|
-
tools_available_names=[tool.name for tool in tools_available],
|
|
804
|
-
max_steps_configured=self.max_steps,
|
|
805
|
-
memory_enabled=self.memory_enabled,
|
|
806
|
-
use_server_manager=self.use_server_manager,
|
|
807
|
-
max_steps_used=max_steps,
|
|
808
|
-
manage_connector=manage_connector,
|
|
809
|
-
external_history_used=external_history is not None,
|
|
810
|
-
steps_taken=steps_taken,
|
|
811
|
-
tools_used_count=len(self.tools_used_names),
|
|
812
|
-
tools_used_names=self.tools_used_names,
|
|
813
|
-
response=result,
|
|
814
|
-
execution_time_ms=execution_time_ms,
|
|
815
|
-
error_type=None if success else "execution_error",
|
|
816
|
-
conversation_history_length=conversation_history_length,
|
|
817
|
-
)
|
|
818
|
-
|
|
819
|
-
# Clean up if necessary (e.g., if not using client-managed sessions)
|
|
820
|
-
if manage_connector and not self.client and initialized_here:
|
|
821
|
-
logger.info("🧹 Closing agent after stream completion")
|
|
822
|
-
await self.close()
|
|
823
|
-
|
|
406
|
+
@telemetry("agent_run")
|
|
824
407
|
async def run(
|
|
825
408
|
self,
|
|
826
409
|
query: str,
|
|
@@ -829,23 +412,15 @@ class MCPAgent:
|
|
|
829
412
|
external_history: list[BaseMessage] | None = None,
|
|
830
413
|
output_schema: type[T] | None = None,
|
|
831
414
|
) -> str | T:
|
|
832
|
-
"""Run a query using
|
|
833
|
-
|
|
834
|
-
This method uses the streaming implementation internally and returns
|
|
835
|
-
the final result after consuming all intermediate steps. If output_schema
|
|
836
|
-
is provided, the agent will be schema-aware and return structured output.
|
|
415
|
+
"""Run a query using LangChain 1.0.0's agent and return the final result.
|
|
837
416
|
|
|
838
417
|
Args:
|
|
839
418
|
query: The query to run.
|
|
840
419
|
max_steps: Optional maximum number of steps to take.
|
|
841
420
|
manage_connector: Whether to handle the connector lifecycle internally.
|
|
842
|
-
If True, this method will connect, initialize, and disconnect from
|
|
843
|
-
the connector automatically. If False, the caller is responsible
|
|
844
|
-
for managing the connector lifecycle.
|
|
845
421
|
external_history: Optional external history to use instead of the
|
|
846
422
|
internal conversation history.
|
|
847
423
|
output_schema: Optional Pydantic BaseModel class for structured output.
|
|
848
|
-
If provided, the agent will attempt to return an instance of this model.
|
|
849
424
|
|
|
850
425
|
Returns:
|
|
851
426
|
The result of running the query as a string, or if output_schema is provided,
|
|
@@ -881,8 +456,8 @@ class MCPAgent:
|
|
|
881
456
|
query, max_steps, manage_connector, external_history, track_execution=False, output_schema=output_schema
|
|
882
457
|
)
|
|
883
458
|
error = None
|
|
884
|
-
steps_taken = 0
|
|
885
459
|
result = None
|
|
460
|
+
steps_taken = 0
|
|
886
461
|
try:
|
|
887
462
|
result, steps_taken = await self._consume_and_return(generator)
|
|
888
463
|
|
|
@@ -982,6 +557,329 @@ class MCPAgent:
|
|
|
982
557
|
|
|
983
558
|
return enhanced_query
|
|
984
559
|
|
|
560
|
+
@telemetry("agent_stream")
|
|
561
|
+
async def stream(
|
|
562
|
+
self,
|
|
563
|
+
query: str,
|
|
564
|
+
max_steps: int | None = None,
|
|
565
|
+
manage_connector: bool = True,
|
|
566
|
+
external_history: list[BaseMessage] | None = None,
|
|
567
|
+
track_execution: bool = True,
|
|
568
|
+
output_schema: type[T] | None = None,
|
|
569
|
+
) -> AsyncGenerator[tuple[AgentAction, str] | str | T, None]:
|
|
570
|
+
"""Async generator using LangChain 1.0.0's create_agent and astream.
|
|
571
|
+
|
|
572
|
+
This method leverages the LangChain 1.0.0 API where create_agent returns
|
|
573
|
+
a CompiledStateGraph that handles the agent loop internally via astream.
|
|
574
|
+
|
|
575
|
+
**Tool Updates with Server Manager:**
|
|
576
|
+
When using server_manager mode, this method handles dynamic tool updates:
|
|
577
|
+
- **Before execution:** Updates are applied immediately to the new stream
|
|
578
|
+
- **During execution:** When tools change, we wait for a "safe restart point"
|
|
579
|
+
(after tool results complete), then interrupt the stream, recreate the agent
|
|
580
|
+
with new tools, and resume execution with accumulated messages.
|
|
581
|
+
- **Safe restart points:** Only restart after tool results to ensure message
|
|
582
|
+
pairs (tool_use + tool_result) are complete, satisfying LLM API requirements.
|
|
583
|
+
- **Max restarts:** Limited to 3 restarts to prevent infinite loops
|
|
584
|
+
|
|
585
|
+
This interrupt-and-restart approach ensures that tools added mid-execution
|
|
586
|
+
(e.g., via connect_to_mcp_server) are immediately available to the agent,
|
|
587
|
+
maintaining the same behavior as the legacy implementation while respecting
|
|
588
|
+
API constraints.
|
|
589
|
+
|
|
590
|
+
Args:
|
|
591
|
+
query: The query to run.
|
|
592
|
+
manage_connector: Whether to handle the connector lifecycle internally.
|
|
593
|
+
external_history: Optional external history to use instead of the
|
|
594
|
+
internal conversation history.
|
|
595
|
+
output_schema: Optional Pydantic BaseModel class for structured output.
|
|
596
|
+
|
|
597
|
+
Yields:
|
|
598
|
+
Intermediate steps and final result from the agent execution.
|
|
599
|
+
"""
|
|
600
|
+
# Delegate to remote agent if in remote mode
|
|
601
|
+
if self._is_remote and self._remote_agent:
|
|
602
|
+
async for item in self._remote_agent.stream(query, max_steps, external_history, output_schema):
|
|
603
|
+
yield item
|
|
604
|
+
return
|
|
605
|
+
|
|
606
|
+
initialized_here = False
|
|
607
|
+
start_time = time.time()
|
|
608
|
+
success = False
|
|
609
|
+
final_output = None
|
|
610
|
+
steps_taken = 0
|
|
611
|
+
|
|
612
|
+
try:
|
|
613
|
+
# 1. Initialize if needed
|
|
614
|
+
if manage_connector and not self._initialized:
|
|
615
|
+
await self.initialize()
|
|
616
|
+
initialized_here = True
|
|
617
|
+
elif not self._initialized and self.auto_initialize:
|
|
618
|
+
await self.initialize()
|
|
619
|
+
initialized_here = True
|
|
620
|
+
|
|
621
|
+
if not self._agent_executor:
|
|
622
|
+
raise RuntimeError("MCP agent failed to initialize")
|
|
623
|
+
|
|
624
|
+
# Check for tool updates before starting execution (if using server manager)
|
|
625
|
+
if self.use_server_manager and self.server_manager:
|
|
626
|
+
current_tools = self.server_manager.tools
|
|
627
|
+
current_tool_names = {tool.name for tool in current_tools}
|
|
628
|
+
existing_tool_names = {tool.name for tool in self._tools}
|
|
629
|
+
|
|
630
|
+
if current_tool_names != existing_tool_names:
|
|
631
|
+
logger.info(
|
|
632
|
+
f"🔄 Tools changed before execution, updating agent. New tools: {', '.join(current_tool_names)}"
|
|
633
|
+
)
|
|
634
|
+
self._tools = current_tools
|
|
635
|
+
# Regenerate system message with ALL current tools
|
|
636
|
+
await self._create_system_message_from_tools(self._tools)
|
|
637
|
+
# Recreate the agent executor with the new tools and system message
|
|
638
|
+
self._agent_executor = self._create_agent()
|
|
639
|
+
|
|
640
|
+
# 2. Build inputs for the agent
|
|
641
|
+
history_to_use = external_history if external_history is not None else self._conversation_history
|
|
642
|
+
|
|
643
|
+
# Convert messages to format expected by LangChain agent
|
|
644
|
+
langchain_history = []
|
|
645
|
+
for msg in history_to_use:
|
|
646
|
+
if isinstance(msg, HumanMessage | AIMessage):
|
|
647
|
+
langchain_history.append(msg)
|
|
648
|
+
|
|
649
|
+
inputs = {"messages": [*langchain_history, HumanMessage(content=query)]}
|
|
650
|
+
|
|
651
|
+
display_query = query[:50].replace("\n", " ") + "..." if len(query) > 50 else query.replace("\n", " ")
|
|
652
|
+
logger.info(f"💬 Received query: '{display_query}'")
|
|
653
|
+
logger.info("🏁 Starting agent execution")
|
|
654
|
+
|
|
655
|
+
# 3. Stream using the built-in astream from CompiledStateGraph
|
|
656
|
+
# The agent graph handles the loop internally
|
|
657
|
+
# With dynamic tool reload: if tools change mid-execution, we interrupt and restart
|
|
658
|
+
max_restarts = 3 # Prevent infinite restart loops
|
|
659
|
+
restart_count = 0
|
|
660
|
+
accumulated_messages = list(langchain_history) + [HumanMessage(content=query)]
|
|
661
|
+
pending_tool_calls = {} # Map tool_call_id -> AgentAction
|
|
662
|
+
|
|
663
|
+
while restart_count <= max_restarts:
|
|
664
|
+
# Update inputs with accumulated messages
|
|
665
|
+
inputs = {"messages": accumulated_messages}
|
|
666
|
+
should_restart = False
|
|
667
|
+
|
|
668
|
+
async for chunk in self._agent_executor.astream(
|
|
669
|
+
inputs,
|
|
670
|
+
stream_mode="updates", # Get updates as they happen
|
|
671
|
+
config={"callbacks": self.callbacks},
|
|
672
|
+
):
|
|
673
|
+
# chunk is a dict with node names as keys
|
|
674
|
+
# The agent node will have 'messages' with the AI response
|
|
675
|
+
# The tools node will have 'messages' with tool calls and results
|
|
676
|
+
|
|
677
|
+
for node_name, node_output in chunk.items():
|
|
678
|
+
logger.debug(f"📦 Node '{node_name}' output: {node_output}")
|
|
679
|
+
|
|
680
|
+
# Extract messages from the node output and accumulate them
|
|
681
|
+
if node_output is not None and "messages" in node_output:
|
|
682
|
+
messages = node_output["messages"]
|
|
683
|
+
if not isinstance(messages, list):
|
|
684
|
+
messages = [messages]
|
|
685
|
+
|
|
686
|
+
# Add new messages to accumulated messages for potential restart
|
|
687
|
+
for msg in messages:
|
|
688
|
+
if msg not in accumulated_messages:
|
|
689
|
+
accumulated_messages.append(msg)
|
|
690
|
+
for message in messages:
|
|
691
|
+
# Track tool calls
|
|
692
|
+
if hasattr(message, "tool_calls") and message.tool_calls:
|
|
693
|
+
# Extract text content from message for the log
|
|
694
|
+
log_text = ""
|
|
695
|
+
if hasattr(message, "content"):
|
|
696
|
+
if isinstance(message.content, str):
|
|
697
|
+
log_text = message.content
|
|
698
|
+
elif isinstance(message.content, list):
|
|
699
|
+
# Extract text blocks from content array
|
|
700
|
+
text_parts = [
|
|
701
|
+
block.get("text", "") if isinstance(block, dict) else str(block)
|
|
702
|
+
for block in message.content
|
|
703
|
+
if isinstance(block, dict) and block.get("type") == "text"
|
|
704
|
+
]
|
|
705
|
+
log_text = "\n".join(text_parts)
|
|
706
|
+
|
|
707
|
+
for tool_call in message.tool_calls:
|
|
708
|
+
tool_name = tool_call.get("name", "unknown")
|
|
709
|
+
tool_input = tool_call.get("args", {})
|
|
710
|
+
tool_call_id = tool_call.get("id")
|
|
711
|
+
|
|
712
|
+
action = AgentAction(tool=tool_name, tool_input=tool_input, log=log_text)
|
|
713
|
+
if tool_call_id:
|
|
714
|
+
pending_tool_calls[tool_call_id] = action
|
|
715
|
+
|
|
716
|
+
self.tools_used_names.append(tool_name)
|
|
717
|
+
steps_taken += 1
|
|
718
|
+
|
|
719
|
+
tool_input_str = str(tool_input)
|
|
720
|
+
if len(tool_input_str) > 100:
|
|
721
|
+
tool_input_str = tool_input_str[:97] + "..."
|
|
722
|
+
logger.info(f"🔧 Tool call: {tool_name} with input: {tool_input_str}")
|
|
723
|
+
|
|
724
|
+
# Track tool results and yield AgentStep
|
|
725
|
+
if hasattr(message, "type") and message.type == "tool":
|
|
726
|
+
observation = message.content
|
|
727
|
+
tool_call_id = getattr(message, "tool_call_id", None)
|
|
728
|
+
|
|
729
|
+
if tool_call_id and tool_call_id in pending_tool_calls:
|
|
730
|
+
action = pending_tool_calls.pop(tool_call_id)
|
|
731
|
+
yield (action, str(observation))
|
|
732
|
+
|
|
733
|
+
observation_str = str(observation)
|
|
734
|
+
if len(observation_str) > 100:
|
|
735
|
+
observation_str = observation_str[:97] + "..."
|
|
736
|
+
observation_str = observation_str.replace("\n", " ")
|
|
737
|
+
logger.info(f"📄 Tool result: {observation_str}")
|
|
738
|
+
|
|
739
|
+
# --- Check for tool updates after tool results (safe restart point) ---
|
|
740
|
+
if self.use_server_manager and self.server_manager:
|
|
741
|
+
current_tools = self.server_manager.tools
|
|
742
|
+
current_tool_names = {tool.name for tool in current_tools}
|
|
743
|
+
existing_tool_names = {tool.name for tool in self._tools}
|
|
744
|
+
|
|
745
|
+
if current_tool_names != existing_tool_names:
|
|
746
|
+
logger.info(
|
|
747
|
+
f"🔄 Tools changed during execution. "
|
|
748
|
+
f"New tools: {', '.join(current_tool_names)}"
|
|
749
|
+
)
|
|
750
|
+
self._tools = current_tools
|
|
751
|
+
# Regenerate system message with ALL current tools
|
|
752
|
+
await self._create_system_message_from_tools(self._tools)
|
|
753
|
+
# Recreate the agent executor with the new tools and system message
|
|
754
|
+
self._agent_executor = self._create_agent()
|
|
755
|
+
|
|
756
|
+
# Set restart flag - safe to restart now after tool results
|
|
757
|
+
should_restart = True
|
|
758
|
+
restart_count += 1
|
|
759
|
+
logger.info(
|
|
760
|
+
f"🔃 Restarting execution with updated tools "
|
|
761
|
+
f"(restart {restart_count}/{max_restarts})"
|
|
762
|
+
)
|
|
763
|
+
break # Break out of the message loop
|
|
764
|
+
|
|
765
|
+
# Track final AI message (without tool calls = final response)
|
|
766
|
+
if isinstance(message, AIMessage) and not getattr(message, "tool_calls", None):
|
|
767
|
+
final_output = self._normalize_output(message.content)
|
|
768
|
+
logger.info("✅ Agent finished with output")
|
|
769
|
+
|
|
770
|
+
# Break out of node loop if restarting
|
|
771
|
+
if should_restart:
|
|
772
|
+
break
|
|
773
|
+
|
|
774
|
+
# Break out of chunk loop if restarting
|
|
775
|
+
if should_restart:
|
|
776
|
+
break
|
|
777
|
+
|
|
778
|
+
# Check if we should restart or if execution completed
|
|
779
|
+
if not should_restart:
|
|
780
|
+
# Execution completed successfully without tool changes
|
|
781
|
+
break
|
|
782
|
+
|
|
783
|
+
# If we've hit max restarts, log warning and continue
|
|
784
|
+
if restart_count > max_restarts:
|
|
785
|
+
logger.warning(f"⚠️ Max restarts ({max_restarts}) reached. Continuing with current tools.")
|
|
786
|
+
break
|
|
787
|
+
|
|
788
|
+
# 4. Update conversation history
|
|
789
|
+
if self.memory_enabled:
|
|
790
|
+
self.add_to_history(HumanMessage(content=query))
|
|
791
|
+
if final_output:
|
|
792
|
+
self.add_to_history(AIMessage(content=final_output))
|
|
793
|
+
|
|
794
|
+
# 5. Handle structured output if requested
|
|
795
|
+
if output_schema and final_output:
|
|
796
|
+
try:
|
|
797
|
+
logger.info("🔧 Attempting structured output...")
|
|
798
|
+
structured_llm = self.llm.with_structured_output(output_schema)
|
|
799
|
+
|
|
800
|
+
# Get schema description
|
|
801
|
+
schema_fields = []
|
|
802
|
+
for field_name, field_info in output_schema.model_fields.items():
|
|
803
|
+
description = getattr(field_info, "description", "") or field_name
|
|
804
|
+
required = not hasattr(field_info, "default") or field_info.default is None
|
|
805
|
+
schema_fields.append(
|
|
806
|
+
f"- {field_name}: {description} " + ("(required)" if required else "(optional)")
|
|
807
|
+
)
|
|
808
|
+
schema_description = "\n".join(schema_fields)
|
|
809
|
+
|
|
810
|
+
structured_result = await self._attempt_structured_output(
|
|
811
|
+
final_output, structured_llm, output_schema, schema_description
|
|
812
|
+
)
|
|
813
|
+
|
|
814
|
+
if self.memory_enabled:
|
|
815
|
+
self.add_to_history(AIMessage(content=f"Structured result: {structured_result}"))
|
|
816
|
+
|
|
817
|
+
logger.info("✅ Structured output successful")
|
|
818
|
+
success = True
|
|
819
|
+
yield structured_result
|
|
820
|
+
return
|
|
821
|
+
except Exception as e:
|
|
822
|
+
logger.error(f"❌ Structured output failed: {e}")
|
|
823
|
+
raise RuntimeError(f"Failed to generate structured output: {str(e)}") from e
|
|
824
|
+
|
|
825
|
+
# 6. Yield final result
|
|
826
|
+
logger.info(f"🎉 Agent execution complete in {time.time() - start_time:.2f} seconds")
|
|
827
|
+
success = True
|
|
828
|
+
yield final_output or "No output generated"
|
|
829
|
+
|
|
830
|
+
except Exception as e:
|
|
831
|
+
logger.error(f"❌ Error running query: {e}")
|
|
832
|
+
if initialized_here and manage_connector:
|
|
833
|
+
logger.info("🧹 Cleaning up resources after error")
|
|
834
|
+
await self.close()
|
|
835
|
+
raise
|
|
836
|
+
|
|
837
|
+
finally:
|
|
838
|
+
# Track comprehensive execution data
|
|
839
|
+
execution_time_ms = int((time.time() - start_time) * 1000)
|
|
840
|
+
|
|
841
|
+
server_count = 0
|
|
842
|
+
if self.client:
|
|
843
|
+
server_count = len(self.client.get_all_active_sessions())
|
|
844
|
+
elif self.connectors:
|
|
845
|
+
server_count = len(self.connectors)
|
|
846
|
+
|
|
847
|
+
conversation_history_length = len(self._conversation_history) if self.memory_enabled else 0
|
|
848
|
+
|
|
849
|
+
# Safely access _tools in case initialization failed
|
|
850
|
+
tools_available = getattr(self, "_tools", [])
|
|
851
|
+
|
|
852
|
+
if track_execution:
|
|
853
|
+
self.telemetry.track_agent_execution(
|
|
854
|
+
execution_method="stream",
|
|
855
|
+
query=query,
|
|
856
|
+
success=success,
|
|
857
|
+
model_provider=self._model_provider,
|
|
858
|
+
model_name=self._model_name,
|
|
859
|
+
server_count=server_count,
|
|
860
|
+
server_identifiers=[connector.public_identifier for connector in self.connectors],
|
|
861
|
+
total_tools_available=len(tools_available),
|
|
862
|
+
tools_available_names=[tool.name for tool in tools_available],
|
|
863
|
+
max_steps_configured=self.max_steps,
|
|
864
|
+
memory_enabled=self.memory_enabled,
|
|
865
|
+
use_server_manager=self.use_server_manager,
|
|
866
|
+
max_steps_used=max_steps,
|
|
867
|
+
manage_connector=manage_connector,
|
|
868
|
+
external_history_used=external_history is not None,
|
|
869
|
+
steps_taken=steps_taken,
|
|
870
|
+
tools_used_count=len(self.tools_used_names),
|
|
871
|
+
tools_used_names=self.tools_used_names,
|
|
872
|
+
response=final_output,
|
|
873
|
+
execution_time_ms=execution_time_ms,
|
|
874
|
+
error_type=None if success else "execution_error",
|
|
875
|
+
conversation_history_length=conversation_history_length,
|
|
876
|
+
)
|
|
877
|
+
|
|
878
|
+
# Clean up if necessary
|
|
879
|
+
if manage_connector and not self.client and initialized_here:
|
|
880
|
+
logger.info("🧹 Closing agent after stream completion")
|
|
881
|
+
await self.close()
|
|
882
|
+
|
|
985
883
|
async def _generate_response_chunks_async(
|
|
986
884
|
self,
|
|
987
885
|
query: str,
|
|
@@ -1011,19 +909,21 @@ class MCPAgent:
|
|
|
1011
909
|
raise RuntimeError("MCP agent failed to initialise – call initialise() first?")
|
|
1012
910
|
|
|
1013
911
|
# 2. Build inputs --------------------------------------------------------
|
|
1014
|
-
|
|
1015
|
-
self._agent_executor.max_iterations = effective_max_steps
|
|
912
|
+
self.max_steps = max_steps or self.max_steps
|
|
1016
913
|
|
|
914
|
+
# 3. Build inputs --------------------------------------------------------
|
|
1017
915
|
history_to_use = external_history if external_history is not None else self._conversation_history
|
|
1018
916
|
inputs = {"input": query, "chat_history": history_to_use}
|
|
1019
917
|
|
|
1020
918
|
# 3. Stream & diff -------------------------------------------------------
|
|
1021
|
-
async for event in self._agent_executor.astream_events(inputs):
|
|
919
|
+
async for event in self._agent_executor.astream_events(inputs, config={"callbacks": self.callbacks}):
|
|
1022
920
|
if event.get("event") == "on_chain_end":
|
|
1023
921
|
output = event["data"]["output"]
|
|
1024
922
|
if isinstance(output, list):
|
|
1025
923
|
for message in output:
|
|
1026
|
-
|
|
924
|
+
# Filter out ToolMessage (equivalent to old ToolAgentAction)
|
|
925
|
+
# to avoid adding intermediate tool execution details to history
|
|
926
|
+
if isinstance(message, BaseMessage) and not isinstance(message, ToolMessage):
|
|
1027
927
|
self.add_to_history(message)
|
|
1028
928
|
yield event
|
|
1029
929
|
|
|
@@ -1038,6 +938,7 @@ class MCPAgent:
|
|
|
1038
938
|
logger.info("🧹 Closing agent after generator completion")
|
|
1039
939
|
await self.close()
|
|
1040
940
|
|
|
941
|
+
@telemetry("agent_stream_events")
|
|
1041
942
|
async def stream_events(
|
|
1042
943
|
self,
|
|
1043
944
|
query: str,
|
|
@@ -1049,7 +950,7 @@ class MCPAgent:
|
|
|
1049
950
|
|
|
1050
951
|
Example::
|
|
1051
952
|
|
|
1052
|
-
async for chunk in agent.
|
|
953
|
+
async for chunk in agent.stream("hello"):
|
|
1053
954
|
print(chunk, end="|", flush=True)
|
|
1054
955
|
"""
|
|
1055
956
|
start_time = time.time()
|