tuff

harness.py

@@ -7,15 +7,21 @@ import uuid
 from dataclasses import dataclass
 from datetime import UTC, datetime
 from pathlib import Path
-from typing import Any
+from typing import TYPE_CHECKING, Any, cast
 
 from openai import OpenAI
 
+if TYPE_CHECKING:
+    from openai.types.chat import (
+        ChatCompletionMessageFunctionToolCall,
+        ChatCompletionMessageParam,
+        ChatCompletionToolUnionParam,
+    )
 DEFAULT_PROMPT_PATH = Path("prompt.md")
 DEFAULT_SESSION_PATH = Path("session/default.jsonl")
-DEFAULT_MODEL_ENV = "GROQ_MODEL_STRING"
-DEFAULT_BASE_URL_ENV = "GROQ_BASE_URL"
-DEFAULT_API_KEY_ENV = "GROQ_API_KEY"
+MODEL_ENV_NAMES = ("CEREBRAS_MODEL_STRING", "GROQ_MODEL_STRING")
+BASE_URL_ENV_NAMES = ("CEREBRAS_BASE_URL", "GROQ_BASE_URL")
+API_KEY_ENV_NAMES = ("CEREBRAS_API_KEY", "GROQ_API_KEY")
 
 
 def utc_now_iso() -> str:
@@ -78,7 +84,7 @@ def append_message(path: Path, role: str, content: str, **extra: Any) -> None:
     append_jsonl(path, payload)
 
 
-def shell_tool_schema() -> list[dict[str, Any]]:
+def shell_tool_schema() -> "list[ChatCompletionToolUnionParam]":
     return [
         {
             "type": "function",
@@ -103,12 +109,20 @@ def shell_tool_schema() -> list[dict[str, Any]]:
 
 def build_client() -> OpenAI:
     load_dotenv()
-    base_url = os.getenv(DEFAULT_BASE_URL_ENV)
-    api_key = os.getenv(DEFAULT_API_KEY_ENV)
+    base_url = next(
+        (os.getenv(name) for name in BASE_URL_ENV_NAMES if os.getenv(name)), None
+    )
+    api_key = next(
+        (os.getenv(name) for name in API_KEY_ENV_NAMES if os.getenv(name)), None
+    )
     if not base_url:
-        raise RuntimeError(f"Missing {DEFAULT_BASE_URL_ENV} in environment or .env")
+        raise RuntimeError(
+            f"Missing one of {', '.join(BASE_URL_ENV_NAMES)} in environment or .env"
+        )
     if not api_key:
-        raise RuntimeError(f"Missing {DEFAULT_API_KEY_ENV} in environment or .env")
+        raise RuntimeError(
+            f"Missing one of {', '.join(API_KEY_ENV_NAMES)} in environment or .env"
+        )
     return OpenAI(base_url=base_url, api_key=api_key)
 
 
@@ -147,7 +161,9 @@ def load_prompt(prompt_path: Path) -> str:
     return prompt_path.read_text(encoding="utf-8").strip()
 
 
-def messages_from_session(session_path: Path, prompt_text: str) -> list[dict[str, Any]]:
+def messages_from_session(
+    session_path: Path, prompt_text: str
+) -> list[ChatCompletionMessageParam]:
     messages: list[dict[str, Any]] = []
     if prompt_text:
         messages.append({"role": "system", "content": prompt_text})
@@ -171,7 +187,7 @@ def messages_from_session(session_path: Path, prompt_text: str) -> list[dict[str
                 }
             )
 
-    return messages
+    return cast(list[ChatCompletionMessageParam], messages)
 
 
 def tool_result_content(payload: dict[str, Any]) -> str:
@@ -195,9 +211,14 @@ class TuffHarness:
         self.session_path = session_path
         self.prompt_path = prompt_path
         self.client = build_client()
-        self.model_name = os.getenv(DEFAULT_MODEL_ENV)
-        if not self.model_name:
-            raise RuntimeError(f"Missing {DEFAULT_MODEL_ENV} in environment or .env")
+        model_name = next(
+            (os.getenv(name) for name in MODEL_ENV_NAMES if os.getenv(name)), None
+        )
+        if not model_name:
+            raise RuntimeError(
+                f"Missing one of {', '.join(MODEL_ENV_NAMES)} in environment or .env"
+            )
+        self.model_name = model_name
 
     def run_turn(self, max_steps: int = 8) -> TurnResult:
         prompt_text = load_prompt(self.prompt_path)
@@ -236,7 +257,10 @@ class TuffHarness:
 
             assistant = response.choices[0].message
             assistant_content = assistant.content or ""
-            assistant_tool_calls = list(assistant.tool_calls or [])
+            assistant_tool_calls = cast(
+                "list[ChatCompletionMessageFunctionToolCall]",
+                list(assistant.tool_calls or []),
+            )
 
             append_message(
                 self.session_path,
@@ -256,6 +280,8 @@ class TuffHarness:
             )
 
             messages.append(
+                cast(
+                    "ChatCompletionMessageParam",
                     {
                         "role": "assistant",
                         "content": assistant_content,
@@ -270,7 +296,8 @@ class TuffHarness:
                             }
                             for tool_call in assistant_tool_calls
                         ],
-                }
+                    },
+                )
             )
 
             if not assistant_tool_calls:

main.py

@@ -1,10 +1,11 @@
 from __future__ import annotations
 
 import json
+import os
 import time
 from argparse import ArgumentParser
 from pathlib import Path
-from typing import Callable
+from typing import TYPE_CHECKING, Callable
 
 from watchdog.events import FileSystemEvent, FileSystemEventHandler
 from watchdog.observers import Observer
@@ -17,6 +18,9 @@ from harness import (
     read_jsonl,
 )
 
+if TYPE_CHECKING:
+    from watchdog.observers.api import BaseObserver
+
 
 def build_cli() -> ArgumentParser:
     cli = ArgumentParser(description="tuff harness")
@@ -133,7 +137,7 @@ class SessionHandler(FileSystemEventHandler):
     def on_any_event(self, event: FileSystemEvent) -> None:
         if event.is_directory:
             return
-        src = Path(event.src_path).resolve()
+        src = Path(os.fsdecode(event.src_path)).resolve()
         if src in self.target_paths:
             self.on_change()
 
@@ -157,7 +161,7 @@ def chat_loop(
     print("chat mode on. type and press enter. /exit to quit, /run to force a turn.")
     while True:
         try:
-            raw = input("you> ")
+            raw = input("USER: ")
         except EOFError:
             break
         except KeyboardInterrupt:
@@ -200,7 +204,7 @@ def main() -> None:
     # Print any existing transcript first.
     print_new_results(session_path, seen_ids)
 
-    observer: Observer | None = None
+    observer: "BaseObserver | None" = None
     if args.watch:
         observer = Observer()
         handler = SessionHandler(