Skip to content

Comments

feat(agent): Add LangGraph intent router PoC for multi-source retrieval (#42)#43

Open
saurabhhhcodes wants to merge 2 commits intokubeflow:mainfrom
saurabhhhcodes:feat/langgraph-intent-router-poc
Open

feat(agent): Add LangGraph intent router PoC for multi-source retrieval (#42)#43
saurabhhhcodes wants to merge 2 commits intokubeflow:mainfrom
saurabhhhcodes:feat/langgraph-intent-router-poc

Conversation

@saurabhhhcodes
Copy link

Summary

Resolves #42.

This PR introduces server/agent_router.py, a focused LangGraph routing layer that classifies user intent and dispatches queries to the correct Milvus partition (docs / issues / platform) — the logical next step once PR #12 (Milvus partitions) lands.

Scoped strictly to agent logic — no Terraform, no KServe manifests, no ingestion pipelines. This keeps the PR small and fast to review.


Architecture

user_query → intent_router → docs        → END
                           → issues      → END
                           → platform    → END

On empty retrieval: → increment_retries → intent_router (max 2 loops)

Key Design Decisions

Concern Approach
Deterministic routing with_structured_output(RouteQuery) + Pydantic strict validation — invalid partition names raise ValueError
Self-correction loop Conditional edge re-routes back to intent_router on empty retrieval (capped at MAX_RETRIES)
Graceful degradation Partition fallback to full-collection search if named Milvus partition doesn't yet exist (safe on current main)
State immutability Every node returns {**state, ...overrides} — no in-place mutation

Tests

36 real unit tests — zero mocks, all passing locally:

pytest tests/test_agent_router.py -v
# 36 passed, 1 warning in 8.09s

Coverage: RouteQuery validation, AgentState defaults, _format_hits, route_to_partition, handle_empty_retrieval, increment_retries, and full graph topology via build_graph().

Note: A real bug was caught and fixed during testing: route_to_partition was using .get("intent", "docs") which does not fire on an empty-string intent. Fixed to state.get("intent") or "docs".

Next Steps (post-review)

Once the routing architecture is validated here, the follow-up PR will wire the retrieval nodes into the actual Milvus partitions from PR #12.

…al (kubeflow#42)

Introduces server/agent_router.py, a LangGraph-powered routing layer
that classifies user intent and dispatches queries to the correct
Milvus partition (docs / issues / platform) rather than doing a blind
full-collection vector search on every request.

Key design decisions:
- AgentState TypedDict tracks question, intent, context, citations,
  error, and retries across every graph node.
- Intent classification uses ChatOpenAI.with_structured_output() with
  a strict Pydantic schema (RouteQuery) — invalid datasource values
  raise a ValueError so hallucinated routes are caught at the boundary.
- Self-correction loop: if a partition returns an empty result set the
  graph increments a retry counter and routes back to intent_router
  (capped at MAX_RETRIES=2) so the LLM can reconsider which partition
  to try next instead of replying with empty context.
- Graceful partition fallback: if the named Milvus partition does not
  yet exist (pre-PR kubeflow#12) the node falls back to a full-collection
  search, keeping the code runnable on the current main branch.

Resolves kubeflow#42
Co-authored-by: saurabhhhcodes <saurabhhhcodes@users.noreply.github.com>

Signed-off-by: Saurabh Kumar Bajpai <saurabhbajpai03@outlook.com>
…assing)

Tests cover:
- RouteQuery Pydantic schema validation (valid/invalid/case-sensitive)
- AgentState structure and defaults
- _format_hits: dedup, empty URL skip, score formatting
- route_to_partition: all intents + empty/missing intent fallback
- handle_empty_retrieval: retry logic and MAX_RETRIES boundary
- increment_retries: counter increment, state immutability
- build_graph: compile, node presence, topology, retry→router edge

Also fixes route_to_partition to handle empty-string intent
(use 'or' instead of .get default, which does not fire on falsy values)

Signed-off-by: Saurabh Kumar Bajpai <saurabhbajpai03@outlook.com>
@google-oss-prow
Copy link

[APPROVALNOTIFIER] This PR is NOT APPROVED

This pull-request has been approved by:
Once this PR has been reviewed and has the lgtm label, please assign franciscojavierarceo for approval. For more information see the Kubernetes Code Review Process.

The full list of commands accepted by this bot can be found here.

Details Needs approval from an approver in each of these files:

Approvers can indicate their approval by writing /approve in a comment
Approvers can cancel approval by writing /approve cancel in a comment

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Implement Agentic Query Router (LangGraph) for multi-source retrieval

1 participant