FIRE-815 | Bugfix | Reduce rogue import time on startup

[yuval-qf]

No description provided.

[coderabbitai]

Summary by CodeRabbit

• New Features

  • Exposed additional modules at package level for easier imports: evaluator policy_evaluation, server models, and common agent_sessions/update_checker.
  • Added a theme accessor function to retrieve the UI theme on demand.

• Refactor

  • Broad lazy-loading of heavy dependencies across agents, services, server, and UI to reduce startup time.
  • Unified and streamlined package init imports without changing public behavior.
  • Updated type hints to forward references to avoid runtime imports.
  • Minor session handling/logging adjustments in generic agent execution flow (no behavior change).

Walkthrough

The PR defers imports across multiple modules using TYPE_CHECKING guards and local (in-function) imports, updates type hints to string-based forward references, adds a few package-level re-exports, and adjusts GenericAgentExecutor’s flow for session resolution and part conversions. No new features; primary behavioral change is import timing and minor executor control-flow refinement.

Changes

Cohort / File(s)	Summary
Examples: T‑Shirt store agent lazy imports examples/tshirt_store_agent/__main__.py, examples/tshirt_store_agent/tshirt_store_agent.py, examples/tshirt_store_agent/tshirt_store_agent_executor.py	Moved ADK-related imports into functions; added TYPE_CHECKING blocks and forward-ref annotations; updated executor method signatures to quoted types; no runtime logic change beyond import timing.
Common: Agent model/session + executor updates rogue/common/agent_model_wrapper.py, rogue/common/agent_sessions.py, rogue/common/generic_agent_executor.py, rogue/common/__init__.py	Deferred heavy imports; added TYPE_CHECKING and forward refs; updated return/type annotations to strings; GenericAgentExecutor now resolves session before processing and uses GenAIPart conversions; re-exported agent_sessions and update_checker in package init.
Evaluator agent modules rogue/evaluator_agent/__init__.py, rogue/evaluator_agent/evaluator_agent.py, rogue/evaluator_agent/policy_evaluation.py, rogue/evaluator_agent/run_evaluator_agent.py	Added policy_evaluation export; converted several signatures to quoted types; moved ADK/litellm imports inside functions; removed a helper building an agent card; preserved overall runtime flow.
Prompt injection evaluator rogue/prompt_injection_evaluator/run_prompt_injection_evaluator.py	Deferred datasets and litellm imports into functions; no signature changes.
Server package inits and services rogue/server/__init__.py, rogue/server/api/__init__.py, rogue/server/services/__init__.py, rogue/server/services/interviewer_service.py, rogue/server/services/llm_service.py	Added models to server exports; simplified api init import syntax; reordered services export; moved litellm imports inside service methods.
UI package init and app/theme rogue/ui/__init__.py, rogue/ui/app.py, rogue/ui/config/theme.py	Consolidated imports in init; lazily import gradio in app; replaced theme object with get_theme() factory returning ThemeClass via local imports.
UI components: lazy Gradio + forward types rogue/ui/components/config_screen.py, .../interviewer.py, .../report_generator.py, .../scenario_generator.py, .../scenario_runner.py	Added TYPE_CHECKING guards; moved gradio imports inside functions; switched parameter/return annotations to forward-referenced "State", "Tabs", "JSON", "Markdown".

Sequence Diagram(s)

sequenceDiagram
  autonumber
  actor Client
  participant API as API Layer
  participant Exec as GenericAgentExecutor
  participant Sess as SessionService
  participant Runner as ADK Runner
  participant GenAI as GenAI SDK

  Client->>API: Request (context, event_queue)
  API->>Exec: execute(context, event_queue)
  rect rgb(245,248,255)
    note right of Exec: Lazy import GenAI types
    Exec->>Sess: _upsert_session(context.session_id?)
    Sess-->>Exec: session_id
  end
  Exec->>Exec: convert_a2a_parts_to_genai()
  Exec->>Runner: run(model, content, session_id)
  alt streaming updates
    Runner-->>Exec: Event (update)
    Exec-->>API: forward update
  end
  Runner-->>Exec: Event (final)
  Exec->>Exec: convert_genai_parts_to_a2a()
  Exec-->>API: final event
  API-->>Client: Response stream

Loading

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~25 minutes

Possibly related PRs

• FIRE-751 - report to BFF #37 — Touches examples/tshirt_store_agent/tshirt_store_agent_executor.py with formatting edits; overlaps on the same file updated here.

Suggested reviewers

• drorIvry
• osher-qualifire

Poem

I hopped through imports, light on my feet,
Delaying the heavy, making start-up neat.
Sessions aligned, parts neatly convert,
UI calls Gradio only when alert.
Burrows run faster, ears all a-tilt—
Lazy loads done, with zero guilt. 🐇✨

Pre-merge checks and finishing touches

❌ Failed checks (1 warning, 1 inconclusive)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 11.63% which is insufficient. The required threshold is 80.00%.	You can run @coderabbitai generate docstrings to improve docstring coverage.
Description Check	❓ Inconclusive	The pull request has no description, so it provides no insight into the changes made, making it impossible for reviewers to quickly understand the intent or scope of the modifications.	Please add a concise pull request description that summarizes the purpose of the changes and highlights key modifications to help reviewers evaluate the update.

✅ Passed checks (1 passed)

Check name	Status	Explanation
Title Check	✅ Passed	The title succinctly identifies the primary purpose of the changes—reducing Rogue’s import startup time—and includes the relevant bug identifier, making it clear and directly tied to the modifications in this pull request.

✨ Finishing touches

• 📝 Generate docstrings

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment
• Commit unit tests in branch bugfix/FIRE-815-rogue-import-startup-time

---

📜 Recent review details

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 7c24824 and 9e86c35.

📒 Files selected for processing (25)

• examples/tshirt_store_agent/__main__.py (2 hunks)
• examples/tshirt_store_agent/tshirt_store_agent.py (2 hunks)
• examples/tshirt_store_agent/tshirt_store_agent_executor.py (1 hunks)
• rogue/common/__init__.py (1 hunks)
• rogue/common/agent_model_wrapper.py (1 hunks)
• rogue/common/agent_sessions.py (1 hunks)
• rogue/common/generic_agent_executor.py (7 hunks)
• rogue/evaluator_agent/__init__.py (1 hunks)
• rogue/evaluator_agent/evaluator_agent.py (7 hunks)
• rogue/evaluator_agent/policy_evaluation.py (1 hunks)
• rogue/evaluator_agent/run_evaluator_agent.py (3 hunks)
• rogue/prompt_injection_evaluator/run_prompt_injection_evaluator.py (2 hunks)
• rogue/server/__init__.py (1 hunks)
• rogue/server/api/__init__.py (1 hunks)
• rogue/server/services/__init__.py (1 hunks)
• rogue/server/services/interviewer_service.py (1 hunks)
• rogue/server/services/llm_service.py (3 hunks)
• rogue/ui/__init__.py (1 hunks)
• rogue/ui/app.py (1 hunks)
• rogue/ui/components/config_screen.py (1 hunks)
• rogue/ui/components/interviewer.py (1 hunks)
• rogue/ui/components/report_generator.py (3 hunks)
• rogue/ui/components/scenario_generator.py (1 hunks)
• rogue/ui/components/scenario_runner.py (3 hunks)
• rogue/ui/config/theme.py (1 hunks)

🧰 Additional context used

📓 Path-based instructions (1)

**/*.py

📄 CodeRabbit inference engine (AGENTS.md)

**/*.py: Format Python code with Black
Ensure code passes flake8 linting
Run mypy with the repository configuration for static typing
Run Bandit security checks using .bandit.yaml configuration
Use isort import conventions for import ordering
Add type hints to all function signatures
Follow PEP 8 naming (snake_case for variables/functions, PascalCase for classes)
Use try/except around code that may raise exceptions

Files:

• rogue/ui/config/theme.py
• rogue/server/services/__init__.py
• rogue/evaluator_agent/policy_evaluation.py
• examples/tshirt_store_agent/tshirt_store_agent.py
• rogue/common/agent_model_wrapper.py
• rogue/ui/components/config_screen.py
• rogue/server/services/llm_service.py
• rogue/evaluator_agent/run_evaluator_agent.py
• rogue/server/__init__.py
• rogue/ui/components/scenario_generator.py
• rogue/server/api/__init__.py
• rogue/evaluator_agent/__init__.py
• rogue/ui/app.py
• rogue/server/services/interviewer_service.py
• rogue/common/__init__.py
• rogue/common/agent_sessions.py
• rogue/ui/components/report_generator.py
• rogue/ui/__init__.py
• examples/tshirt_store_agent/__main__.py
• rogue/ui/components/interviewer.py
• rogue/ui/components/scenario_runner.py
• rogue/prompt_injection_evaluator/run_prompt_injection_evaluator.py
• examples/tshirt_store_agent/tshirt_store_agent_executor.py
• rogue/common/generic_agent_executor.py
• rogue/evaluator_agent/evaluator_agent.py

🧬 Code graph analysis (11)

rogue/ui/components/config_screen.py (2)

sdks/python/rogue_sdk/types.py (2)

• AgentConfig (65-91)
• AuthType (22-42)

rogue/common/workdir_utils.py (1)

• dump_config (24-39)

rogue/server/services/llm_service.py (1)

sdks/python/rogue_sdk/types.py (5)

• EvaluationResults (192-208)
• Scenario (94-133)
• Scenarios (136-154)
• ScenarioType (45-49)
• StructuredSummary (444-450)

rogue/evaluator_agent/run_evaluator_agent.py (2)

rogue/common/generic_agent_executor.py (1)

• _run_agent (37-46)

examples/tshirt_store_agent/tshirt_store_agent_executor.py (1)

• _run_agent (38-47)

rogue/ui/components/scenario_generator.py (2)

sdks/python/rogue_sdk/sdk.py (1)

• RogueSDK (34-391)

rogue/common/workdir_utils.py (2)

• dump_business_context (8-13)
• dump_scenarios (16-21)

rogue/server/api/__init__.py (4)

sdks/python/rogue_sdk/client.py (1)

• health (87-90)

sdks/python/rogue_sdk/sdk.py (1)

• health (56-58)

packages/sdk/src/sdk.ts (1)

• health (68-70)

packages/sdk/src/client.ts (1)

• health (84-86)

rogue/ui/app.py (1)

rogue/ui/config/theme.py (1)

• get_theme (7-56)

rogue/ui/components/report_generator.py (2)

sdks/python/rogue_sdk/types.py (1)

• EvaluationResults (192-208)

rogue/server/services/api_format_service.py (1)

• convert_with_structured_summary (116-141)

rogue/ui/components/interviewer.py (1)

rogue/common/workdir_utils.py (1)

• dump_business_context (8-13)

rogue/prompt_injection_evaluator/run_prompt_injection_evaluator.py (1)

sdks/python/rogue_sdk/types.py (1)

• get_auth_header (30-42)

examples/tshirt_store_agent/tshirt_store_agent_executor.py (2)

rogue/common/generic_agent_executor.py (1)

• _run_agent (37-46)

rogue/evaluator_agent/run_evaluator_agent.py (1)

• _run_agent (18-72)

rogue/common/generic_agent_executor.py (2)

examples/tshirt_store_agent/tshirt_store_agent_executor.py (6)

• _run_agent (38-47)
• convert_a2a_parts_to_genai (143-145)
• cancel (110-112)
• convert_a2a_part_to_genai (148-169)
• convert_genai_parts_to_a2a (172-179)
• convert_genai_part_to_a2a (182-206)

rogue/evaluator_agent/run_evaluator_agent.py (1)

• _run_agent (18-72)

🔇 Additional comments (35)

rogue/server/services/interviewer_service.py (1)

68-69: LGTM—lazy import correctly defers litellm loading. Confirm that CI’s mypy run still passes with this pattern.

rogue/server/services/llm_service.py (2)

5-11: LGTM!

Clean consolidation of imports from rogue_sdk.types. This improves readability and aligns with import organization best practices.

---

176-177: Approve lazy import; verify types manually
The runtime import of litellm.completion defers startup cost and is caught by existing exception handlers if missing.

Optional: extract the duplicated import in generate_scenarios (176-177) and generate_summary_from_results (231-232) into a helper if you prefer less repetition.

Please run mypy (or your CI type checks) locally to confirm no type‐checking errors arise from these runtime imports.

rogue/evaluator_agent/policy_evaluation.py (1)

144-145: LGTM! Lazy import reduces startup time.

The local import of litellm.completion is appropriate for deferring the heavy import until actually needed. The comment clearly explains the rationale.

rogue/server/services/__init__.py (1)

8-8: LGTM! Import reordering has no functional impact.

The reordering of scenario_evaluation_service in the import list does not affect functionality.

rogue/evaluator_agent/__init__.py (1)

1-1: LGTM! Package-level export added.

Adding policy_evaluation to the package exports makes it accessible at rogue.evaluator_agent.policy_evaluation, which is appropriate for the new module.

rogue/prompt_injection_evaluator/run_prompt_injection_evaluator.py (2)

80-81: LGTM! Lazy import reduces startup time.

The local import of litellm.completion appropriately defers the heavy import until needed. The comment clearly explains the rationale.

---

125-129: LGTM! Lazy import reduces startup time.

The local import of load_dataset appropriately defers the heavy import until needed, and the usage on line 129 correctly references the locally imported function.

rogue/ui/components/interviewer.py (3)

2-2: LGTM! TYPE_CHECKING guard avoids runtime import.

The TYPE_CHECKING import and guarded State/Tabs imports correctly defer gradio imports to type-checking time only, reducing runtime import overhead.

Also applies to: 9-10

---

14-15: LGTM! Forward-referenced types support lazy loading.

The string-based forward references for State and Tabs parameters correctly support the lazy import pattern while maintaining type-checking capabilities.

---

17-18: LGTM! Lazy import reduces startup time.

The local import of gradio appropriately defers the heavy import until the function is called. The comment clearly explains the rationale.

rogue/ui/components/scenario_runner.py (3)

3-3: LGTM! TYPE_CHECKING guard avoids runtime import.

The TYPE_CHECKING import and guarded State/Tabs imports correctly defer gradio imports to type-checking time only.

Also applies to: 19-20

---

47-47: LGTM! Forward-referenced types support lazy loading.

The string-based forward references for State and Tabs parameters correctly support the lazy import pattern.

---

48-49: LGTM! Lazy import reduces startup time.

The local import of gradio appropriately defers the heavy import until needed.

rogue/common/agent_model_wrapper.py (3)

2-2: LGTM! TYPE_CHECKING guard avoids runtime import.

The TYPE_CHECKING import and guarded BaseLlm import correctly defer google.adk.models imports to type-checking time only.

Also applies to: 6-7

---

14-14: LGTM! Forward-referenced return type supports lazy loading.

The string-based forward reference for the return type correctly supports the lazy import pattern. The @lru_cache decorator works correctly with forward-referenced types since it doesn't inspect type annotations at runtime.

---

15-17: LGTM! Lazy imports reduce startup time.

The local imports of LLMRegistry and LiteLlm appropriately defer the heavy google.adk imports until needed. The comment clearly explains the rationale.

rogue/ui/components/scenario_generator.py (3)

2-2: LGTM! TYPE_CHECKING guard avoids runtime import.

The TYPE_CHECKING import and guarded State/Tabs imports correctly defer gradio imports to type-checking time only.

Also applies to: 9-10

---

13-13: LGTM! Forward-referenced types support lazy loading.

The string-based forward references for State and Tabs parameters correctly support the lazy import pattern.

---

14-15: LGTM! Lazy import reduces startup time.

The local import of gradio appropriately defers the heavy import until needed.

rogue/server/__init__.py (1)

8-8: LGTM! Clean package export expansion.

The addition of models to the export list properly exposes the new models namespace, consistent with the modular package structure.

rogue/common/agent_sessions.py (2)

1-5: LGTM! Proper lazy import pattern.

The TYPE_CHECKING guard correctly defers the google.adk.sessions import to type-checking time only, reducing runtime import overhead while preserving type safety.

---

8-12: LGTM! Forward references applied correctly.

The string-based type annotations for session_service and return type properly reference types that are only imported under TYPE_CHECKING, maintaining type safety without runtime import cost.

examples/tshirt_store_agent/__main__.py (1)

23-27: LGTM! Function-scoped imports reduce startup time.

Moving the google.adk imports into main() properly defers the heavy import overhead until the function is actually executed. The explanatory comment provides good context for future maintainers.

rogue/ui/__init__.py (1)

1-1: LGTM! Clean import consolidation.

Combining the four separate import statements into a single line improves readability without altering the public API surface.

examples/tshirt_store_agent/tshirt_store_agent.py (2)

2-5: LGTM! TYPE_CHECKING guard for lazy import.

The TYPE_CHECKING block properly exposes LlmAgent for type checking without incurring the runtime import cost.

---

76-80: LGTM! Function-scoped imports with clear rationale.

Moving the google.adk imports into the function body defers the heavy imports until create_tshirt_store_agent() is called. The forward-referenced return type and explanatory comment provide good context.

rogue/ui/config/theme.py (2)

1-4: LGTM! TYPE_CHECKING guard for theme type.

The TYPE_CHECKING block properly exposes ThemeClass for type hints without importing gradio at module load time.

---

7-56: LGTM! Lazy theme construction reduces startup time.

The get_theme() function properly defers the gradio import and theme construction until the function is called, significantly reducing module-level import overhead. The explanatory comment and forward-referenced return type provide good context.

rogue/ui/app.py (2)

15-15: LGTM! Import updated for lazy theme loading.

The import change from theme to get_theme properly aligns with the lazy theme construction pattern introduced in rogue/ui/config/theme.py.

---

18-22: LGTM! Lazy gradio import and theme initialization.

Moving the gradio import into get_app() and calling get_theme() properly defers both the gradio import and theme construction until the function is invoked. The explanatory comment provides good context.

rogue/evaluator_agent/run_evaluator_agent.py (3)

3-3: LGTM! TYPE_CHECKING guards for ADK types.

The TYPE_CHECKING block properly exposes Runner and Session for type checking without importing them at runtime, consistent with the lazy import pattern used throughout this PR.

Also applies to: 13-15

---

18-22: LGTM! Forward-referenced type hints.

The string-based type annotations for agent_runner and session parameters correctly reference types that are only imported under TYPE_CHECKING, maintaining type safety without runtime import cost.

---

85-87: LGTM! Function-scoped imports reduce startup overhead.

Moving the google.adk imports into arun_evaluator_agent() properly defers the heavy import overhead until the function is executed. The explanatory comment provides clear context for the lazy-loading strategy.

rogue/common/generic_agent_executor.py (1)

89-108: Confirm availability of google.genai.types.UserContent.

UserContent only exists in recent google.genai builds; older versions ship just types.Content. If our environment is still pinned to an earlier release, this lazy import will raise at runtime. Please confirm the deployed version exposes UserContent, or keep using types.Content(role="user", …) as a safe fallback.

---

_{Comment @coderabbitai help to get the list of available commands and usage tips.}