sourcebot 0.1.0__py3-none-any.whl

sourcebot/runtime/dag/scheduler/dag_scheduler.py

@@ -0,0 +1,319 @@
+# sourcebot/runtime/dag/scheduler/dag_scheduler.py
+import asyncio
+import json
+import logging
+from pathlib import Path
+from collections import deque
+from datetime import datetime
+from sourcebot.runtime.dag.scheduler.task_graph import TaskGraph
+from sourcebot.runtime.dag.scheduler.state_store import StateStore
+from sourcebot.runtime.dag.scheduler.retry_policy import RetryPolicy
+from sourcebot.runtime.dag.scheduler.run_store import RunStore
+from sourcebot.llm.core.message import Message
+
+logger = logging.getLogger(__name__)
+
+
+class DAGScheduler:
+
+    def __init__(
+        self,
+        agent_factory,
+        workspace,
+        runs_dir="runs",
+        max_concurrent=4,
+        retry_policy=None,
+        
+    ):
+
+        self.agent_factory = agent_factory
+
+        self.runs_dir = Path(workspace)/runs_dir
+
+        self.max_concurrent = max_concurrent
+
+        self.retry_policy = retry_policy or RetryPolicy()
+
+        self.semaphore = asyncio.Semaphore(max_concurrent)
+
+        self.run_store = RunStore(self.runs_dir)
+
+    # Run DAG
+    async def run(self, tasks):
+
+        self.run_id, self.run_dir = self.run_store.create_run(tasks)
+
+        logger.info(f"DAG RUN {self.run_id}")
+
+        self.state = StateStore(self.run_dir)
+
+        graph = TaskGraph(tasks)
+
+        result = await self._execute(graph)
+
+        if result["failed"]:
+            self.run_store.update_status(self.run_dir, "failed")
+        else:
+            self.run_store.update_status(self.run_dir, "completed")
+
+        return result
+
+    # Resume run
+    async def resume(self, run_id):
+
+        self.run_id = run_id
+        self.run_dir = self.runs_dir / run_id
+
+        tasks = self.run_store.load_dag(self.run_dir)
+
+        self.state = StateStore(self.run_dir)
+
+        graph = TaskGraph(tasks)
+
+        logger.info(f"Resuming {run_id}")
+
+        return await self._execute(graph)
+
+
+    # Replay failed tasks
+    async def replay_failed(self, run_id):
+
+        self.run_id = run_id
+        self.run_dir = self.runs_dir / run_id
+
+        tasks = self.run_store.load_dag(self.run_dir)
+
+        self.state = StateStore(self.run_dir)
+
+        for t in list(self.state.state.keys()):
+
+            if self.state.state[t]["status"] == "failed":
+                self.state.state[t]["status"] = "pending"
+
+        self.state.save()
+
+        graph = TaskGraph(tasks)
+
+        logger.info(f"Replay failed tasks {run_id}")
+
+        return await self._execute(graph)
+
+
+    # Core execution
+    async def _execute(self, graph):
+
+        ready = deque()
+
+        running = {}
+
+        future_map = {}
+
+        completed = set()
+        failed = set()
+
+        for tid, deg in graph.in_degree.items():
+
+            state = self.state.status(tid)
+
+            if state == "completed":
+                completed.add(tid)
+                continue
+
+            if deg == 0:
+                ready.append(tid)
+
+        while ready or running:
+
+            while ready and len(running) < self.max_concurrent:
+
+                tid = ready.popleft()
+
+                task = graph.tasks[tid]
+
+                fut = asyncio.create_task(
+                    self._execute_task(task)
+                )
+
+                running[tid] = fut
+                future_map[fut] = tid
+
+            if not running:
+                break
+
+            done, _ = await asyncio.wait(
+                running.values(),
+                return_when=asyncio.FIRST_COMPLETED
+            )
+
+            for fut in done:
+
+                tid = future_map[fut]
+
+                try:
+
+                    result = fut.result()
+
+                    if result["success"]:
+
+                        completed.add(tid)
+
+                        for child in graph.children[tid]:
+
+                            graph.in_degree[child] -= 1
+
+                            if graph.in_degree[child] == 0:
+                                ready.append(child)
+
+                    else:
+                        failed.add(tid)
+
+                except Exception as e:
+
+                    failed.add(tid)
+                    logger.error(f"{tid} crashed: {e}")
+
+                del running[tid]
+
+        return {
+            "completed": list(completed),
+            "failed": list(failed)
+        }
+
+
+    # Task execution
+    async def _execute_task(self, task):
+
+        task_id = task["id"]
+
+        log_file = self.run_dir / "tasks" / f"{task_id}.log"
+
+        attempt = self.state.attempts(task_id) + 1
+
+        while True:
+
+            try:
+
+                self.state.update(
+                    task_id,
+                    status = "running",
+                    attempts = attempt,
+                    started = str(datetime.utcnow())
+                )
+
+                async with self.semaphore:
+
+                    agent = self.agent_factory.build_sub_agent(
+                        task_id = task_id,
+                        task_description = task["description"]
+                    )
+                    messages = self._build_messages(task)
+                    result, _, _ = await agent.run(messages)
+
+                self._write_log(log_file, result)
+
+                if self._is_success(result):
+
+                    self.state.update(
+                        task_id,
+                        status = "completed",
+                        result = result
+                    )
+
+                    return {"success": True}
+
+                raise RuntimeError("task failed")
+
+            except Exception as e:
+
+                self._write_log(log_file, str(e))
+
+                if not self.retry_policy.should_retry(attempt):
+
+                    self.state.update(
+                        task_id,
+                        status = "failed",
+                        error = str(e)
+                    )
+
+                    return {"success": False}
+
+                delay = self.retry_policy.get_delay(attempt)
+
+                logger.warning(
+                    f"{task_id} retry in {delay}s"
+                )
+
+                await asyncio.sleep(delay)
+
+                attempt += 1
+
+    # =================================
+    # Helpers
+    # =================================
+
+    def _write_log(self, file: Path, data):
+
+        file.parent.mkdir(parents=True, exist_ok=True)
+
+        with open(file, "a") as f:
+
+            f.write(
+                f"\n[{datetime.utcnow()}]\n"
+            )
+
+            if isinstance(data, str):
+                f.write(data + "\n")
+            else:
+                f.write(json.dumps(data, indent=2) + "\n")
+
+    def _is_success(self, result):
+
+        if isinstance(result, dict):
+            return result.get("success", True)
+
+        return bool(result)
+
+    def _build_messages(self, task):
+        return [
+            Message(role = "system", content = f"You are executing task {task['id']}"),
+            Message(role = "user", content = self.build_task_description(task)),
+        ]
+
+    def build_task_description(self, task):
+        """Combine the task and context into a complete description."""
+        description = task.get("description", "")
+        context = task.get("context", {})
+        
+        context_parts = []
+        
+        if context.get("rules"):
+            rules = "\n  • ".join(context["rules"])
+            context_parts.append(f"Rules:\n  • {rules}")
+        
+        if context.get("skills"):
+            skills = ", ".join(context["skills"])
+            context_parts.append(f"Required skills: {skills}")
+        
+        if context.get("environment"):
+            env = context["environment"]
+            if env.get("required_tools"):
+                tools = ", ".join(env["required_tools"])
+                context_parts.append(f"Required tools: {tools}")
+            if env.get("working_dir"):
+                context_parts.append(f"Working directory: {env['working_dir']}")
+        
+        if context.get("inherited_context"):
+            inherited = context["inherited_context"]
+            inherited_items = [f"{k}: {v}" for k, v in inherited.items()]
+            context_parts.append(f"Context: {', '.join(inherited_items)}")
+        
+        if context_parts:
+            full_description = f"{description}\n\nContext Information:\n" + "\n".join(context_parts)
+        else:
+            full_description = description
+        
+        return full_description
+
+
+    def save_dag(self, plan_tasks):
+        return self.run_store.create_run(plan_tasks)

sourcebot/runtime/dag/scheduler/retry_policy.py

@@ -0,0 +1,27 @@
+# sourcebot/runtime/dag/scheduler/retry_policy.py
+import random
+class RetryPolicy:
+
+    def __init__(
+        self,
+        max_attempts: int = 3,
+        base_delay: int = 2,
+        max_delay: int = 30,
+        jitter: bool = True,
+    ):
+        self.max_attempts = max_attempts
+        self.base_delay = base_delay
+        self.max_delay = max_delay
+        self.jitter = jitter
+
+    def should_retry(self, attempt: int) -> bool:
+        return attempt < self.max_attempts
+
+    def get_delay(self, attempt: int) -> float:
+
+        delay = min(self.base_delay * (2 ** (attempt - 1)), self.max_delay)
+
+        if self.jitter:
+            delay = delay * (0.5 + random.random())
+
+        return delay

sourcebot/runtime/dag/scheduler/run_store.py

@@ -0,0 +1,58 @@
+# sourcebot/runtime/dag/scheduler/run_store.py
+import json
+import uuid
+from pathlib import Path
+from datetime import datetime
+
+
+class RunStore:
+
+    def __init__(self, runs_dir: Path):
+
+        self.runs_dir = runs_dir
+        self.runs_dir.mkdir(parents=True, exist_ok=True)
+
+    def create_run(self, tasks):
+
+        run_id = f"run_{uuid.uuid4().hex[:8]}"
+
+        run_dir = self.runs_dir / run_id
+        run_dir.mkdir(parents=True)
+
+        (run_dir / "tasks").mkdir()
+
+        (run_dir / "dag.json").write_text(
+            json.dumps(tasks, indent=2)
+        )
+
+        meta = {
+            "run_id": run_id,
+            "created_at": str(datetime.utcnow()),
+            "status": "running"
+        }
+
+        (run_dir / "run_meta.json").write_text(
+            json.dumps(meta, indent=2)
+        )
+
+        return run_id, run_dir
+
+    def load_dag(self, run_dir):
+
+        file = run_dir / "dag.json"
+
+        if not file.exists():
+            raise RuntimeError("dag.json missing")
+
+        return json.loads(file.read_text())
+
+    def update_status(self, run_dir, status):
+
+        meta_file = run_dir / "run_meta.json"
+
+        meta = json.loads(meta_file.read_text())
+
+        meta["status"] = status
+        meta["updated_at"] = str(datetime.utcnow())
+
+        meta_file.write_text(json.dumps(meta, indent=2))

sourcebot/runtime/dag/scheduler/state_store.py

@@ -0,0 +1,40 @@
+# sourcebot/runtime/dag/scheduler/state_store.py
+import json
+from pathlib import Path
+from datetime import datetime
+
+
+class StateStore:
+
+    def __init__(self, run_dir: Path):
+
+        self.file = run_dir / "state.json"
+
+        if self.file.exists():
+            self.state = json.loads(self.file.read_text())
+        else:
+            self.state = {}
+
+    def save(self):
+
+        self.file.write_text(
+            json.dumps(self.state, indent=2)
+        )
+
+    def status(self, task_id):
+
+        return self.state.get(task_id, {}).get("status")
+
+    def attempts(self, task_id):
+
+        return self.state.get(task_id, {}).get("attempts", 0)
+
+    def update(self, task_id, **fields):
+
+        entry = self.state.setdefault(task_id, {})
+
+        entry.update(fields)
+
+        entry["updated_at"] = str(datetime.utcnow())
+
+        self.save()

sourcebot/runtime/dag/scheduler/task_graph.py

@@ -0,0 +1,29 @@
+# sourcebot/runtime/dag/scheduler/task_graph.py
+from collections import defaultdict
+
+class TaskGraph:
+
+    def __init__(self, tasks):
+
+        self.tasks = {t["id"]: t for t in tasks}
+
+        self.children = defaultdict(list)
+        self.in_degree = {}
+
+        for t in tasks:
+
+            tid = t["id"]
+            deps = t.get("depends_on", [])
+
+            self.in_degree[tid] = len(deps)
+
+            for d in deps:
+                self.children[d].append(tid)
+
+    def roots(self):
+
+        return [
+            tid
+            for tid, deg in self.in_degree.items()
+            if deg == 0
+        ]

sourcebot/runtime/init_system.py

@@ -0,0 +1,182 @@
+# sourcebot/runtime/init_system.py 
+from rich.console import Console
+import sys
+from pathlib import Path
+from sourcebot.context import ContextBuilder
+from sourcebot.config import ConfigManager
+from sourcebot.security.policy import SecurityPolicy
+from sourcebot.bus import EventBus
+from sourcebot.context import MessageBuilder
+from sourcebot.conversation.service import ConversationService
+from sourcebot.session import JsonlSessionRepository, SessionService
+from sourcebot.runtime.dag.scheduler import DAGScheduler
+from sourcebot.runtime.dag.planner import DAGPlanner
+from sourcebot.memory import FileMemoryStore, LLMConsolidator, MemoryService, WindowMemoryPolicy
+from sourcebot.runtime.agent import AgentFactory
+from sourcebot.tools import ToolRegistry, ShellTool, SkillListTool, SkillDetailTool, RuleListTool, RuleDetailTool
+from sourcebot.storage import SkillStorage, RulesLoader
+from sourcebot.docker_sandbox import DockerSandbox
+from sourcebot.llm import LLMClientFactory
+import logging
+logger = logging.getLogger(__name__)
+
+class InitSystem:
+    def __init__(self):
+        self.console = Console(
+                    force_terminal=sys.stdout.isatty(),
+                    soft_wrap=True
+                ) 
+        
+        config_manager = ConfigManager()
+
+        # Config 
+        app_root_path = config_manager.app_root_path
+        skill_dir_path = app_root_path/"skills"
+
+        # Workspace config 
+        workspace_config = config_manager.load_workspace_config() 
+
+        # Workspace config error
+        if workspace_config is None:
+            self.console.print("❌ Error: Workspace configuration not found")
+            self.console.print("Please run the initialization command in the root directory of the workspace first:")
+            self.console.print("[cyan]sourcebot init_workspace[/cyan]")
+            self.console.print("Or ensure the current directory is a valid Sourcebot workspace"
+            )
+            
+            sys.exit(1)
+        # Agent config 
+        main_provider_name = workspace_config.models["main_agent"].provider
+        main_model_name = workspace_config.models["main_agent"].model
+        sub_provider_name = workspace_config.models["sub_agent"].provider
+        sub_model_name = workspace_config.models["sub_agent"].model
+
+        # Host working path 
+        host_workspace = Path(workspace_config.project_root)
+
+        # Docker working path 
+        # The Docker working directory remains unchanged
+        docker_workspace = Path("/workspace")
+        # Skill storage 
+        skill_storage = SkillStorage(
+            root_skills = skill_dir_path,
+            builtin_skills = host_workspace/"skills" # The work directory skill path should not appear in the prompt.
+            )
+
+        # Rules loader 
+        rules_loader = RulesLoader(app_root_path)
+
+        # Context builder 
+        # 🔴 Note that the host machine's workspace address should not appear in the prompt; it should always be docker_workspace.
+        context_builder = ContextBuilder(
+            workspace = docker_workspace,  
+            skill_storage = skill_storage,
+            rules_loader = rules_loader,
+            )
+
+        # Message builder 
+        message_builder = MessageBuilder(context_builder)
+
+        # Security policy 
+        # TODO: Security policy not yet implemented
+        security_policy = SecurityPolicy()
+
+        # Docker sandbox 
+        self.docker_sandbox = DockerSandbox(
+            image = workspace_config.docker_image, 
+            host_workspace = str(host_workspace)
+            )
+        # ====================
+        # Tool registry 
+        # ====================
+        tools = ToolRegistry()
+        # Skill
+        tools.register(
+            SkillListTool(
+                skill_storage = skill_storage,
+            )
+        )
+        tools.register(
+            SkillDetailTool(
+                skill_storage = skill_storage,
+                host_workspace = host_workspace,
+            )
+        )
+        # Rule
+        tools.register(
+            RuleListTool(
+                rules_loader = rules_loader,
+            )
+        )
+        tools.register(
+            RuleDetailTool(
+                rules_loader = rules_loader,
+            )
+        )
+        # Shell 
+        tools.register(
+            ShellTool(
+                sandbox = self.docker_sandbox,
+                timeout = 100,
+            )
+        )
+
+        # Message bus 
+        # TODO: Not yet enabled
+        self.bus = EventBus()
+
+        # Agent factory 
+        agent_factory = AgentFactory(
+            bus = self.bus,
+            tools = tools,
+            message_builder = message_builder,
+            main_provider_name = main_provider_name,
+            main_model_name = main_model_name,
+            config_manager = config_manager,
+            policy = security_policy,
+            max_iterations = 40,
+            sub_provider_name = sub_provider_name,
+            sub_model_name = sub_model_name,
+        )
+        # DAG planner 
+        main_llm = LLMClientFactory.create_client(
+                config_manager = config_manager,
+                provider_name = main_provider_name,
+                model_name = main_model_name
+            )
+        self.dag_planner = DAGPlanner(main_llm, context_builder)
+        # DAG scheduler 
+        self.dag_scheduler = DAGScheduler(
+            agent_factory = agent_factory,
+            workspace = host_workspace
+            )
+
+
+        # Main agent runtime 
+        main_agent = agent_factory.build_main_agent()
+
+        # Memory service 
+        file_memory_store = FileMemoryStore(host_workspace)
+        window_memory_policy = WindowMemoryPolicy()    
+        llm_consolidator = LLMConsolidator(main_llm)
+        memory_service = MemoryService(
+            store = file_memory_store, 
+            policy = window_memory_policy,
+            consolidator = llm_consolidator
+            )
+
+        # Session service 
+        session_repository = JsonlSessionRepository(host_workspace)
+        session_service = SessionService(session_repository)
+        
+        # Conversation service 
+        self.conversation_service = ConversationService(
+            session_service = session_service,
+            memory_service = memory_service,
+            agent = main_agent,
+            message_builder = message_builder,
+            bus = self.bus
+        )
+
+
+

sourcebot/runtime/tool_executor.py

@@ -0,0 +1,30 @@
+# sourcebot/runtime/tool_executor.py
+
+import asyncio
+class ToolExecutor:
+
+    def __init__(self, tools, timeout = 60, retries = 2):
+        self.tools = tools
+        self.timeout = timeout
+        self.retries = retries
+
+    async def execute(self, name: str, args):
+
+        last_error = None
+
+        for attempt in range(self.retries + 1):
+
+            try:
+
+                return await asyncio.wait_for(
+                    self.tools.execute(name, args),
+                    timeout = self.timeout
+                )
+
+            except asyncio.TimeoutError:
+                last_error = f"Tool {name} timed out"
+
+            except Exception as e:
+                last_error = str(e)
+
+        return f"Tool failed after retries: {last_error}"

sourcebot/security/policy.py

@@ -0,0 +1,23 @@
+# domain/security/policy.py
+from typing import Any, Dict, List
+RULES_DB = [
+    {"pattern": "eval(", "safe_replacement": "# eval removed for safety"}
+]
+
+class SecurityPolicy:
+    # TODO: To be implemented
+    def apply_policy(self, code: str) -> str:
+        safe_code = code
+        for rule in RULES_DB:
+            if rule["pattern"] in code:
+                safe_code = safe_code.replace(rule["pattern"], rule["safe_replacement"])
+        return safe_code
+    async def before_llm(self, messages: List[Dict]) -> None:
+        pass
+
+    async def before_tool(self, tool_name: str, args: Dict[str, Any]) -> None:
+        pass
+
+    async def after_tool(self, tool_name: str, result: str) -> None:
+        pass
+