Remove max_concurrency, not supported (yet) when we .map() in the beta Graph API

This commit is contained in:
Yiorgis Gozadinos 2025-11-04 12:30:58 +02:00
parent 38861da549
commit 6858760bca
No known key found for this signature in database
9 changed files with 21 additions and 33 deletions

View file

@ -59,7 +59,7 @@ stateDiagram-v2
Key nodes:
- **Plan**: Decomposes the question into focused sub-questions
- **Search (batched)**: Answers sub-questions in parallel batches (respects max_concurrency)
- **Search (parallel)**: Answers all sub-questions in parallel
- **Decision**: Evaluates if we have sufficient information or need another iteration
- **Synthesize**: Generates the final comprehensive answer
@ -69,7 +69,7 @@ Key differences from Research:
- **Direct answers**: Returns just the answer (not a full research report)
- **Question-focused**: Optimized for answering specific questions, not open-ended research
- **Supports citations**: Can include inline source citations like `[document.md]`
- **Configurable iterations**: Control max_iterations (default: 2) and max_concurrency (default: 3)
- **Configurable iterations**: Control max_iterations (default: 2)
CLI usage:
@ -99,8 +99,7 @@ async with HaikuRAG(path_to_db) as client:
state = DeepQAState(
context=context,
max_sub_questions=3,
max_iterations=2,
max_concurrency=3
max_iterations=2
)
deps = DeepQADeps(client=client)
@ -176,7 +175,6 @@ async with HaikuRAG(path_to_db) as client:
context=ResearchContext(original_question=question),
max_iterations=2,
confidence_threshold=0.8,
max_concurrency=2,
)
deps = ResearchDeps(client=client)
@ -211,7 +209,6 @@ async with HaikuRAG(path_to_db) as client:
context=ResearchContext(original_question=question),
max_iterations=2,
confidence_threshold=0.8,
max_concurrency=2,
)
deps = ResearchDeps(client=client)

View file

@ -242,7 +242,6 @@ class HaikuRAGApp:
question: str,
max_iterations: int = 3,
confidence_threshold: float = 0.8,
max_concurrency: int = 1,
verbose: bool = False,
):
"""Run research via the pydantic-graph pipeline (default)."""
@ -262,7 +261,6 @@ class HaikuRAGApp:
context=context,
max_iterations=max_iterations,
confidence_threshold=confidence_threshold,
max_concurrency=max_concurrency,
)
deps = ResearchDeps(
client=client, console=self.console if verbose else None

View file

@ -301,11 +301,6 @@ def research(
"--confidence-threshold",
help="Minimum confidence (0-1) to stop",
),
max_concurrency: int = typer.Option(
1,
"--max-concurrency",
help="Max concurrent searches per iteration (planned)",
),
db: Path = typer.Option(
Config.storage.data_dir / "haiku.rag.lancedb",
"--db",
@ -325,7 +320,6 @@ def research(
question=question,
max_iterations=max_iterations,
confidence_threshold=confidence_threshold,
max_concurrency=max_concurrency,
verbose=verbose,
)
)

View file

@ -219,7 +219,6 @@ def create_mcp_server(db_path: Path) -> FastMCP:
question: str,
max_iterations: int = 3,
confidence_threshold: float = 0.8,
max_concurrency: int = 1,
) -> ResearchReport | None:
"""Run multi-agent research to investigate a complex question.
@ -230,7 +229,6 @@ def create_mcp_server(db_path: Path) -> FastMCP:
question: The research question to investigate.
max_iterations: Maximum search/analyze iterations (default: 3).
confidence_threshold: Minimum confidence score (0-1) to stop early (default: 0.8).
max_concurrency: Maximum concurrent searches per iteration (default: 1).
Returns:
A research report with findings, or None if an error occurred.
@ -249,7 +247,6 @@ def create_mcp_server(db_path: Path) -> FastMCP:
context=ResearchContext(original_question=question),
max_iterations=max_iterations,
confidence_threshold=confidence_threshold,
max_concurrency=max_concurrency,
)
deps = ResearchDeps(client=rag)

View file

@ -149,13 +149,16 @@ def build_deep_qa_graph(
async def get_batch(
ctx: StepContext[DeepQAState, DeepQADeps, None | bool],
) -> list[str] | None:
"""Get next batch of questions from state."""
"""Get all remaining questions for this iteration."""
state = ctx.state
take = max(1, state.max_concurrency)
batch: list[str] = []
while state.context.sub_questions and len(batch) < take:
batch.append(state.context.sub_questions.pop(0))
return batch if batch else None
if not state.context.sub_questions:
return None
# Take ALL remaining questions - max_concurrency controls parallel execution within .map()
batch = list(state.context.sub_questions)
state.context.sub_questions.clear()
return batch
@g.step
async def decide(

View file

@ -21,5 +21,4 @@ class DeepQAState:
context: DeepQAContext
max_sub_questions: int = 3
max_iterations: int = 2
max_concurrency: int = 1
iterations: int = 0

View file

@ -159,13 +159,16 @@ def build_research_graph(
async def get_batch(
ctx: StepContext[ResearchState, ResearchDeps, None | bool],
) -> list[str] | None:
"""Get next batch of questions from state."""
"""Get all remaining questions for this iteration."""
state = ctx.state
take = max(1, state.max_concurrency)
batch: list[str] = []
while state.context.sub_questions and len(batch) < take:
batch.append(state.context.sub_questions.pop(0))
return batch if batch else None
if not state.context.sub_questions:
return None
# Take ALL remaining questions and process them in parallel
batch = list(state.context.sub_questions)
state.context.sub_questions.clear()
return batch
@g.step
async def analyze_insights(

View file

@ -26,7 +26,6 @@ class ResearchState:
context: ResearchContext
iterations: int = 0
max_iterations: int = 3
max_concurrency: int = 1
confidence_threshold: float = 0.8
last_eval: EvaluationResult | None = None
last_analysis: InsightAnalysis | None = None

View file

@ -15,7 +15,6 @@ class ResearchStateSnapshot:
sub_questions: list[str]
iterations: int
max_iterations: int
max_concurrency: int
confidence_threshold: float
pending_sub_questions: int
answered_questions: int
@ -38,7 +37,6 @@ class ResearchStateSnapshot:
sub_questions=list(context.sub_questions),
iterations=state.iterations,
max_iterations=state.max_iterations,
max_concurrency=state.max_concurrency,
confidence_threshold=state.confidence_threshold,
pending_sub_questions=len(context.sub_questions),
answered_questions=len(context.qa_responses),