wehub-resource-sync
334 lines
10 KiB
Python
334 lines
10 KiB
Python
"""Example tool functions for Omnigent agents.
|
|
|
|
These are plain Python functions that can be referenced from YAML agent
|
|
definitions via the ``callable`` field on a FunctionTool.
|
|
"""
|
|
|
|
import time
|
|
from datetime import datetime, timedelta, timezone
|
|
|
|
from omnigent.policies.schema import PolicyEvent, PolicyResponse
|
|
|
|
|
|
def web_search(query: str) -> dict:
|
|
"""Return deterministic fake web results for examples and tests."""
|
|
query_text = str(query).strip()
|
|
if not query_text:
|
|
return {"results": []}
|
|
return {
|
|
"results": [
|
|
{
|
|
"title": f"Overview of {query_text}",
|
|
"url": f"https://example.com/search?q={query_text.replace(' ', '+')}",
|
|
"snippet": f"Public web summary for {query_text}.",
|
|
},
|
|
{
|
|
"title": f"Recent news about {query_text}",
|
|
"url": f"https://news.example.com/{query_text.replace(' ', '-').lower()}",
|
|
"snippet": f"Synthetic news result related to {query_text}.",
|
|
},
|
|
]
|
|
}
|
|
|
|
|
|
def summarize(text: str, max_words: int = 30) -> dict:
|
|
"""Return a deterministic fake summary of *text*.
|
|
|
|
Mirrors :func:`web_search` in shape — a plain callable with a
|
|
deterministic payload, suitable as the ``callable:`` target on
|
|
a YAML ``type: function`` tool. No real LLM or NLP; the
|
|
"summary" is the first ``max_words`` whitespace-delimited
|
|
tokens of the input plus an ellipsis when truncated. Useful
|
|
for examples that pair search → summarize, where the agent
|
|
needs *some* tool to call but the test doesn't care about
|
|
summary quality.
|
|
|
|
:param text: The text to summarize, e.g. a search snippet.
|
|
:param max_words: Truncation cap, e.g. ``30``. Values <= 0
|
|
are coerced to 1 so the summary is never empty.
|
|
:returns: Dict with ``summary`` (the truncated text),
|
|
``input_words`` (token count of the input), and
|
|
``truncated`` (whether the cap was hit).
|
|
"""
|
|
body = str(text).strip()
|
|
cap = max(1, int(max_words)) if isinstance(max_words, int) else 30
|
|
words = body.split()
|
|
truncated = len(words) > cap
|
|
head = words[:cap]
|
|
summary = " ".join(head) + ("…" if truncated else "")
|
|
return {
|
|
"summary": summary,
|
|
"input_words": len(words),
|
|
"truncated": truncated,
|
|
}
|
|
|
|
|
|
def read_internal_doc(doc_id: str) -> dict:
|
|
"""Return a deterministic confidential document payload."""
|
|
doc_key = str(doc_id).strip() or "unknown"
|
|
return {
|
|
"doc_id": doc_key,
|
|
"classification": "confidential",
|
|
"content": f"CONFIDENTIAL: internal notes for {doc_key}.",
|
|
}
|
|
|
|
|
|
def run_shell(command: str) -> dict:
|
|
"""Return deterministic fake shell output for policy examples."""
|
|
command_text = str(command).strip()
|
|
if command_text == "ls":
|
|
return {"stdout": "README.md\nexamples\nomnigent\ntests", "stderr": "", "exit_code": 0}
|
|
return {
|
|
"stdout": f"simulated shell output for: {command_text}",
|
|
"stderr": "",
|
|
"exit_code": 0,
|
|
}
|
|
|
|
|
|
def write_file(path: str, content: str) -> dict:
|
|
"""Pretend to write a file and report what would have happened."""
|
|
return {
|
|
"ok": True,
|
|
"path": str(path),
|
|
"bytes_written": len(str(content).encode("utf-8")),
|
|
}
|
|
|
|
|
|
def get_current_time(timezone_name: str = "UTC") -> str:
|
|
"""Return the current date/time in the given timezone.
|
|
|
|
Supports a handful of common timezone names for the example.
|
|
A production implementation would use ``zoneinfo`` or ``pytz``.
|
|
"""
|
|
offsets = {
|
|
"UTC": 0,
|
|
"US/Eastern": -5,
|
|
"EST": -5,
|
|
"EDT": -4,
|
|
"US/Central": -6,
|
|
"CST": -6,
|
|
"CDT": -5,
|
|
"US/Pacific": -8,
|
|
"PST": -8,
|
|
"PDT": -7,
|
|
"Europe/London": 0,
|
|
"GMT": 0,
|
|
"BST": 1,
|
|
"Europe/Paris": 1,
|
|
"CET": 1,
|
|
"CEST": 2,
|
|
"Asia/Tokyo": 9,
|
|
"JST": 9,
|
|
"Asia/Shanghai": 8,
|
|
"CST_China": 8,
|
|
}
|
|
offset = offsets.get(timezone_name)
|
|
if offset is None:
|
|
return f"Unknown timezone: {timezone_name}. Supported: {', '.join(sorted(offsets.keys()))}"
|
|
tz = timezone(timedelta(hours=offset))
|
|
now = datetime.now(tz)
|
|
return now.strftime(f"%Y-%m-%d %H:%M:%S {timezone_name}")
|
|
|
|
|
|
def calculate(expression: str) -> str:
|
|
"""Safely evaluate a mathematical expression and return the result.
|
|
|
|
Only supports basic arithmetic for safety.
|
|
"""
|
|
allowed = set("0123456789+-*/().% ")
|
|
if not all(c in allowed for c in expression):
|
|
return (
|
|
"Error: expression contains disallowed characters. Only basic arithmetic is supported."
|
|
)
|
|
try:
|
|
result = eval(expression, {"__builtins__": {}}, {})
|
|
return str(result)
|
|
except Exception as exc:
|
|
return f"Error evaluating '{expression}': {exc}"
|
|
|
|
|
|
def sleep(seconds: float):
|
|
if seconds < 0 or seconds > 10:
|
|
raise ValueError("Sleep time must be between 0 and 10 seconds")
|
|
time.sleep(seconds)
|
|
|
|
|
|
def sleep_tool(seconds: float) -> dict[str, float]:
|
|
"""
|
|
Sleep for ``seconds`` then return a marker payload.
|
|
|
|
Plain callable used by ``agent_with_policies.yaml`` and
|
|
``agent_with_tools.yaml`` (and the matching tests). The LLM
|
|
can call this synchronously for short sleeps or dispatch it
|
|
via ``sys_call_async`` for backgrounded / cancellable sleeps;
|
|
cancellation in the latter case flows through
|
|
``sys_cancel_task`` → DBOS workflow cancel → SIGINT to the
|
|
runner subprocess (no per-tool cancel hook needed).
|
|
|
|
Replaces the legacy ``SleepToolRunner`` runner-protocol class
|
|
that this module previously exported. The runner protocol was
|
|
redundant once ``sys_call_async`` shipped — the framework
|
|
handles cancellable async dispatch uniformly for any plain
|
|
callable, so we don't need a separate tool-shape contract for
|
|
it. See ``designs/SERVER_HARNESS_CONTRACT.md`` §"Async work +
|
|
inbox" for the rationale.
|
|
|
|
:param seconds: Sleep duration in seconds, e.g. ``3.5``. Must
|
|
be between 0 and 10 inclusive — values outside the range
|
|
raise ``ValueError`` (matching the policies test that
|
|
expects long-sleep refusal).
|
|
:returns: Dict with ``slept`` set to the sleep duration so
|
|
callers can verify the sleep ran to completion.
|
|
:raises ValueError: When ``seconds`` is negative or greater
|
|
than 10.
|
|
"""
|
|
if seconds < 0 or seconds > 10:
|
|
raise ValueError("Sleep time must be between 0 and 10 seconds")
|
|
time.sleep(seconds)
|
|
return {"slept": float(seconds)}
|
|
|
|
|
|
_ALLOW: PolicyResponse = {"result": "ALLOW"}
|
|
|
|
|
|
def block_division(event: PolicyEvent) -> PolicyResponse:
|
|
"""Block calculate tool calls whose expression contains division.
|
|
|
|
Returns ALLOW for events this policy does not cover. The callable
|
|
self-selects which events to act on without needing an ``on:``
|
|
field in the YAML.
|
|
|
|
Used by ``agent_with_policies.yaml`` to verify that ``type: function``
|
|
handler tools go through Omnigent policy enforcement the same way MCP tools do.
|
|
|
|
:param event: V0 event dict with ``type``, ``target``, ``data``,
|
|
``context`` keys.
|
|
:returns: V0 decision dict.
|
|
"""
|
|
if event.get("type") != "tool_call":
|
|
return _ALLOW
|
|
|
|
data = event.get("data")
|
|
tool_name: str = data.get("name", "") if isinstance(data, dict) else ""
|
|
if tool_name != "calculate":
|
|
return _ALLOW
|
|
|
|
args = data.get("arguments")
|
|
expression: str = args.get("expression", "") if isinstance(args, dict) else ""
|
|
if "/" in expression:
|
|
return {
|
|
"result": "DENY",
|
|
"reason": "Division expressions are denied by policy.",
|
|
}
|
|
|
|
return _ALLOW
|
|
|
|
|
|
def block_long_sleep(event: PolicyEvent) -> PolicyResponse:
|
|
"""Block sleep tool calls longer than 5 seconds.
|
|
|
|
Returns ALLOW for events this policy does not cover. The callable
|
|
self-selects which events to act on without needing an ``on:``
|
|
field in the YAML.
|
|
|
|
:param event: V0 event dict with ``type``, ``target``, ``data``,
|
|
``context`` keys.
|
|
:returns: V0 decision dict.
|
|
"""
|
|
if event.get("type") != "tool_call":
|
|
return _ALLOW
|
|
|
|
data = event.get("data")
|
|
# Match bare name "sleep" or namespaced "sleep-server__sleep".
|
|
tool_name: str = data.get("name", "") if isinstance(data, dict) else ""
|
|
if not tool_name.endswith("sleep"):
|
|
return _ALLOW
|
|
|
|
args = data.get("arguments")
|
|
if not isinstance(args, dict):
|
|
return _ALLOW
|
|
|
|
seconds = args.get("seconds", 0)
|
|
try:
|
|
seconds_value = float(seconds)
|
|
except (TypeError, ValueError):
|
|
return _ALLOW
|
|
|
|
if seconds_value > 5:
|
|
return {
|
|
"result": "DENY",
|
|
"reason": "Sleep calls longer than 5 seconds are denied by policy.",
|
|
}
|
|
|
|
return _ALLOW
|
|
|
|
|
|
_CANADA_KEYWORDS = frozenset(
|
|
{
|
|
"canada",
|
|
"canadian",
|
|
"toronto",
|
|
"vancouver",
|
|
"montreal",
|
|
"ottawa",
|
|
"quebec",
|
|
"ontario",
|
|
"alberta",
|
|
"calgary",
|
|
"edmonton",
|
|
"winnipeg",
|
|
"halifax",
|
|
"victoria",
|
|
}
|
|
)
|
|
|
|
|
|
def block_canada_input(event):
|
|
"""Block user input that mentions Canada or Canadian places.
|
|
|
|
Returns ALLOW for non-request events and requests that don't
|
|
mention Canada-related keywords.
|
|
|
|
:param event: V0 event dict with ``type``, ``target``, ``data``,
|
|
``context`` keys.
|
|
:returns: V0 decision dict.
|
|
"""
|
|
if event.get("type") != "request":
|
|
return _ALLOW
|
|
|
|
data = event.get("data")
|
|
text = data if isinstance(data, str) else str(data)
|
|
lower = text.lower()
|
|
if any(kw in lower for kw in _CANADA_KEYWORDS):
|
|
return {
|
|
"result": "DENY",
|
|
"reason": "Canada-related topics are denied.",
|
|
}
|
|
|
|
return _ALLOW
|
|
|
|
|
|
def block_canada_output(event):
|
|
"""Block assistant output that mentions Canada or Canadian places.
|
|
|
|
Returns ALLOW for non-response events and responses that don't
|
|
mention Canada-related keywords.
|
|
|
|
:param event: V0 event dict with ``type``, ``target``, ``data``,
|
|
``context`` keys.
|
|
:returns: V0 decision dict.
|
|
"""
|
|
if event.get("type") != "response":
|
|
return _ALLOW
|
|
|
|
data = event.get("data")
|
|
text = data if isinstance(data, str) else str(data)
|
|
lower = text.lower()
|
|
if any(kw in lower for kw in _CANADA_KEYWORDS):
|
|
return {
|
|
"result": "DENY",
|
|
"reason": "Canada-related topics are denied.",
|
|
}
|
|
|
|
return _ALLOW
|