feat: webgrep

3 files changed, 120 insertions, 38 deletions

diff --git a/actions.py b/actions.py
index 589dd51..a6051ed 100644
--- a/actions.py
+++ b/actions.py
@@ -2,7 +2,9 @@ from pathlib import Path
 from typing import Optional, Callable, Dict
 from dataclasses import dataclass
 import subprocess
+import re
 from llm import AgentConfig, JSON
+from web_utils import fetch_and_markdownify, FetchSuccess, FetchError
 
 
 @dataclass
@@ -108,16 +110,42 @@ def _shell(args: JSON, config: AgentConfig) -> JSON:
 
 def _mypy(args: JSON, config: AgentConfig) -> JSON:
     '{"action": "mypy", "path": PATH} - run mypy type checking on PATH'
-    try:
-        res = subprocess.run(
-            ["mypy", args["path"]],
-            cwd=config.workdir,
-            capture_output=True,
-            text=True,
-        )
-        return {"status": "success", "stdout": res.stdout, "stderr": res.stderr}
-    except Exception as e:
-        return {"status": "error", "reason": f"An unexpected error occurred: {e}"}
+
+    def inner(path: Path):
+        try:
+            res = subprocess.run(
+                ["mypy", str(path)], cwd=config.workdir, capture_output=True, text=True
+            )
+            if res.returncode == 0:
+                return {"status": "success", "stdout": res.stdout}
+            else:
+                return {"status": "error", "stderr": res.stderr}
+        except Exception as e:
+            return {"status": "error", "reason": f"An unexpected error occurred: {e}"}
+
+    return _path_helper(args["path"], inner, cwd=config.workdir)
+
+
+def _webgrep(args: JSON, config: AgentConfig) -> JSON:
+    '{"action": "webgrep", "url": URL, "keywords": KEYWORD_LIST} - find case-insensitive matches in markdownified content of URL for keywords in KEYWORD_LIST plus some context on both sides'
+
+    match fetch_and_markdownify(args["url"]):
+        case FetchError(e):
+            return {"status": "error", "reason": e}
+        case FetchSuccess(content):
+            matches = map(
+                lambda m: content[
+                    max(0, m.start() - config.search_padding) : min(
+                        len(content), m.end() + config.search_padding
+                    )
+                ],
+                re.finditer(
+                    "|".join(map(lambda kw: re.escape(kw), args["keywords"])),
+                    content,
+                    re.IGNORECASE,
+                ),
+            )
+            return {"status": "success", "matches": list(matches)}
 
 
 ls = Action("ls", _ls.__doc__ or "", _ls)
@@ -135,4 +163,15 @@ shell = Action(
     _shell,
     formatter=lambda args: f"shell command `{args['command']}`",
 )
-mypy = Action("mypy", _mypy.__doc__ or "", _mypy)
+mypy = Action(
+    "mypy",
+    _mypy.__doc__ or "",
+    _mypy,
+    formatter=lambda args: f"type check {args['path']}",
+)
+webgrep = Action(
+    "webgrep",
+    _webgrep.__doc__ or "",
+    _webgrep,
+    formatter=lambda args: f"search {args['url']}",
+)
diff --git a/llm.py b/llm.py
index bea4536..9b6208d 100755
--- a/llm.py
+++ b/llm.py
@@ -2,7 +2,7 @@
 import os
 from pathlib import Path
 from openai import OpenAI
-from typing import Optional, List, Dict, TypeVar, TypeAlias
+from typing import Optional, List, Dict, TypeVar, TypeAlias, Callable
 import argparse
 import re
 import subprocess
@@ -29,6 +29,7 @@ class AgentConfig:
     actions: List["Action"]
     trustlevels: Dict[str, int]
     trustlevel: int
+    search_padding: int = 200
 
 
 client = OpenAI(
@@ -46,28 +47,6 @@ def direct_query(query: str):
     return res.output_text
 
 
-# def reason_query(input):
-#     # TODO some models may natively support reasoning
-
-#     reasoning = direct_query(input + ({"role": "system",
-#                                        "content": "Reason about the user's query step by step.\n\nAn answer will be generated based on your analysis."},))
-
-#     summary = direct_query(input + ({"role": "assistant", "content": reasoning},
-#                                     {"role": "system",
-#                                      "content": "Formulate your previous reasoning as a concise answer to the user."}))
-
-#     return summary
-
-# # kinda funky
-# # def flexible_query(input):
-
-# #     reasoning_required = direct_query(input[-1:] + ({"role": "system", "content": "Does answering the user's question require reasoning? (y|N)"},))
-# #     if reasoning_required == 'y':
-# #         return reason_query(input)
-# #     else:
-# #         return direct_query(input)
-
-
 def chat(content, state=(), query_method=direct_query, role="user"):
 
     state = state + ({"role": role, "content": content},)
@@ -144,16 +123,17 @@ def main():
 
     config = AgentConfig(
         workdir=args.workdir,
-        actions=[ls, read, write, tldr, shell, mypy],
+        actions=[ls, read, write, tldr, shell, mypy, webgrep],
         trustlevels={
             "ls": 0,
             "read": 0,
             "tldr": 0,
             "mypy": 0,
-            "write": 1,
-            "shell": 2,
+            "webgrep": 1,
+            "write": 2,
+            "shell": 3,
         },
-        trustlevel=1,
+        trustlevel=2,
     )
 
     state = (
@@ -172,6 +152,8 @@ Once you are finished with your task, give a response with none of the above JSO
         },
     )
 
+    print(state)
+
     while True:
         query = input("> ").strip()
 
diff --git a/web_utils.py b/web_utils.py
new file mode 100644
index 0000000..91561b5
--- /dev/null
+++ b/web_utils.py
@@ -0,0 +1,61 @@
+import urllib.request
+import subprocess
+import re
+from typing import Dict, Union
+from dataclasses import dataclass
+
+
+@dataclass
+class FetchSuccess:
+    """Successful fetch result containing the markdown content."""
+
+    content: str
+
+
+@dataclass
+class FetchError:
+    """Failed fetch result containing the error message."""
+
+    error: str
+
+
+FetchResult = Union[FetchSuccess, FetchError]
+
+# Memoization cache for web pages
+_cache: Dict[str, FetchResult] = {}
+
+
+def html_to_markdown(html: str) -> str:
+    """Converts HTML to markdown using pandoc with GFM format."""
+    try:
+        # Use gfm (GitHub Flavored Markdown) for better compatibility and robust conversion
+        process = subprocess.run(
+            ["pandoc", "-f", "html", "-t", "gfm", "-o", "-"],
+            input=html,
+            capture_output=True,
+            text=True,
+            check=True,
+            encoding="utf-8",
+        )
+        return process.stdout.strip()
+    except subprocess.CalledProcessError as e:
+        raise RuntimeError(f"Pandoc error: {e}")
+    except Exception as e:
+        raise RuntimeError(f"Unexpected error during markdown conversion: {e}")
+
+
+def fetch_and_markdownify(url: str) -> FetchResult:
+    """Fetches a web page and transforms it to markdown, with memoization."""
+    if url in _cache:
+        return _cache[url]
+
+    try:
+        req = urllib.request.Request(url, headers={"User-Agent": "Mozilla/5.0"})
+        with urllib.request.urlopen(req, timeout=10) as response:
+            html = response.read().decode("utf-8")
+            markdown = html_to_markdown(html)
+            res = FetchSuccess(content=markdown)
+            _cache[url] = res
+            return res
+    except Exception as e:
+        return FetchError(error=str(e))