ursa-ai 0.9.1__py3-none-any.whl

ursa/tools/run_command.py

@@ -0,0 +1,27 @@
+import subprocess
+
+from langchain_core.tools import tool
+
+
+@tool
+def run_cmd(query: str, workspace_dir: str) -> str:
+    """Run command from commandline in the directory workspace_dir"""
+
+    print("RUNNING: ", query)
+    print(
+        "DANGER DANGER DANGER - THERE IS NO GUARDRAIL FOR SAFETY IN THIS IMPLEMENTATION - DANGER DANGER DANGER"
+    )
+    process = subprocess.Popen(
+        query.split(" "),
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+        text=True,
+        cwd=workspace_dir,
+    )
+
+    stdout, stderr = process.communicate(timeout=600)
+
+    print("STDOUT: ", stdout)
+    print("STDERR: ", stderr)
+
+    return f"STDOUT: {stdout} and STDERR: {stderr}"

ursa/tools/write_code.py

@@ -0,0 +1,42 @@
+import os
+
+from langchain_core.tools import tool
+
+
+@tool
+def write_python(code: str, filename: str, workspace_dir: str) -> str:
+    """
+    Writes code to a file in the given workspace.
+
+    Args:
+        code: The code to write
+        filename: the filename to write
+
+    Returns:
+        File writing status: string
+    """
+    print("Writing filename ", filename)
+    try:
+        # Extract code if wrapped in markdown code blocks
+        if "```" in code:
+            code_parts = code.split("```")
+            if len(code_parts) >= 3:
+                # Extract the actual code
+                if "\n" in code_parts[1]:
+                    code = "\n".join(code_parts[1].strip().split("\n")[1:])
+                else:
+                    code = code_parts[2].strip()
+
+        # Write code to a file
+        code_file = os.path.join(workspace_dir, filename)
+
+        with open(code_file, "w") as f:
+            f.write(code)
+        print(f"Written code to file: {code_file}")
+
+        return f"File {filename} written successfully."
+
+    except Exception as e:
+        print(f"Error generating code: {str(e)}")
+        # Return minimal code that prints the error
+        return f"Failed to write {filename} successfully."

ursa/util/diff_renderer.py

@@ -0,0 +1,128 @@
+import difflib
+import re
+from dataclasses import dataclass
+
+from rich.console import Console, ConsoleOptions, RenderResult
+from rich.syntax import Syntax
+from rich.text import Text
+
+# unified diff hunk header regex
+_HUNK_RE = re.compile(r"^@@ -(\d+)(?:,\d+)? \+(\d+)(?:,\d+)? @@")
+
+
+@dataclass
+class _LineStyle:
+    prefix: str
+    bg: str
+
+
+_STYLE = {
+    "add": _LineStyle("+ ", "on #003000"),
+    "del": _LineStyle("- ", "on #300000"),
+    "ctx": _LineStyle("  ", "on grey15"),
+}
+
+
+class DiffRenderer:
+    """Renderable diff—`console.print(DiffRenderer(...))`"""
+
+    def __init__(self, content: str, updated: str, filename: str):
+        # total lines in each version
+        self._old_total = len(content.splitlines())
+        self._new_total = len(updated.splitlines())
+
+        # number of digits in the largest count
+        self._num_width = len(str(max(self._old_total, self._new_total))) + 2
+
+        # get the diff
+        self._diff_lines = list(
+            difflib.unified_diff(
+                content.splitlines(),
+                updated.splitlines(),
+                fromfile=f"{filename} (original)",
+                tofile=f"{filename} (modified)",
+                lineterm="",
+            )
+        )
+
+        # get syntax style
+        try:
+            self._lexer_name = Syntax.guess_lexer(filename, updated)
+        except Exception:
+            self._lexer_name = "text"
+
+    def __rich_console__(
+        self, console: Console, opts: ConsoleOptions
+    ) -> RenderResult:
+        old_line = new_line = None
+        width = console.width
+
+        for raw in self._diff_lines:
+            # grab line numbers from hunk header
+            if m := _HUNK_RE.match(raw):
+                old_line, new_line = map(int, m.groups())
+                # build a marker
+                n = self._num_width
+                tick_col = "." * (n - 1)
+                indent_ticks = f" {tick_col} {tick_col}"
+                # pad to the indent width
+                full_indent = indent_ticks.ljust(2 * n + 3)
+                yield Text(
+                    f"{full_indent}{raw}".ljust(width), style="white on grey30"
+                )
+                continue
+
+            # skip header lines
+            if raw.startswith(("---", "+++")):
+                continue
+
+            # split the line
+            if raw.startswith("+"):
+                style = _STYLE["add"]
+                code = raw[1:]
+            elif raw.startswith("-"):
+                style = _STYLE["del"]
+                code = raw[1:]
+            else:
+                style = _STYLE["ctx"]
+                code = raw[1:] if raw.startswith(" ") else raw
+
+            # compute line numbers
+            if raw.startswith("+"):
+                old_num, new_num = None, new_line
+                new_line += 1
+            elif raw.startswith("-"):
+                old_num, new_num = old_line, None
+                old_line += 1
+            else:
+                old_num, new_num = old_line, new_line
+                old_line += 1
+                new_line += 1
+
+            old_str = str(old_num) if old_num is not None else " "
+            new_str = str(new_num) if new_num is not None else " "
+
+            # Syntax-highlight the code part
+            syntax = Syntax(
+                code, self._lexer_name, line_numbers=False, word_wrap=False
+            )
+            text_code: Text = syntax.highlight(code)
+            if text_code.plain.endswith("\n"):
+                text_code = text_code[:-1]
+            # apply background
+            text_code.stylize(style.bg)
+
+            # line numbers + code
+            nums = Text(
+                f"{old_str:>{self._num_width}}{new_str:>{self._num_width}} ",
+                style=f"white {style.bg}",
+            )
+            diff_mark = Text(style.prefix, style=f"bright_white {style.bg}")
+            line_text = nums + diff_mark + text_code
+
+            # pad to console width
+            pad_len = width - line_text.cell_len
+            if pad_len > 0:
+                line_text.append(" " * pad_len, style=style.bg)
+
+            yield line_text

ursa/util/helperFunctions.py

@@ -0,0 +1,142 @@
+from __future__ import annotations
+
+import json
+import uuid
+from typing import Any, Callable, Dict, Iterable, List, Union
+
+from langchain_core.messages import AIMessage, BaseMessage, ToolMessage
+from langchain_core.runnables import Runnable
+
+
+# --- if you already have your own versions, reuse them ---
+def _parse_args(v: Any) -> Dict[str, Any]:
+    if v is None:
+        return {}
+    if isinstance(v, dict):
+        return v
+    if isinstance(v, str):
+        try:
+            return json.loads(v)
+        except Exception:
+            return {"_raw": v}
+    return {"_raw": v}
+
+
+def extract_tool_calls(msg: AIMessage) -> List[Dict[str, Any]]:
+    # Prefer normalized field
+    if msg.tool_calls:
+        out = []
+        for tc in msg.tool_calls:
+            name = getattr(tc, "name", None) or tc.get("name")
+            args = getattr(tc, "args", None) or tc.get("args")
+            call_id = getattr(tc, "id", None) or tc.get("id")
+            out.append({"name": name, "args": _parse_args(args), "id": call_id})
+        return out
+
+    # Fallbacks (OpenAI raw payloads)
+    ak = msg.additional_kwargs or {}
+    if ak.get("tool_calls"):
+        out = []
+        for tc in ak["tool_calls"]:
+            fn = tc.get("function", {}) or {}
+            out.append({
+                "name": fn.get("name"),
+                "args": _parse_args(fn.get("arguments")),
+                "id": tc.get("id"),
+            })
+        return out
+
+    if ak.get("function_call"):
+        fn = ak["function_call"]
+        return [
+            {
+                "name": fn.get("name"),
+                "args": _parse_args(fn.get("arguments")),
+                "id": None,
+            }
+        ]
+    return []
+
+
+# -----------------------------------------------------------------------------
+
+
+ToolRegistry = Dict[str, Union[Runnable, Callable[..., Any]]]
+
+
+def _stringify_output(x: Any) -> str:
+    if isinstance(x, str):
+        return x
+    try:
+        return json.dumps(x, ensure_ascii=False)
+    except Exception:
+        return str(x)
+
+
+def _invoke_tool(
+    tool: Union[Runnable, Callable[..., Any]], args: Dict[str, Any]
+) -> Any:
+    # Runnable (LangChain tools & chains)
+    if isinstance(tool, Runnable):
+        return tool.invoke(args)
+    # Plain callable
+    try:
+        return tool(**args)
+    except TypeError:
+        # Some tools expect a single positional payload
+        return tool(args)
+
+
+def run_tool_calls(
+    ai_msg: AIMessage,
+    tools: Union[ToolRegistry, Iterable[Union[Runnable, Callable[..., Any]]]],
+) -> List[BaseMessage]:
+    """
+    Args:
+        ai_msg: The LLM's AIMessage containing tool calls.
+        tools: Either a dict {name: tool} or an iterable of tools (must have `.name`
+               for mapping). Each tool can be a Runnable or a plain callable.
+
+    Returns:
+        out: list[BaseMessage] to feed back to the model
+    """
+    # Build a name->tool map
+    if isinstance(tools, dict):
+        registry: ToolRegistry = tools  # type: ignore
+    else:
+        registry = {}
+        for t in tools:
+            name = getattr(t, "name", None) or getattr(t, "__name__", None)
+            if not name:
+                raise ValueError(f"Tool {t!r} has no discoverable name.")
+            registry[name] = t  # type: ignore
+
+    calls = extract_tool_calls(ai_msg)
+
+    if not calls:
+        return []
+
+    out: List[BaseMessage] = []
+    for call in calls:
+        name = call.get("name")
+        args = call.get("args", {}) or {}
+        call_id = call.get("id") or f"call_{uuid.uuid4().hex}"
+
+        # 1) the AIMessage that generated the call
+        out.append(ai_msg)
+
+        # 2) the ToolMessage with the execution result (or error)
+        if name not in registry:
+            content = f"ERROR: unknown tool '{name}'."
+        else:
+            try:
+                result = _invoke_tool(registry[name], args)
+                content = _stringify_output(result)
+            except Exception as e:
+                content = f"ERROR: {type(e).__name__}: {e}"
+
+        out.append(
+            ToolMessage(content=content, tool_call_id=call_id, name=name)
+        )
+
+    return out