| Author | SHA1 | Message | Date |
|---|---|---|---|
|
|
c9a89ae444 | fix | 1 year ago |
|
|
329951d295 | wip | 1 year ago |
|
|
ca57fecaab | Add ToolCallEvent and log it for all built-in tools. | 1 year ago |
| @@ -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 | ||||
| @@ -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: | ||||
| @@ -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) | ||||
| @@ -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, | |||||
| ) | ) | ||||
| ) | ) | ||||
| @@ -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), | |||||
| ) | |||||
| @@ -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, | |||||
| ) | ) | ||||
| ) | ) | ||||
| @@ -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, | |||||
| ) | ) | ||||
| ) | ) | ||||
| @@ -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) | ||||
| @@ -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""" | ||||
| @@ -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": | ||||
| @@ -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": | ||||
| @@ -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 | |||||
| @@ -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: | ||||
| @@ -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 | ||||
| @@ -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 | |||||
| @@ -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 | ||||
| @@ -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()) | ||||
| @@ -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 | ||||
| @@ -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 | ||||