synarmo 0.1.1__py3-none-any.whl

synarmo/__init__.py

@@ -0,0 +1,11 @@
+from synarmo.config import SynarmoConfig
+from synarmo.engine import SynarmoEngine, predict, suggest
+from synarmo.suggestions import Suggestion
+
+__all__ = [
+    "Suggestion",
+    "SynarmoConfig",
+    "SynarmoEngine",
+    "predict",
+    "suggest",
+]

synarmo/autocomplete_eval.py

@@ -0,0 +1,140 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any
+
+
+@dataclass(frozen=True, slots=True)
+class LogprobToken:
+    text: str
+    logprob: float
+
+
+@dataclass(frozen=True, slots=True)
+class AutocompleteCandidate:
+    text: str
+    starter: str
+    rest: str
+    logprob: float
+
+
+@dataclass(frozen=True, slots=True)
+class AutocompleteEvaluation:
+    context: str
+    prompt: str
+    candidates: list[AutocompleteCandidate]
+    top_tokens: list[LogprobToken] = field(default_factory=list)
+
+
+def build_autocomplete_prompt(context: str, typed_text: str) -> str:
+    context_lines = [
+        line
+        for line in context.rstrip().splitlines()
+        if not line.lower().startswith("current typed text:")
+    ]
+    context_block = "\n".join(context_lines).strip()
+    if context_block:
+        context_block = f"Context:\n{context_block}\n\n"
+
+    return f"""{context_block}Continue the message with only the next few words.
+Do not repeat labels or earlier text.
+Message:
+{typed_text}"""
+
+
+def normalize_candidate(text: str, *, max_words: int) -> str:
+    candidate = text.replace("\r", "").replace("\n", " ").strip()
+    candidate = candidate.removeprefix(":").strip()
+    candidate = candidate.strip("\"'` ")
+    words = candidate.split()
+    if max_words > 0:
+        words = words[:max_words]
+    return " ".join(words).strip(" ,;:")
+
+
+def extract_top_logprobs(probe: dict[str, Any]) -> dict[str, float]:
+    choices = probe.get("choices")
+    if not choices:
+        return {}
+
+    logprobs = choices[0].get("logprobs")
+    if not logprobs:
+        return {}
+
+    top_logprobs = logprobs.get("top_logprobs")
+    if not top_logprobs:
+        return {}
+
+    raw_top_logprobs = top_logprobs[0]
+    if not isinstance(raw_top_logprobs, dict):
+        return {}
+
+    return raw_top_logprobs
+
+
+def evaluate_with_llama(
+    llm: Any,
+    *,
+    context: str,
+    typed_text: str,
+    choices: int = 3,
+    max_tokens: int = 5,
+    max_words: int = 1,
+    temperature: float = 0.5,
+    top_p: float = 0.95,
+    logprob_pool: int = 24,
+) -> AutocompleteEvaluation:
+    prompt = build_autocomplete_prompt(context, typed_text)
+    probe = llm(
+        prompt=prompt,
+        max_tokens=1,
+        temperature=temperature,
+        top_p=top_p,
+        logprobs=logprob_pool,
+        echo=False,
+    )
+    raw_top_logprobs = extract_top_logprobs(probe)
+    ranked = sorted(raw_top_logprobs.items(), key=lambda item: item[1], reverse=True)
+    top_tokens = [LogprobToken(text=text, logprob=float(logprob)) for text, logprob in ranked]
+
+    seen_words: set[str] = set()
+    starters: list[LogprobToken] = []
+    for token_text, logprob in ranked:
+        stripped = token_text.strip()
+        if not stripped:
+            continue
+        first_word = stripped.split()[0].lower()
+        if first_word in seen_words:
+            continue
+        seen_words.add(first_word)
+        starters.append(LogprobToken(text=token_text, logprob=float(logprob)))
+        if len(starters) >= choices:
+            break
+
+    candidates: list[AutocompleteCandidate] = []
+    for starter in starters:
+        response = llm(
+            prompt=prompt + starter.text,
+            max_tokens=max(max_tokens - 1, 1),
+            temperature=0.0,
+            top_p=1.0,
+            stop=["\n", ".", "!", "?"],
+            echo=False,
+        )
+        rest = str(response["choices"][0]["text"])
+        candidate = normalize_candidate(starter.text + rest, max_words=max_words)
+        candidates.append(
+            AutocompleteCandidate(
+                text=candidate,
+                starter=starter.text,
+                rest=rest,
+                logprob=starter.logprob,
+            )
+        )
+
+    return AutocompleteEvaluation(
+        context=context,
+        prompt=prompt,
+        candidates=candidates,
+        top_tokens=top_tokens,
+    )

synarmo/cli.py

@@ -0,0 +1,147 @@
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+from synarmo.engine import SynarmoEngine
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(prog="synarmo")
+    subcommands = parser.add_subparsers(dest="command", required=True)
+
+    suggest_parser = subcommands.add_parser("suggest", help="Generate local type-ahead suggestions.")
+    suggest_parser.add_argument("text")
+    suggest_parser.add_argument("--context", default="")
+    suggest_parser.add_argument("--profile", default="default")
+    suggest_parser.add_argument("--backend", choices=["mock", "llama-cpp"], default="mock")
+    suggest_parser.add_argument("--model-path", type=Path)
+    suggest_parser.add_argument("--max-suggestions", type=int)
+
+    model_ensure_parser = subcommands.add_parser(
+        "model-ensure",
+        help="Download or verify the configured local inference model.",
+    )
+    model_ensure_parser.add_argument("--profile", default="default")
+    model_ensure_parser.add_argument("--backend", choices=["mock", "llama-cpp"], default="llama-cpp")
+    model_ensure_parser.add_argument("--model-path", type=Path)
+
+    compose_parser = subcommands.add_parser(
+        "compose",
+        help="Interactively choose suggestions and continue predicting.",
+    )
+    compose_parser.add_argument("text")
+    compose_parser.add_argument("--context", default="")
+    compose_parser.add_argument("--profile", default="default")
+    compose_parser.add_argument("--backend", choices=["mock", "llama-cpp"], default="mock")
+    compose_parser.add_argument("--model-path", type=Path)
+    compose_parser.add_argument("--max-suggestions", type=int)
+
+    serve_parser = subcommands.add_parser("serve", help="Run the local REST/WebSocket service.")
+    serve_parser.add_argument(
+        "--host",
+        default="127.0.0.1",
+        help="Bind address for the service. Use 0.0.0.0 to listen on all interfaces.",
+    )
+    serve_parser.add_argument("--port", type=int, default=8765)
+    serve_parser.add_argument("--profile", default="default")
+    serve_parser.add_argument("--backend", choices=["mock", "llama-cpp"], default="mock")
+    serve_parser.add_argument("--model-path", type=Path)
+    serve_parser.add_argument("--max-suggestions", type=int)
+    return parser
+
+
+def app() -> None:
+    parser = build_parser()
+    args = parser.parse_args()
+
+    if args.command == "suggest":
+        engine = _load_engine(args)
+        for item in engine.suggest(text=args.text, context=args.context):
+            print(item.text)
+        return
+
+    if args.command == "model-ensure":
+        engine = _load_engine(args)
+        print(_format_runtime_diagnostics(engine))
+        print(f"Model ready for {args.backend}")
+        return
+
+    if args.command == "compose":
+        engine = _load_engine(args)
+        _compose(engine, text=args.text, context=args.context)
+        return
+
+    if args.command == "serve":
+        try:
+            import uvicorn
+        except ImportError as exc:
+            raise SystemExit("Install service extras first: pip install synarmo[service]") from exc
+
+        from synarmo.service.app import create_app
+
+        engine = _load_engine(args)
+        print(_format_runtime_diagnostics(engine))
+        uvicorn.run(create_app(engine), host=args.host, port=args.port)
+
+
+def _load_engine(args: argparse.Namespace) -> SynarmoEngine:
+    options: dict[str, object] = {}
+    max_suggestions = getattr(args, "max_suggestions", None)
+    if max_suggestions is not None:
+        options["max_suggestions"] = max_suggestions
+
+    return SynarmoEngine.load(
+        profile=args.profile,
+        backend=args.backend,
+        model_path=args.model_path,
+        **options,
+    )
+
+
+def _format_runtime_diagnostics(engine: SynarmoEngine) -> str:
+    diagnostics = engine.runtime_diagnostics()
+    parts = [
+        f"backend={diagnostics.get('backend', '')}",
+        f"model={diagnostics.get('model', '')}",
+        f"n_gpu_layers={diagnostics.get('n_gpu_layers', '')}",
+    ]
+    if "requested_gpu_layers" in diagnostics:
+        parts.append(f"requested_gpu_layers={diagnostics['requested_gpu_layers']}")
+    if "model_layers" in diagnostics:
+        parts.append(f"model_layers={diagnostics['model_layers']}")
+    if "gpu_offload_supported" in diagnostics:
+        parts.append(f"gpu_offload_supported={diagnostics['gpu_offload_supported']}")
+    if "llama_verbose" in diagnostics:
+        parts.append(f"llama_verbose={diagnostics['llama_verbose']}")
+    return "Synarmo runtime: " + " ".join(parts)
+
+
+def _compose(engine: SynarmoEngine, *, text: str, context: str) -> None:
+    current = text
+    while True:
+        print(f"\n{current}")
+        suggestions = engine.suggest(text=current, context=context)
+        if not suggestions:
+            print("No suggestions.")
+            return
+
+        for index, item in enumerate(suggestions, start=1):
+            print(f"{index}. {item.text}")
+
+        choice = input(
+            f"Choose 1-{len(suggestions)}, enter custom text, or q to quit: "
+        ).strip()
+        if choice.lower() in {"q", "quit", "exit"}:
+            return
+        if not choice:
+            continue
+        if choice.isdigit() and 1 <= int(choice) <= len(suggestions):
+            selected = suggestions[int(choice) - 1].text
+        else:
+            selected = choice
+        current = f"{current.rstrip()} {selected.strip()}".strip()
+
+
+if __name__ == "__main__":
+    app()

synarmo/config.py

@@ -0,0 +1,169 @@
+from __future__ import annotations
+
+import os
+from dataclasses import dataclass, field, replace
+from pathlib import Path
+from typing import Literal
+
+BackendName = Literal["mock", "llama-cpp"]
+
+DEFAULT_MODELS_CACHE = Path("~/models/synarmo")
+ENV_FILE = ".env"
+LOCAL_MODELS_CACHE_ENV = "LOCAL_MODELS_CACHE"
+MODEL_ENV = "SYNARMO_MODEL"
+MODEL_PATH_ENV = "SYNARMO_MODEL_PATH"
+MODEL_REPO_ID_ENV = "SYNARMO_MODEL_REPO_ID"
+MAX_SUGGESTIONS_ENV = "SYNARMO_MAX_SUGGESTIONS"
+MAX_TOKENS_ENV = "SYNARMO_MAX_TOKENS"
+MAX_SUGGESTION_WORDS_ENV = "SYNARMO_MAX_SUGGESTION_WORDS"
+TEMPERATURE_ENV = "SYNARMO_TEMPERATURE"
+TOP_P_ENV = "SYNARMO_TOP_P"
+LOGPROB_POOL_ENV = "SYNARMO_LOGPROB_POOL"
+N_GPU_LAYERS_ENV = "SYNARMO_N_GPU_LAYERS"
+LLAMA_VERBOSE_ENV = "SYNARMO_LLAMA_VERBOSE"
+
+
+def load_env_file(path: str | Path = ENV_FILE) -> None:
+    env_path = Path(path)
+    if not env_path.exists():
+        return
+
+    for raw_line in env_path.read_text(encoding="utf-8").splitlines():
+        line = raw_line.strip()
+        if not line or line.startswith("#") or "=" not in line:
+            continue
+
+        key, value = line.split("=", 1)
+        key = key.strip()
+        value = value.strip().strip("'\"")
+        if key and key not in os.environ:
+            os.environ[key] = value
+
+
+def configured_models_cache() -> Path:
+    value = os.getenv(LOCAL_MODELS_CACHE_ENV)
+    return Path(value).expanduser() if value else DEFAULT_MODELS_CACHE.expanduser()
+
+
+def configured_model_path(model_path: str | Path | None = None) -> Path | None:
+    if model_path:
+        return Path(model_path).expanduser()
+
+    value = _env_model()
+    if value is None:
+        return None
+
+    value = value.expanduser()
+    if value.is_absolute():
+        return value
+    return configured_models_cache() / value
+
+
+def configured_model_filename(model_path: str | Path | None = None) -> str | None:
+    if model_path:
+        return Path(model_path).name
+
+    value = _env_model()
+    return value.name if value else None
+
+
+def configured_model_repo_id() -> str | None:
+    value = os.getenv(MODEL_REPO_ID_ENV)
+    return value.strip() if value and value.strip() else None
+
+
+def _env_model() -> Path | None:
+    value = os.getenv(MODEL_ENV) or os.getenv(MODEL_PATH_ENV)
+    return Path(value) if value else None
+
+
+def configured_max_suggestions() -> int:
+    value = os.getenv(MAX_SUGGESTIONS_ENV)
+    return int(value) if value else 3
+
+
+def configured_max_tokens() -> int:
+    value = os.getenv(MAX_TOKENS_ENV)
+    return int(value) if value else 5
+
+
+def configured_max_suggestion_words() -> int:
+    value = os.getenv(MAX_SUGGESTION_WORDS_ENV)
+    return int(value) if value else 4
+
+
+def configured_temperature() -> float:
+    value = os.getenv(TEMPERATURE_ENV)
+    return float(value) if value else 0.25
+
+
+def configured_top_p() -> float:
+    value = os.getenv(TOP_P_ENV)
+    return float(value) if value else 0.95
+
+
+def configured_logprob_pool() -> int:
+    value = os.getenv(LOGPROB_POOL_ENV)
+    return int(value) if value else 24
+
+
+def configured_n_gpu_layers() -> int:
+    value = os.getenv(N_GPU_LAYERS_ENV)
+    return int(value) if value else 0
+
+
+def configured_llama_verbose() -> bool:
+    value = os.getenv(LLAMA_VERBOSE_ENV)
+    return value is not None and value.strip().lower() in {"1", "true", "yes", "on"}
+
+
+@dataclass(slots=True)
+class SynarmoConfig:
+    backend: BackendName = "mock"
+    model_path: Path | None = None
+    model_repo_id: str | None = None
+    model_filename: str | None = None
+    models_cache_dir: Path = field(default_factory=configured_models_cache)
+    profile: str = "default"
+    max_suggestions: int = field(default_factory=configured_max_suggestions)
+    max_latency_ms: int = 100
+    context_window: int = 2048
+    style_adaptation: bool = True
+    temperature: float = field(default_factory=configured_temperature)
+    top_p: float = field(default_factory=configured_top_p)
+    max_tokens: int = field(default_factory=configured_max_tokens)
+    max_suggestion_words: int = field(default_factory=configured_max_suggestion_words)
+    logprob_pool: int = field(default_factory=configured_logprob_pool)
+    n_gpu_layers: int = field(default_factory=configured_n_gpu_layers)
+    llama_verbose: bool = field(default_factory=configured_llama_verbose)
+    stop: list[str] = field(default_factory=lambda: ["\n\n"])
+    profiles_dir: Path = Path("profiles")
+
+    def __post_init__(self) -> None:
+        self.models_cache_dir = self.models_cache_dir.expanduser()
+        if self.model_path is not None:
+            self.model_path = self.model_path.expanduser()
+        if not 1 <= self.max_suggestions <= 10:
+            raise ValueError("max_suggestions must be between 1 and 10")
+        if self.max_latency_ms < 1:
+            raise ValueError("max_latency_ms must be positive")
+        if self.context_window < 128:
+            raise ValueError("context_window must be at least 128")
+        if not 0.0 <= self.temperature <= 2.0:
+            raise ValueError("temperature must be between 0.0 and 2.0")
+        if not 0.0 < self.top_p <= 1.0:
+            raise ValueError("top_p must be greater than 0.0 and at most 1.0")
+        if not 1 <= self.max_tokens <= 128:
+            raise ValueError("max_tokens must be between 1 and 128")
+        if not 1 <= self.max_suggestion_words <= 20:
+            raise ValueError("max_suggestion_words must be between 1 and 20")
+        if not 1 <= self.logprob_pool <= 50:
+            raise ValueError("logprob_pool must be between 1 and 50")
+        if self.n_gpu_layers < -1:
+            raise ValueError("n_gpu_layers must be -1 or greater")
+
+    def resolved_profile_dir(self) -> Path:
+        return self.profiles_dir.expanduser().resolve() / self.profile
+
+    def copy_with(self, **updates: object) -> "SynarmoConfig":
+        return replace(self, **updates)

synarmo/context.py

@@ -0,0 +1,23 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from synarmo.memory import UserMemory
+
+
+@dataclass(slots=True)
+class ContextAssembler:
+    max_chars: int = 4000
+
+    def assemble(self, *, text: str, context: str | None, memory: UserMemory) -> str:
+        parts: list[str] = []
+        if memory.style_summary:
+            parts.append(f"User style: {memory.style_summary}")
+        if memory.preferences:
+            prefs = "; ".join(f"{key}: {value}" for key, value in sorted(memory.preferences.items()))
+            parts.append(f"Known preferences: {prefs}")
+        if context:
+            parts.append(f"Current context: {context.strip()}")
+        parts.append(f"Current typed text: {text.strip()}")
+        assembled = "\n".join(part for part in parts if part)
+        return assembled[-self.max_chars :]