aixtools 0.1.0__py3-none-any.whl

aixtools/testing/aix_test_model.py

@@ -0,0 +1,147 @@
+"""
+Test model implementation for AI agent testing with predefined responses.
+"""
+
+import asyncio
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
+from types import AsyncGeneratorType
+
+from pydantic import BaseModel
+from pydantic_ai.messages import ModelMessage, ModelResponse, TextPart, ToolCallPart
+from pydantic_ai.models import Model, ModelRequestParameters, StreamedResponse
+from pydantic_ai.models.function import _estimate_usage
+from pydantic_ai.models.test import TestStreamedResponse
+from pydantic_ai.settings import ModelSettings
+
+from ..utils.utils import async_iter
+
+FINAL_RESULT_TOOL_NAME = "final_result"
+
+
+def final_result_tool(result: BaseModel | dict) -> ToolCallPart:
+    """Create a ToolCallPart for the final result."""
+    if isinstance(result, BaseModel):
+        args = result.model_dump()
+    else:
+        args = result
+    return ToolCallPart(tool_name=FINAL_RESULT_TOOL_NAME, args=args)
+
+
+class AixTestModel(Model):
+    """
+    Test model, returns a specified list of answers, including messages or tool calls
+    This is used for testing the agent and model interaction with the rest of the system.
+
+    responses: Is a list of strings or ToolCallPart objects that the model will return in order.
+
+    Note: The agent will continue to invoke 'request()' until it returns a txt (i.e. not a ToolCallPart).
+
+    Example: Unstructured output (text)
+    ```
+        model = AixTestModel(
+                    responses=[
+                        ToolCallPart(tool_name='send_message_to_user', args={'message': 'First, let me say hi...'}),
+                        "Please invoke the agent again to continue the conversation....",
+                        # ---------- The first time you invoke the agent, it will stop here ----------
+
+                        ToolCallPart(tool_name='send_message_to_user', args={'message': 'Hi there again!'}),
+                        "The 10th prime number is 29.",
+                        # ---------- The second time you invoke the agent, it will stop here ----------
+
+                        # If you invoke the agent again, it will continue raising an exception
+                        # because there are no more responses
+                    ]
+                )
+    ```
+
+    Example structured output:
+    ```
+        # Define a model for the final result
+        class MyResult(BaseModel):
+            text: str
+
+        model = AixTestModel(
+                    responses=[
+                        ToolCallPart(tool_name='send_message_to_user', args={'message': 'First, let me say hi...'}),
+                        ToolCallPart(tool_name='send_message_to_user', args={'message': 'Hi there again!'}),
+                        final_result_tool(MyResult(text='The 10th prime is 29')),
+                    ]
+                )
+    ```
+    """
+
+    def __init__(  # pylint: disable=super-init-not-called
+        self,
+        responses: list[str | TextPart | ToolCallPart] | AsyncGeneratorType,
+        sleep_time: float | None = None,
+    ):
+        self.response_iter = responses(self) if callable(responses) else async_iter(responses)
+        self.messages: list[ModelMessage] = None
+        self.sleep_time = sleep_time
+        self.last_model_request_parameters: ModelRequestParameters | None = None
+
+    @property
+    def last_message(self) -> ModelResponse | None:
+        """Return the last response."""
+        return self.messages[-1] if self.messages else None
+
+    @property
+    def last_message_part(self) -> TextPart | ToolCallPart | None:
+        """Return the last part of the response."""
+        return self.last_message.parts[-1] if self.last_message and self.last_message.parts else None
+
+    async def request(
+        self,
+        messages: list[ModelMessage],
+        model_settings: ModelSettings | None,
+        model_request_parameters: ModelRequestParameters,
+    ) -> ModelResponse:
+        self.last_model_request_parameters = model_request_parameters
+        model_response = await self._request(messages, model_settings, model_request_parameters)
+        model_response.usage = _estimate_usage([*messages, model_response])
+        return model_response
+
+    @asynccontextmanager
+    async def request_stream(
+        self,
+        messages: list[ModelMessage],
+        model_settings: ModelSettings | None,
+        model_request_parameters: ModelRequestParameters,
+    ) -> AsyncIterator[StreamedResponse]:
+        model_response = await self._request(messages, model_settings, model_request_parameters)
+        yield TestStreamedResponse(_model_name=self.model_name, _structured_response=model_response, _messages=messages)
+
+    @property
+    def model_name(self) -> str:
+        return self.__class__.__name__.lower()
+
+    @property
+    def system(self) -> str:
+        """The system / model provider."""
+        return "test_system"
+
+    async def _request(
+        self,
+        messages: list[ModelMessage],
+        model_settings: ModelSettings | None,  # pylint: disable=unused-argument
+        model_request_parameters: ModelRequestParameters,  # pylint: disable=unused-argument
+    ) -> ModelResponse:
+        self.messages = messages
+        res = await anext(self.response_iter, None)
+        assert res, "No more responses available."
+        if callable(res):
+            res = res(self, messages)
+        match res:
+            case str():
+                return ModelResponse(parts=[TextPart(res)], model_name=self.model_name)
+            case ToolCallPart():
+                return ModelResponse(parts=[res], model_name=self.model_name)
+            case TextPart():
+                return ModelResponse(parts=[res], model_name=self.model_name)
+            case Exception():
+                raise res
+            case _:
+                raise ValueError(f"Invalid response type: {type(res)}, response: {res}")
+        if self.sleep_time:
+            await asyncio.sleep(self.sleep_time)

aixtools/testing/mock_tool.py

@@ -0,0 +1,66 @@
+"""
+Mock tool implementation for testing agent interactions.
+"""
+
+import functools
+import inspect
+from typing import Any, Awaitable, Callable, List, TypeVar, Union
+
+T = TypeVar("T")
+
+
+def mock_tool(func: Callable[..., Union[T, Awaitable[T]]], return_values: List[Any]) -> Callable[..., Any]:
+    """
+    Creates a mock version of the provided function that returns values from a predefined list.
+    Supports both synchronous and asynchronous functions.
+
+    Args:
+        func: The original function to mock (can be sync or async)
+        return_values: A list of values to be returned sequentially on each call
+
+    Returns:
+        A new function with the same signature and docstring as the original function,
+        but returns values from the provided list sequentially.
+        If the original function is async, the mock will also be async.
+
+    Raises:
+        IndexError: When the mock function is called more times than there are return values
+    """
+    # Make a copy of the return values to avoid modifying the original list
+    values = return_values.copy()
+
+    # Check if the function is asynchronous
+    is_async = inspect.iscoroutinefunction(func)
+
+    if is_async:
+        # Create an async wrapper for async functions
+        @functools.wraps(func)
+        async def async_mock_wrapper(*args, **kwargs):  # pylint: disable=unused-argument
+            # Check if we have any return values left
+            if not values:
+                raise IndexError(
+                    f"No more mock return values available for {func.__name__}. "
+                    f"The function has been called more times than the number of provided return values."
+                )
+
+            # Return and remove the first value from our list
+            return values.pop(0)
+
+        # Return the async wrapper function
+        return async_mock_wrapper
+
+    # Create a sync wrapper for sync functions
+    @functools.wraps(func)
+    def sync_mock_wrapper(*args, **kwargs):  # pylint: disable=unused-argument
+        # Check if we have any return values left
+        if not values:
+            raise IndexError(
+                f"No more mock return values available for {func.__name__}. "
+                f"The function has been called more times than the number of provided return values."
+            )
+
+        # Return and remove the first value from our list
+        return values.pop(0)
+
+    # Return the sync wrapper function
+    return sync_mock_wrapper

aixtools/testing/model_patch_cache.py

@@ -0,0 +1,279 @@
+"""
+This module is useful for testing, so that we can run tests without having to call the real function.
+The cached values are stored in `cache_file`.
+
+In `learn=False` mode:
+    - The model_request_cache_wrapper first tries to answer from the cached value.
+    - If the cached value does not exist, it raises an exception.
+
+In `learn=True` mode:
+    - The model_request_cache_wrapper first tries to answer from the cached value.
+    - If the cached value does not exist, it invokes the real function
+      and adds the new result to `cache_file` for future use.
+
+Values are saved to pickle files, but since objects could be non-picklable, we use `safe_deepcopy_for_cache`.
+"""
+
+import datetime
+import functools
+import hashlib
+import json
+import pickle
+import uuid
+from contextlib import asynccontextmanager
+from pathlib import Path
+from typing import Any, Dict, Mapping, Sequence, Set, Tuple, Type
+
+from aixtools.logging.log_objects import safe_deepcopy
+from aixtools.logging.logging_config import get_logger
+from aixtools.model_patch.model_patch import get_request_fn, get_request_stream_fn, model_patch
+
+logger = get_logger(__name__)
+
+
+class CacheKeyError(Exception):
+    """Exception raised when a key is not found in the cache."""
+
+
+def safe_deepcopy_for_cache_key(obj: Any, normalize_types: Set[Type] = None) -> Any:
+    """
+    A modified version of safe_deepcopy that normalizes or skips fields based on their types.
+    This is useful for generating cache keys where some fields (like timestamps, UUIDs)
+    should be normalized to ensure consistent cache hits.
+
+    Args:
+        obj: The object to copy
+        normalize_types: Set of types to normalize (replace with placeholders)
+
+    Returns:
+        A deep copy of the object with normalized values for specified types
+    """
+    # Default types to normalize if none provided
+    if normalize_types is None:
+        normalize_types = {datetime.datetime, datetime.date, datetime.time, uuid.UUID}
+
+    # Check if the object is of a type that should be normalized
+    if any(isinstance(obj, t) for t in normalize_types):
+        # Replace with a placeholder indicating the type
+        return f"<{type(obj).__name__}>"
+
+    # Check if the object is primitive (int, float, str, bool)
+    if isinstance(obj, (int, float, str, bool)):
+        return obj
+
+    # Handle mappings (dict and other mapping types)
+    if isinstance(obj, Mapping):
+        return {k: safe_deepcopy_for_cache_key(v, normalize_types) for k, v in obj.items()}
+
+    # Handle sequences (list, tuple, and other sequence types) but not strings
+    if isinstance(obj, Sequence) and not isinstance(obj, str):
+        return [safe_deepcopy_for_cache_key(item, normalize_types) for item in obj]
+
+    # Handle objects with __dict__ attribute (custom classes)
+    if hasattr(obj, "__dict__"):
+        # For objects, we create a dictionary representation
+        result = {}
+        for attr, value in vars(obj).items():
+            result[attr] = safe_deepcopy_for_cache_key(value, normalize_types)
+        return result
+
+    # For other types, return a string representation
+    return f"<{type(obj).__name__}>"
+
+
+def generate_cache_key(method_name: str, args: Tuple, kwargs: Dict) -> str:
+    """
+    Generate a unique cache key based on the method name and its arguments.
+    Uses safe_deepcopy_for_cache to normalize values that change frequently.
+
+    Args:
+        method_name: Name of the method being called
+        args: Positional arguments to the method
+        kwargs: Keyword arguments to the method
+
+    Returns:
+        A string hash that uniquely identifies this method call
+    """
+    # Normalize the arguments and kwargs
+    normalized_args = safe_deepcopy_for_cache_key(args)
+    normalized_kwargs = safe_deepcopy_for_cache_key(kwargs)
+
+    # Create a dictionary with the normalized information
+    key_dict = {"method_name": method_name, "args": normalized_args, "kwargs": normalized_kwargs}
+
+    # Convert to a stable string representation and hash it
+    key_str = json.dumps(key_dict, sort_keys=True, default=str)
+    return hashlib.md5(key_str.encode()).hexdigest()
+
+
+def save_to_cache(cache_file: Path, key: str, value: Any) -> None:
+    """
+    Save a value to the cache file using the given key.
+
+    Args:
+        cache_file: Path to the cache file
+        key: The cache key
+        value: The value to cache (will be deep-copied to ensure it's pickable)
+    """
+    # Create parent directories if they don't exist
+    cache_file.parent.mkdir(parents=True, exist_ok=True)
+
+    # Load existing cache
+    cache = {}
+    if cache_file.exists():
+        try:
+            with open(cache_file, "rb") as f:
+                cache = pickle.load(f)
+        except (pickle.PickleError, EOFError):
+            # If the file is corrupted, start with an empty cache
+            cache = {}
+
+    # Make a safe copy of the value and add it to the cache
+    safe_value = safe_deepcopy(value)
+    cache[key] = safe_value
+    logger.debug("Cache updated: %s -> %r", key, safe_value)
+
+    # Save the updated cache
+    with open(cache_file, "wb") as f:
+        pickle.dump(cache, f)
+
+
+def get_from_cache(cache_file: Path, key: str) -> Any:
+    """
+    Retrieve a value from the cache file using the given key.
+
+    Args:
+        cache_file: Path to the cache file
+        key: The cache key
+
+    Returns:
+        The cached value
+
+    Raises:
+        CacheKeyError: If the key is not found in the cache
+    """
+    if not cache_file.exists():
+        raise CacheKeyError(f"Cache file {cache_file} does not exist")
+    try:
+        with open(cache_file, "rb") as f:
+            cache = pickle.load(f)
+    except (pickle.PickleError, EOFError) as e:
+        raise CacheKeyError(f"Cache file {cache_file} is corrupted") from e
+    if key not in cache:
+        raise CacheKeyError(f"Key {key} not found in cache")
+    return cache[key]
+
+
+def request_cache(fn, cache_file: Path, learn=False):
+    """
+    model_request_cache_wrapper for async method calls that uses a cache.
+
+    In learn=False mode:
+    - Tries to answer from cached value
+    - If the cached value does not exist, raises an exception
+
+    In learn=True mode:
+    - Tries to answer from cached value
+    - If the cached value does not exist, invokes the real function and adds the result to cache
+
+    Args:
+        fn: The async function to wrap
+        cache_file: Path to the cache file
+        learn: Whether to learn new responses (True) or only use cached ones (False)
+    """
+
+    @functools.wraps(fn)
+    async def model_request_cache_wrapper(*args, **kwargs):
+        # Generate a unique cache key for this request
+        cache_key = generate_cache_key(fn.__name__, args, kwargs)
+        try:
+            # Try to get the result from cache
+            result = get_from_cache(cache_file, cache_key)
+            logger.debug("Cache hit for %s with key %s", fn.__name__, cache_key)
+            return result
+        except CacheKeyError as e:
+            # If not in cache and learn=False, raise an exception
+            if not learn:
+                raise CacheKeyError(f"No cached response for {fn.__name__} with key {cache_key} and learn=False") from e
+            # If learn=True, invoke the original method
+            result = await fn(*args, **kwargs)
+            # Save the result to cache
+            save_to_cache(cache_file, cache_key, result)
+            return result
+
+    return model_request_cache_wrapper
+
+
+def request_stream_cache(fn, cache_file: Path, learn=False):
+    """
+    model_request_cache_wrapper for async streaming method calls that uses a cache.
+
+    Similar to request_cache, but handles streaming responses by caching all items
+    and then replaying them when retrieved from cache.
+
+    Args:
+        fn: The async context manager function to wrap
+        cache_file: Path to the cache file
+        learn: Whether to learn new responses (True) or only use cached ones (False)
+    """
+
+    @functools.wraps(fn)
+    @asynccontextmanager
+    async def model_request_stream_cache_wrapper(*args, **kwargs):
+        # Generate a unique cache key for this request
+        cache_key = generate_cache_key(fn.__name__, args, kwargs)
+
+        try:
+            # Try to get the cached items
+            cached_items = get_from_cache(cache_file, cache_key)
+
+            # Define a generator that yields the cached items
+            async def cached_gen():
+                for item in cached_items:
+                    yield item
+
+            yield cached_gen()
+
+        except CacheKeyError as e:
+            # If not in cache and learn=False, raise an exception
+            if not learn:
+                raise CacheKeyError(f"No cached stream for {fn.__name__} with key {cache_key} and learn=False") from e
+
+            # If learn=True, invoke the original method
+            async with fn(*args, **kwargs) as stream:
+                # Collect all items to save to cache
+                all_items = []
+
+                async def gen():
+                    item_num = 0
+                    async for item in stream:
+                        all_items.append(item)
+                        item_num += 1
+                        yield item
+
+                # Yield the generator
+                gen_instance = gen()
+                yield gen_instance
+
+                # After the context manager exits, save all items to cache
+                # We need to make sure all items have been consumed
+                try:
+                    async for _ in gen_instance:
+                        pass
+                except StopAsyncIteration:
+                    pass
+
+                # Save the collected items to cache
+                save_to_cache(cache_file, cache_key, all_items)
+
+    return model_request_stream_cache_wrapper
+
+
+def model_patch_cache(model, cache_file: Path, learn=False):
+    """Patch model with methods for caching requests and responses"""
+    logger.debug("Using cache file: %s", cache_file)
+    return model_patch(
+        model,
+        request_method=request_cache(get_request_fn(model), cache_file, learn),
+        request_stream_method=request_stream_cache(get_request_stream_fn(model), cache_file, learn),
+    )

aixtools/tools/doctor/__init__.py

@@ -0,0 +1,3 @@
+from .tool_doctor import tool_doctor
+
+__all__ = ["tool_doctor"]

aixtools/tools/doctor/tool_doctor.py

@@ -0,0 +1,61 @@
+from aixtools.agents import get_agent, run_agent
+from aixtools.tools.doctor.tool_recommendation import ToolRecommendation
+
+TOOL_DOCTOR_PROMPT = """
+## Tool doctor
+You are helping to debug common errors in tools and tool definitions.
+
+Given the tools, for each tool you will give feedback about the tool
+definition.
+
+1. Name: Check if the tool name is descriptive and follows naming conventions.
+2. Description: Ensure the tool's description is clear and provides enough
+   detail
+   so that users understand its purpose and functionality.
+3. Return type: Check if the return type is specified and matches the tool's
+   functionality.
+4. Arguments: Verify that the tool arguments are well-defined and include types
+   and descriptions. Ensure that argument names are descriptive and follow
+   naming conventions.
+5. Look for any missing or redundant information in the tool definition.
+
+Some rules:
+- Ignore a tool called 'final_result'.
+- Do not suggest change if the tool is already well-defined.
+- Don't be nitpicking, focus on significant improvements that can be made to
+  the tool definition.
+- Don't suggest trivial improvements or changes for things that are
+  self-evident or already well-defined.
+"""
+
+
+async def tool_doctor_single(tool) -> ToolRecommendation:
+    """Run the tool doctor agent to analyze a single tool"""
+    agent = get_agent(tools=[tool], output_type=ToolRecommendation)
+    ret = await run_agent(agent, TOOL_DOCTOR_PROMPT, log_model_requests=True, verbose=True, debug=True)
+    return ret[0]  # type: ignore
+
+
+async def tool_doctor_multiple(tools: list) -> list[ToolRecommendation]:
+    """Run the tool doctor agent to analyze the tools"""
+    agent = get_agent(tools=tools, output_type=list[ToolRecommendation])
+    ret = await run_agent(agent, TOOL_DOCTOR_PROMPT)
+    return ret[0]  # type: ignore
+
+
+async def tool_doctor(tools: list, max_tools_per_batch=5, verbose=True) -> list[ToolRecommendation]:
+    """Run the tool doctor agent to analyze tools and give recommendations."""
+    # Split tools into batches if they exceed the max_tools_per_batch limit
+    results = []
+    for i in range(0, len(tools), max_tools_per_batch):
+        batch = tools[i : i + max_tools_per_batch]
+        batch_num = i // max_tools_per_batch + 1
+        tool_names = [t.__name__ for t in batch]
+        print(f"Processing batch {batch_num} with {len(batch)} tools: {tool_names}")
+        ret = await tool_doctor_multiple(batch)
+        results.extend(ret)
+    # Print results if verbose
+    if verbose:
+        for r in results:
+            print(r)
+    return results

aixtools/tools/doctor/tool_recommendation.py

@@ -0,0 +1,44 @@
+from pydantic import BaseModel
+
+
+class ArgumentRecommendation(BaseModel):
+    """A recommendation for an argument"""
+
+    arg_name: str
+    arg_type: str
+    arg_description_improvement: str
+
+    def __str__(self):
+        return f"{self.arg_name} ({self.arg_type}): {self.arg_description_improvement}"
+
+
+class ToolRecommendation(BaseModel):
+    """A recomendation for a tool"""
+
+    name: str
+    needs_improvement: bool
+    description_improvement: str
+    arguments: list[ArgumentRecommendation]
+    return_description_improvement: str
+
+    def signature(self):
+        """Generate the function signature for the tool"""
+        signature = f"def {self.name}("
+        if self.arguments:
+            signature += ", ".join([f"{arg.arg_name}: {arg.arg_type}" for arg in self.arguments])
+        signature += ")"
+        return signature
+
+    def __str__(self):
+        out = f"Tool name: '{self.name}'\n"
+        out += f"  - Signature: {self.signature()}\n"
+        out += f"  - Needs improvement: {self.needs_improvement}\n"
+        if self.description_improvement:
+            out += f"  - Description improvement: {self.description_improvement}\n"
+        if self.return_description_improvement:
+            out += f"  - Return description improvement: {self.return_description_improvement}\n"
+        if self.arguments:
+            out += "  - Arguments:\n"
+            for arg in self.arguments:
+                out += f"    - {arg}\n"
+        return out

aixtools/utils/__init__.py

@@ -0,0 +1,35 @@
+"""
+Utils package initialization.
+"""
+
+from aixtools.logging.logging_config import get_logger  # pylint: disable=import-error
+from aixtools.utils import config
+from aixtools.utils.enum_with_description import EnumWithDescription
+from aixtools.utils.persisted_dict import PersistedDict
+from aixtools.utils.utils import (
+    escape_backticks,
+    escape_newline,
+    find_file,
+    prepend_all_lines,
+    remove_quotes,
+    tabit,
+    to_str,
+    tripple_quote_strip,
+    truncate,
+)
+
+__all__ = [
+    "config",
+    "PersistedDict",
+    "EnumWithDescription",
+    "escape_newline",
+    "escape_backticks",
+    "find_file",
+    "get_logger",
+    "prepend_all_lines",
+    "remove_quotes",
+    "tabit",
+    "to_str",
+    "truncate",
+    "tripple_quote_strip",
+]