EL comments addressed

backend/onyx/agents/agent_search/deep_search_a/main/nodes/answer_comparison.py

@@ -35,10 +35,9 @@ def answer_comparison(state: MainState, config: RunnableConfig) -> AnswerCompari
     # no need to stream this
     resp = model.invoke(msg)
 
-    if isinstance(resp.content, str) and "yes" in resp.content.lower():
+    refined_answer_improvement = (
-        refined_answer_improvement = True
+        isinstance(resp.content, str) and "yes" in resp.content.lower()
-    else:
+    )
-        refined_answer_improvement = False
 
     dispatch_custom_event(
         "refined_answer_improvement",

backend/onyx/agents/agent_search/deep_search_a/main/nodes/generate_initial_answer.py

@@ -66,7 +66,6 @@ def generate_initial_answer(
     history = build_history_prompt(agent_a_config.prompt_builder)
     date_str = get_today_prompt()
 
-    # sub_question_docs = state.documents
     sub_question_docs = state.context_documents
     all_original_question_documents = state.all_original_question_documents
 

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

@@ -48,43 +48,25 @@ class RefinedAgentEndStats(BaseModel):
     agent_refined_metrics: AgentRefinedMetrics = AgentRefinedMetrics()
 
 
-class BaseDecompUpdateBase(BaseModel):
+class BaseDecompUpdate(RefinedAgentStartStats, RefinedAgentEndStats):
     agent_start_time: datetime = datetime.now()
     initial_decomp_questions: list[str] = []
 
 
-class AnswerComparisonBase(BaseModel):
+class AnswerComparison(LoggerUpdate):
     refined_answer_improvement: bool = False
 
 
-class AnswerComparison(AnswerComparisonBase, LoggerUpdate):
+class RoutingDecision(LoggerUpdate):
-    pass
-
-
-class RoutingDecisionBase(BaseModel):
     routing: str = ""
     sample_doc_str: str = ""
 
 
-class RoutingDecision(RoutingDecisionBase, LoggerUpdate):
-    pass
-
-
-class LoggingUpdate(BaseModel):
-    log_messages: list[str] = []
-
-
-class BaseDecompUpdate(
-    RefinedAgentStartStats, RefinedAgentEndStats, BaseDecompUpdateBase
-):
-    pass
-
-
 class InitialAnswerBASEUpdate(BaseModel):
     initial_base_answer: str = ""
 
 
-class InitialAnswerUpdateBase(BaseModel):
+class InitialAnswerUpdate(LoggerUpdate):
     initial_answer: str = ""
     initial_agent_stats: InitialAgentResultStats | None = None
     generated_sub_questions: list[str] = []
@@ -92,29 +74,21 @@ class InitialAnswerUpdateBase(BaseModel):
     agent_base_metrics: AgentBaseMetrics | None = None
 
 
-class InitialAnswerUpdate(InitialAnswerUpdateBase, LoggerUpdate):
+class RefinedAnswerUpdate(RefinedAgentEndStats):
-    pass
-
-
-class RefinedAnswerUpdateBase(BaseModel):
     refined_answer: str = ""
     refined_agent_stats: RefinedAgentStats | None = None
     refined_answer_quality: bool = False
 
 
-class RefinedAnswerUpdate(RefinedAgentEndStats, RefinedAnswerUpdateBase):
+class InitialAnswerQualityUpdate(LoggerUpdate):
-    pass
-
-
-class InitialAnswerQualityUpdate(LoggingUpdate):
     initial_answer_quality: bool = False
 
 
-class RequireRefinedAnswerUpdate(LoggingUpdate):
+class RequireRefinedAnswerUpdate(LoggerUpdate):
     require_refined_answer: bool = True
 
 
-class DecompAnswersUpdate(LoggingUpdate):
+class DecompAnswersUpdate(LoggerUpdate):
     documents: Annotated[list[InferenceSection], dedup_inference_sections] = []
     context_documents: Annotated[list[InferenceSection], dedup_inference_sections] = []
     decomp_answer_results: Annotated[
@@ -122,12 +96,12 @@ class DecompAnswersUpdate(LoggingUpdate):
     ] = []
 
 
-class FollowUpDecompAnswersUpdate(LoggingUpdate):
+class FollowUpDecompAnswersUpdate(LoggerUpdate):
     refined_documents: Annotated[list[InferenceSection], dedup_inference_sections] = []
     refined_decomp_answer_results: Annotated[list[QuestionAnswerResults], add] = []
 
 
-class ExpandedRetrievalUpdate(LoggingUpdate):
+class ExpandedRetrievalUpdate(LoggerUpdate):
     all_original_question_documents: Annotated[
         list[InferenceSection], dedup_inference_sections
     ]
@@ -135,26 +109,16 @@ class ExpandedRetrievalUpdate(LoggingUpdate):
     original_question_retrieval_stats: AgentChunkStats = AgentChunkStats()
 
 
-class EntityTermExtractionUpdateBase(LoggingUpdate):
+class EntityTermExtractionUpdate(LoggerUpdate):
     entity_retlation_term_extractions: EntityRelationshipTermExtraction = (
         EntityRelationshipTermExtraction()
     )
 
 
-class EntityTermExtractionUpdate(EntityTermExtractionUpdateBase, LoggerUpdate):
+class FollowUpSubQuestionsUpdate(RefinedAgentStartStats):
-    pass
-
-
-class FollowUpSubQuestionsUpdateBase(BaseModel):
     refined_sub_questions: dict[int, FollowUpSubQuestion] = {}
 
 
-class FollowUpSubQuestionsUpdate(
-    RefinedAgentStartStats, FollowUpSubQuestionsUpdateBase
-):
-    pass
-
-
 ## Graph Input State
 ## Graph Input State
 
@@ -169,22 +133,21 @@ class MainInput(CoreState):
 class MainState(
     # This includes the core state
     MainInput,
-    LoggerUpdate,
+    BaseDecompUpdate,
-    BaseDecompUpdateBase,
+    InitialAnswerUpdate,
-    InitialAnswerUpdateBase,
     InitialAnswerBASEUpdate,
     DecompAnswersUpdate,
     ExpandedRetrievalUpdate,
-    EntityTermExtractionUpdateBase,
+    EntityTermExtractionUpdate,
     InitialAnswerQualityUpdate,
     RequireRefinedAnswerUpdate,
-    FollowUpSubQuestionsUpdateBase,
+    FollowUpSubQuestionsUpdate,
     FollowUpDecompAnswersUpdate,
-    RefinedAnswerUpdateBase,
+    RefinedAnswerUpdate,
     RefinedAgentStartStats,
     RefinedAgentEndStats,
-    RoutingDecisionBase,
+    RoutingDecision,
-    AnswerComparisonBase,
+    AnswerComparison,
 ):
     # expanded_retrieval_result: Annotated[list[ExpandedRetrievalResult], add]
     base_raw_search_result: Annotated[list[ExpandedRetrievalResult], add]

backend/onyx/agents/agent_search/shared_graph_utils/prompts.py

@@ -724,7 +724,7 @@ is close to the information it supports. If you have multiple citations that sup
 as [[D1]]()[[D3]](), or [[D2]]()[[D4]](), etc.
 Feel free to also cite sub-questions in addition to documents, but make sure that you have documents cited with the sub-question
 citation. If you want to cite both a document and a sub-question, please use [[D1]]()[[Q3]](), or [[D2]]()[[D7]]()[[Q4]](), etc.
-Again, please do NEVER cite sub-questions without a document citation!
+Again, please NEVER cite sub-questions without a document citation!
 Proper citations are very important for the user!
 
 IMPORTANT RULES:
@@ -818,7 +818,7 @@ Please provide inline citations to documents in the format [[D1]](), [[D2]](), [
 citations, please cite for example as [[D1]]()[[D3]](), or [[D2]]()[[D4]](), etc.
 Feel free to also cite sub-questions in addition to documents, but make sure that you have documents cited with the sub-question
 citation. If you want to cite both a document and a sub-question, please use [[D1]]()[[Q3]](), or [[D2]]()[[D7]]()[[Q4]](), etc.
-Again, please do NEVER cite sub-questions without a document citation!
+Again, please NEVER cite sub-questions without a document citation!
 Proper citations are very important for the user!\n\n
 
 {history}
@@ -969,7 +969,7 @@ the refined answer is substantially better than the initial answer. Better could
  Here is the refined answer:
  {refined_answer}
 
- With these criteriain mind, is the refined answer substantially better than the initial answer?
+ With these criteria in mind, is the refined answer substantially better than the initial answer?
 
  Please answer with a simple 'yes' or 'no'.
     """