Compare commits

...

3 Commits

Author SHA1 Message Date
  Eric Zhu c9a89ae444 fix 1 year ago
  Eric Zhu 329951d295 wip 1 year ago
  Eric Zhu ca57fecaab Add ToolCallEvent and log it for all built-in tools. 1 year ago
19 changed files with 639 additions and 86 deletions
Unified View
  1. +48
    -3
      python/packages/autogen-core/src/autogen_core/logging.py
  2. +19
    -1
      python/packages/autogen-core/src/autogen_core/tools/_function_tool.py
  3. +8
    -2
      python/packages/autogen-core/tests/test_tool_agent.py
  4. +0
    -8
      python/packages/autogen-ext/src/autogen_ext/models/anthropic/_anthropic_client.py
  5. +394
    -0
      python/packages/autogen-ext/src/autogen_ext/models/llama_cpp/_llama_cpp_completion_client.py
  6. +0
    -8
      python/packages/autogen-ext/src/autogen_ext/models/ollama/_ollama_client.py
  7. +0
    -8
      python/packages/autogen-ext/src/autogen_ext/models/openai/_openai_client.py
  8. +1
    -1
      python/packages/autogen-ext/src/autogen_ext/models/replay/_replay_chat_completion_client.py
  9. +17
    -2
      python/packages/autogen-ext/src/autogen_ext/tools/code_execution/_code_execution.py
  10. +16
    -4
      python/packages/autogen-ext/src/autogen_ext/tools/graphrag/_global_search.py
  11. +17
    -4
      python/packages/autogen-ext/src/autogen_ext/tools/graphrag/_local_search.py
  12. +12
    -3
      python/packages/autogen-ext/src/autogen_ext/tools/http/_http_tool.py
  13. +20
    -1
      python/packages/autogen-ext/src/autogen_ext/tools/langchain/_langchain_adapter.py
  14. +13
    -1
      python/packages/autogen-ext/src/autogen_ext/tools/mcp/_base.py
  15. +22
    -12
      python/packages/autogen-ext/tests/tools/graphrag/test_graphrag_tools.py
  16. +9
    -4
      python/packages/autogen-ext/tests/tools/http/test_http_tool.py
  17. +8
    -4
      python/packages/autogen-ext/tests/tools/test_langchain_tools.py
  18. +27
    -16
      python/packages/autogen-ext/tests/tools/test_mcp_tools.py
  19. +8
    -4
      python/packages/autogen-ext/tests/tools/test_python_code_executor_tool.py

+ 48
- 3
python/packages/autogen-core/src/autogen_core/logging.py View File

@@ -3,6 +3,7 @@ from enum import Enum
from typing import Any, Dict, cast from typing import Any, Dict, cast


from ._agent_id import AgentId from ._agent_id import AgentId
from ._message_handler_context import MessageHandlerContext
from ._topic import TopicId from ._topic import TopicId




@@ -14,7 +15,6 @@ class LLMCallEvent:
response: Dict[str, Any], response: Dict[str, Any],
prompt_tokens: int, prompt_tokens: int,
completion_tokens: int, completion_tokens: int,
agent_id: AgentId | None = None,
**kwargs: Any, **kwargs: Any,
) -> None: ) -> None:
"""To be used by model clients to log the call to the LLM. """To be used by model clients to log the call to the LLM.
@@ -24,7 +24,6 @@ class LLMCallEvent:
response (Dict[str, Any]): The response of the call. Must be json serializable. response (Dict[str, Any]): The response of the call. Must be json serializable.
prompt_tokens (int): Number of tokens used in the prompt. prompt_tokens (int): Number of tokens used in the prompt.
completion_tokens (int): Number of tokens used in the completion. completion_tokens (int): Number of tokens used in the completion.
agent_id (AgentId | None, optional): The agent id of the model. Defaults to None.


Example: Example:


@@ -43,8 +42,11 @@ class LLMCallEvent:
self.kwargs["response"] = response self.kwargs["response"] = response
self.kwargs["prompt_tokens"] = prompt_tokens self.kwargs["prompt_tokens"] = prompt_tokens
self.kwargs["completion_tokens"] = completion_tokens self.kwargs["completion_tokens"] = completion_tokens
try:
agent_id = MessageHandlerContext.agent_id()
except RuntimeError:
agent_id = None
self.kwargs["agent_id"] = None if agent_id is None else str(agent_id) self.kwargs["agent_id"] = None if agent_id is None else str(agent_id)
self.kwargs["type"] = "LLMCall"


@property @property
def prompt_tokens(self) -> int: def prompt_tokens(self) -> int:
@@ -59,6 +61,49 @@ class LLMCallEvent:
return json.dumps(self.kwargs) return json.dumps(self.kwargs)




class ToolCallEvent:
def __init__(
self,
*,
tool_name: str,
arguments: Dict[str, Any],
result: Any = None,
**kwargs: Any,
) -> None:
"""Used by subclasses of :class:`~autogen_core.tools.BaseTool` to log executions of tools.

Args:
tool_name (str): The name of the tool.
arguments (Dict[str, Any]): The arguments of the tool. Must be json serializable.
result (Any, optional): The result of the tool. Must be json serializable. Defaults to None.

Example:

.. code-block:: python

from autogen_core import EVENT_LOGGER_NAME
from autogen_core.logging import ToolCallEvent

logger = logging.getLogger(EVENT_LOGGER_NAME)
logger.info(ToolCallEvent(tool_name="Tool1", call_id="123", arguments={"arg1": "value1"}))

"""
self.kwargs = kwargs
self.kwargs["type"] = "ToolCall"
self.kwargs["tool_name"] = tool_name
self.kwargs["arguments"] = arguments
self.kwargs["result"] = result
try:
agent_id = MessageHandlerContext.agent_id()
except RuntimeError:
agent_id = None
self.kwargs["agent_id"] = None if agent_id is None else str(agent_id)

# This must output the event in a json serializable format
def __str__(self) -> str:
return json.dumps(self.kwargs)


class MessageKind(Enum): class MessageKind(Enum):
DIRECT = 1 DIRECT = 1
PUBLISH = 2 PUBLISH = 2


+ 19
- 1
python/packages/autogen-core/src/autogen_core/tools/_function_tool.py View File

@@ -1,5 +1,6 @@
import asyncio import asyncio
import functools import functools
import logging
import warnings import warnings
from textwrap import dedent from textwrap import dedent
from typing import Any, Callable, Sequence from typing import Any, Callable, Sequence
@@ -7,15 +8,18 @@ from typing import Any, Callable, Sequence
from pydantic import BaseModel from pydantic import BaseModel
from typing_extensions import Self from typing_extensions import Self


from .. import CancellationToken
from .. import EVENT_LOGGER_NAME, CancellationToken, MessageHandlerContext
from .._component_config import Component from .._component_config import Component
from .._function_utils import ( from .._function_utils import (
args_base_model_from_signature, args_base_model_from_signature,
get_typed_signature, get_typed_signature,
) )
from ..code_executor._func_with_reqs import Import, import_to_str, to_code from ..code_executor._func_with_reqs import Import, import_to_str, to_code
from ..logging import ToolCallEvent
from ._base import BaseTool from ._base import BaseTool


logger = logging.getLogger(EVENT_LOGGER_NAME)



class FunctionToolConfig(BaseModel): class FunctionToolConfig(BaseModel):
"""Configuration for a function tool.""" """Configuration for a function tool."""
@@ -129,6 +133,20 @@ class FunctionTool(BaseTool[BaseModel, BaseModel], Component[FunctionToolConfig]
cancellation_token.link_future(future) cancellation_token.link_future(future)
result = await future result = await future


# If we are running in the context of a handler we can get the agent_id
try:
agent_id = MessageHandlerContext.agent_id()
except RuntimeError:
agent_id = None
# Log the function call.
event = ToolCallEvent(
tool_name=self.name,
arguments=args.model_dump(),
result=self.return_value_as_string(result) if result is not None else None,
agent_id=agent_id,
)
logger.info(event)

return result return result


def _to_config(self) -> FunctionToolConfig: def _to_config(self) -> FunctionToolConfig:


+ 8
- 2
python/packages/autogen-core/tests/test_tool_agent.py View File

@@ -1,9 +1,10 @@
import asyncio import asyncio
import json import json
import logging
from typing import Any, AsyncGenerator, List, Mapping, Optional, Sequence, Union from typing import Any, AsyncGenerator, List, Mapping, Optional, Sequence, Union


import pytest import pytest
from autogen_core import AgentId, CancellationToken, FunctionCall, SingleThreadedAgentRuntime
from autogen_core import EVENT_LOGGER_NAME, AgentId, CancellationToken, FunctionCall, SingleThreadedAgentRuntime
from autogen_core.models import ( from autogen_core.models import (
AssistantMessage, AssistantMessage,
ChatCompletionClient, ChatCompletionClient,
@@ -25,6 +26,8 @@ from autogen_core.tool_agent import (
) )
from autogen_core.tools import FunctionTool, Tool, ToolSchema from autogen_core.tools import FunctionTool, Tool, ToolSchema


logging.getLogger(EVENT_LOGGER_NAME).setLevel(logging.INFO)



def _pass_function(input: str) -> str: def _pass_function(input: str) -> str:
return "pass" return "pass"
@@ -40,7 +43,7 @@ async def _async_sleep_function(input: str) -> str:




@pytest.mark.asyncio @pytest.mark.asyncio
async def test_tool_agent() -> None:
async def test_tool_agent(caplog: pytest.LogCaptureFixture) -> None:
runtime = SingleThreadedAgentRuntime() runtime = SingleThreadedAgentRuntime()
await ToolAgent.register( await ToolAgent.register(
runtime, runtime,
@@ -63,6 +66,9 @@ async def test_tool_agent() -> None:
) )
assert result == FunctionExecutionResult(call_id="1", content="pass", is_error=False, name="pass") assert result == FunctionExecutionResult(call_id="1", content="pass", is_error=False, name="pass")


# Check log.
assert any(("ToolCall" in record.message and str(agent) in record.message) for record in caplog.records)

# Test raise function # Test raise function
with pytest.raises(ToolExecutionException): with pytest.raises(ToolExecutionException):
await runtime.send_message(FunctionCall(id="2", arguments=json.dumps({"input": "raise"}), name="raise"), agent) await runtime.send_message(FunctionCall(id="2", arguments=json.dumps({"input": "raise"}), name="raise"), agent)


+ 0
- 8
python/packages/autogen-ext/src/autogen_ext/models/anthropic/_anthropic_client.py View File

@@ -44,7 +44,6 @@ from autogen_core import (
Component, Component,
FunctionCall, FunctionCall,
Image, Image,
MessageHandlerContext,
) )
from autogen_core.logging import LLMCallEvent from autogen_core.logging import LLMCallEvent
from autogen_core.models import ( from autogen_core.models import (
@@ -503,19 +502,12 @@ class BaseAnthropicChatCompletionClient(ChatCompletionClient):
completion_tokens=result.usage.output_tokens, completion_tokens=result.usage.output_tokens,
) )


# Log the event if in a handler context
try:
agent_id = MessageHandlerContext.agent_id()
except RuntimeError:
agent_id = None

logger.info( logger.info(
LLMCallEvent( LLMCallEvent(
messages=cast(Dict[str, Any], anthropic_messages), messages=cast(Dict[str, Any], anthropic_messages),
response=result.model_dump(), response=result.model_dump(),
prompt_tokens=usage.prompt_tokens, prompt_tokens=usage.prompt_tokens,
completion_tokens=usage.completion_tokens, completion_tokens=usage.completion_tokens,
agent_id=agent_id,
) )
) )




+ 394
- 0
python/packages/autogen-ext/src/autogen_ext/models/llama_cpp/_llama_cpp_completion_client.py View File

@@ -0,0 +1,394 @@
import logging # added import
import re
from typing import Any, AsyncGenerator, Dict, List, Literal, Mapping, Optional, Sequence, TypedDict, Union, cast

from autogen_core import EVENT_LOGGER_NAME, CancellationToken, FunctionCall
from autogen_core.logging import LLMCallEvent
from autogen_core.models import (
AssistantMessage,
ChatCompletionClient,
CreateResult,
FinishReasons,
FunctionExecutionResultMessage,
LLMMessage,
ModelInfo,
RequestUsage,
SystemMessage,
UserMessage,
validate_model_info,
)
from autogen_core.tools import Tool, ToolSchema
from llama_cpp import (
ChatCompletionFunctionParameters,
ChatCompletionRequestAssistantMessage,
ChatCompletionRequestFunctionMessage,
ChatCompletionRequestSystemMessage,
ChatCompletionRequestToolMessage,
ChatCompletionRequestUserMessage,
ChatCompletionTool,
ChatCompletionToolFunction,
Llama,
llama_chat_format,
)
from typing_extensions import Unpack

logger = logging.getLogger(EVENT_LOGGER_NAME) # initialize logger


def normalize_stop_reason(stop_reason: str | None) -> FinishReasons:
if stop_reason is None:
return "unknown"

# Convert to lower case
stop_reason = stop_reason.lower()

KNOWN_STOP_MAPPINGS: Dict[str, FinishReasons] = {
"stop": "stop",
"length": "length",
"content_filter": "content_filter",
"function_calls": "function_calls",
"end_turn": "stop",
"tool_calls": "function_calls",
}

return KNOWN_STOP_MAPPINGS.get(stop_reason, "unknown")


def normalize_name(name: str) -> str:
"""
LLMs sometimes ask functions while ignoring their own format requirements, this function should be used to replace invalid characters with "_".

Prefer _assert_valid_name for validating user configuration or input
"""
return re.sub(r"[^a-zA-Z0-9_-]", "_", name)[:64]


def assert_valid_name(name: str) -> str:
"""
Ensure that configured names are valid, raises ValueError if not.

For munging LLM responses use _normalize_name to ensure LLM specified names don't break the API.
"""
if not re.match(r"^[a-zA-Z0-9_-]+$", name):
raise ValueError(f"Invalid name: {name}. Only letters, numbers, '_' and '-' are allowed.")
if len(name) > 64:
raise ValueError(f"Invalid name: {name}. Name must be less than 64 characters.")
return name


def convert_tools(
tools: Sequence[Tool | ToolSchema],
) -> List[ChatCompletionTool]:
result: List[ChatCompletionTool] = []
for tool in tools:
if isinstance(tool, Tool):
tool_schema = tool.schema
else:
assert isinstance(tool, dict)
tool_schema = tool

result.append(
ChatCompletionTool(
type="function",
function=ChatCompletionToolFunction(
name=tool_schema["name"],
description=(tool_schema["description"] if "description" in tool_schema else ""),
parameters=(
cast(ChatCompletionFunctionParameters, tool_schema["parameters"])
if "parameters" in tool_schema
else {}
),
),
)
)
# Check if all tools have valid names.
for tool_param in result:
assert_valid_name(tool_param["function"]["name"])
return result


class LlamaCppParams(TypedDict, total=False):
# from_pretrained parameters:
repo_id: Optional[str]
filename: Optional[str]
additional_files: Optional[List[Any]]
local_dir: Optional[str]
local_dir_use_symlinks: Union[bool, Literal["auto"]]
cache_dir: Optional[str]
# __init__ parameters:
model_path: str
n_gpu_layers: int
split_mode: int
main_gpu: int
tensor_split: Optional[List[float]]
rpc_servers: Optional[str]
vocab_only: bool
use_mmap: bool
use_mlock: bool
kv_overrides: Optional[Dict[str, Union[bool, int, float, str]]]
seed: int
n_ctx: int
n_batch: int
n_ubatch: int
n_threads: Optional[int]
n_threads_batch: Optional[int]
rope_scaling_type: Optional[int]
pooling_type: int
rope_freq_base: float
rope_freq_scale: float
yarn_ext_factor: float
yarn_attn_factor: float
yarn_beta_fast: float
yarn_beta_slow: float
yarn_orig_ctx: int
logits_all: bool
embedding: bool
offload_kqv: bool
flash_attn: bool
no_perf: bool
last_n_tokens_size: int
lora_base: Optional[str]
lora_scale: float
lora_path: Optional[str]
numa: Union[bool, int]
chat_format: Optional[str]
chat_handler: Optional[llama_chat_format.LlamaChatCompletionHandler]
draft_model: Optional[Any] # LlamaDraftModel not exposed by llama_cpp
tokenizer: Optional[Any] # BaseLlamaTokenizer not exposed by llama_cpp
type_k: Optional[int]
type_v: Optional[int]
spm_infill: bool
verbose: bool


class LlamaCppChatCompletionClient(ChatCompletionClient):
"""Chat completion client for LlamaCpp models.
To use this client, you must install the `llama-cpp` extra:

.. code-block:: bash

pip install "autogen-ext[llama-cpp]"

This client allows you to interact with LlamaCpp models, either by specifying a local model path or by downloading a model from Hugging Face Hub.

Args:
model_path (optional, str): The path to the LlamaCpp model file. Required if repo_id and filename are not provided.
repo_id (optional, str): The Hugging Face Hub repository ID. Required if model_path is not provided.
filename (optional, str): The filename of the model within the Hugging Face Hub repository. Required if model_path is not provided.
n_gpu_layers (optional, int): The number of layers to put on the GPU.
n_ctx (optional, int): The context size.
n_batch (optional, int): The batch size.
verbose (optional, bool): Whether to print verbose output.
model_info (optional, ModelInfo): The capabilities of the model. Defaults to a ModelInfo instance with function_calling set to True.
**kwargs: Additional parameters to pass to the Llama class.

Examples:

The following code snippet shows how to use the client with a local model file:

.. code-block:: python

llama_client = LlamaCppChatCompletionClient(
model_path="/path/to/your/model.gguf",
result = await llama_client.create([UserMessage(content="What is the capital of France?", source="user")]) # type: ignore
print(result)

The following code snippet shows how to use the client with a model from Hugging Face Hub:

.. code-block:: python

llama_client = LlamaCppChatCompletionClient(
repo_id="TheBloke/Mistral-7B-Instruct-v0.1-GGUF",
filename="mistral-7b-instruct-v0.1.Q4_K_M.gguf",
result = await llama_client.create([UserMessage(content="What is the capital of France?", source="user")]) # type: ignore
print(result)
"""

def __init__(
self,
model_info: Optional[ModelInfo] = None,
**kwargs: Unpack[LlamaCppParams],
) -> None:
"""
Initialize the LlamaCpp client.
"""

if model_info:
validate_model_info(model_info)

if "repo_id" in kwargs and "filename" in kwargs and kwargs["repo_id"] and kwargs["filename"]:
repo_id: str = cast(str, kwargs.pop("repo_id"))
filename: str = cast(str, kwargs.pop("filename"))
pretrained = Llama.from_pretrained(repo_id=repo_id, filename=filename, **kwargs) # type: ignore
assert isinstance(pretrained, Llama)
self.llm = pretrained

elif "model_path" in kwargs:
self.llm = Llama(**kwargs) # pyright: ignore[reportUnknownMemberType]
else:
raise ValueError("Please provide model_path if ... or provide repo_id and filename if ....")
self._total_usage = {"prompt_tokens": 0, "completion_tokens": 0}

async def create(
self,
messages: Sequence[LLMMessage],
*,
tools: Sequence[Tool | ToolSchema] = [],
# None means do not override the default
# A value means to override the client default - often specified in the constructor
json_output: Optional[bool] = None,
extra_create_args: Mapping[str, Any] = {},
cancellation_token: Optional[CancellationToken] = None,
) -> CreateResult:
# Convert LLMMessage objects to dictionaries with 'role' and 'content'
# converted_messages: List[Dict[str, str | Image | list[str | Image] | list[FunctionCall]]] = []
converted_messages: list[
ChatCompletionRequestSystemMessage
| ChatCompletionRequestUserMessage
| ChatCompletionRequestAssistantMessage
| ChatCompletionRequestUserMessage
| ChatCompletionRequestToolMessage
| ChatCompletionRequestFunctionMessage
] = []
for msg in messages:
if isinstance(msg, SystemMessage):
converted_messages.append({"role": "system", "content": msg.content})
elif isinstance(msg, UserMessage) and isinstance(msg.content, str):
converted_messages.append({"role": "user", "content": msg.content})
elif isinstance(msg, AssistantMessage) and isinstance(msg.content, str):
converted_messages.append({"role": "assistant", "content": msg.content})
elif (
isinstance(msg, SystemMessage) or isinstance(msg, UserMessage) or isinstance(msg, AssistantMessage)
) and isinstance(msg.content, list):
raise ValueError("Multi-part messages such as those containing images are currently not supported.")
else:
raise ValueError(f"Unsupported message type: {type(msg)}")

if self.model_info["function_calling"]:
response = self.llm.create_chat_completion(
messages=converted_messages, tools=convert_tools(tools), stream=False
)
else:
response = self.llm.create_chat_completion(messages=converted_messages, stream=False)

if not isinstance(response, dict):
raise ValueError("Unexpected response type from LlamaCpp model.")

self._total_usage["prompt_tokens"] += response["usage"]["prompt_tokens"]
self._total_usage["completion_tokens"] += response["usage"]["completion_tokens"]

# Parse the response
response_tool_calls: ChatCompletionTool | None = None
response_text: str | None = None
if "choices" in response and len(response["choices"]) > 0:
if "message" in response["choices"][0]:
response_text = response["choices"][0]["message"]["content"]
if "tool_calls" in response["choices"][0]:
response_tool_calls = response["choices"][0]["tool_calls"] # type: ignore

content: List[FunctionCall] | str = ""
thought: str | None = None
if response_tool_calls:
content = []
for tool_call in response_tool_calls:
if not isinstance(tool_call, dict):
raise ValueError("Unexpected tool call type from LlamaCpp model.")
content.append(
FunctionCall(
id=tool_call["id"],
arguments=tool_call["function"]["arguments"],
name=normalize_name(tool_call["function"]["name"]),
)
)
if response_text and len(response_text) > 0:
thought = response_text
else:
if response_text:
content = response_text

# Detect tool usage in the response
if not response_tool_calls and not response_text:
logger.debug("DEBUG: No response text found. Returning empty response.")
return CreateResult(
content="", usage=RequestUsage(prompt_tokens=0, completion_tokens=0), finish_reason="stop", cached=False
)

# Create a CreateResult object
if "finish_reason" in response["choices"][0]:
finish_reason = response["choices"][0]["finish_reason"]
else:
finish_reason = "unknown"
if finish_reason not in ("stop", "length", "function_calls", "content_filter", "unknown"):
finish_reason = "unknown"
create_result = CreateResult(
content=content,
thought=thought,
usage=cast(RequestUsage, response["usage"]),
finish_reason=normalize_stop_reason(finish_reason), # type: ignore
cached=False,
)

logger.info(
LLMCallEvent(
messages=cast(Dict[str, Any], messages),
response=create_result.model_dump(),
prompt_tokens=response["usage"]["prompt_tokens"],
completion_tokens=response["usage"]["completion_tokens"],
)
)
return create_result

async def create_stream(
self,
messages: Sequence[LLMMessage],
*,
tools: Sequence[Tool | ToolSchema] = [],
# None means do not override the default
# A value means to override the client default - often specified in the constructor
json_output: Optional[bool] = None,
extra_create_args: Mapping[str, Any] = {},
cancellation_token: Optional[CancellationToken] = None,
) -> AsyncGenerator[Union[str, CreateResult], None]:
raise NotImplementedError("Stream not yet implemented for LlamaCppChatCompletionClient")
yield ""

# Implement abstract methods
def actual_usage(self) -> RequestUsage:
return RequestUsage(
prompt_tokens=self._total_usage.get("prompt_tokens", 0),
completion_tokens=self._total_usage.get("completion_tokens", 0),
)

@property
def capabilities(self) -> ModelInfo:
return self.model_info

def count_tokens(
self,
messages: Sequence[SystemMessage | UserMessage | AssistantMessage | FunctionExecutionResultMessage],
**kwargs: Any,
) -> int:
total = 0
for msg in messages:
# Use the Llama model's tokenizer to encode the content
tokens = self.llm.tokenize(str(msg.content).encode("utf-8"))
total += len(tokens)
return total

@property
def model_info(self) -> ModelInfo:
return ModelInfo(vision=False, json_output=False, family="llama-cpp", function_calling=True)

def remaining_tokens(
self,
messages: Sequence[SystemMessage | UserMessage | AssistantMessage | FunctionExecutionResultMessage],
**kwargs: Any,
) -> int:
used_tokens = self.count_tokens(messages)
return max(self.llm.n_ctx() - used_tokens, 0)

def total_usage(self) -> RequestUsage:
return RequestUsage(
prompt_tokens=self._total_usage.get("prompt_tokens", 0),
completion_tokens=self._total_usage.get("completion_tokens", 0),
)

+ 0
- 8
python/packages/autogen-ext/src/autogen_ext/models/ollama/_ollama_client.py View File

@@ -26,7 +26,6 @@ from autogen_core import (
Component, Component,
FunctionCall, FunctionCall,
Image, Image,
MessageHandlerContext,
) )
from autogen_core.logging import LLMCallEvent from autogen_core.logging import LLMCallEvent
from autogen_core.models import ( from autogen_core.models import (
@@ -483,19 +482,12 @@ class BaseOllamaChatCompletionClient(ChatCompletionClient):
completion_tokens=(result.eval_count if result.eval_count is not None else 0), completion_tokens=(result.eval_count if result.eval_count is not None else 0),
) )


# If we are running in the context of a handler we can get the agent_id
try:
agent_id = MessageHandlerContext.agent_id()
except RuntimeError:
agent_id = None

logger.info( logger.info(
LLMCallEvent( LLMCallEvent(
messages=cast(Dict[str, Any], ollama_messages), messages=cast(Dict[str, Any], ollama_messages),
response=result.model_dump(), response=result.model_dump(),
prompt_tokens=usage.prompt_tokens, prompt_tokens=usage.prompt_tokens,
completion_tokens=usage.completion_tokens, completion_tokens=usage.completion_tokens,
agent_id=agent_id,
) )
) )




+ 0
- 8
python/packages/autogen-ext/src/autogen_ext/models/openai/_openai_client.py View File

@@ -29,7 +29,6 @@ from autogen_core import (
Component, Component,
FunctionCall, FunctionCall,
Image, Image,
MessageHandlerContext,
) )
from autogen_core.logging import LLMCallEvent from autogen_core.logging import LLMCallEvent
from autogen_core.models import ( from autogen_core.models import (
@@ -531,19 +530,12 @@ class BaseOpenAIChatCompletionClient(ChatCompletionClient):
completion_tokens=(result.usage.completion_tokens if result.usage is not None else 0), completion_tokens=(result.usage.completion_tokens if result.usage is not None else 0),
) )


# If we are running in the context of a handler we can get the agent_id
try:
agent_id = MessageHandlerContext.agent_id()
except RuntimeError:
agent_id = None

logger.info( logger.info(
LLMCallEvent( LLMCallEvent(
messages=cast(Dict[str, Any], oai_messages), messages=cast(Dict[str, Any], oai_messages),
response=result.model_dump(), response=result.model_dump(),
prompt_tokens=usage.prompt_tokens, prompt_tokens=usage.prompt_tokens,
completion_tokens=usage.completion_tokens, completion_tokens=usage.completion_tokens,
agent_id=agent_id,
) )
) )




+ 1
- 1
python/packages/autogen-ext/src/autogen_ext/models/replay/_replay_chat_completion_client.py View File

@@ -3,7 +3,6 @@ from __future__ import annotations
import logging import logging
import warnings import warnings
from typing import Any, AsyncGenerator, Dict, List, Mapping, Optional, Sequence, Union from typing import Any, AsyncGenerator, Dict, List, Mapping, Optional, Sequence, Union
from typing_extensions import Self


from autogen_core import EVENT_LOGGER_NAME, CancellationToken, Component from autogen_core import EVENT_LOGGER_NAME, CancellationToken, Component
from autogen_core.models import ( from autogen_core.models import (
@@ -18,6 +17,7 @@ from autogen_core.models import (
) )
from autogen_core.tools import Tool, ToolSchema from autogen_core.tools import Tool, ToolSchema
from pydantic import BaseModel from pydantic import BaseModel
from typing_extensions import Self


logger = logging.getLogger(EVENT_LOGGER_NAME) logger = logging.getLogger(EVENT_LOGGER_NAME)




+ 17
- 2
python/packages/autogen-ext/src/autogen_ext/tools/code_execution/_code_execution.py View File

@@ -1,9 +1,14 @@
from autogen_core import CancellationToken, Component, ComponentModel
import logging

from autogen_core import EVENT_LOGGER_NAME, CancellationToken, Component, ComponentModel
from autogen_core.code_executor import CodeBlock, CodeExecutor from autogen_core.code_executor import CodeBlock, CodeExecutor
from autogen_core.logging import ToolCallEvent
from autogen_core.tools import BaseTool from autogen_core.tools import BaseTool
from pydantic import BaseModel, Field, model_serializer from pydantic import BaseModel, Field, model_serializer
from typing_extensions import Self from typing_extensions import Self


logger = logging.getLogger(EVENT_LOGGER_NAME)



class CodeExecutionInput(BaseModel): class CodeExecutionInput(BaseModel):
code: str = Field(description="The contents of the Python code block that should be executed") code: str = Field(description="The contents of the Python code block that should be executed")
@@ -84,7 +89,17 @@ class PythonCodeExecutionTool(
code_blocks=code_blocks, cancellation_token=cancellation_token code_blocks=code_blocks, cancellation_token=cancellation_token
) )


return CodeExecutionResult(success=result.exit_code == 0, output=result.output)
exec_result = CodeExecutionResult(success=result.exit_code == 0, output=result.output)

# Log the event
event = ToolCallEvent(
tool_name=self.name,
arguments=args.model_dump(),
result=exec_result.model_dump(),
)
logger.info(event)

return exec_result


def _to_config(self) -> PythonCodeExecutionToolConfig: def _to_config(self) -> PythonCodeExecutionToolConfig:
"""Convert current instance to config object""" """Convert current instance to config object"""


+ 16
- 4
python/packages/autogen-ext/src/autogen_ext/tools/graphrag/_global_search.py View File

@@ -1,9 +1,11 @@
# mypy: disable-error-code="no-any-unimported,misc" # mypy: disable-error-code="no-any-unimported,misc"
import logging
from pathlib import Path from pathlib import Path


import pandas as pd import pandas as pd
import tiktoken import tiktoken
from autogen_core import CancellationToken
from autogen_core import EVENT_LOGGER_NAME, CancellationToken
from autogen_core.logging import ToolCallEvent
from autogen_core.tools import BaseTool from autogen_core.tools import BaseTool
from graphrag.config.config_file_loader import load_config_from_file from graphrag.config.config_file_loader import load_config_from_file
from graphrag.query.indexer_adapters import ( from graphrag.query.indexer_adapters import (
@@ -24,6 +26,8 @@ from ._config import MapReduceConfig
_default_context_config = ContextConfig() _default_context_config = ContextConfig()
_default_mapreduce_config = MapReduceConfig() _default_mapreduce_config = MapReduceConfig()


logger = logging.getLogger(EVENT_LOGGER_NAME)



class GlobalSearchToolArgs(BaseModel): class GlobalSearchToolArgs(BaseModel):
query: str = Field(..., description="The user query to perform global search on.") query: str = Field(..., description="The user query to perform global search on.")
@@ -177,9 +181,17 @@ class GlobalSearchTool(BaseTool[GlobalSearchToolArgs, GlobalSearchToolReturn]):
) )


async def run(self, args: GlobalSearchToolArgs, cancellation_token: CancellationToken) -> GlobalSearchToolReturn: async def run(self, args: GlobalSearchToolArgs, cancellation_token: CancellationToken) -> GlobalSearchToolReturn:
result = await self._search_engine.asearch(args.query)
assert isinstance(result.response, str), "Expected response to be a string"
return GlobalSearchToolReturn(answer=result.response)
search_result = await self._search_engine.asearch(args.query)
assert isinstance(search_result.response, str), "Expected response to be a string"
result = GlobalSearchToolReturn(answer=search_result.response)
# Log the event
event = ToolCallEvent(
tool_name=self.name,
arguments=args.model_dump(),
result=result.model_dump(),
)
logger.info(event)
return result


@classmethod @classmethod
def from_settings(cls, settings_path: str | Path) -> "GlobalSearchTool": def from_settings(cls, settings_path: str | Path) -> "GlobalSearchTool":


+ 17
- 4
python/packages/autogen-ext/src/autogen_ext/tools/graphrag/_local_search.py View File

@@ -1,10 +1,12 @@
# mypy: disable-error-code="no-any-unimported,misc" # mypy: disable-error-code="no-any-unimported,misc"
import logging
import os import os
from pathlib import Path from pathlib import Path


import pandas as pd import pandas as pd
import tiktoken import tiktoken
from autogen_core import CancellationToken
from autogen_core import EVENT_LOGGER_NAME, CancellationToken
from autogen_core.logging import ToolCallEvent
from autogen_core.tools import BaseTool from autogen_core.tools import BaseTool
from graphrag.config.config_file_loader import load_config_from_file from graphrag.config.config_file_loader import load_config_from_file
from graphrag.query.indexer_adapters import ( from graphrag.query.indexer_adapters import (
@@ -25,6 +27,8 @@ from ._config import LocalDataConfig as DataConfig
_default_context_config = LocalContextConfig() _default_context_config = LocalContextConfig()
_default_search_config = SearchConfig() _default_search_config = SearchConfig()


logger = logging.getLogger(EVENT_LOGGER_NAME)



class LocalSearchToolArgs(BaseModel): class LocalSearchToolArgs(BaseModel):
query: str = Field(..., description="The user query to perform local search on.") query: str = Field(..., description="The user query to perform local search on.")
@@ -188,9 +192,18 @@ class LocalSearchTool(BaseTool[LocalSearchToolArgs, LocalSearchToolReturn]):
) )


async def run(self, args: LocalSearchToolArgs, cancellation_token: CancellationToken) -> LocalSearchToolReturn: async def run(self, args: LocalSearchToolArgs, cancellation_token: CancellationToken) -> LocalSearchToolReturn:
result = await self._search_engine.asearch(args.query) # type: ignore
assert isinstance(result.response, str), "Expected response to be a string"
return LocalSearchToolReturn(answer=result.response)
search_result = await self._search_engine.asearch(args.query) # type: ignore
assert isinstance(search_result.response, str), "Expected response to be a string"
result = LocalSearchToolReturn(answer=search_result.response)
# Log the tool call event
logger.info(
ToolCallEvent(
tool_name=self.name,
arguments=args.model_dump(),
result=result.model_dump(),
)
)
return result


@classmethod @classmethod
def from_settings(cls, settings_path: str | Path) -> "LocalSearchTool": def from_settings(cls, settings_path: str | Path) -> "LocalSearchTool":


+ 12
- 3
python/packages/autogen-ext/src/autogen_ext/tools/http/_http_tool.py View File

@@ -1,13 +1,17 @@
import logging
import re import re
from typing import Any, Literal, Optional, Type from typing import Any, Literal, Optional, Type


import httpx import httpx
from autogen_core import CancellationToken, Component
from autogen_core import EVENT_LOGGER_NAME, CancellationToken, Component
from autogen_core.logging import ToolCallEvent
from autogen_core.tools import BaseTool from autogen_core.tools import BaseTool
from json_schema_to_pydantic import create_model from json_schema_to_pydantic import create_model
from pydantic import BaseModel, Field from pydantic import BaseModel, Field
from typing_extensions import Self from typing_extensions import Self


logger = logging.getLogger(EVENT_LOGGER_NAME)



class HttpToolConfig(BaseModel): class HttpToolConfig(BaseModel):
name: str name: str
@@ -224,10 +228,15 @@ class HttpTool(BaseTool[BaseModel, Any], Component[HttpToolConfig]):
case _: # Default case POST case _: # Default case POST
response = await client.post(url, headers=self.server_params.headers, json=model_dump) response = await client.post(url, headers=self.server_params.headers, json=model_dump)


result: Any = None
match self.server_params.return_type: match self.server_params.return_type:
case "text": case "text":
return response.text
result = response.text
case "json": case "json":
return response.json()
result = response.json()
case _: case _:
raise ValueError(f"Invalid return type: {self.server_params.return_type}") raise ValueError(f"Invalid return type: {self.server_params.return_type}")
# Log the event
event = ToolCallEvent(tool_name=self.name, arguments=args.model_dump(), result=result)
logger.info(event)
return result

+ 20
- 1
python/packages/autogen-ext/src/autogen_ext/tools/langchain/_langchain_adapter.py View File

@@ -2,15 +2,19 @@ from __future__ import annotations


import asyncio import asyncio
import inspect import inspect
import logging
from typing import TYPE_CHECKING, Any, Callable, Dict, Type, cast from typing import TYPE_CHECKING, Any, Callable, Dict, Type, cast


from autogen_core import CancellationToken
from autogen_core import EVENT_LOGGER_NAME, CancellationToken
from autogen_core.logging import ToolCallEvent
from autogen_core.tools import BaseTool from autogen_core.tools import BaseTool
from pydantic import BaseModel, Field, create_model from pydantic import BaseModel, Field, create_model


if TYPE_CHECKING: if TYPE_CHECKING:
from langchain_core.tools import BaseTool as LangChainTool from langchain_core.tools import BaseTool as LangChainTool


logger = logging.getLogger(EVENT_LOGGER_NAME)



class LangChainToolAdapter(BaseTool[BaseModel, Any]): class LangChainToolAdapter(BaseTool[BaseModel, Any]):
"""Allows you to wrap a LangChain tool and make it available to AutoGen. """Allows you to wrap a LangChain tool and make it available to AutoGen.
@@ -194,6 +198,21 @@ class LangChainToolAdapter(BaseTool[BaseModel, Any]):
# Run in a thread to avoid blocking the event loop # Run in a thread to avoid blocking the event loop
result = await asyncio.to_thread(self._call_sync, kwargs) result = await asyncio.to_thread(self._call_sync, kwargs)


# Log the event
serializable_result: Any = None
if isinstance(result, BaseModel):
serializable_result = result.model_dump()
elif isinstance(result, str):
serializable_result = result
else:
serializable_result = str(result)
event = ToolCallEvent(
tool_name=self.name,
arguments=args.model_dump(),
result=serializable_result,
)
logger.info(event)

return result return result


def _call_sync(self, kwargs: Dict[str, Any]) -> Any: def _call_sync(self, kwargs: Dict[str, Any]) -> Any:


+ 13
- 1
python/packages/autogen-ext/src/autogen_ext/tools/mcp/_base.py View File

@@ -1,7 +1,9 @@
import logging
from abc import ABC from abc import ABC
from typing import Any, Generic, Type, TypeVar from typing import Any, Generic, Type, TypeVar


from autogen_core import CancellationToken
from autogen_core import EVENT_LOGGER_NAME, CancellationToken
from autogen_core.logging import ToolCallEvent
from autogen_core.tools import BaseTool from autogen_core.tools import BaseTool
from json_schema_to_pydantic import create_model from json_schema_to_pydantic import create_model
from mcp import Tool from mcp import Tool
@@ -12,6 +14,8 @@ from ._session import create_mcp_server_session


TServerParams = TypeVar("TServerParams", bound=McpServerParams) TServerParams = TypeVar("TServerParams", bound=McpServerParams)


logger = logging.getLogger(EVENT_LOGGER_NAME)



class McpToolAdapter(BaseTool[BaseModel, Any], ABC, Generic[TServerParams]): class McpToolAdapter(BaseTool[BaseModel, Any], ABC, Generic[TServerParams]):
""" """
@@ -68,6 +72,14 @@ class McpToolAdapter(BaseTool[BaseModel, Any], ABC, Generic[TServerParams]):
if result.isError: if result.isError:
raise Exception(f"MCP tool execution failed: {result.content}") raise Exception(f"MCP tool execution failed: {result.content}")


event = ToolCallEvent(
tool_name=self.name,
arguments=kwargs,
result=result.content,
# TODO: add other relevant fields to the MCP tool event.
)
logger.info(event)

return result.content return result.content
except Exception as e: except Exception as e:
raise Exception(str(e)) from e raise Exception(str(e)) from e


+ 22
- 12
python/packages/autogen-ext/tests/tools/graphrag/test_graphrag_tools.py View File

@@ -91,6 +91,7 @@ async def test_global_search_tool(
entity_df_fixture: pd.DataFrame, entity_df_fixture: pd.DataFrame,
report_df_fixture: pd.DataFrame, report_df_fixture: pd.DataFrame,
entity_embedding_fixture: pd.DataFrame, entity_embedding_fixture: pd.DataFrame,
caplog: pytest.LogCaptureFixture,
) -> None: ) -> None:
# Create a temporary directory to simulate the data config # Create a temporary directory to simulate the data config
with tempfile.TemporaryDirectory() as tempdir: with tempfile.TemporaryDirectory() as tempdir:
@@ -120,12 +121,16 @@ async def test_global_search_tool(


global_search_tool = GlobalSearchTool(token_encoder=token_encoder, llm=llm, data_config=data_config) global_search_tool = GlobalSearchTool(token_encoder=token_encoder, llm=llm, data_config=data_config)


# Example of running the tool and checking the result
query = "What is the overall sentiment of the community reports?"
cancellation_token = CancellationToken()
result = await global_search_tool.run_json(args={"query": query}, cancellation_token=cancellation_token)
assert isinstance(result, GlobalSearchToolReturn)
assert isinstance(result.answer, str)
with caplog.at_level("INFO"):
# Example of running the tool and checking the result
query = "What is the overall sentiment of the community reports?"
cancellation_token = CancellationToken()
result = await global_search_tool.run_json(args={"query": query}, cancellation_token=cancellation_token)
assert isinstance(result, GlobalSearchToolReturn)
assert isinstance(result.answer, str)

# Check if the log contains the expected message
assert result.answer in caplog.text




@pytest.mark.asyncio @pytest.mark.asyncio
@@ -135,6 +140,7 @@ async def test_local_search_tool(
text_unit_df_fixture: pd.DataFrame, text_unit_df_fixture: pd.DataFrame,
entity_embedding_fixture: pd.DataFrame, entity_embedding_fixture: pd.DataFrame,
monkeypatch: pytest.MonkeyPatch, monkeypatch: pytest.MonkeyPatch,
caplog: pytest.LogCaptureFixture,
) -> None: ) -> None:
# Create a temporary directory to simulate the data config # Create a temporary directory to simulate the data config
with tempfile.TemporaryDirectory() as tempdir: with tempfile.TemporaryDirectory() as tempdir:
@@ -176,9 +182,13 @@ async def test_local_search_tool(
token_encoder=token_encoder, llm=llm, embedder=embedder, data_config=data_config token_encoder=token_encoder, llm=llm, embedder=embedder, data_config=data_config
) )


# Example of running the tool and checking the result
query = "What are the relationships between Dr. Becher and the station-master?"
cancellation_token = CancellationToken()
result = await local_search_tool.run_json(args={"query": query}, cancellation_token=cancellation_token)
assert isinstance(result, LocalSearchToolReturn)
assert isinstance(result.answer, str)
with caplog.at_level("INFO"):
# Example of running the tool and checking the result
query = "What are the relationships between Dr. Becher and the station-master?"
cancellation_token = CancellationToken()
result = await local_search_tool.run_json(args={"query": query}, cancellation_token=cancellation_token)
assert isinstance(result, LocalSearchToolReturn)
assert isinstance(result.answer, str)

# Check if the log contains the expected message
assert result.answer in caplog.text

+ 9
- 4
python/packages/autogen-ext/tests/tools/http/test_http_tool.py View File

@@ -1,4 +1,5 @@
import json import json
import logging


import httpx import httpx
import pytest import pytest
@@ -45,12 +46,16 @@ def test_component_base_class(test_config: ComponentModel) -> None:




@pytest.mark.asyncio @pytest.mark.asyncio
async def test_post_request(test_config: ComponentModel, test_server: None) -> None:
async def test_post_request(test_config: ComponentModel, test_server: None, caplog: pytest.LogCaptureFixture) -> None:
tool = HttpTool.load_component(test_config) tool = HttpTool.load_component(test_config)
result = await tool.run_json({"query": "test query", "value": 42}, CancellationToken())


assert isinstance(result, str)
assert json.loads(result)["result"] == "Received: test query with value 42"
with caplog.at_level(logging.INFO):
result = await tool.run_json({"query": "test query", "value": 42}, CancellationToken())

assert isinstance(result, str)
assert json.loads(result)["result"] == "Received: test query with value 42"

assert "Received: test query with value 42" in caplog.text




@pytest.mark.asyncio @pytest.mark.asyncio


+ 8
- 4
python/packages/autogen-ext/tests/tools/test_langchain_tools.py View File

@@ -1,3 +1,4 @@
import logging
from typing import Optional, Type, cast from typing import Optional, Type, cast


import pytest import pytest
@@ -42,7 +43,7 @@ class CustomCalculatorTool(LangChainTool):




@pytest.mark.asyncio @pytest.mark.asyncio
async def test_langchain_tool_adapter() -> None:
async def test_langchain_tool_adapter(caplog: pytest.LogCaptureFixture) -> None:
# Create a LangChain tool # Create a LangChain tool
langchain_tool = add # type: ignore langchain_tool = add # type: ignore


@@ -66,9 +67,12 @@ async def test_langchain_tool_adapter() -> None:
assert set(schema["parameters"]["required"]) == {"a", "b"} assert set(schema["parameters"]["required"]) == {"a", "b"}
assert len(schema["parameters"]["properties"]) == 2 assert len(schema["parameters"]["properties"]) == 2


# Test run method
result = await adapter.run_json({"a": 2, "b": 3}, CancellationToken())
assert result == 5
# Check log.
with caplog.at_level(logging.INFO):
# Test run method
result = await adapter.run_json({"a": 2, "b": 3}, CancellationToken())
assert result == 5
assert str(result) in caplog.text


# Test that the adapter's run method can be called multiple times # Test that the adapter's run method can be called multiple times
result = await adapter.run_json({"a": 5, "b": 7}, CancellationToken()) result = await adapter.run_json({"a": 5, "b": 7}, CancellationToken())


+ 27
- 16
python/packages/autogen-ext/tests/tools/test_mcp_tools.py View File

@@ -1,3 +1,4 @@
import logging
from unittest.mock import AsyncMock, MagicMock from unittest.mock import AsyncMock, MagicMock


import pytest import pytest
@@ -109,6 +110,7 @@ async def test_mcp_tool_execution(
mock_tool_response: MagicMock, mock_tool_response: MagicMock,
cancellation_token: CancellationToken, cancellation_token: CancellationToken,
monkeypatch: pytest.MonkeyPatch, monkeypatch: pytest.MonkeyPatch,
caplog: pytest.LogCaptureFixture,
) -> None: ) -> None:
"""Test that adapter properly executes tools through ClientSession.""" """Test that adapter properly executes tools through ClientSession."""
mock_context = AsyncMock() mock_context = AsyncMock()
@@ -120,15 +122,19 @@ async def test_mcp_tool_execution(


mock_session.call_tool.return_value = mock_tool_response mock_session.call_tool.return_value = mock_tool_response


adapter = StdioMcpToolAdapter(server_params=sample_server_params, tool=sample_tool)
result = await adapter.run(
args=create_model(sample_tool.inputSchema)(**{"test_param": "test"}),
cancellation_token=cancellation_token,
)
with caplog.at_level(logging.INFO):
adapter = StdioMcpToolAdapter(server_params=sample_server_params, tool=sample_tool)
result = await adapter.run(
args=create_model(sample_tool.inputSchema)(**{"test_param": "test"}),
cancellation_token=cancellation_token,
)

assert result == mock_tool_response.content
mock_session.initialize.assert_called_once()
mock_session.call_tool.assert_called_once()


assert result == mock_tool_response.content
mock_session.initialize.assert_called_once()
mock_session.call_tool.assert_called_once()
# Check log.
assert "test_output" in caplog.text




@pytest.mark.asyncio @pytest.mark.asyncio
@@ -206,6 +212,7 @@ async def test_sse_tool_execution(
sample_sse_tool: Tool, sample_sse_tool: Tool,
mock_sse_session: AsyncMock, mock_sse_session: AsyncMock,
monkeypatch: pytest.MonkeyPatch, monkeypatch: pytest.MonkeyPatch,
caplog: pytest.LogCaptureFixture,
) -> None: ) -> None:
"""Test that SSE adapter properly executes tools through ClientSession.""" """Test that SSE adapter properly executes tools through ClientSession."""
params = SseServerParams(url="http://test-url") params = SseServerParams(url="http://test-url")
@@ -219,15 +226,19 @@ async def test_sse_tool_execution(
lambda *args, **kwargs: mock_context, # type: ignore lambda *args, **kwargs: mock_context, # type: ignore
) )


adapter = SseMcpToolAdapter(server_params=params, tool=sample_sse_tool)
result = await adapter.run(
args=create_model(sample_sse_tool.inputSchema)(**{"test_param": "test"}),
cancellation_token=CancellationToken(),
)
with caplog.at_level(logging.INFO):
adapter = SseMcpToolAdapter(server_params=params, tool=sample_sse_tool)
result = await adapter.run(
args=create_model(sample_sse_tool.inputSchema)(**{"test_param": "test"}),
cancellation_token=CancellationToken(),
)

assert result == mock_sse_session.call_tool.return_value.content
mock_sse_session.initialize.assert_called_once()
mock_sse_session.call_tool.assert_called_once()


assert result == mock_sse_session.call_tool.return_value.content
mock_sse_session.initialize.assert_called_once()
mock_sse_session.call_tool.assert_called_once()
# Check log.
assert "test_output" in caplog.text




@pytest.mark.asyncio @pytest.mark.asyncio


+ 8
- 4
python/packages/autogen-ext/tests/tools/test_python_code_executor_tool.py View File

@@ -1,3 +1,4 @@
import logging
import tempfile import tempfile


import pytest import pytest
@@ -7,7 +8,7 @@ from autogen_ext.tools.code_execution import CodeExecutionInput, PythonCodeExecu




@pytest.mark.asyncio @pytest.mark.asyncio
async def test_python_code_execution_tool() -> None:
async def test_python_code_execution_tool(caplog: pytest.LogCaptureFixture) -> None:
"""Test basic functionality of PythonCodeExecutionTool.""" """Test basic functionality of PythonCodeExecutionTool."""
# Create a temporary directory for the executor # Create a temporary directory for the executor
with tempfile.TemporaryDirectory() as temp_dir: with tempfile.TemporaryDirectory() as temp_dir:
@@ -15,9 +16,12 @@ async def test_python_code_execution_tool() -> None:
executor = LocalCommandLineCodeExecutor(work_dir=temp_dir) executor = LocalCommandLineCodeExecutor(work_dir=temp_dir)
tool = PythonCodeExecutionTool(executor=executor) tool = PythonCodeExecutionTool(executor=executor)


# Test simple code execution
code = "print('hello world!')"
result = await tool.run(args=CodeExecutionInput(code=code), cancellation_token=CancellationToken())
with caplog.at_level(logging.INFO):
# Test simple code execution
code = "print('hello world!')"
result = await tool.run(args=CodeExecutionInput(code=code), cancellation_token=CancellationToken())
# Check log output
assert "hello world!" in caplog.text


# Verify successful execution # Verify successful execution
assert result.success is True assert result.success is True


Loading…
Cancel
Save