diff options
Diffstat (limited to 'llm.py')
| -rwxr-xr-x | llm.py | 183 |
1 files changed, 183 insertions, 0 deletions
@@ -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() + |