feat: socks5 proxy

1 files changed, 28 insertions, 8 deletions

diff --git a/llm.py b/llm.py
index ca61b48..9abfe12 100755
--- a/llm.py
+++ b/llm.py
@@ -1,7 +1,8 @@
 #!/bin/python3
 import os
 from pathlib import Path
-from openai import OpenAI
+import httpx2
+from openai import OpenAI, DefaultHttpx2Client
 from openai.types.responses import EasyInputMessageParam as Message
 from typing import Optional, List, Dict, TypeVar, TypeAlias, Callable, Tuple, Literal
 import argparse
@@ -39,13 +40,14 @@ class AgentConfig:
     trustlevels: Dict[str, int]
     trustlevel: int
     query_method: QueryMethod
+    socks5_proxy: Optional[str] = None
     search_padding: int = 200
 
 
 def estimate_tokens(state: DiscussionState) -> int:
     """Simple heuristic to estimate tokens in the conversation state."""
     text = "".join(str(m["content"]) for m in state)
-    return len(text) // 3.5
+    return int(len(text) / 3.5)
 
 
 def mk_direct_query(
@@ -187,13 +189,24 @@ def main():
         default=2,
         help="The trust level of the agent (default: 2)",
     )
+    parser.add_argument(
+        "--socks5-proxy",
+        type=str,
+        default=os.environ.get("LLM_SOCKS5_PROXY"),
+        help=(
+            "Route API queries (and webgrep fetches) through a SOCKS5 proxy, e.g. "
+            "socks5://host:1080. Use socks5h:// to resolve DNS through the proxy; "
+            "user:pass@ credentials in the URL are supported. "
+            "(default: LLM_SOCKS5_PROXY environment variable)"
+        ),
+    )
     args = parser.parse_args()
 
-    client = OpenAI(base_url=args.endpoint, api_key=args.key)
+    http_client: Optional[httpx2.Client] = (
+        DefaultHttpx2Client(proxy=args.socks5_proxy) if args.socks5_proxy else None
+    )
 
-    # We need a version of direct_query that doesn't trigger squashing to use inside the squasher
-    # to avoid infinite recursion.
-    raw_direct_query = mk_direct_query(client, args.model, squasher=None)
+    client = OpenAI(base_url=args.endpoint, api_key=args.key, http_client=http_client)
 
     def squash_history(state: DiscussionState) -> DiscussionState:
         if len(state) <= 1 + args.min_history:
@@ -203,7 +216,7 @@ def main():
                 state[-1],
             )
         else:
-            (summary, _) = raw_direct_query(
+            (summary, _) = mk_direct_query(client, args.model, squasher=None)(
                 "Please summarize the preceding conversation history concisely, preserving all key technical decisions, progress made, and current goals.",
                 state[1 : -args.min_history],
                 role="system",
@@ -233,6 +246,7 @@ def main():
             squasher=squash_history,
             max_token_count=args.model_context_size,
         ),
+        socks5_proxy=args.socks5_proxy,
     )
 
     state = (
@@ -252,8 +266,14 @@ Once you are finished with your task, give a response with none of the above JSO
     )
 
     while True:
-        query = input("> ").strip()
+        try:
+            query = input("> ").strip()
+        except KeyboardInterrupt:
+            break
 
+        # Don't include this in the try above st. if the user
+        # interrupts the agent, it's possible to find out what the
+        # program was up to at the moment.
         (response, state) = query_agent(query, config, state)
 
         print(response)