diff options
Diffstat (limited to 'shitoutcode/__main__.py')
-rw-r--r-- | shitoutcode/__main__.py | 54 |
1 files changed, 54 insertions, 0 deletions
diff --git a/shitoutcode/__main__.py b/shitoutcode/__main__.py new file mode 100644 index 0000000..09ca715 --- /dev/null +++ b/shitoutcode/__main__.py @@ -0,0 +1,54 @@ +import os +import random +import sys +from shitoutcode import Env, LLMAPIException, SourceFile, gen_rand_srccode + +ARGV0 = "shitoutcode" + +def loadEnv_from_env () -> Env: + def wrap_default (val: str, t: type, dv = None): + if val is None: + return dv + return t(val) + + ret = Env() + ret.model = wrap_default(os.getenv("LLM_MODEL"), str, "gpt-4o-mini") + ret.output_tokens_min = wrap_default(os.getenv("OUT_TOKENS_MIN"), int, ret.output_tokens_min) + ret.output_tokens_max = wrap_default(os.getenv("OUT_TOKENS_MAX"), int, ret.output_tokens_max) + ret.temp = wrap_default(os.getenv("LLM_TEMP"), float) + ret.seed = wrap_default(os.getenv("LLM_SEED"), int) + ret.max_tokens = wrap_default(os.getenv("LLM_SEED"), int, ret.max_tokens) + + return ret + +def getPathLine (x: SourceFile) -> str: + if x.category: + return x.category + os.path.sep + x.name + return x.name + +def pickRandLang () -> str: + thelist = [ + "C", + "C++", + "Javascript", + "Java" + ] + r = random.randint(0, len(thelist) - 1) + + return thelist[r] + +env = loadEnv_from_env() +lang = sys.argv[1] if len(sys.argv) > 1 else pickRandLang() +extra_prompt = sys.argv[2] if len(sys.argv) > 2 else None + +try: + result = gen_rand_srccode(lang, extra_prompt, env) + if result is SourceFile: + print(getPathLine(result)) + print() + print(result.contents) + else: + print(result) +except LLMAPIException as e: + sys.stderr.write(ARGV0 + ": model gave up: " + str(e) + os.linesep) + exit(1) |