versionhq 1.1.9.14__py3-none-any.whl → 1.1.10.2__py3-none-any.whl

versionhq/tool/__init__.py

@@ -1,56 +0,0 @@
-from enum import Enum
-
-DEFAULT_AUTH_SCHEME = "OAUTH2"
-
-class ComposioAuthScheme(str, Enum):
-    OAUTH2 = "OAUTH2"
-    BEARER_TOKEN = "BEARER_TOKEN"
-    API_KEY = "API_KEY"
-
-
-class ComposioAppName(str, Enum):
-    """
-    Enum to store app names that we can connect via Composio as data pipelines or destination services.
-    """
-
-    SALESFORCE = "salesforce"
-    AIRTABLE = "airtable"
-    MAILCHIMP = "mailchimp"
-    HUBSPOT = "hubspot"
-    KLAVIYO = "klaviyo"
-    GOOGLESHEET = "googlesheets"
-    GMAIL = "gmail"
-    FACEBOOK = "facebook"
-    TWITTER = "twitter"
-    TWITTER_MEDIA = "twitter_media"
-    LINKEDIN = "linkedin"
-
-
-composio_app_set = [
-    (ComposioAppName.SALESFORCE, ComposioAuthScheme.OAUTH2),
-    (ComposioAppName.AIRTABLE, ComposioAuthScheme.OAUTH2, ComposioAuthScheme.API_KEY, ComposioAuthScheme.BEARER_TOKEN),
-    (ComposioAppName.MAILCHIMP, ComposioAuthScheme.OAUTH2),
-    (ComposioAppName.HUBSPOT, ComposioAuthScheme.OAUTH2, ComposioAuthScheme.BEARER_TOKEN),
-    (ComposioAppName.KLAVIYO, ComposioAuthScheme.OAUTH2, ComposioAuthScheme.API_KEY),
-    (ComposioAppName.GOOGLESHEET, ComposioAuthScheme.OAUTH2),
-    (ComposioAppName.GMAIL, ComposioAuthScheme.OAUTH2, ComposioAuthScheme.BEARER_TOKEN),
-    (ComposioAppName.TWITTER, ComposioAuthScheme.OAUTH2),
-    (ComposioAppName.TWITTER_MEDIA, ComposioAuthScheme.OAUTH2),
-    (ComposioAppName.FACEBOOK, ComposioAuthScheme.OAUTH2),
-    (ComposioAppName.LINKEDIN, ComposioAuthScheme.OAUTH2),
-]
-
-class ComposioStatus(str, Enum):
-    INITIATED = "INITIATED"
-    ACTIVE = "ACTIVE"
-    FAILED = "FAILED"
-
-
-
-
-class ComposioAction(str, Enum):
-    """
-    Enum to store composio's action that can be called via `Actions.xxx`
-    """
-    # HUBSPOT_INITIATE_DATA_IMPORT_PROCESS = "hubspot_initate_date_import_process"
-    HUBSPOT_CREATE_PIPELINE_STAGE = "hubspot_create_pipeline_stage"

versionhq/tool/cache_handler.py

@@ -0,0 +1,40 @@
+from typing import Any, Dict, Optional
+
+from pydantic import BaseModel, PrivateAttr, Field, InstanceOf
+
+
+class CacheHandler(BaseModel):
+    """
+    A class to add or read cache
+    """
+
+    _cache: Dict[str, Any] = PrivateAttr(default_factory=dict)
+
+    def add(self, tool_name: str, input: str, output: Any) -> None:
+        self._cache[f"{tool_name}-{input}"] = output
+
+    def read(self, tool_name: str, input: str) -> Optional[str]:
+        return self._cache.get(f"{tool_name}-{input}")
+
+
+
+class CacheTool(BaseModel):
+    """
+    A cache tool to read the cached result.
+    """
+
+    name: str = "Cache Tool"
+    cache_handler: InstanceOf[CacheHandler] = Field(default_factory=CacheHandler)
+
+    def read_cache(self, key):
+        split = key.split("tool:")
+        tool = split[1].split("|input:")[0].strip()
+        tool_input = split[1].split("|input:")[1].strip()
+        return self.cache_handler.read(tool, tool_input)
+
+    def tool(self):
+        return Tool(
+            func=self.read_cache,
+            name=self.name,
+            description="Read from cache"
+        )

versionhq/tool/composio_tool.py

@@ -11,9 +11,10 @@ from pydantic_core import PydanticCustomError
 from composio import ComposioToolSet
 from composio_langchain import action
 
-from versionhq.tool import ComposioAppName, ComposioAuthScheme, composio_app_set, ComposioStatus, ComposioAction
+from versionhq.tool.composio_tool_vars import ComposioAppName, ComposioAuthScheme, composio_app_set, ComposioStatus, ComposioAction
+from versionhq.tool.cache_handler import CacheHandler
 from versionhq._utils.logger import Logger
-from versionhq._utils.cache_handler import CacheHandler
+
 
 load_dotenv(override=True)
 

versionhq/tool/composio_tool_vars.py

@@ -0,0 +1,56 @@
+from enum import Enum
+
+DEFAULT_AUTH_SCHEME = "OAUTH2"
+
+class ComposioAuthScheme(str, Enum):
+    OAUTH2 = "OAUTH2"
+    BEARER_TOKEN = "BEARER_TOKEN"
+    API_KEY = "API_KEY"
+
+
+class ComposioAppName(str, Enum):
+    """
+    Enum to store app names that we can connect via Composio as data pipelines or destination services.
+    """
+
+    SALESFORCE = "salesforce"
+    AIRTABLE = "airtable"
+    MAILCHIMP = "mailchimp"
+    HUBSPOT = "hubspot"
+    KLAVIYO = "klaviyo"
+    GOOGLESHEET = "googlesheets"
+    GMAIL = "gmail"
+    FACEBOOK = "facebook"
+    TWITTER = "twitter"
+    TWITTER_MEDIA = "twitter_media"
+    LINKEDIN = "linkedin"
+
+
+composio_app_set = [
+    (ComposioAppName.SALESFORCE, ComposioAuthScheme.OAUTH2),
+    (ComposioAppName.AIRTABLE, ComposioAuthScheme.OAUTH2, ComposioAuthScheme.API_KEY, ComposioAuthScheme.BEARER_TOKEN),
+    (ComposioAppName.MAILCHIMP, ComposioAuthScheme.OAUTH2),
+    (ComposioAppName.HUBSPOT, ComposioAuthScheme.OAUTH2, ComposioAuthScheme.BEARER_TOKEN),
+    (ComposioAppName.KLAVIYO, ComposioAuthScheme.OAUTH2, ComposioAuthScheme.API_KEY),
+    (ComposioAppName.GOOGLESHEET, ComposioAuthScheme.OAUTH2),
+    (ComposioAppName.GMAIL, ComposioAuthScheme.OAUTH2, ComposioAuthScheme.BEARER_TOKEN),
+    (ComposioAppName.TWITTER, ComposioAuthScheme.OAUTH2),
+    (ComposioAppName.TWITTER_MEDIA, ComposioAuthScheme.OAUTH2),
+    (ComposioAppName.FACEBOOK, ComposioAuthScheme.OAUTH2),
+    (ComposioAppName.LINKEDIN, ComposioAuthScheme.OAUTH2),
+]
+
+class ComposioStatus(str, Enum):
+    INITIATED = "INITIATED"
+    ACTIVE = "ACTIVE"
+    FAILED = "FAILED"
+
+
+
+
+class ComposioAction(str, Enum):
+    """
+    Enum to store composio's action that can be called via `Actions.xxx`
+    """
+    # HUBSPOT_INITIATE_DATA_IMPORT_PROCESS = "hubspot_initate_date_import_process"
+    HUBSPOT_CREATE_PIPELINE_STAGE = "hubspot_create_pipeline_stage"

versionhq/tool/decorator.py

@@ -10,12 +10,11 @@ def tool(*args):
     """
 
     def create_tool(tool_name: str) -> Callable:
-
-        def _make_tool(f: Callable) -> Tool:
-            if f.__doc__ is None:
+        def _make_tool(func: Callable) -> Tool:
+            if func.__doc__ is None:
                 raise ValueError("Function must have a docstring")
 
-            if f.__annotations__ is None:
+            if func.__annotations__ is None:
                 raise ValueError("Function must have type annotations")
 
             class_name = "".join(tool_name.split()).title()
@@ -24,11 +23,11 @@ def tool(*args):
                 (BaseModel,),
                 {
                     "__annotations__": {
-                        k: v for k, v in f.__annotations__.items() if k != "return"
+                        k: v for k, v in func.__annotations__.items() if k != "return"
                     },
                 },
             )
-            return Tool(name=tool_name, function=f, args_schema=args_schema)
+            return Tool(name=tool_name, func=func, args_schema=args_schema)
 
         return _make_tool
 

versionhq/tool/model.py

@@ -1,10 +1,13 @@
 from abc import ABC, abstractmethod
 from inspect import signature
-from typing import Any, Dict, Callable, Type, Optional, get_args, get_origin
+from typing import Any, Dict, Callable, Type, Optional, get_args, get_origin, get_type_hints
 from typing_extensions import Self
-from pydantic import InstanceOf, BaseModel, ConfigDict, Field, field_validator, model_validator
+from pydantic import InstanceOf, BaseModel, ConfigDict, Field, field_validator, model_validator, PrivateAttr, create_model
+from pydantic_core import PydanticCustomError
 
-from versionhq._utils.cache_handler import CacheHandler
+from versionhq.llm.llm_vars import SchemaType
+from versionhq.tool.cache_handler import CacheHandler
+from versionhq._utils.logger import Logger
 
 
 class BaseTool(ABC, BaseModel):
@@ -12,16 +15,27 @@ class BaseTool(ABC, BaseModel):
     Abstract class for Tool class.
     """
 
-    class _ArgsSchemaPlaceholder(BaseModel):
+    class ArgsSchemaPlaceholder(BaseModel):
         pass
 
-    args_schema: Type[BaseModel] = Field(default_factory=_ArgsSchemaPlaceholder)
+    _logger: Logger = PrivateAttr(default_factory=lambda: Logger(verbose=True))
+
+    object_type: str = Field(default="function")
+    name: str = Field(default=None)
+    description: str = Field(default=None)
+    properties: Dict[str, Any] = Field(default_factory=dict, description="for llm func calling")
+    args_schema: Type[BaseModel] = Field(default_factory=ArgsSchemaPlaceholder)
+
+    tool_handler: Optional[Dict[str, Any] | Any] = Field(default=None, description="store tool_handler to record the tool usage")
+    should_cache: bool = Field(default=True, description="whether the tool usage should be cached")
+    cache_function: Callable = lambda _args=None, _result=None: True
+    cache_handler: Optional[InstanceOf[CacheHandler]] = Field(default=None)
 
 
     @field_validator("args_schema", mode="before")
     @classmethod
     def _default_args_schema(cls, v: Type[BaseModel]) -> Type[BaseModel]:
-        if not isinstance(v, cls._ArgsSchemaPlaceholder):
+        if not isinstance(v, cls.ArgsSchemaPlaceholder):
             return v
 
         return type(
@@ -31,21 +45,29 @@ class BaseTool(ABC, BaseModel):
         )
 
 
-    @abstractmethod
-    def _run(self, *args: Any, **kwargs: Any,) -> Any:
-        """any handling"""
-
-
-
-class Tool(BaseTool):
-    name: str = Field(default=None)
-    goal: str = Field(default=None)
-    function: Callable = Field(default=None)
-    tool_handler: Optional[Dict[str, Any] | Any] = Field(default=None, description="store tool_handler to record the usage of this tool")
-    should_cache: bool = Field(default=True, description="whether the tool usage should be cached")
-    cache_function: Callable = lambda _args=None, _result=None: True
-    cache_handler: Optional[InstanceOf[CacheHandler]] = Field(default=None)
-
+    @field_validator("properties", mode="before")
+    @classmethod
+    def _default_properties(cls, v: Dict[str, Any]) -> Dict[str, Any]:
+        p, r = dict(), list()
+        for k, v in cls._run.__annotations__.items():
+            if k != "return":
+                p.update({ k: { "type": SchemaType(type(v)).convert(), "name": k, }} )
+                r.append(k)
+
+        return {
+            "type": cls.object_type,
+            "function": {
+                "name": cls.name.replace(" ", "_"),
+                "description": cls.description,
+                "parameters": {
+                    "type": "object",
+                    "properties": p,
+                    "required": r,
+                    "additionalProperties": False
+                },
+                "strict": True
+            }
+        }
 
     @model_validator(mode="after")
     def set_up_tool_handler(self) -> Self:
@@ -59,13 +81,9 @@ class Tool(BaseTool):
 
         return self
 
-
-    @model_validator(mode="after")
-    def set_up_function(self) -> Self:
-        if self.function is None:
-            self.function = self._run
-            self._set_args_schema_from_func()
-        return self
+    @abstractmethod
+    def _run(self, *args: Any, **kwargs: Any,) -> Any:
+        """any handling"""
 
 
     @staticmethod
@@ -95,108 +113,236 @@ class Tool(BaseTool):
                 import json
                 raw_args = json.loads(raw_args)
             except json.JSONDecodeError as e:
-                raise ValueError(f"Failed to parse arguments as JSON: {e}")
+                raise ValueError(f"Failed to parse arguments as JSON: {str(e)}")
 
         try:
             validated_args = self.args_schema.model_validate(raw_args)
             return validated_args.model_dump()
 
         except Exception as e:
-            raise ValueError(f"Arguments validation failed: {e}")
+            raise ValueError(f"Arguments validation failed: {str(e)}")
 
 
-    def _set_args_schema_from_func(self):
-        class_name = f"{self.__class__.__name__}Schema"
-        self.args_schema = type(
-            class_name,
-            (BaseModel,),
-            { "__annotations__": {
-                k: v for k, v in self._run.__annotations__.items() if k != "return"
-            } },
-        )
+    def _create_schema(self) -> type[BaseModel]:
+        """
+        Create a Pydantic schema from a function's signature
+        """
+        import inspect
+
+        sig = inspect.signature(self.func)
+        type_hints = get_type_hints(self.func)
+        fields = {}
+        for param_name, param in sig.parameters.items():
+            if param_name in ("self", "cls"):
+                continue
+
+            annotation = type_hints.get(param_name, Any)
+            default = ... if param.default == param.empty else param.default
+            fields[param_name] = (annotation, Field(default=default))
+
+        schema_name = f"{self.func.__name__.title()}Schema"
+        return create_model(schema_name, **fields)
+
+
+
+class Tool(BaseTool):
+    func: Callable = Field(default=None)
+
+
+    @model_validator(mode="after")
+    def validate_func(self) -> Self:
+        if not self.func and not self._run:
+            self._logger.log(level="error", message=f"Tool must have a function", color="red")
+            raise PydanticCustomError("function_missing", f"Function is missing in the tool.", {})
+
+        elif self.func and not isinstance(self.func, Callable):
+            self._logger.log(level="error", message=f"The tool is missing a valid function", color="red")
+            raise PydanticCustomError("invalid_function", f"The value in the function field must be callable.", {})
+
+        else:
+            try:
+                self.args_schema = self._create_schema_from_function()
+                self._validate_function_signature()
+
+            except Exception as e:
+                self._logger.log(level="error", message=f"The tool is missing a valid function: {str(e)}", color="red")
+                raise PydanticCustomError("invalid_function", f"Invalid function: {str(e)}", {})
+
+        return self
+
+
+    @model_validator(mode="after")
+    def set_up_name(self) -> Self:
+        if not self.name:
+            self.name = self.func.__name__ if self.func else ""
+
+        return self
+
+
+    @model_validator(mode="after")
+    def set_up_description(self) -> Self:
+        if not self.description:
+            if not self.args_schema:
+                self.args_schema = self._default_args_schema(self)
+
+            args_schema = {
+                name: {
+                    "description": field.description,
+                    "type": self._get_arg_annotations(field.annotation),
+                }
+                for name, field in self.args_schema.model_fields.items()
+            }
+            self.description = f"Tool: {self.name}\nArgs: {args_schema}"
+
+        return self
+
+
+    @model_validator(mode="after")
+    def set_up_args_schema(self) -> Self:
+        """
+        Set up args schema based on the given function.
+        """
+        if self.func:
+            self.args_schema = self._create_schema_from_function()
+        return self
+
+
+    @model_validator(mode="after")
+    def set_up_func_calling_properties(self) -> Self:
+        """
+        Format function_calling params from args_schema.
+        """
+
+        p, r = dict(), list()
+        if self.args_schema:
+            for name, field in self.args_schema.model_fields.items():
+                if name != "kwargs" and name != "args":
+                    p.update(
+                        {
+                            name: {
+                                "description": field.description if field.description else "",
+                                "type": SchemaType(self._get_arg_annotations(field.annotation)).convert(),
+                            }
+                        }
+                    )
+                    r.append(name)
+
+            properties = {
+                "type": self.object_type,
+                "function": {
+                    "name": self.name.replace(" ", "_"),
+                    "description": self.description if self.description else "a tool function to execute",
+                    "parameters": {
+                        "type": "object",
+                        "properties": p,
+                        "required": r,
+                        "additionalProperties": False
+                    },
+                    "strict": True,
+                },
+            }
+            self.properties = properties
         return self
 
 
+    def _create_schema_from_function(self) -> type[BaseModel]:
+        """
+        Create a Pydantic schema from a function's signature
+        """
+        import inspect
+
+        sig = inspect.signature(self.func)
+        type_hints = get_type_hints(self.func)
+        fields = {}
+        for param_name, param in sig.parameters.items():
+            if param_name in ("self", "cls"):
+                continue
+
+            annotation = type_hints.get(param_name, Any)
+            default = ... if param.default == param.empty else param.default
+            fields[param_name] = (annotation, Field(default=default))
+
+        schema_name = f"{self.func.__name__.title()}Schema"
+        return create_model(schema_name, **fields)
+
+
+    def _validate_function_signature(self) -> None:
+        """
+        Validate that the function signature matches the args schema.
+        """
+
+        import inspect
+
+        sig = inspect.signature(self.func)
+        schema_fields = self.args_schema.model_fields
+
+        for param_name, param in sig.parameters.items():
+            if param_name in ("self", "cls"):
+                continue
+
+            if param.kind in (inspect.Parameter.VAR_KEYWORD, inspect.Parameter.VAR_POSITIONAL):
+                continue
+
+            if param.default == inspect.Parameter.empty:
+                if param_name not in schema_fields:
+                    raise ValueError(f"Required function parameter '{param_name}' not found in args_schema")
+
+
     def _run(self, *args: Any, **kwargs: Any) -> Any:
-        return self.run(*args, **kwargs)
+        return self.func(*args, **kwargs)
 
 
-    def run(self, *args, **kwargs) -> Any:
+    def _handle_toolset(self, params: Dict[str, Any] = None) -> Any:
         """
-        Use tool and record its usage if should_cache is True.
+        Read the cache from the ToolHandler instance or execute _run() method.
         """
-        from versionhq.tool.tool_handler import ToolHandler
 
-        result = None
-        tool_set = ToolSet(tool=self, kwargs={})
+        from versionhq.tool.tool_handler import ToolHandler
 
-        if self.function:
-            result = self.function(*args, **kwargs)
+        if not self.args_schema:
+            self.args_schema = self._create_schema_from_function()
 
-        else:
-            acceptable_args = self.args_schema.model_json_schema()["properties"].keys()
-            acceptable_kwargs = { k: v for k, v in kwargs.items() if k in acceptable_args }
-            tool_set = ToolSet(tool=self, kwargs=acceptable_kwargs)
+        result = None
+        acceptable_args = self.args_schema.model_json_schema()["properties"].keys()
+        acceptable_kwargs = { k: v for k, v in params.items() if k in acceptable_args } if params else dict()
+        parsed_kwargs = self._parse_args(raw_args=acceptable_kwargs)
+        tool_set = ToolSet(tool=self, kwargs=acceptable_kwargs)
 
-            if self.tool_handler:
-                if self.tool_handler.has_called_before(tool_set):
-                    self.tool_handler.error = "Agent execution error"
+        if self.tool_handler and isinstance(self.tool_handler, ToolHandler):
+            if self.tool_handler.has_called_before(tool_set):
+                self.tool_handler.error = "Agent execution error"
 
-                elif self.tool_handler.cache:
-                    result = self.tools_handler.cache.read(tool=tool_set.tool.name, input=tool_set.kwargs)
-                    if result is None:
-                        parsed_kwargs = self._parse_args(raw_args=acceptable_kwargs)
-                        result = self.function(**parsed_kwargs) if self.function else None
+            elif self.tool_handler.cache:
+                result = self.tool_handler.cache.read(tool_name=tool_set.tool.name, input=str(tool_set.kwargs))
+                if not result:
+                    result = self.func(**parsed_kwargs)
 
             else:
-                tool_handler = ToolHandler(last_used_tool=tool_set, cache_handler=self.cache_handler, should_cache=self.should_cache)
-                self.tool_handler = tool_handler
-                parsed_kwargs = self._parse_args(raw_args=acceptable_kwargs)
-                result = self.function(**parsed_kwargs) if self.function else None
+                result = self.func(**parsed_kwargs)
+
+        else:
+            tool_handler = ToolHandler(last_used_tool=tool_set, cache_handler=self.cache_handler, should_cache=self.should_cache)
+            self.tool_handler = tool_handler
+            result = self.func(**parsed_kwargs)
 
 
         if self.should_cache is True:
-            self.tool_handler.record_last_tool_used(tool_set, result, self.should_cache)
+            self.tool_handler.record_last_tool_used(last_used_tool=tool_set, output=result, should_cache=self.should_cache)
 
         return result
 
 
-    @property
-    def description(self) -> str:
-        args_schema = {
-            name: {
-                "description": field.description,
-                "type": self._get_arg_annotations(field.annotation),
-            }
-            for name, field in self.args_schema.model_fields.items()
-        }
-
-        return f"Tool Name: {self.name}\nTool Arguments: {args_schema}\nGoal: {self.goal}"
+    def run(self, params: Dict[str, Any] = None) -> Any:
+        """
+        Execute a tool using a toolset and cached tools
+        """
+        result = self._handle_toolset(params)
+        return result
 
 
 class ToolSet(BaseModel):
     """
-    Store the tool called and any kwargs used.
+    Store the tool called and any kwargs used. (The tool name and kwargs will be stored in the cache.)
     """
-    tool: InstanceOf[Tool] | Any = Field(..., description="store the tool instance to be called.")
-    kwargs: Optional[Dict[str, Any]] = Field(..., description="kwargs passed to the tool")
-
-
-class InstructorToolSet(BaseModel):
-    tool: InstanceOf[Tool] | Any = Field(..., description="store the tool instance to be called.")
+    tool: InstanceOf[Tool] | Type[Tool] = Field(..., description="store the tool instance to be called.")
     kwargs: Optional[Dict[str, Any]] = Field(..., description="kwargs passed to the tool")
-
-
-class CacheTool(BaseModel):
-    """
-    Default tools to hit the cache.
-    """
-
-    name: str = "Hit Cache"
-    cache_handler: CacheHandler = Field(default_factory=CacheHandler)
-
-    def hit_cache(self, key):
-        split = key.split("tool:")
-        tool = split[1].split("|input:")[0].strip()
-        tool_input = split[1].split("|input:")[1].strip()
-        return self.cache_handler.read(tool, tool_input)

versionhq/tool/tool_handler.py

@@ -1,8 +1,8 @@
 from typing import Optional, Any
 from pydantic import InstanceOf
 
-from versionhq.tool.model import ToolSet, InstructorToolSet
-from versionhq._utils.cache_handler import CacheHandler
+from versionhq.tool.model import ToolSet
+from versionhq.tool.cache_handler import CacheHandler, CacheTool
 
 
 class ToolHandler:
@@ -10,34 +10,26 @@ class ToolHandler:
     Record the tool usage by ToolSet instance with cache and error recording.
     """
 
-    last_used_tool: InstanceOf[ToolSet] | InstanceOf[InstructorToolSet]
-    cache: Optional[CacheHandler]
+    last_used_tool: InstanceOf[ToolSet]
+    cache: InstanceOf[CacheHandler] = CacheHandler()
     error: Optional[str]
     should_cache: bool
 
-    def __init__(
-        self,
-        last_used_tool: InstanceOf[ToolSet] | InstanceOf[InstructorToolSet] = None,
-        cache_handler: Optional[CacheHandler] = None,
-        should_cache: bool = True
-    ):
-        self.cache = cache_handler
+    def __init__(self, last_used_tool: InstanceOf[ToolSet] = None, cache_handler: InstanceOf[CacheHandler] = None, should_cache: bool = True):
         self.last_used_tool = last_used_tool
+        self.cache = cache_handler if cache_handler else CacheHandler()
         self.should_cache = should_cache
 
 
-    def record_last_tool_used(
-        self, last_used_tool: InstanceOf[ToolSet] | InstanceOf[InstructorToolSet], output: str, should_cache: bool = True
-    ) -> Any:
-        from versionhq.tool.model import CacheTool
-
+    def record_last_tool_used(self, last_used_tool: InstanceOf[ToolSet], output: str, should_cache: bool = True) -> None:
         self.last_used_tool = last_used_tool
 
-        if self.cache and should_cache and last_used_tool.tool.name != CacheTool().name:
-            self.cache.add(tool=last_used_tool.tool.name, input=last_used_tool.kwargs, output=output)
+        if should_cache:
+            self.cache = CacheHandler()
+            self.cache.add(tool_name=last_used_tool.tool.name, input=str(last_used_tool.kwargs), output=output)
 
 
-    def has_called_before(self, tool_set: ToolSet = None) -> bool:
+    def has_called_before(self, tool_set: InstanceOf[ToolSet] = None) -> bool:
         if tool_set is None or not self.last_used_tool:
             return False