cleanup

backend/ee/onyx/server/query_and_chat/models.py

@@ -125,10 +125,10 @@ class OneShotQARequest(ChunkContext):
     # will also disable Thread-based Rewording if specified
     query_override: str | None = None
 
-    # If True, skips generative an AI response to the search query
+    # If True, skips generating an AI response to the search query
     skip_gen_ai_answer_generation: bool = False
 
-    # If True, uses pro search instead of basic search
+    # If True, uses agentic search instead of basic search
     use_agentic_search: bool = False
 
     @model_validator(mode="after")

backend/onyx/agents/agent_search/basic/states.py

@@ -12,9 +12,8 @@ from onyx.agents.agent_search.orchestration.states import ToolChoiceUpdate
 # If you are using a value from the config and realize it needs to change,
 # you should add it to the state and use/update the version in the state.
 
+
 ## Graph Input State
-
-
 class BasicInput(BaseModel):
     # Langgraph needs a nonempty input, but we pass in all static
     # data through a RunnableConfig.
@@ -22,18 +21,11 @@ class BasicInput(BaseModel):
 
 
 ## Graph Output State
-
-
 class BasicOutput(TypedDict):
     tool_call_chunk: AIMessageChunk
 
 
-## Update States
-
-
 ## Graph State
-
-
 class BasicState(
     BasicInput,
     ToolChoiceInput,

backend/onyx/agents/agent_search/basic/utils.py

@@ -26,7 +26,6 @@ def process_llm_stream(
     displayed_search_results: list[LlmDoc] | None = None,
 ) -> AIMessageChunk:
     tool_call_chunk = AIMessageChunk(content="")
-    # for response in response_handler_manager.handle_llm_response(stream):
 
     if final_search_results and displayed_search_results:
         answer_handler: AnswerResponseHandler = CitationResponseHandler(

backend/onyx/agents/agent_search/deep_search/initial/generate_initial_answer/edges.py

@@ -25,16 +25,6 @@ def parallelize_initial_sub_question_answering(
     """
     edge_start_time = datetime.now()
     if len(state.initial_sub_questions) > 0:
-        # sub_question_record_ids = [subq_record.id for subq_record in state["sub_question_records"]]
-        # if len(state["sub_question_records"]) == 0:
-        #     if state["config"].use_persistence:
-        #         raise ValueError("No sub-questions found for initial decompozed questions")
-        #     else:
-        #         # in this case, we are doing retrieval on the original question.
-        #         # to make all the logic consistent, we create a new sub-question
-        #         # with the same content as the original question
-        #         sub_question_record_ids = [1] * len(state["initial_decomp_questions"])
-
         return [
             Send(
                 "answer_query_subgraph",

backend/onyx/agents/agent_search/deep_search/initial/generate_initial_answer/graph_builder.py

@@ -83,11 +83,6 @@ def generate_initial_answer_graph_builder(test_mode: bool = False) -> StateGraph
         end_key="generate_initial_answer",
     )
 
-    # graph.add_edge(
-    #     start_key="retrieval_consolidation",
-    #     end_key="generate_initial_answer",
-    # )
-
     graph.add_edge(
         start_key="generate_initial_answer",
         end_key="validate_initial_answer",

backend/onyx/agents/agent_search/deep_search/initial/generate_sub_answers/edges.py

@@ -23,16 +23,6 @@ def parallelize_initial_sub_question_answering(
     """
     edge_start_time = datetime.now()
     if len(state.initial_sub_questions) > 0:
-        # sub_question_record_ids = [subq_record.id for subq_record in state["sub_question_records"]]
-        # if len(state["sub_question_records"]) == 0:
-        #     if state["config"].use_persistence:
-        #         raise ValueError("No sub-questions found for initial decompozed questions")
-        #     else:
-        #         # in this case, we are doing retrieval on the original question.
-        #         # to make all the logic consistent, we create a new sub-question
-        #         # with the same content as the original question
-        #         sub_question_record_ids = [1] * len(state["initial_decomp_questions"])
-
         return [
             Send(
                 "answer_sub_question_subgraphs",

backend/onyx/agents/agent_search/deep_search/initial/generate_sub_answers/graph_builder.py

@@ -58,30 +58,11 @@ def generate_sub_answers_graph_builder() -> StateGraph:
         action=format_initial_sub_answers,
     )
 
-    ### Add edges ###
-
-    # raph.add_edge(start_key=START, end_key="base_raw_search_subgraph")
-
-    # graph.add_edge(
-    #     start_key="start_agent_search",
-    #     end_key="extract_entity_term",
-    # )
-
     graph.add_edge(
         start_key=START,
         end_key="decompose_orig_question",
     )
 
-    # graph.add_edge(
-    #     start_key="LLM",
-    #     end_key=END,
-    # )
-
-    # graph.add_edge(
-    #     start_key=START,
-    #     end_key="initial_sub_question_creation",
-    # )
-
     graph.add_conditional_edges(
         source="decompose_orig_question",
         path=parallelize_initial_sub_question_answering,
@@ -97,14 +78,4 @@ def generate_sub_answers_graph_builder() -> StateGraph:
         end_key=END,
     )
 
-    # graph.add_edge(
-    #     start_key="generate_refined_answer",
-    #     end_key="check_refined_answer",
-    # )
-
-    # graph.add_edge(
-    #     start_key="check_refined_answer",
-    #     end_key=END,
-    # )
-
     return graph

backend/onyx/agents/agent_search/deep_search/initial/retrieve_orig_question_docs/states.py

@@ -5,19 +5,13 @@ from onyx.agents.agent_search.deep_search.shared.expanded_retrieval.states impor
     ExpandedRetrievalInput,
 )
 
-## Update States
-
 
 ## Graph Input State
-
-
 class BaseRawSearchInput(ExpandedRetrievalInput):
     pass
 
 
 ## Graph Output State
-
-
 class BaseRawSearchOutput(OrigQuestionRetrievalUpdate):
     """
     This is a list of results even though each call of this subgraph only returns one result.
@@ -29,8 +23,6 @@ class BaseRawSearchOutput(OrigQuestionRetrievalUpdate):
 
 
 ## Graph State
-
-
 class BaseRawSearchState(
     BaseRawSearchInput, BaseRawSearchOutput, OrigQuestionRetrievalUpdate
 ):

backend/onyx/agents/agent_search/deep_search/main/edges.py

@@ -48,16 +48,6 @@ def parallelize_initial_sub_question_answering(
 ) -> list[Send | Hashable]:
     edge_start_time = datetime.now()
     if len(state.initial_sub_questions) > 0:
-        # sub_question_record_ids = [subq_record.id for subq_record in state["sub_question_records"]]
-        # if len(state["sub_question_records"]) == 0:
-        #     if state["config"].use_persistence:
-        #         raise ValueError("No sub-questions found for initial decompozed questions")
-        #     else:
-        #         # in this case, we are doing retrieval on the original question.
-        #         # to make all the logic consistent, we create a new sub-question
-        #         # with the same content as the original question
-        #         sub_question_record_ids = [1] * len(state["initial_decomp_questions"])
-
         return [
             Send(
                 "answer_query_subgraph",

backend/onyx/agents/agent_search/deep_search/main/states.py

@@ -34,11 +34,8 @@ from onyx.agents.agent_search.shared_graph_utils.operators import (
 )
 from onyx.context.search.models import InferenceSection
 
+
 ### States ###
-
-## Update States
-
-
 class LoggerUpdate(BaseModel):
     log_messages: Annotated[list[str], add] = []
 
@@ -147,15 +144,11 @@ class RefinedQuestionDecompositionUpdate(RefinedAgentStartStats, LoggerUpdate):
 
 
 ## Graph Input State
-
-
 class MainInput(CoreState):
     pass
 
 
 ## Graph State
-
-
 class MainState(
     # This includes the core state
     MainInput,
@@ -181,7 +174,5 @@ class MainState(
 
 
 ## Graph Output State - presently not used
-
-
 class MainOutput(TypedDict):
     log_messages: list[str]

backend/onyx/agents/agent_search/deep_search/refinement/consolidate_sub_answers/models.py

@@ -1,3 +0,0 @@
-### Models ###
-# class AnswerRetrievalStats(BaseModel):
-#     answer_retrieval_stats: dict[str, float | int]

backend/onyx/agents/agent_search/deep_search/shared/expanded_retrieval/nodes/format_results.py

@@ -37,8 +37,8 @@ def format_results(
     query_info = get_query_info(state.query_retrieval_results)
 
     graph_config = cast(GraphConfig, config["metadata"]["config"])
-    # main question docs will be sent later after aggregation and deduping with sub-question docs
 
+    # Main question docs will be sent later after aggregation and deduping with sub-question docs
     reranked_documents = state.reranked_documents
 
     if not (level == 0 and question_num == 0):
@@ -79,8 +79,6 @@ def format_results(
 
     if sub_question_retrieval_stats is None:
         sub_question_retrieval_stats = AgentChunkRetrievalStats()
-    # else:
-    #    sub_question_retrieval_stats = [sub_question_retrieval_stats]
 
     return ExpandedRetrievalUpdate(
         expanded_retrieval_result=QuestionRetrievalResult(

backend/onyx/agents/agent_search/models.py

@@ -17,11 +17,6 @@ class GraphInputs(BaseModel):
     """Input data required for the graph execution"""
 
     search_request: SearchRequest
-    # contains message history for the current chat session
-    # has the following (at most one is non-None)
-    # TODO: unify this into a single message history
-    # message_history: list[PreviousMessage] | None = None
-    # single_message_history: str | None = None
     prompt_builder: AnswerPromptBuilder
     files: list[InMemoryChatFile] | None = None
     structured_response_format: dict | None = None

backend/onyx/agents/agent_search/orchestration/nodes/llm_tool_choice.py

@@ -68,7 +68,7 @@ def llm_tool_choice(
         if chosen_tool_and_args:
             tool, tool_args = chosen_tool_and_args
 
-    # If we have a tool and tool args, we are redy to request a tool call.
+    # If we have a tool and tool args, we are ready to request a tool call.
     # This only happens if the tool call was forced or we are using a non-tool calling LLM.
     if tool and tool_args:
         return ToolChoiceUpdate(