diff --git a/examples/ag-ui-research/backend/agent.py b/examples/ag-ui-research/backend/agent.py index 0182f1e6..3dce81ac 100644 --- a/examples/ag-ui-research/backend/agent.py +++ b/examples/ag-ui-research/backend/agent.py @@ -33,6 +33,7 @@ class AgentDeps: client: HaikuRAG agui_emitter: "AGUIEmitter[ResearchState, ResearchReport] | None" = None + search_filter: str | None = None model = get_model(Config.research.model, Config) @@ -74,6 +75,7 @@ async def run_research(ctx: RunContext[AgentDeps], question: str) -> str: graph = build_research_graph(Config) context = ResearchContext(original_question=question) state = ResearchState.from_config(context=context, config=Config) + state.search_filter = ctx.deps.search_filter graph_deps = ResearchDeps( client=ctx.deps.client, diff --git a/examples/ag-ui-research/backend/main.py b/examples/ag-ui-research/backend/main.py index 48593415..3909c496 100644 --- a/examples/ag-ui-research/backend/main.py +++ b/examples/ag-ui-research/backend/main.py @@ -92,8 +92,19 @@ async def stream_research_agent(request: Request) -> StreamingResponse: effective_db_path = Path(effective_db_path) client = get_client(effective_db_path) + # Build search filter from document IDs if provided + document_ids = input_data.state.get("documentFilter") + search_filter = None + if document_ids: + ids_str = ", ".join(f"'{id}'" for id in document_ids) + search_filter = f"id IN ({ids_str})" + # Create agent dependencies with shared emitter - agent_deps = AgentDeps(client=client, agui_emitter=emitter) + agent_deps = AgentDeps( + client=client, + agui_emitter=emitter, + search_filter=search_filter, + ) # Start run with empty initial state emitter.start_run(