| 1 | from abc import abstractmethod |
| 2 | from dataclasses import dataclass |
| 3 | from typing import Any |
| 4 | |
| 5 | from agent import Agent, LoopData |
| 6 | from helpers.extension import call_extensions_async |
| 7 | from helpers.print_style import PrintStyle |
| 8 | from helpers.strings import sanitize_string |
| 9 | |
| 10 | |
| 11 | @dataclass |
| 12 | class Response: |
| 13 | message:str |
| 14 | break_loop: bool |
| 15 | additional: dict[str, Any] | None = None |
| 16 | |
| 17 | class Tool: |
| 18 | |
| 19 | def __init__(self, agent: Agent, name: str, method: str | None, args: dict[str,str], message: str, loop_data: LoopData | None, **kwargs) -> None: |
| 20 | self.agent = agent |
| 21 | self.name = name |
| 22 | self.method = method |
| 23 | self.args = args |
| 24 | self.loop_data = loop_data |
| 25 | self.message = message |
| 26 | self.progress: str = "" |
| 27 | |
| 28 | @abstractmethod |
| 29 | async def execute(self,**kwargs) -> Response: |
| 30 | pass |
| 31 | |
| 32 | async def set_progress(self, content: str | None): |
| 33 | ctx = {"content": content or ""} |
| 34 | await call_extensions_async("tool_output_update", self.agent, ctx=ctx) |
| 35 | self.progress = ctx["content"] |
| 36 | |
| 37 | def add_progress(self, content: str | None): |
| 38 | if not content: |
| 39 | return |
| 40 | self.progress += content |
| 41 | |
| 42 | async def before_execution(self, **kwargs): |
| 43 | PrintStyle(font_color="#1B4F72", padding=True, background_color="white", bold=True).print(f"{self.agent.agent_name}: Using tool '{self.name}'") |
| 44 | self.log = self.get_log_object() |
| 45 | if self.args and isinstance(self.args, dict): |
| 46 | for key, value in self.args.items(): |
| 47 | ctx = {"content": str(value) if not isinstance(value, str) else value} |
| 48 | await call_extensions_async("tool_output_update", self.agent, ctx=ctx) |
| 49 | display_value = ctx["content"] |
| 50 | PrintStyle(font_color="#85C1E9", bold=True).stream(self.nice_key(key)+": ") |
| 51 | PrintStyle(font_color="#85C1E9", padding=isinstance(value,str) and "\n" in value).stream(display_value) |
| 52 | PrintStyle().print() |
| 53 | |
| 54 | async def after_execution(self, response: Response, **kwargs): |
| 55 | text = sanitize_string(response.message.strip()) |
| 56 | self.agent.hist_add_tool_result(self.name, text, id=self.log.id, **(response.additional or {})) |
| 57 | PrintStyle(font_color="#1B4F72", background_color="white", padding=True, bold=True).print(f"{self.agent.agent_name}: Response from tool '{self.name}'") |
| 58 | PrintStyle(font_color="#85C1E9").print(text) |
| 59 | self.log.update(content=text) |
| 60 | |
| 61 | def get_log_object(self): |
| 62 | import uuid |
| 63 | pre_id = str(uuid.uuid4()) |
| 64 | if self.method: |
| 65 | heading = f"icon://construction {self.agent.agent_name}: Using tool '{self.name}:{self.method}'" |
| 66 | else: |
| 67 | heading = f"icon://construction {self.agent.agent_name}: Using tool '{self.name}'" |
| 68 | return self.agent.context.log.log(type="tool", heading=heading, content="", kvps=self.args, _tool_name=self.name, id=pre_id) |
| 69 | |
| 70 | def nice_key(self, key:str): |
| 71 | words = key.split('_') |
| 72 | words = [words[0].capitalize()] + [word.lower() for word in words[1:]] |
| 73 | result = ' '.join(words) |
| 74 | return result |