tools

dsi_search_tools

load_dsi_tool(db_path, run_path='', master_db_folder='')

Load a DSI object from the path and add information to the context for the llm to use.

Parameters:

Name	Type	Description	Default
db_path	str	the path to the DSI object to load	required
run_path	str	the path this code is being run from	''
master_db_folder	str	the folder containing the master database, used to resolve relative paths when loading new databases	''

Returns:

Name	Type	Description
str	str	message indicating success or failure

Source code in src/ursa/tools/dsi_search_tools.py

@tool
def load_dsi_tool(
    db_path: str, run_path: str = "", master_db_folder: str = ""
) -> str:
    """Load a DSI object from the path and add information to the context for the llm to use.

    Args:
        db_path (str): the path to the DSI object to load
        run_path (str): the path this code is being run from
        master_db_folder (str): the folder containing the master database, used to resolve relative paths when loading new databases

    Returns:
        str: message indicating success or failure
    """
    if no_DSI:
        return "Optional dependence [dsi] not installed. Tool will not work."

    master_database_previously_set = True
    if master_db_folder == "":
        master_database_previously_set = False
        # the ai is loading the master database for the first time
        master_database_path, master_db_folder = _get_db_abs_path(
            db_path, run_path
        )
        data_path = master_database_path.strip()
    else:
        p = Path(db_path).expanduser()
        if not p.is_absolute():
            _db_path = str(Path(master_db_folder).expanduser() / p)
        else:
            _db_path = str(p)

        data_path = _db_path.strip()

    if not _check_db_valid(data_path):
        return f"Failed to load DSI database at: {data_path}. Please check the path and ensure it points to a valid DSI .db file."

    try:
        _, _db_schema, _db_description = _get_db_info(data_path)
        _current_db_abs_path = data_path
        schema_text = json.dumps(_db_schema, indent=2, ensure_ascii=False)

        if master_database_previously_set is False:
            return f"""
- Current working database path (current_db_abs_path): {_current_db_abs_path}
- Master database path (master_database_path): {master_database_path}
- Master database folder (master_db_folder): {master_db_folder}
- Current database schema: {schema_text}
- Database description: {_db_description}
""".strip()
        else:
            return f"""
- Current working database path (current_db_abs_path): {_current_db_abs_path}
- Current database schema: {schema_text}
- Database description: {_db_description}
""".strip()

    except Exception as e:  # noqa: BLE001
        return f"Failed to load database information: {e}"

query_dsi_tool(query_str, db_path)

Execute a SQL query on a DSI object

Arg

query_str (str): the SQL query to run on DSI object db_path (str): the absolute path to the DSI database to query

Returns:

Name	Type	Description
collection	dict	the results of the query

Source code in src/ursa/tools/dsi_search_tools.py

@tool
def query_dsi_tool(query_str: str, db_path: str) -> dict:
    """Execute a SQL query on a DSI object

    Arg:
        query_str (str): the SQL query to run on DSI object
        db_path (str): the absolute path to the DSI database to query

    Returns:
        collection: the results of the query
    """
    if no_DSI:
        return {
            "error": "Optional dependency [dsi] not installed. Tool will not work."
        }

    # print(f"query {query_str}, db_path: {db_path}")

    _store = None
    try:
        with redirect_stdout(_NULL), redirect_stderr(_NULL):
            _store = DSI(db_path, check_same_thread=False)
            df = _store.query(query_str, collection=True)

        if df is None:
            return {}
        return df.to_dict(orient="records")

    except Exception:  # noqa: BLE001
        return {}

    finally:
        if _store is not None:
            try:
                with redirect_stdout(_NULL), redirect_stderr(_NULL):
                    _store.close()
            except Exception as e:  # noqa: BLE001
                print(f"Error closing DSI store: {e}")

feasibility_checker

heuristic_feasibility_check(constraints, variable_name, variable_type, variable_bounds, samples=10000)

A tool for checking feasibility of the constraints.

Parameters:

Name	Type	Description	Default
constraints	Annotated[list[str], "List of strings like 'x0+x1<=5'"]	list of strings like 'x0 + x1 <= 5', etc.	required
variable_name	Annotated[list[str], "List of strings like 'x0', 'x1', etc."]	list of strings containing variable names used in constraint expressions.	required
variable_type	Annotated[list[str], "List of strings like 'real', 'integer', 'boolean', etc."]	list of strings like 'real', 'integer', 'boolean', etc.	required
variable_bounds	Annotated[list[list[float]], "List of (lower bound, upper bound) tuples for x0, x1, ...'"]	list of (lower, upper) tuples for x0, x1, etc.	required
samples	Annotated[int, 'Number of random sample. Default 10000']	number of random samples, default value 10000	10000

Returns:

Type	Description
tuple[str]	A string indicating whether a feasible solution was found.

Source code in src/ursa/tools/feasibility_checker.py

@tool(parse_docstring=True)
def heuristic_feasibility_check(
    constraints: Annotated[list[str], "List of strings like 'x0+x1<=5'"],
    variable_name: Annotated[
        list[str], "List of strings like 'x0', 'x1', etc."
    ],
    variable_type: Annotated[
        list[str], "List of strings like 'real', 'integer', 'boolean', etc."
    ],
    variable_bounds: Annotated[
        list[list[float]],
        "List of (lower bound, upper bound) tuples for x0, x1, ...'",
    ],
    samples: Annotated[int, "Number of random sample. Default 10000"] = 10000,
) -> tuple[str]:
    """
    A tool for checking feasibility of the constraints.

    Args:
        constraints: list of strings like 'x0 + x1 <= 5', etc.
        variable_name: list of strings containing variable names used in constraint expressions.
        variable_type: list of strings like 'real', 'integer', 'boolean', etc.
        variable_bounds: list of (lower, upper) tuples for x0, x1, etc.
        samples: number of random samples, default value 10000

    Returns:
        A string indicating whether a feasible solution was found.
    """

    symbols = sp.symbols(variable_name)

    # Build a dict mapping each name to its Symbol, for parsing
    locals_map = {name: sym for name, sym in zip(variable_name, symbols)}

    # Parse constraints into Sympy Boolean expressions
    parsed_constraints = []
    try:
        for expr in constraints:
            parsed = parse_expr(
                expr,
                local_dict=locals_map,
                transformations=standard_transformations,
                evaluate=False,
            )
            parsed_constraints.append(parsed)
    except Exception as e:
        return f"Error parsing constraints: {e}"

    # Sampling loop
    n = len(parsed_constraints)
    funcs = [
        sp.lambdify(symbols, c, modules=["math", "numpy"])
        for c in parsed_constraints
    ]
    constraint_satisfied = np.zeros(n, dtype=int)
    for _ in range(samples):
        point = {}
        for i, sym in enumerate(symbols):
            typ = variable_type[i].lower()
            low, high = variable_bounds[i]
            if typ == "integer":
                value = random.randint(int(low), int(high))
            elif typ in ("real", "continuous"):
                value = random.uniform(low, high)
            elif typ in ("boolean", "logical"):
                value = random.choice([False, True])
            else:
                raise ValueError(
                    f"Unknown type {variable_type[i]} for variable {variable_name[i]}"
                )
            point[sym] = value

        # Evaluate all constraints at this point
        try:
            vals = [point[s] for s in symbols]
            cons_satisfaction = [
                bool(np.asarray(f(*vals)).all()) for f in funcs
            ]
            if all(cons_satisfaction):
                # Found a feasible point
                readable = {str(k): round(v, 3) for k, v in point.items()}
                return f"Feasible solution found: {readable}"
            else:
                constraint_satisfied += np.array(cons_satisfaction)
        except Exception as e:
            return f"Error evaluating constraint at point {point}: {e}"

    rates = constraint_satisfied / samples  # fraction satisfied per constraint
    order = np.argsort(rates)  # lowest (most violated) first

    lines = []
    for rank, idx in enumerate(order, start=1):
        expr_text = constraints[
            idx
        ]  # use the original string; easier to read than str(sympy_expr)
        sat = constraint_satisfied[idx]
        lines.append(
            f"[C{idx + 1}] {expr_text} — satisfied {sat:,}/{samples:,} ({sat / samples:.1%}), "
            f"violated {1 - sat / samples:.1%}"
        )

    return (
        f"No feasible solution found after {samples:,} samples. Most violated constraints (low→high satisfaction):\n "
        + "\n  ".join(lines)
    )

feasibility_tools

Unified feasibility checker with heuristic pre-check and exact auto-routing.

Backends (imported lazily and used only if available): - PySMT (cvc5/msat/yices/z3) for SMT-style logic, disjunctions, and nonlinear constructs. - OR-Tools CP-SAT for strictly linear integer/boolean instances with integer coefficients. - OR-Tools CBC (pywraplp) for linear MILP/LP (mixed real + integer, or pure LP). - SciPy HiGHS (linprog) for pure continuous LP feasibility.

Install any subset you need

pip install pysmt && pysmt-install --cvc5 # or --z3/--msat/--yices pip install ortools pip install scipy pip install numpy

This file exposes a single LangChain tool: feasibility_check_auto.

feasibility_check_auto(constraints, variable_name, variable_type, variable_bounds, prefer_smt_solver='cvc5', heuristic_enabled=True, heuristic_first=True, heuristic_samples=2000, heuristic_seed=None, heuristic_unbounded_radius_real=1000.0, heuristic_unbounded_radius_int=10 ** 6, numeric_tolerance=1e-08)

Unified feasibility checker with heuristic pre-check and exact auto-routing.

Performs an optional randomized feasibility search. If no witness is found (or the heuristic is disabled), the function auto-routes to an exact backend based on the detected problem structure (PySMT for SMT/logic/nonlinear, OR-Tools CP-SAT for linear integer/boolean, OR-Tools CBC for MILP/LP, or SciPy HiGHS for pure LP).

Parameters:

Name	Type	Description	Default
constraints	Annotated[list[str], "Constraint strings like 'x0 + 2*x1 <= 5' or '(x0<=3) | (x1>=2)'"]	Constraint strings such as "x0 + 2*x1 <= 5" or "(x0<=3) | (x1>=2)".	required
variable_name	Annotated[list[str], ['x0', 'x1', ...]]	Variable names, e.g., ["x0", "x1"].	required
variable_type	Annotated[list[str], ['real' | 'integer' | 'boolean', ...]]	Variable types aligned with variable_name. Each must be one of "real", "integer", or "boolean".	required
variable_bounds	Annotated[list[list[Optional[float]]], '[(low, high), ...] (use None for unbounded)']	Per-variable [low, high] bounds aligned with variable_name. Use None to denote an unbounded side.	required
prefer_smt_solver	Annotated[str, "SMT backend if needed: 'cvc5'|'msat'|'yices'|'z3'"]	SMT backend name used by PySMT ("cvc5", "msat", "yices", or "z3").	'cvc5'
heuristic_enabled	Annotated[bool, 'Run a fast randomized search first?']	Whether to run the heuristic sampler.	True
heuristic_first	Annotated[bool, 'Try heuristic before exact routing']	If True, run the heuristic before exact routing; if False, run it after.	True
heuristic_samples	Annotated[int, 'Samples for heuristic search']	Number of heuristic samples.	2000
heuristic_seed	Annotated[Optional[int], 'Seed for reproducibility']	Random seed for reproducibility.	None
heuristic_unbounded_radius_real	Annotated[float, 'Sampling range for unbounded real vars']	Sampling radius for unbounded real variables.	1000.0
heuristic_unbounded_radius_int	Annotated[int, 'Sampling range for unbounded integer vars']	Sampling radius for unbounded integer variables.	10 ** 6
numeric_tolerance	Annotated[float, 'Tolerance for relational checks (Eq/Lt/Le/etc.)']	Tolerance used in relational checks (e.g., Eq, Lt, Le).	1e-08

Returns:

Type	Description
str	A message indicating the chosen backend and the feasibility result. On success,
str	includes an example model (assignment). On infeasibility, includes a short
str	diagnostic or solver status.

Raises:

Type	Description
ValueError	If constraints cannot be parsed or an unsupported variable type is provided.

Source code in src/ursa/tools/feasibility_tools.py

@tool(parse_docstring=True)
def feasibility_check_auto(
    constraints: Annotated[
        list[str],
        "Constraint strings like 'x0 + 2*x1 <= 5' or '(x0<=3) | (x1>=2)'",
    ],
    variable_name: Annotated[list[str], "['x0','x1',...]"],
    variable_type: Annotated[list[str], "['real'|'integer'|'boolean', ...]"],
    variable_bounds: Annotated[
        list[list[Optional[float]]],
        "[(low, high), ...] (use None for unbounded)",
    ],
    prefer_smt_solver: Annotated[
        str, "SMT backend if needed: 'cvc5'|'msat'|'yices'|'z3'"
    ] = "cvc5",
    heuristic_enabled: Annotated[
        bool, "Run a fast randomized search first?"
    ] = True,
    heuristic_first: Annotated[
        bool, "Try heuristic before exact routing"
    ] = True,
    heuristic_samples: Annotated[int, "Samples for heuristic search"] = 2000,
    heuristic_seed: Annotated[Optional[int], "Seed for reproducibility"] = None,
    heuristic_unbounded_radius_real: Annotated[
        float, "Sampling range for unbounded real vars"
    ] = 1e3,
    heuristic_unbounded_radius_int: Annotated[
        int, "Sampling range for unbounded integer vars"
    ] = 10**6,
    numeric_tolerance: Annotated[
        float, "Tolerance for relational checks (Eq/Lt/Le/etc.)"
    ] = 1e-8,
) -> str:
    """Unified feasibility checker with heuristic pre-check and exact auto-routing.

    Performs an optional randomized feasibility search. If no witness is found (or the
    heuristic is disabled), the function auto-routes to an exact backend based on the
    detected problem structure (PySMT for SMT/logic/nonlinear, OR-Tools CP-SAT for
    linear integer/boolean, OR-Tools CBC for MILP/LP, or SciPy HiGHS for pure LP).

    Args:
        constraints: Constraint strings such as "x0 + 2*x1 <= 5" or "(x0<=3) | (x1>=2)".
        variable_name: Variable names, e.g., ["x0", "x1"].
        variable_type: Variable types aligned with `variable_name`. Each must be one of
            "real", "integer", or "boolean".
        variable_bounds: Per-variable [low, high] bounds aligned with `variable_name`.
            Use None to denote an unbounded side.
        prefer_smt_solver: SMT backend name used by PySMT ("cvc5", "msat", "yices", or "z3").
        heuristic_enabled: Whether to run the heuristic sampler.
        heuristic_first: If True, run the heuristic before exact routing; if False, run it after.
        heuristic_samples: Number of heuristic samples.
        heuristic_seed: Random seed for reproducibility.
        heuristic_unbounded_radius_real: Sampling radius for unbounded real variables.
        heuristic_unbounded_radius_int: Sampling radius for unbounded integer variables.
        numeric_tolerance: Tolerance used in relational checks (e.g., Eq, Lt, Le).

    Returns:
        A message indicating the chosen backend and the feasibility result. On success,
        includes an example model (assignment). On infeasibility, includes a short
        diagnostic or solver status.

    Raises:
        ValueError: If constraints cannot be parsed or an unsupported variable type is provided.
    """
    # 1) Parse
    try:
        symbols, sympy_cons = _parse_constraints(constraints, variable_name)
    except Exception as e:
        return f"Parse error: {e}"

    # 2) Heuristic (optional)
    if heuristic_enabled and heuristic_first:
        try:
            h_model = _heuristic_feasible(
                sympy_cons,
                symbols,
                variable_name,
                variable_type,
                variable_bounds,
                samples=heuristic_samples,
                seed=heuristic_seed,
                tol=numeric_tolerance,
                unbounded_radius_real=heuristic_unbounded_radius_real,
                unbounded_radius_int=heuristic_unbounded_radius_int,
            )
            if h_model is not None:
                return f"[backend=heuristic] Feasible (sampled witness). Example solution: {h_model}"
        except Exception:
            # Ignore heuristic issues and continue to exact route
            pass

    # 3) Classify & route
    info = _classify(sympy_cons, symbols, variable_type)

    # SMT needed or nonlinear / non-conj
    if info["requires_smt"] or not info["all_linear"]:
        res = _solve_with_pysmt(
            sympy_cons,
            symbols,
            variable_name,
            variable_type,
            variable_bounds,
            solver_name=prefer_smt_solver,
        )
        # Optional heuristic after exact if requested
        if (
            heuristic_enabled
            and not heuristic_first
            and any(
                kw in res.lower()
                for kw in ("unknown", "not installed", "unsupported", "failed")
            )
        ):
            h_model = _heuristic_feasible(
                sympy_cons,
                symbols,
                variable_name,
                variable_type,
                variable_bounds,
                samples=heuristic_samples,
                seed=heuristic_seed,
                tol=numeric_tolerance,
                unbounded_radius_real=heuristic_unbounded_radius_real,
                unbounded_radius_int=heuristic_unbounded_radius_int,
            )
            if h_model is not None:
                return f"[backend=heuristic] Feasible (sampled witness). Example solution: {h_model}"
        return res

    # Linear-only path: collect atomic conjuncts
    conjuncts: list[sp.Expr] = []
    for c in sympy_cons:
        atoms, _ = _flatten_conjunction(c)
        conjuncts.extend(atoms)

    has_int, has_bool, has_real = (
        info["has_int"],
        info["has_bool"],
        info["has_real"],
    )

    # Pure LP (continuous only)
    if not has_int and not has_bool and has_real:
        res = _solve_with_highs_lp(
            conjuncts, symbols, variable_name, variable_bounds
        )
        if "not installed" in res.lower():
            res = _solve_with_cbc_milp(
                conjuncts,
                symbols,
                variable_name,
                variable_type,
                variable_bounds,
            )
        if (
            heuristic_enabled
            and not heuristic_first
            and any(kw in res.lower() for kw in ("failed", "unknown"))
        ):
            h_model = _heuristic_feasible(
                sympy_cons,
                symbols,
                variable_name,
                variable_type,
                variable_bounds,
                samples=heuristic_samples,
                seed=heuristic_seed,
                tol=numeric_tolerance,
                unbounded_radius_real=heuristic_unbounded_radius_real,
                unbounded_radius_int=heuristic_unbounded_radius_int,
            )
            if h_model is not None:
                return f"[backend=heuristic] Feasible (sampled witness). Example solution: {h_model}"
        return res

    # All integer/boolean → CP-SAT first (if integer coefficients), else CBC MILP
    if (has_int or has_bool) and not has_real:
        res = _solve_with_cpsat_integer_boolean(
            conjuncts, symbols, variable_name, variable_type, variable_bounds
        )
        if (
            any(
                kw in res
                for kw in (
                    "routing to MILP/LP",
                    "handles linear conjunctions only",
                )
            )
            or "not installed" in res.lower()
        ):
            res = _solve_with_cbc_milp(
                conjuncts,
                symbols,
                variable_name,
                variable_type,
                variable_bounds,
            )
        return res

    # Mixed reals + integers → CBC MILP
    res = _solve_with_cbc_milp(
        conjuncts, symbols, variable_name, variable_type, variable_bounds
    )

    # Optional heuristic after exact (if backend missing/failing)
    if (
        heuristic_enabled
        and not heuristic_first
        and any(
            kw in res.lower() for kw in ("not installed", "failed", "status:")
        )
    ):
        h_model = _heuristic_feasible(
            sympy_cons,
            symbols,
            variable_name,
            variable_type,
            variable_bounds,
            samples=heuristic_samples,
            seed=heuristic_seed,
            tol=numeric_tolerance,
            unbounded_radius_real=heuristic_unbounded_radius_real,
            unbounded_radius_int=heuristic_unbounded_radius_int,
        )
        if h_model is not None:
            return f"[backend=heuristic] Feasible (sampled witness). Example solution: {h_model}"

    return res

fm_base_tool

TorchModuleTool

Bases: BaseModel, Generic[Input, Output, ModelInput, ModelOutput]

A helper class for exposing a PyTorch model as an MCP tool for inference. Provides default methods for running the following pipeline:

1. Preprocess a sequence of Inputs into a ModelInput
2. Pass ModelInput through the PyTorch model getting ModelOutput
3. Postprocess ModelOutput into a suitable sequence of Outputs

Complex models (i.e. multi-GPU) may not be fully supported by this class.

Source code in src/ursa/tools/fm_base_tool.py

class TorchModuleTool(
    BaseModel, Generic[Input, Output, ModelInput, ModelOutput]
):
    """
    A helper class for exposing a PyTorch model as an MCP tool for inference.
    Provides default methods for running the following pipeline:

    1. Preprocess a sequence of `Inputs` into a `ModelInput`
    2. Pass `ModelInput` through the PyTorch model getting `ModelOutput`
    3. Postprocess `ModelOutput` into a suitable sequence of `Outputs`

    Complex models (i.e. multi-GPU) may not be fully supported by this class.
    """

    model_config = ConfigDict(arbitrary_types_allowed=True)

    fm: torch.nn.Module
    """ The underlying PyTorch model used for inference """

    name: str = None
    """ A short name for the foundation model """

    description: str
    """ What the foundation model does / how to use it """

    args_schema: type[Input]
    """ The input schema for the model """

    output_schema: type[Output]
    """ The output_schema for the model """

    batch_size: int = 1
    """ Inputs to the model will be batched into set of at most this size """

    device: torch.device = Field(default_factory=default_device)
    """ The accelerator on which the model is placed """

    def preprocess(self, input: Sequence[Input]) -> ModelInput:
        """
        Convert tool input into the form accepted by the model
        The input will be of type `list[args_schema]` with a length
        of `batch_size`

        Defaults to `torch.data.default_collate`
        """

        return default_collate(list(input))

    def _forward(self, model_inputs: ModelInput) -> ModelOutput:
        """Process a batch of observations with the model"""
        return self.fm(model_inputs).to("cpu")

    def postprocess(self, model_output: ModelOutput) -> Iterable[Output]:
        """Postprocess the model's raw output into a relevant tool output format"""
        yield from model_output

    @classmethod
    def from_pretrained(
        cls, pretrained_model_name_or_path: str, **kwargs
    ) -> "TorchModuleTool":
        """Instantiate tool from a pretrained checkpoint either on disk,
        or automatically downloaded from a repository (i.e. HuggingFace)
        """
        raise NotImplementedError()

    def model_post_init(self, __context) -> None:
        # Move the model to the indicated device
        self.fm = self.fm.to(self.device)

        # Default to the class name
        if self.name is None:
            self.name = self.__class__.__name__

    @final
    def batch(self, inputs: list[Input], **kwargs) -> list[Output]:
        return list(self.batch_as_completed(inputs, **kwargs))

    @final
    def batch_as_completed(
        self,
        inputs: list[Input],
        max_concurency: int | None = None,
    ) -> Iterable[Output]:
        n = max_concurency or self.batch_size
        for batch in batched(inputs, n=n):
            from torch import inference_mode

            with inference_mode():
                batch = self.preprocess(batch)
                y = self._forward(batch)
                yield from self.postprocess(y)

    @final
    def __call__(self, input: Input):
        with torch.inference_mode():
            batch = self.preprocess([input])
            y = self._forward(batch)
            return next(iter(self.postprocess(y)))

    @final
    def __to_fastmcp(self) -> FastMCPTool:
        field_definitions = {
            field: (field_info.annotation, field_info)
            for field, field_info in self.args_schema.model_fields.items()
        }
        arg_model = create_model(
            f"{self.name}arguments",
            **field_definitions,
            __base__=ArgModelBase,
        )
        fn_metadata = FuncMetadata(
            arg_model=arg_model,
            output_model=self.output_schema,
            output_schema=self.output_schema.model_json_schema(),
        )

        async def fn(**input) -> self.output_schema:
            x = self.args_schema(**input)
            return self(x)

        return FastMCPTool(
            fn=fn,
            name=self.name,
            description=self.description,
            parameters=self.args_schema.model_json_schema(),
            fn_metadata=fn_metadata,
            is_async=True,
        )

    @final
    def add_to_fastmcp(self, server: FastMCP) -> FastMCPTool:
        """Add `self` as a tool to `server`"""
        fasttool = self.__to_fastmcp()

        if fasttool.name not in server._tool_manager._tools:
            server._tool_manager._tools[fasttool.name] = fasttool

        elif server._tool_manager.warn_on_duplicate_tools:
            logging.warning(f"Tool already exists: {fasttool.name}")

        return fasttool

args_schema instance-attribute

The input schema for the model

batch_size = 1 class-attribute instance-attribute

Inputs to the model will be batched into set of at most this size

description instance-attribute

What the foundation model does / how to use it

device = Field(default_factory=default_device) class-attribute instance-attribute

The accelerator on which the model is placed

fm instance-attribute

The underlying PyTorch model used for inference

name = None class-attribute instance-attribute

A short name for the foundation model

output_schema instance-attribute

The output_schema for the model

add_to_fastmcp(server)

Add self as a tool to server

Source code in src/ursa/tools/fm_base_tool.py

@final
def add_to_fastmcp(self, server: FastMCP) -> FastMCPTool:
    """Add `self` as a tool to `server`"""
    fasttool = self.__to_fastmcp()

    if fasttool.name not in server._tool_manager._tools:
        server._tool_manager._tools[fasttool.name] = fasttool

    elif server._tool_manager.warn_on_duplicate_tools:
        logging.warning(f"Tool already exists: {fasttool.name}")

    return fasttool

from_pretrained(pretrained_model_name_or_path, **kwargs) classmethod

Instantiate tool from a pretrained checkpoint either on disk, or automatically downloaded from a repository (i.e. HuggingFace)

Source code in src/ursa/tools/fm_base_tool.py

@classmethod
def from_pretrained(
    cls, pretrained_model_name_or_path: str, **kwargs
) -> "TorchModuleTool":
    """Instantiate tool from a pretrained checkpoint either on disk,
    or automatically downloaded from a repository (i.e. HuggingFace)
    """
    raise NotImplementedError()

postprocess(model_output)

Postprocess the model's raw output into a relevant tool output format

Source code in src/ursa/tools/fm_base_tool.py

def postprocess(self, model_output: ModelOutput) -> Iterable[Output]:
    """Postprocess the model's raw output into a relevant tool output format"""
    yield from model_output

preprocess(input)

Convert tool input into the form accepted by the model The input will be of type list[args_schema] with a length of batch_size

Defaults to torch.data.default_collate

Source code in src/ursa/tools/fm_base_tool.py

def preprocess(self, input: Sequence[Input]) -> ModelInput:
    """
    Convert tool input into the form accepted by the model
    The input will be of type `list[args_schema]` with a length
    of `batch_size`

    Defaults to `torch.data.default_collate`
    """

    return default_collate(list(input))

git_tools

git_add(runtime, repo_path=None, pathspecs=None)

Stage files for commit using git add.

Source code in src/ursa/tools/git_tools.py

@tool
def git_add(
    runtime: ToolRuntime[AgentContext],
    repo_path: AsciiStr | None = None,
    pathspecs: list[AsciiStr] | None = None,
) -> str:
    """Stage files for commit using git add."""
    repo = _repo_path(repo_path, runtime)
    if not pathspecs:
        return _format_result("", "No pathspecs provided to git_add")
    return _run_git(repo, ["add", "--", *list(pathspecs)])

git_commit(runtime, message, repo_path=None)

Create a git commit with the provided message.

Source code in src/ursa/tools/git_tools.py

@tool
def git_commit(
    runtime: ToolRuntime[AgentContext],
    message: AsciiStr,
    repo_path: AsciiStr | None = None,
) -> str:
    """Create a git commit with the provided message."""
    repo = _repo_path(repo_path, runtime)
    if not message.strip():
        return _format_result("", "Commit message must not be empty")
    return _run_git(repo, ["commit", "--message", message])

git_create_branch(runtime, branch, repo_path=None)

Create a branch without switching to it.

Source code in src/ursa/tools/git_tools.py

@tool
def git_create_branch(
    runtime: ToolRuntime[AgentContext],
    branch: AsciiStr,
    repo_path: AsciiStr | None = None,
) -> str:
    """Create a branch without switching to it."""
    repo = _repo_path(repo_path, runtime)
    err = _check_ref_format(repo, branch)
    if err:
        return _format_result("", err)
    return _run_git(repo, ["branch", branch])

git_diff(runtime, repo_path=None, staged=False, pathspecs=None)

Return git diff for a repository inside the workspace.

Source code in src/ursa/tools/git_tools.py

@tool
def git_diff(
    runtime: ToolRuntime[AgentContext],
    repo_path: AsciiStr | None = None,
    staged: bool = False,
    pathspecs: list[AsciiStr] | None = None,
) -> str:
    """Return git diff for a repository inside the workspace."""
    repo = _repo_path(repo_path, runtime)
    args = ["diff"]
    if staged:
        args.append("--staged")
    if pathspecs:
        args.append("--")
        args.extend(list(pathspecs))
    return _run_git(repo, args)

git_log(runtime, repo_path=None, limit=20)

Return recent git log entries for a repository.

Source code in src/ursa/tools/git_tools.py

@tool
def git_log(
    runtime: ToolRuntime[AgentContext],
    repo_path: AsciiStr | None = None,
    limit: int = 20,
) -> str:
    """Return recent git log entries for a repository."""
    repo = _repo_path(repo_path, runtime)
    limit = max(1, int(limit))
    return _run_git(repo, ["log", f"-n{limit}", "--oneline", "--decorate"])

git_ls_files(runtime, repo_path=None, pathspecs=None)

List tracked files, optionally filtered by pathspecs.

Source code in src/ursa/tools/git_tools.py

@tool
def git_ls_files(
    runtime: ToolRuntime[AgentContext],
    repo_path: AsciiStr | None = None,
    pathspecs: list[AsciiStr] | None = None,
) -> str:
    """List tracked files, optionally filtered by pathspecs."""
    repo = _repo_path(repo_path, runtime)
    args = ["ls-files"]
    if pathspecs:
        args.append("--")
        args.extend(list(pathspecs))
    return _run_git(repo, args)

git_status(runtime, repo_path=None)

Return git status for a repository inside the workspace.

Parameters:

Name	Type	Description	Default
repo_path	AsciiStr | None	Path to repository relative to workspace. If None, uses workspace root. Recommended to always specify a repo_path to avoid large untracked file lists.	None

Source code in src/ursa/tools/git_tools.py

@tool
def git_status(
    runtime: ToolRuntime[AgentContext],
    repo_path: AsciiStr | None = None,
) -> str:
    """Return git status for a repository inside the workspace.

    Args:
        repo_path: Path to repository relative to workspace. If None, uses workspace root.
                   Recommended to always specify a repo_path to avoid large untracked file lists.
    """
    repo = _repo_path(repo_path, runtime)

    # Warn if using workspace root without explicit repo_path
    workspace = Path(runtime.context.workspace).absolute()
    if repo_path is None and repo == workspace:
        return (
            "WARNING: git_status called on workspace root without specifying repo_path. "
            "This may show many untracked files. "
            "Please specify a specific repository path (e.g., repo_path='my-project'). "
            "Use list_directory to see available repositories in the workspace first."
        )

    result = _run_git(repo, ["status", "-sb"])

    # Limit output size for very large untracked file lists
    if len(result) > 10000:
        lines = result.split("\n")
        if len(lines) > 100:
            return (
                f"Git status output too large ({len(lines)} lines). "
                f"Showing first 50 and last 50 lines:\n"
                f"{''.join(lines[:50])}\n"
                f"... ({len(lines) - 100} lines omitted) ...\n"
                f"{''.join(lines[-50:])}\n"
                f"Consider using git_status on a specific subdirectory."
            )

    return result

git_switch(runtime, branch, repo_path=None, create=False)

Switch branches using git switch (optionally create).

Source code in src/ursa/tools/git_tools.py

@tool
def git_switch(
    runtime: ToolRuntime[AgentContext],
    branch: AsciiStr,
    repo_path: AsciiStr | None = None,
    create: bool = False,
) -> str:
    """Switch branches using git switch (optionally create)."""
    repo = _repo_path(repo_path, runtime)
    err = _check_ref_format(repo, branch)
    if err:
        return _format_result("", err)
    args = ["switch"]
    if create:
        args.append("-c")
    args.append(branch)
    return _run_git(repo, args)

go_tools

go_build(runtime, repo_path=None)

Build a Go module using go build ./...

Source code in src/ursa/tools/go_tools.py

@tool
def go_build(
    runtime: ToolRuntime[AgentContext],
    repo_path: AsciiStr | None = None,
) -> str:
    """Build a Go module using go build ./..."""
    repo = _repo_path(repo_path, runtime)
    try:
        result = subprocess.run(
            ["go", "build", "./..."],
            text=True,
            capture_output=True,
            timeout=GO_BUILD_TIMEOUT,
            cwd=repo,
            check=False,
        )
    except subprocess.TimeoutExpired:
        return _format_result(
            "",
            f"go build timed out after {GO_BUILD_TIMEOUT}s (5 minutes). "
            "Large builds may need to be run in smaller chunks.",
        )
    except Exception as exc:
        return _format_result("", f"Error running go build: {exc}")
    return _format_result(result.stdout, result.stderr)

go_mod_tidy(runtime, repo_path=None)

Clean up and validate Go module dependencies using go mod tidy.

Source code in src/ursa/tools/go_tools.py

@tool
def go_mod_tidy(
    runtime: ToolRuntime[AgentContext],
    repo_path: AsciiStr | None = None,
) -> str:
    """Clean up and validate Go module dependencies using go mod tidy."""
    repo = _repo_path(repo_path, runtime)
    try:
        result = subprocess.run(
            ["go", "mod", "tidy"],
            text=True,
            capture_output=True,
            timeout=GO_ANALYSIS_TIMEOUT,
            cwd=repo,
            check=False,
        )
    except subprocess.TimeoutExpired:
        return _format_result(
            "", f"go mod tidy timed out after {GO_ANALYSIS_TIMEOUT}s."
        )
    except Exception as exc:
        return _format_result("", f"Error running go mod tidy: {exc}")
    return _format_result(result.stdout, result.stderr)

go_test(runtime, repo_path=None, verbose=True)

Run Go tests using go test ./...

Source code in src/ursa/tools/go_tools.py

@tool
def go_test(
    runtime: ToolRuntime[AgentContext],
    repo_path: AsciiStr | None = None,
    verbose: bool = True,
) -> str:
    """Run Go tests using go test ./..."""
    repo = _repo_path(repo_path, runtime)
    args = ["go", "test"]
    if verbose:
        args.append("-v")
    args.append("./...")
    try:
        result = subprocess.run(
            args,
            text=True,
            capture_output=True,
            timeout=GO_TEST_TIMEOUT,
            cwd=repo,
            check=False,
        )
    except subprocess.TimeoutExpired:
        return _format_result(
            "",
            f"go test timed out after {GO_TEST_TIMEOUT}s (10 minutes). "
            "Large test suites may need to run selectively.",
        )
    except Exception as exc:
        return _format_result("", f"Error running go test: {exc}")
    return _format_result(result.stdout, result.stderr)

go_vet(runtime, repo_path=None)

Run Go vet for code pattern analysis using go vet ./...

Source code in src/ursa/tools/go_tools.py

@tool
def go_vet(
    runtime: ToolRuntime[AgentContext],
    repo_path: AsciiStr | None = None,
) -> str:
    """Run Go vet for code pattern analysis using go vet ./..."""
    repo = _repo_path(repo_path, runtime)
    try:
        result = subprocess.run(
            ["go", "vet", "./..."],
            text=True,
            capture_output=True,
            timeout=GO_ANALYSIS_TIMEOUT,
            cwd=repo,
            check=False,
        )
    except subprocess.TimeoutExpired:
        return _format_result(
            "", f"go vet timed out after {GO_ANALYSIS_TIMEOUT}s."
        )
    except Exception as exc:
        return _format_result("", f"Error running go vet: {exc}")
    return _format_result(result.stdout, result.stderr)

gofmt_files(runtime, paths, repo_path=None)

Format Go files in-place using gofmt.

Source code in src/ursa/tools/go_tools.py

@tool
def gofmt_files(
    runtime: ToolRuntime[AgentContext],
    paths: list[AsciiStr],
    repo_path: AsciiStr | None = None,
) -> str:
    """Format Go files in-place using gofmt."""
    if not paths:
        return _format_result("", "No paths provided to gofmt_files")
    if any(not str(p).endswith(".go") for p in paths):
        return _format_result("", "gofmt_files only accepts .go files")
    repo = _repo_path(repo_path, runtime)
    try:
        result = subprocess.run(
            ["gofmt", "-w", *list(paths)],
            text=True,
            capture_output=True,
            timeout=GO_FORMAT_TIMEOUT,
            cwd=repo,
            check=False,
        )
    except Exception as exc:
        return _format_result("", f"Error running gofmt: {exc}")
    return _format_result(result.stdout, result.stderr)

golangci_lint(runtime, repo_path=None)

Run golangci-lint on the repository.

Automatically detects and uses .golangci.yml if present, otherwise uses sensible defaults.

Source code in src/ursa/tools/go_tools.py

@tool
def golangci_lint(
    runtime: ToolRuntime[AgentContext],
    repo_path: AsciiStr | None = None,
) -> str:
    """Run golangci-lint on the repository.

    Automatically detects and uses .golangci.yml if present,
    otherwise uses sensible defaults.
    """
    from ursa.tools.git_tools import GIT_TIMEOUT

    repo = _repo_path(repo_path, runtime)

    # Check if golangci-lint is installed
    try:
        subprocess.run(
            ["golangci-lint", "--version"],
            text=True,
            capture_output=True,
            timeout=GIT_TIMEOUT,
            check=False,
        )
    except FileNotFoundError:
        return _format_result(
            "",
            "Error: golangci-lint is not installed. "
            "Install it with: go install github.com/golangci/golangci-lint/cmd/golangci-lint@latest",
        )
    except Exception as exc:
        return _format_result("", f"Error checking golangci-lint: {exc}")

    # Run golangci-lint
    config_file = repo / ".golangci.yml"
    args = ["golangci-lint", "run"]
    if config_file.exists():
        args.extend(["--config", str(config_file)])

    try:
        result = subprocess.run(
            args,
            text=True,
            capture_output=True,
            timeout=LINT_TIMEOUT,
            cwd=repo,
            check=False,
        )
    except subprocess.TimeoutExpired:
        return _format_result(
            "", f"golangci-lint timed out after {LINT_TIMEOUT}s (3 minutes)."
        )
    except Exception as exc:
        return _format_result("", f"Error running golangci-lint: {exc}")

    return _format_result(result.stdout, result.stderr)

read_file_tool

download_file_tool(url, output_path, runtime)

Download a file from a URL and save it locally.

Arg

url (str): a string containing the URL of the file to download. output_path (str): the local path where the file should be saved. Default to '.' for workspace root

Returns: Confirmation message with the saved file path.

Source code in src/ursa/tools/read_file_tool.py

@tool
def download_file_tool(
    url: Annotated[AsciiStr, "web link for the file"],
    output_path: Annotated[
        str, "local path to save the file within the workspace"
    ],
    runtime: ToolRuntime[AgentContext],
) -> str:
    """Download a file from a URL and save it locally.

    Arg:
        url (str): a string containing the URL of the file to download.
        output_path (str): the local path where the file should be saved.
                           Default to '.' for workspace root
    Returns:
        Confirmation message with the saved file path.
    """

    try:
        # Download
        response = requests.get(url, stream=True)
        response.raise_for_status()

        # Ensure directory exists
        output_path = runtime.context.workspace / Path(output_path)
        output_path.parent.mkdir(parents=True, exist_ok=True)

        # Write file to disk
        with open(output_path, "wb") as f:
            f.writelines(response.iter_content(chunk_size=8192))

        return f"File successfully downloaded to: {output_path}"

    except Exception as e:  # noqa: BLE001
        return f"Error downloading file: {e!s}"

read_file(filename, runtime)

Read a file from the workspace.

• If filename ends with .pdf, extract text from the PDF.
• If extracted text is very small (likely scanned), optionally run OCR to add a text layer.
• Otherwise read as UTF-8 text.

Parameters:

Name	Type	Description	Default
filename	str	File name relative to the workspace directory.	required

Returns:

Type	Description
str	Extracted text content.

Source code in src/ursa/tools/read_file_tool.py

@tool
def read_file(filename: str, runtime: ToolRuntime[AgentContext]) -> str:
    """Read a file from the workspace.

    - If filename ends with .pdf, extract text from the PDF.
    - If extracted text is very small (likely scanned), optionally run OCR to add a text layer.
    - Otherwise read as UTF-8 text.

    Args:
        filename: File name relative to the workspace directory.

    Returns:
        Extracted text content.
    """
    full_filename = runtime.context.workspace.joinpath(filename)

    print("[READING]:", full_filename)
    # Move all the reading to a function in the parse util
    text = read_text_from_file(full_filename)
    return text

run_command_tool

run_command(query, runtime)

Execute a shell command in the workspace and return its combined output.

Runs the specified command using subprocess.run in the given workspace directory, captures stdout and stderr, enforces a maximum character budget, and formats both streams into a single string. KeyboardInterrupt during execution is caught and reported.

Parameters:

Name	Type	Description	Default
query	AsciiStr	The shell command to execute.	required

Returns:

Type	Description
str	A formatted string with "STDOUT:" followed by the truncated stdout and
str	"STDERR:" followed by the truncated stderr.

Source code in src/ursa/tools/run_command_tool.py

@tool
def run_command(query: AsciiStr, runtime: ToolRuntime[AgentContext]) -> str:
    """Execute a shell command in the workspace and return its combined output.

    Runs the specified command using subprocess.run in the given workspace
    directory, captures stdout and stderr, enforces a maximum character budget,
    and formats both streams into a single string. KeyboardInterrupt during
    execution is caught and reported.

    Args:
        query: The shell command to execute.

    Returns:
        A formatted string with "STDOUT:" followed by the truncated stdout and
        "STDERR:" followed by the truncated stderr.
    """
    workspace_dir = Path(runtime.context.workspace)
    if runtime.store is not None:
        search_results = runtime.store.search(
            ("workspace", "file_edit"), limit=1000
        )
        edited_files = [item.key for item in search_results]
    else:
        edited_files = []

    if runtime.store is not None:
        search_results = runtime.store.search(
            ("workspace", "safe_codes"), limit=1000
        )
        safe_codes = [item.key for item in search_results]
    else:
        safe_codes = []

    llm = runtime.context.llm
    safety_result = llm.with_structured_output(SafetyAssessment).invoke(
        get_safety_prompt(query, safe_codes, edited_files)
    )

    if not safety_result["is_safe"]:
        tool_response = f"[UNSAFE] That command `{query}` was deemed unsafe and cannot be run.\nFor reason: {safety_result['reason']}"
        console.print(
            "[bold red][WARNING][/bold red] Command deemed unsafe:",
            query,
        )
        # Also surface the model's rationale for transparency.
        console.print("[bold red][WARNING][/bold red] REASON:", tool_response)
        return tool_response
    else:
        console.print(
            f"[green]Command passed safety check:[/green] {query}\nFor reason: {safety_result['reason']}"
        )

    print("RUNNING: ", query)

    try:
        result = subprocess.run(
            query,
            text=True,
            shell=True,
            timeout=60000,
            capture_output=True,
            cwd=workspace_dir,
        )
        stdout, stderr = result.stdout, result.stderr
    except KeyboardInterrupt:
        print("Keyboard Interrupt of command: ", query)
        stdout, stderr = "", "KeyboardInterrupt:"

    # Fit BOTH streams under a single overall cap
    stdout_fit, stderr_fit = _fit_streams_to_budget(
        stdout or "", stderr or "", runtime.context.tool_character_limit
    )

    print("STDOUT: ", stdout_fit)
    print("STDERR: ", stderr_fit)

    return f"STDOUT:\n{stdout_fit}\nSTDERR:\n{stderr_fit}"

search_tools

run_arxiv_search(prompt, query, runtime, max_results=3)

Search ArXiv for the first 'max_results' papers and summarize them in the context of the user prompt

Parameters:

Name	Type	Description	Default
prompt	str	string describing the information the agent is interested in from arxiv papers	required
query	str	1 and 8 word search query for the Arxiv search API to find papers relevant to the prompt	required
max_results	int	integer number of papers to return (defaults 3). Request fewer if searching for something very specific or a larger number if broadly searching for information. Do not exceeed 10.	3

Source code in src/ursa/tools/search_tools.py

@tool
def run_arxiv_search(
    prompt: str, query: str, runtime: ToolRuntime, max_results: int = 3
):
    """
    Search ArXiv for the first 'max_results' papers and summarize them in the context
    of the user prompt

    Arguments:
        prompt:
            string describing the information the agent is interested in from arxiv papers
        query:
            1 and 8 word search query for the Arxiv search API to find papers relevant to the prompt
        max_results:
            integer number of papers to return (defaults 3). Request fewer if searching for something
            very specific or a larger number if broadly searching for information. Do not exceeed 10.
    """
    try:
        agent = ArxivAgent(
            llm=runtime.context.llm,
            summarize=True,
            process_images=False,
            max_results=max_results,
            workspace=runtime.context.workspace,
            # rag_embedding=self.embedding,
            database_path=Path("./arxiv_downloaded"),
            summaries_path=Path("./arxiv_summaries"),
            download=True,
        )
        console.print(f"[bold cyan]Searching ArXiv for: [default]{query}")
        assert isinstance(query, str)

        arxiv_result = asyncio.run(
            agent.ainvoke(
                arxiv_search_query=query,
                context=prompt,
            )
        )["final_summary"]

        console.print(
            Panel(
                f"{arxiv_result}",
                title=f"[bold cyan on black]ArXiv summary for {query}",
                border_style="cyan on black",
                style="cyan on black",
            )
        )
        return f"[ArXiv Agent Output]:\n {arxiv_result}"
    except Exception as e:
        return f"Unexpected error while running ArxivAgent: {e}"

run_osti_search(prompt, query, runtime, max_results=3)

Search OSTI.gov for the first 'max_results' papers and summarize them in the context of the user prompt

Parameters:

Name	Type	Description	Default
prompt	str	string describing the information the agent is interested in from arxiv papers	required
query	str	1 and 8 word search query for the OSTI.gov search API to find papers relevant to the prompt	required
max_results	int	integer number of papers to return (defaults 3). Request fewer if searching for something very specific or a larger number if broadly searching for information. Do not exceeed 10.	3

Source code in src/ursa/tools/search_tools.py

@tool
def run_osti_search(
    prompt: str,
    query: str,
    runtime: ToolRuntime,
    max_results: int = 3,
):
    """
    Search OSTI.gov for the first 'max_results' papers and summarize them in the context
    of the user prompt

    Arguments:
        prompt:
            string describing the information the agent is interested in from arxiv papers
        query:
            1 and 8 word search query for the OSTI.gov search API to find papers relevant to the prompt
        max_results:
            integer number of papers to return (defaults 3). Request fewer if searching for something
            very specific or a larger number if broadly searching for information. Do not exceeed 10.
    """
    max_results
    try:
        agent = OSTIAgent(
            llm=runtime.context.llm,
            summarize=True,
            process_images=False,
            max_results=max_results,
            workspace=runtime.context.workspace,
            # rag_embedding=self.embedding,
            database_path=Path("./osti_downloaded_papers"),
            summaries_path=Path("./osti_generated_summaries"),
            vectorstore_path=Path("./osti_vectorstores"),
            download=True,
        )
        console.print(f"[bold cyan]Searching OSTI.gov for: [default]{query}")
        assert isinstance(query, str)

        osti_result = asyncio.run(
            agent.ainvoke(
                query=query,
                context=prompt,
            )
        )["final_summary"]

        console.print(
            Panel(
                f"[cyan on black]{osti_result}",
                title=f"[bold cyan on black]OSTI.gov summary for {query}",
                border_style="cyan on black",
                style="cyan on black",
            )
        )
        return f"[OSTI Agent Output]:\n {osti_result}"
    except Exception as e:
        return f"Unexpected error while running OSTIAgent: {e}"

run_web_search(prompt, query, runtime, max_results=3)

Search the internet for the first 'max_results' pages and summarize them in the context of the user prompt

Parameters:

Name	Type	Description	Default
prompt	str	string describing the information the agent is interested in from websites	required
query	str	1 and 8 word search query for the web search engines to find papers relevant to the prompt	required
max_results	int	integer number of pages to return (defaults 3). Request fewer if searching for something very specific or a larger number if broadly searching for information. Do not exceeed 10.	3

Source code in src/ursa/tools/search_tools.py

@tool
def run_web_search(
    prompt: str,
    query: str,
    runtime: ToolRuntime,
    max_results: int = 3,
):
    """
    Search the internet for the first 'max_results' pages and summarize them in the context
    of the user prompt

    Arguments:
        prompt:
            string describing the information the agent is interested in from websites
        query:
            1 and 8 word search query for the web search engines to find papers relevant to the prompt
        max_results:
            integer number of pages to return (defaults 3). Request fewer if searching for something
            very specific or a larger number if broadly searching for information. Do not exceeed 10.
    """
    try:
        agent = WebSearchAgent(
            llm=runtime.context.llm,
            summarize=True,
            process_images=False,
            max_results=max_results,
            workspace=runtime.context.workspace,
            # rag_embedding=self.embedding,
            database_path=Path("./web_downloads"),
            summaries_path=Path("./web_summaries"),
            download=True,
        )
        console.print(f"[bold cyan]Searching Web for: [default]{query}")
        assert isinstance(query, str)

        web_result = asyncio.run(
            agent.ainvoke(
                query=query,
                context=prompt,
            )
        )["final_summary"]

        console.print(
            Panel(
                f"{web_result}",
                title=f"[bold cyan on black]Web summary for {query}",
                border_style="cyan on black",
                style="cyan on black",
            )
        )
        return f"[Web Search Agent Output]:\n {web_result}"
    except Exception as e:
        return f"Unexpected error while running WebSearchAgent: {e}"

write_code_tool

edit_code(old_code, new_code, filename, runtime, repo_path=None)

Replace the first occurrence of old_code with new_code in filename.

Parameters:

Name	Type	Description	Default
old_code	str	Code fragment to search for.	required
new_code	str	Replacement fragment.	required
filename	AsciiStr	Target file inside the workspace.	required
repo_path	AsciiStr | None	Optional repo path - if provided, file must be within this repo.	None

Returns:

Type	Description
str	Success / failure message.

Source code in src/ursa/tools/write_code_tool.py

@tool
def edit_code(
    old_code: str,
    new_code: str,
    filename: AsciiStr,
    runtime: ToolRuntime[AgentContext],
    repo_path: AsciiStr | None = None,
) -> str:
    """Replace the **first** occurrence of *old_code* with *new_code* in *filename*.

    Args:
        old_code: Code fragment to search for.
        new_code: Replacement fragment.
        filename: Target file inside the workspace.
        repo_path: Optional repo path - if provided, file must be within this repo.

    Returns:
        Success / failure message.
    """
    workspace_dir = runtime.context.workspace
    console.print("[cyan]Editing file:[/cyan]", filename)

    # Validate file path
    repo = None
    if repo_path:
        repo, error = _resolve_repo_dir(
            repo_path,
            workspace_dir,
            "edit",
            filename,
        )
        if error:
            return error

    code_file, error = _validate_file_path(
        filename,
        workspace_dir,
        repo,
    )
    if error:
        console.print(
            f"[bold bright_white on red] :heavy_multiplication_x: [/] [red]{error}[/]"
        )
        return f"Failed to edit {filename}: {error}"
    if code_file is None:
        return f"Failed to edit {filename}: Invalid file path."

    try:
        content = read_text_file(code_file)
    except FileNotFoundError:
        console.print(
            "[bold bright_white on red] :heavy_multiplication_x: [/] "
            "[red]File not found:[/]",
        )
        return f"Failed: {filename} not found."
    except ValueError as exc:
        return f"Failed to edit {filename}: {exc}"
    except OSError as exc:
        return f"Failed to edit {filename}: Could not read file: {exc}"

    # Clean up markdown fences
    old_code_clean = old_code
    new_code_clean = new_code

    if old_code_clean not in content:
        console.print(
            "[yellow] ⚠️ 'old_code' not found in file'; no changes made.[/]"
        )
        return f"No changes made to {filename}: 'old_code' not found in file."

    updated = content.replace(old_code_clean, new_code_clean, 1)

    console.print(
        Panel(
            DiffRenderer(content, updated, filename),
            title="Diff Preview",
            border_style="cyan",
        )
    )

    try:
        with open(code_file, "w", encoding="utf-8") as f:
            f.write(updated)
    except Exception as exc:
        console.print(
            "[bold bright_white on red] :heavy_multiplication_x: [/] "
            "[red]Failed to write file:[/]",
            exc,
        )
        return f"Failed to edit {filename}: {exc}"

    console.print(
        f"[bold bright_white on green] :heavy_check_mark: [/] "
        f"[green]File updated:[/] {code_file}"
    )

    # Record the edit operation
    if (store := runtime.store) is not None:
        store.put(
            ("workspace", "file_edit"),
            filename,
            {
                "modified": time.time(),
                "tool_call_id": runtime.tool_call_id,
                "thread_id": runtime.config.get("metadata", {}).get(
                    "thread_id", None
                ),
            },
        )
    return f"File {filename} updated successfully."

write_code(code, filename, runtime)

Write source code to a file

Records successful file edits to the graph's store

Parameters:

Name	Type	Description	Default
code	str	The source code content to be written to disk.	required
filename	AsciiStr	Name of the target file (including its extension).	required

Source code in src/ursa/tools/write_code_tool.py

@tool(description="Write source code to a file")
def write_code(
    code: str,
    filename: AsciiStr,
    runtime: ToolRuntime[AgentContext],
) -> str:
    """Write source code to a file

    Records successful file edits to the graph's store

    Args:
        code: The source code content to be written to disk.
        filename: Name of the target file (including its extension).

    """
    return _write_code_file(code, filename, runtime)

write_code_with_repo(code, filename, runtime, repo_path)

Write source code to a file constrained to a repository path.

Parameters:

Name	Type	Description	Default
code	str	The source code content to be written to disk.	required
filename	AsciiStr	Name of the target file (including its extension).	required
repo_path	AsciiStr	Repo path - file must resolve within this directory.	required

Source code in src/ursa/tools/write_code_tool.py

@tool(description="Write source code to a file within a repository boundary")
def write_code_with_repo(
    code: str,
    filename: AsciiStr,
    runtime: ToolRuntime[AgentContext],
    repo_path: AsciiStr,
) -> str:
    """Write source code to a file constrained to a repository path.

    Args:
        code: The source code content to be written to disk.
        filename: Name of the target file (including its extension).
        repo_path: Repo path - file must resolve within this directory.
    """
    workspace_dir = runtime.context.workspace
    repo, error = _resolve_repo_dir(repo_path, workspace_dir, "write", filename)
    if error:
        return error

    return _write_code_file(code, filename, runtime, repo)