snapshot of https://github.com/ActivityWatch/aw-client.git (master, 2026-09-16; blobs > 131072B omitted)

examples/suggest_categories_gpt.py

@@ -0,0 +1,213 @@
+"""
+Uses GPT-3 to suggest new categories.
+
+Builds on suggest_categories.py for basic operations, like getting events.
+"""
+
+from datetime import datetime, timezone
+from typing import Any, Dict, List, Tuple
+
+from aw_core import Event
+from aw_transform.classify import Rule, categorize
+
+from suggest_categories import example_categories, get_events
+
+Category = Tuple[List[str], Dict[str, Any]]
+
+
+def prompt_preamble(categories: List[Category]) -> str:
+    categories_str = "\n\n".join(
+        [
+            f" - Category: {' > '.join(name)}\n   Regex: {rule['regex']}"
+            for name, rule in categories
+        ]
+    )
+
+    prompt = f"""We will classify window titles into user-defined categories defined by regular expressions.
+If a suitable one doesn't exists, we will create one with a suitable regex.
+
+Existing categories:
+
+{categories_str}
+
+---
+
+What category should "ActivityWatch - wwww.github.com" be in?
+Category: Work > ActivityWatch
+
+What category should "reddit: the front page of the internet" be in?
+New Category: Media > Social > Reddit
+Regex: reddit
+
+What category should "Twitter" be in?
+New Category: Media > Social > Twitter
+Regex: Twitter
+
+What category should "Demis Hassabis: DeepMind - AI, Superintelligence & the Future of Humanity | Lex Fridman Podcast - YouTube - Mozilla Firefox" be in?
+New Category: Media > Video > YouTube
+Regex: YouTube
+
+What category should "Tweetdeck" be in?
+Modify Category: Media > Social > Twitter
+Append Regex: Tweetdeck
+
+What category should "cloudflare-ipfs.com | 524: A timeout occurred - cloudflare-ipfs.com - Mozilla Firefox" be in?
+Skip: No suitable category found or to suggest, best left as uncategorized.
+
+What category should "Mozilla Firefox" be in?
+Skip: No suitable category found or to suggest, best left as uncategorized.
+
+What category should "RimWorld" be in?
+New Category: Games > RimWorld
+Regex: RimWorld
+
+What category should "Minecraft" be in?
+New Category: Games > Minecraft
+Regex: Minecraft
+
+What category should "Free Porn Videos & Sex Movies - Porno, XXX, Porn Tube | Pornhub — Mozilla Firefox" be in?
+New Category: Media > Porn
+Regex: Pornhub"""
+
+    return prompt
+
+
+def process_prompt(prompt, categories, quiet=False) -> List[Category]:
+    """processes the prompt preamble for categories created/modified in the prompt"""
+    for entry in prompt.split("---", 1)[1].split("\n\n"):
+        if not entry.strip():
+            continue
+        # FIXME: Will break if string contains double-quotes
+        title = entry.split("\n", 1)[0].split('"', 2)[1]
+        response = entry.strip().split("\n", 1)[1]
+        categories = parse_gpt_response(response, categories, title=title, quiet=quiet)
+    return categories
+
+
+def gpt_suggest(event: Event, categories: List[Category]) -> List[Category]:
+    """
+    Use OpenAI GPT-3 to suggest new categories.
+
+    Takes an uncategorized event and the current categories, and returns a list of suggested categories.
+    """
+    import os
+    import openai
+    from copy import deepcopy
+
+    openai.api_key = os.getenv("OPENAI_API_KEY")
+
+    title = event.data["title"]
+    categories = deepcopy(categories)
+
+    prompt = prompt_preamble(categories)
+    categories = process_prompt(prompt, categories, quiet=True)
+
+    prompt = f"""{prompt}
+
+What category should "{title}" be in?"""
+
+    response = openai.Completion.create(
+        model="text-davinci-002",
+        prompt=prompt,
+        temperature=0,
+        max_tokens=64,
+        top_p=1.0,
+        frequency_penalty=0.0,
+        presence_penalty=0.0,
+    )
+    text = response["choices"][0]["text"]
+    print("-" * 80)
+    # print(prompt + text)
+    # print("> " + text.strip().replace("\n", "\n> "))
+
+    return parse_gpt_response(text, categories, title)
+
+
+def check_is_category(text: str, category: Category) -> bool:
+    name, rule = category
+    event = categorize(
+        [Event(timestamp=datetime.now(tz=timezone.utc), data={"title": text})],
+        [(list(name), Rule(rule))],
+    )[0]
+    return event.data["$category"] == list(category[0])
+
+
+def parse_gpt_response(text: str, categories: List[Category], title=None, quiet=False):
+    category_names = [tuple(name) for name, _ in categories]
+
+    line1, *lines = text.strip().split("\n")
+    if line1.startswith("Category:"):
+        # chose existing category
+        cat_name = tuple(line1.split(":", 1)[1].strip().split(" > "))
+        if not quiet:
+            print(f"Chose existing category {cat_name}  (title: {title})")
+        if cat_name not in [tuple(name[: len(cat_name)]) for name in category_names]:
+            print(f"No category named {cat_name} found, skipping")
+    elif line1.startswith("New Category:"):
+        line2 = lines[0]
+        category = line1.split(":", 1)[1].strip().split(" > ")
+        regex = line2.strip().split(":", 1)[1].strip()
+        if not quiet:
+            print(f"Added category {category} with regex {regex}  (title: {title})")
+        cat: Category = (
+            category,
+            {"type": "regex", "regex": regex, "ignore_case": True},
+        )
+        if title and not check_is_category(title, cat):
+            print(
+                f"Bad suggested regex '{cat[1]['regex']}'. Title '{title}' does not match category {category}."
+            )
+        else:
+            categories.append(cat)
+    elif line1.startswith("Modify Category:"):
+        line2 = lines[0]
+        category = line1.split(":", 1)[1].strip().split(" > ")
+        assert line2.startswith("Append Regex:")
+        regex = line2.strip().split(":", 1)[1].strip()
+        # get existing category
+        for name, rule in categories:
+            if name == category:
+                rule["regex"] += "|" + regex
+                if not quiet:
+                    print(f"Appended {regex} to {category}")
+                break
+    elif line1.startswith("Skip:"):
+        pass
+    else:
+        print(f"Unknown response: '{text.strip()}'")
+
+    return categories
+
+
+def main():
+    categories = example_categories()
+    events = get_events(categories)
+
+    events_by_dur = sorted(events, key=lambda e: e.duration, reverse=True)
+    for event in events_by_dur[:100]:
+        # re-categorize and skip event if it is already categorized by a new rule
+        event, *_ = categorize(
+            [event], [(list(name), Rule(rule)) for name, rule in categories]
+        )
+        if list(event.data["$category"]) != ["Uncategorized"]:
+            continue
+
+        categories = gpt_suggest(event, categories)
+
+
+def test_parse_gpt_response():
+    categories = example_categories()
+    prompt = prompt_preamble(categories)
+    categories = process_prompt(prompt, categories)
+
+    cat_twitter = list(
+        (name, rule)
+        for name, rule in categories
+        if tuple(name) == ("Media", "Social", "Twitter")
+    )
+    assert len(cat_twitter) == 1
+    assert cat_twitter[0][1]["regex"] == "Twitter|Tweetdeck"
+
+
+if __name__ == "__main__":
+    main()