fix: better configuration and typing

1 files changed, 57 insertions, 24 deletions

diff --git a/llm.py b/llm.py
index 9b6208d..db72ba3 100755
--- a/llm.py
+++ b/llm.py
@@ -2,7 +2,8 @@
 import os
 from pathlib import Path
 from openai import OpenAI
-from typing import Optional, List, Dict, TypeVar, TypeAlias, Callable
+from openai.types.responses import EasyInputMessageParam as Message
+from typing import Optional, List, Dict, TypeVar, TypeAlias, Callable, Tuple, Literal
 import argparse
 import re
 import subprocess
@@ -12,6 +13,8 @@ from dataclasses import dataclass
 from actions import *
 
 T = TypeVar("T")
+A = TypeVar("A")
+B = TypeVar("B")
 
 import logging
 
@@ -20,6 +23,12 @@ logging.basicConfig(filename="thoughts.log", level=logging.INFO)
 
 JSON: TypeAlias = Dict[str, "JSON"] | List["JSON"] | str | int | float | bool | None
 
+DiscussionState: TypeAlias = Tuple[Message, ...]
+DiscussionM: TypeAlias = Tuple[str, DiscussionState]
+QueryMethod: TypeAlias = Callable[
+    [str, DiscussionState], DiscussionM
+]  # DiscussionM -> DiscussionM if multiple arguments are interpreted as product types
+
 
 @dataclass
 class AgentConfig:
@@ -29,33 +38,30 @@ class AgentConfig:
     actions: List["Action"]
     trustlevels: Dict[str, int]
     trustlevel: int
+    query_method: QueryMethod
     search_padding: int = 200
 
 
-client = OpenAI(
-    api_key=os.environ["AALTO_LLM_KEY"],
-    base_url="https://llm-gateway.k8s.aalto.fi/api/v1",
-)
-
-
-def direct_query(query: str):
-    logger.info(f"send: {query}")
-    res = client.responses.create(
-        model="RedHatAI/gemma-4-31B-it-FP8-Dynamic", input=query
-    )
-    logger.info(f"receive: {res.output_text}")
-    return res.output_text
+def mk_direct_query(client: OpenAI, model_name: str) -> QueryMethod:
+    def direct_query(
+        query: str,
+        state: DiscussionState,
+        role: Literal["user", "system", "developer"] = "user",
+    ) -> DiscussionM:
+        logger.info(f"send: {query}")
 
+        state = state + ({"role": role, "content": query},)
 
-def chat(content, state=(), query_method=direct_query, role="user"):
+        res = client.responses.create(model=model_name, input=list(state))
 
-    state = state + ({"role": role, "content": content},)
+        logger.info(f"receive: {res.output_text}")
 
-    response = query_method(state)
+        return (
+            res.output_text,
+            state + ({"role": "assistant", "content": res.output_text},),
+        )
 
-    state = state + ({"role": "assistant", "content": response},)
-
-    return (response, state)
+    return direct_query
 
 
 def first(l: List[T], p: Callable[[T], bool]) -> Optional[T]:
@@ -66,7 +72,7 @@ def query_agent(query: str, config: AgentConfig, state=()):
 
     state = state + ({"role": "user", "content": query},)
 
-    (res, state) = chat(
+    (res, state) = config.query_method(
         f"Based on the current state, what should you do next?\n"
         "Respond JSON indicating an action, or respond in natural language if you are finished.",
         state,
@@ -119,6 +125,32 @@ def main():
         default=Path.cwd(),
         help="The directory the agent has access to (default: current directory)",
     )
+    parser.add_argument(
+        "--endpoint",
+        type=str,
+        default=os.environ.get(
+            "LLM_API_ENDPOINT", "https://llm-gateway.k8s.aalto.fi/api/v1"
+        ),
+        help="The API endpoint for the LLM",
+    )
+    parser.add_argument(
+        "--key",
+        type=str,
+        default=os.environ.get("LLM_API_KEY"),
+        help="The key for the LLM API endpoint",
+    )
+    parser.add_argument(
+        "--model",
+        type=str,
+        default="RedHatAI/gemma-4-31B-it-FP8-Dynamic",
+        help="The model name to use",
+    )
+    parser.add_argument(
+        "--trust-level",
+        type=int,
+        default=2,
+        help="The trust level of the agent (default: 2)",
+    )
     args = parser.parse_args()
 
     config = AgentConfig(
@@ -133,7 +165,10 @@ def main():
             "write": 2,
             "shell": 3,
         },
-        trustlevel=2,
+        trustlevel=args.trust_level,
+        query_method=mk_direct_query(
+            OpenAI(base_url=args.endpoint, api_key=args.key), args.model
+        ),
     )
 
     state = (
@@ -152,8 +187,6 @@ Once you are finished with your task, give a response with none of the above JSO
         },
     )
 
-    print(state)
-
     while True:
         query = input("> ").strip()