[AIC-py] usability fixes for llama cookbook

cookbooks/Cli-Mate/cli-mate.py

@@ -1,8 +1,13 @@
+from dataclasses import dataclass
+from textwrap import dedent
 import warnings
 
+from result import Err, Ok, Result
+
 warnings.filterwarnings("ignore")
 
 import asyncio
+from asyncio import AbstractEventLoop
 import argparse
 import signal
 import sys
@@ -81,12 +86,110 @@ async def mod_code(
     return 0
 
 
+@dataclass
+class Help:
+    pass
+
+
+@dataclass
+class Run:
+    user_input: str
+
+
+class Reload:
+    # TODO
+    # if should_reload and source_code_file is not None:
+    # user_input = deprefix(user_input.strip(), "reload")
+    # with open(source_code_file.strip(), "r", encoding="utf8") as file:
+    #     source_code = file.read()
+    #     llm_input = f"QUERY ABOUT SOURCE CODE:\n{user_input}\nSOURCE CODE:\n```{source_code}\n```"
+    pass
+
+
+class Clear:
+    pass
+
+
+class Pass:
+    pass
+
+
+class MultilineToggle:
+    pass
+
+
+Command = Pass | Help | Run | Reload | Clear | MultilineToggle
+
+
+def _get_command(user_input: str) -> Command:
+    normed = user_input.strip().lower()
+    if normed in ["h", "help", "?"]:
+        return Help()
+    elif normed in ["r", "reload"]:
+        return Reload()
+    elif normed in ["c", "clear"]:
+        return Clear()
+    elif normed in ["m", "multiline"]:
+        return MultilineToggle()
+    else:
+        return Run(user_input=user_input)
+
+
+def _print_help():
+    print(
+        dedent(
+            """
+        Exit loop: Ctrl-D
+        Toggle multiline input mode: m or multiline
+        Clear screen: c or clear
+        Reload source code: r or reload
+            """
+        )
+    )
+
+
+async def _run_llm(
+    runtime: AIConfigRuntime, llm_input: str
+) -> Result[ExecuteResult, str]:
+    # Dynamically generate the prompt name and prompt object
+    new_prompt_name = f"prompt{len(runtime.prompts)+1}"  # Prompt{number of prompts}
+    new_prompt = Prompt(name=new_prompt_name, input=llm_input)
+
+    # Add the new prompt and run the model
+    runtime.add_prompt(new_prompt.name, new_prompt)
+
+    def callback(delta: Any, _: Any, __: int):
+        if state["interrupt"]:
+            raise InterruptException()
+
+        print(delta.get("content", ""), end="", flush=True)
+
+    options = InferenceOptions(stream=True, stream_callback=callback)
+    state["interrupt"] = False
+    try:
+        result = await runtime.run(new_prompt_name, {}, options=options)
+        print(flush=True)
+        return Ok(result)
+    except InterruptException:
+        return Err("interrupted")
+
+
+async def _get_raw_input(
+    event_loop: AbstractEventLoop, session: PromptSession[str], is_multiline: bool
+) -> str:
+    def _prompt():  # type: ignore
+        return session.prompt(
+            "> ",
+            multiline=is_multiline,
+        )
+
+    return await event_loop.run_in_executor(None, _prompt)
+
+
 async def loop(aiconfig_path: str, source_code_file: str | None):
     runtime = AIConfigRuntime.load(aiconfig_path)
     event_loop = asyncio.get_event_loop()
-
     session = PromptSession()
-
     state["interrupt"] = False
 
     def signal_handler(_: int, __: FrameType | None):
@@ -95,52 +198,47 @@ def signal_handler(_: int, __: FrameType | None):
 
     signal.signal(signal.SIGINT, signal_handler)
 
-    i = 0
+    is_multiline = False
+    print("Enter 'h', 'help', or '?' for help.", flush=True)
     while True:
         try:
-            user_input = await event_loop.run_in_executor(
-                None, session.prompt, "Query: [ctrl-D to exit] "
-            )
+            raw_input = await _get_raw_input(event_loop, session, is_multiline)
         except KeyboardInterrupt:
             continue
         except EOFError:
             print("Exiting")
             break
 
-        if user_input.strip() == "":
-            continue
-
-        should_reload = user_input.strip().startswith("reload") or i == 0
-        if should_reload and source_code_file is not None:
-            user_input = deprefix(user_input.strip(), "reload")
-            with open(source_code_file.strip(), "r", encoding="utf8") as file:
-                source_code = file.read()
-                prompt = f"QUERY ABOUT SOURCE CODE:\n{user_input}\nSOURCE CODE:\n```{source_code}\n```"
-        else:
-            prompt = user_input
-
-        # Dynamically generate the prompt name and prompt object
-        new_prompt_name = f"prompt{len(runtime.prompts)+1}"  # Prompt{number of prompts}
-        new_prompt = Prompt(name=new_prompt_name, input=prompt)
-
-        # Add the new prompt and run the model
-        runtime.add_prompt(new_prompt.name, new_prompt)
-
-        def callback(delta: Any, _: Any, __: int):
-            if state["interrupt"]:
-                raise InterruptException()
-
-            print(delta.get("content", ""), end="", flush=True)
-
-        options = InferenceOptions(stream=True, stream_callback=callback)
-        state["interrupt"] = False
-        try:
-            result = await runtime.run(new_prompt_name, {}, options=options)
-            # print(f"{result=}")
-            print(flush=True)
-            i += 1
-        except InterruptException:
-            continue
+        command = _get_command(raw_input)
+
+        match command:
+            case Pass():
+                pass
+            case Help():
+                _print_help()
+            case MultilineToggle():
+                is_multiline = not is_multiline
+                print(f"Multiline input mode: {'on' if is_multiline else 'off'}")
+                if is_multiline:
+                    print("Hit option-enter to submit.")
+            case Run(user_input=user_input):
+                prompt = f"""
+                    INSTRUCTIONS: respond to the following query as concisely as possible.
+                    Do not output more tokens than necessary.
+                    QUERY: {user_input}
+                    """
+                llm_res = await _run_llm(runtime, prompt)
+                match llm_res:
+                    case Ok(_):
+                        # TODO somethign with res?
+                        pass
+                    case Err(msg):
+                        print(msg)
+            case Reload():
+                # TODO
+                pass
+            case Clear():
+                print("\033c", end="")
 
 
 async def main():

cookbooks/llama/python/README.md

@@ -0,0 +1,16 @@
+Example app using llama aiconfig extension.
+
+1. Install (example using anaconda)
+   `conda create -n aiconfig-llama-cookbook`
+   `conda activate aiconfig-llama-cookbook`
+   `conda install pip`
+   `pip install -r python/requirements.txt`
+
+2. Download a model, e.g.
+   `curl -L https://huggingface.co/TheBloke/Llama-2-7B-Chat-GGUF/resolve/main/llama-2-7b-chat.Q4_K_M.gguf --output ./models/llama-2-7b-chat.Q4_K_M.gguf`
+3. cd into cookbook root dir
+   `$ pwd`
+   aiconfig/cookbooks/llama
+4. Create an AIConfig like this: https://github.com/lastmile-ai/aiconfig/blob/e92e5a3c80b9c2b74a9432f0441318a951d54d0c/cookbooks/llama/llama-aiconfig.json
+5. Run with your local paths:
+   `python python/ask_llama.py --aiconfig-path='../llama/llama-aiconfig.json' --model-path='../../models/llama-2-7b-chat.Q4_K_M.gguf' 2> ask-llama.err`

cookbooks/llama/python/ask_llama.py

@@ -1,15 +1,33 @@
 import asyncio
+import sys
 
 from aiconfig.model_parser import InferenceOptions
 from llama import LlamaModelParser
 
 from aiconfig import AIConfigRuntime
+import argparse
 
 
 async def main():
-    llama_model_parser = LlamaModelParser(
-        model_path="models/llama-2-7b-chat.Q4_K_M.gguf"
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--aiconfig-path",
+        type=str,
+        required=True,
+        help="Relative or absolute path to aiconfig json, e.g. cookbooks/llama/llama-aiconfig.json",
     )
+    parser.add_argument(
+        "--model-path",
+        type=str,
+        required=True,
+        help="Relative or absolute path to model",
+    )
+    args = parser.parse_args()
+    return await run(args.aiconfig_path, args.model_path)
+
+
+async def run(aiconfig_path: str, model_path: str):
+    llama_model_parser = LlamaModelParser(model_path=model_path)
 
     for lm in [
         "llama-2-7b-chat",
@@ -18,7 +36,7 @@ async def main():
     ]:
         AIConfigRuntime.register_model_parser(llama_model_parser, lm)
 
-    config = AIConfigRuntime.load("cookbooks/llama/llama-aiconfig.json")
+    config = AIConfigRuntime.load(aiconfig_path)
 
     def stream_callback(data, accumulated_message, index):
         print(data, end="", flush=True)
@@ -38,9 +56,9 @@ def stream_callback(data, accumulated_message, index):
     await config.run("prompt13b", params={}, options=inference_options)
 
     print("\n\nRunning prompt13b_code...")
-    code_res = await config.run("prompt13b_code", params={}, options=inference_options)
-    print(f"\n\n\n\nCode response:\n{code_res}")
+    await config.run("prompt13b_code", params={}, options=inference_options)
 
 
 if __name__ == "__main__":
-    asyncio.run(main())
+    res = asyncio.run(main())
+    sys.exit(res)

cookbooks/llama/python/pyproject.toml

@@ -0,0 +1,43 @@
+[build-system]
+requires = ["setuptools", "wheel"]
+
+[project]
+name = "python-aiconfig-llama-cookbook"
+version = "0.0.1"
+authors = [
+    { name="Jonathan Lessinger", email="[email protected]" },
+]
+description = "LLama cookbook using AIConfig Library"
+readme = "README.md"
+requires-python = ">=3.7"
+classifiers = [
+  "Intended Audience :: Developers",
+  "Programming Language :: Python :: 3",
+]
+dynamic = ["dependencies"]
+
+[tool.setuptools.dynamic]
+dependencies = {file = ["requirements.txt"]}
+
+[project.urls]
+"Homepage" = "https://github.com/lastmile-ai/aiconfig"
+"Bug Tracker" = "https://github.com/lastmile-ai/aiconfig/issues"
+
+# Black formatting
+[tool.black]
+line-length = 99
+include = '\.pyi?$'
+exclude = '''
+/(
+      .eggs         # exclude a few common directories in the
+    | .git          # root of the project
+    | .hg
+    | .mypy_cache
+    | .tox
+    | venv
+    | _build
+    | buck-out
+    | build
+    | dist
+  )/
+'''

cookbooks/llama/python/requirements.txt

@@ -0,0 +1,2 @@
+python-aiconfig-llama
+python-aiconfig

extensions/llama/python/README.md

@@ -0,0 +1,5 @@
+AIConfig Model Parser for LLama cpp (Python bindings).
+
+Usage: See cookbook:
+
+https://github.com/lastmile-ai/aiconfig/blob/c64224ed48ccb7f8cbd2d3a1b2e8bd250aeb9ff2/cookbooks/llama/python/ask_llama.py#L4

extensions/llama/python/pyproject.toml

@@ -0,0 +1,43 @@
+[build-system]
+requires = ["setuptools", "wheel"]
+
+[project]
+name = "python-aiconfig-llama"
+version = "0.0.1"
+authors = [
+    { name="Jonathan Lessinger", email="[email protected]" },
+]
+description = "LLama extension for AIConfig Library"
+readme = "README.md"
+requires-python = ">=3.7"
+classifiers = [
+  "Intended Audience :: Developers",
+  "Programming Language :: Python :: 3",
+]
+dynamic = ["dependencies"]
+
+[tool.setuptools.dynamic]
+dependencies = {file = ["requirements.txt"]}
+
+[project.urls]
+"Homepage" = "https://github.com/lastmile-ai/aiconfig"
+"Bug Tracker" = "https://github.com/lastmile-ai/aiconfig/issues"
+
+# Black formatting
+[tool.black]
+line-length = 99
+include = '\.pyi?$'
+exclude = '''
+/(
+      .eggs         # exclude a few common directories in the
+    | .git          # root of the project
+    | .hg
+    | .mypy_cache
+    | .tox
+    | venv
+    | _build
+    | buck-out
+    | build
+    | dist
+  )/
+'''