ag2 0.9.6__py3-none-any.whl → 0.9.8.post1__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 ag2 might be problematic. Click here for more details.
- {ag2-0.9.6.dist-info → ag2-0.9.8.post1.dist-info}/METADATA +102 -75
- ag2-0.9.8.post1.dist-info/RECORD +387 -0
- autogen/__init__.py +1 -2
- autogen/_website/generate_api_references.py +4 -5
- autogen/_website/generate_mkdocs.py +9 -15
- autogen/_website/notebook_processor.py +13 -14
- autogen/_website/process_notebooks.py +10 -10
- autogen/_website/utils.py +5 -4
- autogen/agentchat/agent.py +13 -13
- autogen/agentchat/assistant_agent.py +7 -6
- autogen/agentchat/contrib/agent_eval/agent_eval.py +3 -3
- autogen/agentchat/contrib/agent_eval/critic_agent.py +3 -3
- autogen/agentchat/contrib/agent_eval/quantifier_agent.py +3 -3
- autogen/agentchat/contrib/agent_eval/subcritic_agent.py +3 -3
- autogen/agentchat/contrib/agent_optimizer.py +3 -3
- autogen/agentchat/contrib/capabilities/generate_images.py +11 -11
- autogen/agentchat/contrib/capabilities/teachability.py +15 -15
- autogen/agentchat/contrib/capabilities/transforms.py +17 -18
- autogen/agentchat/contrib/capabilities/transforms_util.py +5 -5
- autogen/agentchat/contrib/capabilities/vision_capability.py +4 -3
- autogen/agentchat/contrib/captainagent/agent_builder.py +30 -30
- autogen/agentchat/contrib/captainagent/captainagent.py +22 -21
- autogen/agentchat/contrib/captainagent/tool_retriever.py +2 -3
- autogen/agentchat/contrib/gpt_assistant_agent.py +9 -9
- autogen/agentchat/contrib/graph_rag/document.py +3 -3
- autogen/agentchat/contrib/graph_rag/falkor_graph_query_engine.py +3 -3
- autogen/agentchat/contrib/graph_rag/falkor_graph_rag_capability.py +6 -6
- autogen/agentchat/contrib/graph_rag/graph_query_engine.py +3 -3
- autogen/agentchat/contrib/graph_rag/neo4j_graph_query_engine.py +5 -11
- autogen/agentchat/contrib/graph_rag/neo4j_graph_rag_capability.py +6 -6
- autogen/agentchat/contrib/graph_rag/neo4j_native_graph_query_engine.py +7 -7
- autogen/agentchat/contrib/graph_rag/neo4j_native_graph_rag_capability.py +6 -6
- autogen/agentchat/contrib/img_utils.py +1 -1
- autogen/agentchat/contrib/llamaindex_conversable_agent.py +11 -11
- autogen/agentchat/contrib/llava_agent.py +18 -4
- autogen/agentchat/contrib/math_user_proxy_agent.py +11 -11
- autogen/agentchat/contrib/multimodal_conversable_agent.py +8 -8
- autogen/agentchat/contrib/qdrant_retrieve_user_proxy_agent.py +6 -5
- autogen/agentchat/contrib/rag/chromadb_query_engine.py +22 -26
- autogen/agentchat/contrib/rag/llamaindex_query_engine.py +14 -17
- autogen/agentchat/contrib/rag/mongodb_query_engine.py +27 -37
- autogen/agentchat/contrib/rag/query_engine.py +7 -5
- autogen/agentchat/contrib/retrieve_assistant_agent.py +5 -5
- autogen/agentchat/contrib/retrieve_user_proxy_agent.py +8 -7
- autogen/agentchat/contrib/society_of_mind_agent.py +15 -14
- autogen/agentchat/contrib/swarm_agent.py +76 -98
- autogen/agentchat/contrib/text_analyzer_agent.py +7 -7
- autogen/agentchat/contrib/vectordb/base.py +10 -18
- autogen/agentchat/contrib/vectordb/chromadb.py +2 -1
- autogen/agentchat/contrib/vectordb/couchbase.py +18 -20
- autogen/agentchat/contrib/vectordb/mongodb.py +6 -5
- autogen/agentchat/contrib/vectordb/pgvectordb.py +40 -41
- autogen/agentchat/contrib/vectordb/qdrant.py +5 -5
- autogen/agentchat/contrib/web_surfer.py +20 -19
- autogen/agentchat/conversable_agent.py +311 -295
- autogen/agentchat/group/context_str.py +1 -3
- autogen/agentchat/group/context_variables.py +15 -25
- autogen/agentchat/group/group_tool_executor.py +10 -10
- autogen/agentchat/group/group_utils.py +15 -15
- autogen/agentchat/group/guardrails.py +7 -7
- autogen/agentchat/group/handoffs.py +19 -36
- autogen/agentchat/group/multi_agent_chat.py +7 -7
- autogen/agentchat/group/on_condition.py +4 -7
- autogen/agentchat/group/on_context_condition.py +4 -7
- autogen/agentchat/group/patterns/auto.py +8 -7
- autogen/agentchat/group/patterns/manual.py +7 -6
- autogen/agentchat/group/patterns/pattern.py +13 -12
- autogen/agentchat/group/patterns/random.py +3 -3
- autogen/agentchat/group/patterns/round_robin.py +3 -3
- autogen/agentchat/group/reply_result.py +2 -4
- autogen/agentchat/group/speaker_selection_result.py +5 -5
- autogen/agentchat/group/targets/group_chat_target.py +7 -6
- autogen/agentchat/group/targets/group_manager_target.py +4 -4
- autogen/agentchat/group/targets/transition_target.py +2 -1
- autogen/agentchat/groupchat.py +58 -61
- autogen/agentchat/realtime/experimental/audio_adapters/twilio_audio_adapter.py +4 -4
- autogen/agentchat/realtime/experimental/audio_adapters/websocket_audio_adapter.py +4 -4
- autogen/agentchat/realtime/experimental/clients/gemini/client.py +7 -7
- autogen/agentchat/realtime/experimental/clients/oai/base_client.py +8 -8
- autogen/agentchat/realtime/experimental/clients/oai/rtc_client.py +6 -6
- autogen/agentchat/realtime/experimental/clients/realtime_client.py +10 -9
- autogen/agentchat/realtime/experimental/realtime_agent.py +10 -9
- autogen/agentchat/realtime/experimental/realtime_observer.py +3 -3
- autogen/agentchat/realtime/experimental/realtime_swarm.py +44 -44
- autogen/agentchat/user_proxy_agent.py +10 -9
- autogen/agentchat/utils.py +3 -3
- autogen/agents/contrib/time/time_reply_agent.py +6 -5
- autogen/agents/contrib/time/time_tool_agent.py +2 -1
- autogen/agents/experimental/deep_research/deep_research.py +3 -3
- autogen/agents/experimental/discord/discord.py +2 -2
- autogen/agents/experimental/document_agent/chroma_query_engine.py +29 -44
- autogen/agents/experimental/document_agent/docling_doc_ingest_agent.py +9 -14
- autogen/agents/experimental/document_agent/document_agent.py +15 -16
- autogen/agents/experimental/document_agent/document_conditions.py +3 -3
- autogen/agents/experimental/document_agent/document_utils.py +5 -9
- autogen/agents/experimental/document_agent/inmemory_query_engine.py +14 -20
- autogen/agents/experimental/document_agent/parser_utils.py +4 -4
- autogen/agents/experimental/document_agent/url_utils.py +14 -23
- autogen/agents/experimental/reasoning/reasoning_agent.py +33 -33
- autogen/agents/experimental/slack/slack.py +2 -2
- autogen/agents/experimental/telegram/telegram.py +2 -3
- autogen/agents/experimental/websurfer/websurfer.py +4 -4
- autogen/agents/experimental/wikipedia/wikipedia.py +5 -7
- autogen/browser_utils.py +8 -8
- autogen/cache/abstract_cache_base.py +5 -5
- autogen/cache/cache.py +12 -12
- autogen/cache/cache_factory.py +4 -4
- autogen/cache/cosmos_db_cache.py +9 -9
- autogen/cache/disk_cache.py +6 -6
- autogen/cache/in_memory_cache.py +4 -4
- autogen/cache/redis_cache.py +4 -4
- autogen/code_utils.py +18 -18
- autogen/coding/base.py +6 -6
- autogen/coding/docker_commandline_code_executor.py +9 -9
- autogen/coding/func_with_reqs.py +7 -6
- autogen/coding/jupyter/base.py +3 -3
- autogen/coding/jupyter/docker_jupyter_server.py +3 -4
- autogen/coding/jupyter/import_utils.py +3 -3
- autogen/coding/jupyter/jupyter_client.py +5 -5
- autogen/coding/jupyter/jupyter_code_executor.py +3 -4
- autogen/coding/jupyter/local_jupyter_server.py +2 -6
- autogen/coding/local_commandline_code_executor.py +8 -7
- autogen/coding/markdown_code_extractor.py +1 -2
- autogen/coding/utils.py +1 -2
- autogen/doc_utils.py +3 -2
- autogen/environments/docker_python_environment.py +19 -29
- autogen/environments/python_environment.py +8 -17
- autogen/environments/system_python_environment.py +3 -4
- autogen/environments/venv_python_environment.py +8 -12
- autogen/environments/working_directory.py +1 -2
- autogen/events/agent_events.py +106 -109
- autogen/events/base_event.py +6 -5
- autogen/events/client_events.py +15 -14
- autogen/events/helpers.py +1 -1
- autogen/events/print_event.py +4 -5
- autogen/fast_depends/_compat.py +10 -15
- autogen/fast_depends/core/build.py +17 -36
- autogen/fast_depends/core/model.py +64 -113
- autogen/fast_depends/dependencies/model.py +2 -1
- autogen/fast_depends/dependencies/provider.py +3 -2
- autogen/fast_depends/library/model.py +4 -4
- autogen/fast_depends/schema.py +7 -7
- autogen/fast_depends/use.py +17 -25
- autogen/fast_depends/utils.py +10 -30
- autogen/formatting_utils.py +6 -6
- autogen/graph_utils.py +1 -4
- autogen/import_utils.py +13 -13
- autogen/interop/crewai/crewai.py +2 -2
- autogen/interop/interoperable.py +2 -2
- autogen/interop/langchain/langchain_chat_model_factory.py +3 -2
- autogen/interop/langchain/langchain_tool.py +2 -6
- autogen/interop/litellm/litellm_config_factory.py +6 -7
- autogen/interop/pydantic_ai/pydantic_ai.py +4 -7
- autogen/interop/registry.py +2 -1
- autogen/io/base.py +5 -5
- autogen/io/run_response.py +33 -32
- autogen/io/websockets.py +6 -5
- autogen/json_utils.py +1 -2
- autogen/llm_config/__init__.py +11 -0
- autogen/llm_config/client.py +58 -0
- autogen/llm_config/config.py +384 -0
- autogen/llm_config/entry.py +154 -0
- autogen/logger/base_logger.py +4 -3
- autogen/logger/file_logger.py +2 -1
- autogen/logger/logger_factory.py +2 -2
- autogen/logger/logger_utils.py +2 -2
- autogen/logger/sqlite_logger.py +3 -2
- autogen/math_utils.py +4 -5
- autogen/mcp/__main__.py +6 -6
- autogen/mcp/helpers.py +4 -4
- autogen/mcp/mcp_client.py +170 -29
- autogen/mcp/mcp_proxy/fastapi_code_generator_helpers.py +3 -4
- autogen/mcp/mcp_proxy/mcp_proxy.py +23 -26
- autogen/mcp/mcp_proxy/operation_grouping.py +4 -5
- autogen/mcp/mcp_proxy/operation_renaming.py +6 -10
- autogen/mcp/mcp_proxy/security.py +2 -3
- autogen/messages/agent_messages.py +96 -98
- autogen/messages/base_message.py +6 -5
- autogen/messages/client_messages.py +15 -14
- autogen/messages/print_message.py +4 -5
- autogen/oai/__init__.py +1 -2
- autogen/oai/anthropic.py +42 -41
- autogen/oai/bedrock.py +68 -57
- autogen/oai/cerebras.py +26 -25
- autogen/oai/client.py +118 -138
- autogen/oai/client_utils.py +3 -3
- autogen/oai/cohere.py +34 -11
- autogen/oai/gemini.py +40 -17
- autogen/oai/gemini_types.py +11 -12
- autogen/oai/groq.py +22 -10
- autogen/oai/mistral.py +17 -11
- autogen/oai/oai_models/__init__.py +14 -2
- autogen/oai/oai_models/_models.py +2 -2
- autogen/oai/oai_models/chat_completion.py +13 -14
- autogen/oai/oai_models/chat_completion_message.py +11 -9
- autogen/oai/oai_models/chat_completion_message_tool_call.py +26 -3
- autogen/oai/oai_models/chat_completion_token_logprob.py +3 -4
- autogen/oai/oai_models/completion_usage.py +8 -9
- autogen/oai/ollama.py +22 -10
- autogen/oai/openai_responses.py +40 -17
- autogen/oai/openai_utils.py +159 -85
- autogen/oai/together.py +29 -14
- autogen/retrieve_utils.py +6 -7
- autogen/runtime_logging.py +5 -4
- autogen/token_count_utils.py +7 -4
- autogen/tools/contrib/time/time.py +0 -1
- autogen/tools/dependency_injection.py +5 -6
- autogen/tools/experimental/browser_use/browser_use.py +10 -10
- autogen/tools/experimental/code_execution/python_code_execution.py +5 -7
- autogen/tools/experimental/crawl4ai/crawl4ai.py +12 -15
- autogen/tools/experimental/deep_research/deep_research.py +9 -8
- autogen/tools/experimental/duckduckgo/duckduckgo_search.py +5 -11
- autogen/tools/experimental/firecrawl/firecrawl_tool.py +98 -115
- autogen/tools/experimental/google/authentication/credentials_local_provider.py +1 -1
- autogen/tools/experimental/google/drive/drive_functions.py +4 -4
- autogen/tools/experimental/google/drive/toolkit.py +5 -5
- autogen/tools/experimental/google_search/google_search.py +5 -5
- autogen/tools/experimental/google_search/youtube_search.py +5 -5
- autogen/tools/experimental/messageplatform/discord/discord.py +8 -12
- autogen/tools/experimental/messageplatform/slack/slack.py +14 -20
- autogen/tools/experimental/messageplatform/telegram/telegram.py +8 -12
- autogen/tools/experimental/perplexity/perplexity_search.py +18 -29
- autogen/tools/experimental/reliable/reliable.py +68 -74
- autogen/tools/experimental/searxng/searxng_search.py +20 -19
- autogen/tools/experimental/tavily/tavily_search.py +12 -19
- autogen/tools/experimental/web_search_preview/web_search_preview.py +13 -7
- autogen/tools/experimental/wikipedia/wikipedia.py +7 -10
- autogen/tools/function_utils.py +7 -7
- autogen/tools/tool.py +6 -5
- autogen/types.py +2 -2
- autogen/version.py +1 -1
- ag2-0.9.6.dist-info/RECORD +0 -421
- autogen/llm_config.py +0 -385
- {ag2-0.9.6.dist-info → ag2-0.9.8.post1.dist-info}/WHEEL +0 -0
- {ag2-0.9.6.dist-info → ag2-0.9.8.post1.dist-info}/licenses/LICENSE +0 -0
- {ag2-0.9.6.dist-info → ag2-0.9.8.post1.dist-info}/licenses/NOTICE.md +0 -0
|
@@ -0,0 +1,154 @@
|
|
|
1
|
+
# Copyright (c) 2023 - 2025, AG2ai, Inc., AG2ai open-source projects maintainers and core contributors
|
|
2
|
+
#
|
|
3
|
+
# SPDX-License-Identifier: Apache-2.0
|
|
4
|
+
|
|
5
|
+
import re
|
|
6
|
+
from abc import ABC, abstractmethod
|
|
7
|
+
from collections.abc import Iterable, Mapping
|
|
8
|
+
from typing import Any
|
|
9
|
+
|
|
10
|
+
from httpx import Client as httpxClient
|
|
11
|
+
from pydantic import BaseModel, ConfigDict, Field, HttpUrl, SecretStr, ValidationInfo, field_serializer, field_validator
|
|
12
|
+
from typing_extensions import Required, TypedDict
|
|
13
|
+
|
|
14
|
+
from .client import ModelClient
|
|
15
|
+
|
|
16
|
+
|
|
17
|
+
class LLMConfigEntryDict(TypedDict, total=False):
|
|
18
|
+
api_type: Required[str]
|
|
19
|
+
model: str
|
|
20
|
+
max_tokens: int | None
|
|
21
|
+
top_p: float | None
|
|
22
|
+
temperature: float | None
|
|
23
|
+
|
|
24
|
+
api_key: SecretStr | None
|
|
25
|
+
api_version: str | None
|
|
26
|
+
base_url: HttpUrl | None
|
|
27
|
+
voice: str | None
|
|
28
|
+
http_client: httpxClient | None
|
|
29
|
+
model_client_cls: str | None
|
|
30
|
+
response_format: str | dict[str, Any] | BaseModel | type[BaseModel] | None
|
|
31
|
+
default_headers: Mapping[str, Any] | None
|
|
32
|
+
tags: list[str]
|
|
33
|
+
|
|
34
|
+
|
|
35
|
+
class ApplicationConfig(BaseModel):
|
|
36
|
+
max_tokens: int | None = Field(default=None, ge=0)
|
|
37
|
+
top_p: float | None = Field(default=None, gt=0, lt=1)
|
|
38
|
+
temperature: float | None = Field(default=None, ge=0, le=1)
|
|
39
|
+
|
|
40
|
+
@field_validator("top_p", mode="before")
|
|
41
|
+
@classmethod
|
|
42
|
+
def check_top_p(cls, v: float | None, info: ValidationInfo) -> float | None:
|
|
43
|
+
if v is not None and info.data.get("temperature") is not None:
|
|
44
|
+
raise ValueError("temperature and top_p cannot be set at the same time.")
|
|
45
|
+
return v
|
|
46
|
+
|
|
47
|
+
@field_validator("temperature", mode="before")
|
|
48
|
+
@classmethod
|
|
49
|
+
def check_temperature(cls, v: float | None, info: ValidationInfo) -> float | None:
|
|
50
|
+
if v is not None and info.data.get("top_p") is not None:
|
|
51
|
+
raise ValueError("temperature and top_p cannot be set at the same time.")
|
|
52
|
+
return v
|
|
53
|
+
|
|
54
|
+
|
|
55
|
+
class LLMConfigEntry(ApplicationConfig, ABC):
|
|
56
|
+
api_type: str
|
|
57
|
+
model: str = Field(..., min_length=1)
|
|
58
|
+
|
|
59
|
+
api_key: SecretStr | None = None
|
|
60
|
+
api_version: str | None = None
|
|
61
|
+
base_url: HttpUrl | None = None
|
|
62
|
+
voice: str | None = None
|
|
63
|
+
model_client_cls: str | None = None
|
|
64
|
+
http_client: httpxClient | None = None
|
|
65
|
+
response_format: str | dict[str, Any] | BaseModel | type[BaseModel] | None = None
|
|
66
|
+
default_headers: Mapping[str, Any] | None = None
|
|
67
|
+
tags: list[str] = Field(default_factory=list)
|
|
68
|
+
|
|
69
|
+
# Following field is configuration for pydantic to disallow extra fields
|
|
70
|
+
model_config = ConfigDict(extra="allow", arbitrary_types_allowed=True)
|
|
71
|
+
|
|
72
|
+
def apply_application_config(self, application_config: ApplicationConfig) -> "LLMConfigEntry":
|
|
73
|
+
"""Apply application level configurations."""
|
|
74
|
+
# TODO: should create a new instance instead of mutating current one
|
|
75
|
+
self.max_tokens = self.max_tokens or application_config.max_tokens
|
|
76
|
+
self.top_p = self.top_p or application_config.top_p
|
|
77
|
+
self.temperature = self.temperature or application_config.temperature
|
|
78
|
+
return self
|
|
79
|
+
|
|
80
|
+
@abstractmethod
|
|
81
|
+
def create_client(self) -> "ModelClient": ...
|
|
82
|
+
|
|
83
|
+
@field_validator("base_url", mode="before")
|
|
84
|
+
@classmethod
|
|
85
|
+
def check_base_url(cls, v: HttpUrl | str | None, info: ValidationInfo) -> str | None:
|
|
86
|
+
if v is None: # Handle None case explicitly
|
|
87
|
+
return None
|
|
88
|
+
if not str(v).startswith("https://") and not str(v).startswith("http://"):
|
|
89
|
+
return f"http://{str(v)}"
|
|
90
|
+
return str(v)
|
|
91
|
+
|
|
92
|
+
@field_serializer("base_url", when_used="unless-none") # Ensure serializer also respects None
|
|
93
|
+
def serialize_base_url(self, v: HttpUrl | None) -> str | None:
|
|
94
|
+
return str(v) if v is not None else None
|
|
95
|
+
|
|
96
|
+
@field_serializer("api_key", when_used="unless-none")
|
|
97
|
+
def serialize_api_key(self, v: SecretStr) -> str:
|
|
98
|
+
return v.get_secret_value()
|
|
99
|
+
|
|
100
|
+
def model_dump(self, *args: Any, exclude_none: bool = True, **kwargs: Any) -> dict[str, Any]:
|
|
101
|
+
return BaseModel.model_dump(self, exclude_none=exclude_none, *args, **kwargs)
|
|
102
|
+
|
|
103
|
+
def model_dump_json(self, *args: Any, exclude_none: bool = True, **kwargs: Any) -> str:
|
|
104
|
+
return BaseModel.model_dump_json(self, exclude_none=exclude_none, *args, **kwargs)
|
|
105
|
+
|
|
106
|
+
def get(self, key: str, default: Any | None = None) -> Any:
|
|
107
|
+
val = getattr(self, key, default)
|
|
108
|
+
if isinstance(val, SecretStr):
|
|
109
|
+
return val.get_secret_value()
|
|
110
|
+
return val
|
|
111
|
+
|
|
112
|
+
def __getitem__(self, key: str) -> Any:
|
|
113
|
+
try:
|
|
114
|
+
val = getattr(self, key)
|
|
115
|
+
if isinstance(val, SecretStr):
|
|
116
|
+
return val.get_secret_value()
|
|
117
|
+
return val
|
|
118
|
+
except AttributeError:
|
|
119
|
+
raise KeyError(f"Key '{key}' not found in {self.__class__.__name__}")
|
|
120
|
+
|
|
121
|
+
def __setitem__(self, key: str, value: Any) -> None:
|
|
122
|
+
setattr(self, key, value)
|
|
123
|
+
|
|
124
|
+
def __contains__(self, key: str) -> bool:
|
|
125
|
+
return hasattr(self, key)
|
|
126
|
+
|
|
127
|
+
def items(self) -> Iterable[tuple[str, Any]]:
|
|
128
|
+
d = self.model_dump()
|
|
129
|
+
return d.items()
|
|
130
|
+
|
|
131
|
+
def keys(self) -> Iterable[str]:
|
|
132
|
+
d = self.model_dump()
|
|
133
|
+
return d.keys()
|
|
134
|
+
|
|
135
|
+
def values(self) -> Iterable[Any]:
|
|
136
|
+
d = self.model_dump()
|
|
137
|
+
return d.values()
|
|
138
|
+
|
|
139
|
+
def __repr__(self) -> str:
|
|
140
|
+
# Override to eliminate none values from the repr
|
|
141
|
+
d = self.model_dump()
|
|
142
|
+
r = [f"{k}={repr(v)}" for k, v in d.items()]
|
|
143
|
+
|
|
144
|
+
s = f"{self.__class__.__name__}({', '.join(r)})"
|
|
145
|
+
|
|
146
|
+
# Replace any keys ending with '_key' or '_token' values with stars for security
|
|
147
|
+
# This regex will match any key ending with '_key' or '_token' and its value, and replace the value with stars
|
|
148
|
+
# It also captures the type of quote used (single or double) and reuses it in the replacement
|
|
149
|
+
s = re.sub(r'(\w+_(key|token)\s*=\s*)([\'"]).*?\3', r"\1\3**********\3", s, flags=re.IGNORECASE)
|
|
150
|
+
|
|
151
|
+
return s
|
|
152
|
+
|
|
153
|
+
def __str__(self) -> str:
|
|
154
|
+
return repr(self)
|
autogen/logger/base_logger.py
CHANGED
|
@@ -9,7 +9,8 @@ from __future__ import annotations
|
|
|
9
9
|
import sqlite3
|
|
10
10
|
import uuid
|
|
11
11
|
from abc import ABC, abstractmethod
|
|
12
|
-
from
|
|
12
|
+
from collections.abc import Callable
|
|
13
|
+
from typing import TYPE_CHECKING, Any, TypeVar
|
|
13
14
|
|
|
14
15
|
if TYPE_CHECKING:
|
|
15
16
|
from openai import AzureOpenAI, OpenAI
|
|
@@ -18,8 +19,8 @@ if TYPE_CHECKING:
|
|
|
18
19
|
from .. import Agent, ConversableAgent, OpenAIWrapper
|
|
19
20
|
|
|
20
21
|
F = TypeVar("F", bound=Callable[..., Any])
|
|
21
|
-
ConfigItem = dict[str,
|
|
22
|
-
LLMConfig = dict[str,
|
|
22
|
+
ConfigItem = dict[str, str | list[str]]
|
|
23
|
+
LLMConfig = dict[str, None | float | int | ConfigItem | list[ConfigItem]]
|
|
23
24
|
|
|
24
25
|
|
|
25
26
|
class BaseLogger(ABC):
|
autogen/logger/file_logger.py
CHANGED
|
@@ -11,7 +11,8 @@ import logging
|
|
|
11
11
|
import os
|
|
12
12
|
import threading
|
|
13
13
|
import uuid
|
|
14
|
-
from
|
|
14
|
+
from collections.abc import Callable
|
|
15
|
+
from typing import TYPE_CHECKING, Any, TypeVar
|
|
15
16
|
|
|
16
17
|
from ..doc_utils import export_module
|
|
17
18
|
from .base_logger import BaseLogger, LLMConfig
|
autogen/logger/logger_factory.py
CHANGED
|
@@ -4,7 +4,7 @@
|
|
|
4
4
|
#
|
|
5
5
|
# Portions derived from https://github.com/microsoft/autogen are under the MIT License.
|
|
6
6
|
# SPDX-License-Identifier: MIT
|
|
7
|
-
from typing import Any, Literal
|
|
7
|
+
from typing import Any, Literal
|
|
8
8
|
|
|
9
9
|
from ..doc_utils import export_module
|
|
10
10
|
from .base_logger import BaseLogger
|
|
@@ -20,7 +20,7 @@ class LoggerFactory:
|
|
|
20
20
|
|
|
21
21
|
@staticmethod
|
|
22
22
|
def get_logger(
|
|
23
|
-
logger_type: Literal["sqlite", "file"] = "sqlite", config:
|
|
23
|
+
logger_type: Literal["sqlite", "file"] = "sqlite", config: dict[str, Any] | None = None
|
|
24
24
|
) -> BaseLogger:
|
|
25
25
|
"""Factory method to create logger objects.
|
|
26
26
|
|
autogen/logger/logger_utils.py
CHANGED
|
@@ -7,7 +7,7 @@
|
|
|
7
7
|
import inspect
|
|
8
8
|
from datetime import datetime, timezone
|
|
9
9
|
from pathlib import Path, PurePath
|
|
10
|
-
from typing import Any
|
|
10
|
+
from typing import Any
|
|
11
11
|
|
|
12
12
|
__all__ = ("get_current_ts", "to_dict")
|
|
13
13
|
|
|
@@ -22,7 +22,7 @@ def get_current_ts() -> str:
|
|
|
22
22
|
|
|
23
23
|
|
|
24
24
|
def to_dict(
|
|
25
|
-
obj:
|
|
25
|
+
obj: int | float | str | bool | dict[Any, Any] | list[Any] | tuple[Any, ...] | Any,
|
|
26
26
|
exclude: tuple[str, ...] = (),
|
|
27
27
|
no_recursive: tuple[Any, ...] = (),
|
|
28
28
|
) -> Any:
|
autogen/logger/sqlite_logger.py
CHANGED
|
@@ -12,7 +12,8 @@ import os
|
|
|
12
12
|
import sqlite3
|
|
13
13
|
import threading
|
|
14
14
|
import uuid
|
|
15
|
-
from
|
|
15
|
+
from collections.abc import Callable
|
|
16
|
+
from typing import TYPE_CHECKING, Any, TypeVar
|
|
16
17
|
|
|
17
18
|
from ..doc_utils import export_module
|
|
18
19
|
from .base_logger import BaseLogger, LLMConfig
|
|
@@ -293,7 +294,7 @@ class SqliteLogger(BaseLogger):
|
|
|
293
294
|
client_id,
|
|
294
295
|
wrapper_id,
|
|
295
296
|
self.session_id,
|
|
296
|
-
json.dumps(request),
|
|
297
|
+
json.dumps(to_dict(request)),
|
|
297
298
|
response_messages,
|
|
298
299
|
is_cached,
|
|
299
300
|
cost,
|
autogen/math_utils.py
CHANGED
|
@@ -4,10 +4,9 @@
|
|
|
4
4
|
#
|
|
5
5
|
# Portions derived from https://github.com/microsoft/autogen are under the MIT License.
|
|
6
6
|
# SPDX-License-Identifier: MIT
|
|
7
|
-
from typing import Optional
|
|
8
7
|
|
|
9
8
|
|
|
10
|
-
def remove_boxed(string: str) ->
|
|
9
|
+
def remove_boxed(string: str) -> str | None:
|
|
11
10
|
"""Source: https://github.com/hendrycks/math
|
|
12
11
|
Extract the text within a \\boxed`{...}` environment.
|
|
13
12
|
|
|
@@ -28,7 +27,7 @@ def remove_boxed(string: str) -> Optional[str]:
|
|
|
28
27
|
return None
|
|
29
28
|
|
|
30
29
|
|
|
31
|
-
def last_boxed_only_string(string: str) ->
|
|
30
|
+
def last_boxed_only_string(string: str) -> str | None:
|
|
32
31
|
"""Source: https://github.com/hendrycks/math
|
|
33
32
|
Extract the last \\boxed`{...}` or \\fbox`{...}` element from a string.
|
|
34
33
|
"""
|
|
@@ -239,7 +238,7 @@ def _strip_string(string: str) -> str:
|
|
|
239
238
|
return string
|
|
240
239
|
|
|
241
240
|
|
|
242
|
-
def get_answer(solution:
|
|
241
|
+
def get_answer(solution: str | None) -> str | None:
|
|
243
242
|
if solution is None:
|
|
244
243
|
return None
|
|
245
244
|
last_boxed = last_boxed_only_string(solution)
|
|
@@ -251,7 +250,7 @@ def get_answer(solution: Optional[str]) -> Optional[str]:
|
|
|
251
250
|
return answer
|
|
252
251
|
|
|
253
252
|
|
|
254
|
-
def is_equiv(str1:
|
|
253
|
+
def is_equiv(str1: str | None, str2: str | None) -> float:
|
|
255
254
|
"""Returns (as a float) whether two strings containing math are equivalent up to differences of formatting in
|
|
256
255
|
- units
|
|
257
256
|
- fractions
|
autogen/mcp/__main__.py
CHANGED
|
@@ -2,7 +2,7 @@
|
|
|
2
2
|
#
|
|
3
3
|
# SPDX-License-Identifier: Apache-2.0
|
|
4
4
|
import logging
|
|
5
|
-
from typing import Annotated, Literal
|
|
5
|
+
from typing import Annotated, Literal
|
|
6
6
|
|
|
7
7
|
from .. import __version__
|
|
8
8
|
from ..import_utils import optional_import_block, require_optional_import
|
|
@@ -27,7 +27,7 @@ def create_typer_app() -> "typer.Typer":
|
|
|
27
27
|
@app.callback()
|
|
28
28
|
def callback(
|
|
29
29
|
version: Annotated[
|
|
30
|
-
|
|
30
|
+
bool | None,
|
|
31
31
|
typer.Option("--version", help="Show the version and exit.", callback=version_callback),
|
|
32
32
|
] = None,
|
|
33
33
|
) -> None:
|
|
@@ -41,19 +41,19 @@ def create_typer_app() -> "typer.Typer":
|
|
|
41
41
|
@app.command()
|
|
42
42
|
def create(
|
|
43
43
|
openapi_specification: Annotated[
|
|
44
|
-
|
|
44
|
+
str | None,
|
|
45
45
|
"Specification of the OpenAPI to use for the proxy generation.",
|
|
46
46
|
] = None,
|
|
47
47
|
openapi_url: Annotated[
|
|
48
|
-
|
|
48
|
+
str | None,
|
|
49
49
|
"URL to the OpenAPI specification to use for the proxy generation.",
|
|
50
50
|
] = None,
|
|
51
51
|
client_source_path: Annotated[
|
|
52
|
-
|
|
52
|
+
str | None,
|
|
53
53
|
"Path to the generated proxy client source code.",
|
|
54
54
|
] = None,
|
|
55
55
|
server_url: Annotated[
|
|
56
|
-
|
|
56
|
+
str | None,
|
|
57
57
|
"Comma-separated list of server URLs to use for the proxy generation.",
|
|
58
58
|
] = None,
|
|
59
59
|
configuration_type: Annotated[
|
autogen/mcp/helpers.py
CHANGED
|
@@ -5,21 +5,21 @@ import asyncio
|
|
|
5
5
|
import os
|
|
6
6
|
import signal
|
|
7
7
|
from asyncio.subprocess import PIPE, Process, create_subprocess_exec
|
|
8
|
+
from collections.abc import AsyncGenerator
|
|
8
9
|
from contextlib import asynccontextmanager
|
|
9
|
-
from typing import AsyncGenerator, Dict, Optional
|
|
10
10
|
|
|
11
11
|
|
|
12
12
|
@asynccontextmanager
|
|
13
13
|
async def run_streamable_http_client(
|
|
14
|
-
*, mcp_server_path: str, env_vars:
|
|
14
|
+
*, mcp_server_path: str, env_vars: dict[str, str] | None = None, startup_wait_secs: float = 5.0
|
|
15
15
|
) -> AsyncGenerator[Process, None]:
|
|
16
|
-
"""
|
|
17
|
-
Async context manager to run a Python subprocess for streamable-http with custom env vars.
|
|
16
|
+
"""Async context manager to run a Python subprocess for streamable-http with custom env vars.
|
|
18
17
|
|
|
19
18
|
Args:
|
|
20
19
|
mcp_server_path: Path to the Python script to run.
|
|
21
20
|
env_vars: Environment variables to export to the subprocess.
|
|
22
21
|
startup_wait_secs: Time to wait for the server to start (in seconds).
|
|
22
|
+
|
|
23
23
|
Yields:
|
|
24
24
|
An asyncio.subprocess.Process object.
|
|
25
25
|
"""
|
autogen/mcp/mcp_client.py
CHANGED
|
@@ -3,38 +3,153 @@
|
|
|
3
3
|
# SPDX-License-Identifier: Apache-2.0
|
|
4
4
|
|
|
5
5
|
|
|
6
|
-
import
|
|
7
|
-
from
|
|
6
|
+
from collections.abc import AsyncIterator
|
|
7
|
+
from contextlib import AbstractAsyncContextManager, AsyncExitStack, asynccontextmanager
|
|
8
|
+
from datetime import datetime, timedelta
|
|
8
9
|
from pathlib import Path
|
|
9
|
-
from typing import Annotated, Any,
|
|
10
|
+
from typing import Annotated, Any, Literal, Protocol, cast
|
|
10
11
|
|
|
11
12
|
import anyio
|
|
12
|
-
from
|
|
13
|
+
from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
|
|
14
|
+
from mcp.client.session import ClientSession
|
|
15
|
+
from mcp.client.sse import sse_client
|
|
16
|
+
from mcp.client.stdio import StdioServerParameters, stdio_client
|
|
17
|
+
from pydantic import AnyUrl, BaseModel, Field
|
|
13
18
|
|
|
14
19
|
from ..doc_utils import export_module
|
|
15
20
|
from ..import_utils import optional_import_block, require_optional_import
|
|
16
21
|
from ..tools import Tool, Toolkit
|
|
17
22
|
|
|
18
23
|
with optional_import_block():
|
|
19
|
-
from mcp import
|
|
20
|
-
from mcp.types import
|
|
21
|
-
|
|
22
|
-
ReadResourceResult,
|
|
23
|
-
ResourceTemplate,
|
|
24
|
-
TextContent,
|
|
25
|
-
)
|
|
26
|
-
from mcp.types import (
|
|
27
|
-
Tool as MCPTool,
|
|
28
|
-
)
|
|
24
|
+
from mcp.shared.message import SessionMessage
|
|
25
|
+
from mcp.types import CallToolResult, ReadResourceResult, ResourceTemplate, TextContent
|
|
26
|
+
from mcp.types import Tool as MCPTool
|
|
29
27
|
|
|
30
28
|
__all__ = ["ResultSaved", "create_toolkit"]
|
|
31
29
|
|
|
30
|
+
# Type definitions
|
|
31
|
+
EncodingErrorHandlerType = Literal["strict", "ignore", "replace"]
|
|
32
|
+
|
|
33
|
+
# Default constants
|
|
34
|
+
DEFAULT_TEXT_ENCODING = "utf-8"
|
|
35
|
+
DEFAULT_TEXT_ENCODING_ERROR_HANDLER: EncodingErrorHandlerType = "strict"
|
|
36
|
+
DEFAULT_HTTP_REQUEST_TIMEOUT = 5
|
|
37
|
+
DEFAULT_SSE_EVENT_READ_TIMEOUT = 60 * 5
|
|
38
|
+
DEFAULT_STREAMABLE_HTTP_REQUEST_TIMEOUT = timedelta(seconds=30)
|
|
39
|
+
DEFAULT_STREAMABLE_HTTP_SSE_EVENT_READ_TIMEOUT = timedelta(seconds=60 * 5)
|
|
40
|
+
|
|
41
|
+
|
|
42
|
+
class SessionConfigProtocol(Protocol):
|
|
43
|
+
"""Protocol for session configuration classes that can create MCP sessions."""
|
|
44
|
+
|
|
45
|
+
server_name: str
|
|
46
|
+
|
|
47
|
+
@asynccontextmanager
|
|
48
|
+
async def create_session(self, exit_stack: AsyncExitStack) -> AsyncIterator[ClientSession]:
|
|
49
|
+
"""Create a session using the given exit stack."""
|
|
50
|
+
try:
|
|
51
|
+
yield cast(ClientSession, None) # placeholder yield to satisfy AsyncIterator type
|
|
52
|
+
except Exception:
|
|
53
|
+
raise NotImplementedError
|
|
54
|
+
|
|
55
|
+
|
|
56
|
+
class BasicSessionConfig(BaseModel):
|
|
57
|
+
"""Basic session configuration."""
|
|
58
|
+
|
|
59
|
+
server_name: str = Field(..., description="Name of the server")
|
|
60
|
+
|
|
61
|
+
async def initialize(
|
|
62
|
+
self,
|
|
63
|
+
client: AbstractAsyncContextManager[
|
|
64
|
+
tuple[
|
|
65
|
+
MemoryObjectReceiveStream[SessionMessage | Exception],
|
|
66
|
+
MemoryObjectSendStream[SessionMessage],
|
|
67
|
+
]
|
|
68
|
+
],
|
|
69
|
+
exit_stack: AsyncExitStack,
|
|
70
|
+
) -> ClientSession:
|
|
71
|
+
"""Initialize the session."""
|
|
72
|
+
reader, writer = await exit_stack.enter_async_context(client)
|
|
73
|
+
session = cast(
|
|
74
|
+
ClientSession,
|
|
75
|
+
await exit_stack.enter_async_context(ClientSession(reader, writer)),
|
|
76
|
+
)
|
|
77
|
+
return session
|
|
78
|
+
|
|
79
|
+
|
|
80
|
+
class SseConfig(BasicSessionConfig):
|
|
81
|
+
"""Configuration for a single SSE MCP server."""
|
|
82
|
+
|
|
83
|
+
url: str = Field(..., description="URL of the SSE server")
|
|
84
|
+
headers: dict[str, Any] | None = Field(default=None, description="HTTP headers to send to the SSE endpoint")
|
|
85
|
+
timeout: float = Field(default=DEFAULT_HTTP_REQUEST_TIMEOUT, description="HTTP timeout")
|
|
86
|
+
sse_read_timeout: float = Field(default=DEFAULT_SSE_EVENT_READ_TIMEOUT, description="SSE read timeout")
|
|
87
|
+
|
|
88
|
+
@asynccontextmanager
|
|
89
|
+
async def create_session(self, exit_stack: AsyncExitStack) -> AsyncIterator[ClientSession]:
|
|
90
|
+
"""
|
|
91
|
+
Create a new session to an MCP server using SSE transport.
|
|
92
|
+
|
|
93
|
+
Args:
|
|
94
|
+
exit_stack: AsyncExitStack for managing async resources
|
|
95
|
+
|
|
96
|
+
Yields:
|
|
97
|
+
ClientSession: The MCP client session
|
|
98
|
+
"""
|
|
99
|
+
# Create and store the connection
|
|
100
|
+
client = sse_client(self.url, self.headers, self.timeout, self.sse_read_timeout)
|
|
101
|
+
yield await self.initialize(client, exit_stack)
|
|
102
|
+
|
|
103
|
+
|
|
104
|
+
class StdioConfig(BasicSessionConfig):
|
|
105
|
+
"""Configuration for a single stdio MCP server."""
|
|
106
|
+
|
|
107
|
+
command: str = Field(..., description="Command to execute")
|
|
108
|
+
args: list[str] = Field(..., description="Arguments for the command")
|
|
109
|
+
transport: Literal["stdio"] = Field(default="stdio", description="Transport type")
|
|
110
|
+
environment: dict[str, str] | None = Field(default=None, description="Environment variables")
|
|
111
|
+
working_dir: str | Path | None = Field(default=None, description="Working directory")
|
|
112
|
+
encoding: str = Field(default=DEFAULT_TEXT_ENCODING, description="Character encoding")
|
|
113
|
+
encoding_error_handler: EncodingErrorHandlerType = Field(
|
|
114
|
+
default=DEFAULT_TEXT_ENCODING_ERROR_HANDLER, description="How to handle encoding errors"
|
|
115
|
+
)
|
|
116
|
+
session_options: dict[str, Any] | None = Field(default=None, description="Additional session options")
|
|
117
|
+
|
|
118
|
+
@asynccontextmanager
|
|
119
|
+
async def create_session(self, exit_stack: AsyncExitStack) -> AsyncIterator[ClientSession]:
|
|
120
|
+
"""
|
|
121
|
+
Create a new session to an MCP server using stdio transport.
|
|
122
|
+
|
|
123
|
+
Args:
|
|
124
|
+
exit_stack: AsyncExitStack for managing async resources
|
|
125
|
+
|
|
126
|
+
Yields:
|
|
127
|
+
ClientSession: The MCP client session
|
|
128
|
+
"""
|
|
129
|
+
client = stdio_client(
|
|
130
|
+
StdioServerParameters(
|
|
131
|
+
command=self.command,
|
|
132
|
+
args=self.args,
|
|
133
|
+
env=self.environment,
|
|
134
|
+
encoding=self.encoding,
|
|
135
|
+
encoding_error_handler=self.encoding_error_handler,
|
|
136
|
+
)
|
|
137
|
+
)
|
|
138
|
+
yield await self.initialize(client, exit_stack)
|
|
139
|
+
|
|
140
|
+
|
|
141
|
+
class MCPConfig(BaseModel):
|
|
142
|
+
"""Configuration for multiple MCP sessions using stdio transport."""
|
|
143
|
+
|
|
144
|
+
# we should use final classes to allow pydantic to validate the type
|
|
145
|
+
servers: list[SseConfig | StdioConfig] = Field(..., description="List of stdio & sse server configurations")
|
|
146
|
+
|
|
32
147
|
|
|
33
148
|
class MCPClient:
|
|
34
149
|
@staticmethod
|
|
35
150
|
def _convert_call_tool_result( # type: ignore[no-any-unimported]
|
|
36
151
|
call_tool_result: "CallToolResult", # type: ignore[no-any-unimported]
|
|
37
|
-
) -> tuple[
|
|
152
|
+
) -> tuple[str | list[str], Any]:
|
|
38
153
|
text_contents: list[TextContent] = [] # type: ignore[no-any-unimported]
|
|
39
154
|
non_text_contents = []
|
|
40
155
|
for content in call_tool_result.content:
|
|
@@ -43,7 +158,7 @@ class MCPClient:
|
|
|
43
158
|
else:
|
|
44
159
|
non_text_contents.append(content)
|
|
45
160
|
|
|
46
|
-
tool_content:
|
|
161
|
+
tool_content: str | list[str] = [content.text for content in text_contents]
|
|
47
162
|
if len(text_contents) == 1:
|
|
48
163
|
tool_content = tool_content[0]
|
|
49
164
|
|
|
@@ -65,7 +180,7 @@ class MCPClient:
|
|
|
65
180
|
|
|
66
181
|
async def call_tool( # type: ignore[no-any-unimported]
|
|
67
182
|
**arguments: dict[str, Any],
|
|
68
|
-
) -> tuple[
|
|
183
|
+
) -> tuple[str | list[str], Any]:
|
|
69
184
|
call_tool_result = await session.call_tool(tool.name, arguments)
|
|
70
185
|
return MCPClient._convert_call_tool_result(call_tool_result)
|
|
71
186
|
|
|
@@ -83,7 +198,7 @@ class MCPClient:
|
|
|
83
198
|
cls,
|
|
84
199
|
resource_template: Any,
|
|
85
200
|
session: "ClientSession",
|
|
86
|
-
resource_download_folder:
|
|
201
|
+
resource_download_folder: Path | None,
|
|
87
202
|
**kwargs: Any,
|
|
88
203
|
) -> Tool:
|
|
89
204
|
if not isinstance(resource_template, ResourceTemplate):
|
|
@@ -97,8 +212,8 @@ Here is the correct format for the URI template:
|
|
|
97
212
|
{mcp_resource.uriTemplate}
|
|
98
213
|
"""
|
|
99
214
|
|
|
100
|
-
async def call_resource(uri: Annotated[str, uri_description]) ->
|
|
101
|
-
result = await session.read_resource(uri)
|
|
215
|
+
async def call_resource(uri: Annotated[str, uri_description]) -> ReadResourceResult | ResultSaved: # type: ignore[no-any-unimported]
|
|
216
|
+
result = await session.read_resource(AnyUrl(uri))
|
|
102
217
|
|
|
103
218
|
if not resource_download_folder:
|
|
104
219
|
return result
|
|
@@ -131,7 +246,7 @@ Here is the correct format for the URI template:
|
|
|
131
246
|
*,
|
|
132
247
|
use_mcp_tools: bool,
|
|
133
248
|
use_mcp_resources: bool,
|
|
134
|
-
resource_download_folder:
|
|
249
|
+
resource_download_folder: Path | None,
|
|
135
250
|
) -> Toolkit: # type: ignore[no-any-unimported]
|
|
136
251
|
"""Load all available MCP tools and convert them to AG2 Toolkit."""
|
|
137
252
|
all_ag2_tools: list[Tool] = []
|
|
@@ -156,10 +271,7 @@ Here is the correct format for the URI template:
|
|
|
156
271
|
return Toolkit(tools=all_ag2_tools)
|
|
157
272
|
|
|
158
273
|
@classmethod
|
|
159
|
-
def get_unsupported_reason(cls) ->
|
|
160
|
-
if sys.version_info < (3, 10):
|
|
161
|
-
return "This submodule is only supported for Python versions 3.10 and above"
|
|
162
|
-
|
|
274
|
+
def get_unsupported_reason(cls) -> str | None:
|
|
163
275
|
with optional_import_block() as result:
|
|
164
276
|
import mcp # noqa: F401
|
|
165
277
|
|
|
@@ -169,13 +281,42 @@ Here is the correct format for the URI template:
|
|
|
169
281
|
return None
|
|
170
282
|
|
|
171
283
|
|
|
284
|
+
class MCPClientSessionManager:
|
|
285
|
+
"""
|
|
286
|
+
A class to manage MCP client sessions.
|
|
287
|
+
"""
|
|
288
|
+
|
|
289
|
+
def __init__(self) -> None:
|
|
290
|
+
"""Initialize the MCP client session manager."""
|
|
291
|
+
self.exit_stack = AsyncExitStack()
|
|
292
|
+
self.sessions: dict[str, ClientSession] = {}
|
|
293
|
+
|
|
294
|
+
@asynccontextmanager
|
|
295
|
+
async def open_session(
|
|
296
|
+
self,
|
|
297
|
+
config: SessionConfigProtocol,
|
|
298
|
+
) -> AsyncIterator[ClientSession]:
|
|
299
|
+
"""Open a new session to an MCP server based on configuration.
|
|
300
|
+
|
|
301
|
+
Args:
|
|
302
|
+
config: SessionConfigProtocol object containing session configuration
|
|
303
|
+
|
|
304
|
+
Yields:
|
|
305
|
+
ClientSession: The MCP client session
|
|
306
|
+
"""
|
|
307
|
+
async with config.create_session(self.exit_stack) as session:
|
|
308
|
+
await session.initialize()
|
|
309
|
+
self.sessions[config.server_name] = session
|
|
310
|
+
yield session
|
|
311
|
+
|
|
312
|
+
|
|
172
313
|
@export_module("autogen.mcp")
|
|
173
314
|
async def create_toolkit(
|
|
174
315
|
session: "ClientSession",
|
|
175
316
|
*,
|
|
176
317
|
use_mcp_tools: bool = True,
|
|
177
318
|
use_mcp_resources: bool = True,
|
|
178
|
-
resource_download_folder:
|
|
319
|
+
resource_download_folder: Path | str | None = None,
|
|
179
320
|
) -> Toolkit: # type: ignore[no-any-unimported]
|
|
180
321
|
"""Create a toolkit from the MCP client session.
|
|
181
322
|
|
|
@@ -184,12 +325,12 @@ async def create_toolkit(
|
|
|
184
325
|
use_mcp_tools (bool): Whether to include MCP tools in the toolkit.
|
|
185
326
|
use_mcp_resources (bool): Whether to include MCP resources in the toolkit.
|
|
186
327
|
resource_download_folder (Optional[Union[Path, str]]): The folder to download files to.
|
|
328
|
+
|
|
187
329
|
Returns:
|
|
188
330
|
Toolkit: The toolkit containing the converted tools.
|
|
189
331
|
"""
|
|
190
|
-
if resource_download_folder:
|
|
191
|
-
|
|
192
|
-
resource_download_folder = Path(resource_download_folder)
|
|
332
|
+
if resource_download_folder is not None:
|
|
333
|
+
resource_download_folder = Path(resource_download_folder)
|
|
193
334
|
await anyio.to_thread.run_sync(lambda: resource_download_folder.mkdir(parents=True, exist_ok=True))
|
|
194
335
|
|
|
195
336
|
return await MCPClient.load_mcp_toolkit(
|
|
@@ -4,7 +4,6 @@
|
|
|
4
4
|
from collections.abc import Iterator
|
|
5
5
|
from contextlib import contextmanager
|
|
6
6
|
from functools import cached_property
|
|
7
|
-
from typing import Optional, Union
|
|
8
7
|
|
|
9
8
|
from ...import_utils import optional_import_block
|
|
10
9
|
|
|
@@ -24,7 +23,7 @@ __all__ = ["SUCCESFUL_IMPORT", "patch_get_parameter_type"]
|
|
|
24
23
|
@contextmanager
|
|
25
24
|
def patch_get_parameter_type() -> Iterator[None]:
|
|
26
25
|
class ArgumentWithDescription(Argument): # type: ignore[misc]
|
|
27
|
-
description:
|
|
26
|
+
description: str | None = None
|
|
28
27
|
|
|
29
28
|
@cached_property
|
|
30
29
|
def argument(self) -> str:
|
|
@@ -43,10 +42,10 @@ def patch_get_parameter_type() -> Iterator[None]:
|
|
|
43
42
|
|
|
44
43
|
def get_parameter_type(
|
|
45
44
|
self: OpenAPIParser,
|
|
46
|
-
parameters:
|
|
45
|
+
parameters: ReferenceObject | ParameterObject,
|
|
47
46
|
snake_case: bool,
|
|
48
47
|
path: list[str],
|
|
49
|
-
) ->
|
|
48
|
+
) -> Argument | None:
|
|
50
49
|
# get the original argument
|
|
51
50
|
argument = original_get_parameter_type(self, parameters, snake_case, path)
|
|
52
51
|
|