initial commit

1 files changed, 183 insertions, 0 deletions

diff --git a/llm.py b/llm.py
new file mode 100755
index 0000000..f43fe2e
--- /dev/null
+++ b/llm.py
@@ -0,0 +1,183 @@
+#!/bin/python3
+import os
+from pathlib import Path
+from openai import OpenAI
+from typing import Optional
+import argparse
+import subprocess
+import readline
+import json
+
+import logging
+logger = logging.getLogger(__name__)
+logging.basicConfig(filename="thoughts.log",
+                    level=logging.INFO)
+
+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 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},)
+    
+    response = query_method(state)
+
+    state = state + ({"role": "assistant", "content": response},)
+    
+    return (response, state)
+
+def strip_enclosing(s: str, start="```json", end="```") -> str:
+    s = s.strip()
+    return s[len(start):-len(end)] if s[:len(start)] == start and s[-len(end):] == end else s
+
+
+def path_checked(path: Path, cwd = Path.cwd()) -> Optional[Path]:
+    """Resolves `path`, returning `None` if it's not a subdirectory of `cwd`."""
+    path = path.resolve()
+    cwd = cwd.resolve()
+    return path if path.is_relative_to(cwd) else None
+
+
+def inner_query_agent(actions: list, query: str, state = ()):
+
+    state = state + ({"role": "user", "content": query},)
+
+    (res, state) = chat(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)
+
+    try:
+        action = json.loads(strip_enclosing(res))
+    except json.JSONDecodeError:
+        return (res, state)
+
+    for f in actions:
+        result_json = f(action)
+        if result_json is not None:
+            (res, state) = inner_query_agent(actions, json.dumps(result_json), state)
+            return (res, state)
+
+    logger.info(f"valid json '{res}' but no actions matched")
+
+    return (res, state)
+
+
+def query_agent(query: str, state = ()):
+
+    def path_helper(path_str: str, callback):
+        path =  path_checked(Path(path_str))
+        if path is None:
+            return {"status": "error", "reason": f"Permission denied: {path_str}"}
+        else:
+            try:
+                return {"status": "success", **callback(path)}
+            except Exception as e:
+                return {"status": "error", "reason": f"Failure: {e}"}
+
+    def ls(args):
+        '{"action": "ls", "path": PATH} - list files at PATH'
+        def inner(path: Path):
+            return {"content":  list(map(lambda s: str(s), path.iterdir()))}
+        match args:
+            case {"action": "ls", "path": path_str}:
+                return path_helper(path_str, inner)
+
+    def read(args):
+        '{"action": "read", "path": PATH} - read file at PATH'
+        def inner(path: Path):
+            contents = path.read_text(encoding="utf-8")
+            return {"content": contents}
+        match args:
+            case {"action": "read", "path": path_str}:
+                return path_helper(path_str, inner)
+
+    def write(args):
+        '{"action": "write", "path": PATH, "content": CONTENT} - write CONTENT into file at PATH'
+        def inner(path: Path, text: str):
+            n_bytes = path.write_text(text)
+            return {"bytes": n_bytes}
+        match args:
+            case {"action": "write", "path": path_str, "content": content}:
+                return path_helper(path_str, lambda path: inner(path, content))
+
+    
+    def tldr(args):
+        '{"action": "tldr", "program": PROGRAM} - read tldr docs for PROGRAM'
+        match args:
+            case {"action": "tldr", "program": program_str}:
+                try:
+                    # Run man command and capture output
+                    # -P cat is used to avoid paging and just print the text
+                    result = subprocess.run(
+                        ["tldr", program_str],
+                        capture_output=True,
+                        text=True,
+                        check=True
+                    )
+                    return {"status": "success", "content": result.stdout}
+                except subprocess.CalledProcessError:
+                    return {"status": "error", "reason": f"Tldr for '{program_str}' not found."}
+                except Exception as e:
+                    return {"status": "error", "reason": f"An unexpected error occurred: {e}"}
+
+    actions = [ls, read, write, tldr]
+
+    state = state + ({"role": "system",
+                      "content": f"""You are an autonomous coding agent that can help with programming tasks.
+You can explore repositories, read files, search for information, and make changes to code.
+
+You have access to the following actions.
+{'\n'.join(map(lambda f: '* ' + f.__doc__, actions))}
+* RESPONSE - send non-json text RESPONSE to the user once the task is complete
+"""},)
+
+    state = state + ({"role": "system", "content": "Analyze the following user-given task and break it down into steps."},)
+
+    return inner_query_agent(actions, query, state)
+    
+
+def main():
+
+    state = ()
+
+    while True:
+        
+        query = input("> ").strip()
+
+        (response, state) = query_agent(query, state)
+        
+        print(response)
+        
+
+if __name__ == "__main__":
+    main()
+