vllmd 0.2.0__py3-none-any.whl

src/vllmd/__init__.py

@@ -0,0 +1,5 @@
+"""vllmd — run local models via vLLM in Docker containers."""
+
+from .runner import RunConfig, logs, start, status, stop
+
+__all__ = ["RunConfig", "logs", "start", "status", "stop"]

src/vllmd/cli.py

@@ -0,0 +1,292 @@
+"""CLI entry point."""
+
+from __future__ import annotations
+
+import subprocess
+from pathlib import Path
+from typing import Annotated
+
+import typer
+from rich.console import Console
+from rich.table import Table
+
+from .runner import (
+    RunConfig,
+    build_docker_run_cmd,
+    list_containers,
+    logs,
+    status,
+    stop,
+    stop_all,
+    wait_ready,
+)
+from .sessions.cli import session_app
+from .vectordb.cli import db_app
+
+app = typer.Typer(
+    name="vllmd",
+    help="Run local models via vLLM in Docker containers.",
+    no_args_is_help=True,
+)
+app.add_typer(db_app, name="db")
+app.add_typer(session_app, name="session")
+console = Console()
+
+_NAME_HELP = "Container name (default: vllmd-<model-dir-name>)"
+
+
+@app.command()
+def run(
+    model: Annotated[
+        Path,
+        typer.Option("--model", "-m", help="Path to the model directory on disk"),
+    ],
+    port: Annotated[
+        int,
+        typer.Option("--port", "-p", help="Host port to expose the vLLM API on"),
+    ] = 8000,
+    name: Annotated[
+        str | None,
+        typer.Option("--name", "-n", help=_NAME_HELP),
+    ] = None,
+    gpu: Annotated[
+        bool,
+        typer.Option("--gpu/--no-gpu", help="Pass --gpus all to the container"),
+    ] = True,
+    dtype: Annotated[
+        str,
+        typer.Option("--dtype", help="Model dtype (auto, float16, bfloat16, float32)"),
+    ] = "auto",
+    max_model_len: Annotated[
+        int | None,
+        typer.Option("--max-model-len", help="Override max context length"),
+    ] = None,
+    detach: Annotated[
+        bool,
+        typer.Option(
+            "--detach",
+            "-d",
+            help="Start container in background; wait for API ready, then return.",
+        ),
+    ] = False,
+    wait: Annotated[
+        bool,
+        typer.Option(
+            "--wait/--no-wait",
+            help="With --detach: wait for the API to be ready before returning.",
+        ),
+    ] = True,
+    extra: Annotated[
+        list[str] | None,
+        typer.Argument(help="Extra args forwarded verbatim to vLLM"),
+    ] = None,
+) -> None:
+    """Start a vLLM container serving MODEL on PORT."""
+    config = RunConfig(
+        model_path=model,
+        port=port,
+        name=name,
+        gpu=gpu,
+        dtype=dtype,
+        max_model_len=max_model_len,
+        extra_args=extra or [],
+    )
+
+    model_path = model.resolve()
+    console.print(f"[bold]Starting vLLM container[/bold] '{config.container_name}'")
+    console.print(f"  Model:    {model_path}")
+    console.print(f"  Model ID: {config.model_id}")
+    console.print(f"  Port:     {port}")
+    console.print(f"  GPU:      {'yes' if gpu else 'no'}")
+    console.print()
+
+    docker_cmd = build_docker_run_cmd(config)
+
+    if detach:
+        subprocess.Popen(
+            docker_cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL
+        )
+        if wait:
+            console.print("[dim]Waiting for API to become ready…[/dim]")
+            if wait_ready(config):
+                console.print(f"[green]Ready.[/green] Endpoint: {config.endpoint}")
+                console.print(f"  Model ID: [bold]{config.model_id}[/bold]")
+            else:
+                console.print(
+                    "[yellow]Timed out waiting for API. "
+                    "Container may still be loading.[/yellow]"
+                )
+                console.print(f"  Endpoint: {config.endpoint}")
+        else:
+            console.print(f"Container started. Endpoint: {config.endpoint}")
+    else:
+        try:
+            subprocess.run(docker_cmd, check=True)
+        except FileNotFoundError as e:
+            console.print(f"[red]{e}[/red]")
+            raise typer.Exit(1) from e
+        except subprocess.CalledProcessError as e:
+            console.print(f"[red]Docker exited with code {e.returncode}[/red]")
+            raise typer.Exit(e.returncode) from e
+
+
+@app.command(name="stop")
+def stop_cmd(
+    name: Annotated[
+        str | None,
+        typer.Option("--name", "-n", help="Container name to stop"),
+    ] = None,
+    all_containers: Annotated[
+        bool,
+        typer.Option("--all", "-a", help="Stop all vllmd-managed containers"),
+    ] = False,
+) -> None:
+    """Stop one or all vllmd containers."""
+    if all_containers:
+        stopped = stop_all()
+        if stopped:
+            for n in stopped:
+                console.print(f"[green]Stopped '{n}'.[/green]")
+        else:
+            console.print("[yellow]No running vllmd containers found.[/yellow]")
+        return
+
+    if name is None:
+        running = list_containers()
+        if len(running) == 1:
+            name = running[0]["name"]
+        elif len(running) == 0:
+            console.print("[yellow]No running vllmd containers found.[/yellow]")
+            raise typer.Exit(1)
+        else:
+            console.print(
+                "[red]Multiple containers running — specify --name or use --all:[/red]"
+            )
+            for c in running:
+                console.print(f"  {c['name']}")
+            raise typer.Exit(1)
+
+    try:
+        stop(name)
+        console.print(f"[green]Stopped '{name}'.[/green]")
+    except RuntimeError as e:
+        console.print(f"[red]{e}[/red]")
+        raise typer.Exit(1) from e
+    except subprocess.CalledProcessError as e:
+        console.print(f"[red]Docker exited with code {e.returncode}[/red]")
+        raise typer.Exit(e.returncode) from e
+
+
+@app.command(name="ps")
+def ps_cmd() -> None:
+    """List all running vllmd-managed containers."""
+    containers = list_containers()
+    if not containers:
+        console.print("[dim]No vllmd containers running.[/dim]")
+        return
+
+    table = Table("Name", "Model", "Port", "Endpoint", "Status", show_header=True)
+    for c in containers:
+        table.add_row(
+            c["name"],
+            c["model_id"],
+            str(c["port"] or "?"),
+            c["endpoint"],
+            c["status"],
+        )
+    console.print(table)
+
+
+@app.command(name="status")
+def status_cmd(
+    name: Annotated[
+        str | None,
+        typer.Option("--name", "-n", help="Container name (omit to show all)"),
+    ] = None,
+) -> None:
+    """Show container and API status. Omit --name to show all containers."""
+    if name is None:
+        # Show summary table for all managed containers
+        containers = list_containers()
+        if not containers:
+            console.print("[dim]No vllmd containers running.[/dim]")
+            raise typer.Exit(1)
+
+        table = Table("Name", "Model", "Port", "API", "Status", show_header=True)
+        all_healthy = True
+        for c in containers:
+            info = status(c["name"])
+            api_str = (
+                "[green]healthy[/green]"
+                if info["api_healthy"]
+                else "[yellow]unreachable[/yellow]"
+            )
+            if not info["api_healthy"]:
+                all_healthy = False
+            table.add_row(
+                c["name"],
+                c["model_id"],
+                str(c["port"] or "?"),
+                api_str,
+                c["status"],
+            )
+        console.print(table)
+        if not all_healthy:
+            raise typer.Exit(1)
+        return
+
+    info = status(name)
+
+    table = Table(show_header=False, box=None, padding=(0, 2))
+    table.add_column("Key", style="bold")
+    table.add_column("Value")
+
+    running_str = "[green]running[/green]" if info["running"] else "[red]stopped[/red]"
+    api_str = (
+        "[green]healthy[/green]"
+        if info["api_healthy"]
+        else "[yellow]unreachable[/yellow]"
+    )
+
+    table.add_row("Container", running_str)
+    table.add_row("API", api_str)
+    if info["container"]:
+        table.add_row("Started", info["container"].get("StartedAt", "—"))
+
+    console.print(table)
+    if not info["running"]:
+        raise typer.Exit(1)
+
+
+@app.command(name="logs")
+def logs_cmd(
+    name: Annotated[
+        str | None,
+        typer.Option(
+            "--name", "-n", help="Container name (auto-resolved if only one is running)"
+        ),
+    ] = None,
+    follow: Annotated[
+        bool,
+        typer.Option("--follow", "-f", help="Follow log output"),
+    ] = False,
+) -> None:
+    """Print logs from a vllmd container."""
+    if name is None:
+        running = list_containers()
+        if len(running) == 1:
+            name = running[0]["name"]
+        elif len(running) == 0:
+            console.print("[yellow]No running vllmd containers found.[/yellow]")
+            raise typer.Exit(1)
+        else:
+            console.print("[red]Multiple containers running — specify --name:[/red]")
+            for c in running:
+                console.print(f"  {c['name']}")
+            raise typer.Exit(1)
+
+    try:
+        logs(name, follow=follow)
+    except subprocess.CalledProcessError as e:
+        console.print(f"[red]Docker exited with code {e.returncode}[/red]")
+        raise typer.Exit(e.returncode) from e

src/vllmd/runner.py

@@ -0,0 +1,239 @@
+"""Docker management for vLLM model containers."""
+
+from __future__ import annotations
+
+import contextlib
+import json
+import re
+import subprocess
+import time
+import urllib.request
+from dataclasses import dataclass, field
+from pathlib import Path
+
+VLLM_IMAGE = "vllm/vllm-openai:latest"
+HEALTH_TIMEOUT = 300
+HEALTH_INTERVAL = 3
+
+MANAGED_LABEL = "com.vllmd.managed"
+MODEL_LABEL = "com.vllmd.model"
+MODEL_PATH_LABEL = "com.vllmd.model_path"
+
+
+@dataclass
+class RunConfig:
+    model_path: Path
+    port: int = 8000
+    name: str | None = None  # None → derived from model dir name
+    gpu: bool = True
+    dtype: str = "auto"
+    max_model_len: int | None = None
+    extra_args: list[str] = field(default_factory=list)
+
+    @property
+    def model_id(self) -> str:
+        return self.model_path.resolve().name
+
+    @property
+    def container_name(self) -> str:
+        return self.name or f"vllmd-{self.model_id}"
+
+    @property
+    def endpoint(self) -> str:
+        return f"http://localhost:{self.port}"
+
+
+def _docker(*args: str, capture: bool = False) -> subprocess.CompletedProcess:
+    return subprocess.run(
+        ["docker", *args],
+        check=True,
+        capture_output=capture,
+        text=True,
+    )
+
+
+def _container_exists(name: str) -> bool:
+    result = subprocess.run(
+        ["docker", "ps", "-a", "--filter", f"name=^{name}$", "--format", "{{.Names}}"],
+        capture_output=True,
+        text=True,
+    )
+    return name in result.stdout.splitlines()
+
+
+def _wait_ready(endpoint: str, timeout: int = HEALTH_TIMEOUT) -> bool:
+    deadline = time.monotonic() + timeout
+    while time.monotonic() < deadline:
+        try:
+            with urllib.request.urlopen(f"{endpoint}/v1/models", timeout=5):
+                return True
+        except Exception:
+            time.sleep(HEALTH_INTERVAL)
+    return False
+
+
+def _parse_host_port(ports_str: str) -> int | None:
+    """Extract the host port from a Ports string like '0.0.0.0:8001->8000/tcp'."""
+    m = re.search(r":(\d+)->8000", ports_str)
+    return int(m.group(1)) if m else None
+
+
+def _parse_labels(labels_str: str) -> dict[str, str]:
+    """Parse Docker's comma-separated 'key=value,key=value' label string."""
+    result: dict[str, str] = {}
+    for part in labels_str.split(","):
+        part = part.strip()
+        if "=" in part:
+            k, _, v = part.partition("=")
+            result[k.strip()] = v.strip()
+    return result
+
+
+def build_docker_run_cmd(config: RunConfig) -> list[str]:
+    model_path = config.model_path.resolve()
+    cmd = [
+        "docker",
+        "run",
+        "--rm",
+        "--name",
+        config.container_name,
+        "-p",
+        f"{config.port}:8000",
+        "-v",
+        f"{model_path}:/model:ro",
+        f"--label={MANAGED_LABEL}=true",
+        f"--label={MODEL_LABEL}={config.model_id}",
+        f"--label={MODEL_PATH_LABEL}={model_path}",
+    ]
+    if config.gpu:
+        cmd += ["--gpus", "all"]
+    cmd += [
+        VLLM_IMAGE,
+        "--model",
+        "/model",
+        "--served-model-name",
+        config.model_id,
+        "--dtype",
+        config.dtype,
+        "--host",
+        "0.0.0.0",
+        "--port",
+        "8000",
+    ]
+    if config.max_model_len is not None:
+        cmd += ["--max-model-len", str(config.max_model_len)]
+    cmd += config.extra_args
+    return cmd
+
+
+def start(config: RunConfig) -> None:
+    """Start a vLLM container (foreground, blocking)."""
+    model_path = config.model_path.resolve()
+    if not model_path.exists():
+        raise FileNotFoundError(f"Model path not found: {model_path}")
+    if _container_exists(config.container_name):
+        raise RuntimeError(
+            f"Container '{config.container_name}' already exists. "
+            "Stop it first or use a different --name."
+        )
+    subprocess.run(build_docker_run_cmd(config), check=True)
+
+
+def stop(name: str) -> None:
+    """Stop and remove a named container."""
+    if not _container_exists(name):
+        raise RuntimeError(f"No container named '{name}' found.")
+    _docker("stop", name)
+
+
+def stop_all() -> list[str]:
+    """Stop all vllmd-managed containers. Returns list of stopped names."""
+    containers = list_containers()
+    stopped = []
+    for c in containers:
+        _docker("stop", c["name"])
+        stopped.append(c["name"])
+    return stopped
+
+
+def list_containers() -> list[dict]:
+    """Return info for all running vllmd-managed containers."""
+    result = subprocess.run(
+        [
+            "docker",
+            "ps",
+            "--filter",
+            f"label={MANAGED_LABEL}=true",
+            "--format",
+            "{{json .}}",
+        ],
+        capture_output=True,
+        text=True,
+    )
+    containers = []
+    for line in result.stdout.strip().splitlines():
+        if not line:
+            continue
+        data = json.loads(line)
+        labels = _parse_labels(data.get("Labels", ""))
+        host_port = _parse_host_port(data.get("Ports", ""))
+        model_id = labels.get(MODEL_LABEL, "?")
+        model_path = labels.get(MODEL_PATH_LABEL, "?")
+        containers.append(
+            {
+                "name": data["Names"],
+                "model_id": model_id,
+                "model_path": model_path,
+                "port": host_port,
+                "endpoint": f"http://localhost:{host_port}" if host_port else "?",
+                "status": data.get("Status", "?"),
+            }
+        )
+    return containers
+
+
+def status(name: str) -> dict:
+    """Return container state and API health for a named container."""
+    result = subprocess.run(
+        ["docker", "inspect", name, "--format", "{{json .State}}"],
+        capture_output=True,
+        text=True,
+    )
+    if result.returncode != 0:
+        return {"running": False, "api_healthy": False, "container": None}
+
+    state = json.loads(result.stdout.strip())
+    running = state.get("Running", False)
+    api_healthy = False
+
+    if running:
+        port_result = subprocess.run(
+            ["docker", "port", name, "8000"],
+            capture_output=True,
+            text=True,
+        )
+        port = 8000
+        if port_result.returncode == 0:
+            binding = port_result.stdout.strip().split(":")[-1]
+            with contextlib.suppress(ValueError):
+                port = int(binding)
+        with contextlib.suppress(Exception):
+            url = f"http://localhost:{port}/v1/models"
+            with urllib.request.urlopen(url, timeout=3):
+                api_healthy = True
+
+    return {"running": running, "api_healthy": api_healthy, "container": state}
+
+
+def wait_ready(config: RunConfig) -> bool:
+    """Block until the vLLM API is reachable, or timeout."""
+    return _wait_ready(config.endpoint)
+
+
+def logs(name: str, follow: bool = False) -> None:
+    """Stream or print container logs."""
+    cmd = ["logs"]
+    if follow:
+        cmd.append("-f")
+    cmd.append(name)
+    _docker(*cmd)

src/vllmd/sessions/__init__.py

@@ -0,0 +1,5 @@
+"""Session management for persistent model conversations."""
+
+from .session import Message, Session
+
+__all__ = ["Message", "Session"]

src/vllmd/sessions/chat.py

@@ -0,0 +1,123 @@
+"""Chat function: assembles context and calls the vLLM completions endpoint."""
+
+from __future__ import annotations
+
+import json
+import urllib.request
+from pathlib import Path
+
+from ..vectordb.embeddings import embed
+from ..vectordb.store import COLLECTION_CODE, COLLECTION_DOCUMENTS, VectorStore
+from .session import Session
+
+MAX_HISTORY = 20
+N_CONTEXT_CHUNKS = 3
+MAX_TOKENS = 2048
+TIMEOUT = 120
+
+
+def chat(
+    session: Session,
+    user_message: str,
+    *,
+    max_history: int = MAX_HISTORY,
+    n_context: int = N_CONTEXT_CHUNKS,
+    max_tokens: int = MAX_TOKENS,
+) -> str:
+    """Send *user_message* within *session*, returning the assistant reply.
+
+    Retrieves semantic context from the session's vector store if an
+    embedding_model is configured. Falls back silently to history-only context
+    if the embedding endpoint is unavailable.
+    """
+    messages: list[dict] = []
+
+    if session.system_prompt:
+        messages.append({"role": "system", "content": session.system_prompt})
+
+    # Semantic context retrieval (skip entirely if no embedding model configured)
+    context_text = (
+        _retrieve_context(session, user_message, n_context)
+        if session.embedding_model
+        else ""
+    )
+    if context_text:
+        messages.append(
+            {"role": "system", "content": f"Relevant context:\n\n{context_text}"}
+        )
+
+    # Recent sequential history
+    for msg in session.messages[-max_history:]:
+        messages.append({"role": msg.role, "content": msg.content})
+
+    messages.append({"role": "user", "content": user_message})
+
+    response = _complete(session.endpoint, session.model_id, messages, max_tokens)
+
+    # Persist the exchange
+    from .session import Message
+
+    session.messages.append(Message(role="user", content=user_message))
+    session.messages.append(Message(role="assistant", content=response))
+
+    # Store embeddings in vector DB for future retrieval
+    if session.embedding_model:
+        _store_history(session, user_message, response)
+
+    return response
+
+
+def retrieve_context(session: Session, query: str, n: int = N_CONTEXT_CHUNKS) -> str:
+    """Public helper to show what context would be injected for *query*."""
+    return _retrieve_context(session, query, n)
+
+
+# ------------------------------------------------------------------
+# Internals
+# ------------------------------------------------------------------
+
+
+def _retrieve_context(session: Session, query: str, n: int) -> str:
+    if not session.embedding_model:
+        return ""
+    try:
+        store = VectorStore(Path(session.db_path))
+        query_vec = embed(session.endpoint, session.embedding_model, [query])[0]
+        chunks: list[str] = []
+        for collection in (COLLECTION_DOCUMENTS, COLLECTION_CODE):
+            results = store.search(query_vec, collection, n_results=n)
+            chunks.extend(r["content"] for r in results)
+        if not chunks:
+            return ""
+        return "\n\n---\n\n".join(chunks[:n])
+    except Exception:
+        return ""
+
+
+def _complete(
+    endpoint: str, model_id: str, messages: list[dict], max_tokens: int
+) -> str:
+    payload = json.dumps(
+        {"model": model_id, "messages": messages, "max_tokens": max_tokens}
+    ).encode()
+    req = urllib.request.Request(
+        f"{endpoint.rstrip('/')}/v1/chat/completions",
+        data=payload,
+        headers={"Content-Type": "application/json"},
+    )
+    with urllib.request.urlopen(req, timeout=TIMEOUT) as resp:
+        data = json.loads(resp.read())
+    return data["choices"][0]["message"]["content"]
+
+
+def _store_history(session: Session, user_message: str, response: str) -> None:
+    try:
+        store = VectorStore(Path(session.db_path))
+
+        def embedder(texts: list[str]) -> list[list[float]]:
+            return embed(session.endpoint, session.embedding_model, texts)
+
+        store.add_history(session.id, "user", user_message, embedder)
+        store.add_history(session.id, "assistant", response, embedder)
+    except Exception:
+        pass