oshell 0.1.1__py3-none-any.whl

oshell/config.py

@@ -0,0 +1,91 @@
+"""Configuration loading and persistence for oshell.
+
+Settings are resolved in this order (highest priority first):
+1. Command-line flags
+2. Environment variables
+3. Config file (~/.config/oshell/config.json on Linux/macOS)
+4. Built-in defaults
+"""
+
+from __future__ import annotations
+
+import json
+import os
+from dataclasses import asdict, dataclass
+from pathlib import Path
+
+from platformdirs import user_config_dir
+
+APP_NAME = "oshell"
+
+CONFIG_DIR = Path(user_config_dir(APP_NAME))
+CONFIG_PATH = CONFIG_DIR / "config.json"
+
+
+@dataclass
+class Config:
+    """Resolved runtime configuration."""
+
+    provider: str = "ollama"  # ollama | openai | anthropic | groq | gemini
+    model: str = "llama3.2"
+    system_prompt: str = "You are a helpful, concise assistant."
+    persona: str = ""  # named preset; overrides system_prompt when set
+    temperature: float = 0.7
+
+    # Provider endpoints / credentials
+    ollama_host: str = "http://localhost:11434"
+    openai_base_url: str = "https://api.openai.com/v1"
+    openai_api_key: str = ""
+    anthropic_api_key: str = ""
+    groq_api_key: str = ""
+    gemini_api_key: str = ""
+
+    # Retrieval (chat-with-your-files)
+    embed_model: str = "nomic-embed-text"  # ollama embedding model
+
+    # Media generation
+    image_model: str = "gpt-image-1"
+    image_size: str = "1024x1024"
+    video_model: str = "minimax/video-01"
+    replicate_api_token: str = ""
+    media_output_dir: str = ""  # empty -> ~/oshell/media
+
+    @classmethod
+    def load(cls) -> "Config":
+        """Load config from file + environment, applying defaults."""
+        data: dict = {}
+        if CONFIG_PATH.exists():
+            try:
+                data = json.loads(CONFIG_PATH.read_text(encoding="utf-8"))
+            except (json.JSONDecodeError, OSError):
+                data = {}
+
+        cfg = cls(**{k: v for k, v in data.items() if k in cls.__annotations__})
+
+        # Environment overrides
+        cfg.openai_api_key = os.environ.get("OPENAI_API_KEY", cfg.openai_api_key)
+        if os.environ.get("OPENAI_BASE_URL"):
+            cfg.openai_base_url = os.environ["OPENAI_BASE_URL"]
+        if os.environ.get("OLLAMA_HOST"):
+            cfg.ollama_host = os.environ["OLLAMA_HOST"]
+        if os.environ.get("OSHELL_PROVIDER"):
+            cfg.provider = os.environ["OSHELL_PROVIDER"]
+        if os.environ.get("OSHELL_MODEL"):
+            cfg.model = os.environ["OSHELL_MODEL"]
+        if os.environ.get("REPLICATE_API_TOKEN"):
+            cfg.replicate_api_token = os.environ["REPLICATE_API_TOKEN"]
+        cfg.anthropic_api_key = os.environ.get(
+            "ANTHROPIC_API_KEY", cfg.anthropic_api_key
+        )
+        cfg.groq_api_key = os.environ.get("GROQ_API_KEY", cfg.groq_api_key)
+        cfg.gemini_api_key = os.environ.get("GEMINI_API_KEY", cfg.gemini_api_key)
+
+        return cfg
+
+    def save(self) -> Path:
+        """Persist current config to disk."""
+        CONFIG_DIR.mkdir(parents=True, exist_ok=True)
+        CONFIG_PATH.write_text(
+            json.dumps(asdict(self), indent=2), encoding="utf-8"
+        )
+        return CONFIG_PATH

oshell/history.py

@@ -0,0 +1,57 @@
+"""Conversation history persistence (JSONL session files)."""
+
+from __future__ import annotations
+
+import json
+import time
+from pathlib import Path
+from typing import List
+
+from platformdirs import user_data_dir
+
+from .config import APP_NAME
+from .providers import Message
+
+HISTORY_DIR = Path(user_data_dir(APP_NAME)) / "sessions"
+
+
+def _ensure_dir() -> None:
+    HISTORY_DIR.mkdir(parents=True, exist_ok=True)
+
+
+def new_session_path() -> Path:
+    """Return a fresh, timestamped session file path."""
+    _ensure_dir()
+    stamp = time.strftime("%Y%m%d-%H%M%S")
+    return HISTORY_DIR / f"chat-{stamp}.jsonl"
+
+
+def append(path: Path, message: Message) -> None:
+    """Append a single message to a session file."""
+    _ensure_dir()
+    with path.open("a", encoding="utf-8") as fh:
+        fh.write(json.dumps(message) + "\n")
+
+
+def load(path: Path) -> List[Message]:
+    """Load all messages from a session file."""
+    if not path.exists():
+        return []
+    messages: List[Message] = []
+    with path.open(encoding="utf-8") as fh:
+        for line in fh:
+            line = line.strip()
+            if line:
+                messages.append(json.loads(line))
+    return messages
+
+
+def list_sessions() -> List[Path]:
+    """Return saved sessions, newest first."""
+    if not HISTORY_DIR.exists():
+        return []
+    return sorted(
+        HISTORY_DIR.glob("chat-*.jsonl"),
+        key=lambda p: p.stat().st_mtime,
+        reverse=True,
+    )

oshell/media.py

@@ -0,0 +1,188 @@
+"""Media generation backends: images (OpenAI) and video (Replicate).
+
+Both return the path(s) to the saved file(s). Outputs are written to a
+timestamped folder so repeated runs never clobber each other.
+"""
+
+from __future__ import annotations
+
+import base64
+import time
+from dataclasses import dataclass
+from pathlib import Path
+from typing import List
+
+import httpx
+from platformdirs import user_pictures_dir
+
+from .config import APP_NAME, Config
+
+# How long to wait for slow video jobs (seconds).
+VIDEO_POLL_TIMEOUT = 600
+VIDEO_POLL_INTERVAL = 3
+
+
+@dataclass
+class MediaResult:
+    """Where the generated media landed, plus the prompt actually used."""
+
+    paths: List[Path]
+    prompt: str
+
+
+def _output_dir(cfg: Config, kind: str) -> Path:
+    base = (
+        Path(cfg.media_output_dir)
+        if cfg.media_output_dir
+        else Path(user_pictures_dir()) / APP_NAME
+    )
+    stamp = time.strftime("%Y%m%d-%H%M%S")
+    target = base / kind / stamp
+    target.mkdir(parents=True, exist_ok=True)
+    return target
+
+
+def _slug(text: str, limit: int = 40) -> str:
+    keep = [c if c.isalnum() else "-" for c in text.lower()]
+    slug = "".join(keep).strip("-")
+    while "--" in slug:
+        slug = slug.replace("--", "-")
+    return slug[:limit] or "media"
+
+
+class ImageGenerator:
+    """Generates images with the OpenAI Images API (gpt-image-1 / DALL·E 3)."""
+
+    def __init__(self, cfg: Config) -> None:
+        if not cfg.openai_api_key:
+            raise ValueError(
+                "OpenAI API key not set. Export OPENAI_API_KEY or run "
+                "`oshell config set openai_api_key <key>`."
+            )
+        self.base_url = cfg.openai_base_url.rstrip("/")
+        self.api_key = cfg.openai_api_key
+        self.model = cfg.image_model
+        self.size = cfg.image_size
+        self.cfg = cfg
+
+    def generate(self, prompt: str, n: int = 1) -> MediaResult:
+        headers = {"Authorization": f"Bearer {self.api_key}"}
+        payload = {
+            "model": self.model,
+            "prompt": prompt,
+            "n": n,
+            "size": self.size,
+        }
+        resp = httpx.post(
+            f"{self.base_url}/images/generations",
+            json=payload,
+            headers=headers,
+            timeout=180,
+        )
+        resp.raise_for_status()
+        data = resp.json().get("data", [])
+
+        out_dir = _output_dir(self.cfg, "images")
+        slug = _slug(prompt)
+        paths: List[Path] = []
+        for i, item in enumerate(data):
+            suffix = "" if len(data) == 1 else f"-{i + 1}"
+            dest = out_dir / f"{slug}{suffix}.png"
+            if item.get("b64_json"):
+                dest.write_bytes(base64.b64decode(item["b64_json"]))
+            elif item.get("url"):
+                img = httpx.get(item["url"], timeout=120)
+                img.raise_for_status()
+                dest.write_bytes(img.content)
+            else:
+                continue
+            paths.append(dest)
+
+        if not paths:
+            raise RuntimeError("Image API returned no usable data.")
+        return MediaResult(paths=paths, prompt=prompt)
+
+
+class VideoGenerator:
+    """Generates video via Replicate's prediction API.
+
+    Works with any text-to-video model on Replicate (default
+    ``minimax/video-01``); just change ``video_model`` in config.
+    """
+
+    def __init__(self, cfg: Config) -> None:
+        if not cfg.replicate_api_token:
+            raise ValueError(
+                "Replicate token not set. Export REPLICATE_API_TOKEN or run "
+                "`oshell config set replicate_api_token <token>`."
+            )
+        self.token = cfg.replicate_api_token
+        self.model = cfg.video_model
+        self.cfg = cfg
+
+    def _headers(self) -> dict:
+        return {
+            "Authorization": f"Bearer {self.token}",
+            "Content-Type": "application/json",
+        }
+
+    def generate(self, prompt: str) -> MediaResult:
+        # Use the model-scoped predictions endpoint so we don't need a version hash.
+        url = f"https://api.replicate.com/v1/models/{self.model}/predictions"
+        resp = httpx.post(
+            url,
+            json={"input": {"prompt": prompt}},
+            headers={**self._headers(), "Prefer": "wait"},
+            timeout=120,
+        )
+        resp.raise_for_status()
+        prediction = resp.json()
+
+        prediction = self._await_completion(prediction)
+        output = prediction.get("output")
+        urls = self._collect_urls(output)
+        if not urls:
+            raise RuntimeError(
+                f"Video job finished with status '{prediction.get('status')}' "
+                "but produced no output URL."
+            )
+
+        out_dir = _output_dir(self.cfg, "videos")
+        slug = _slug(prompt)
+        paths: List[Path] = []
+        for i, video_url in enumerate(urls):
+            suffix = "" if len(urls) == 1 else f"-{i + 1}"
+            dest = out_dir / f"{slug}{suffix}.mp4"
+            data = httpx.get(video_url, timeout=300)
+            data.raise_for_status()
+            dest.write_bytes(data.content)
+            paths.append(dest)
+        return MediaResult(paths=paths, prompt=prompt)
+
+    def _await_completion(self, prediction: dict) -> dict:
+        terminal = {"succeeded", "failed", "canceled"}
+        deadline = time.time() + VIDEO_POLL_TIMEOUT
+        while prediction.get("status") not in terminal:
+            if time.time() > deadline:
+                raise TimeoutError("Video generation timed out.")
+            poll_url = prediction.get("urls", {}).get("get")
+            if not poll_url:
+                break
+            time.sleep(VIDEO_POLL_INTERVAL)
+            r = httpx.get(poll_url, headers=self._headers(), timeout=60)
+            r.raise_for_status()
+            prediction = r.json()
+
+        if prediction.get("status") == "failed":
+            raise RuntimeError(f"Video generation failed: {prediction.get('error')}")
+        return prediction
+
+    @staticmethod
+    def _collect_urls(output) -> List[str]:
+        if not output:
+            return []
+        if isinstance(output, str):
+            return [output]
+        if isinstance(output, list):
+            return [u for u in output if isinstance(u, str)]
+        return []

oshell/media_agent.py

@@ -0,0 +1,55 @@
+"""The media agent: turns a short brief into rich media.
+
+It uses the chat LLM to *enhance* a user's brief into a detailed, model-ready
+prompt (composition, lighting, style, camera, mood…), then hands that prompt
+to the image or video generator. This is the "sophisticated" part — you write
+"a fox in a forest" and get a cinematic, well-specified prompt.
+"""
+
+from __future__ import annotations
+
+from .media import ImageGenerator, MediaResult, VideoGenerator
+from .providers import Message, Provider
+
+_IMAGE_ENHANCE = """\
+You are a prompt engineer for text-to-image models. Rewrite the user's brief \
+into ONE vivid, detailed image prompt. Include subject, composition, lighting, \
+color palette, art style, and mood. Keep it under 80 words. Output ONLY the \
+prompt text — no quotes, no preamble, no explanation."""
+
+_VIDEO_ENHANCE = """\
+You are a prompt engineer for text-to-video models. Rewrite the user's brief \
+into ONE detailed video prompt describing the scene, subject motion, camera \
+movement, lighting, and mood. Keep it under 80 words. Output ONLY the prompt \
+text — no quotes, no preamble, no explanation."""
+
+
+class MediaAgent:
+    """Enhances a brief with the LLM, then generates media."""
+
+    def __init__(self, provider: Provider) -> None:
+        self.provider = provider
+
+    def enhance(self, brief: str, kind: str) -> str:
+        """Expand ``brief`` into a detailed prompt for the given media kind."""
+        system = _IMAGE_ENHANCE if kind == "image" else _VIDEO_ENHANCE
+        messages: list[Message] = [
+            {"role": "system", "content": system},
+            {"role": "user", "content": brief},
+        ]
+        enhanced = "".join(self.provider.stream_chat(messages)).strip()
+        # Strip stray surrounding quotes some models add.
+        enhanced = enhanced.strip('"').strip()
+        return enhanced or brief
+
+    def make_image(
+        self, brief: str, generator: ImageGenerator, n: int = 1, enhance: bool = True
+    ) -> MediaResult:
+        prompt = self.enhance(brief, "image") if enhance else brief
+        return generator.generate(prompt, n=n)
+
+    def make_video(
+        self, brief: str, generator: VideoGenerator, enhance: bool = True
+    ) -> MediaResult:
+        prompt = self.enhance(brief, "video") if enhance else brief
+        return generator.generate(prompt)

oshell/personas.py

@@ -0,0 +1,90 @@
+"""Named personas (system-prompt presets).
+
+Built-in personas live here; users can add their own with
+``oshell persona add <name> "<system prompt>"`` which writes to a small JSON
+file next to the config.
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Dict
+
+from .config import CONFIG_DIR
+
+PERSONA_PATH = CONFIG_DIR / "personas.json"
+
+BUILTIN: Dict[str, str] = {
+    "default": "You are a helpful, concise assistant.",
+    "reviewer": (
+        "You are a meticulous senior code reviewer. Point out bugs, security "
+        "issues, edge cases, and style problems. Be direct and specific, cite "
+        "line-level concerns, and suggest concrete fixes."
+    ),
+    "teacher": (
+        "You are a patient programming teacher. Explain concepts step by step "
+        "with simple analogies and short examples. Check understanding and "
+        "avoid jargon unless you define it."
+    ),
+    "shell": (
+        "You are a command-line expert. Prefer giving the exact command(s) to "
+        "run, with a one-line explanation. Assume a competent user."
+    ),
+    "rubber-duck": (
+        "You are a rubber-duck debugging partner. Ask probing questions that "
+        "help the user reason through their problem rather than giving the "
+        "answer outright."
+    ),
+    "concise": "Answer in as few words as possible. No preamble. No filler.",
+    "pirate": "You are a witty pirate. Answer correctly, but talk like a pirate.",
+}
+
+
+def _load_custom() -> Dict[str, str]:
+    if not PERSONA_PATH.exists():
+        return {}
+    try:
+        return json.loads(PERSONA_PATH.read_text(encoding="utf-8"))
+    except (json.JSONDecodeError, OSError):
+        return {}
+
+
+def all_personas() -> Dict[str, str]:
+    """Built-ins merged with user-defined personas (custom wins)."""
+    merged = dict(BUILTIN)
+    merged.update(_load_custom())
+    return merged
+
+
+def get(name: str) -> str | None:
+    """Return the system prompt for ``name``, or None if unknown."""
+    return all_personas().get(name)
+
+
+def add(name: str, prompt: str) -> Path:
+    """Create or update a custom persona."""
+    custom = _load_custom()
+    custom[name] = prompt
+    CONFIG_DIR.mkdir(parents=True, exist_ok=True)
+    PERSONA_PATH.write_text(json.dumps(custom, indent=2), encoding="utf-8")
+    return PERSONA_PATH
+
+
+def remove(name: str) -> bool:
+    """Delete a custom persona. Returns True if it existed."""
+    custom = _load_custom()
+    if name not in custom:
+        return False
+    del custom[name]
+    PERSONA_PATH.write_text(json.dumps(custom, indent=2), encoding="utf-8")
+    return True
+
+
+def resolve_system_prompt(cfg) -> str:
+    """Return the active system prompt, honoring a selected persona."""
+    if cfg.persona:
+        prompt = get(cfg.persona)
+        if prompt:
+            return prompt
+    return cfg.system_prompt