Claude commited on
Commit
e82d9c9
Β·
1 Parent(s): 2cbcbfd

refactor: Organize orchestrators into dedicated package

Browse files

- Create src/orchestrators/ package with proper module structure
- base.py: Shared protocols (SearchHandlerProtocol, JudgeHandlerProtocol)
- simple.py: Basic search-judge loop orchestrator (was orchestrator.py)
- advanced.py: Multi-agent MS Agent Framework orchestrator (was orchestrator_magentic.py)
- hierarchical.py: Sub-iteration middleware orchestrator
- factory.py: Factory pattern for orchestrator creation
- __init__.py: Clean public API with facade pattern

Design principles applied:
- Single Responsibility (SRP): Each file has one clear job
- Interface Segregation (ISP): Protocols in base.py
- Factory Pattern (GoF): Centralized creation logic
- Facade Pattern: Clean public imports via __init__.py

All imports updated across:
- src/app.py, src/agents/
- tests/unit/, tests/e2e/, tests/integration/
- examples/orchestrator_demo/

All 147 tests pass, linting and typecheck clean.

examples/orchestrator_demo/run_agent.py CHANGED
@@ -23,7 +23,7 @@ import os
23
  import sys
24
 
25
  from src.agent_factory.judges import JudgeHandler
26
- from src.orchestrator import Orchestrator
27
  from src.tools.clinicaltrials import ClinicalTrialsTool
28
  from src.tools.europepmc import EuropePMCTool
29
  from src.tools.pubmed import PubMedTool
 
23
  import sys
24
 
25
  from src.agent_factory.judges import JudgeHandler
26
+ from src.orchestrators import Orchestrator
27
  from src.tools.clinicaltrials import ClinicalTrialsTool
28
  from src.tools.europepmc import EuropePMCTool
29
  from src.tools.pubmed import PubMedTool
examples/orchestrator_demo/run_magentic.py CHANGED
@@ -17,7 +17,7 @@ import os
17
  import sys
18
 
19
  from src.agent_factory.judges import JudgeHandler
20
- from src.orchestrator_factory import create_orchestrator
21
  from src.tools.clinicaltrials import ClinicalTrialsTool
22
  from src.tools.europepmc import EuropePMCTool
23
  from src.tools.pubmed import PubMedTool
 
17
  import sys
18
 
19
  from src.agent_factory.judges import JudgeHandler
20
+ from src.orchestrators import create_orchestrator
21
  from src.tools.clinicaltrials import ClinicalTrialsTool
22
  from src.tools.europepmc import EuropePMCTool
23
  from src.tools.pubmed import PubMedTool
src/agents/judge_agent.py CHANGED
@@ -12,7 +12,7 @@ from agent_framework import (
12
  Role,
13
  )
14
 
15
- from src.orchestrator import JudgeHandlerProtocol
16
  from src.utils.models import Evidence, JudgeAssessment
17
 
18
 
 
12
  Role,
13
  )
14
 
15
+ from src.orchestrators import JudgeHandlerProtocol
16
  from src.utils.models import Evidence, JudgeAssessment
17
 
18
 
src/agents/search_agent.py CHANGED
@@ -10,7 +10,7 @@ from agent_framework import (
10
  Role,
11
  )
12
 
13
- from src.orchestrator import SearchHandlerProtocol
14
  from src.utils.models import Citation, Evidence, SearchResult
15
 
16
  if TYPE_CHECKING:
 
10
  Role,
11
  )
12
 
13
+ from src.orchestrators import SearchHandlerProtocol
14
  from src.utils.models import Citation, Evidence, SearchResult
15
 
16
  if TYPE_CHECKING:
src/app.py CHANGED
@@ -11,7 +11,7 @@ from pydantic_ai.providers.anthropic import AnthropicProvider
11
  from pydantic_ai.providers.openai import OpenAIProvider
12
 
13
  from src.agent_factory.judges import HFInferenceJudgeHandler, JudgeHandler, MockJudgeHandler
14
- from src.orchestrator_factory import create_orchestrator
15
  from src.tools.clinicaltrials import ClinicalTrialsTool
16
  from src.tools.europepmc import EuropePMCTool
17
  from src.tools.openalex import OpenAlexTool
 
11
  from pydantic_ai.providers.openai import OpenAIProvider
12
 
13
  from src.agent_factory.judges import HFInferenceJudgeHandler, JudgeHandler, MockJudgeHandler
14
+ from src.orchestrators import create_orchestrator
15
  from src.tools.clinicaltrials import ClinicalTrialsTool
16
  from src.tools.europepmc import EuropePMCTool
17
  from src.tools.openalex import OpenAlexTool
src/orchestrators/__init__.py ADDED
@@ -0,0 +1,71 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Orchestrators package - provides different orchestration strategies.
2
+
3
+ This package implements the Strategy Pattern, allowing the application
4
+ to switch between different orchestration approaches:
5
+
6
+ - Simple: Basic search-judge loop using pydantic-ai (free tier compatible)
7
+ - Advanced: Multi-agent coordination using Microsoft Agent Framework
8
+ - Hierarchical: Sub-iteration middleware with fine-grained control
9
+
10
+ Usage:
11
+ from src.orchestrators import create_orchestrator, Orchestrator
12
+
13
+ # Auto-detect mode based on available API keys
14
+ orchestrator = create_orchestrator(search_handler, judge_handler)
15
+
16
+ # Or explicitly specify mode
17
+ orchestrator = create_orchestrator(mode="advanced", api_key="sk-...")
18
+
19
+ Protocols:
20
+ from src.orchestrators import SearchHandlerProtocol, JudgeHandlerProtocol
21
+
22
+ Design Patterns Applied:
23
+ - Factory Pattern: create_orchestrator() creates appropriate orchestrator
24
+ - Strategy Pattern: Different orchestrators implement different strategies
25
+ - Facade Pattern: This __init__.py provides a clean public API
26
+ """
27
+
28
+ # Protocols (Interface Segregation Principle)
29
+ from src.orchestrators.base import JudgeHandlerProtocol, SearchHandlerProtocol
30
+
31
+ # Factory (creational pattern)
32
+ from src.orchestrators.factory import create_orchestrator
33
+
34
+ # Orchestrators (Strategy Pattern implementations)
35
+ from src.orchestrators.simple import Orchestrator
36
+
37
+ # Lazy imports for optional dependencies
38
+ # These are not imported at module level to avoid breaking simple mode
39
+ # when agent-framework-core is not installed
40
+
41
+
42
+ def get_advanced_orchestrator() -> type:
43
+ """Get the AdvancedOrchestrator class (requires agent-framework-core)."""
44
+ from src.orchestrators.advanced import AdvancedOrchestrator
45
+
46
+ return AdvancedOrchestrator
47
+
48
+
49
+ def get_hierarchical_orchestrator() -> type:
50
+ """Get the HierarchicalOrchestrator class (requires agent-framework-core)."""
51
+ from src.orchestrators.hierarchical import HierarchicalOrchestrator
52
+
53
+ return HierarchicalOrchestrator
54
+
55
+
56
+ # Backwards compatibility aliases
57
+ # TODO: Remove after migration period
58
+ def get_magentic_orchestrator() -> type:
59
+ """Deprecated: Use get_advanced_orchestrator() instead."""
60
+ return get_advanced_orchestrator()
61
+
62
+
63
+ __all__ = [
64
+ "JudgeHandlerProtocol",
65
+ "Orchestrator",
66
+ "SearchHandlerProtocol",
67
+ "create_orchestrator",
68
+ "get_advanced_orchestrator",
69
+ "get_hierarchical_orchestrator",
70
+ "get_magentic_orchestrator",
71
+ ]
src/{orchestrator_magentic.py β†’ orchestrators/advanced.py} RENAMED
@@ -1,4 +1,18 @@
1
- """Magentic-based orchestrator using ChatAgent pattern."""
 
 
 
 
 
 
 
 
 
 
 
 
 
 
2
 
3
  import asyncio
4
  from collections.abc import AsyncGenerator
@@ -32,12 +46,18 @@ if TYPE_CHECKING:
32
  logger = structlog.get_logger()
33
 
34
 
35
- class MagenticOrchestrator:
36
  """
37
- Magentic-based orchestrator using ChatAgent pattern.
38
 
39
  Each agent has an internal LLM that understands natural language
40
  instructions from the manager and can call tools appropriately.
 
 
 
 
 
 
41
  """
42
 
43
  def __init__(
@@ -90,7 +110,7 @@ class MagenticOrchestrator:
90
  return None
91
 
92
  def _build_workflow(self) -> Any:
93
- """Build the Magentic workflow with ChatAgent participants."""
94
  # Create agents with internal LLMs
95
  search_agent = create_search_agent(self._chat_client)
96
  judge_agent = create_judge_agent(self._chat_client)
@@ -122,7 +142,7 @@ class MagenticOrchestrator:
122
 
123
  async def run(self, query: str) -> AsyncGenerator[AgentEvent, None]:
124
  """
125
- Run the Magentic workflow.
126
 
127
  Args:
128
  query: User's research question
@@ -130,11 +150,11 @@ class MagenticOrchestrator:
130
  Yields:
131
  AgentEvent objects for real-time UI updates
132
  """
133
- logger.info("Starting Magentic orchestrator", query=query)
134
 
135
  yield AgentEvent(
136
  type="started",
137
- message=f"Starting research (Magentic mode): {query}",
138
  iteration=0,
139
  )
140
 
@@ -221,7 +241,7 @@ The final output should be a structured research report."""
221
  )
222
 
223
  except Exception as e:
224
- logger.error("Magentic workflow failed", error=str(e))
225
  yield AgentEvent(
226
  type="error",
227
  message=f"Workflow error: {e!s}",
@@ -317,3 +337,8 @@ The final output should be a structured research report."""
317
  )
318
 
319
  return None
 
 
 
 
 
 
1
+ """Advanced Orchestrator using Microsoft Agent Framework.
2
+
3
+ This orchestrator uses the ChatAgent pattern from Microsoft's agent-framework-core
4
+ package for multi-agent coordination. It provides richer orchestration capabilities
5
+ including specialized agents (Search, Hypothesis, Judge, Report) coordinated by
6
+ a manager agent.
7
+
8
+ Note: Previously named 'orchestrator_magentic.py' - renamed to eliminate confusion
9
+ with the 'magentic' PyPI package (which is a different library).
10
+
11
+ Design Patterns:
12
+ - Mediator: Manager agent coordinates between specialized agents
13
+ - Strategy: Different agents implement different strategies for their tasks
14
+ - Observer: Event stream allows UI to observe progress
15
+ """
16
 
17
  import asyncio
18
  from collections.abc import AsyncGenerator
 
46
  logger = structlog.get_logger()
47
 
48
 
49
+ class AdvancedOrchestrator:
50
  """
51
+ Advanced orchestrator using Microsoft Agent Framework ChatAgent pattern.
52
 
53
  Each agent has an internal LLM that understands natural language
54
  instructions from the manager and can call tools appropriately.
55
+
56
+ This orchestrator provides:
57
+ - Multi-agent coordination (Search, Hypothesis, Judge, Report)
58
+ - Manager agent for workflow orchestration
59
+ - Streaming events for real-time UI updates
60
+ - Configurable timeouts and round limits
61
  """
62
 
63
  def __init__(
 
110
  return None
111
 
112
  def _build_workflow(self) -> Any:
113
+ """Build the workflow with ChatAgent participants."""
114
  # Create agents with internal LLMs
115
  search_agent = create_search_agent(self._chat_client)
116
  judge_agent = create_judge_agent(self._chat_client)
 
142
 
143
  async def run(self, query: str) -> AsyncGenerator[AgentEvent, None]:
144
  """
145
+ Run the workflow.
146
 
147
  Args:
148
  query: User's research question
 
150
  Yields:
151
  AgentEvent objects for real-time UI updates
152
  """
153
+ logger.info("Starting Advanced orchestrator", query=query)
154
 
155
  yield AgentEvent(
156
  type="started",
157
+ message=f"Starting research (Advanced mode): {query}",
158
  iteration=0,
159
  )
160
 
 
241
  )
242
 
243
  except Exception as e:
244
+ logger.error("Workflow failed", error=str(e))
245
  yield AgentEvent(
246
  type="error",
247
  message=f"Workflow error: {e!s}",
 
337
  )
338
 
339
  return None
340
+
341
+
342
+ # Backwards compatibility alias
343
+ # TODO: Remove after all imports are updated
344
+ MagenticOrchestrator = AdvancedOrchestrator
src/orchestrators/base.py ADDED
@@ -0,0 +1,52 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Base protocols and shared types for orchestrators.
2
+
3
+ This module defines the interfaces that orchestrators depend on,
4
+ following the Interface Segregation Principle (ISP) and
5
+ Dependency Inversion Principle (DIP).
6
+ """
7
+
8
+ from typing import Protocol
9
+
10
+ from src.utils.models import Evidence, JudgeAssessment, SearchResult
11
+
12
+
13
+ class SearchHandlerProtocol(Protocol):
14
+ """Protocol for search handler.
15
+
16
+ Defines the interface for executing searches across biomedical databases.
17
+ Implementations include SearchHandler (scatter-gather across PubMed,
18
+ ClinicalTrials.gov, Europe PMC).
19
+ """
20
+
21
+ async def execute(self, query: str, max_results_per_tool: int = 10) -> SearchResult:
22
+ """Execute a search query.
23
+
24
+ Args:
25
+ query: The search query string
26
+ max_results_per_tool: Maximum results to fetch per search tool
27
+
28
+ Returns:
29
+ SearchResult containing evidence and metadata
30
+ """
31
+ ...
32
+
33
+
34
+ class JudgeHandlerProtocol(Protocol):
35
+ """Protocol for judge handler.
36
+
37
+ Defines the interface for assessing evidence quality and sufficiency.
38
+ Implementations include JudgeHandler (pydantic-ai), HFInferenceJudgeHandler,
39
+ and MockJudgeHandler.
40
+ """
41
+
42
+ async def assess(self, question: str, evidence: list[Evidence]) -> JudgeAssessment:
43
+ """Assess whether collected evidence is sufficient.
44
+
45
+ Args:
46
+ question: The original research question
47
+ evidence: List of evidence items to assess
48
+
49
+ Returns:
50
+ JudgeAssessment with sufficiency determination and next steps
51
+ """
52
+ ...
src/{orchestrator_factory.py β†’ orchestrators/factory.py} RENAMED
@@ -1,24 +1,37 @@
1
- """Factory for creating orchestrators."""
 
 
 
 
 
 
 
 
 
2
 
3
  from typing import Any, Literal
4
 
5
  import structlog
6
 
7
- from src.orchestrator import JudgeHandlerProtocol, Orchestrator, SearchHandlerProtocol
 
8
  from src.utils.config import settings
9
  from src.utils.models import OrchestratorConfig
10
 
11
  logger = structlog.get_logger()
12
 
13
 
14
- def _get_magentic_orchestrator_class() -> Any:
15
- """Import MagenticOrchestrator lazily to avoid hard dependency."""
 
 
 
16
  try:
17
- from src.orchestrator_magentic import MagenticOrchestrator
18
 
19
- return MagenticOrchestrator
20
  except ImportError as e:
21
- logger.error("Failed to import MagenticOrchestrator", error=str(e))
22
  raise ValueError(
23
  "Advanced mode requires agent-framework-core. Please install it or use mode='simple'."
24
  ) from e
@@ -28,33 +41,46 @@ def create_orchestrator(
28
  search_handler: SearchHandlerProtocol | None = None,
29
  judge_handler: JudgeHandlerProtocol | None = None,
30
  config: OrchestratorConfig | None = None,
31
- mode: Literal["simple", "magentic", "advanced"] | None = None,
32
  api_key: str | None = None,
33
  ) -> Any:
34
  """
35
  Create an orchestrator instance.
36
 
 
 
 
 
37
  Args:
38
  search_handler: The search handler (required for simple mode)
39
  judge_handler: The judge handler (required for simple mode)
40
  config: Optional configuration
41
- mode: "simple", "magentic", "advanced" or None (auto-detect)
 
42
  api_key: Optional API key for advanced mode (OpenAI)
43
 
44
  Returns:
45
  Orchestrator instance
 
 
 
 
46
  """
47
  effective_mode = _determine_mode(mode, api_key)
48
  logger.info("Creating orchestrator", mode=effective_mode)
49
 
50
  if effective_mode == "advanced":
51
- orchestrator_cls = _get_magentic_orchestrator_class()
52
  return orchestrator_cls(
53
  max_rounds=config.max_iterations if config else 10,
54
  api_key=api_key,
55
- timeout_seconds=settings.magentic_timeout,
56
  )
57
 
 
 
 
 
 
58
  # Simple mode requires handlers
59
  if search_handler is None or judge_handler is None:
60
  raise ValueError("Simple mode requires search_handler and judge_handler")
@@ -67,10 +93,24 @@ def create_orchestrator(
67
 
68
 
69
  def _determine_mode(explicit_mode: str | None, api_key: str | None) -> str:
70
- """Determine which mode to use."""
 
 
 
 
 
 
 
 
 
 
 
 
71
  if explicit_mode:
72
  if explicit_mode in ("magentic", "advanced"):
73
  return "advanced"
 
 
74
  return "simple"
75
 
76
  # Auto-detect: advanced if paid API key available
 
1
+ """Factory for creating orchestrators.
2
+
3
+ Implements the Factory Pattern (GoF) for creating the appropriate
4
+ orchestrator based on configuration and available credentials.
5
+
6
+ Design Principles:
7
+ - Open/Closed: Easy to add new orchestrator types without modifying existing code
8
+ - Dependency Inversion: Returns protocol-compatible objects, not concrete types
9
+ - Single Responsibility: Only handles orchestrator creation logic
10
+ """
11
 
12
  from typing import Any, Literal
13
 
14
  import structlog
15
 
16
+ from src.orchestrators.base import JudgeHandlerProtocol, SearchHandlerProtocol
17
+ from src.orchestrators.simple import Orchestrator
18
  from src.utils.config import settings
19
  from src.utils.models import OrchestratorConfig
20
 
21
  logger = structlog.get_logger()
22
 
23
 
24
+ def _get_advanced_orchestrator_class() -> Any:
25
+ """Import AdvancedOrchestrator lazily to avoid hard dependency.
26
+
27
+ This allows the simple mode to work without agent-framework-core installed.
28
+ """
29
  try:
30
+ from src.orchestrators.advanced import AdvancedOrchestrator
31
 
32
+ return AdvancedOrchestrator
33
  except ImportError as e:
34
+ logger.error("Failed to import AdvancedOrchestrator", error=str(e))
35
  raise ValueError(
36
  "Advanced mode requires agent-framework-core. Please install it or use mode='simple'."
37
  ) from e
 
41
  search_handler: SearchHandlerProtocol | None = None,
42
  judge_handler: JudgeHandlerProtocol | None = None,
43
  config: OrchestratorConfig | None = None,
44
+ mode: Literal["simple", "magentic", "advanced", "hierarchical"] | None = None,
45
  api_key: str | None = None,
46
  ) -> Any:
47
  """
48
  Create an orchestrator instance.
49
 
50
+ This factory automatically selects the appropriate orchestrator based on:
51
+ 1. Explicit mode parameter (if provided)
52
+ 2. Available API keys (auto-detection)
53
+
54
  Args:
55
  search_handler: The search handler (required for simple mode)
56
  judge_handler: The judge handler (required for simple mode)
57
  config: Optional configuration
58
+ mode: "simple", "magentic", "advanced", "hierarchical" or None (auto-detect)
59
+ Note: "magentic" is an alias for "advanced" (kept for backwards compatibility)
60
  api_key: Optional API key for advanced mode (OpenAI)
61
 
62
  Returns:
63
  Orchestrator instance
64
+
65
+ Raises:
66
+ ValueError: If required handlers are missing for simple mode
67
+ ValueError: If advanced mode is requested but dependencies are missing
68
  """
69
  effective_mode = _determine_mode(mode, api_key)
70
  logger.info("Creating orchestrator", mode=effective_mode)
71
 
72
  if effective_mode == "advanced":
73
+ orchestrator_cls = _get_advanced_orchestrator_class()
74
  return orchestrator_cls(
75
  max_rounds=config.max_iterations if config else 10,
76
  api_key=api_key,
 
77
  )
78
 
79
+ if effective_mode == "hierarchical":
80
+ from src.orchestrators.hierarchical import HierarchicalOrchestrator
81
+
82
+ return HierarchicalOrchestrator()
83
+
84
  # Simple mode requires handlers
85
  if search_handler is None or judge_handler is None:
86
  raise ValueError("Simple mode requires search_handler and judge_handler")
 
93
 
94
 
95
  def _determine_mode(explicit_mode: str | None, api_key: str | None) -> str:
96
+ """Determine which mode to use.
97
+
98
+ Priority:
99
+ 1. Explicit mode parameter
100
+ 2. Auto-detect based on available API keys
101
+
102
+ Args:
103
+ explicit_mode: Mode explicitly requested by caller
104
+ api_key: API key provided by caller
105
+
106
+ Returns:
107
+ Effective mode string: "simple", "advanced", or "hierarchical"
108
+ """
109
  if explicit_mode:
110
  if explicit_mode in ("magentic", "advanced"):
111
  return "advanced"
112
+ if explicit_mode == "hierarchical":
113
+ return "hierarchical"
114
  return "simple"
115
 
116
  # Auto-detect: advanced if paid API key available
src/orchestrators/hierarchical.py ADDED
@@ -0,0 +1,133 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Hierarchical Orchestrator using middleware and sub-teams.
2
+
3
+ This orchestrator implements a hierarchical team pattern where sub-teams
4
+ can be composed and coordinated through middleware. It provides more
5
+ granular control over the research workflow compared to the simple
6
+ orchestrator.
7
+
8
+ Design Patterns:
9
+ - Composite: Teams can contain sub-teams
10
+ - Chain of Responsibility: Middleware processes requests in sequence
11
+ - Template Method: SubIterationMiddleware defines the iteration skeleton
12
+ """
13
+
14
+ import asyncio
15
+ from collections.abc import AsyncGenerator
16
+
17
+ import structlog
18
+
19
+ from src.agents.judge_agent_llm import LLMSubIterationJudge
20
+ from src.agents.magentic_agents import create_search_agent
21
+ from src.middleware.sub_iteration import SubIterationMiddleware, SubIterationTeam
22
+ from src.services.embeddings import get_embedding_service
23
+ from src.state import init_magentic_state
24
+ from src.utils.models import AgentEvent
25
+
26
+ logger = structlog.get_logger()
27
+
28
+
29
+ class ResearchTeam(SubIterationTeam):
30
+ """Adapts ChatAgent to SubIterationTeam protocol.
31
+
32
+ This adapter allows the search agent to be used within the
33
+ sub-iteration middleware framework.
34
+ """
35
+
36
+ def __init__(self) -> None:
37
+ self.agent = create_search_agent()
38
+
39
+ async def execute(self, task: str) -> str:
40
+ """Execute a research task.
41
+
42
+ Args:
43
+ task: The research task description
44
+
45
+ Returns:
46
+ Text response from the agent
47
+ """
48
+ response = await self.agent.run(task)
49
+ if response.messages:
50
+ for msg in reversed(response.messages):
51
+ if msg.role == "assistant" and msg.text:
52
+ return str(msg.text)
53
+ return "No response from agent."
54
+
55
+
56
+ class HierarchicalOrchestrator:
57
+ """Orchestrator that uses hierarchical teams and sub-iterations.
58
+
59
+ This orchestrator provides:
60
+ - Sub-iteration middleware for fine-grained control
61
+ - LLM-based judge for sub-iteration decisions
62
+ - Event-driven architecture for UI updates
63
+ """
64
+
65
+ def __init__(self) -> None:
66
+ """Initialize the hierarchical orchestrator."""
67
+ self.team = ResearchTeam()
68
+ self.judge = LLMSubIterationJudge()
69
+ self.middleware = SubIterationMiddleware(self.team, self.judge, max_iterations=5)
70
+
71
+ async def run(self, query: str) -> AsyncGenerator[AgentEvent, None]:
72
+ """Run the hierarchical workflow.
73
+
74
+ Args:
75
+ query: User's research question
76
+
77
+ Yields:
78
+ AgentEvent objects for real-time UI updates
79
+ """
80
+ logger.info("Starting hierarchical orchestrator", query=query)
81
+
82
+ try:
83
+ service = get_embedding_service()
84
+ init_magentic_state(service)
85
+ except Exception as e:
86
+ logger.warning(
87
+ "Embedding service initialization failed, using default state",
88
+ error=str(e),
89
+ )
90
+ init_magentic_state()
91
+
92
+ yield AgentEvent(type="started", message=f"Starting research: {query}")
93
+
94
+ queue: asyncio.Queue[AgentEvent | None] = asyncio.Queue()
95
+
96
+ async def event_callback(event: AgentEvent) -> None:
97
+ await queue.put(event)
98
+
99
+ task_future = asyncio.create_task(self.middleware.run(query, event_callback))
100
+
101
+ while not task_future.done():
102
+ get_event = asyncio.create_task(queue.get())
103
+ done, _ = await asyncio.wait(
104
+ {task_future, get_event}, return_when=asyncio.FIRST_COMPLETED
105
+ )
106
+
107
+ if get_event in done:
108
+ event = get_event.result()
109
+ if event:
110
+ yield event
111
+ else:
112
+ get_event.cancel()
113
+
114
+ # Process remaining events
115
+ while not queue.empty():
116
+ ev = queue.get_nowait()
117
+ if ev:
118
+ yield ev
119
+
120
+ try:
121
+ result, assessment = await task_future
122
+
123
+ assessment_text = assessment.reasoning if assessment else "None"
124
+ yield AgentEvent(
125
+ type="complete",
126
+ message=(
127
+ f"Research complete.\n\nResult:\n{result}\n\nAssessment:\n{assessment_text}"
128
+ ),
129
+ data={"assessment": assessment.model_dump() if assessment else None},
130
+ )
131
+ except Exception as e:
132
+ logger.error("Orchestrator failed", error=str(e))
133
+ yield AgentEvent(type="error", message=f"Orchestrator failed: {e}")
src/{orchestrator.py β†’ orchestrators/simple.py} RENAMED
@@ -1,11 +1,20 @@
1
- """Orchestrator - the agent loop connecting Search and Judge."""
 
 
 
 
 
 
 
 
2
 
3
  import asyncio
4
  from collections.abc import AsyncGenerator
5
- from typing import Any, Protocol
6
 
7
  import structlog
8
 
 
9
  from src.utils.config import settings
10
  from src.utils.models import (
11
  AgentEvent,
@@ -18,23 +27,13 @@ from src.utils.models import (
18
  logger = structlog.get_logger()
19
 
20
 
21
- class SearchHandlerProtocol(Protocol):
22
- """Protocol for search handler."""
23
-
24
- async def execute(self, query: str, max_results_per_tool: int = 10) -> SearchResult: ...
25
-
26
-
27
- class JudgeHandlerProtocol(Protocol):
28
- """Protocol for judge handler."""
29
-
30
- async def assess(self, question: str, evidence: list[Evidence]) -> JudgeAssessment: ...
31
-
32
-
33
  class Orchestrator:
34
  """
35
- The agent orchestrator - runs the Search -> Judge -> Loop cycle.
36
 
37
  This is a generator-based design that yields events for real-time UI updates.
 
 
38
  """
39
 
40
  def __init__(
 
1
+ """Simple Orchestrator - the basic agent loop connecting Search and Judge.
2
+
3
+ This orchestrator uses a simple loop pattern with pydantic-ai for structured
4
+ LLM outputs. It works with free tier (HuggingFace Inference) or paid APIs
5
+ (OpenAI, Anthropic).
6
+
7
+ Design Pattern: Template Method - defines the skeleton of the search-judge loop
8
+ while allowing handlers to implement specific behaviors.
9
+ """
10
 
11
  import asyncio
12
  from collections.abc import AsyncGenerator
13
+ from typing import Any
14
 
15
  import structlog
16
 
17
+ from src.orchestrators.base import JudgeHandlerProtocol, SearchHandlerProtocol
18
  from src.utils.config import settings
19
  from src.utils.models import (
20
  AgentEvent,
 
27
  logger = structlog.get_logger()
28
 
29
 
 
 
 
 
 
 
 
 
 
 
 
 
30
  class Orchestrator:
31
  """
32
+ The simple agent orchestrator - runs the Search -> Judge -> Loop cycle.
33
 
34
  This is a generator-based design that yields events for real-time UI updates.
35
+ Uses pydantic-ai for structured LLM outputs without requiring the full
36
+ Microsoft Agent Framework.
37
  """
38
 
39
  def __init__(
tests/e2e/test_advanced_mode.py CHANGED
@@ -6,7 +6,7 @@ import pytest
6
  agent_framework = pytest.importorskip("agent_framework")
7
  from agent_framework import MagenticAgentMessageEvent, MagenticFinalResultEvent
8
 
9
- from src.orchestrator_magentic import MagenticOrchestrator
10
 
11
 
12
  class MockChatMessage:
@@ -24,7 +24,7 @@ async def test_advanced_mode_completes_mocked():
24
  """Verify Advanced mode runs without crashing (mocked workflow)."""
25
 
26
  # Initialize orchestrator (mocking requirements check)
27
- with patch("src.orchestrator_magentic.check_magentic_requirements"):
28
  orchestrator = MagenticOrchestrator(max_rounds=5)
29
 
30
  # Mock the workflow
@@ -51,7 +51,7 @@ async def test_advanced_mode_completes_mocked():
51
  # _init_embedding_service: Avoids loading embeddings
52
  with (
53
  patch.object(orchestrator, "_build_workflow", return_value=mock_workflow),
54
- patch("src.orchestrator_magentic.init_magentic_state"),
55
  patch.object(orchestrator, "_init_embedding_service", return_value=None),
56
  ):
57
  events = []
 
6
  agent_framework = pytest.importorskip("agent_framework")
7
  from agent_framework import MagenticAgentMessageEvent, MagenticFinalResultEvent
8
 
9
+ from src.orchestrators.advanced import AdvancedOrchestrator as MagenticOrchestrator
10
 
11
 
12
  class MockChatMessage:
 
24
  """Verify Advanced mode runs without crashing (mocked workflow)."""
25
 
26
  # Initialize orchestrator (mocking requirements check)
27
+ with patch("src.orchestrators.advanced.check_magentic_requirements"):
28
  orchestrator = MagenticOrchestrator(max_rounds=5)
29
 
30
  # Mock the workflow
 
51
  # _init_embedding_service: Avoids loading embeddings
52
  with (
53
  patch.object(orchestrator, "_build_workflow", return_value=mock_workflow),
54
+ patch("src.orchestrators.advanced.init_magentic_state"),
55
  patch.object(orchestrator, "_init_embedding_service", return_value=None),
56
  ):
57
  events = []
tests/e2e/test_simple_mode.py CHANGED
@@ -1,6 +1,6 @@
1
  import pytest
2
 
3
- from src.orchestrator import Orchestrator
4
  from src.utils.models import OrchestratorConfig
5
 
6
 
 
1
  import pytest
2
 
3
+ from src.orchestrators import Orchestrator
4
  from src.utils.models import OrchestratorConfig
5
 
6
 
tests/integration/test_dual_mode_e2e.py CHANGED
@@ -6,7 +6,7 @@ import pytest
6
 
7
  pytestmark = [pytest.mark.integration, pytest.mark.slow]
8
 
9
- from src.orchestrator_factory import create_orchestrator
10
  from src.utils.models import Citation, Evidence, OrchestratorConfig
11
 
12
 
@@ -65,17 +65,18 @@ async def test_advanced_mode_explicit_instantiation():
65
  MagenticOrchestrator can be instantiated when explicitly requested.
66
  The settings patch ensures any internal checks pass.
67
  """
68
- with patch("src.orchestrator_factory.settings") as mock_settings:
69
  # Settings patch ensures factory checks pass (even though mode is explicit)
70
  mock_settings.has_openai_key = True
71
 
72
  with patch("src.agents.magentic_agents.OpenAIChatClient"):
73
  # Mock agent creation to avoid real API calls during init
74
  with (
75
- patch("src.orchestrator_magentic.create_search_agent"),
76
- patch("src.orchestrator_magentic.create_judge_agent"),
77
- patch("src.orchestrator_magentic.create_hypothesis_agent"),
78
- patch("src.orchestrator_magentic.create_report_agent"),
 
79
  ):
80
  # Explicit mode="advanced" - tests the explicit path, not auto-detect
81
  orch = create_orchestrator(mode="advanced")
 
6
 
7
  pytestmark = [pytest.mark.integration, pytest.mark.slow]
8
 
9
+ from src.orchestrators import create_orchestrator
10
  from src.utils.models import Citation, Evidence, OrchestratorConfig
11
 
12
 
 
65
  MagenticOrchestrator can be instantiated when explicitly requested.
66
  The settings patch ensures any internal checks pass.
67
  """
68
+ with patch("src.orchestrators.factory.settings") as mock_settings:
69
  # Settings patch ensures factory checks pass (even though mode is explicit)
70
  mock_settings.has_openai_key = True
71
 
72
  with patch("src.agents.magentic_agents.OpenAIChatClient"):
73
  # Mock agent creation to avoid real API calls during init
74
  with (
75
+ patch("src.orchestrators.advanced.check_magentic_requirements"),
76
+ patch("src.orchestrators.advanced.create_search_agent"),
77
+ patch("src.orchestrators.advanced.create_judge_agent"),
78
+ patch("src.orchestrators.advanced.create_hypothesis_agent"),
79
+ patch("src.orchestrators.advanced.create_report_agent"),
80
  ):
81
  # Explicit mode="advanced" - tests the explicit path, not auto-detect
82
  orch = create_orchestrator(mode="advanced")
tests/unit/test_magentic_fix.py CHANGED
@@ -9,7 +9,7 @@ pytest.importorskip("agent_framework")
9
 
10
  from agent_framework import MagenticFinalResultEvent # noqa: E402
11
 
12
- from src.orchestrator_magentic import MagenticOrchestrator # noqa: E402
13
 
14
 
15
  class MockChatMessage:
@@ -39,7 +39,7 @@ class MockChatMessage:
39
  @pytest.fixture
40
  def mock_magentic_requirements():
41
  """Mock the API key check so tests run in CI without OPENAI_API_KEY."""
42
- with patch("src.orchestrator_magentic.check_magentic_requirements"):
43
  yield
44
 
45
 
@@ -74,12 +74,12 @@ class TestMagenticFixes:
74
  # Also verify it's used in _build_workflow
75
  # Mock all the agent creation and OpenAI client calls
76
  with (
77
- patch("src.orchestrator_magentic.create_search_agent") as mock_search,
78
- patch("src.orchestrator_magentic.create_judge_agent") as mock_judge,
79
- patch("src.orchestrator_magentic.create_hypothesis_agent") as mock_hypo,
80
- patch("src.orchestrator_magentic.create_report_agent") as mock_report,
81
- patch("src.orchestrator_magentic.OpenAIChatClient") as mock_client,
82
- patch("src.orchestrator_magentic.MagenticBuilder") as mock_builder,
83
  ):
84
  # Setup mocks
85
  mock_search.return_value = MagicMock()
 
9
 
10
  from agent_framework import MagenticFinalResultEvent # noqa: E402
11
 
12
+ from src.orchestrators.advanced import AdvancedOrchestrator as MagenticOrchestrator # noqa: E402
13
 
14
 
15
  class MockChatMessage:
 
39
  @pytest.fixture
40
  def mock_magentic_requirements():
41
  """Mock the API key check so tests run in CI without OPENAI_API_KEY."""
42
+ with patch("src.orchestrators.advanced.check_magentic_requirements"):
43
  yield
44
 
45
 
 
74
  # Also verify it's used in _build_workflow
75
  # Mock all the agent creation and OpenAI client calls
76
  with (
77
+ patch("src.orchestrators.advanced.create_search_agent") as mock_search,
78
+ patch("src.orchestrators.advanced.create_judge_agent") as mock_judge,
79
+ patch("src.orchestrators.advanced.create_hypothesis_agent") as mock_hypo,
80
+ patch("src.orchestrators.advanced.create_report_agent") as mock_report,
81
+ patch("src.orchestrators.advanced.OpenAIChatClient") as mock_client,
82
+ patch("src.orchestrators.advanced.MagenticBuilder") as mock_builder,
83
  ):
84
  # Setup mocks
85
  mock_search.return_value = MagicMock()
tests/unit/test_magentic_termination.py CHANGED
@@ -5,7 +5,7 @@ from unittest.mock import MagicMock, patch
5
  import pytest
6
  from agent_framework import MagenticAgentMessageEvent
7
 
8
- from src.orchestrator_magentic import MagenticOrchestrator
9
  from src.utils.models import AgentEvent
10
 
11
  # Skip tests if agent_framework is not installed
@@ -25,7 +25,7 @@ class MockChatMessage:
25
  @pytest.fixture
26
  def mock_magentic_requirements():
27
  """Mock requirements check."""
28
- with patch("src.orchestrator_magentic.check_magentic_requirements"):
29
  yield
30
 
31
 
 
5
  import pytest
6
  from agent_framework import MagenticAgentMessageEvent
7
 
8
+ from src.orchestrators.advanced import AdvancedOrchestrator as MagenticOrchestrator
9
  from src.utils.models import AgentEvent
10
 
11
  # Skip tests if agent_framework is not installed
 
25
  @pytest.fixture
26
  def mock_magentic_requirements():
27
  """Mock requirements check."""
28
+ with patch("src.orchestrators.advanced.check_magentic_requirements"):
29
  yield
30
 
31
 
tests/unit/test_orchestrator.py CHANGED
@@ -4,7 +4,7 @@ from unittest.mock import AsyncMock
4
 
5
  import pytest
6
 
7
- from src.orchestrator import Orchestrator
8
  from src.utils.models import (
9
  AgentEvent,
10
  AssessmentDetails,
 
4
 
5
  import pytest
6
 
7
+ from src.orchestrators import Orchestrator
8
  from src.utils.models import (
9
  AgentEvent,
10
  AssessmentDetails,
tests/unit/test_orchestrator_factory.py CHANGED
@@ -6,19 +6,18 @@ import pytest
6
 
7
  pytestmark = pytest.mark.unit
8
 
9
- from src.orchestrator import Orchestrator
10
- from src.orchestrator_factory import create_orchestrator
11
 
12
 
13
  @pytest.fixture
14
  def mock_settings():
15
- with patch("src.orchestrator_factory.settings", autospec=True) as mock_settings:
16
  yield mock_settings
17
 
18
 
19
  @pytest.fixture
20
  def mock_magentic_cls():
21
- with patch("src.orchestrator_factory._get_magentic_orchestrator_class") as mock:
22
  # The mock returns a class (callable), which returns an instance
23
  mock_class = MagicMock()
24
  mock.return_value = mock_class
 
6
 
7
  pytestmark = pytest.mark.unit
8
 
9
+ from src.orchestrators import Orchestrator, create_orchestrator
 
10
 
11
 
12
  @pytest.fixture
13
  def mock_settings():
14
+ with patch("src.orchestrators.factory.settings", autospec=True) as mock_settings:
15
  yield mock_settings
16
 
17
 
18
  @pytest.fixture
19
  def mock_magentic_cls():
20
+ with patch("src.orchestrators.factory._get_advanced_orchestrator_class") as mock:
21
  # The mock returns a class (callable), which returns an instance
22
  mock_class = MagicMock()
23
  mock.return_value = mock_class
tests/unit/test_streaming_fix.py CHANGED
@@ -59,9 +59,9 @@ async def test_streaming_events_are_buffered_not_spammed():
59
  assert len(results) > 0, "Should have yielded results"
60
 
61
  # Check that we see the accumulated message
62
- assert any(
63
- "πŸ“‘ **STREAMING**: This is a test" in r for r in results
64
- ), "Buffer didn't accumulate correctly"
65
 
66
  # The critical check for the "Spam" bug:
67
  # In the spam bug, the output grew like:
 
59
  assert len(results) > 0, "Should have yielded results"
60
 
61
  # Check that we see the accumulated message
62
+ assert any("πŸ“‘ **STREAMING**: This is a test" in r for r in results), (
63
+ "Buffer didn't accumulate correctly"
64
+ )
65
 
66
  # The critical check for the "Spam" bug:
67
  # In the spam bug, the output grew like:
tests/unit/test_ui_elements.py CHANGED
@@ -6,17 +6,17 @@ from src.app import create_demo
6
  def test_examples_include_advanced_mode():
7
  """Verify that one example entry uses 'advanced' mode."""
8
  demo, _ = create_demo()
9
- assert any(
10
- "advanced" == example[1] for example in demo.examples
11
- ), "Expected at least one example to be 'advanced' mode"
12
 
13
 
14
  def test_accordion_label_updated():
15
  """Verify the accordion label reflects the new, concise text."""
16
  _, accordion = create_demo()
17
- assert (
18
- accordion.label == "βš™οΈ Mode & API Key (Free tier works!)"
19
- ), "Accordion label not updated to 'βš™οΈ Mode & API Key (Free tier works!)'"
20
 
21
 
22
  def test_orchestrator_mode_info_text_updated():
@@ -25,9 +25,9 @@ def test_orchestrator_mode_info_text_updated():
25
  # Assuming additional_inputs is a list and the Radio is the first element
26
  orchestrator_radio = demo.additional_inputs[0]
27
  expected_info = "⚑ Simple: Free/OpenAI/Anthropic | πŸ”¬ Advanced: OpenAI only"
28
- assert isinstance(
29
- orchestrator_radio, gr.Radio
30
- ), "Expected first additional input to be gr.Radio"
31
- assert (
32
- orchestrator_radio.info == expected_info
33
- ), "Orchestrator Mode info text not updated correctly"
 
6
  def test_examples_include_advanced_mode():
7
  """Verify that one example entry uses 'advanced' mode."""
8
  demo, _ = create_demo()
9
+ assert any("advanced" == example[1] for example in demo.examples), (
10
+ "Expected at least one example to be 'advanced' mode"
11
+ )
12
 
13
 
14
  def test_accordion_label_updated():
15
  """Verify the accordion label reflects the new, concise text."""
16
  _, accordion = create_demo()
17
+ assert accordion.label == "βš™οΈ Mode & API Key (Free tier works!)", (
18
+ "Accordion label not updated to 'βš™οΈ Mode & API Key (Free tier works!)'"
19
+ )
20
 
21
 
22
  def test_orchestrator_mode_info_text_updated():
 
25
  # Assuming additional_inputs is a list and the Radio is the first element
26
  orchestrator_radio = demo.additional_inputs[0]
27
  expected_info = "⚑ Simple: Free/OpenAI/Anthropic | πŸ”¬ Advanced: OpenAI only"
28
+ assert isinstance(orchestrator_radio, gr.Radio), (
29
+ "Expected first additional input to be gr.Radio"
30
+ )
31
+ assert orchestrator_radio.info == expected_info, (
32
+ "Orchestrator Mode info text not updated correctly"
33
+ )