highperfocused/danswer
1
0
Fork 0
mirror of https://github.com/danswer-ai/danswer.git synced 2025-05-18 07:40:05 +02:00
Code Issues Packages Projects Releases Wiki Activity

renamings and consolidation of formatting nodes in orig question retrieval

Browse Source
This commit is contained in:
joachim-danswer 2025-01-30 13:20:06 -08:00 committed by Evan Lohn
parent 3b13380051
commit ceaaa05af0
9 changed files with 52 additions and 96 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
backend/onyx/agents/agent_search/deep_search_a/initial/generate_initial_answer/graph_builder.py
View File

@ -20,7 +20,7 @@ from onyx.agents.agent_search.deep_search_a.initial.generate_initial_answer.stat
from onyx.agents.agent_search.deep_search_a.initial.generate_initial_answer.states import (
SearchSQState,
)
from onyx.agents.agent_search.deep_search_a.initial.retrieve_orig_question_documents.graph_builder import (
from onyx.agents.agent_search.deep_search_a.initial.retrieve_orig_question_docs.graph_builder import (
base_raw_search_graph_builder,
)
from onyx.utils.logger import setup_logger

4
backend/onyx/agents/agent_search/deep_search_a/initial/generate_initial_answer/nodes/ingest_retrieved_documents.py
View File

@ -1,6 +1,6 @@
from datetime import datetime
from onyx.agents.agent_search.deep_search_a.initial.retrieve_orig_question_documents.states import (
from onyx.agents.agent_search.deep_search_a.initial.retrieve_orig_question_docs.states import (
BaseRawSearchOutput,
)
from onyx.agents.agent_search.deep_search_a.main.operations import logger
@ -10,7 +10,7 @@ from onyx.agents.agent_search.deep_search_a.main.states import (
from onyx.agents.agent_search.shared_graph_utils.models import AgentChunkStats
def ingest_initial_base_retrieval(
def ingest_retrieved_documents(
state: BaseRawSearchOutput,
) -> ExpandedRetrievalUpdate:
now_start = datetime.now()

48
backend/onyx/agents/agent_search/deep_search_a/initial/retrieve_orig_question_documents/graph_builder.py → backend/onyx/agents/agent_search/deep_search_a/initial/retrieve_orig_question_docs/graph_builder.py
View File

@ -2,22 +2,19 @@ from langgraph.graph import END
from langgraph.graph import START
from langgraph.graph import StateGraph
from onyx.agents.agent_search.deep_search_a.initial.retrieve_orig_question_documents.nodes.format_raw_search_results import (
format_raw_search_results,
from onyx.agents.agent_search.deep_search_a.initial.retrieve_orig_question_docs.nodes.format_orig_question_search_input import (
format_orig_question_search_input,
)
from onyx.agents.agent_search.deep_search_a.initial.retrieve_orig_question_documents.nodes.generate_raw_search_data import (
generate_raw_search_data,
from onyx.agents.agent_search.deep_search_a.initial.retrieve_orig_question_docs.nodes.format_orig_question_search_output import (
format_orig_question_search_output,
)
from onyx.agents.agent_search.deep_search_a.initial.retrieve_orig_question_documents.nodes.ingest_initial_base_retrieval import (
ingest_initial_base_retrieval,
)
from onyx.agents.agent_search.deep_search_a.initial.retrieve_orig_question_documents.states import (
from onyx.agents.agent_search.deep_search_a.initial.retrieve_orig_question_docs.states import (
BaseRawSearchInput,
)
from onyx.agents.agent_search.deep_search_a.initial.retrieve_orig_question_documents.states import (
from onyx.agents.agent_search.deep_search_a.initial.retrieve_orig_question_docs.states import (
BaseRawSearchOutput,
)
from onyx.agents.agent_search.deep_search_a.initial.retrieve_orig_question_documents.states import (
from onyx.agents.agent_search.deep_search_a.initial.retrieve_orig_question_docs.states import (
BaseRawSearchState,
)
from onyx.agents.agent_search.deep_search_a.shared.expanded_retrieval.graph_builder import (
@ -35,8 +32,8 @@ def base_raw_search_graph_builder() -> StateGraph:
### Add nodes ###
graph.add_node(
node="generate_raw_search_data",
action=generate_raw_search_data,
node="format_orig_question_search_input",
action=format_orig_question_search_input,
)
expanded_retrieval = expanded_retrieval_graph_builder().compile()
@ -45,40 +42,25 @@ def base_raw_search_graph_builder() -> StateGraph:
action=expanded_retrieval,
)
graph.add_node(
node="format_raw_search_results",
action=format_raw_search_results,
)
graph.add_node(
node="ingest_initial_base_retrieval",
action=ingest_initial_base_retrieval,
node="format_orig_question_search_output",
action=format_orig_question_search_output,
)
### Add edges ###
graph.add_edge(start_key=START, end_key="generate_raw_search_data")
graph.add_edge(start_key=START, end_key="format_orig_question_search_input")
graph.add_edge(
start_key="generate_raw_search_data",
start_key="format_orig_question_search_input",
end_key="expanded_retrieval_base_search",
)
graph.add_edge(
start_key="expanded_retrieval_base_search",
end_key="format_raw_search_results",
)
# graph.add_edge(
# start_key="expanded_retrieval_base_search",
# end_key=END,
# )
graph.add_edge(
start_key="format_raw_search_results",
end_key="ingest_initial_base_retrieval",
end_key="format_orig_question_search_output",
)
graph.add_edge(
start_key="ingest_initial_base_retrieval",
start_key="format_orig_question_search_output",
end_key=END,
)

0
backend/onyx/agents/agent_search/deep_search_a/initial/retrieve_orig_question_documents/models.py → backend/onyx/agents/agent_search/deep_search_a/initial/retrieve_orig_question_docs/models.py
View File

2
backend/onyx/agents/agent_search/deep_search_a/initial/retrieve_orig_question_documents/nodes/generate_raw_search_data.py → backend/onyx/agents/agent_search/deep_search_a/initial/retrieve_orig_question_docs/nodes/format_orig_question_search_input.py
View File

@ -12,7 +12,7 @@ from onyx.utils.logger import setup_logger
logger = setup_logger()
def generate_raw_search_data(
def format_orig_question_search_input(
state: CoreState, config: RunnableConfig
) -> ExpandedRetrievalInput:
logger.debug("generate_raw_search_data")

33
backend/onyx/agents/agent_search/deep_search_a/initial/retrieve_orig_question_docs/nodes/format_orig_question_search_output.py Normal file
View File

@ -0,0 +1,33 @@
from onyx.agents.agent_search.deep_search_a.main.states import ExpandedRetrievalUpdate
from onyx.agents.agent_search.deep_search_a.shared.expanded_retrieval.states import (
ExpandedRetrievalOutput,
)
from onyx.agents.agent_search.shared_graph_utils.models import AgentChunkStats
from onyx.utils.logger import setup_logger
logger = setup_logger()
def format_orig_question_search_output(
state: ExpandedRetrievalOutput,
) -> ExpandedRetrievalUpdate:
# return BaseRawSearchOutput(
# base_expanded_retrieval_result=state.expanded_retrieval_result,
# # base_retrieval_results=[state.expanded_retrieval_result],
# # base_search_documents=[],
# )
sub_question_retrieval_stats = (
state.expanded_retrieval_result.sub_question_retrieval_stats
)
if sub_question_retrieval_stats is None:
sub_question_retrieval_stats = AgentChunkStats()
else:
sub_question_retrieval_stats = sub_question_retrieval_stats
return ExpandedRetrievalUpdate(
original_question_retrieval_results=state.expanded_retrieval_result.expanded_queries_results,
all_original_question_documents=state.expanded_retrieval_result.context_documents,
original_question_retrieval_stats=sub_question_retrieval_stats,
log_messages=[],
)

0
backend/onyx/agents/agent_search/deep_search_a/initial/retrieve_orig_question_documents/states.py → backend/onyx/agents/agent_search/deep_search_a/initial/retrieve_orig_question_docs/states.py
View File

18
backend/onyx/agents/agent_search/deep_search_a/initial/retrieve_orig_question_documents/nodes/format_raw_search_results.py
View File

@ -1,18 +0,0 @@
from onyx.agents.agent_search.deep_search_a.initial.retrieve_orig_question_documents.states import (
BaseRawSearchOutput,
)
from onyx.agents.agent_search.deep_search_a.shared.expanded_retrieval.states import (
ExpandedRetrievalOutput,
)
from onyx.utils.logger import setup_logger
logger = setup_logger()
def format_raw_search_results(state: ExpandedRetrievalOutput) -> BaseRawSearchOutput:
logger.debug("format_raw_search_results")
return BaseRawSearchOutput(
base_expanded_retrieval_result=state.expanded_retrieval_result,
# base_retrieval_results=[state.expanded_retrieval_result],
# base_search_documents=[],
)

41
backend/onyx/agents/agent_search/deep_search_a/initial/retrieve_orig_question_documents/nodes/ingest_initial_base_retrieval.py
View File

@ -1,41 +0,0 @@
from datetime import datetime
from onyx.agents.agent_search.deep_search_a.initial.retrieve_orig_question_documents.states import (
BaseRawSearchOutput,
)
from onyx.agents.agent_search.deep_search_a.main.operations import logger
from onyx.agents.agent_search.deep_search_a.main.states import (
ExpandedRetrievalUpdate,
)
from onyx.agents.agent_search.shared_graph_utils.models import AgentChunkStats
def ingest_initial_base_retrieval(
state: BaseRawSearchOutput,
) -> ExpandedRetrievalUpdate:
now_start = datetime.now()
logger.info(f"--------{now_start}--------INGEST INITIAL RETRIEVAL---")
sub_question_retrieval_stats = (
state.base_expanded_retrieval_result.sub_question_retrieval_stats
)
if sub_question_retrieval_stats is None:
sub_question_retrieval_stats = AgentChunkStats()
else:
sub_question_retrieval_stats = sub_question_retrieval_stats
now_end = datetime.now()
logger.debug(
f"--------{now_end}--{now_end - now_start}--------INGEST INITIAL RETRIEVAL END---"
)
return ExpandedRetrievalUpdate(
original_question_retrieval_results=state.base_expanded_retrieval_result.expanded_queries_results,
all_original_question_documents=state.base_expanded_retrieval_result.context_documents,
original_question_retrieval_stats=sub_question_retrieval_stats,
log_messages=[
f"{now_start} -- Main - Ingestion base retrieval, Time taken: {now_end - now_start}"
],
)