dragon-cli 0.2.0__py3-none-any.whl

dragon_cli/__init__.py

@@ -0,0 +1,5 @@
+__version__ = "0.2.0"
+
+from ._cuda import preload_cuda_runtime
+
+preload_cuda_runtime()

dragon_cli/__main__.py

@@ -0,0 +1,4 @@
+from .cli import main
+
+if __name__ == "__main__":
+    main()

dragon_cli/_cuda.py

@@ -0,0 +1,42 @@
+"""Preload CUDA runtime shared libraries shipped via pip wheels.
+
+The GPU-enabled llama-cpp-python wheel links against libcudart/libcublas/
+libnvrtc, which the pip ``nvidia-*-cu12`` packages install under
+``site-packages/nvidia/*/lib`` (not on the system library path). Loading them
+here by absolute path with ``RTLD_GLOBAL`` lets libllama.so resolve them
+without LD_LIBRARY_PATH. No-op when the libs are already resolvable or absent
+(e.g. the CPU-only wheel).
+
+Kept dependency-free (stdlib only) so it can run at package import time
+without risking circular imports.
+"""
+
+import ctypes
+import glob
+import os
+import sys
+
+_CUDA_SONAMES = (
+    "libcudart.so.12",
+    "libcublas.so.12",
+    "libcublasLt.so.12",
+    "libnvrtc.so.12",
+)
+
+
+def preload_cuda_runtime() -> None:
+    loaded: set[str] = set()
+    for entry in sys.path:
+        nvidia_dir = os.path.join(entry, "nvidia")
+        if not entry or not os.path.isdir(nvidia_dir):
+            continue
+        for soname in _CUDA_SONAMES:
+            if soname in loaded:
+                continue
+            for match in glob.glob(os.path.join(nvidia_dir, "*", "lib", soname)):
+                try:
+                    ctypes.CDLL(match, mode=ctypes.RTLD_GLOBAL)
+                    loaded.add(soname)
+                    break
+                except OSError:
+                    continue

dragon_cli/app.py

@@ -0,0 +1,123 @@
+"""Async orchestration: model selection, MCP startup, and the chat session."""
+
+from pathlib import Path
+
+from .chat import ChatSession, chat_loop
+from .config import C_RED, C_RESET
+from .mcp import MCPManager
+from .models import MODEL_CATALOG, available_ram_gb, download_model, list_catalog, load_llm
+from .storage import ConversationStore
+
+
+def _gpu_layers_arg(args) -> int | None:
+    """Translate ``--n-gpu-layers`` into a ``load_llm`` override.
+
+    The CLI default is -1 (auto). Any value >= 0 is an explicit override
+    (0 = CPU only, >0 = exact layer count). None means "let the runtime
+    auto-decide".
+    """
+    n = getattr(args, "n_gpu_layers", -1)
+    return n if n is not None and n >= 0 else None
+
+
+def choose_model(args) -> dict | None:
+    # Handle --auto flag for automatic model selection
+    if getattr(args, 'auto', False):
+        from .recommendation import ModelRecommender
+        recommender = ModelRecommender()
+        best_model = recommender.recommend_best(priority=getattr(args, 'priority', 'balanced'))
+        
+        if best_model:
+            print(f"\n{C_RESET}Auto-selected model based on your hardware:")
+            print(f"  → {best_model['name']}")
+            reason = recommender.get_recommendation_reason(best_model)
+            print(f"  {C_DIM}Why: {reason}{C_RESET}\n")
+            return best_model
+        else:
+            print(f"{C_RED}[error]{C_RESET} No compatible models found for your hardware.")
+            # Fall back to manual selection
+            pass
+
+    # Skip the plain catalog print when a model was already chosen upstream
+    # (e.g. by the rich interactive menu), so the user doesn't see it twice.
+    if not args.model:
+        list_catalog()
+    if args.model:
+        choice = args.model
+    else:
+        try:
+            choice = int(input("\nChoose a model number: ").strip())
+        except (ValueError, EOFError, KeyboardInterrupt):
+            print("Invalid choice.")
+            return None
+    if not (1 <= choice <= len(MODEL_CATALOG)):
+        print("Number not in the catalog.")
+        return None
+    model = MODEL_CATALOG[choice - 1]
+    ram = available_ram_gb()
+    if 0 < ram < model["ram_min_gb"]:
+        ok = input(f"{C_RED}[warning]{C_RESET} Model needs ~{model['ram_min_gb']} GB, "
+                   f"you have {ram:.1f} GB free. Continue? (y/N) ").strip().lower()
+        if ok != "y":
+            return None
+    return model
+
+
+def print_conversations(store: ConversationStore) -> None:
+    conversations = store.list()
+    if not conversations:
+        print("No conversations yet.")
+        return
+    for c in conversations:
+        print(f"  #{c.id}  {c.title}  ({c.message_count} msgs, updated {c.updated_at})")
+
+
+async def run(args) -> None:
+    if args.serve:
+        model = choose_model(args)
+        if not model:
+            return
+        path = download_model(model)
+        llm = load_llm(path, model["ctx"], model_info=model,
+                       gpu_layers=_gpu_layers_arg(args),
+                       force_cpu=getattr(args, "cpu", False),
+                       force_gpu=getattr(args, "gpu", False))
+        from .server import run_server
+        await run_server(llm, model["name"], args.host, args.port)
+        return
+
+    store = ConversationStore(Path(args.db))
+
+    if args.list_conversations:
+        print_conversations(store)
+        store.close()
+        return
+
+    model = choose_model(args)
+    if not model:
+        store.close()
+        return
+
+    path = download_model(model)
+    llm = load_llm(path, model["ctx"], model_info=model,
+                   gpu_layers=_gpu_layers_arg(args),
+                   force_cpu=getattr(args, "cpu", False),
+                   force_gpu=getattr(args, "gpu", False))
+
+    mcp = MCPManager()
+    if not args.no_mcp:
+        await mcp.start()
+
+    session = ChatSession(llm, mcp, store, model["name"])
+    if args.resume is not None:
+        if not session.resume(args.resume):
+            print(f"{C_RED}[error]{C_RESET} No conversation with id {args.resume}, starting a new one.")
+            session.start_new()
+    else:
+        session.start_new()
+
+    try:
+        await chat_loop(session)
+    finally:
+        await mcp.stop()
+        store.close()

dragon_cli/banner.py

@@ -0,0 +1,96 @@
+"""ASCII dragon banner shown at startup."""
+
+BANNER = r'''
+    _,_
+  /7/Y/^\
+  vuVV|C)|                        __ _
+    \|^ /                       .'  Y '>,
+    )| \)                      / _   _   \
+   //)|\\                      )(_) (_)(|}
+  / ^| \ \                     {  4A   } /
+ //^| || \\                     \uLuJJ/\l
+>//||| |\\\|                    |3    p)/
+| """""  7/>l__ _____ ____      /nnm_n//
+L>_   _-< 7/|_-__,__-)\,__)(".  \_>-<_/D
+)D" Y "c)  9)       //V     \_"-._.__G<D     __
+\  ||    . _/       (       \ "-._"> G<D-_.-"_=)_
+ ) __   / /         "      _/.__ />_c__.-__<"/ ( \
+ | | |  |(|               < "-._"> _.G_.___)\   \7\
+  \"=" // |              (,"-.__.|\ \<.__.-" )   \ \
+   '---'  |              |,"-.__"| \!"-.__.-".\   \ \
+     |_;._/              (_"-.__"'\ \"-.__.-".|    \_\
+     )(" V                \"-.__"'|\ \-.__.-".)     \ \
+        (                  "-.__'"\_\ \.__.-"./      \ l
+         )                  ".__"">>G\ \__.-">        V )
+  AI PROVIDERS                 ""  G<\ \_.-"        ./7
+                                     G<\ \          ///
+                                ___  G<D\ \        /7^
+                              /".-_"-"G"<\ \__    ///))
+                              ||,-_ "     '.\_"- /// (\
+                              | (_)  _   "/ __7 |X/   V
+                             /\__""-/ "\ _""c_"GBV
+                            _)(  ""'    "--.c__>D'
+                           /  /             )E_>"
+                         _/  (             |  \()
+                        / \ /              |  |
+                        /\\|               |  |
+                       / '((               |  |
+                      /  / )\              \  |
+                     /  y  \y              |Y |
+                    /  /    (              |  |
+                   L ."     |              |  /
+                  | \(                     |  |
+                   \_|                     |  |
+                   |  \                    { "|
+                   | ||                    |  |
+                   |x||                    \_,/
+                   } ||                    / \'
+                   | ||                    |_/
+                   | (|                    | }\
+                   | ||                    } ||
+                   | ||                    | ||
+                   | ||                    |\||
+                   / ||                    | ||
+                   | ||                    ( |!
+                   | |/                    ) ||
+                 _/   \                    | }|
+             _.-"_ ( )|   jjs              ! ||
+          c_"-_-"_    )                    | ||
+           c,-_-'_--""                     { ||
+           "C(_/"                          \ /|
+                                           (! )
+                                           /| \
+                                          /  |(
+                                         /7||\\
+                                        ()U cUu"
+'''
+
+SUBTITLE = "  DRAGON CLI — local LLM + MCP  |  /help for commands"
+
+# Compact mascot icon shown as a splash on the interactive menu. Raw string so
+# the backslashes in the art are preserved verbatim.
+ICON = r""",--.
+                          {    }
+                          K,   }
+                         /  `Y`
+                    _   /   /
+                   {_'-K.__/
+                     `/-.__L._
+                     /  ' /`\_}
+                    /  ' /     -AI PROVIDERS-
+            ____   /  ' /
+     ,-'~~~~    ~~/  ' /_
+   ,'             ``~~~%%',
+  (                     %  Y
+ {                      %% I
+{      -                 %  `.
+|       ',                %  )
+|        |   ,..__      __. Y
+|    .,_./  Y ' / ^Y   J   )|
+\           |' /   |   |   ||
+ \          L_/    . _ (_,.'(
+  \,   ,      ^^""' / |      )
+    \_  \          /,L]     /
+      '-_`-,       ` `   ./`
+         `-(_            )
+             ^^\..___,.--`"""

dragon_cli/chat/__init__.py

@@ -0,0 +1,4 @@
+from .session import ChatSession, chat_loop
+from .streaming import generate_streaming
+
+__all__ = ["ChatSession", "chat_loop", "generate_streaming"]

dragon_cli/chat/commands.py

@@ -0,0 +1,175 @@
+"""Slash-command dispatch for the chat loop: help, model info, and
+conversation management (list/new/resume/rename/delete/search/export)."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import TYPE_CHECKING, Callable
+
+from ..config import C_DIM, C_GREEN, C_RED, C_RESET, MCP_CONFIG_FILE
+from ..models import list_catalog
+
+if TYPE_CHECKING:
+    from .session import ChatSession
+
+HELP_TEXT = """Commands:
+  /help              show this help
+  /models            list the model catalog
+  /tools             list connected MCP tools
+  /mcp               show MCP server connection status
+  /clear             clear the current in-memory history
+  /new [title]       start a new conversation (saves and switches)
+  /list              list saved conversations
+  /resume <id>       resume a saved conversation
+  /rename <id> <title>  rename a saved conversation
+  /delete <id>       delete a saved conversation
+  /search <text>     search saved messages for text
+  /export <id> [md|json]  export a conversation to a file
+  /exit, /quit       end the program"""
+
+
+def cmd_help(session: "ChatSession", arg: str) -> None:
+    print(HELP_TEXT)
+
+
+def cmd_models(session: "ChatSession", arg: str) -> None:
+    list_catalog()
+
+
+def cmd_tools(session: "ChatSession", arg: str) -> None:
+    if not session.mcp.tools:
+        print("No MCP tools connected.")
+        return
+    for name, info in session.mcp.tools.items():
+        print(f"  - {name}: {info['description'][:90]}")
+
+
+def cmd_mcp(session: "ChatSession", arg: str) -> None:
+    servers = list(session.mcp.sessions.keys()) or "none"
+    print(f"Connected servers: {servers}")
+    print(f"Config: {MCP_CONFIG_FILE.resolve()}")
+
+
+def cmd_clear(session: "ChatSession", arg: str) -> None:
+    session.clear()
+    print("History cleared.")
+
+
+def cmd_new(session: "ChatSession", arg: str) -> None:
+    session.start_new(arg.strip() or None)
+    print(f"{C_GREEN}[ok]{C_RESET} Started conversation #{session.conv_id}.")
+
+
+def cmd_list(session: "ChatSession", arg: str) -> None:
+    conversations = session.store.list()
+    if not conversations:
+        print("No saved conversations yet.")
+        return
+    for c in conversations:
+        marker = " *" if c.id == session.conv_id else ""
+        print(f"  #{c.id}{marker}  {c.title}  ({c.message_count} msgs, updated {c.updated_at})")
+
+
+def cmd_resume(session: "ChatSession", arg: str) -> None:
+    try:
+        conv_id = int(arg.strip())
+    except ValueError:
+        print("Usage: /resume <id>")
+        return
+    if session.resume(conv_id):
+        print(f"{C_GREEN}[ok]{C_RESET} Resumed conversation #{conv_id}.")
+    else:
+        print(f"{C_RED}[error]{C_RESET} No conversation with id {conv_id}.")
+
+
+def cmd_rename(session: "ChatSession", arg: str) -> None:
+    parts = arg.strip().split(maxsplit=1)
+    if len(parts) != 2 or not parts[0].isdigit():
+        print("Usage: /rename <id> <new title>")
+        return
+    conv_id, title = int(parts[0]), parts[1]
+    if session.store.rename(conv_id, title):
+        print(f"{C_GREEN}[ok]{C_RESET} Renamed #{conv_id} to '{title}'.")
+    else:
+        print(f"{C_RED}[error]{C_RESET} No conversation with id {conv_id}.")
+
+
+def cmd_delete(session: "ChatSession", arg: str) -> None:
+    if not arg.strip().isdigit():
+        print("Usage: /delete <id>")
+        return
+    conv_id = int(arg.strip())
+    confirm = input(f"Delete conversation #{conv_id}? This cannot be undone. (y/N) ").strip().lower()
+    if confirm != "y":
+        print("Cancelled.")
+        return
+    if session.store.delete(conv_id):
+        print(f"{C_GREEN}[ok]{C_RESET} Deleted #{conv_id}.")
+        if conv_id == session.conv_id:
+            session.start_new()
+    else:
+        print(f"{C_RED}[error]{C_RESET} No conversation with id {conv_id}.")
+
+
+def cmd_search(session: "ChatSession", arg: str) -> None:
+    text = arg.strip()
+    if not text:
+        print("Usage: /search <text>")
+        return
+    results = session.store.search(text)
+    if not results:
+        print("No matches.")
+        return
+    for conv_id, title, role, content in results[:20]:
+        snippet = content.replace("\n", " ")[:80]
+        print(f"  #{conv_id} {title} [{role}]: {snippet}")
+
+
+def cmd_export(session: "ChatSession", arg: str) -> None:
+    parts = arg.strip().split()
+    if not parts or not parts[0].isdigit():
+        print("Usage: /export <id> [md|json]")
+        return
+    conv_id = int(parts[0])
+    fmt = parts[1] if len(parts) > 1 else "md"
+    if fmt not in ("md", "json"):
+        print("Format must be 'md' or 'json'.")
+        return
+    path = Path(f"conversation_{conv_id}.{fmt}")
+    result = session.store.export(conv_id, path, fmt)
+    if result is None:
+        print(f"{C_RED}[error]{C_RESET} No conversation with id {conv_id}.")
+    else:
+        print(f"{C_GREEN}[ok]{C_RESET} Exported to {result.resolve()}")
+
+
+def cmd_exit(session: "ChatSession", arg: str) -> bool:
+    return False
+
+
+COMMANDS: dict[str, Callable[["ChatSession", str], bool | None]] = {
+    "help": cmd_help,
+    "models": cmd_models,
+    "tools": cmd_tools,
+    "mcp": cmd_mcp,
+    "clear": cmd_clear,
+    "new": cmd_new,
+    "list": cmd_list,
+    "resume": cmd_resume,
+    "rename": cmd_rename,
+    "delete": cmd_delete,
+    "search": cmd_search,
+    "export": cmd_export,
+    "exit": cmd_exit,
+    "quit": cmd_exit,
+}
+
+
+def dispatch(session: "ChatSession", raw_input: str) -> bool | None:
+    """Runs a slash command. Returns False if the chat loop should stop."""
+    name, _, arg = raw_input[1:].partition(" ")
+    handler = COMMANDS.get(name.lower())
+    if handler is None:
+        print(f"{C_DIM}Unknown command '{raw_input.split()[0]}'. Try /help.{C_RESET}")
+        return True
+    return handler(session, arg)

dragon_cli/chat/session.py

@@ -0,0 +1,97 @@
+"""Chat session state: history, current conversation, and the main loop."""
+
+import asyncio
+
+from ..config import C_DIM, C_GREEN, C_YELLOW, C_RESET
+from ..mcp import MCPManager, extract_tool_call
+from ..storage import ConversationStore
+from .commands import dispatch
+from .streaming import generate_streaming
+
+SYSTEM_PROMPT = "You are a helpful assistant running locally on the user's machine."
+
+
+class ChatSession:
+    """Bundles the LLM, MCP manager, conversation store, and in-memory history."""
+
+    def __init__(self, llm, mcp: MCPManager, store: ConversationStore, model_name: str):
+        self.llm = llm
+        self.mcp = mcp
+        self.store = store
+        self.model_name = model_name
+        self.conv_id: int | None = None
+        self.history: list[dict] = []
+
+    def _system_message(self) -> str:
+        return SYSTEM_PROMPT + self.mcp.tools_prompt()
+
+    def start_new(self, title: str | None = None) -> None:
+        title = title or "New conversation"
+        self.conv_id = self.store.create(title, self.model_name)
+        self.history = [{"role": "system", "content": self._system_message()}]
+        self.store.append_message(self.conv_id, "system", self.history[0]["content"])
+
+    def resume(self, conv_id: int) -> bool:
+        found = self.store.get(conv_id)
+        if found is None:
+            return False
+        conv, messages = found
+        self.conv_id = conv.id
+        self.history = [{"role": m.role, "content": m.content} for m in messages]
+        if not self.history or self.history[0]["role"] != "system":
+            self.history.insert(0, {"role": "system", "content": self._system_message()})
+        return True
+
+    def clear(self) -> None:
+        self.history = self.history[:1]
+
+    async def turn(self, user_input: str) -> None:
+        self.history.append({"role": "user", "content": user_input})
+        self.store.append_message(self.conv_id, "user", user_input)
+
+        for _ in range(5):
+            response = await asyncio.to_thread(generate_streaming, self.llm, self.history)
+            self.history.append({"role": "assistant", "content": response})
+            self.store.append_message(self.conv_id, "assistant", response)
+
+            call = extract_tool_call(response)
+            if not call or not self.mcp.tools:
+                break
+            tool_name, arguments = call
+            print(f"{C_DIM}[tool] running {tool_name}({arguments}){C_RESET}")
+            result = await self.mcp.call(tool_name, arguments)
+            print(f"{C_DIM}[tool] result: {result[:200]}...{C_RESET}")
+            followup = (
+                f"[TOOL RESULT for {tool_name}]:\n{result}\n\n"
+                f"Now answer the user's original question based on this."
+            )
+            self.history.append({"role": "user", "content": followup})
+            self.store.append_message(self.conv_id, "user", followup)
+
+
+async def chat_loop(session: ChatSession) -> None:
+    if session.conv_id is None:
+        session.start_new()
+
+    backend = getattr(session.llm, "exec_mode", "CPU")
+    n_tools = len(session.mcp.tools)
+    tools_msg = f" | {n_tools} MCP tool{'s' if n_tools != 1 else ''}" if n_tools else " | MCP off"
+    print(f"\n{C_GREEN}=== Chat started — {session.model_name} on {backend}{tools_msg} "
+          f"— /help for commands, /exit to quit ==={C_RESET}\n")
+    while True:
+        try:
+            user_input = (await asyncio.to_thread(input, f"{C_YELLOW}You:{C_RESET} ")).strip()
+        except (EOFError, KeyboardInterrupt):
+            break
+        if not user_input:
+            continue
+
+        if user_input.startswith("/"):
+            keep_going = dispatch(session, user_input)
+            if keep_going is False:
+                break
+            continue
+
+        await session.turn(user_input)
+
+    print("Shutting down. See you next time!")

dragon_cli/chat/streaming.py

@@ -0,0 +1,18 @@
+"""Streaming text generation via llama.cpp."""
+
+from ..config import C_CYAN, C_RESET
+
+
+def generate_streaming(llm, history: list[dict]) -> str:
+    print(f"{C_CYAN}Model:{C_RESET} ", end="", flush=True)
+    response = ""
+    stream = llm.create_chat_completion(
+        messages=history, max_tokens=1024, temperature=0.7, stream=True,
+    )
+    for chunk in stream:
+        token = chunk["choices"][0]["delta"].get("content", "")
+        if token:
+            print(token, end="", flush=True)
+            response += token
+    print()
+    return response