main_mock.py

#!/usr/bin/env python3
import json
import os
import sys
import traceback
from pathlib import Path
from typing import Optional
from unittest.mock import patch

from src.core.action import TaskCreateAction
from src.core.action.actions import ReportAction
from src.core.action.handlers import ReportActionHandler
from src.core.agent.agent import Agent
from src.core.agent.subagent_task import AgentTask
from src.core.backend.command_env_executor import get_docker_executor
from src.core.bash.factory import get_bash_handlers
from src.core.context import ContextStore
from src.core.file import get_file_handlers
from src.core.llm import LlmConfig
from src.core.middleware import (
    LoggingMiddleware,
    ErrorRecoveryMiddleware,
    ActionOutputTruncationMiddleware,
    TracingMiddleware,
    SubagentTaskBootstrapMiddleware,
    SubagentTurnCompletionMiddleware,
)
from src.core.orchestrator.orchestrator_session_history_middleware import OrchestratorSessionHistoryMiddleware
from src.core.orchestrator.orchestrator_session_prompt_middleware import OrchestratorSessionPromptMiddleware
from src.core.orchestrator.session_history import SessionHistory
from src.core.orchestrator.turn_history import TurnHistory
from src.core.task import create_task_manager, TaskStore
from src.core.task.create_task_handler import CreateTaskActionHandler
from src.core.task.subagent_luncher import AgentLauncher
from src.ext.subagent_report import SubagentReportMiddleware
from src.misc import pretty_log, PrettyLogger
from src.system_msgs.system_msg_loader import load_orchestrator_system_message, load_explorer_system_message, load_coder_system_message

# Add src to path
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
LLM_RESPONSES_DIR = Path(__file__).parent.parent / "llm_responses"

ACTION_OUTPUT_MAX_CHARS = 1_000

task_instruction = (
    """Create and run a server on port 3000 that has a single GET endpoint: /fib.

It should expect a query param /fib?n={some number} and return the nth Fibonacci number as a JSON object with a key 'result'.

If the query param is not provided, it should return a 400 Bad Request error.
If the query param is not an integer, it should return a 400 Bad Request error.

Automatically choose to install any dependencies if it is required to develop the server
"""
)


def load_recorded_responses(responses_dir: Path) -> list[str]:
    """Load all recorded LLM responses sorted by filename (chronological order)."""
    response_files = sorted(responses_dir.glob("response_*.json"))
    if not response_files:
        raise FileNotFoundError(f"No response files found in {responses_dir}")

    responses = []
    for filepath in response_files:
        with open(filepath, "r", encoding="utf-8") as f:
            data = json.load(f)
        responses.append(data["content"])

    pretty_log.info(f"Loaded {len(responses)} recorded LLM responses from {responses_dir}")
    return responses


def create_sequential_mock(responses: list[str]):
    """Create a mock function that returns responses sequentially."""
    call_index = [0]

    def mock_get_llm_response(messages, llm_config, api_base=None, max_retries=10):
        idx = call_index[0]
        if idx >= len(responses):
            raise RuntimeError(
                f"Ran out of recorded responses: requested call #{idx + 1} "
                f"but only {len(responses)} responses available"
            )
        content = responses[idx]
        call_index[0] += 1
        pretty_log.debug(f"[MOCK] Returning recorded response #{idx + 1}/{len(responses)}")
        return content

    return mock_get_llm_response


def initialize_orchestrator_and_run_task():
    """Initialize the orchestrator agent and run the task."""
    responses = load_recorded_responses(LLM_RESPONSES_DIR)
    mock_fn = create_sequential_mock(responses)
    pretty_log.section_header("Initializing Code Assistant")
    pretty_log.info("User input: " + task_instruction)

    # ("anthropic/claude-sonnet-4-20250514", 0.1),
    # ("openrouter/qwen/qwen3-coder", 0.1),
    llm_config = LlmConfig(
        model="openai/gpt-4.1-2025-04-14",
        temperature=1,
        max_tokens=2000,
    )
    this_dir_path: Path = Path(__file__).parent.resolve()
    logging_dir = Path(this_dir_path) / "tracing_logs"
    subagents = get_subagents(llm_config, logging_dir)
    context_store = ContextStore()
    task_store = TaskStore()
    task_manager = create_task_manager(task_store, context_store)
    agent_launcher = AgentLauncher(task_manager, context_store, subagents)
    create_task_handler = CreateTaskActionHandler(task_manager, agent_launcher)
    session_history = SessionHistory(
        task_store=task_store,
        context_store=context_store,
        turn_history=TurnHistory()
    )

    actions = {
        TaskCreateAction: create_task_handler.handle,
    }
    orchestrator_agent = Agent(
        agent_name="orchestrator",
        system_prompt=load_orchestrator_system_message(),
        actions=actions,
        llm_config=llm_config,
        middlewares=[
            OrchestratorSessionPromptMiddleware(session_history, load_orchestrator_system_message()),
            OrchestratorSessionHistoryMiddleware(session_history)
        ]
    )
    with patch("src.core.agent.agent.get_llm_response", mock_fn):
        result = orchestrator_agent.run_task(AgentTask(
            task_id="",
            instruction=task_instruction,
            agent_name="orchestrator",
        ), max_turns=5)
        pretty_log.info(f"task result: {result}")

    return "SUCCESS"


def get_subagents(llm_config: LlmConfig, logging_dir: Optional[Path] = None) -> dict[str, Agent]:
    executor = get_docker_executor()
    bash_actions = get_bash_handlers(executor)
    files_actions = get_file_handlers(executor)
    bash_actions[ReportAction] = ReportActionHandler().handle
    subagent_middlewares = [
        SubagentTaskBootstrapMiddleware(),
        LoggingMiddleware(),
        ErrorRecoveryMiddleware(),
        ActionOutputTruncationMiddleware(max_chars=ACTION_OUTPUT_MAX_CHARS),
        TracingMiddleware(logging_dir),
        SubagentReportMiddleware(),
        SubagentTurnCompletionMiddleware(),
    ]
    subagents = {
        "explorer": Agent(
            system_prompt=load_explorer_system_message(),
            actions=files_actions | bash_actions,
            agent_name="explorer",
            llm_config=llm_config,
            middlewares=subagent_middlewares,
        ),
        "coder": Agent(
            system_prompt=load_coder_system_message(),
            actions=files_actions | bash_actions,
            agent_name="coder",
            llm_config=llm_config,
            middlewares=subagent_middlewares,
        ),
    }
    return subagents


def main():
    if not os.getenv("LITE_LLM_API_KEY") and not os.getenv("LITELLM_API_KEY"):
        pretty_log.error("Environment variable LITE_LLM_API_KEY or LITELLM_API_KEY is required to run the test.")
        return

    results = []
    try:
        result = initialize_orchestrator_and_run_task()
        results.append(("Test 1", "SUCCESS", result))
    except Exception as e:
        results.append(("Test 1", "FAILED", f"{e}\n{traceback.format_exc()}"))

    pretty_log.section_header("TASK RESULTS")
    for test_name, status, details in results:
        pretty_log.info(f"{test_name}: {status}")
        if status == "FAILED":
            pretty_log.info(f"Details: {details}")


PrettyLogger.PRINT_DEBUG = True
if __name__ == "__main__":
    main()