uprobe 0.1.0__py3-none-any.whl

uprobe/__init__.py

@@ -0,0 +1,18 @@
+"""
+U-Probe package.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+__version__ = "0.1.0"
+
+__all__ = ["UProbeAPI", "__version__"]
+
+def __getattr__(name: str) -> Any:
+    """Lazy attribute access for optional heavy imports."""
+    if name == "UProbeAPI":
+        from uprobe.core.api import UProbeAPI
+        return UProbeAPI
+    raise AttributeError(f"module 'uprobe' has no attribute {name!r}")

uprobe/__main__.py

@@ -0,0 +1,17 @@
+
+"""
+Main entry point for U-Probe CLI.
+"""
+import sys
+import os
+from pathlib import Path
+
+# Add project root to sys.path for development
+project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+sys.path.append(project_root)
+
+from uprobe.core.cli import main
+
+
+if __name__ == '__main__':
+    main()

uprobe/core/.readthedocs.yaml

@@ -0,0 +1,53 @@
+# Read the Docs configuration file for U-Probe
+# See https://docs.readthedocs.io/en/stable/config-file/v2.html for details
+
+version: 2
+
+# Set the OS, Python version and other tools you might need
+build:
+  os: ubuntu-22.04
+  tools:
+    python: "3.8"
+    # You can also specify other tools and their versions:
+    # nodejs: "16"
+    # rust: "1.55"
+    # golang: "1.17"
+
+# Build documentation in the docs/ directory with Sphinx
+sphinx:
+  configuration: docs/source/conf.py
+  # You can configure Sphinx to use a different builder, for instance use the dirhtml builder for simpler URLs
+  # builder: "dirhtml"
+  # Fail on all warnings
+  fail_on_warning: false
+
+# Optionally build your docs in additional formats such as PDF
+formats:
+  - pdf
+  - epub
+
+# Optional but recommended, declare the Python requirements required
+# to build your documentation
+# See https://docs.readthedocs.io/en/stable/guides/reproducible-builds.html
+python:
+  install:
+    - method: pip
+      path: .
+      extra_requirements:
+        - docs
+    - requirements: docs/requirements.txt
+
+# Set the version of Python and other tools you might need
+# build:
+#   os: ubuntu-22.04
+#   tools:
+#     python: "3.8"
+#   jobs:
+#     post_create_environment:
+#       # Install poetry
+#       - pip install poetry
+#       # Tell poetry to not use a virtual environment
+#       - poetry config virtualenvs.create false
+#     post_install:
+#       # Install dependencies
+#       - poetry install --with docs

uprobe/core/__init__.py

@@ -0,0 +1,26 @@
+"""
+U-Probe package.
+
+Note:
+This package exposes `UProbeAPI`, but we avoid importing it eagerly to keep
+lightweight entry points (e.g. `python -m uprobe.core.agent.repl_bootstrap`) working
+even in environments where optional runtime dependencies are not installed yet.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from uprobe import __version__
+
+__all__ = ["UProbeAPI"]
+
+
+def __getattr__(name: str) -> Any:
+    """Lazy attribute access for optional heavy imports."""
+    if name == "UProbeAPI":
+        from .api import UProbeAPI  # Local import to avoid import-time deps
+
+        return UProbeAPI
+    raise AttributeError(f"module 'uprobe.core' has no attribute {name!r}")
+

uprobe/core/agent/__init__.py

@@ -0,0 +1 @@
+

uprobe/core/agent/repl_bootstrap.py

@@ -0,0 +1,193 @@
+"""
+Bootstrap Pantheon REPL in a workspace with the U-Probe team template.
+
+This module intentionally bypasses the existing U-Probe CLI and the
+`uprobe.core.agent.api.UProbeAgentAPI` workflow. Instead, it:
+1) Copies `uprobe_team.md` into `<workspace>/.pantheon/teams/`
+2) Launches `python -m pantheon.repl --template <copied_md>`
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+import shutil
+import subprocess
+import sys
+from pathlib import Path
+
+
+def _build_arg_parser() -> argparse.ArgumentParser:
+    """Build CLI argument parser."""
+    parser = argparse.ArgumentParser(
+        prog="python -m uprobe.core.agent.repl_bootstrap",
+        description=(
+            "Copy uprobe_team.md into <workspace>/.pantheon/teams/ and launch "
+            "Pantheon REPL with --template so uprobe_team is loaded by default."
+        ),
+    )
+    parser.add_argument(
+        "--workspace",
+        type=str,
+        default=str(Path.cwd()),
+        help="Workspace directory to install templates into (default: CWD).",
+    )
+    parser.add_argument(
+        "--force",
+        action="store_true",
+        help="Overwrite existing <workspace>/.pantheon/teams/uprobe_team.md.",
+    )
+
+    # Pass-through options for pantheon.repl (keep it minimal but useful).
+    parser.add_argument(
+        "--memory-dir",
+        type=str,
+        default=None,
+        help="Pantheon memory dir (passed to pantheon.repl).",
+    )
+    parser.add_argument(
+        "--log-level",
+        type=str,
+        default=None,
+        help="Log level for REPL (DEBUG, INFO, WARNING, ERROR, CRITICAL).",
+    )
+    parser.add_argument(
+        "--quiet",
+        action="store_true",
+        help="Disable console logging in REPL (passed to pantheon.repl).",
+    )
+    parser.add_argument(
+        "--resync",
+        action="store_true",
+        help=(
+            "Force pantheon.repl to resync templates "
+            "(cleans skills/agents/teams/prompts dirs)."
+        ),
+    )
+    parser.add_argument(
+        "--chat-id",
+        type=str,
+        default=None,
+        help="Resume a specific chat ID (passed to pantheon.repl).",
+    )
+    parser.add_argument(
+        "repl_args",
+        nargs=argparse.REMAINDER,
+        help=(
+            "Additional arguments passed to pantheon.repl after '--'. "
+            "Example: -- --log-level DEBUG"
+        ),
+    )
+    return parser
+
+
+def _normalize_repl_args(remainder: list[str]) -> list[str]:
+    """Normalize argparse REMAINDER (strip leading '--' if present)."""
+    if remainder and remainder[0] == "--":
+        return remainder[1:]
+    return remainder
+
+
+def _install_team_template(workspace: Path, force: bool) -> Path:
+    """Copy uprobe_team.md into <workspace>/.pantheon/teams/ and return dest."""
+    src = Path(__file__).resolve().parent / "templates" / "uprobe_team.md"
+    if not src.exists():
+        raise FileNotFoundError(f"Team template not found: {src}")
+
+    dest_dir = workspace / ".pantheon" / "teams"
+    dest_dir.mkdir(parents=True, exist_ok=True)
+
+    dest = dest_dir / "uprobe_team.md"
+    if dest.exists() and not force:
+        return dest
+
+    shutil.copy2(src, dest)
+    return dest
+
+
+def _install_protocol_template(workspace: Path, force: bool) -> Path:
+    """Copy DEFAULT_PROTOCOL.yaml into <workspace>/ and return dest."""
+    src = Path(__file__).resolve().parent / "templates" / "DEFAULT_PROTOCOL.yaml"
+    if not src.exists():
+        raise FileNotFoundError(f"Protocol template not found: {src}")
+
+    dest = workspace / "DEFAULT_PROTOCOL.yaml"
+    # Skip copying DEFAULT_PROTOCOL.yaml to avoid cluttering the workspace
+    # if dest.exists() and not force:
+    #     return dest
+    #
+    # shutil.copy2(src, dest)
+    return dest
+
+
+def _launch_repl(
+    workspace: Path,
+    template_path: Path,
+    memory_dir: str | None,
+    log_level: str | None,
+    quiet: bool,
+    resync: bool,
+    chat_id: str | None,
+    extra_args: list[str],
+) -> int:
+    """Launch pantheon.repl with --template and return exit code."""
+    # Ensure `pantheon` can be imported even when running in an arbitrary
+    # workspace directory during development (repo not installed site-wide).
+    repo_root = Path(__file__).resolve().parents[3]
+    env = os.environ.copy()
+    pythonpath = env.get("PYTHONPATH", "")
+    env["PYTHONPATH"] = (
+        f"{repo_root}{os.pathsep}{pythonpath}" if pythonpath else str(repo_root)
+    )
+
+    cmd: list[str] = [
+        sys.executable,
+        "-m",
+        "pantheon.repl",
+        "--template",
+        str(template_path),
+        "--workspace",
+        str(workspace),
+    ]
+    if memory_dir:
+        cmd += ["--memory-dir", memory_dir]
+    if chat_id:
+        cmd += ["--chat-id", chat_id]
+    if log_level:
+        cmd += ["--log-level", log_level]
+    if quiet:
+        cmd += ["--quiet"]
+    if resync:
+        cmd += ["--resync"]
+    if extra_args:
+        cmd += extra_args
+
+    completed = subprocess.run(cmd, cwd=str(workspace), env=env)
+    return int(completed.returncode)
+
+
+def main(argv: list[str] | None = None) -> int:
+    """Main entry point."""
+    parser = _build_arg_parser()
+    args = parser.parse_args(argv)
+
+    workspace = Path(args.workspace).expanduser().resolve()
+    template_path = _install_team_template(workspace=workspace, force=bool(args.force))
+    # _install_protocol_template(workspace=workspace, force=bool(args.force))
+
+    extra_args = _normalize_repl_args(list(args.repl_args or []))
+    return _launch_repl(
+        workspace=workspace,
+        template_path=template_path,
+        memory_dir=args.memory_dir,
+        log_level=args.log_level,
+        quiet=bool(args.quiet),
+        resync=bool(args.resync),
+        chat_id=args.chat_id,
+        extra_args=extra_args,
+    )
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())
+

uprobe/core/agent/session_manager.py

@@ -0,0 +1,207 @@
+import os
+import uuid
+import shutil
+import logging
+from pathlib import Path
+from typing import Dict, Any, Optional, List, BinaryIO
+
+from pantheon.chatroom import ChatRoom
+from pantheon.factory.template_io import UnifiedMarkdownParser
+from pantheon.utils.vision import parse_image_mentions
+
+logger = logging.getLogger(__name__)
+
+
+class AgentSessionManager:
+    """Manages agent sessions using Pantheon ChatRoom."""
+
+    def __init__(self, workspace_root: Path | None = None):
+        if workspace_root is None:
+            workspace_root = Path(__file__).resolve().parents[3]
+        self.workspace_root = workspace_root
+        # Use absolute path to the template file in the source code
+        self.template_path = Path(__file__).resolve().parent / "templates" / "uprobe_team.md"
+        self.protocol_template_path = Path(__file__).resolve().parent / "templates" / "DEFAULT_PROTOCOL.yaml"
+        self.pantheon_memory_dir = self.workspace_root / ".pantheon" / "memory"
+        self.output_dir = Path(os.environ.get("UPROBE_OUTPUT_DIR", str(self.workspace_root / "outputs")))
+        self.output_dir.mkdir(parents=True, exist_ok=True)
+        self.chatroom = ChatRoom(memory_dir=str(self.pantheon_memory_dir), workspace_path=str(self.workspace_root))
+        self.sessions: Dict[str, Dict[str, Any]] = {}
+
+    async def initialize(self):
+        """Initialize the ChatRoom and background services."""
+        self._ensure_team_template_registered()
+        # self._ensure_protocol_template_installed()
+        
+        # Set pantheon log level to DEBUG to print full agent execution output
+        from pantheon.utils.log import set_level
+        set_level("DEBUG")
+        
+        await self.chatroom.run_setup()
+
+    def _load_team_template(self, model: Optional[str]) -> Dict[str, Any]:
+        """Load and configure the team template."""
+        # Template is already ensured in initialize(), but we resolve the path again to be safe
+        team_template_path = self.workspace_root / ".pantheon" / "teams" / "uprobe_team.md"
+        if not team_template_path.exists():
+             # Fallback if initialize wasn't called or file was deleted
+             team_template_path = self._ensure_team_template_registered()
+             
+        parser = UnifiedMarkdownParser()
+        team_config = parser.parse_file(team_template_path)
+        team_dict = team_config.to_dict()
+        team_dict["type"] = "team"
+        team_dict["source_path"] = str(team_template_path)
+        if model:
+            for agent_cfg in team_dict.get("agents", []):
+                if isinstance(agent_cfg, dict):
+                    agent_cfg["model"] = model
+        return team_dict
+
+    def _ensure_team_template_registered(self) -> Path:
+        """Ensure uprobe_team.md is registered under .pantheon/teams."""
+        dest_dir = self.workspace_root / ".pantheon" / "teams"
+        dest_dir.mkdir(parents=True, exist_ok=True)
+        dest = dest_dir / "uprobe_team.md"
+        if not dest.exists():
+            shutil.copy2(self.template_path, dest)
+        return dest
+
+    def _ensure_protocol_template_installed(self) -> Path:
+        """Ensure DEFAULT_PROTOCOL.yaml is present in the workspace root."""
+        dest = self.workspace_root / "DEFAULT_PROTOCOL.yaml"
+        # Skip copying DEFAULT_PROTOCOL.yaml to avoid cluttering the workspace
+        # if not dest.exists() and self.protocol_template_path.exists():
+        #     shutil.copy2(self.protocol_template_path, dest)
+        #     logger.info(f"Installed protocol template to {dest}")
+        return dest
+
+    def _build_chat_message(self, content: str) -> List[dict]:
+        """Build chat message with image support."""
+        return parse_image_mentions(content)
+
+    async def upload_file(self, session_id: str, file_obj: BinaryIO, filename: str, upload_dir: Optional[Path] = None) -> Dict[str, Any]:
+        """Save an uploaded file and return its metadata."""
+        if upload_dir is None:
+            upload_dir = self.output_dir / "uploads" / session_id
+        upload_dir.mkdir(parents=True, exist_ok=True)
+        file_id = str(uuid.uuid4())
+        safe_filename = Path(filename).name
+        file_path = upload_dir / safe_filename
+        with open(file_path, "wb") as f:
+            shutil.copyfileobj(file_obj, f)
+        size = file_path.stat().st_size
+        logger.info(f"Saved uploaded file to {file_path}")
+        if session_id in self.sessions:
+            if "uploads" not in self.sessions[session_id]:
+                self.sessions[session_id]["uploads"] = {}
+            self.sessions[session_id]["uploads"][file_id] = {"path": str(file_path), "filename": safe_filename, "size": size}
+        return {"id": file_id, "filename": safe_filename, "path": str(file_path), "size": size}
+
+    async def create_session(self, model: str = "gpt-4", api_key: Optional[str] = None, api_base: Optional[str] = None, proxy: Optional[str] = None) -> str:
+        """Start a new agent session."""
+        if api_key:
+            os.environ["OPENAI_API_KEY"] = api_key
+        if api_base:
+            os.environ["OPENAI_API_BASE"] = api_base
+        if proxy:
+            os.environ["http_proxy"] = proxy
+            os.environ["https_proxy"] = proxy
+        await self.initialize()
+        create_res = await self.chatroom.create_chat("http-session")
+        if not create_res.get("success"):
+            raise RuntimeError(f"Failed to create chat: {create_res.get('message')}")
+        chat_id = create_res["chat_id"]
+        try:
+            template = self._load_team_template(model)
+            setup_res = await self.chatroom.setup_team_for_chat(chat_id, template)
+            if not setup_res.get("success"):
+                raise RuntimeError(f"Failed to setup team: {setup_res.get('message')}")
+        except Exception as e:
+            await self.chatroom.delete_chat(chat_id)
+            raise e
+        self.sessions[chat_id] = {"chat_id": chat_id, "model": model}
+        logger.info(f"Started agent session {chat_id} with model={model}")
+        return chat_id
+
+    async def chat(self, session_id: str, content: str, attachment_ids: Optional[List[str]] = None, process_step_message=None) -> Dict[str, Any]:
+        """Send a message to an active session."""
+        if session_id not in self.sessions:
+            pass
+        if attachment_ids is None:
+            attachment_ids = []
+        if attachment_ids and session_id in self.sessions:
+            uploads = self.sessions[session_id].get("uploads", {})
+            attached_files = []
+            for aid in attachment_ids:
+                if aid in uploads:
+                    attached_files.append(uploads[aid])
+            if attached_files:
+                file_context = "\n\nUser uploaded the following files:\n"
+                for f in attached_files:
+                    file_context += f"- {f['filename']} (Path: {f['path']})\n"
+                content += file_context
+        chat_message = self._build_chat_message(content)
+        response = await self.chatroom.chat(session_id, chat_message, process_step_message=process_step_message)
+        return response
+
+    async def rewind_and_rerun(self, session_id: str, user_turn_index: int, content: str, attachment_ids: Optional[List[str]] = None, process_step_message=None) -> Dict[str, Any]:
+        """Rewind chat history to a user turn and rerun from edited content."""
+        if session_id not in self.sessions:
+            pass
+        if attachment_ids is None:
+            attachment_ids = []
+        try:
+            await self.chatroom.stop_chat(session_id)
+        except Exception:
+            pass
+        memory = self.chatroom.memory_manager.get_memory(session_id)
+        user_turns = memory.get_user_turns()
+        if user_turn_index < 0 or user_turn_index >= len(user_turns):
+            raise ValueError("Invalid user_turn_index")
+        target_index = user_turns[user_turn_index][0]
+        memory.revert_to_message(target_index)
+        try:
+            await memory.flush()
+        except Exception:
+            pass
+        return await self.chat(session_id, content, attachment_ids, process_step_message=process_step_message)
+
+    async def delete_upload(self, session_id: str, attachment_id: str) -> bool:
+        """Delete an uploaded file by attachment id."""
+        uploads = self.sessions.get(session_id, {}).get("uploads", {})
+        meta = uploads.pop(attachment_id, None)
+        if not meta:
+            return False
+        file_path = Path(meta.get("path", ""))
+        if file_path.exists():
+            file_path.unlink()
+        return True
+
+    async def stop_session(self, session_id: str) -> bool:
+        """Stop and clean up a session."""
+        if session_id in self.sessions:
+            del self.sessions[session_id]
+        try:
+            await self.chatroom.stop_chat(session_id)
+        except Exception:
+            pass
+        result = await self.chatroom.delete_chat(session_id)
+        return result.get("success", False)
+
+
+_INSTANCE: Optional[AgentSessionManager] = None
+
+
+def get_session_manager(workspace_root: Path | None = None, output_dir: Path | None = None, memory_dir: Path | None = None) -> AgentSessionManager:
+    """Get or create the global session manager instance."""
+    global _INSTANCE
+    if _INSTANCE is None:
+        _INSTANCE = AgentSessionManager(workspace_root=workspace_root)
+        if output_dir:
+            _INSTANCE.output_dir = output_dir
+            _INSTANCE.output_dir.mkdir(parents=True, exist_ok=True)
+        if memory_dir:
+            _INSTANCE.pantheon_memory_dir = memory_dir
+            _INSTANCE.chatroom = ChatRoom(memory_dir=str(_INSTANCE.pantheon_memory_dir), workspace_path=str(_INSTANCE.workspace_root))
+    return _INSTANCE

uprobe/core/agent/templates/DEFAULT_PROTOCOL.yaml

@@ -0,0 +1,55 @@
+name: # protocol name
+
+description: Protocol for designing probe type xxx from species xxx
+
+genome:   # genome name
+
+extracts:
+  target_region:
+    source:   # genome | exon | CDS | UTR
+    length:   # int
+    overlap:   # int
+
+targets:
+  - xxx1: # chr:start-end | gene_name | name: string
+  - xxx2: # chr:start-end | gene_name | name: string
+
+encoding:
+  xxx1:
+    BC1: xxx
+    BC2: xxx
+  xxx2:
+    BC1: xxx
+    BC2: xxx
+
+probes:
+
+# NOTE:
+# - `attributes` and `post_process` can be omitted.
+# - If they are missing or empty, the CLI will auto-generate defaults
+#   based on DNA (source=genome) vs RNA (source=exon/CDS/UTR) mode.
+attributes:
+
+
+post_process:
+  filters:
+
+  # clean
+  avoid_otp:
+
+  # equal_space
+  equal_space:
+
+  # remove overlap
+  remove_overlap:
+    location_interval: 0
+
+  sorts:
+
+
+
+summary:
+  report_name: dna_report  # template name  # class: dna_report or rna_report
+  attributes:
+
+