Skip to content

Commit 9428294

Browse files
trial-danswertrial-danswer
committed
Seeing outputs from multiple sources.
1 parent 264878a commit 9428294

File tree

11 files changed

+940
-10
lines changed

11 files changed

+940
-10
lines changed

‎backend/onyx/agents/agent_search/kb_search/nodes/b2p_process_individual_deep_search.py‎

Lines changed: 13 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -169,13 +169,20 @@ def process_individual_deep_search(
169169

170170
logger.debug("DivCon Step A2 - Object Source Research - completed for an object")
171171

172+
if state.research_object_results:
173+
research_object_results = state.research_object_results
174+
else:
175+
research_object_results = []
176+
177+
research_object_results.append(
178+
{
179+
"object": object.replace("::", ":: ").capitalize(),
180+
"results": object_research_results,
181+
}
182+
)
183+
172184
return ResearchObjectUpdate(
173-
research_object_results=[
174-
{
175-
"object": object.replace("::", ":: ").capitalize(),
176-
"results": object_research_results,
177-
}
178-
],
185+
research_object_results=research_object_results,
179186
log_messages=[
180187
get_langgraph_node_log_string(
181188
graph_component="main",

‎backend/onyx/agents/agent_search/naomi/README.md‎

Lines changed: 113 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,113 @@
1+
# Naomi Orchestration Graph
2+
3+
The Naomi orchestration graph is designed to execute both basic and kb_search graphs with a decision node that determines which graph to run based on the current execution stage.
4+
5+
## Overview
6+
7+
The graph follows this flow:
8+
1. **Decision Node**: Determines which stage to execute (BASIC, KB_SEARCH, or COMPLETE)
9+
2. **Basic Graph Execution**: Runs the basic search graph and stores results
10+
3. **KB Search Graph Execution**: Runs the knowledge graph search and stores results
11+
4. **Finalization**: Combines results from both graphs into a final answer
12+
5. **Loop Back**: Returns to decision node to determine next steps
13+
14+
## Architecture
15+
16+
### State Management
17+
- `NaomiState`: Main state that tracks current stage and stores results from both graphs
18+
- `ExecutionStage`: Enum defining the three stages (BASIC, KB_SEARCH, COMPLETE)
19+
- Results from each graph are stored separately and combined at the end
20+
21+
### Nodes
22+
- `decision_node`: Determines which graph to execute next
23+
- `execute_basic_graph`: Runs the basic search graph
24+
- `execute_kb_search_graph`: Runs the knowledge graph search
25+
- `finalize_results`: Combines results and creates final answer
26+
27+
### Conditional Edges
28+
- `route_after_decision`: Routes to appropriate execution node based on current stage
29+
- `should_continue`: Determines if graph should continue or end
30+
31+
## Usage
32+
33+
### Basic Usage
34+
```python
35+
from onyx.agents.agent_search.naomi import naomi_graph_builder, NaomiInput
36+
37+
# Build the graph
38+
graph = naomi_graph_builder()
39+
compiled_graph = graph.compile()
40+
41+
# Create input
42+
input_data = NaomiInput()
43+
44+
# Execute with proper config
45+
result = compiled_graph.invoke(input_data, config={"metadata": {"config": config}})
46+
47+
# Access results
48+
final_answer = result.get("final_answer")
49+
basic_results = result.get("basic_results")
50+
kb_search_results = result.get("kb_search_results")
51+
```
52+
53+
### Testing
54+
Run the test script to see the graph in action:
55+
```bash
56+
cd backend/onyx/agents/agent_search/naomi
57+
python test_naomi.py
58+
```
59+
60+
## Execution Flow
61+
62+
1. **Start**: Graph begins with decision node
63+
2. **Stage Check**: Decision node checks current stage
64+
3. **Graph Execution**:
65+
- If BASIC stage: Execute basic graph
66+
- If KB_SEARCH stage: Execute kb_search graph
67+
- If COMPLETE stage: Finalize results
68+
4. **Result Storage**: Results are stored in state
69+
5. **Loop**: Return to decision node
70+
6. **Completion**: When COMPLETE stage has results, graph ends
71+
72+
## Customization
73+
74+
### Modifying Decision Logic
75+
Edit the `decision_node` function in `nodes.py` to implement custom decision logic based on:
76+
- Query complexity
77+
- Previous results quality
78+
- User preferences
79+
- Performance requirements
80+
81+
### Adding New Stages
82+
1. Add new stage to `ExecutionStage` enum in `states.py`
83+
2. Update decision logic in `nodes.py`
84+
3. Add routing logic in `conditional_edges.py`
85+
4. Update graph builder in `graph_builder.py`
86+
87+
### Config Integration
88+
The current implementation is simplified. In production, you'll need to:
89+
- Pass proper config with LLMs and database session
90+
- Handle authentication and permissions
91+
- Implement proper error handling and retry logic
92+
- Add monitoring and logging
93+
94+
## Dependencies
95+
96+
- `langgraph`: For graph construction and execution
97+
- `pydantic`: For state validation
98+
- `onyx.agents.agent_search.basic`: Basic search graph
99+
- `onyx.agents.agent_search.kb_search`: Knowledge graph search
100+
- `onyx.utils.logger`: For logging
101+
102+
## File Structure
103+
104+
```
105+
naomi/
106+
├── __init__.py # Package initialization
107+
├── states.py # State definitions
108+
├── nodes.py # Node functions
109+
├── conditional_edges.py # Conditional edge logic
110+
├── graph_builder.py # Main graph builder
111+
├── test_naomi.py # Test script
112+
└── README.md # This file
113+
```

‎backend/onyx/agents/agent_search/naomi/basic_graph_builder.py‎

Lines changed: 101 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,101 @@
1+
from langgraph.graph import END
2+
from langgraph.graph import START
3+
from langgraph.graph import StateGraph
4+
5+
from onyx.agents.agent_search.basic.states import BasicInput
6+
from onyx.agents.agent_search.basic.states import BasicState
7+
from onyx.agents.agent_search.naomi.states import BasicShortOutput
8+
from onyx.agents.agent_search.orchestration.nodes.call_tool import call_tool
9+
from onyx.agents.agent_search.orchestration.nodes.choose_tool import choose_tool
10+
from onyx.agents.agent_search.orchestration.nodes.prepare_tool_input import (
11+
prepare_tool_input,
12+
)
13+
from onyx.utils.logger import setup_logger
14+
15+
logger = setup_logger()
16+
17+
18+
def basic_graph_builder() -> StateGraph:
19+
graph = StateGraph(
20+
state_schema=BasicState, input=BasicInput, output=BasicShortOutput
21+
)
22+
23+
### Add nodes ###
24+
25+
graph.add_node(
26+
node="prepare_tool_input",
27+
action=prepare_tool_input,
28+
)
29+
30+
graph.add_node(
31+
node="choose_tool",
32+
action=choose_tool,
33+
)
34+
35+
graph.add_node(
36+
node="call_tool",
37+
action=call_tool,
38+
)
39+
40+
graph.add_node(
41+
node="return_output",
42+
action=return_output,
43+
)
44+
45+
### Add edges ###
46+
47+
graph.add_edge(start_key=START, end_key="prepare_tool_input")
48+
49+
graph.add_edge(start_key="prepare_tool_input", end_key="choose_tool")
50+
51+
graph.add_conditional_edges(
52+
"choose_tool", should_continue, ["call_tool", "return_output"]
53+
)
54+
55+
graph.add_edge(
56+
start_key="call_tool",
57+
end_key="return_output",
58+
)
59+
60+
graph.add_edge(
61+
start_key="return_output",
62+
end_key=END,
63+
)
64+
65+
return graph
66+
67+
68+
def should_continue(state: BasicState) -> str:
69+
return (
70+
# If there are no tool calls, basic graph already streamed the answer
71+
END
72+
if state.tool_choice is None
73+
else "call_tool"
74+
)
75+
76+
77+
def return_output(state: BasicState) -> BasicShortOutput:
78+
return BasicShortOutput(
79+
tool_call_output=state.tool_call_output,
80+
tool_choice=state.tool_choice,
81+
)
82+
83+
84+
if __name__ == "__main__":
85+
from onyx.db.engine.sql_engine import get_session_with_current_tenant
86+
from onyx.context.search.models import SearchRequest
87+
from onyx.llm.factory import get_default_llms
88+
from onyx.agents.agent_search.shared_graph_utils.utils import get_test_config
89+
90+
graph = basic_graph_builder()
91+
compiled_graph = graph.compile()
92+
input = BasicInput(unused=True)
93+
primary_llm, fast_llm = get_default_llms()
94+
with get_session_with_current_tenant() as db_session:
95+
config, _ = get_test_config(
96+
db_session=db_session,
97+
primary_llm=primary_llm,
98+
fast_llm=fast_llm,
99+
search_request=SearchRequest(query="How does onyx use FastAPI?"),
100+
)
101+
compiled_graph.invoke(input, config={"metadata": {"config": config}})

‎backend/onyx/agents/agent_search/naomi/conditional_edges.py‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,16 @@
1+
from onyx.agents.agent_search.naomi.states import ExecutionStage
2+
from onyx.agents.agent_search.naomi.states import NaomiState
3+
4+
5+
def route_after_decision(state: NaomiState) -> str:
6+
"""
7+
Route to the appropriate node based on the current stage after decision node.
8+
"""
9+
if state.current_stage == ExecutionStage.BASIC:
10+
return "execute_basic_graph"
11+
if state.current_stage == ExecutionStage.KB_SEARCH:
12+
return "execute_kb_search_graph"
13+
elif state.current_stage == ExecutionStage.COMPLETE:
14+
return "finalize_results"
15+
else:
16+
raise ValueError(f"Invalid stage: {state.current_stage}")

‎backend/onyx/agents/agent_search/naomi/graph_builder.py‎

Lines changed: 113 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,113 @@
1+
from langgraph.graph import END
2+
from langgraph.graph import START
3+
from langgraph.graph import StateGraph
4+
5+
from onyx.agents.agent_search.naomi.conditional_edges import (
6+
route_after_decision,
7+
)
8+
from onyx.agents.agent_search.naomi.nodes.nodes import decision_node
9+
from onyx.agents.agent_search.naomi.nodes.nodes import execute_basic_graph
10+
from onyx.agents.agent_search.naomi.nodes.nodes import execute_kb_search_graph
11+
from onyx.agents.agent_search.naomi.nodes.nodes import finalize_results
12+
from onyx.agents.agent_search.naomi.states import NaomiInput
13+
from onyx.agents.agent_search.naomi.states import NaomiState
14+
from onyx.utils.logger import setup_logger
15+
16+
logger = setup_logger()
17+
18+
19+
def naomi_graph_builder() -> StateGraph:
20+
"""
21+
LangGraph graph builder for the naomi orchestration process.
22+
This graph orchestrates both basic and kb_search graphs with a decision node.
23+
"""
24+
25+
graph = StateGraph(
26+
state_schema=NaomiState,
27+
input=NaomiInput,
28+
)
29+
30+
### Add nodes ###
31+
32+
graph.add_node(
33+
"decision_node",
34+
decision_node,
35+
)
36+
37+
graph.add_node(
38+
"execute_basic_graph",
39+
execute_basic_graph,
40+
)
41+
42+
graph.add_node(
43+
"execute_kb_search_graph",
44+
execute_kb_search_graph,
45+
)
46+
47+
graph.add_node(
48+
"finalize_results",
49+
finalize_results,
50+
)
51+
52+
### Add edges ###
53+
54+
# Start with decision node
55+
graph.add_edge(start_key=START, end_key="decision_node")
56+
57+
# Decision node routes to appropriate execution node
58+
graph.add_conditional_edges(
59+
"decision_node",
60+
route_after_decision,
61+
["execute_basic_graph", "execute_kb_search_graph", "finalize_results"],
62+
)
63+
64+
# After executing basic graph, go back to decision node
65+
graph.add_edge(
66+
start_key="execute_basic_graph",
67+
end_key="decision_node",
68+
)
69+
70+
# After executing kb_search graph, go back to decision node
71+
graph.add_edge(
72+
start_key="execute_kb_search_graph",
73+
end_key="decision_node",
74+
)
75+
76+
# After finalizing results, check if we should continue or end
77+
graph.add_edge(start_key="finalize_results", end_key=END)
78+
79+
return graph
80+
81+
82+
if __name__ == "__main__":
83+
# Test the graph
84+
from onyx.db.engine.sql_engine import get_session_with_current_tenant
85+
from onyx.context.search.models import SearchRequest
86+
from onyx.llm.factory import get_default_llms
87+
from onyx.agents.agent_search.shared_graph_utils.utils import get_test_config
88+
89+
graph = naomi_graph_builder()
90+
compiled_graph = graph.compile()
91+
92+
# Create test input
93+
input_data = NaomiInput()
94+
95+
# Get LLMs and config
96+
primary_llm, fast_llm = get_default_llms()
97+
98+
with get_session_with_current_tenant() as db_session:
99+
config, _ = get_test_config(
100+
db_session=db_session,
101+
primary_llm=primary_llm,
102+
fast_llm=fast_llm,
103+
search_request=SearchRequest(query="How does onyx use FastAPI?"),
104+
)
105+
106+
# Execute the graph
107+
result = compiled_graph.invoke(
108+
input_data, config={"metadata": {"config": config}}
109+
)
110+
111+
print("Final Answer:", result.get("final_answer", ""))
112+
print("Basic Results:", result.get("basic_results", {}))
113+
print("KB Search Results:", result.get("kb_search_results", {}))

0 commit comments

Comments
 (0)