oshell 0.1.1__py3-none-any.whl

oshell/providers.py

@@ -0,0 +1,275 @@
+"""LLM provider layer.
+
+Both providers expose a single ``stream_chat`` method that yields text
+chunks as they arrive, so the CLI can render tokens live.
+"""
+
+from __future__ import annotations
+
+import json
+from typing import Iterator, List, Protocol
+
+import httpx
+
+from .config import Config
+
+Message = dict  # {"role": "user" | "assistant" | "system", "content": str}
+
+
+class Provider(Protocol):
+    """Common interface every LLM backend implements."""
+
+    name: str
+
+    def stream_chat(self, messages: List[Message]) -> Iterator[str]:
+        """Yield response text chunks for the given conversation."""
+        ...
+
+    def list_models(self) -> List[str]:
+        """Return available model names (best effort)."""
+        ...
+
+
+class OllamaProvider:
+    """Talks to a local Ollama server (https://ollama.com)."""
+
+    name = "ollama"
+
+    def __init__(self, cfg: Config) -> None:
+        self.host = cfg.ollama_host.rstrip("/")
+        self.model = cfg.model
+        self.temperature = cfg.temperature
+
+    def stream_chat(self, messages: List[Message]) -> Iterator[str]:
+        payload = {
+            "model": self.model,
+            "messages": messages,
+            "stream": True,
+            "options": {"temperature": self.temperature},
+        }
+        with httpx.stream(
+            "POST", f"{self.host}/api/chat", json=payload, timeout=None
+        ) as resp:
+            resp.raise_for_status()
+            for line in resp.iter_lines():
+                if not line:
+                    continue
+                data = json.loads(line)
+                if data.get("done"):
+                    break
+                chunk = data.get("message", {}).get("content", "")
+                if chunk:
+                    yield chunk
+
+    def list_models(self) -> List[str]:
+        try:
+            resp = httpx.get(f"{self.host}/api/tags", timeout=10)
+            resp.raise_for_status()
+            return [m["name"] for m in resp.json().get("models", [])]
+        except httpx.HTTPError:
+            return []
+
+
+class OpenAICompatibleProvider:
+    """Base for any service speaking the OpenAI Chat Completions API.
+
+    OpenAI, Groq, and Gemini (via its OpenAI-compatible endpoint) all share
+    this exact wire format, so they differ only in base URL, key, and the
+    error message shown when the key is missing.
+    """
+
+    name = "openai"
+    _key_hint = "OPENAI_API_KEY"
+
+    def __init__(self, base_url: str, api_key: str, model: str, temperature: float):
+        if not api_key:
+            raise ValueError(
+                f"{self.name.title()} API key not set. Export {self._key_hint} or run "
+                f"`oshell config set {self.name}_api_key <key>`."
+            )
+        self.base_url = base_url.rstrip("/")
+        self.api_key = api_key
+        self.model = model
+        self.temperature = temperature
+
+    def stream_chat(self, messages: List[Message]) -> Iterator[str]:
+        headers = {"Authorization": f"Bearer {self.api_key}"}
+        payload = {
+            "model": self.model,
+            "messages": messages,
+            "stream": True,
+            "temperature": self.temperature,
+        }
+        with httpx.stream(
+            "POST",
+            f"{self.base_url}/chat/completions",
+            json=payload,
+            headers=headers,
+            timeout=None,
+        ) as resp:
+            resp.raise_for_status()
+            for line in resp.iter_lines():
+                if not line or not line.startswith("data: "):
+                    continue
+                data_str = line[len("data: ") :]
+                if data_str.strip() == "[DONE]":
+                    break
+                try:
+                    data = json.loads(data_str)
+                except json.JSONDecodeError:
+                    continue
+                choices = data.get("choices") or []
+                if not choices:
+                    continue
+                chunk = choices[0].get("delta", {}).get("content")
+                if chunk:
+                    yield chunk
+
+    def list_models(self) -> List[str]:
+        headers = {"Authorization": f"Bearer {self.api_key}"}
+        try:
+            resp = httpx.get(
+                f"{self.base_url}/models", headers=headers, timeout=10
+            )
+            resp.raise_for_status()
+            return sorted(m["id"] for m in resp.json().get("data", []))
+        except httpx.HTTPError:
+            return []
+
+
+class OpenAIProvider(OpenAICompatibleProvider):
+    """OpenAI Chat Completions API."""
+
+    name = "openai"
+    _key_hint = "OPENAI_API_KEY"
+
+    def __init__(self, cfg: Config) -> None:
+        super().__init__(
+            cfg.openai_base_url, cfg.openai_api_key, cfg.model, cfg.temperature
+        )
+
+
+class GroqProvider(OpenAICompatibleProvider):
+    """Groq — extremely fast inference, OpenAI-compatible API."""
+
+    name = "groq"
+    _key_hint = "GROQ_API_KEY"
+
+    def __init__(self, cfg: Config) -> None:
+        model = cfg.model if cfg.model else "llama-3.3-70b-versatile"
+        super().__init__(
+            "https://api.groq.com/openai/v1",
+            cfg.groq_api_key,
+            model,
+            cfg.temperature,
+        )
+
+
+class GeminiProvider(OpenAICompatibleProvider):
+    """Google Gemini via its OpenAI-compatible endpoint."""
+
+    name = "gemini"
+    _key_hint = "GEMINI_API_KEY"
+
+    def __init__(self, cfg: Config) -> None:
+        model = cfg.model if cfg.model else "gemini-1.5-flash"
+        super().__init__(
+            "https://generativelanguage.googleapis.com/v1beta/openai",
+            cfg.gemini_api_key,
+            model,
+            cfg.temperature,
+        )
+
+
+class AnthropicProvider:
+    """Anthropic Claude via the native Messages API (SSE streaming)."""
+
+    name = "anthropic"
+
+    def __init__(self, cfg: Config) -> None:
+        if not cfg.anthropic_api_key:
+            raise ValueError(
+                "Anthropic API key not set. Export ANTHROPIC_API_KEY or run "
+                "`oshell config set anthropic_api_key <key>`."
+            )
+        self.api_key = cfg.anthropic_api_key
+        self.model = cfg.model if cfg.model else "claude-3-5-sonnet-latest"
+        self.temperature = cfg.temperature
+
+    @staticmethod
+    def _split(messages: List[Message]):
+        """Anthropic takes the system prompt separately from the turns."""
+        system = ""
+        turns: List[Message] = []
+        for m in messages:
+            if m["role"] == "system":
+                system = m["content"]
+            else:
+                turns.append({"role": m["role"], "content": m["content"]})
+        return system, turns
+
+    def stream_chat(self, messages: List[Message]) -> Iterator[str]:
+        system, turns = self._split(messages)
+        headers = {
+            "x-api-key": self.api_key,
+            "anthropic-version": "2023-06-01",
+            "content-type": "application/json",
+        }
+        payload = {
+            "model": self.model,
+            "messages": turns,
+            "max_tokens": 4096,
+            "temperature": self.temperature,
+            "stream": True,
+        }
+        if system:
+            payload["system"] = system
+        with httpx.stream(
+            "POST",
+            "https://api.anthropic.com/v1/messages",
+            json=payload,
+            headers=headers,
+            timeout=None,
+        ) as resp:
+            resp.raise_for_status()
+            for line in resp.iter_lines():
+                if not line or not line.startswith("data: "):
+                    continue
+                try:
+                    data = json.loads(line[len("data: ") :])
+                except json.JSONDecodeError:
+                    continue
+                if data.get("type") == "content_block_delta":
+                    chunk = data.get("delta", {}).get("text")
+                    if chunk:
+                        yield chunk
+
+    def list_models(self) -> List[str]:
+        headers = {
+            "x-api-key": self.api_key,
+            "anthropic-version": "2023-06-01",
+        }
+        try:
+            resp = httpx.get(
+                "https://api.anthropic.com/v1/models", headers=headers, timeout=10
+            )
+            resp.raise_for_status()
+            return sorted(m["id"] for m in resp.json().get("data", []))
+        except httpx.HTTPError:
+            return []
+
+
+def get_provider(cfg: Config) -> Provider:
+    """Factory that returns the configured provider instance."""
+    providers = {
+        "openai": OpenAIProvider,
+        "ollama": OllamaProvider,
+        "anthropic": AnthropicProvider,
+        "groq": GroqProvider,
+        "gemini": GeminiProvider,
+    }
+    cls = providers.get(cfg.provider)
+    if cls is None:
+        valid = ", ".join(sorted(providers))
+        raise ValueError(f"Unknown provider: {cfg.provider!r} (use one of: {valid})")
+    return cls(cfg)

oshell/retrieval.py

@@ -0,0 +1,180 @@
+"""Chat-with-your-files: a tiny local RAG index.
+
+Embeddings come from Ollama (default ``nomic-embed-text``) or OpenAI, so we
+keep the "no heavy dependencies" promise — vectors are stored as JSON and
+similarity search is plain Python. Good enough for personal knowledge bases.
+"""
+
+from __future__ import annotations
+
+import json
+import math
+import time
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Iterable, List, Tuple
+
+import httpx
+from platformdirs import user_data_dir
+
+from .config import APP_NAME, Config
+
+INDEX_DIR = Path(user_data_dir(APP_NAME)) / "index"
+INDEX_PATH = INDEX_DIR / "store.json"
+
+TEXT_SUFFIXES = {
+    ".txt", ".md", ".markdown", ".rst", ".py", ".js", ".ts", ".tsx", ".jsx",
+    ".json", ".yaml", ".yml", ".toml", ".cfg", ".ini", ".html", ".css",
+    ".java", ".go", ".rs", ".rb", ".php", ".c", ".cpp", ".h", ".hpp", ".sh",
+    ".sql", ".csv",
+}
+MAX_FILE_BYTES = 500_000
+CHUNK_CHARS = 1200
+CHUNK_OVERLAP = 150
+
+
+@dataclass
+class Chunk:
+    source: str
+    text: str
+    vector: List[float]
+
+
+def _chunk_text(text: str) -> Iterable[str]:
+    step = CHUNK_CHARS - CHUNK_OVERLAP
+    for start in range(0, len(text), step):
+        piece = text[start : start + CHUNK_CHARS].strip()
+        if piece:
+            yield piece
+
+
+def _cosine(a: List[float], b: List[float]) -> float:
+    dot = sum(x * y for x, y in zip(a, b))
+    na = math.sqrt(sum(x * x for x in a))
+    nb = math.sqrt(sum(y * y for y in b))
+    if na == 0 or nb == 0:
+        return 0.0
+    return dot / (na * nb)
+
+
+class Embedder:
+    """Embeds text via Ollama or OpenAI, matching the configured provider."""
+
+    def __init__(self, cfg: Config) -> None:
+        self.cfg = cfg
+        # Use OpenAI embeddings only when explicitly on OpenAI with a key.
+        self.use_openai = cfg.provider == "openai" and bool(cfg.openai_api_key)
+
+    def embed(self, text: str) -> List[float]:
+        if self.use_openai:
+            return self._embed_openai(text)
+        return self._embed_ollama(text)
+
+    def _embed_ollama(self, text: str) -> List[float]:
+        host = self.cfg.ollama_host.rstrip("/")
+        # Newer Ollama exposes /api/embed; older ones use /api/embeddings.
+        try:
+            resp = httpx.post(
+                f"{host}/api/embed",
+                json={"model": self.cfg.embed_model, "input": text},
+                timeout=120,
+            )
+            if resp.status_code == 404:
+                raise httpx.HTTPStatusError("no /api/embed", request=resp.request, response=resp)
+            resp.raise_for_status()
+            data = resp.json()
+            embeddings = data.get("embeddings")
+            if embeddings:
+                return embeddings[0]
+            if data.get("embedding"):
+                return data["embedding"]
+        except httpx.HTTPStatusError:
+            pass
+
+        resp = httpx.post(
+            f"{host}/api/embeddings",
+            json={"model": self.cfg.embed_model, "prompt": text},
+            timeout=120,
+        )
+        resp.raise_for_status()
+        return resp.json()["embedding"]
+
+    def _embed_openai(self, text: str) -> List[float]:
+        resp = httpx.post(
+            f"{self.cfg.openai_base_url.rstrip('/')}/embeddings",
+            json={"model": "text-embedding-3-small", "input": text},
+            headers={"Authorization": f"Bearer {self.cfg.openai_api_key}"},
+            timeout=120,
+        )
+        resp.raise_for_status()
+        return resp.json()["data"][0]["embedding"]
+
+
+def _iter_files(root: Path) -> Iterable[Path]:
+    for path in sorted(root.rglob("*")):
+        if not path.is_file():
+            continue
+        if path.suffix.lower() not in TEXT_SUFFIXES:
+            continue
+        if any(part in {".git", "node_modules", "__pycache__", ".venv"} for part in path.parts):
+            continue
+        try:
+            if path.stat().st_size > MAX_FILE_BYTES:
+                continue
+        except OSError:
+            continue
+        yield path
+
+
+def build_index(root: Path, cfg: Config, progress=None) -> int:
+    """Index every text file under ``root``. Returns the chunk count."""
+    embedder = Embedder(cfg)
+    chunks: List[dict] = []
+    for path in _iter_files(root):
+        try:
+            text = path.read_text(encoding="utf-8", errors="ignore")
+        except OSError:
+            continue
+        rel = str(path.relative_to(root))
+        for piece in _chunk_text(text):
+            vector = embedder.embed(piece)
+            chunks.append({"source": rel, "text": piece, "vector": vector})
+            if progress:
+                progress(rel)
+
+    INDEX_DIR.mkdir(parents=True, exist_ok=True)
+    INDEX_PATH.write_text(
+        json.dumps(
+            {
+                "root": str(root.resolve()),
+                "created": time.time(),
+                "provider": "openai" if embedder.use_openai else "ollama",
+                "chunks": chunks,
+            }
+        ),
+        encoding="utf-8",
+    )
+    return len(chunks)
+
+
+def index_exists() -> bool:
+    return INDEX_PATH.exists()
+
+
+def search(query: str, cfg: Config, top_k: int = 5) -> List[Tuple[str, str, float]]:
+    """Return the top_k (source, text, score) matches for ``query``."""
+    if not INDEX_PATH.exists():
+        raise FileNotFoundError(
+            "No index found. Run `oshell index <folder>` first."
+        )
+    store = json.loads(INDEX_PATH.read_text(encoding="utf-8"))
+    chunks = store.get("chunks", [])
+    if not chunks:
+        return []
+
+    qvec = Embedder(cfg).embed(query)
+    scored = [
+        (c["source"], c["text"], _cosine(qvec, c["vector"])) for c in chunks
+    ]
+    scored.sort(key=lambda x: x[2], reverse=True)
+    return scored[:top_k]

oshell/storyboard.py

@@ -0,0 +1,163 @@
+"""Storyboard mode: brief → planned scenes → image per scene → stitched video.
+
+The LLM acts as a director: it breaks a one-line brief into a sequence of
+scenes (each with an image prompt and a caption). We generate an image for
+every scene, then stitch them into an MP4 slideshow with ffmpeg.
+
+ffmpeg is only required for the final stitch — the per-scene images are saved
+regardless, so the command still produces useful output if ffmpeg is missing.
+"""
+
+from __future__ import annotations
+
+import json
+import re
+import shutil
+import subprocess
+import time
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Callable, List, Optional
+
+from .config import APP_NAME, Config
+from .media import ImageGenerator, _output_dir, _slug
+from .providers import Message, Provider
+
+_DIRECTOR_SYSTEM = """\
+You are a creative director planning a short visual story. Given a brief, break \
+it into a sequence of distinct scenes. Respond with ONLY a JSON array; each \
+element is an object: {{"prompt": "<detailed text-to-image prompt with \
+composition, lighting, style and mood>", "caption": "<short on-screen \
+caption>"}}. Produce exactly {n} scenes. No markdown, no code fences, no extra \
+text."""
+
+
+@dataclass
+class Scene:
+    prompt: str
+    caption: str
+    image: Optional[Path] = None
+
+
+@dataclass
+class StoryboardResult:
+    scenes: List[Scene]
+    video: Optional[Path]
+    image_dir: Path
+
+
+def _parse_scenes(raw: str, fallback_brief: str, n: int) -> List[Scene]:
+    text = raw.strip()
+    if text.startswith("```"):
+        text = re.sub(r"^```[a-zA-Z]*\n?", "", text)
+        text = re.sub(r"\n?```$", "", text).strip()
+    # Pull out the first JSON array if the model added prose.
+    start = text.find("[")
+    end = text.rfind("]")
+    if start != -1 and end != -1:
+        text = text[start : end + 1]
+    try:
+        data = json.loads(text)
+    except json.JSONDecodeError:
+        data = []
+
+    scenes: List[Scene] = []
+    for item in data:
+        if isinstance(item, dict) and item.get("prompt"):
+            scenes.append(
+                Scene(prompt=str(item["prompt"]), caption=str(item.get("caption", "")))
+            )
+    # Fallback: if parsing failed, make simple scenes from the brief.
+    if not scenes:
+        scenes = [
+            Scene(prompt=f"{fallback_brief}, scene {i + 1}", caption="")
+            for i in range(n)
+        ]
+    return scenes
+
+
+class StoryboardAgent:
+    """Plans scenes with the LLM and renders them into a video."""
+
+    def __init__(self, provider: Provider, image_gen: ImageGenerator, cfg: Config):
+        self.provider = provider
+        self.image_gen = image_gen
+        self.cfg = cfg
+
+    def plan(self, brief: str, n: int) -> List[Scene]:
+        system = _DIRECTOR_SYSTEM.format(n=n)
+        messages: List[Message] = [
+            {"role": "system", "content": system},
+            {"role": "user", "content": brief},
+        ]
+        raw = "".join(self.provider.stream_chat(messages))
+        return _parse_scenes(raw, brief, n)
+
+    def render(
+        self,
+        brief: str,
+        n: int = 4,
+        seconds_per_scene: float = 2.5,
+        on_scene: Optional[Callable[[int, Scene], None]] = None,
+    ) -> StoryboardResult:
+        scenes = self.plan(brief, n)
+
+        out_dir = _output_dir(self.cfg, "storyboards")
+        for i, scene in enumerate(scenes):
+            result = self.image_gen.generate(scene.prompt, n=1)
+            # Move/rename the generated image into the storyboard folder, ordered.
+            src = result.paths[0]
+            dest = out_dir / f"scene-{i + 1:02d}.png"
+            dest.write_bytes(src.read_bytes())
+            scene.image = dest
+            if on_scene:
+                on_scene(i, scene)
+
+        video = self._stitch(scenes, out_dir, seconds_per_scene)
+        return StoryboardResult(scenes=scenes, video=video, image_dir=out_dir)
+
+    @staticmethod
+    def ffmpeg_available() -> bool:
+        return shutil.which("ffmpeg") is not None
+
+    def _stitch(
+        self, scenes: List[Scene], out_dir: Path, seconds_per_scene: float
+    ) -> Optional[Path]:
+        images = [s.image for s in scenes if s.image]
+        if not images or not self.ffmpeg_available():
+            return None
+
+        # Build a concat demuxer list with a per-image duration.
+        list_path = out_dir / "scenes.txt"
+        lines = []
+        for img in images:
+            lines.append(f"file '{img.as_posix()}'")
+            lines.append(f"duration {seconds_per_scene}")
+        # The concat demuxer needs the last file repeated (no trailing duration).
+        lines.append(f"file '{images[-1].as_posix()}'")
+        list_path.write_text("\n".join(lines), encoding="utf-8")
+
+        video_path = out_dir / "storyboard.mp4"
+        cmd = [
+            "ffmpeg",
+            "-y",
+            "-f",
+            "concat",
+            "-safe",
+            "0",
+            "-i",
+            str(list_path),
+            "-vf",
+            "scale=1024:-2:force_original_aspect_ratio=decrease,"
+            "pad=1024:1024:(ow-iw)/2:(oh-ih)/2,format=yuv420p",
+            "-r",
+            "30",
+            str(video_path),
+        ]
+        try:
+            proc = subprocess.run(
+                cmd, capture_output=True, text=True, timeout=300
+            )
+        except (subprocess.TimeoutExpired, OSError):
+            return None
+        return video_path if proc.returncode == 0 and video_path.exists() else None