{answer}
+
+ ZenML Deep Research pipeline flow
+
+ ZenML Deep Research pipeline flow
+
+ Sample report generated by the Deep Research Agent
+tags for paragraphs and organize the content +logically with clear transitions between the different aspects outlined above.\n\nIMPORTANT: Do NOT include any headings like "Conclusion",
The research comprehensively demonstrates that while LLMOps shares foundational principles with MLOps, it is not merely a direct subset but rather a specialized and +significantly extended operational framework necessitated by the unique characteristics of Large Language Models (LLMs). Overarching themes reveal that the sheer scale of LLMs, their emergent +behaviors, the criticality of prompt engineering, and the continuous need for pre-training and fine-tuning fundamentally alter the traditional MLOps lifecycle, tooling requirements, and +operational challenges. Findings consistently highlight that LLMs demand specialized deployment strategies, advanced monitoring for unpredictable outputs, and dedicated workflows for prompt +management and RAG. These elements interrelate, with the scale driving the need for efficient serving frameworks, emergent behaviors necessitating sophisticated observability, and continuous +adaptation requiring specialized fine-tuning tools, all converging to form a distinct operational paradigm.
\n\nIn direct response to the main query, the research definitively establishes +that LLMOps is a specialized domain within MLOps, rather than a completely different discipline. It builds upon MLOps' core tenets of data management, model deployment, and monitoring but +introduces a new layer of complexity and specific requirements. The evidence strongly indicates that while MLOps provides the groundwork, the unique attributes of LLMs—such as their computational +demands, the nuances of prompt engineering, and the imperative for continuous adaptation—mandate a distinct set of tools, practices, and expertise. What remains less certain is the precise +boundary where MLOps ends and LLMOps begins, as many MLOps platforms are evolving to incorporate LLM-specific functionalities.
\n\nDespite the comprehensive overview, several uncertainties +and information gaps persist. The research could benefit from more detailed real-world case studies illustrating the practical implementation of LLMOps across diverse industries, particularly +concerning cost optimization and long-term maintenance strategies. There's also a need for deeper exploration into specific frameworks for accountability and bias mitigation beyond technical +fixes, addressing the broader societal implications. Furthermore, the geopolitical dimensions of LLM development and deployment, including data sovereignty and international regulatory trends, +remain largely unexplored. Future research should focus on these areas to provide a more holistic understanding of LLMOps' practical, ethical, and global landscape.
\n\nThe findings carry +significant implications for organizations and practitioners. They underscore the necessity of investing in specialized LLMOps tools and expertise, recognizing that traditional MLOps approaches +alone are insufficient for effectively managing LLMs in production. For stakeholders, this means adapting existing MLOps teams and infrastructure, or building new capabilities, to address the +unique challenges of LLMs, including their computational intensity, the need for robust prompt engineering, and continuous model adaptation. The insights also highlight the critical importance of +integrating ethical considerations and responsible AI practices directly into the operational framework from the outset.
", + 'role': 'assistant', + 'tool_calls': None, + 'function_call': None + }, + usage=Usage(input=4598, output=573, total=5171, unit={prompt.description}
' if prompt.description else ""} + +{prompt.description}
' if prompt.description else ""} + +{state.main_query}
" + else: + html += "No main query specified
" + + html += """ +Found {len(results)} results
+{info.synthesized_answer}
+ """ + + if info.key_sources: + html += """ +Key Sources:
+{state.viewpoint_analysis.perspective_gaps}
+ """ + + if state.viewpoint_analysis.integrative_insights: + html += f""" +{state.viewpoint_analysis.integrative_insights}
+ """ + + html += """ +Improvements Made:
+Final HTML report is available but not displayed here. View the HTML artifact to see the complete report.
+No additional queries recommended
+ """ + + html += """ +No critique summary available
+ """ + + html += """ +No additional questions identified
+ """ + + html += """ +Main Query: {}
+Current Stage: {}
+Sub-questions: {}
+Search Results: {} queries with results
+Synthesized Info: {} topics synthesized
+| Model | +Input Tokens | +Output Tokens | +Total Tokens | +Cost | +
|---|---|---|---|---|
| {model} | +{tokens.get("input_tokens", 0):,} | +{tokens.get("output_tokens", 0):,} | +{tokens.get("total_tokens", 0):,} | +${cost:.4f} | +
| Prompt Type | +Total Cost | +Calls | +Avg $/Call | +% of Total | +Input Tokens | +Output Tokens | +
|---|---|---|---|---|---|---|
| {prompt_type_display} | +${metric.total_cost:.4f} | +{metric.call_count} | +${metric.avg_cost_per_call:.4f} | +{metric.percentage_of_total_cost:.1f}% | +{metric.input_tokens:,} | +{metric.output_tokens:,} | +
+ """ + import json + + html += json.dumps(metadata.trace_metadata, indent=2) + html += """ ++ """ + + html += """ +
{code_content} blocks)
+ parts = []
+ in_pre = False
+ for line in formatted_text.split("\n"):
+ if "" in line:
+ in_pre = True
+ parts.append(line)
+ elif "
" in line:
+ in_pre = False
+ parts.append(line)
+ elif in_pre:
+ # Inside a code block, preserve newlines
+ parts.append(line)
+ else:
+ # Outside code blocks, convert newlines to
+ parts.append(line + "
")
+
+ return "".join(parts)
+
+
+def generate_executive_summary(
+ state: ResearchState,
+ prompts_bundle: PromptsBundle,
+ llm_model: str = "sambanova/DeepSeek-R1-Distill-Llama-70B",
+ langfuse_project_name: str = "deep-research",
+) -> str:
+ """Generate an executive summary using LLM based on research findings.
+
+ Args:
+ state: The current research state
+ prompts_bundle: Bundle containing all prompts for the pipeline
+ llm_model: The model to use for generation
+ langfuse_project_name: Name of the Langfuse project for tracking
+
+ Returns:
+ HTML formatted executive summary
+ """
+ logger.info("Generating executive summary using LLM")
+
+ # Prepare the context with all research findings
+ context = f"Main Research Query: {state.main_query}\n\n"
+
+ # Add synthesized findings for each sub-question
+ for i, sub_question in enumerate(state.sub_questions, 1):
+ info = state.enhanced_info.get(
+ sub_question
+ ) or state.synthesized_info.get(sub_question)
+ if info:
+ context += f"Sub-question {i}: {sub_question}\n"
+ context += f"Answer Summary: {info.synthesized_answer[:500]}...\n"
+ context += f"Confidence: {info.confidence_level}\n"
+ context += f"Key Sources: {', '.join(info.key_sources[:3]) if info.key_sources else 'N/A'}\n\n"
+
+ # Add viewpoint analysis insights if available
+ if state.viewpoint_analysis:
+ context += "Key Areas of Agreement:\n"
+ for agreement in state.viewpoint_analysis.main_points_of_agreement[:3]:
+ context += f"- {agreement}\n"
+ context += "\nKey Tensions:\n"
+ for tension in state.viewpoint_analysis.areas_of_tension[:2]:
+ context += f"- {tension.topic}\n"
+
+ # Get the executive summary prompt
+ try:
+ executive_summary_prompt = prompts_bundle.get_prompt_content(
+ "executive_summary_prompt"
+ )
+ logger.info("Successfully retrieved executive_summary_prompt")
+ except Exception as e:
+ logger.error(f"Failed to get executive_summary_prompt: {e}")
+ logger.info(
+ f"Available prompts: {list(prompts_bundle.list_all_prompts().keys())}"
+ )
+ return generate_fallback_executive_summary(state)
+
+ try:
+ # Call LLM to generate executive summary
+ result = run_llm_completion(
+ prompt=context,
+ system_prompt=executive_summary_prompt,
+ model=llm_model,
+ temperature=0.7,
+ max_tokens=800,
+ project=langfuse_project_name,
+ tags=["executive_summary_generation"],
+ )
+
+ if result:
+ content = remove_reasoning_from_output(result)
+ # Clean up the HTML
+ content = extract_html_from_content(content)
+ logger.info("Successfully generated LLM-based executive summary")
+ return content
+ else:
+ logger.warning("Failed to generate executive summary via LLM")
+ return generate_fallback_executive_summary(state)
+
+ except Exception as e:
+ logger.error(f"Error generating executive summary: {e}")
+ return generate_fallback_executive_summary(state)
+
+
+def generate_introduction(
+ state: ResearchState,
+ prompts_bundle: PromptsBundle,
+ llm_model: str = "sambanova/DeepSeek-R1-Distill-Llama-70B",
+ langfuse_project_name: str = "deep-research",
+) -> str:
+ """Generate an introduction using LLM based on research query and sub-questions.
+
+ Args:
+ state: The current research state
+ prompts_bundle: Bundle containing all prompts for the pipeline
+ llm_model: The model to use for generation
+ langfuse_project_name: Name of the Langfuse project for tracking
+
+ Returns:
+ HTML formatted introduction
+ """
+ logger.info("Generating introduction using LLM")
+
+ # Prepare the context
+ context = f"Main Research Query: {state.main_query}\n\n"
+ context += "Sub-questions being explored:\n"
+ for i, sub_question in enumerate(state.sub_questions, 1):
+ context += f"{i}. {sub_question}\n"
+
+ # Get the introduction prompt
+ try:
+ introduction_prompt = prompts_bundle.get_prompt_content(
+ "introduction_prompt"
+ )
+ logger.info("Successfully retrieved introduction_prompt")
+ except Exception as e:
+ logger.error(f"Failed to get introduction_prompt: {e}")
+ logger.info(
+ f"Available prompts: {list(prompts_bundle.list_all_prompts().keys())}"
+ )
+ return generate_fallback_introduction(state)
+
+ try:
+ # Call LLM to generate introduction
+ result = run_llm_completion(
+ prompt=context,
+ system_prompt=introduction_prompt,
+ model=llm_model,
+ temperature=0.7,
+ max_tokens=600,
+ project=langfuse_project_name,
+ tags=["introduction_generation"],
+ )
+
+ if result:
+ content = remove_reasoning_from_output(result)
+ # Clean up the HTML
+ content = extract_html_from_content(content)
+ logger.info("Successfully generated LLM-based introduction")
+ return content
+ else:
+ logger.warning("Failed to generate introduction via LLM")
+ return generate_fallback_introduction(state)
+
+ except Exception as e:
+ logger.error(f"Error generating introduction: {e}")
+ return generate_fallback_introduction(state)
+
+
+def generate_fallback_executive_summary(state: ResearchState) -> str:
+ """Generate a fallback executive summary when LLM fails."""
+ summary = f"This report examines the question: {html.escape(state.main_query)}
"
+ summary += f"The research explored {len(state.sub_questions)} key dimensions of this topic, "
+ summary += "synthesizing findings from multiple sources to provide a comprehensive analysis.
"
+
+ # Add confidence overview
+ confidence_counts = {"high": 0, "medium": 0, "low": 0}
+ for info in state.enhanced_info.values():
+ level = info.confidence_level.lower()
+ if level in confidence_counts:
+ confidence_counts[level] += 1
+
+ summary += f"Overall confidence in findings: {confidence_counts['high']} high, "
+ summary += f"{confidence_counts['medium']} medium, {confidence_counts['low']} low.
"
+
+ return summary
+
+
+def generate_fallback_introduction(state: ResearchState) -> str:
+ """Generate a fallback introduction when LLM fails."""
+ intro = f"This report addresses the research query: {html.escape(state.main_query)}
"
+ intro += f"The research was conducted by breaking down the main query into {len(state.sub_questions)} "
+ intro += (
+ "sub-questions to explore different aspects of the topic in depth. "
+ )
+ intro += "Each sub-question was researched independently, with findings synthesized from various sources.
"
+ return intro
+
+
+def generate_conclusion(
+ state: ResearchState,
+ prompts_bundle: PromptsBundle,
+ llm_model: str = "sambanova/DeepSeek-R1-Distill-Llama-70B",
+ langfuse_project_name: str = "deep-research",
+) -> str:
+ """Generate a comprehensive conclusion using LLM based on all research findings.
+
+ Args:
+ state: The ResearchState containing all research findings
+ prompts_bundle: Bundle containing all prompts for the pipeline
+ llm_model: The model to use for conclusion generation
+
+ Returns:
+ str: HTML-formatted conclusion content
+ """
+ logger.info("Generating comprehensive conclusion using LLM")
+
+ # Prepare input data for conclusion generation
+ conclusion_input = {
+ "main_query": state.main_query,
+ "sub_questions": state.sub_questions,
+ "enhanced_info": {},
+ }
+
+ # Include enhanced information for each sub-question
+ for question in state.sub_questions:
+ if question in state.enhanced_info:
+ info = state.enhanced_info[question]
+ conclusion_input["enhanced_info"][question] = {
+ "synthesized_answer": info.synthesized_answer,
+ "confidence_level": info.confidence_level,
+ "information_gaps": info.information_gaps,
+ "key_sources": info.key_sources,
+ "improvements": getattr(info, "improvements", []),
+ }
+ elif question in state.synthesized_info:
+ # Fallback to synthesized info if enhanced info not available
+ info = state.synthesized_info[question]
+ conclusion_input["enhanced_info"][question] = {
+ "synthesized_answer": info.synthesized_answer,
+ "confidence_level": info.confidence_level,
+ "information_gaps": info.information_gaps,
+ "key_sources": info.key_sources,
+ "improvements": [],
+ }
+
+ # Include viewpoint analysis if available
+ if state.viewpoint_analysis:
+ conclusion_input["viewpoint_analysis"] = {
+ "main_points_of_agreement": state.viewpoint_analysis.main_points_of_agreement,
+ "areas_of_tension": [
+ {"topic": tension.topic, "viewpoints": tension.viewpoints}
+ for tension in state.viewpoint_analysis.areas_of_tension
+ ],
+ "perspective_gaps": state.viewpoint_analysis.perspective_gaps,
+ "integrative_insights": state.viewpoint_analysis.integrative_insights,
+ }
+
+ # Include reflection metadata if available
+ if state.reflection_metadata:
+ conclusion_input["reflection_metadata"] = {
+ "critique_summary": state.reflection_metadata.critique_summary,
+ "additional_questions_identified": state.reflection_metadata.additional_questions_identified,
+ "improvements_made": state.reflection_metadata.improvements_made,
+ }
+
+ try:
+ # Get the prompt from the bundle
+ conclusion_prompt = prompts_bundle.get_prompt_content(
+ "conclusion_generation_prompt"
+ )
+
+ # Generate conclusion using LLM
+ conclusion_html = run_llm_completion(
+ prompt=json.dumps(conclusion_input, indent=2),
+ system_prompt=conclusion_prompt,
+ model=llm_model,
+ clean_output=True,
+ max_tokens=1500, # Sufficient for comprehensive conclusion
+ project=langfuse_project_name,
+ )
+
+ # Clean up any formatting issues
+ conclusion_html = conclusion_html.strip()
+
+ # Remove any h2 tags with "Conclusion" text that LLM might have added
+ # Since we already have a Conclusion header in the template
+ conclusion_html = re.sub(
+ r"]*>\s*Conclusion\s*\s*",
+ "",
+ conclusion_html,
+ flags=re.IGNORECASE,
+ )
+ conclusion_html = re.sub(
+ r"]*>\s*Conclusion\s*\s*",
+ "",
+ conclusion_html,
+ flags=re.IGNORECASE,
+ )
+
+ # Also remove plain text "Conclusion" at the start if it exists
+ conclusion_html = re.sub(
+ r"^Conclusion\s*\n*",
+ "",
+ conclusion_html.strip(),
+ flags=re.IGNORECASE,
+ )
+
+ if not conclusion_html.startswith(""):
+ # Wrap in paragraph tags if not already formatted
+ conclusion_html = f"
{conclusion_html}
"
+
+ logger.info("Successfully generated LLM-based conclusion")
+ return conclusion_html
+
+ except Exception as e:
+ logger.warning(f"Failed to generate LLM conclusion: {e}")
+ # Return a basic fallback conclusion
+ return f"""This report has explored {html.escape(state.main_query)} through a structured research approach, examining {len(state.sub_questions)} focused sub-questions and synthesizing information from diverse sources. The findings provide a comprehensive understanding of the topic, highlighting key aspects, perspectives, and current knowledge.
+ While some information gaps remain, as noted in the respective sections, this research provides a solid foundation for understanding the topic and its implications.
"""
+
+
+def generate_report_from_template(
+ state: ResearchState,
+ prompts_bundle: PromptsBundle,
+ llm_model: str = "sambanova/DeepSeek-R1-Distill-Llama-70B",
+ langfuse_project_name: str = "deep-research",
+) -> str:
+ """Generate a final HTML report from a static template.
+
+ Instead of using an LLM to generate HTML, this function uses predefined HTML
+ templates and populates them with data from the research state.
+
+ Args:
+ state: The current research state
+ prompts_bundle: Bundle containing all prompts for the pipeline
+ llm_model: The model to use for conclusion generation
+
+ Returns:
+ str: The HTML content of the report
+ """
+ logger.info(
+ f"Generating templated HTML report for query: {state.main_query}"
+ )
+
+ # Generate table of contents for sub-questions
+ sub_questions_toc = ""
+ for i, question in enumerate(state.sub_questions, 1):
+ safe_id = f"question-{i}"
+ sub_questions_toc += (
+ f'{html.escape(question)} \n'
+ )
+
+ # Add viewpoint analysis to TOC if available
+ additional_sections_toc = ""
+ if state.viewpoint_analysis:
+ additional_sections_toc += (
+ 'Viewpoint Analysis \n'
+ )
+
+ # Generate HTML for sub-questions
+ sub_questions_html = ""
+ all_sources = set()
+
+ for i, question in enumerate(state.sub_questions, 1):
+ info = state.enhanced_info.get(question, None)
+
+ # Skip if no information is available
+ if not info:
+ continue
+
+ # Process confidence level
+ confidence = info.confidence_level.lower()
+ confidence_upper = info.confidence_level.upper()
+
+ # Process key sources
+ key_sources_html = ""
+ if info.key_sources:
+ all_sources.update(info.key_sources)
+ sources_list = "\n".join(
+ [
+ f'{html.escape(source)} '
+ if source.startswith(("http://", "https://"))
+ else f"{html.escape(source)} "
+ for source in info.key_sources
+ ]
+ )
+ key_sources_html = f"""
+
+ Key Sources
+
+ {sources_list}
+
+
+ """
+
+ # Process information gaps
+ info_gaps_html = ""
+ if info.information_gaps:
+ info_gaps_html = f"""
+
+ Information Gaps
+ {format_text_with_code_blocks(info.information_gaps)}
+
+ """
+
+ # Determine confidence icon based on level
+ confidence_icon = "🔴" # Default (low)
+ if confidence_upper == "HIGH":
+ confidence_icon = "🟢"
+ elif confidence_upper == "MEDIUM":
+ confidence_icon = "🟡"
+
+ # Format the subquestion section using the template
+ sub_question_html = SUB_QUESTION_TEMPLATE.format(
+ index=i,
+ question=html.escape(question),
+ confidence=confidence,
+ confidence_upper=confidence_upper,
+ confidence_icon=confidence_icon,
+ answer=format_text_with_code_blocks(info.synthesized_answer),
+ info_gaps_html=info_gaps_html,
+ key_sources_html=key_sources_html,
+ )
+
+ sub_questions_html += sub_question_html
+
+ # Generate viewpoint analysis HTML if available
+ viewpoint_analysis_html = ""
+ if state.viewpoint_analysis:
+ # Format points of agreement
+ agreements_html = ""
+ for point in state.viewpoint_analysis.main_points_of_agreement:
+ agreements_html += f"{html.escape(point)} \n"
+
+ # Format areas of tension
+ tensions_html = ""
+ for tension in state.viewpoint_analysis.areas_of_tension:
+ viewpoints_html = ""
+ for title, content in tension.viewpoints.items():
+ # Create category-specific styling
+ category_class = f"category-{title.lower()}"
+ category_title = title.capitalize()
+
+ viewpoints_html += f"""
+
+ {category_title}
+ {html.escape(content)}
+
+ """
+
+ tensions_html += f"""
+
+ {html.escape(tension.topic)}
+
+ {viewpoints_html}
+
+
+ """
+
+ # Format the viewpoint analysis section using the template
+ viewpoint_analysis_html = VIEWPOINT_ANALYSIS_TEMPLATE.format(
+ agreements_html=agreements_html,
+ tensions_html=tensions_html,
+ perspective_gaps=format_text_with_code_blocks(
+ state.viewpoint_analysis.perspective_gaps
+ ),
+ integrative_insights=format_text_with_code_blocks(
+ state.viewpoint_analysis.integrative_insights
+ ),
+ )
+
+ # Generate references HTML
+ references_html = ""
+ if all_sources:
+ for source in sorted(all_sources):
+ if source.startswith(("http://", "https://")):
+ references_html += f'- {html.escape(source)}
\n'
+ else:
+ references_html += f"- {html.escape(source)}
\n"
+ else:
+ references_html += (
+ "- No external sources were referenced in this research.
"
+ )
+ references_html += "
"
+
+ # Generate dynamic executive summary using LLM
+ logger.info("Generating dynamic executive summary...")
+ executive_summary = generate_executive_summary(
+ state, prompts_bundle, llm_model, langfuse_project_name
+ )
+ logger.info(
+ f"Executive summary generated: {len(executive_summary)} characters"
+ )
+
+ # Generate dynamic introduction using LLM
+ logger.info("Generating dynamic introduction...")
+ introduction_html = generate_introduction(
+ state, prompts_bundle, llm_model, langfuse_project_name
+ )
+ logger.info(f"Introduction generated: {len(introduction_html)} characters")
+
+ # Generate comprehensive conclusion using LLM
+ conclusion_html = generate_conclusion(
+ state, prompts_bundle, llm_model, langfuse_project_name
+ )
+
+ # Generate complete HTML report
+ html_content = STATIC_HTML_TEMPLATE.format(
+ main_query=html.escape(state.main_query),
+ sub_questions_toc=sub_questions_toc,
+ additional_sections_toc=additional_sections_toc,
+ executive_summary=executive_summary,
+ introduction_html=introduction_html,
+ num_sub_questions=len(state.sub_questions),
+ sub_questions_html=sub_questions_html,
+ viewpoint_analysis_html=viewpoint_analysis_html,
+ conclusion_html=conclusion_html,
+ references_html=references_html,
+ )
+
+ return html_content
+
+
+def _generate_fallback_report(state: ResearchState) -> str:
+ """Generate a minimal fallback report when the main report generation fails.
+
+ This function creates a simplified HTML report with a consistent structure when
+ the main report generation process encounters an error. The HTML includes:
+ - A header section with the main research query
+ - An error notice
+ - Introduction section
+ - Individual sections for each sub-question with available answers
+ - A references section if sources are available
+
+ Args:
+ state: The current research state containing query and answer information
+
+ Returns:
+ str: A basic HTML report with a standard research report structure
+ """
+ # Create a simple HTML structure with embedded CSS for styling
+ html = f"""
+
+
+
+
+
+
+
+
+ Research Report: {state.main_query}
+
+
+ Note: This is a fallback report generated due to an error in the report generation process.
+
+
+
+
+ Table of Contents
+
+ - Introduction
+"""
+
+ # Add TOC entries for each sub-question
+ for i, sub_question in enumerate(state.sub_questions):
+ safe_id = f"question-{i + 1}"
+ html += f' - {sub_question}
\n'
+
+ html += """ - References
+
+
+
+
+
+ Executive Summary
+ This report presents findings related to the main research query. It explores multiple aspects of the topic through structured sub-questions and synthesizes information from various sources.
+
+
+
+ Introduction
+ This report addresses the research query: "{state.main_query}"
+ The analysis is structured around {len(state.sub_questions)} sub-questions that explore different dimensions of this topic.
+
+"""
+
+ # Add each sub-question and its synthesized information
+ for i, sub_question in enumerate(state.sub_questions):
+ safe_id = f"question-{i + 1}"
+ info = state.enhanced_info.get(sub_question, None)
+
+ if not info:
+ # Try to get from synthesized info if not in enhanced info
+ info = state.synthesized_info.get(sub_question, None)
+
+ if info:
+ answer = info.synthesized_answer
+ confidence = info.confidence_level
+
+ # Add appropriate confidence class
+ confidence_class = ""
+ if confidence == "high":
+ confidence_class = "confidence-high"
+ elif confidence == "medium":
+ confidence_class = "confidence-medium"
+ elif confidence == "low":
+ confidence_class = "confidence-low"
+
+ html += f"""
+
+ {i + 1}. {sub_question}
+ Confidence Level: {confidence.upper()}
+
+ {answer}
+
+ """
+
+ # Add information gaps if available
+ if hasattr(info, "information_gaps") and info.information_gaps:
+ html += f"""
+
+ Information Gaps
+ {info.information_gaps}
+
+ """
+
+ # Add improvements if available
+ if hasattr(info, "improvements") and info.improvements:
+ html += """
+
+ Improvements Made
+
+ """
+
+ for improvement in info.improvements:
+ html += f" - {improvement}
\n"
+
+ html += """
+
+
+ """
+
+ # Add key sources if available
+ if hasattr(info, "key_sources") and info.key_sources:
+ html += """
+
+ Key Sources
+
+ """
+
+ for source in info.key_sources:
+ html += f" - {source}
\n"
+
+ html += """
+
+
+ """
+
+ html += """
+
+ """
+ else:
+ html += f"""
+
+ {i + 1}. {sub_question}
+ No information available for this question.
+
+ """
+
+ # Add conclusion section
+ html += """
+
+ Conclusion
+ This report has explored the research query through multiple sub-questions, providing synthesized information based on available sources. While limitations exist in some areas, the report provides a structured analysis of the topic.
+
+ """
+
+ # Add sources if available
+ sources_set = set()
+ for info in state.enhanced_info.values():
+ if info.key_sources:
+ sources_set.update(info.key_sources)
+
+ if sources_set:
+ html += """
+
+ References
+
+ """
+
+ for source in sorted(sources_set):
+ html += f" - {source}
\n"
+
+ html += """
+
+
+ """
+ else:
+ html += """
+
+ References
+ No references available.
+
+ """
+
+ # Close the HTML structure
+ html += """
+
+
+
+ """
+
+ return html
+
+
+@step(
+ output_materializers={
+ "state": ResearchStateMaterializer,
+ }
+)
+def pydantic_final_report_step(
+ state: ResearchState,
+ prompts_bundle: PromptsBundle,
+ use_static_template: bool = True,
+ llm_model: str = "sambanova/DeepSeek-R1-Distill-Llama-70B",
+ langfuse_project_name: str = "deep-research",
+) -> Tuple[
+ Annotated[ResearchState, "state"],
+ Annotated[HTMLString, "report_html"],
+]:
+ """Generate the final research report in HTML format using Pydantic models.
+
+ This step uses the Pydantic models and materializers to generate a final
+ HTML report and return both the updated state and the HTML report as
+ separate artifacts.
+
+ Args:
+ state: The current research state (Pydantic model)
+ prompts_bundle: Bundle containing all prompts for the pipeline
+ use_static_template: Whether to use a static template instead of LLM generation
+ llm_model: The model to use for report generation with provider prefix
+
+ Returns:
+ A tuple containing the updated research state and the HTML report
+ """
+ start_time = time.time()
+ logger.info("Generating final research report using Pydantic models")
+
+ if use_static_template:
+ # Use the static HTML template approach
+ logger.info("Using static HTML template for report generation")
+ html_content = generate_report_from_template(
+ state, prompts_bundle, llm_model, langfuse_project_name
+ )
+
+ # Update the state with the final report HTML
+ state.set_final_report(html_content)
+
+ # Collect metadata about the report
+ execution_time = time.time() - start_time
+
+ # Count sources
+ all_sources = set()
+ for info in state.enhanced_info.values():
+ if info.key_sources:
+ all_sources.update(info.key_sources)
+
+ # Count confidence levels
+ confidence_distribution = {"high": 0, "medium": 0, "low": 0}
+ for info in state.enhanced_info.values():
+ level = info.confidence_level.lower()
+ if level in confidence_distribution:
+ confidence_distribution[level] += 1
+
+ # Log metadata
+ log_metadata(
+ metadata={
+ "report_generation": {
+ "execution_time_seconds": execution_time,
+ "generation_method": "static_template",
+ "llm_model": llm_model,
+ "report_length_chars": len(html_content),
+ "num_sub_questions": len(state.sub_questions),
+ "num_sources": len(all_sources),
+ "has_viewpoint_analysis": bool(state.viewpoint_analysis),
+ "has_reflection": bool(state.reflection_metadata),
+ "confidence_distribution": confidence_distribution,
+ "fallback_report": False,
+ }
+ }
+ )
+
+ # Log model metadata for cross-pipeline tracking
+ log_metadata(
+ metadata={
+ "research_quality": {
+ "confidence_distribution": confidence_distribution,
+ }
+ },
+ infer_model=True,
+ )
+
+ # Log artifact metadata for the HTML report
+ log_metadata(
+ metadata={
+ "html_report_characteristics": {
+ "size_bytes": len(html_content.encode("utf-8")),
+ "has_toc": "toc" in html_content.lower(),
+ "has_executive_summary": "executive summary"
+ in html_content.lower(),
+ "has_conclusion": "conclusion" in html_content.lower(),
+ "has_references": "references" in html_content.lower(),
+ }
+ },
+ infer_artifact=True,
+ artifact_name="report_html",
+ )
+
+ logger.info(
+ "Final research report generated successfully with static template"
+ )
+ return state, HTMLString(html_content)
+
+ # Otherwise use the LLM-generated approach
+ # Convert Pydantic model to dict for LLM input
+ report_input = {
+ "main_query": state.main_query,
+ "sub_questions": state.sub_questions,
+ "synthesized_information": state.enhanced_info,
+ }
+
+ if state.viewpoint_analysis:
+ report_input["viewpoint_analysis"] = state.viewpoint_analysis
+
+ if state.reflection_metadata:
+ report_input["reflection_metadata"] = state.reflection_metadata
+
+ # Generate the report
+ try:
+ logger.info(f"Calling {llm_model} to generate final report")
+
+ # Get the prompt from the bundle
+ report_prompt = prompts_bundle.get_prompt_content(
+ "report_generation_prompt"
+ )
+
+ # Use the utility function to run LLM completion
+ html_content = run_llm_completion(
+ prompt=json.dumps(report_input),
+ system_prompt=report_prompt,
+ model=llm_model,
+ clean_output=False, # Don't clean in case of breaking HTML formatting
+ max_tokens=4000, # Increased token limit for detailed report generation
+ project=langfuse_project_name,
+ )
+
+ # Clean up any JSON wrapper or other artifacts
+ html_content = remove_reasoning_from_output(html_content)
+
+ # Process the HTML content to remove code block markers and fix common issues
+ html_content = clean_html_output(html_content)
+
+ # Basic validation of HTML content
+ if not html_content.strip().startswith("<"):
+ logger.warning(
+ "Generated content does not appear to be valid HTML"
+ )
+ # Try to extract HTML if it might be wrapped in code blocks or JSON
+ html_content = extract_html_from_content(html_content)
+
+ # Update the state with the final report HTML
+ state.set_final_report(html_content)
+
+ # Collect metadata about the report
+ execution_time = time.time() - start_time
+
+ # Count sources
+ all_sources = set()
+ for info in state.enhanced_info.values():
+ if info.key_sources:
+ all_sources.update(info.key_sources)
+
+ # Count confidence levels
+ confidence_distribution = {"high": 0, "medium": 0, "low": 0}
+ for info in state.enhanced_info.values():
+ level = info.confidence_level.lower()
+ if level in confidence_distribution:
+ confidence_distribution[level] += 1
+
+ # Log metadata
+ log_metadata(
+ metadata={
+ "report_generation": {
+ "execution_time_seconds": execution_time,
+ "generation_method": "llm_generated",
+ "llm_model": llm_model,
+ "report_length_chars": len(html_content),
+ "num_sub_questions": len(state.sub_questions),
+ "num_sources": len(all_sources),
+ "has_viewpoint_analysis": bool(state.viewpoint_analysis),
+ "has_reflection": bool(state.reflection_metadata),
+ "confidence_distribution": confidence_distribution,
+ "fallback_report": False,
+ }
+ }
+ )
+
+ # Log model metadata for cross-pipeline tracking
+ log_metadata(
+ metadata={
+ "research_quality": {
+ "confidence_distribution": confidence_distribution,
+ }
+ },
+ infer_model=True,
+ )
+
+ logger.info("Final research report generated successfully")
+ return state, HTMLString(html_content)
+
+ except Exception as e:
+ logger.error(f"Error generating final report: {e}")
+ # Generate a minimal fallback report
+ fallback_html = _generate_fallback_report(state)
+
+ # Process the fallback HTML to ensure it's clean
+ fallback_html = clean_html_output(fallback_html)
+
+ # Update the state with the fallback report
+ state.set_final_report(fallback_html)
+
+ # Collect metadata about the fallback report
+ execution_time = time.time() - start_time
+
+ # Count sources
+ all_sources = set()
+ for info in state.enhanced_info.values():
+ if info.key_sources:
+ all_sources.update(info.key_sources)
+
+ # Count confidence levels
+ confidence_distribution = {"high": 0, "medium": 0, "low": 0}
+ for info in state.enhanced_info.values():
+ level = info.confidence_level.lower()
+ if level in confidence_distribution:
+ confidence_distribution[level] += 1
+
+ # Log metadata for fallback report
+ log_metadata(
+ metadata={
+ "report_generation": {
+ "execution_time_seconds": execution_time,
+ "generation_method": "fallback",
+ "llm_model": llm_model,
+ "report_length_chars": len(fallback_html),
+ "num_sub_questions": len(state.sub_questions),
+ "num_sources": len(all_sources),
+ "has_viewpoint_analysis": bool(state.viewpoint_analysis),
+ "has_reflection": bool(state.reflection_metadata),
+ "confidence_distribution": confidence_distribution,
+ "fallback_report": True,
+ "error_message": str(e),
+ }
+ }
+ )
+
+ # Log model metadata for cross-pipeline tracking
+ log_metadata(
+ metadata={
+ "research_quality": {
+ "confidence_distribution": confidence_distribution,
+ }
+ },
+ infer_model=True,
+ )
+
+ return state, HTMLString(fallback_html)
diff --git a/deep_research/steps/query_decomposition_step.py b/deep_research/steps/query_decomposition_step.py
new file mode 100644
index 00000000..78e50b9f
--- /dev/null
+++ b/deep_research/steps/query_decomposition_step.py
@@ -0,0 +1,174 @@
+import logging
+import time
+from typing import Annotated
+
+from materializers.pydantic_materializer import ResearchStateMaterializer
+from utils.llm_utils import get_structured_llm_output
+from utils.prompt_models import PromptsBundle
+from utils.pydantic_models import ResearchState
+from zenml import log_metadata, step
+
+logger = logging.getLogger(__name__)
+
+
+@step(output_materializers=ResearchStateMaterializer)
+def initial_query_decomposition_step(
+ state: ResearchState,
+ prompts_bundle: PromptsBundle,
+ llm_model: str = "sambanova/DeepSeek-R1-Distill-Llama-70B",
+ max_sub_questions: int = 8,
+ langfuse_project_name: str = "deep-research",
+) -> Annotated[ResearchState, "updated_state"]:
+ """Break down a complex research query into specific sub-questions.
+
+ Args:
+ state: The current research state
+ prompts_bundle: Bundle containing all prompts for the pipeline
+ llm_model: The reasoning model to use with provider prefix
+ max_sub_questions: Maximum number of sub-questions to generate
+
+ Returns:
+ Updated research state with sub-questions
+ """
+ start_time = time.time()
+ logger.info(f"Decomposing research query: {state.main_query}")
+
+ # Get the prompt from the bundle
+ system_prompt = prompts_bundle.get_prompt_content(
+ "query_decomposition_prompt"
+ )
+
+ try:
+ # Call OpenAI API to decompose the query
+ updated_system_prompt = (
+ system_prompt
+ + f"\nPlease generate at most {max_sub_questions} sub-questions."
+ )
+ logger.info(
+ f"Calling {llm_model} to decompose query into max {max_sub_questions} sub-questions"
+ )
+
+ # Define fallback questions
+ fallback_questions = [
+ {
+ "sub_question": f"What is {state.main_query}?",
+ "reasoning": "Basic understanding of the topic",
+ },
+ {
+ "sub_question": f"What are the key aspects of {state.main_query}?",
+ "reasoning": "Exploring important dimensions",
+ },
+ {
+ "sub_question": f"What are the implications of {state.main_query}?",
+ "reasoning": "Understanding broader impact",
+ },
+ ]
+
+ # Use utility function to get structured output
+ decomposed_questions = get_structured_llm_output(
+ prompt=state.main_query,
+ system_prompt=updated_system_prompt,
+ model=llm_model,
+ fallback_response=fallback_questions,
+ project=langfuse_project_name,
+ )
+
+ # Extract just the sub-questions
+ sub_questions = [
+ item.get("sub_question")
+ for item in decomposed_questions
+ if "sub_question" in item
+ ]
+
+ # Limit to max_sub_questions
+ sub_questions = sub_questions[:max_sub_questions]
+
+ logger.info(f"Generated {len(sub_questions)} sub-questions")
+ for i, question in enumerate(sub_questions, 1):
+ logger.info(f" {i}. {question}")
+
+ # Update the state with the new sub-questions
+ state.update_sub_questions(sub_questions)
+
+ # Log step metadata
+ execution_time = time.time() - start_time
+ log_metadata(
+ metadata={
+ "query_decomposition": {
+ "execution_time_seconds": execution_time,
+ "num_sub_questions": len(sub_questions),
+ "llm_model": llm_model,
+ "max_sub_questions_requested": max_sub_questions,
+ "fallback_used": False,
+ "main_query_length": len(state.main_query),
+ "sub_questions": sub_questions,
+ }
+ }
+ )
+
+ # Log model metadata for cross-pipeline tracking
+ log_metadata(
+ metadata={
+ "research_scope": {
+ "num_sub_questions": len(sub_questions),
+ }
+ },
+ infer_model=True,
+ )
+
+ # Log artifact metadata for the output state
+ log_metadata(
+ metadata={
+ "state_characteristics": {
+ "total_sub_questions": len(state.sub_questions),
+ "has_search_results": bool(state.search_results),
+ "has_synthesized_info": bool(state.synthesized_info),
+ }
+ },
+ infer_artifact=True,
+ )
+
+ return state
+
+ except Exception as e:
+ logger.error(f"Error decomposing query: {e}")
+ # Return fallback questions in the state
+ fallback_questions = [
+ f"What is {state.main_query}?",
+ f"What are the key aspects of {state.main_query}?",
+ f"What are the implications of {state.main_query}?",
+ ]
+ fallback_questions = fallback_questions[:max_sub_questions]
+ logger.info(f"Using {len(fallback_questions)} fallback questions:")
+ for i, question in enumerate(fallback_questions, 1):
+ logger.info(f" {i}. {question}")
+ state.update_sub_questions(fallback_questions)
+
+ # Log metadata for fallback scenario
+ execution_time = time.time() - start_time
+ log_metadata(
+ metadata={
+ "query_decomposition": {
+ "execution_time_seconds": execution_time,
+ "num_sub_questions": len(fallback_questions),
+ "llm_model": llm_model,
+ "max_sub_questions_requested": max_sub_questions,
+ "fallback_used": True,
+ "error_message": str(e),
+ "main_query_length": len(state.main_query),
+ "sub_questions": fallback_questions,
+ }
+ }
+ )
+
+ # Log model metadata for cross-pipeline tracking
+ log_metadata(
+ metadata={
+ "research_scope": {
+ "num_sub_questions": len(fallback_questions),
+ }
+ },
+ infer_model=True,
+ )
+
+ return state
diff --git a/deep_research/tests/__init__.py b/deep_research/tests/__init__.py
new file mode 100644
index 00000000..6206856b
--- /dev/null
+++ b/deep_research/tests/__init__.py
@@ -0,0 +1 @@
+"""Test package for ZenML Deep Research project."""
diff --git a/deep_research/tests/conftest.py b/deep_research/tests/conftest.py
new file mode 100644
index 00000000..b972a5e1
--- /dev/null
+++ b/deep_research/tests/conftest.py
@@ -0,0 +1,11 @@
+"""Test configuration for pytest.
+
+This file sets up the proper Python path for importing modules in tests.
+"""
+
+import os
+import sys
+
+# Add the project root directory to the Python path
+project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
+sys.path.insert(0, project_root)
diff --git a/deep_research/tests/test_approval_utils.py b/deep_research/tests/test_approval_utils.py
new file mode 100644
index 00000000..fe859b0f
--- /dev/null
+++ b/deep_research/tests/test_approval_utils.py
@@ -0,0 +1,150 @@
+"""Unit tests for approval utility functions."""
+
+from utils.approval_utils import (
+ calculate_estimated_cost,
+ format_approval_request,
+ format_critique_summary,
+ format_query_list,
+ parse_approval_response,
+ summarize_research_progress,
+)
+from utils.pydantic_models import ResearchState, SynthesizedInfo
+
+
+def test_parse_approval_responses():
+ """Test parsing different approval responses."""
+ queries = ["query1", "query2", "query3"]
+
+ # Test approve all
+ decision = parse_approval_response("APPROVE ALL", queries)
+ assert decision.approved == True
+ assert decision.selected_queries == queries
+ assert decision.approval_method == "APPROVE_ALL"
+
+ # Test skip
+ decision = parse_approval_response(
+ "skip", queries
+ ) # Test case insensitive
+ assert decision.approved == False
+ assert decision.selected_queries == []
+ assert decision.approval_method == "SKIP"
+
+ # Test selection
+ decision = parse_approval_response("SELECT 1,3", queries)
+ assert decision.approved == True
+ assert decision.selected_queries == ["query1", "query3"]
+ assert decision.approval_method == "SELECT_SPECIFIC"
+
+ # Test invalid selection
+ decision = parse_approval_response("SELECT invalid", queries)
+ assert decision.approved == False
+ assert decision.approval_method == "PARSE_ERROR"
+
+ # Test out of range indices
+ decision = parse_approval_response("SELECT 1,5,10", queries)
+ assert decision.approved == True
+ assert decision.selected_queries == ["query1"] # Only valid indices
+ assert decision.approval_method == "SELECT_SPECIFIC"
+
+ # Test unknown response
+ decision = parse_approval_response("maybe later", queries)
+ assert decision.approved == False
+ assert decision.approval_method == "UNKNOWN_RESPONSE"
+
+
+def test_format_approval_request():
+ """Test formatting of approval request messages."""
+ message = format_approval_request(
+ main_query="Test query",
+ progress_summary={
+ "completed_count": 5,
+ "avg_confidence": 0.75,
+ "low_confidence_count": 2,
+ },
+ critique_points=[
+ {"issue": "Missing data", "importance": "high"},
+ {"issue": "Minor gap", "importance": "low"},
+ ],
+ proposed_queries=["query1", "query2"],
+ )
+
+ assert "Test query" in message
+ assert "5" in message
+ assert "0.75" in message
+ assert "2 queries" in message
+ assert "approve" in message.lower()
+ assert "reject" in message.lower()
+ assert "Missing data" in message
+
+
+def test_summarize_research_progress():
+ """Test research progress summarization."""
+ state = ResearchState(
+ main_query="test",
+ synthesized_info={
+ "q1": SynthesizedInfo(
+ synthesized_answer="a1", confidence_level="high"
+ ),
+ "q2": SynthesizedInfo(
+ synthesized_answer="a2", confidence_level="medium"
+ ),
+ "q3": SynthesizedInfo(
+ synthesized_answer="a3", confidence_level="low"
+ ),
+ "q4": SynthesizedInfo(
+ synthesized_answer="a4", confidence_level="low"
+ ),
+ },
+ )
+
+ summary = summarize_research_progress(state)
+
+ assert summary["completed_count"] == 4
+ # (1.0 + 0.5 + 0.0 + 0.0) / 4 = 1.5 / 4 = 0.375, rounded to 0.38
+ assert summary["avg_confidence"] == 0.38
+ assert summary["low_confidence_count"] == 2
+
+
+def test_format_critique_summary():
+ """Test critique summary formatting."""
+ # Test with no critiques
+ result = format_critique_summary([])
+ assert result == "No critical issues identified."
+
+ # Test with few critiques
+ critiques = [{"issue": "Issue 1"}, {"issue": "Issue 2"}]
+ result = format_critique_summary(critiques)
+ assert "- Issue 1" in result
+ assert "- Issue 2" in result
+ assert "more issues" not in result
+
+ # Test with many critiques
+ critiques = [{"issue": f"Issue {i}"} for i in range(5)]
+ result = format_critique_summary(critiques)
+ assert "- Issue 0" in result
+ assert "- Issue 1" in result
+ assert "- Issue 2" in result
+ assert "- Issue 3" not in result
+ assert "... and 2 more issues" in result
+
+
+def test_format_query_list():
+ """Test query list formatting."""
+ # Test empty list
+ result = format_query_list([])
+ assert result == "No queries proposed."
+
+ # Test with queries
+ queries = ["Query A", "Query B", "Query C"]
+ result = format_query_list(queries)
+ assert "1. Query A" in result
+ assert "2. Query B" in result
+ assert "3. Query C" in result
+
+
+def test_calculate_estimated_cost():
+ """Test cost estimation."""
+ assert calculate_estimated_cost([]) == 0.0
+ assert calculate_estimated_cost(["q1"]) == 0.01
+ assert calculate_estimated_cost(["q1", "q2", "q3"]) == 0.03
+ assert calculate_estimated_cost(["q1"] * 10) == 0.10
diff --git a/deep_research/tests/test_prompt_loader.py b/deep_research/tests/test_prompt_loader.py
new file mode 100644
index 00000000..fdc7b7da
--- /dev/null
+++ b/deep_research/tests/test_prompt_loader.py
@@ -0,0 +1,108 @@
+"""Unit tests for prompt loader utilities."""
+
+import pytest
+from utils.prompt_loader import get_prompt_for_step, load_prompts_bundle
+from utils.prompt_models import PromptsBundle
+
+
+class TestPromptLoader:
+ """Test cases for prompt loader functions."""
+
+ def test_load_prompts_bundle(self):
+ """Test loading prompts bundle from prompts.py."""
+ bundle = load_prompts_bundle(pipeline_version="2.0.0")
+
+ # Check it returns a PromptsBundle
+ assert isinstance(bundle, PromptsBundle)
+
+ # Check pipeline version
+ assert bundle.pipeline_version == "2.0.0"
+
+ # Check all core prompts are loaded
+ assert bundle.search_query_prompt is not None
+ assert bundle.query_decomposition_prompt is not None
+ assert bundle.synthesis_prompt is not None
+ assert bundle.viewpoint_analysis_prompt is not None
+ assert bundle.reflection_prompt is not None
+ assert bundle.additional_synthesis_prompt is not None
+ assert bundle.conclusion_generation_prompt is not None
+
+ # Check prompts have correct metadata
+ assert bundle.search_query_prompt.name == "search_query_prompt"
+ assert (
+ bundle.search_query_prompt.description
+ == "Generates effective search queries from sub-questions"
+ )
+ assert bundle.search_query_prompt.version == "1.0.0"
+ assert "search" in bundle.search_query_prompt.tags
+
+ # Check that actual prompt content is loaded
+ assert "search query" in bundle.search_query_prompt.content.lower()
+ assert "json schema" in bundle.search_query_prompt.content.lower()
+
+ def test_load_prompts_bundle_default_version(self):
+ """Test loading prompts bundle with default version."""
+ bundle = load_prompts_bundle()
+ assert bundle.pipeline_version == "1.0.0"
+
+ def test_get_prompt_for_step(self):
+ """Test getting prompt content for specific steps."""
+ bundle = load_prompts_bundle()
+
+ # Test valid step names
+ test_cases = [
+ ("query_decomposition", "query_decomposition_prompt"),
+ ("search_query_generation", "search_query_prompt"),
+ ("synthesis", "synthesis_prompt"),
+ ("viewpoint_analysis", "viewpoint_analysis_prompt"),
+ ("reflection", "reflection_prompt"),
+ ("additional_synthesis", "additional_synthesis_prompt"),
+ ("conclusion_generation", "conclusion_generation_prompt"),
+ ]
+
+ for step_name, expected_prompt_attr in test_cases:
+ content = get_prompt_for_step(bundle, step_name)
+ expected_content = getattr(bundle, expected_prompt_attr).content
+ assert content == expected_content
+
+ def test_get_prompt_for_step_invalid(self):
+ """Test getting prompt for invalid step name."""
+ bundle = load_prompts_bundle()
+
+ with pytest.raises(
+ ValueError, match="No prompt mapping found for step: invalid_step"
+ ):
+ get_prompt_for_step(bundle, "invalid_step")
+
+ def test_all_prompts_have_content(self):
+ """Test that all loaded prompts have non-empty content."""
+ bundle = load_prompts_bundle()
+
+ all_prompts = bundle.list_all_prompts()
+ for name, prompt in all_prompts.items():
+ assert prompt.content, f"Prompt {name} has empty content"
+ assert (
+ len(prompt.content) > 50
+ ), f"Prompt {name} content seems too short"
+
+ def test_all_prompts_have_descriptions(self):
+ """Test that all loaded prompts have descriptions."""
+ bundle = load_prompts_bundle()
+
+ all_prompts = bundle.list_all_prompts()
+ for name, prompt in all_prompts.items():
+ assert prompt.description, f"Prompt {name} has no description"
+ assert (
+ len(prompt.description) > 10
+ ), f"Prompt {name} description seems too short"
+
+ def test_all_prompts_have_tags(self):
+ """Test that all loaded prompts have at least one tag."""
+ bundle = load_prompts_bundle()
+
+ all_prompts = bundle.list_all_prompts()
+ for name, prompt in all_prompts.items():
+ assert prompt.tags, f"Prompt {name} has no tags"
+ assert (
+ len(prompt.tags) >= 1
+ ), f"Prompt {name} should have at least one tag"
diff --git a/deep_research/tests/test_prompt_models.py b/deep_research/tests/test_prompt_models.py
new file mode 100644
index 00000000..f8b0519f
--- /dev/null
+++ b/deep_research/tests/test_prompt_models.py
@@ -0,0 +1,183 @@
+"""Unit tests for prompt models and utilities."""
+
+import pytest
+from utils.prompt_models import PromptsBundle, PromptTemplate
+
+
+class TestPromptTemplate:
+ """Test cases for PromptTemplate model."""
+
+ def test_prompt_template_creation(self):
+ """Test creating a prompt template with all fields."""
+ prompt = PromptTemplate(
+ name="test_prompt",
+ content="This is a test prompt",
+ description="A test prompt for unit testing",
+ version="1.0.0",
+ tags=["test", "unit"],
+ )
+
+ assert prompt.name == "test_prompt"
+ assert prompt.content == "This is a test prompt"
+ assert prompt.description == "A test prompt for unit testing"
+ assert prompt.version == "1.0.0"
+ assert prompt.tags == ["test", "unit"]
+
+ def test_prompt_template_minimal(self):
+ """Test creating a prompt template with minimal fields."""
+ prompt = PromptTemplate(
+ name="minimal_prompt", content="Minimal content"
+ )
+
+ assert prompt.name == "minimal_prompt"
+ assert prompt.content == "Minimal content"
+ assert prompt.description == ""
+ assert prompt.version == "1.0.0"
+ assert prompt.tags == []
+
+
+class TestPromptsBundle:
+ """Test cases for PromptsBundle model."""
+
+ @pytest.fixture
+ def sample_prompts(self):
+ """Create sample prompts for testing."""
+ return {
+ "search_query_prompt": PromptTemplate(
+ name="search_query_prompt",
+ content="Search query content",
+ description="Generates search queries",
+ tags=["search"],
+ ),
+ "query_decomposition_prompt": PromptTemplate(
+ name="query_decomposition_prompt",
+ content="Query decomposition content",
+ description="Decomposes queries",
+ tags=["analysis"],
+ ),
+ "synthesis_prompt": PromptTemplate(
+ name="synthesis_prompt",
+ content="Synthesis content",
+ description="Synthesizes information",
+ tags=["synthesis"],
+ ),
+ "viewpoint_analysis_prompt": PromptTemplate(
+ name="viewpoint_analysis_prompt",
+ content="Viewpoint analysis content",
+ description="Analyzes viewpoints",
+ tags=["analysis"],
+ ),
+ "reflection_prompt": PromptTemplate(
+ name="reflection_prompt",
+ content="Reflection content",
+ description="Reflects on research",
+ tags=["reflection"],
+ ),
+ "additional_synthesis_prompt": PromptTemplate(
+ name="additional_synthesis_prompt",
+ content="Additional synthesis content",
+ description="Additional synthesis",
+ tags=["synthesis"],
+ ),
+ "conclusion_generation_prompt": PromptTemplate(
+ name="conclusion_generation_prompt",
+ content="Conclusion generation content",
+ description="Generates conclusions",
+ tags=["report"],
+ ),
+ }
+
+ def test_prompts_bundle_creation(self, sample_prompts):
+ """Test creating a prompts bundle."""
+ bundle = PromptsBundle(**sample_prompts)
+
+ assert bundle.search_query_prompt.name == "search_query_prompt"
+ assert (
+ bundle.query_decomposition_prompt.name
+ == "query_decomposition_prompt"
+ )
+ assert bundle.pipeline_version == "1.0.0"
+ assert isinstance(bundle.created_at, str)
+ assert bundle.custom_prompts == {}
+
+ def test_prompts_bundle_with_custom_prompts(self, sample_prompts):
+ """Test creating a prompts bundle with custom prompts."""
+ custom_prompt = PromptTemplate(
+ name="custom_prompt",
+ content="Custom prompt content",
+ description="A custom prompt",
+ )
+
+ bundle = PromptsBundle(
+ **sample_prompts, custom_prompts={"custom_prompt": custom_prompt}
+ )
+
+ assert "custom_prompt" in bundle.custom_prompts
+ assert bundle.custom_prompts["custom_prompt"].name == "custom_prompt"
+
+ def test_get_prompt_by_name(self, sample_prompts):
+ """Test retrieving prompts by name."""
+ bundle = PromptsBundle(**sample_prompts)
+
+ # Test getting a core prompt
+ prompt = bundle.get_prompt_by_name("search_query_prompt")
+ assert prompt is not None
+ assert prompt.name == "search_query_prompt"
+
+ # Test getting a non-existent prompt
+ prompt = bundle.get_prompt_by_name("non_existent")
+ assert prompt is None
+
+ def test_get_prompt_by_name_custom(self, sample_prompts):
+ """Test retrieving custom prompts by name."""
+ custom_prompt = PromptTemplate(
+ name="custom_prompt", content="Custom content"
+ )
+
+ bundle = PromptsBundle(
+ **sample_prompts, custom_prompts={"custom_prompt": custom_prompt}
+ )
+
+ prompt = bundle.get_prompt_by_name("custom_prompt")
+ assert prompt is not None
+ assert prompt.name == "custom_prompt"
+
+ def test_list_all_prompts(self, sample_prompts):
+ """Test listing all prompts."""
+ bundle = PromptsBundle(**sample_prompts)
+
+ all_prompts = bundle.list_all_prompts()
+ assert len(all_prompts) == 7 # 7 core prompts
+ assert "search_query_prompt" in all_prompts
+ assert "conclusion_generation_prompt" in all_prompts
+
+ def test_list_all_prompts_with_custom(self, sample_prompts):
+ """Test listing all prompts including custom ones."""
+ custom_prompt = PromptTemplate(
+ name="custom_prompt", content="Custom content"
+ )
+
+ bundle = PromptsBundle(
+ **sample_prompts, custom_prompts={"custom_prompt": custom_prompt}
+ )
+
+ all_prompts = bundle.list_all_prompts()
+ assert len(all_prompts) == 8 # 7 core + 1 custom
+ assert "custom_prompt" in all_prompts
+
+ def test_get_prompt_content(self, sample_prompts):
+ """Test getting prompt content by type."""
+ bundle = PromptsBundle(**sample_prompts)
+
+ content = bundle.get_prompt_content("search_query_prompt")
+ assert content == "Search query content"
+
+ content = bundle.get_prompt_content("synthesis_prompt")
+ assert content == "Synthesis content"
+
+ def test_get_prompt_content_invalid(self, sample_prompts):
+ """Test getting prompt content with invalid type."""
+ bundle = PromptsBundle(**sample_prompts)
+
+ with pytest.raises(AttributeError):
+ bundle.get_prompt_content("invalid_prompt_type")
diff --git a/deep_research/tests/test_pydantic_final_report_step.py b/deep_research/tests/test_pydantic_final_report_step.py
new file mode 100644
index 00000000..c0f13530
--- /dev/null
+++ b/deep_research/tests/test_pydantic_final_report_step.py
@@ -0,0 +1,167 @@
+"""Tests for the Pydantic-based final report step.
+
+This module contains tests for the Pydantic-based implementation of
+final_report_step, which uses the new Pydantic models and materializers.
+"""
+
+from typing import Dict, List
+
+import pytest
+from steps.pydantic_final_report_step import pydantic_final_report_step
+from utils.pydantic_models import (
+ ReflectionMetadata,
+ ResearchState,
+ SearchResult,
+ SynthesizedInfo,
+ ViewpointAnalysis,
+ ViewpointTension,
+)
+from zenml.types import HTMLString
+
+
+@pytest.fixture
+def sample_research_state() -> ResearchState:
+ """Create a sample research state for testing."""
+ # Create a basic research state
+ state = ResearchState(main_query="What are the impacts of climate change?")
+
+ # Add sub-questions
+ state.update_sub_questions(["Economic impacts", "Environmental impacts"])
+
+ # Add search results
+ search_results: Dict[str, List[SearchResult]] = {
+ "Economic impacts": [
+ SearchResult(
+ url="https://example.com/economy",
+ title="Economic Impacts of Climate Change",
+ snippet="Overview of economic impacts",
+ content="Detailed content about economic impacts of climate change",
+ )
+ ]
+ }
+ state.update_search_results(search_results)
+
+ # Add synthesized info
+ synthesized_info: Dict[str, SynthesizedInfo] = {
+ "Economic impacts": SynthesizedInfo(
+ synthesized_answer="Climate change will have significant economic impacts...",
+ key_sources=["https://example.com/economy"],
+ confidence_level="high",
+ ),
+ "Environmental impacts": SynthesizedInfo(
+ synthesized_answer="Environmental impacts include rising sea levels...",
+ key_sources=["https://example.com/environment"],
+ confidence_level="high",
+ ),
+ }
+ state.update_synthesized_info(synthesized_info)
+
+ # Add enhanced info (same as synthesized for this test)
+ state.enhanced_info = state.synthesized_info
+
+ # Add viewpoint analysis
+ viewpoint_analysis = ViewpointAnalysis(
+ main_points_of_agreement=[
+ "Climate change is happening",
+ "Action is needed",
+ ],
+ areas_of_tension=[
+ ViewpointTension(
+ topic="Economic policy",
+ viewpoints={
+ "Progressive": "Support carbon taxes and regulations",
+ "Conservative": "Prefer market-based solutions",
+ },
+ )
+ ],
+ perspective_gaps="Indigenous perspectives are underrepresented",
+ integrative_insights="A balanced approach combining regulations and market incentives may be most effective",
+ )
+ state.update_viewpoint_analysis(viewpoint_analysis)
+
+ # Add reflection metadata
+ reflection_metadata = ReflectionMetadata(
+ critique_summary=["Need more sources for economic impacts"],
+ additional_questions_identified=[
+ "How will climate change affect different regions?"
+ ],
+ searches_performed=[
+ "economic impacts of climate change",
+ "regional climate impacts",
+ ],
+ improvements_made=2,
+ )
+ state.reflection_metadata = reflection_metadata
+
+ return state
+
+
+def test_pydantic_final_report_step_returns_tuple():
+ """Test that the step returns a tuple with state and HTML."""
+ # Create a simple state
+ state = ResearchState(main_query="What is climate change?")
+ state.update_sub_questions(["What causes climate change?"])
+
+ # Run the step
+ result = pydantic_final_report_step(state=state)
+
+ # Assert that result is a tuple with 2 elements
+ assert isinstance(result, tuple)
+ assert len(result) == 2
+
+ # Assert first element is ResearchState
+ assert isinstance(result[0], ResearchState)
+
+ # Assert second element is HTMLString
+ assert isinstance(result[1], HTMLString)
+
+
+def test_pydantic_final_report_step_with_complex_state(sample_research_state):
+ """Test that the step handles a complex state properly."""
+ # Run the step with a complex state
+ result = pydantic_final_report_step(state=sample_research_state)
+
+ # Unpack the results
+ updated_state, html_report = result
+
+ # Assert state contains final report HTML
+ assert updated_state.final_report_html != ""
+
+ # Assert HTML report contains key elements
+ html_str = str(html_report)
+ assert "Economic impacts" in html_str
+ assert "Environmental impacts" in html_str
+ assert "Viewpoint Analysis" in html_str
+ assert "Progressive" in html_str
+ assert "Conservative" in html_str
+
+
+def test_pydantic_final_report_step_updates_state():
+ """Test that the step properly updates the state."""
+ # Create an initial state without a final report
+ state = ResearchState(
+ main_query="What is climate change?",
+ sub_questions=["What causes climate change?"],
+ synthesized_info={
+ "What causes climate change?": SynthesizedInfo(
+ synthesized_answer="Climate change is caused by greenhouse gases.",
+ confidence_level="high",
+ )
+ },
+ enhanced_info={
+ "What causes climate change?": SynthesizedInfo(
+ synthesized_answer="Climate change is caused by greenhouse gases.",
+ confidence_level="high",
+ )
+ },
+ )
+
+ # Verify initial state has no report
+ assert state.final_report_html == ""
+
+ # Run the step
+ updated_state, _ = pydantic_final_report_step(state=state)
+
+ # Verify state was updated with a report
+ assert updated_state.final_report_html != ""
+ assert "climate change" in updated_state.final_report_html.lower()
diff --git a/deep_research/tests/test_pydantic_materializer.py b/deep_research/tests/test_pydantic_materializer.py
new file mode 100644
index 00000000..49cb17f8
--- /dev/null
+++ b/deep_research/tests/test_pydantic_materializer.py
@@ -0,0 +1,161 @@
+"""Tests for Pydantic-based materializer.
+
+This module contains tests for the Pydantic-based implementation of
+ResearchStateMaterializer, verifying that it correctly serializes and
+visualizes ResearchState objects.
+"""
+
+import os
+import tempfile
+from typing import Dict, List
+
+import pytest
+from materializers.pydantic_materializer import ResearchStateMaterializer
+from utils.pydantic_models import (
+ ResearchState,
+ SearchResult,
+ SynthesizedInfo,
+ ViewpointAnalysis,
+ ViewpointTension,
+)
+
+
+@pytest.fixture
+def sample_state() -> ResearchState:
+ """Create a sample research state for testing."""
+ # Create a basic research state
+ state = ResearchState(main_query="What are the impacts of climate change?")
+
+ # Add sub-questions
+ state.update_sub_questions(["Economic impacts", "Environmental impacts"])
+
+ # Add search results
+ search_results: Dict[str, List[SearchResult]] = {
+ "Economic impacts": [
+ SearchResult(
+ url="https://example.com/economy",
+ title="Economic Impacts of Climate Change",
+ snippet="Overview of economic impacts",
+ content="Detailed content about economic impacts of climate change",
+ )
+ ]
+ }
+ state.update_search_results(search_results)
+
+ # Add synthesized info
+ synthesized_info: Dict[str, SynthesizedInfo] = {
+ "Economic impacts": SynthesizedInfo(
+ synthesized_answer="Climate change will have significant economic impacts...",
+ key_sources=["https://example.com/economy"],
+ confidence_level="high",
+ )
+ }
+ state.update_synthesized_info(synthesized_info)
+
+ return state
+
+
+def test_materializer_initialization():
+ """Test that the materializer can be initialized."""
+ # Create a temporary directory for artifact storage
+ with tempfile.TemporaryDirectory() as tmpdirname:
+ materializer = ResearchStateMaterializer(uri=tmpdirname)
+ assert materializer is not None
+
+
+def test_materializer_save_and_load(sample_state: ResearchState):
+ """Test saving and loading a state using the materializer."""
+ # Create a temporary directory for artifact storage
+ with tempfile.TemporaryDirectory() as tmpdirname:
+ # Initialize materializer with temporary artifact URI
+ materializer = ResearchStateMaterializer(uri=tmpdirname)
+
+ # Save the state
+ materializer.save(sample_state)
+
+ # Load the state
+ loaded_state = materializer.load(ResearchState)
+
+ # Verify that the loaded state matches the original
+ assert loaded_state.main_query == sample_state.main_query
+ assert loaded_state.sub_questions == sample_state.sub_questions
+ assert len(loaded_state.search_results) == len(
+ sample_state.search_results
+ )
+ assert (
+ loaded_state.get_current_stage()
+ == sample_state.get_current_stage()
+ )
+
+ # Check that key fields were preserved
+ question = "Economic impacts"
+ assert (
+ loaded_state.synthesized_info[question].synthesized_answer
+ == sample_state.synthesized_info[question].synthesized_answer
+ )
+ assert (
+ loaded_state.synthesized_info[question].confidence_level
+ == sample_state.synthesized_info[question].confidence_level
+ )
+
+
+def test_materializer_save_visualizations(sample_state: ResearchState):
+ """Test generating and saving visualizations."""
+ # Create a temporary directory for artifact storage
+ with tempfile.TemporaryDirectory() as tmpdirname:
+ # Initialize materializer with temporary artifact URI
+ materializer = ResearchStateMaterializer(uri=tmpdirname)
+
+ # Generate and save visualizations
+ viz_paths = materializer.save_visualizations(sample_state)
+
+ # Verify visualization file exists
+ html_path = list(viz_paths.keys())[0]
+ assert os.path.exists(html_path)
+
+ # Verify the file has content
+ with open(html_path, "r") as f:
+ content = f.read()
+ # Check for expected elements in the HTML
+ assert "Research State" in content
+ assert sample_state.main_query in content
+ assert "Economic impacts" in content
+
+
+def test_html_generation_stages(sample_state: ResearchState):
+ """Test that HTML visualization reflects the correct research stage."""
+ # Create the materializer
+ with tempfile.TemporaryDirectory() as tmpdirname:
+ materializer = ResearchStateMaterializer(uri=tmpdirname)
+
+ # Generate visualization at initial state
+ html = materializer._generate_visualization_html(sample_state)
+ # Verify stage by checking for expected elements in the HTML
+ assert (
+ "Synthesized Information" in html
+ ) # Should show synthesized info
+
+ # Add viewpoint analysis
+ state_with_viewpoints = sample_state.model_copy(deep=True)
+ viewpoint_analysis = ViewpointAnalysis(
+ main_points_of_agreement=["There will be economic impacts"],
+ areas_of_tension=[
+ ViewpointTension(
+ topic="Job impacts",
+ viewpoints={
+ "Positive": "New green jobs",
+ "Negative": "Job losses",
+ },
+ )
+ ],
+ )
+ state_with_viewpoints.update_viewpoint_analysis(viewpoint_analysis)
+ html = materializer._generate_visualization_html(state_with_viewpoints)
+ assert "Viewpoint Analysis" in html
+ assert "Points of Agreement" in html
+
+ # Add final report
+ state_with_report = state_with_viewpoints.model_copy(deep=True)
+ state_with_report.set_final_report("Final report content")
+ html = materializer._generate_visualization_html(state_with_report)
+ assert "Final Report" in html
diff --git a/deep_research/tests/test_pydantic_models.py b/deep_research/tests/test_pydantic_models.py
new file mode 100644
index 00000000..b900f8d8
--- /dev/null
+++ b/deep_research/tests/test_pydantic_models.py
@@ -0,0 +1,303 @@
+"""Tests for Pydantic model implementations.
+
+This module contains tests for the Pydantic models that validate:
+1. Basic model instantiation
+2. Default values
+3. Serialization and deserialization
+4. Method functionality
+"""
+
+import json
+from typing import Dict, List
+
+from utils.pydantic_models import (
+ ReflectionMetadata,
+ ResearchState,
+ SearchResult,
+ SynthesizedInfo,
+ ViewpointAnalysis,
+ ViewpointTension,
+)
+
+
+def test_search_result_creation():
+ """Test creating a SearchResult model."""
+ # Create with defaults
+ result = SearchResult()
+ assert result.url == ""
+ assert result.content == ""
+ assert result.title == ""
+ assert result.snippet == ""
+
+ # Create with values
+ result = SearchResult(
+ url="https://example.com",
+ content="Example content",
+ title="Example Title",
+ snippet="This is a snippet",
+ )
+ assert result.url == "https://example.com"
+ assert result.content == "Example content"
+ assert result.title == "Example Title"
+ assert result.snippet == "This is a snippet"
+
+
+def test_search_result_serialization():
+ """Test serializing and deserializing a SearchResult."""
+ result = SearchResult(
+ url="https://example.com",
+ content="Example content",
+ title="Example Title",
+ snippet="This is a snippet",
+ )
+
+ # Serialize to dict
+ result_dict = result.model_dump()
+ assert result_dict["url"] == "https://example.com"
+ assert result_dict["content"] == "Example content"
+
+ # Serialize to JSON
+ result_json = result.model_dump_json()
+ result_dict_from_json = json.loads(result_json)
+ assert result_dict_from_json["url"] == "https://example.com"
+
+ # Deserialize from dict
+ new_result = SearchResult.model_validate(result_dict)
+ assert new_result.url == "https://example.com"
+ assert new_result.content == "Example content"
+
+ # Deserialize from JSON
+ new_result_from_json = SearchResult.model_validate_json(result_json)
+ assert new_result_from_json.url == "https://example.com"
+
+
+def test_viewpoint_tension_model():
+ """Test the ViewpointTension model."""
+ # Empty model
+ tension = ViewpointTension()
+ assert tension.topic == ""
+ assert tension.viewpoints == {}
+
+ # With data
+ tension = ViewpointTension(
+ topic="Climate Change Impacts",
+ viewpoints={
+ "Economic": "Focuses on financial costs and benefits",
+ "Environmental": "Emphasizes ecosystem impacts",
+ },
+ )
+ assert tension.topic == "Climate Change Impacts"
+ assert len(tension.viewpoints) == 2
+ assert "Economic" in tension.viewpoints
+
+ # Serialization
+ tension_dict = tension.model_dump()
+ assert tension_dict["topic"] == "Climate Change Impacts"
+ assert len(tension_dict["viewpoints"]) == 2
+
+ # Deserialization
+ new_tension = ViewpointTension.model_validate(tension_dict)
+ assert new_tension.topic == tension.topic
+ assert new_tension.viewpoints == tension.viewpoints
+
+
+def test_synthesized_info_model():
+ """Test the SynthesizedInfo model."""
+ # Default values
+ info = SynthesizedInfo()
+ assert info.synthesized_answer == ""
+ assert info.key_sources == []
+ assert info.confidence_level == "medium"
+ assert info.information_gaps == ""
+ assert info.improvements == []
+
+ # With values
+ info = SynthesizedInfo(
+ synthesized_answer="This is a synthesized answer",
+ key_sources=["https://source1.com", "https://source2.com"],
+ confidence_level="high",
+ information_gaps="Missing some context",
+ improvements=["Add more detail", "Check more sources"],
+ )
+ assert info.synthesized_answer == "This is a synthesized answer"
+ assert len(info.key_sources) == 2
+ assert info.confidence_level == "high"
+
+ # Serialization and deserialization
+ info_dict = info.model_dump()
+ new_info = SynthesizedInfo.model_validate(info_dict)
+ assert new_info.synthesized_answer == info.synthesized_answer
+ assert new_info.key_sources == info.key_sources
+
+
+def test_viewpoint_analysis_model():
+ """Test the ViewpointAnalysis model."""
+ # Create tensions for the analysis
+ tension1 = ViewpointTension(
+ topic="Economic Impact",
+ viewpoints={
+ "Positive": "Creates jobs",
+ "Negative": "Increases inequality",
+ },
+ )
+ tension2 = ViewpointTension(
+ topic="Environmental Impact",
+ viewpoints={
+ "Positive": "Reduces emissions",
+ "Negative": "Land use changes",
+ },
+ )
+
+ # Create the analysis
+ analysis = ViewpointAnalysis(
+ main_points_of_agreement=[
+ "Need for action",
+ "Technological innovation",
+ ],
+ areas_of_tension=[tension1, tension2],
+ perspective_gaps="Missing indigenous perspectives",
+ integrative_insights="Combined economic and environmental approach needed",
+ )
+
+ assert len(analysis.main_points_of_agreement) == 2
+ assert len(analysis.areas_of_tension) == 2
+ assert analysis.areas_of_tension[0].topic == "Economic Impact"
+
+ # Test serialization
+ analysis_dict = analysis.model_dump()
+ assert len(analysis_dict["areas_of_tension"]) == 2
+ assert analysis_dict["areas_of_tension"][0]["topic"] == "Economic Impact"
+
+ # Test deserialization
+ new_analysis = ViewpointAnalysis.model_validate(analysis_dict)
+ assert len(new_analysis.areas_of_tension) == 2
+ assert new_analysis.areas_of_tension[0].topic == "Economic Impact"
+ assert new_analysis.perspective_gaps == analysis.perspective_gaps
+
+
+def test_reflection_metadata_model():
+ """Test the ReflectionMetadata model."""
+ metadata = ReflectionMetadata(
+ critique_summary=["Need more sources", "Missing detailed analysis"],
+ additional_questions_identified=["What about future trends?"],
+ searches_performed=["future climate trends", "economic impacts"],
+ improvements_made=3,
+ error=None,
+ )
+
+ assert len(metadata.critique_summary) == 2
+ assert len(metadata.additional_questions_identified) == 1
+ assert metadata.improvements_made == 3
+ assert metadata.error is None
+
+ # Serialization
+ metadata_dict = metadata.model_dump()
+ assert len(metadata_dict["critique_summary"]) == 2
+ assert metadata_dict["improvements_made"] == 3
+
+ # Deserialization
+ new_metadata = ReflectionMetadata.model_validate(metadata_dict)
+ assert new_metadata.improvements_made == metadata.improvements_made
+ assert new_metadata.critique_summary == metadata.critique_summary
+
+
+def test_research_state_model():
+ """Test the main ResearchState model."""
+ # Create with defaults
+ state = ResearchState()
+ assert state.main_query == ""
+ assert state.sub_questions == []
+ assert state.search_results == {}
+ assert state.get_current_stage() == "empty"
+
+ # Set main query
+ state.main_query = "What are the impacts of climate change?"
+ assert state.get_current_stage() == "initial"
+
+ # Test update methods
+ state.update_sub_questions(
+ ["What are economic impacts?", "What are environmental impacts?"]
+ )
+ assert len(state.sub_questions) == 2
+ assert state.get_current_stage() == "after_query_decomposition"
+
+ # Add search results
+ search_results: Dict[str, List[SearchResult]] = {
+ "What are economic impacts?": [
+ SearchResult(
+ url="https://example.com/economy",
+ title="Economic Impacts",
+ snippet="Overview of economic impacts",
+ content="Detailed content about economic impacts",
+ )
+ ]
+ }
+ state.update_search_results(search_results)
+ assert state.get_current_stage() == "after_search"
+ assert len(state.search_results["What are economic impacts?"]) == 1
+
+ # Add synthesized info
+ synthesized_info: Dict[str, SynthesizedInfo] = {
+ "What are economic impacts?": SynthesizedInfo(
+ synthesized_answer="Economic impacts include job losses and growth opportunities",
+ key_sources=["https://example.com/economy"],
+ confidence_level="high",
+ )
+ }
+ state.update_synthesized_info(synthesized_info)
+ assert state.get_current_stage() == "after_synthesis"
+
+ # Add viewpoint analysis
+ analysis = ViewpointAnalysis(
+ main_points_of_agreement=["Economic changes are happening"],
+ areas_of_tension=[
+ ViewpointTension(
+ topic="Job impacts",
+ viewpoints={
+ "Positive": "New green jobs",
+ "Negative": "Fossil fuel job losses",
+ },
+ )
+ ],
+ )
+ state.update_viewpoint_analysis(analysis)
+ assert state.get_current_stage() == "after_viewpoint_analysis"
+
+ # Add reflection results
+ enhanced_info = {
+ "What are economic impacts?": SynthesizedInfo(
+ synthesized_answer="Enhanced answer with more details",
+ key_sources=[
+ "https://example.com/economy",
+ "https://example.com/new-source",
+ ],
+ confidence_level="high",
+ improvements=["Added more context", "Added more sources"],
+ )
+ }
+ metadata = ReflectionMetadata(
+ critique_summary=["Needed more sources"],
+ improvements_made=2,
+ )
+ state.update_after_reflection(enhanced_info, metadata)
+ assert state.get_current_stage() == "after_reflection"
+
+ # Set final report
+ state.set_final_report("Final report content")
+ assert state.get_current_stage() == "final_report"
+ assert state.final_report_html == "Final report content"
+
+ # Test serialization and deserialization
+ state_dict = state.model_dump()
+ new_state = ResearchState.model_validate(state_dict)
+
+ # Verify key properties were preserved
+ assert new_state.main_query == state.main_query
+ assert len(new_state.sub_questions) == len(state.sub_questions)
+ assert new_state.get_current_stage() == state.get_current_stage()
+ assert new_state.viewpoint_analysis is not None
+ assert len(new_state.viewpoint_analysis.areas_of_tension) == 1
+ assert (
+ new_state.viewpoint_analysis.areas_of_tension[0].topic == "Job impacts"
+ )
+ assert new_state.final_report_html == state.final_report_html
diff --git a/deep_research/utils/__init__.py b/deep_research/utils/__init__.py
new file mode 100644
index 00000000..395e1d67
--- /dev/null
+++ b/deep_research/utils/__init__.py
@@ -0,0 +1,7 @@
+"""
+Utilities package for the ZenML Deep Research project.
+
+This package contains various utility functions and helpers used throughout the project,
+including data models, LLM interaction utilities, search functionality, and common helper
+functions for text processing and state management.
+"""
diff --git a/deep_research/utils/approval_utils.py b/deep_research/utils/approval_utils.py
new file mode 100644
index 00000000..56a8ff91
--- /dev/null
+++ b/deep_research/utils/approval_utils.py
@@ -0,0 +1,159 @@
+"""Utility functions for the human approval process."""
+
+from typing import Any, Dict, List
+
+from utils.pydantic_models import ApprovalDecision
+
+
+def summarize_research_progress(state) -> Dict[str, Any]:
+ """Summarize the current research progress."""
+ completed_count = len(state.synthesized_info)
+ confidence_levels = [
+ info.confidence_level for info in state.synthesized_info.values()
+ ]
+
+ # Calculate average confidence (high=1.0, medium=0.5, low=0.0)
+ confidence_map = {"high": 1.0, "medium": 0.5, "low": 0.0}
+ avg_confidence = sum(
+ confidence_map.get(c, 0.5) for c in confidence_levels
+ ) / max(len(confidence_levels), 1)
+
+ low_confidence_count = sum(1 for c in confidence_levels if c == "low")
+
+ return {
+ "completed_count": completed_count,
+ "avg_confidence": round(avg_confidence, 2),
+ "low_confidence_count": low_confidence_count,
+ }
+
+
+def format_critique_summary(critique_points: List[Dict[str, Any]]) -> str:
+ """Format critique points for display."""
+ if not critique_points:
+ return "No critical issues identified."
+
+ formatted = []
+ for point in critique_points[:3]: # Show top 3
+ issue = point.get("issue", "Unknown issue")
+ formatted.append(f"- {issue}")
+
+ if len(critique_points) > 3:
+ formatted.append(f"- ... and {len(critique_points) - 3} more issues")
+
+ return "\n".join(formatted)
+
+
+def format_query_list(queries: List[str]) -> str:
+ """Format query list for display."""
+ if not queries:
+ return "No queries proposed."
+
+ formatted = []
+ for i, query in enumerate(queries, 1):
+ formatted.append(f"{i}. {query}")
+
+ return "\n".join(formatted)
+
+
+def calculate_estimated_cost(queries: List[str]) -> float:
+ """Calculate estimated cost for additional queries."""
+ # Rough estimate: ~$0.01 per query (including search API + LLM costs)
+ return round(len(queries) * 0.01, 2)
+
+
+def format_approval_request(
+ main_query: str,
+ progress_summary: Dict[str, Any],
+ critique_points: List[Dict[str, Any]],
+ proposed_queries: List[str],
+ timeout: int = 3600,
+) -> str:
+ """Format the approval request message."""
+
+ # High-priority critiques
+ high_priority = [
+ c for c in critique_points if c.get("importance") == "high"
+ ]
+
+ message = f"""📊 **Research Progress Update**
+
+**Main Query:** {main_query}
+
+**Current Status:**
+- Sub-questions analyzed: {progress_summary["completed_count"]}
+- Average confidence: {progress_summary["avg_confidence"]}
+- Low confidence areas: {progress_summary["low_confidence_count"]}
+
+**Key Issues Identified:**
+{format_critique_summary(high_priority or critique_points)}
+
+**Proposed Additional Research** ({len(proposed_queries)} queries):
+{format_query_list(proposed_queries)}
+
+**Estimated Additional Time:** ~{len(proposed_queries) * 2} minutes
+**Estimated Additional Cost:** ~${calculate_estimated_cost(proposed_queries)}
+
+**Response Options:**
+- Reply with `approve`, `yes`, `ok`, or `LGTM` to proceed with all queries
+- Reply with `reject`, `no`, `skip`, or `decline` to finish with current findings
+
+**Timeout:** Response required within {timeout // 60} minutes"""
+
+ return message
+
+
+def parse_approval_response(
+ response: str, proposed_queries: List[str]
+) -> ApprovalDecision:
+ """Parse the approval response from user."""
+
+ response_upper = response.strip().upper()
+
+ if response_upper == "APPROVE ALL":
+ return ApprovalDecision(
+ approved=True,
+ selected_queries=proposed_queries,
+ approval_method="APPROVE_ALL",
+ reviewer_notes=response,
+ )
+
+ elif response_upper == "SKIP":
+ return ApprovalDecision(
+ approved=False,
+ selected_queries=[],
+ approval_method="SKIP",
+ reviewer_notes=response,
+ )
+
+ elif response_upper.startswith("SELECT"):
+ # Parse selection like "SELECT 1,3,5"
+ try:
+ # Extract the part after "SELECT"
+ selection_part = response_upper[6:].strip()
+ indices = [int(x.strip()) - 1 for x in selection_part.split(",")]
+ selected = [
+ proposed_queries[i]
+ for i in indices
+ if 0 <= i < len(proposed_queries)
+ ]
+ return ApprovalDecision(
+ approved=True,
+ selected_queries=selected,
+ approval_method="SELECT_SPECIFIC",
+ reviewer_notes=response,
+ )
+ except Exception as e:
+ return ApprovalDecision(
+ approved=False,
+ selected_queries=[],
+ approval_method="PARSE_ERROR",
+ reviewer_notes=f"Failed to parse: {response} - {str(e)}",
+ )
+
+ else:
+ return ApprovalDecision(
+ approved=False,
+ selected_queries=[],
+ approval_method="UNKNOWN_RESPONSE",
+ reviewer_notes=f"Unknown response: {response}",
+ )
diff --git a/deep_research/utils/helper_functions.py b/deep_research/utils/helper_functions.py
new file mode 100644
index 00000000..256cdd28
--- /dev/null
+++ b/deep_research/utils/helper_functions.py
@@ -0,0 +1,192 @@
+import json
+import logging
+import os
+from json.decoder import JSONDecodeError
+from typing import Any, Dict, Optional
+
+import yaml
+
+logger = logging.getLogger(__name__)
+
+
+def remove_reasoning_from_output(output: str) -> str:
+ """Remove the reasoning portion from LLM output.
+
+ Args:
+ output: Raw output from LLM that may contain reasoning
+
+ Returns:
+ Cleaned output without the reasoning section
+ """
+ if not output:
+ return ""
+
+ if "" in output:
+ return output.split("")[-1].strip()
+ return output.strip()
+
+
+def clean_json_tags(text: str) -> str:
+ """Clean JSON markdown tags from text.
+
+ Args:
+ text: Text with potential JSON markdown tags
+
+ Returns:
+ Cleaned text without JSON markdown tags
+ """
+ if not text:
+ return ""
+
+ cleaned = text.replace("```json\n", "").replace("\n```", "")
+ cleaned = cleaned.replace("```json", "").replace("```", "")
+ return cleaned
+
+
+def clean_markdown_tags(text: str) -> str:
+ """Clean Markdown tags from text.
+
+ Args:
+ text: Text with potential markdown tags
+
+ Returns:
+ Cleaned text without markdown tags
+ """
+ if not text:
+ return ""
+
+ cleaned = text.replace("```markdown\n", "").replace("\n```", "")
+ cleaned = cleaned.replace("```markdown", "").replace("```", "")
+ return cleaned
+
+
+def extract_html_from_content(content: str) -> str:
+ """Attempt to extract HTML content from a response that might be wrapped in other formats.
+
+ Args:
+ content: The content to extract HTML from
+
+ Returns:
+ The extracted HTML, or a basic fallback if extraction fails
+ """
+ if not content:
+ return ""
+
+ # Try to find HTML between tags
+ if "" in content:
+ start = content.find("") + 7 # Include the closing tag
+ return content[start:end]
+
+ # Try to find div class="research-report"
+ if '" in content:
+ start = content.find('")
+ if last_div > start:
+ return content[start : last_div + 6] # Include the closing tag
+
+ # Look for code blocks
+ if "```html" in content and "```" in content:
+ start = content.find("```html") + 7
+ end = content.find("```", start)
+ if end > start:
+ return content[start:end].strip()
+
+ # Look for JSON with an "html" field
+ try:
+ parsed = json.loads(content)
+ if isinstance(parsed, dict) and "html" in parsed:
+ return parsed["html"]
+ except:
+ pass
+
+ # If all extraction attempts fail, return the original content
+ return content
+
+
+def safe_json_loads(json_str: Optional[str]) -> Dict[str, Any]:
+ """Safely parse JSON string.
+
+ Args:
+ json_str: JSON string to parse, can be None.
+
+ Returns:
+ Dict[str, Any]: Parsed JSON as dictionary or empty dict if parsing fails or input is None.
+ """
+ if json_str is None:
+ # Optionally, log a warning here if None input is unexpected for certain call sites
+ # logger.warning("safe_json_loads received None input.")
+ return {}
+ try:
+ return json.loads(json_str)
+ except (
+ JSONDecodeError,
+ TypeError,
+ ): # Catch TypeError if json_str is not a valid type for json.loads
+ # Optionally, log the error and the problematic string (or its beginning)
+ # logger.warning(f"Failed to decode JSON string: '{str(json_str)[:200]}...'", exc_info=True)
+ return {}
+
+
+def load_pipeline_config(config_path: str) -> Dict[str, Any]:
+ """Load pipeline configuration from YAML file.
+
+ This is used only for pipeline-level configuration, not for step parameters.
+ Step parameters should be defined directly in the step functions.
+
+ Args:
+ config_path: Path to the configuration YAML file
+
+ Returns:
+ Pipeline configuration dictionary
+ """
+ # Get absolute path if relative
+ if not os.path.isabs(config_path):
+ base_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+ config_path = os.path.join(base_dir, config_path)
+
+ # Load YAML configuration
+ try:
+ with open(config_path, "r") as f:
+ config = yaml.safe_load(f)
+ return config
+ except Exception as e:
+ logger.error(f"Error loading pipeline configuration: {e}")
+ # Return a minimal default configuration in case of loading error
+ return {
+ "pipeline": {
+ "name": "deep_research_pipeline",
+ "enable_cache": True,
+ },
+ "environment": {
+ "docker": {
+ "requirements": [
+ "openai>=1.0.0",
+ "tavily-python>=0.2.8",
+ "PyYAML>=6.0",
+ "click>=8.0.0",
+ "pydantic>=2.0.0",
+ "typing_extensions>=4.0.0",
+ ]
+ }
+ },
+ "resources": {"cpu": 1, "memory": "4Gi"},
+ "timeout": 3600,
+ }
+
+
+def check_required_env_vars(env_vars: list[str]) -> list[str]:
+ """Check if required environment variables are set.
+
+ Args:
+ env_vars: List of environment variable names to check
+
+ Returns:
+ List of missing environment variables
+ """
+ missing_vars = []
+ for var in env_vars:
+ if not os.environ.get(var):
+ missing_vars.append(var)
+ return missing_vars
diff --git a/deep_research/utils/llm_utils.py b/deep_research/utils/llm_utils.py
new file mode 100644
index 00000000..1f4dc194
--- /dev/null
+++ b/deep_research/utils/llm_utils.py
@@ -0,0 +1,387 @@
+import contextlib
+import json
+import logging
+from typing import Any, Dict, List, Optional
+
+import litellm
+from litellm import completion
+from utils.helper_functions import (
+ clean_json_tags,
+ remove_reasoning_from_output,
+ safe_json_loads,
+)
+from utils.prompts import SYNTHESIS_PROMPT
+from zenml import get_step_context
+
+logger = logging.getLogger(__name__)
+
+# This module uses litellm for all LLM interactions
+# Models are specified with a provider prefix (e.g., "sambanova/DeepSeek-R1-Distill-Llama-70B")
+# ALL model names require a provider prefix (e.g., "sambanova/", "openai/", "anthropic/")
+
+litellm.callbacks = ["langfuse"]
+
+
+def run_llm_completion(
+ prompt: str,
+ system_prompt: str,
+ model: str = "sambanova/Llama-4-Maverick-17B-128E-Instruct",
+ clean_output: bool = True,
+ max_tokens: int = 2000, # Increased default token limit
+ temperature: float = 0.2,
+ top_p: float = 0.9,
+ project: str = "deep-research",
+ tags: Optional[List[str]] = None,
+) -> str:
+ """Run an LLM completion with standard error handling and output cleaning.
+
+ Uses litellm for model inference.
+
+ Args:
+ prompt: User prompt for the LLM
+ system_prompt: System prompt for the LLM
+ model: Model to use for completion (with provider prefix)
+ clean_output: Whether to clean reasoning and JSON tags from output. When True,
+ this removes any reasoning sections marked with tags and strips JSON
+ code block markers.
+ max_tokens: Maximum tokens to generate
+ temperature: Sampling temperature
+ top_p: Top-p sampling value
+ project: Langfuse project name for LLM tracking
+ tags: Optional list of tags for Langfuse tracking. If provided, also converted to trace_metadata format.
+
+ Returns:
+ str: Processed LLM output with optional cleaning applied
+ """
+ try:
+ # Ensure model name has provider prefix
+ if not any(
+ model.startswith(prefix + "/")
+ for prefix in [
+ "sambanova",
+ "openai",
+ "anthropic",
+ "meta",
+ "google",
+ "aws",
+ "openrouter",
+ ]
+ ):
+ # Raise an error if no provider prefix is specified
+ error_msg = f"Model '{model}' does not have a provider prefix. Please specify provider (e.g., 'sambanova/{model}')"
+ logger.error(error_msg)
+ raise ValueError(error_msg)
+
+ # Get pipeline run name and id for trace_name and trace_id if running in a step
+ trace_name = None
+ trace_id = None
+ with contextlib.suppress(RuntimeError):
+ context = get_step_context()
+ trace_name = context.pipeline_run.name
+ trace_id = str(context.pipeline_run.id)
+ # Build metadata dict
+ metadata = {"project": project}
+ if tags is not None:
+ metadata["tags"] = tags
+ # Convert tags to trace_metadata format
+ metadata["trace_metadata"] = {tag: True for tag in tags}
+ if trace_name:
+ metadata["trace_name"] = trace_name
+ if trace_id:
+ metadata["trace_id"] = trace_id
+
+ response = completion(
+ model=model,
+ messages=[
+ {"role": "system", "content": system_prompt},
+ {"role": "user", "content": prompt},
+ ],
+ max_tokens=max_tokens,
+ temperature=temperature,
+ top_p=top_p,
+ metadata=metadata,
+ )
+
+ # Defensive access to content
+ content = None
+ if response and response.choices and len(response.choices) > 0:
+ choice = response.choices[0]
+ if choice and choice.message:
+ content = choice.message.content
+
+ if content is None:
+ logger.warning("LLM response content is missing or empty.")
+ return ""
+
+ if clean_output:
+ content = remove_reasoning_from_output(content)
+ content = clean_json_tags(content)
+
+ return content
+ except Exception as e:
+ logger.error(f"Error in LLM completion: {e}")
+ return ""
+
+
+def get_structured_llm_output(
+ prompt: str,
+ system_prompt: str,
+ model: str = "sambanova/Llama-4-Maverick-17B-128E-Instruct",
+ fallback_response: Optional[Dict[str, Any]] = None,
+ max_tokens: int = 2000, # Increased default token limit for structured outputs
+ temperature: float = 0.2,
+ top_p: float = 0.9,
+ project: str = "deep-research",
+ tags: Optional[List[str]] = None,
+) -> Dict[str, Any]:
+ """Get structured JSON output from an LLM with error handling.
+
+ Uses litellm for model inference.
+
+ Args:
+ prompt: User prompt for the LLM
+ system_prompt: System prompt for the LLM
+ model: Model to use for completion (with provider prefix)
+ fallback_response: Fallback response if parsing fails
+ max_tokens: Maximum tokens to generate
+ temperature: Sampling temperature
+ top_p: Top-p sampling value
+ project: Langfuse project name for LLM tracking
+ tags: Optional list of tags for Langfuse tracking. Defaults to ["structured_llm_output"] if None.
+
+ Returns:
+ Parsed JSON response or fallback
+ """
+ try:
+ # Use provided tags or default to ["structured_llm_output"]
+ if tags is None:
+ tags = ["structured_llm_output"]
+
+ content = run_llm_completion(
+ prompt=prompt,
+ system_prompt=system_prompt,
+ model=model,
+ clean_output=True,
+ max_tokens=max_tokens,
+ temperature=temperature,
+ top_p=top_p,
+ project=project,
+ tags=tags,
+ )
+
+ if not content:
+ logger.warning("Empty content returned from LLM")
+ return fallback_response if fallback_response is not None else {}
+
+ result = safe_json_loads(content)
+
+ if not result and fallback_response is not None:
+ return fallback_response
+
+ return result
+ except Exception as e:
+ logger.error(f"Error processing structured LLM output: {e}")
+ return fallback_response if fallback_response is not None else {}
+
+
+def is_text_relevant(text1: str, text2: str, min_word_length: int = 4) -> bool:
+ """Determine if two pieces of text are relevant to each other.
+
+ Relevance is determined by checking if one text is contained within the other,
+ or if they share significant words (words longer than min_word_length).
+ This is a simple heuristic approach that checks for:
+ 1. Complete containment (one text string inside the other)
+ 2. Shared significant words (words longer than min_word_length)
+
+ Args:
+ text1: First text to compare
+ text2: Second text to compare
+ min_word_length: Minimum length of words to check for shared content
+
+ Returns:
+ bool: True if the texts are deemed relevant to each other based on the criteria
+ """
+ if not text1 or not text2:
+ return False
+
+ return (
+ text1.lower() in text2.lower()
+ or text2.lower() in text1.lower()
+ or any(
+ word
+ for word in text1.lower().split()
+ if len(word) > min_word_length and word in text2.lower()
+ )
+ )
+
+
+def find_most_relevant_string(
+ target: str,
+ options: List[str],
+ model: Optional[str] = "sambanova/Llama-4-Maverick-17B-128E-Instruct",
+ project: str = "deep-research",
+ tags: Optional[List[str]] = None,
+) -> Optional[str]:
+ """Find the most relevant string from a list of options using simple text matching.
+
+ If model is provided, uses litellm to determine relevance.
+
+ Args:
+ target: The target string to find relevance for
+ options: List of string options to check against
+ model: Model to use for matching (with provider prefix)
+ project: Langfuse project name for LLM tracking
+ tags: Optional list of tags for Langfuse tracking. Defaults to ["find_most_relevant_string"] if None.
+
+ Returns:
+ The most relevant string, or None if no relevant options
+ """
+ if not options:
+ return None
+
+ if len(options) == 1:
+ return options[0]
+
+ # If model is provided, use litellm for more accurate matching
+ if model:
+ try:
+ # Ensure model name has provider prefix
+ if not any(
+ model.startswith(prefix + "/")
+ for prefix in [
+ "sambanova",
+ "openai",
+ "anthropic",
+ "meta",
+ "google",
+ "aws",
+ "openrouter",
+ ]
+ ):
+ # Raise an error if no provider prefix is specified
+ error_msg = f"Model '{model}' does not have a provider prefix. Please specify provider (e.g., 'sambanova/{model}')"
+ logger.error(error_msg)
+ raise ValueError(error_msg)
+
+ system_prompt = "You are a research assistant."
+ prompt = f"""Given the text: "{target}"
+Which of the following options is most relevant to this text?
+{options}
+
+Respond with only the exact text of the most relevant option."""
+
+ # Get pipeline run name and id for trace_name and trace_id if running in a step
+ trace_name = None
+ trace_id = None
+ try:
+ context = get_step_context()
+ trace_name = context.pipeline_run.name
+ trace_id = str(context.pipeline_run.id)
+ except RuntimeError:
+ # Not running in a step context
+ pass
+
+ # Use provided tags or default to ["find_most_relevant_string"]
+ if tags is None:
+ tags = ["find_most_relevant_string"]
+
+ # Build metadata dict
+ metadata = {"project": project, "tags": tags}
+ # Convert tags to trace_metadata format
+ metadata["trace_metadata"] = {tag: True for tag in tags}
+ if trace_name:
+ metadata["trace_name"] = trace_name
+ if trace_id:
+ metadata["trace_id"] = trace_id
+
+ response = completion(
+ model=model,
+ messages=[
+ {"role": "system", "content": system_prompt},
+ {"role": "user", "content": prompt},
+ ],
+ max_tokens=100,
+ temperature=0.2,
+ metadata=metadata,
+ )
+
+ answer = response.choices[0].message.content.strip()
+
+ # Check if the answer is one of the options
+ if answer in options:
+ return answer
+
+ # If not an exact match, find the closest one
+ for option in options:
+ if option in answer or answer in option:
+ return option
+
+ except Exception as e:
+ logger.error(f"Error finding relevant string with LLM: {e}")
+
+ # Simple relevance check - find exact matches first
+ for option in options:
+ if target.lower() == option.lower():
+ return option
+
+ # Then check partial matches
+ for option in options:
+ if is_text_relevant(target, option):
+ return option
+
+ # Return the first option as a fallback
+ return options[0]
+
+
+def synthesize_information(
+ synthesis_input: Dict[str, Any],
+ model: str = "sambanova/Llama-4-Maverick-17B-128E-Instruct",
+ system_prompt: Optional[str] = None,
+ project: str = "deep-research",
+ tags: Optional[List[str]] = None,
+) -> Dict[str, Any]:
+ """Synthesize information from search results for a sub-question.
+
+ Uses litellm for model inference.
+
+ Args:
+ synthesis_input: Dictionary with sub-question, search results, and sources
+ model: Model to use (with provider prefix)
+ system_prompt: System prompt for the LLM
+ project: Langfuse project name for LLM tracking
+ tags: Optional list of tags for Langfuse tracking. Defaults to ["information_synthesis"] if None.
+
+ Returns:
+ Dictionary with synthesized information
+ """
+ if system_prompt is None:
+ system_prompt = SYNTHESIS_PROMPT
+
+ sub_question_for_log = synthesis_input.get(
+ "sub_question", "unknown question"
+ )
+
+ # Define the fallback response
+ fallback_response = {
+ "synthesized_answer": f"Synthesis failed for '{sub_question_for_log}'.",
+ "key_sources": synthesis_input.get("sources", [])[:1],
+ "confidence_level": "low",
+ "information_gaps": "An error occurred during the synthesis process.",
+ }
+
+ # Use provided tags or default to ["information_synthesis"]
+ if tags is None:
+ tags = ["information_synthesis"]
+
+ # Use the utility function to get structured output
+ result = get_structured_llm_output(
+ prompt=json.dumps(synthesis_input),
+ system_prompt=system_prompt,
+ model=model,
+ fallback_response=fallback_response,
+ max_tokens=3000, # Increased for more detailed synthesis
+ project=project,
+ tags=tags,
+ )
+
+ return result
diff --git a/deep_research/utils/prompt_loader.py b/deep_research/utils/prompt_loader.py
new file mode 100644
index 00000000..c2b3d2f7
--- /dev/null
+++ b/deep_research/utils/prompt_loader.py
@@ -0,0 +1,136 @@
+"""Utility functions for loading prompts into the PromptsBundle model.
+
+This module provides functions to create PromptsBundle instances from
+the existing prompt definitions in prompts.py.
+"""
+
+from utils import prompts
+from utils.prompt_models import PromptsBundle, PromptTemplate
+
+
+def load_prompts_bundle(pipeline_version: str = "1.2.0") -> PromptsBundle:
+ """Load all prompts from prompts.py into a PromptsBundle.
+
+ Args:
+ pipeline_version: Version of the pipeline using these prompts
+
+ Returns:
+ PromptsBundle containing all prompts
+ """
+ # Create PromptTemplate instances for each prompt
+ search_query_prompt = PromptTemplate(
+ name="search_query_prompt",
+ content=prompts.DEFAULT_SEARCH_QUERY_PROMPT,
+ description="Generates effective search queries from sub-questions",
+ version="1.0.0",
+ tags=["search", "query", "information-gathering"],
+ )
+
+ query_decomposition_prompt = PromptTemplate(
+ name="query_decomposition_prompt",
+ content=prompts.QUERY_DECOMPOSITION_PROMPT,
+ description="Breaks down complex research queries into specific sub-questions",
+ version="1.0.0",
+ tags=["analysis", "decomposition", "planning"],
+ )
+
+ synthesis_prompt = PromptTemplate(
+ name="synthesis_prompt",
+ content=prompts.SYNTHESIS_PROMPT,
+ description="Synthesizes search results into comprehensive answers for sub-questions",
+ version="1.1.0",
+ tags=["synthesis", "integration", "analysis"],
+ )
+
+ viewpoint_analysis_prompt = PromptTemplate(
+ name="viewpoint_analysis_prompt",
+ content=prompts.VIEWPOINT_ANALYSIS_PROMPT,
+ description="Analyzes synthesized answers across different perspectives and viewpoints",
+ version="1.1.0",
+ tags=["analysis", "viewpoint", "perspective"],
+ )
+
+ reflection_prompt = PromptTemplate(
+ name="reflection_prompt",
+ content=prompts.REFLECTION_PROMPT,
+ description="Evaluates research and identifies gaps, biases, and areas for improvement",
+ version="1.0.0",
+ tags=["reflection", "critique", "improvement"],
+ )
+
+ additional_synthesis_prompt = PromptTemplate(
+ name="additional_synthesis_prompt",
+ content=prompts.ADDITIONAL_SYNTHESIS_PROMPT,
+ description="Enhances original synthesis with new information and addresses critique points",
+ version="1.1.0",
+ tags=["synthesis", "enhancement", "integration"],
+ )
+
+ conclusion_generation_prompt = PromptTemplate(
+ name="conclusion_generation_prompt",
+ content=prompts.CONCLUSION_GENERATION_PROMPT,
+ description="Synthesizes all research findings into a comprehensive conclusion",
+ version="1.0.0",
+ tags=["report", "conclusion", "synthesis"],
+ )
+
+ executive_summary_prompt = PromptTemplate(
+ name="executive_summary_prompt",
+ content=prompts.EXECUTIVE_SUMMARY_GENERATION_PROMPT,
+ description="Creates a compelling, insight-driven executive summary",
+ version="1.1.0",
+ tags=["report", "summary", "insights"],
+ )
+
+ introduction_prompt = PromptTemplate(
+ name="introduction_prompt",
+ content=prompts.INTRODUCTION_GENERATION_PROMPT,
+ description="Creates a contextual, engaging introduction",
+ version="1.1.0",
+ tags=["report", "introduction", "context"],
+ )
+
+ # Create and return the bundle
+ return PromptsBundle(
+ search_query_prompt=search_query_prompt,
+ query_decomposition_prompt=query_decomposition_prompt,
+ synthesis_prompt=synthesis_prompt,
+ viewpoint_analysis_prompt=viewpoint_analysis_prompt,
+ reflection_prompt=reflection_prompt,
+ additional_synthesis_prompt=additional_synthesis_prompt,
+ conclusion_generation_prompt=conclusion_generation_prompt,
+ executive_summary_prompt=executive_summary_prompt,
+ introduction_prompt=introduction_prompt,
+ pipeline_version=pipeline_version,
+ )
+
+
+def get_prompt_for_step(bundle: PromptsBundle, step_name: str) -> str:
+ """Get the appropriate prompt content for a specific step.
+
+ Args:
+ bundle: The PromptsBundle containing all prompts
+ step_name: Name of the step requesting the prompt
+
+ Returns:
+ The prompt content string
+
+ Raises:
+ ValueError: If no prompt mapping exists for the step
+ """
+ # Map step names to prompt attributes
+ step_to_prompt_mapping = {
+ "query_decomposition": "query_decomposition_prompt",
+ "search_query_generation": "search_query_prompt",
+ "synthesis": "synthesis_prompt",
+ "viewpoint_analysis": "viewpoint_analysis_prompt",
+ "reflection": "reflection_prompt",
+ "additional_synthesis": "additional_synthesis_prompt",
+ "conclusion_generation": "conclusion_generation_prompt",
+ }
+
+ prompt_attr = step_to_prompt_mapping.get(step_name)
+ if not prompt_attr:
+ raise ValueError(f"No prompt mapping found for step: {step_name}")
+
+ return bundle.get_prompt_content(prompt_attr)
diff --git a/deep_research/utils/prompt_models.py b/deep_research/utils/prompt_models.py
new file mode 100644
index 00000000..c96a7bd0
--- /dev/null
+++ b/deep_research/utils/prompt_models.py
@@ -0,0 +1,123 @@
+"""Pydantic models for prompt tracking and management.
+
+This module contains models for bundling prompts as trackable artifacts
+in the ZenML pipeline, enabling better observability and version control.
+"""
+
+from datetime import datetime
+from typing import Dict, Optional
+
+from pydantic import BaseModel, Field
+
+
+class PromptTemplate(BaseModel):
+ """Represents a single prompt template with metadata."""
+
+ name: str = Field(..., description="Unique identifier for the prompt")
+ content: str = Field(..., description="The actual prompt template content")
+ description: str = Field(
+ "", description="Human-readable description of what this prompt does"
+ )
+ version: str = Field("1.0.0", description="Version of the prompt template")
+ tags: list[str] = Field(
+ default_factory=list, description="Tags for categorizing prompts"
+ )
+
+ model_config = {
+ "extra": "ignore",
+ "frozen": False,
+ "validate_assignment": True,
+ }
+
+
+class PromptsBundle(BaseModel):
+ """Bundle of all prompts used in the research pipeline.
+
+ This model serves as a single artifact that contains all prompts,
+ making them trackable, versionable, and visualizable in the ZenML dashboard.
+ """
+
+ # Core prompts used in the pipeline
+ search_query_prompt: PromptTemplate
+ query_decomposition_prompt: PromptTemplate
+ synthesis_prompt: PromptTemplate
+ viewpoint_analysis_prompt: PromptTemplate
+ reflection_prompt: PromptTemplate
+ additional_synthesis_prompt: PromptTemplate
+ conclusion_generation_prompt: PromptTemplate
+ executive_summary_prompt: PromptTemplate
+ introduction_prompt: PromptTemplate
+
+ # Metadata
+ pipeline_version: str = Field(
+ "1.0.0", description="Version of the pipeline using these prompts"
+ )
+ created_at: str = Field(
+ default_factory=lambda: datetime.now().isoformat(),
+ description="Timestamp when this bundle was created",
+ )
+
+ # Additional prompts can be stored here
+ custom_prompts: Dict[str, PromptTemplate] = Field(
+ default_factory=dict,
+ description="Additional custom prompts not part of the core set",
+ )
+
+ model_config = {
+ "extra": "ignore",
+ "frozen": False,
+ "validate_assignment": True,
+ }
+
+ def get_prompt_by_name(self, name: str) -> Optional[PromptTemplate]:
+ """Retrieve a prompt by its name.
+
+ Args:
+ name: Name of the prompt to retrieve
+
+ Returns:
+ PromptTemplate if found, None otherwise
+ """
+ # Check core prompts
+ for field_name, field_value in self.__dict__.items():
+ if (
+ isinstance(field_value, PromptTemplate)
+ and field_value.name == name
+ ):
+ return field_value
+
+ # Check custom prompts
+ return self.custom_prompts.get(name)
+
+ def list_all_prompts(self) -> Dict[str, PromptTemplate]:
+ """Get all prompts as a dictionary.
+
+ Returns:
+ Dictionary mapping prompt names to PromptTemplate objects
+ """
+ all_prompts = {}
+
+ # Add core prompts
+ for field_name, field_value in self.__dict__.items():
+ if isinstance(field_value, PromptTemplate):
+ all_prompts[field_value.name] = field_value
+
+ # Add custom prompts
+ all_prompts.update(self.custom_prompts)
+
+ return all_prompts
+
+ def get_prompt_content(self, prompt_type: str) -> str:
+ """Get the content of a specific prompt by its type.
+
+ Args:
+ prompt_type: Type of prompt (e.g., 'search_query_prompt', 'synthesis_prompt')
+
+ Returns:
+ The prompt content string
+
+ Raises:
+ AttributeError: If prompt type doesn't exist
+ """
+ prompt = getattr(self, prompt_type)
+ return prompt.content
diff --git a/deep_research/utils/prompts.py b/deep_research/utils/prompts.py
new file mode 100644
index 00000000..e072b238
--- /dev/null
+++ b/deep_research/utils/prompts.py
@@ -0,0 +1,1605 @@
+"""
+Centralized collection of prompts used throughout the deep research pipeline.
+
+This module contains all system prompts used by LLM calls in various steps of the
+research pipeline to ensure consistency and make prompt management easier.
+"""
+
+# Search query generation prompt
+# Used to generate effective search queries from sub-questions
+DEFAULT_SEARCH_QUERY_PROMPT = """
+You are a Deep Research assistant. Given a specific research sub-question, your task is to formulate an effective search
+query that will help find relevant information to answer the question.
+
+A good search query should:
+1. Extract the key concepts from the sub-question
+2. Use precise, specific terminology
+3. Exclude unnecessary words or context
+4. Include alternative terms or synonyms when helpful
+5. Be concise yet comprehensive enough to find relevant results
+
+Format the output in json with the following json schema definition:
+
+
+
+Make sure that the output is a json object with an output json schema defined above.
+Only return the json object, no explanation or additional text.
+"""
+
+# Query decomposition prompt
+# Used to break down complex research queries into specific sub-questions
+QUERY_DECOMPOSITION_PROMPT = """
+You are a Deep Research assistant specializing in research design. You will be given a MAIN RESEARCH QUERY that needs to be explored comprehensively. Your task is to create diverse, insightful sub-questions that explore different dimensions of the topic.
+
+IMPORTANT: The main query should be interpreted as a single research question, not as a noun phrase. For example:
+- If the query is "Is LLMOps a subset of MLOps?", create questions ABOUT LLMOps and MLOps, not questions like "What is 'Is LLMOps a subset of MLOps?'"
+- Focus on the concepts, relationships, and implications within the query
+
+Create sub-questions that explore these DIFFERENT DIMENSIONS:
+
+1. **Definitional/Conceptual**: Define key terms and establish conceptual boundaries
+ Example: "What are the core components and characteristics of LLMOps?"
+
+2. **Comparative/Relational**: Compare and contrast the concepts mentioned
+ Example: "How do the workflows and tooling of LLMOps differ from traditional MLOps?"
+
+3. **Historical/Evolutionary**: Trace development and emergence
+ Example: "How did LLMOps emerge from MLOps practices?"
+
+4. **Structural/Technical**: Examine technical architecture and implementation
+ Example: "What specific tools and platforms are unique to LLMOps?"
+
+5. **Practical/Use Cases**: Explore real-world applications
+ Example: "What are the key use cases that require LLMOps but not traditional MLOps?"
+
+6. **Stakeholder/Industry**: Consider different perspectives and adoption
+ Example: "How are different industries adopting LLMOps vs MLOps?"
+
+7. **Challenges/Limitations**: Identify problems and constraints
+ Example: "What unique challenges does LLMOps face that MLOps doesn't?"
+
+8. **Future/Trends**: Look at emerging developments
+ Example: "How is the relationship between LLMOps and MLOps expected to evolve?"
+
+QUALITY GUIDELINES:
+- Each sub-question must explore a DIFFERENT dimension - no repetitive variations
+- Questions should be specific, concrete, and investigable
+- Mix descriptive ("what/who") with analytical ("why/how") questions
+- Ensure questions build toward answering the main query comprehensively
+- Frame questions to elicit detailed, nuanced responses
+- Consider technical, business, organizational, and strategic aspects
+
+Format the output in json with the following json schema definition:
+
+
+
+Make sure that the output is a json object with an output json schema defined above.
+Only return the json object, no explanation or additional text.
+"""
+
+# Synthesis prompt for individual sub-questions
+# Used to synthesize search results into comprehensive answers for sub-questions
+SYNTHESIS_PROMPT = """
+You are a Deep Research assistant specializing in information synthesis. Given a sub-question and search results, your task is to synthesize the information
+into a comprehensive, accurate, and well-structured answer.
+
+Your synthesis should:
+1. Begin with a direct, concise answer to the sub-question in the first paragraph
+2. Provide detailed evidence and explanation in subsequent paragraphs (at least 3-5 paragraphs total)
+3. Integrate information from multiple sources, citing them within your answer
+4. Acknowledge any conflicting information or contrasting viewpoints you encounter
+5. Use data, statistics, examples, and quotations when available to strengthen your answer
+6. Organize information logically with a clear flow between concepts
+7. Identify key sources that provided the most valuable information (at least 2-3 sources)
+8. Explicitly acknowledge information gaps where the search results were incomplete
+9. Write in plain text format - do NOT use markdown formatting, bullet points, or special characters
+
+Confidence level criteria:
+- HIGH: Multiple high-quality sources provide consistent information, comprehensive coverage of the topic, and few information gaps
+- MEDIUM: Decent sources with some consistency, but notable information gaps or some conflicting information
+- LOW: Limited sources, major information gaps, significant contradictions, or only tangentially relevant information
+
+Information gaps should specifically identify:
+1. Aspects of the question that weren't addressed in the search results
+2. Areas where more detailed or up-to-date information would be valuable
+3. Perspectives or data sources that would complement the existing information
+
+Format the output in json with the following json schema definition:
+
+
+
+Make sure that the output is a json object with an output json schema defined above.
+Only return the json object, no explanation or additional text.
+"""
+
+# Viewpoint analysis prompt for cross-perspective examination
+# Used to analyze synthesized answers across different perspectives and viewpoints
+VIEWPOINT_ANALYSIS_PROMPT = """
+You are a Deep Research assistant specializing in multi-perspective analysis. You will be given a set of synthesized answers
+to sub-questions related to a main research query. Your task is to perform a thorough, nuanced analysis of how different
+perspectives would interpret this information.
+
+Think deeply about the following viewpoint categories and how they would approach the information differently:
+- Scientific: Evidence-based, empirical approach focused on data, research findings, and methodological rigor
+- Political: Power dynamics, governance structures, policy implications, and ideological frameworks
+- Economic: Resource allocation, financial impacts, market dynamics, and incentive structures
+- Social: Cultural norms, community impacts, group dynamics, and public welfare
+- Ethical: Moral principles, values considerations, rights and responsibilities, and normative judgments
+- Historical: Long-term patterns, precedents, contextual development, and evolutionary change
+
+For each synthesized answer, analyze how these different perspectives would interpret the information by:
+
+1. Identifying 5-8 main points of agreement where multiple perspectives align (with specific examples)
+2. Analyzing at least 3-5 areas of tension between perspectives with:
+ - A clear topic title for each tension point
+ - Contrasting interpretations from at least 2-3 different viewpoint categories per tension
+ - Specific examples or evidence showing why these perspectives differ
+ - The nuanced positions of each perspective, not just simplified oppositions
+
+3. Thoroughly examining perspective gaps by identifying:
+ - Which perspectives are underrepresented or missing in the current research
+ - How including these missing perspectives would enrich understanding
+ - Specific questions or dimensions that remain unexplored
+ - Write in plain text format - do NOT use markdown formatting, bullet points, or special characters
+
+4. Developing integrative insights that:
+ - Synthesize across multiple perspectives to form a more complete understanding
+ - Highlight how seemingly contradictory viewpoints can complement each other
+ - Suggest frameworks for reconciling tensions or finding middle-ground approaches
+ - Identify actionable takeaways that incorporate multiple perspectives
+ - Write in plain text format - do NOT use markdown formatting, bullet points, or special characters
+
+Format the output in json with the following json schema definition:
+
+
+
+Make sure that the output is a json object with an output json schema defined above.
+Only return the json object, no explanation or additional text.
+"""
+
+# Reflection prompt for self-critique and improvement
+# Used to evaluate the research and identify gaps, biases, and areas for improvement
+REFLECTION_PROMPT = """
+You are a Deep Research assistant with the ability to critique and improve your own research. You will be given:
+1. The main research query
+2. The sub-questions explored so far
+3. The synthesized information for each sub-question
+4. Any viewpoint analysis performed
+
+Your task is to critically evaluate this research and identify:
+1. Areas where the research is incomplete or has gaps
+2. Questions that are important but not yet answered
+3. Aspects where additional evidence or depth would significantly improve the research
+4. Potential biases or limitations in the current findings
+
+Be constructively critical and identify the most important improvements that would substantially enhance the research.
+
+Format the output in json with the following json schema definition:
+
+
+
+Make sure that the output is a json object with an output json schema defined above.
+Only return the json object, no explanation or additional text.
+"""
+
+# Additional synthesis prompt for incorporating new information
+# Used to enhance original synthesis with new information and address critique points
+ADDITIONAL_SYNTHESIS_PROMPT = """
+You are a Deep Research assistant. You will be given:
+1. The original synthesized information on a research topic
+2. New information from additional research
+3. A critique of the original synthesis
+
+Your task is to enhance the original synthesis by incorporating the new information and addressing the critique.
+The updated synthesis should:
+1. Integrate new information seamlessly
+2. Address gaps identified in the critique
+3. Maintain a balanced, comprehensive, and accurate representation
+4. Preserve the strengths of the original synthesis
+5. Write in plain text format - do NOT use markdown formatting, bullet points, or special characters
+
+Format the output in json with the following json schema definition:
+
+
+
+Make sure that the output is a json object with an output json schema defined above.
+Only return the json object, no explanation or additional text.
+"""
+
+# Final report generation prompt
+# Used to compile a comprehensive HTML research report from all synthesized information
+REPORT_GENERATION_PROMPT = """
+You are a Deep Research assistant responsible for compiling an in-depth, comprehensive research report. You will be given:
+1. The original research query
+2. The sub-questions that were explored
+3. Synthesized information for each sub-question
+4. Viewpoint analysis comparing different perspectives (if available)
+5. Reflection metadata highlighting improvements and limitations
+
+Your task is to create a well-structured, coherent, professional-quality research report with the following features:
+
+EXECUTIVE SUMMARY (250-400 words):
+- Begin with a compelling, substantive executive summary that provides genuine insight
+- Highlight 3-5 key findings or insights that represent the most important discoveries
+- Include brief mention of methodology and limitations
+- Make the summary self-contained so it can be read independently of the full report
+- End with 1-2 sentences on broader implications or applications of the research
+
+INTRODUCTION (200-300 words):
+- Provide relevant background context on the main research query
+- Explain why this topic is significant or worth investigating
+- Outline the methodological approach used (sub-questions, search strategy, synthesis)
+- Preview the overall structure of the report
+
+SUB-QUESTION SECTIONS:
+- For each sub-question, create a dedicated section with:
+ * A descriptive section title (not just repeating the sub-question)
+ * A brief (1 paragraph) overview of key findings for this sub-question
+ * A "Key Findings" box highlighting 3-4 important discoveries for scannable reading
+ * The detailed, synthesized answer with appropriate paragraph breaks, lists, and formatting
+ * Proper citation of sources within the text (e.g., "According to [Source Name]...")
+ * Clear confidence indicator with appropriate styling
+ * Information gaps clearly identified in their own subsection
+ * Complete list of key sources used
+
+VIEWPOINT ANALYSIS SECTION (if available):
+- Create a detailed section that:
+ * Explains the purpose and value of multi-perspective analysis
+ * Presents points of agreement as actionable insights, not just observations
+ * Structures tension areas with clear topic headings and balanced presentation of viewpoints
+ * Uses visual elements (different background colors, icons) to distinguish different perspectives
+ * Integrates perspective gaps and insights into a cohesive narrative
+
+CONCLUSION (300-400 words):
+- Synthesize the overall findings, not just summarizing each section
+- Connect insights from different sub-questions to form higher-level understanding
+- Address the main research query directly with evidence-based conclusions
+- Acknowledge remaining uncertainties and suggestions for further research
+- End with implications or applications of the research findings
+
+OVERALL QUALITY REQUIREMENTS:
+1. Create visually scannable content with clear headings, bullet points, and short paragraphs
+2. Use semantic HTML (h1, h2, h3, p, blockquote, etc.) to create proper document structure
+3. Include a comprehensive table of contents with anchor links to all major sections
+4. Format all sources consistently in the references section with proper linking when available
+5. Use tables, lists, and blockquotes to improve readability and highlight important information
+6. Apply appropriate styling for different confidence levels (high, medium, low)
+7. Ensure proper HTML nesting and structure throughout the document
+8. Balance sufficient detail with clarity and conciseness
+9. Make all text directly actionable and insight-driven, not just descriptive
+
+The report should be formatted in HTML with appropriate headings, paragraphs, citations, and formatting.
+Use semantic HTML (h1, h2, h3, p, blockquote, etc.) to create a structured document.
+Include a table of contents at the beginning with anchor links to each section.
+For citations, use a consistent format and collect them in a references section at the end.
+
+Include this exact CSS stylesheet in your HTML to ensure consistent styling (do not modify it):
+
+```css
+
+```
+
+The HTML structure should follow this pattern:
+
+```html
+
+
+
+
+
+ [CSS STYLESHEET GOES HERE]
+
+
+
+ Research Report: [Main Query]
+
+
+
+ Table of Contents
+
+ - Executive Summary
+ - Introduction
+ [SUB-QUESTIONS LINKS]
+ [ADDITIONAL SECTIONS LINKS]
+ - Conclusion
+ - References
+
+
+
+
+
+ Executive Summary
+ [CONCISE SUMMARY OF KEY FINDINGS]
+
+
+
+
+ Introduction
+ [INTRODUCTION TO THE RESEARCH QUERY]
+ [OVERVIEW OF THE APPROACH AND SUB-QUESTIONS]
+
+
+
+ [FOR EACH SUB-QUESTION]:
+
+ [INDEX]. [SUB-QUESTION TEXT]
+ Confidence Level: [LEVEL]
+
+
+
+ Key Findings
+
+ - [KEY FINDING 1]
+ - [KEY FINDING 2]
+ [...]
+
+
+
+
+ [DETAILED ANSWER]
+
+
+
+
+ Information Gaps
+ [GAPS TEXT]
+
+
+
+
+ Key Sources
+
+ - [SOURCE 1]
+ - [SOURCE 2]
+ [...]
+
+
+
+
+
+
+ Viewpoint Analysis
+
+ Points of Agreement
+
+
+ - [AGREEMENT 1]
+ - [AGREEMENT 2]
+ [...]
+
+
+
+ Areas of Tension
+ [FOR EACH TENSION]:
+
+ [TENSION TOPIC]
+
+ - [VIEWPOINT 1 TITLE]
+ - [VIEWPOINT 1 CONTENT]
+ - [VIEWPOINT 2 TITLE]
+ - [VIEWPOINT 2 CONTENT]
+ [...]
+
+
+
+ Perspective Gaps
+ [PERSPECTIVE GAPS CONTENT]
+
+ Integrative Insights
+ [INTEGRATIVE INSIGHTS CONTENT]
+
+
+
+
+ Conclusion
+ [CONCLUSION TEXT]
+
+
+
+
+ References
+
+ - [REFERENCE 1]
+ - [REFERENCE 2]
+ [...]
+
+
+
+
+
+```
+
+Special instructions:
+1. For each sub-question, display the confidence level with appropriate styling (confidence-high, confidence-medium, or confidence-low)
+2. Extract 2-3 key findings from each answer to create the key-findings box
+3. Format all sources consistently in the references section
+4. Use tables, lists, and blockquotes where appropriate to improve readability
+5. Use the notice classes (info, warning) to highlight important information or limitations
+6. Ensure all sections have proper ID attributes for the table of contents links
+
+Return only the complete HTML code for the report, with no explanations or additional text.
+"""
+
+# Static HTML template for direct report generation without LLM
+STATIC_HTML_TEMPLATE = """
+
+
+
+
+ Research Report: {main_query}
+
+
+
+
+ Research Report: {main_query}
+
+
+
+ Table of Contents
+
+ - Executive Summary
+ - Introduction
+ {sub_questions_toc}
+ {additional_sections_toc}
+ - Conclusion
+ - References
+
+
+
+
+
+ Executive Summary
+ {executive_summary}
+
+
+
+
+ Introduction
+ {introduction_html}
+
+
+
+ {sub_questions_html}
+
+
+ {viewpoint_analysis_html}
+
+
+
+ Conclusion
+ {conclusion_html}
+
+
+
+
+ References
+ {references_html}
+
+
+
+
+"""
+
+# Template for sub-question section in the static HTML report
+SUB_QUESTION_TEMPLATE = """
+
+
+ {index}. {question}
+
+
+ Confidence: {confidence_upper}
+
+
+
+
+ {answer}
+
+
+ {info_gaps_html}
+
+ {key_sources_html}
+
+"""
+
+# Template for viewpoint analysis section in the static HTML report
+VIEWPOINT_ANALYSIS_TEMPLATE = """
+
+ Viewpoint Analysis
+
+
+ Points of Agreement
+
+
+ {agreements_html}
+
+
+
+
+
+ Areas of Tension
+
+ {tensions_html}
+
+
+
+
+ Perspective Gaps
+
+ {perspective_gaps}
+
+
+
+
+ Integrative Insights
+
+ {integrative_insights}
+
+
+
+"""
+
+# Executive Summary generation prompt
+# Used to create a compelling, insight-driven executive summary
+EXECUTIVE_SUMMARY_GENERATION_PROMPT = """
+You are a Deep Research assistant specializing in creating executive summaries. Given comprehensive research findings, your task is to create a compelling executive summary that captures the essence of the research and its key insights.
+
+Your executive summary should:
+
+1. **Opening Statement (1-2 sentences):**
+ - Start with a powerful, direct answer to the main research question
+ - Make it clear and definitive based on the evidence gathered
+
+2. **Key Findings (3-5 bullet points):**
+ - Extract the MOST IMPORTANT discoveries from across all sub-questions
+ - Focus on insights that are surprising, actionable, or paradigm-shifting
+ - Each finding should be specific and evidence-based, not generic
+ - Prioritize findings that directly address the main query
+
+3. **Critical Insights (2-3 sentences):**
+ - Synthesize patterns or themes that emerged across multiple sub-questions
+ - Highlight any unexpected discoveries or counter-intuitive findings
+ - Connect disparate findings to reveal higher-level understanding
+
+4. **Implications (2-3 sentences):**
+ - What do these findings mean for practitioners/stakeholders?
+ - What actions or decisions can be made based on this research?
+ - Why should the reader care about these findings?
+
+5. **Confidence and Limitations (1-2 sentences):**
+ - Briefly acknowledge the overall confidence level of the findings
+ - Note any significant gaps or areas requiring further investigation
+
+IMPORTANT GUIDELINES:
+- Be CONCISE but INSIGHTFUL - every sentence should add value
+- Use active voice and strong, definitive language where evidence supports it
+- Avoid generic statements - be specific to the actual research findings
+- Lead with the most important information
+- Make it self-contained - reader should understand key findings without reading the full report
+- Target length: 250-400 words
+
+Format as well-structured HTML paragraphs using tags and
/- for bullet points.
+"""
+
+# Introduction generation prompt
+# Used to create a contextual, engaging introduction
+INTRODUCTION_GENERATION_PROMPT = """
+You are a Deep Research assistant specializing in creating engaging introductions. Given a research query and the sub-questions explored, your task is to create an introduction that provides context and sets up the reader's expectations.
+
+Your introduction should:
+
+1. **Context and Relevance (2-3 sentences):**
+ - Why is this research question important NOW?
+ - What makes this topic significant or worth investigating?
+ - Connect to current trends, debates, or challenges in the field
+
+2. **Scope and Approach (2-3 sentences):**
+ - What specific aspects of the topic does this research explore?
+ - Briefly mention the key dimensions covered (based on sub-questions)
+ - Explain the systematic approach without being too technical
+
+3. **What to Expect (2-3 sentences):**
+ - Preview the structure of the report
+ - Hint at some of the interesting findings or tensions discovered
+ - Set expectations about the depth and breadth of analysis
+
+IMPORTANT GUIDELINES:
+- Make it engaging - hook the reader's interest from the start
+- Provide real context, not generic statements
+- Connect to why this matters for the reader
+- Keep it concise but informative (200-300 words)
+- Use active voice and clear language
+- Build anticipation for the findings without giving everything away
+
+Format as well-structured HTML paragraphs using
tags. Do NOT include any headings or section titles.
+"""
+
+# Conclusion generation prompt
+# Used to synthesize all research findings into a comprehensive conclusion
+CONCLUSION_GENERATION_PROMPT = """
+You are a Deep Research assistant specializing in synthesizing comprehensive research conclusions. Given all the research findings from a deep research study, your task is to create a thoughtful, evidence-based conclusion that ties together the overall findings.
+
+Your conclusion should:
+
+1. **Synthesis and Integration (150-200 words):**
+ - Connect insights from different sub-questions to form a higher-level understanding
+ - Identify overarching themes and patterns that emerge from the research
+ - Highlight how different findings relate to and support each other
+ - Avoid simply summarizing each section separately
+
+2. **Direct Response to Main Query (100-150 words):**
+ - Address the original research question directly with evidence-based conclusions
+ - State what the research definitively established vs. what remains uncertain
+ - Provide a clear, actionable answer based on the synthesized evidence
+
+3. **Limitations and Future Directions (100-120 words):**
+ - Acknowledge remaining uncertainties and information gaps across all sections
+ - Suggest specific areas where additional research would be most valuable
+ - Identify what types of evidence or perspectives would strengthen the findings
+
+4. **Implications and Applications (80-100 words):**
+ - Explain the practical significance of the research findings
+ - Suggest how the insights might be applied or what they mean for stakeholders
+ - Connect findings to broader contexts or implications
+
+Format your output as a well-structured conclusion section in HTML format with appropriate paragraph breaks and formatting. Use
tags for paragraphs and organize the content logically with clear transitions between the different aspects outlined above.
+
+IMPORTANT: Do NOT include any headings like "Conclusion",
, or tags - the section already has a heading. Start directly with the conclusion content in paragraph form. Just create flowing, well-structured paragraphs that cover all four aspects naturally.
+
+Ensure the conclusion feels cohesive and draws meaningful connections between findings rather than just listing them sequentially.
+"""
diff --git a/deep_research/utils/pydantic_models.py b/deep_research/utils/pydantic_models.py
new file mode 100644
index 00000000..9fca23a3
--- /dev/null
+++ b/deep_research/utils/pydantic_models.py
@@ -0,0 +1,300 @@
+"""Pydantic model definitions for the research pipeline.
+
+This module contains all the Pydantic models that represent the state of the research
+pipeline. These models replace the previous dataclasses implementation and leverage
+Pydantic's validation, serialization, and integration with ZenML.
+"""
+
+import time
+from typing import Any, Dict, List, Optional
+
+from pydantic import BaseModel, Field
+from typing_extensions import Literal
+
+
+class SearchResult(BaseModel):
+ """Represents a search result for a sub-question."""
+
+ url: str = ""
+ content: str = ""
+ title: str = ""
+ snippet: str = ""
+ metadata: Optional[Dict[str, Any]] = Field(default_factory=dict)
+
+ model_config = {
+ "extra": "ignore", # Ignore extra fields during deserialization
+ "frozen": False, # Allow attribute updates
+ "validate_assignment": True, # Validate when attributes are set
+ }
+
+
+class ViewpointTension(BaseModel):
+ """Represents a tension between different viewpoints on a topic."""
+
+ topic: str = ""
+ viewpoints: Dict[str, str] = Field(default_factory=dict)
+
+ model_config = {
+ "extra": "ignore",
+ "frozen": False,
+ "validate_assignment": True,
+ }
+
+
+class SynthesizedInfo(BaseModel):
+ """Represents synthesized information for a sub-question."""
+
+ synthesized_answer: str = ""
+ key_sources: List[str] = Field(default_factory=list)
+ confidence_level: Literal["high", "medium", "low"] = "medium"
+ information_gaps: str = ""
+ improvements: List[str] = Field(default_factory=list)
+
+ model_config = {
+ "extra": "ignore",
+ "frozen": False,
+ "validate_assignment": True,
+ }
+
+
+class ViewpointAnalysis(BaseModel):
+ """Represents the analysis of different viewpoints on the research topic."""
+
+ main_points_of_agreement: List[str] = Field(default_factory=list)
+ areas_of_tension: List[ViewpointTension] = Field(default_factory=list)
+ perspective_gaps: str = ""
+ integrative_insights: str = ""
+
+ model_config = {
+ "extra": "ignore",
+ "frozen": False,
+ "validate_assignment": True,
+ }
+
+
+class ReflectionMetadata(BaseModel):
+ """Metadata about the reflection process."""
+
+ critique_summary: List[str] = Field(default_factory=list)
+ additional_questions_identified: List[str] = Field(default_factory=list)
+ searches_performed: List[str] = Field(default_factory=list)
+ improvements_made: float = Field(
+ default=0
+ ) # Changed from int to float to handle timestamp values
+ error: Optional[str] = None
+
+ model_config = {
+ "extra": "ignore",
+ "frozen": False,
+ "validate_assignment": True,
+ }
+
+
+class ResearchState(BaseModel):
+ """Comprehensive state object for the enhanced research pipeline."""
+
+ # Initial query information
+ main_query: str = ""
+ sub_questions: List[str] = Field(default_factory=list)
+
+ # Information gathering results
+ search_results: Dict[str, List[SearchResult]] = Field(default_factory=dict)
+
+ # Synthesized information
+ synthesized_info: Dict[str, SynthesizedInfo] = Field(default_factory=dict)
+
+ # Viewpoint analysis
+ viewpoint_analysis: Optional[ViewpointAnalysis] = None
+
+ # Reflection results
+ enhanced_info: Dict[str, SynthesizedInfo] = Field(default_factory=dict)
+ reflection_metadata: Optional[ReflectionMetadata] = None
+
+ # Final report
+ final_report_html: str = ""
+
+ # Search cost tracking
+ search_costs: Dict[str, float] = Field(
+ default_factory=dict,
+ description="Total costs by search provider (e.g., {'exa': 0.0, 'tavily': 0.0})",
+ )
+ search_cost_details: List[Dict[str, Any]] = Field(
+ default_factory=list,
+ description="Detailed log of each search with cost information",
+ )
+ # Format: [{"provider": "exa", "query": "...", "cost": 0.0, "timestamp": ..., "step": "...", "sub_question": "..."}]
+
+ model_config = {
+ "extra": "ignore",
+ "frozen": False,
+ "validate_assignment": True,
+ }
+
+ def get_current_stage(self) -> str:
+ """Determine the current stage of research based on filled data."""
+ if self.final_report_html:
+ return "final_report"
+ elif self.enhanced_info:
+ return "after_reflection"
+ elif self.viewpoint_analysis:
+ return "after_viewpoint_analysis"
+ elif self.synthesized_info:
+ return "after_synthesis"
+ elif self.search_results:
+ return "after_search"
+ elif self.sub_questions:
+ return "after_query_decomposition"
+ elif self.main_query:
+ return "initial"
+ else:
+ return "empty"
+
+ def update_sub_questions(self, sub_questions: List[str]) -> None:
+ """Update the sub-questions list."""
+ self.sub_questions = sub_questions
+
+ def update_search_results(
+ self, search_results: Dict[str, List[SearchResult]]
+ ) -> None:
+ """Update the search results."""
+ self.search_results = search_results
+
+ def update_synthesized_info(
+ self, synthesized_info: Dict[str, SynthesizedInfo]
+ ) -> None:
+ """Update the synthesized information."""
+ self.synthesized_info = synthesized_info
+
+ def update_viewpoint_analysis(
+ self, viewpoint_analysis: ViewpointAnalysis
+ ) -> None:
+ """Update the viewpoint analysis."""
+ self.viewpoint_analysis = viewpoint_analysis
+
+ def update_after_reflection(
+ self,
+ enhanced_info: Dict[str, SynthesizedInfo],
+ metadata: ReflectionMetadata,
+ ) -> None:
+ """Update with reflection results."""
+ self.enhanced_info = enhanced_info
+ self.reflection_metadata = metadata
+
+ def set_final_report(self, html: str) -> None:
+ """Set the final report HTML."""
+ self.final_report_html = html
+
+
+class ReflectionOutput(BaseModel):
+ """Output from the reflection generation step."""
+
+ state: ResearchState
+ recommended_queries: List[str] = Field(default_factory=list)
+ critique_summary: List[Dict[str, Any]] = Field(default_factory=list)
+ additional_questions: List[str] = Field(default_factory=list)
+
+ model_config = {
+ "extra": "ignore",
+ "frozen": False,
+ "validate_assignment": True,
+ }
+
+
+class ApprovalDecision(BaseModel):
+ """Approval decision from human reviewer."""
+
+ approved: bool = False
+ selected_queries: List[str] = Field(default_factory=list)
+ approval_method: str = "" # "APPROVE_ALL", "SKIP", "SELECT_SPECIFIC"
+ reviewer_notes: str = ""
+ timestamp: float = Field(default_factory=lambda: time.time())
+
+ model_config = {
+ "extra": "ignore",
+ "frozen": False,
+ "validate_assignment": True,
+ }
+
+
+class PromptTypeMetrics(BaseModel):
+ """Metrics for a specific prompt type."""
+
+ prompt_type: str
+ total_cost: float
+ input_tokens: int
+ output_tokens: int
+ call_count: int
+ avg_cost_per_call: float
+ percentage_of_total_cost: float
+
+ model_config = {
+ "extra": "ignore",
+ "frozen": False,
+ "validate_assignment": True,
+ }
+
+
+class TracingMetadata(BaseModel):
+ """Metadata about token usage, costs, and performance for a pipeline run."""
+
+ # Pipeline information
+ pipeline_run_name: str = ""
+ pipeline_run_id: str = ""
+
+ # Token usage
+ total_input_tokens: int = 0
+ total_output_tokens: int = 0
+ total_tokens: int = 0
+
+ # Cost information
+ total_cost: float = 0.0
+ cost_breakdown_by_model: Dict[str, float] = Field(default_factory=dict)
+
+ # Performance metrics
+ total_latency_seconds: float = 0.0
+ formatted_latency: str = ""
+ observation_count: int = 0
+
+ # Model usage
+ models_used: List[str] = Field(default_factory=list)
+ model_token_breakdown: Dict[str, Dict[str, int]] = Field(
+ default_factory=dict
+ )
+ # Format: {"model_name": {"input_tokens": X, "output_tokens": Y, "total_tokens": Z}}
+
+ # Trace information
+ trace_id: str = ""
+ trace_name: str = ""
+ trace_tags: List[str] = Field(default_factory=list)
+ trace_metadata: Dict[str, Any] = Field(default_factory=dict)
+
+ # Step-by-step breakdown
+ step_costs: Dict[str, float] = Field(default_factory=dict)
+ step_tokens: Dict[str, Dict[str, int]] = Field(default_factory=dict)
+ # Format: {"step_name": {"input_tokens": X, "output_tokens": Y}}
+
+ # Prompt-level metrics
+ prompt_metrics: List[PromptTypeMetrics] = Field(
+ default_factory=list, description="Cost breakdown by prompt type"
+ )
+
+ # Search provider costs
+ search_costs: Dict[str, float] = Field(
+ default_factory=dict, description="Total costs by search provider"
+ )
+ search_queries_count: Dict[str, int] = Field(
+ default_factory=dict,
+ description="Number of queries by search provider",
+ )
+ search_cost_details: List[Dict[str, Any]] = Field(
+ default_factory=list, description="Detailed search cost information"
+ )
+
+ # Timestamp
+ collected_at: float = Field(default_factory=lambda: time.time())
+
+ model_config = {
+ "extra": "ignore",
+ "frozen": False,
+ "validate_assignment": True,
+ }
diff --git a/deep_research/utils/search_utils.py b/deep_research/utils/search_utils.py
new file mode 100644
index 00000000..a9b04d51
--- /dev/null
+++ b/deep_research/utils/search_utils.py
@@ -0,0 +1,721 @@
+import logging
+import os
+from enum import Enum
+from typing import Any, Dict, List, Optional, Union
+
+from tavily import TavilyClient
+
+try:
+ from exa_py import Exa
+
+ EXA_AVAILABLE = True
+except ImportError:
+ EXA_AVAILABLE = False
+ Exa = None
+
+from utils.llm_utils import get_structured_llm_output
+from utils.prompts import DEFAULT_SEARCH_QUERY_PROMPT
+from utils.pydantic_models import SearchResult
+
+logger = logging.getLogger(__name__)
+
+
+class SearchProvider(Enum):
+ TAVILY = "tavily"
+ EXA = "exa"
+ BOTH = "both"
+
+
+class SearchEngineConfig:
+ """Configuration for search engines"""
+
+ def __init__(self):
+ self.tavily_api_key = os.getenv("TAVILY_API_KEY")
+ self.exa_api_key = os.getenv("EXA_API_KEY")
+ self.default_provider = os.getenv("DEFAULT_SEARCH_PROVIDER", "tavily")
+ self.enable_parallel_search = (
+ os.getenv("ENABLE_PARALLEL_SEARCH", "false").lower() == "true"
+ )
+
+
+def get_search_client(provider: Union[str, SearchProvider]) -> Optional[Any]:
+ """Get the appropriate search client based on provider."""
+ if isinstance(provider, str):
+ provider = SearchProvider(provider.lower())
+
+ config = SearchEngineConfig()
+
+ if provider == SearchProvider.TAVILY:
+ if not config.tavily_api_key:
+ raise ValueError("TAVILY_API_KEY environment variable not set")
+ return TavilyClient(api_key=config.tavily_api_key)
+
+ elif provider == SearchProvider.EXA:
+ if not EXA_AVAILABLE:
+ raise ImportError(
+ "exa-py is not installed. Please install it with: pip install exa-py"
+ )
+ if not config.exa_api_key:
+ raise ValueError("EXA_API_KEY environment variable not set")
+ return Exa(config.exa_api_key)
+
+ return None
+
+
+def tavily_search(
+ query: str,
+ include_raw_content: bool = True,
+ max_results: int = 3,
+ cap_content_length: int = 20000,
+) -> Dict[str, Any]:
+ """Perform a search using the Tavily API.
+
+ Args:
+ query: Search query
+ include_raw_content: Whether to include raw content in results
+ max_results: Maximum number of results to return
+ cap_content_length: Maximum length of content to return
+
+ Returns:
+ Dict[str, Any]: Search results from Tavily in the following format:
+ {
+ "query": str, # The original query
+ "results": List[Dict], # List of search result objects
+ "error": str, # Error message (if an error occurred, otherwise omitted)
+ }
+
+ Each result in "results" has the following structure:
+ {
+ "url": str, # URL of the search result
+ "raw_content": str, # Raw content of the page (if include_raw_content=True)
+ "title": str, # Title of the page
+ "snippet": str, # Snippet of the page content
+ }
+ """
+ try:
+ tavily_client = get_search_client(SearchProvider.TAVILY)
+
+ # First try with advanced search
+ results = tavily_client.search(
+ query=query,
+ include_raw_content=include_raw_content,
+ max_results=max_results,
+ search_depth="advanced", # Use advanced search for better results
+ include_domains=[], # No domain restrictions
+ exclude_domains=[], # No exclusions
+ include_answer=False, # We don't need the answer field
+ include_images=False, # We don't need images
+ # Note: 'include_snippets' is not a supported parameter
+ )
+
+ # Check if we got good results (with non-None and non-empty content)
+ if include_raw_content and "results" in results:
+ bad_content_count = sum(
+ 1
+ for r in results["results"]
+ if "raw_content" in r
+ and (
+ r["raw_content"] is None or r["raw_content"].strip() == ""
+ )
+ )
+
+ # If more than half of results have bad content, try a different approach
+ if bad_content_count > len(results["results"]) / 2:
+ logger.warning(
+ f"{bad_content_count}/{len(results['results'])} results have None or empty content. "
+ "Trying to use 'content' field instead of 'raw_content'..."
+ )
+
+ # Try to use the 'content' field which comes by default
+ for result in results["results"]:
+ if (
+ "raw_content" in result
+ and (
+ result["raw_content"] is None
+ or result["raw_content"].strip() == ""
+ )
+ ) and "content" in result:
+ result["raw_content"] = result["content"]
+ logger.info(
+ f"Using 'content' field as 'raw_content' for URL {result.get('url', 'unknown')}"
+ )
+
+ # Re-check after our fix
+ bad_content_count = sum(
+ 1
+ for r in results["results"]
+ if "raw_content" in r
+ and (
+ r["raw_content"] is None
+ or r["raw_content"].strip() == ""
+ )
+ )
+
+ if bad_content_count > 0:
+ logger.warning(
+ f"Still have {bad_content_count}/{len(results['results'])} results with bad content after fixes."
+ )
+
+ # Try alternative approach - search with 'include_answer=True'
+ try:
+ # Search with include_answer=True which may give us better content
+ logger.info(
+ "Trying alternative search with include_answer=True"
+ )
+ alt_results = tavily_client.search(
+ query=query,
+ include_raw_content=include_raw_content,
+ max_results=max_results,
+ search_depth="advanced",
+ include_domains=[],
+ exclude_domains=[],
+ include_answer=True, # Include answer this time
+ include_images=False,
+ )
+
+ # Check if we got any improved content
+ if "results" in alt_results:
+ # Create a merged results set taking the best content
+ for i, result in enumerate(alt_results["results"]):
+ if i < len(results["results"]):
+ if (
+ "raw_content" in result
+ and result["raw_content"]
+ and (
+ results["results"][i].get(
+ "raw_content"
+ )
+ is None
+ or results["results"][i]
+ .get("raw_content", "")
+ .strip()
+ == ""
+ )
+ ):
+ # Replace the bad content with better content from alt_results
+ results["results"][i]["raw_content"] = (
+ result["raw_content"]
+ )
+ logger.info(
+ f"Replaced bad content with better content from alternative search for URL {result.get('url', 'unknown')}"
+ )
+
+ # If answer is available, add it as a special result
+ if "answer" in alt_results and alt_results["answer"]:
+ answer_text = alt_results["answer"]
+ answer_result = {
+ "url": "tavily-generated-answer",
+ "title": "Generated Answer",
+ "raw_content": f"Generated Answer based on search results:\n\n{answer_text}",
+ "content": answer_text,
+ }
+ results["results"].append(answer_result)
+ logger.info(
+ "Added Tavily generated answer as additional search result"
+ )
+
+ except Exception as alt_error:
+ logger.warning(
+ f"Failed to get better results with alternative search: {alt_error}"
+ )
+
+ # Cap content length if specified
+ if cap_content_length > 0 and "results" in results:
+ for result in results["results"]:
+ if "raw_content" in result and result["raw_content"]:
+ result["raw_content"] = result["raw_content"][
+ :cap_content_length
+ ]
+
+ return results
+ except Exception as e:
+ logger.error(f"Error in Tavily search: {e}")
+ # Return an error structure that's compatible with our expected format
+ return {"query": query, "results": [], "error": str(e)}
+
+
+def exa_search(
+ query: str,
+ max_results: int = 3,
+ cap_content_length: int = 20000,
+ search_mode: str = "auto",
+ include_highlights: bool = False,
+) -> Dict[str, Any]:
+ """Perform a search using the Exa API.
+
+ Args:
+ query: Search query
+ max_results: Maximum number of results to return
+ cap_content_length: Maximum length of content to return
+ search_mode: Search mode ("neural", "keyword", or "auto")
+ include_highlights: Whether to include highlights in results
+
+ Returns:
+ Dict[str, Any]: Search results from Exa in a format compatible with Tavily
+ """
+ try:
+ exa_client = get_search_client(SearchProvider.EXA)
+
+ # Configure content options
+ text_options = {"max_characters": cap_content_length}
+
+ kwargs = {
+ "query": query,
+ "num_results": max_results,
+ "type": search_mode, # "neural", "keyword", or "auto"
+ "text": text_options,
+ }
+
+ if include_highlights:
+ kwargs["highlights"] = {
+ "highlights_per_url": 2,
+ "num_sentences": 3,
+ }
+
+ response = exa_client.search_and_contents(**kwargs)
+
+ # Extract cost information
+ exa_cost = 0.0
+ if hasattr(response, "cost_dollars") and hasattr(
+ response.cost_dollars, "total"
+ ):
+ exa_cost = response.cost_dollars.total
+ logger.info(
+ f"Exa search cost for query '{query}': ${exa_cost:.4f}"
+ )
+
+ # Convert to standardized format compatible with Tavily
+ results = {"query": query, "results": [], "exa_cost": exa_cost}
+
+ for r in response.results:
+ result_dict = {
+ "url": r.url,
+ "title": r.title or "",
+ "snippet": "",
+ "raw_content": getattr(r, "text", ""),
+ "content": getattr(r, "text", ""),
+ }
+
+ # Add highlights as snippet if available
+ if hasattr(r, "highlights") and r.highlights:
+ result_dict["snippet"] = " ".join(r.highlights[:1])
+
+ # Store additional metadata
+ result_dict["_metadata"] = {
+ "provider": "exa",
+ "score": getattr(r, "score", None),
+ "published_date": getattr(r, "published_date", None),
+ "author": getattr(r, "author", None),
+ }
+
+ results["results"].append(result_dict)
+
+ return results
+
+ except Exception as e:
+ logger.error(f"Error in Exa search: {e}")
+ return {"query": query, "results": [], "error": str(e)}
+
+
+def unified_search(
+ query: str,
+ provider: Union[str, SearchProvider, None] = None,
+ max_results: int = 3,
+ cap_content_length: int = 20000,
+ search_mode: str = "auto",
+ include_highlights: bool = False,
+ compare_results: bool = False,
+ **kwargs,
+) -> Union[List[SearchResult], Dict[str, List[SearchResult]]]:
+ """Unified search interface supporting multiple providers.
+
+ Args:
+ query: Search query
+ provider: Search provider to use (tavily, exa, both)
+ max_results: Maximum number of results
+ cap_content_length: Maximum content length
+ search_mode: Search mode for Exa ("neural", "keyword", "auto")
+ include_highlights: Include highlights for Exa results
+ compare_results: Return results from both providers separately
+
+ Returns:
+ List[SearchResult] or Dict mapping provider to results (when compare_results=True or provider="both")
+ """
+ # Use default provider if not specified
+ if provider is None:
+ config = SearchEngineConfig()
+ provider = config.default_provider
+
+ # Convert string to enum if needed
+ if isinstance(provider, str):
+ provider = SearchProvider(provider.lower())
+
+ # Handle single provider case
+ if provider == SearchProvider.TAVILY:
+ results = tavily_search(
+ query,
+ max_results=max_results,
+ cap_content_length=cap_content_length,
+ )
+ extracted, cost = extract_search_results(results, provider="tavily")
+ return extracted if not compare_results else {"tavily": extracted}
+
+ elif provider == SearchProvider.EXA:
+ results = exa_search(
+ query=query,
+ max_results=max_results,
+ cap_content_length=cap_content_length,
+ search_mode=search_mode,
+ include_highlights=include_highlights,
+ )
+ extracted, cost = extract_search_results(results, provider="exa")
+ return extracted if not compare_results else {"exa": extracted}
+
+ elif provider == SearchProvider.BOTH:
+ # Run both searches
+ tavily_results = tavily_search(
+ query,
+ max_results=max_results,
+ cap_content_length=cap_content_length,
+ )
+ exa_results = exa_search(
+ query=query,
+ max_results=max_results,
+ cap_content_length=cap_content_length,
+ search_mode=search_mode,
+ include_highlights=include_highlights,
+ )
+
+ # Extract results from both
+ tavily_extracted, tavily_cost = extract_search_results(
+ tavily_results, provider="tavily"
+ )
+ exa_extracted, exa_cost = extract_search_results(
+ exa_results, provider="exa"
+ )
+
+ if compare_results:
+ return {"tavily": tavily_extracted, "exa": exa_extracted}
+ else:
+ # Merge results, interleaving them
+ merged = []
+ max_len = max(len(tavily_extracted), len(exa_extracted))
+ for i in range(max_len):
+ if i < len(tavily_extracted):
+ merged.append(tavily_extracted[i])
+ if i < len(exa_extracted):
+ merged.append(exa_extracted[i])
+ return merged[:max_results] # Limit to requested number
+
+ else:
+ raise ValueError(f"Unknown provider: {provider}")
+
+
+def extract_search_results(
+ search_results: Dict[str, Any], provider: str = "tavily"
+) -> tuple[List[SearchResult], float]:
+ """Extract SearchResult objects from provider-specific API responses.
+
+ Args:
+ search_results: Results from search API
+ provider: Which provider the results came from
+
+ Returns:
+ Tuple of (List[SearchResult], float): List of converted SearchResult objects with standardized fields
+ and the search cost (0.0 if not available).
+ SearchResult is a Pydantic model defined in data_models.py that includes:
+ - url: The URL of the search result
+ - content: The raw content of the page
+ - title: The title of the page
+ - snippet: A brief snippet of the page content
+ """
+ results_list = []
+ search_cost = search_results.get(
+ "exa_cost", 0.0
+ ) # Extract cost if present
+
+ if "results" in search_results:
+ for result in search_results["results"]:
+ if "url" in result:
+ # Get fields with defaults
+ url = result["url"]
+ title = result.get("title", "")
+
+ # Try to extract the best content available:
+ # 1. First try raw_content (if we requested it)
+ # 2. Then try regular content (always available)
+ # 3. Then try to use snippet combined with title
+ # 4. Last resort: use just title
+
+ raw_content = result.get("raw_content", None)
+ regular_content = result.get("content", "")
+ snippet = result.get("snippet", "")
+
+ # Set our final content - prioritize raw_content if available and not None
+ if raw_content is not None and raw_content.strip():
+ content = raw_content
+ # Next best is the regular content field
+ elif regular_content and regular_content.strip():
+ content = regular_content
+ logger.info(
+ f"Using 'content' field for URL {url} because raw_content was not available"
+ )
+ # Try to create a usable content from snippet and title
+ elif snippet:
+ content = f"Title: {title}\n\nContent: {snippet}"
+ logger.warning(
+ f"Using title and snippet as content fallback for {url}"
+ )
+ # Last resort - just use the title
+ elif title:
+ content = (
+ f"Title: {title}\n\nNo content available for this URL."
+ )
+ logger.warning(
+ f"Using only title as content fallback for {url}"
+ )
+ # Nothing available
+ else:
+ content = ""
+ logger.warning(
+ f"No content available for URL {url}, using empty string"
+ )
+
+ # Create SearchResult with provider metadata
+ search_result = SearchResult(
+ url=url,
+ content=content,
+ title=title,
+ snippet=snippet,
+ )
+
+ # Add provider info to metadata if available
+ if "_metadata" in result:
+ search_result.metadata = result["_metadata"]
+ else:
+ search_result.metadata = {"provider": provider}
+
+ results_list.append(search_result)
+
+ # If we got the answer (Tavily specific), add it as a special result
+ if (
+ provider == "tavily"
+ and "answer" in search_results
+ and search_results["answer"]
+ ):
+ answer_text = search_results["answer"]
+ results_list.append(
+ SearchResult(
+ url="tavily-generated-answer",
+ content=f"Generated Answer based on search results:\n\n{answer_text}",
+ title="Tavily Generated Answer",
+ snippet=answer_text[:100] + "..."
+ if len(answer_text) > 100
+ else answer_text,
+ metadata={"provider": "tavily", "type": "generated_answer"},
+ )
+ )
+ logger.info("Added Tavily generated answer as a search result")
+
+ return results_list, search_cost
+
+
+def generate_search_query(
+ sub_question: str,
+ model: str = "sambanova/Llama-4-Maverick-17B-128E-Instruct",
+ system_prompt: Optional[str] = None,
+ project: str = "deep-research",
+) -> Dict[str, Any]:
+ """Generate an optimized search query for a sub-question.
+
+ Uses litellm for model inference via get_structured_llm_output.
+
+ Args:
+ sub_question: The sub-question to generate a search query for
+ model: Model to use (with provider prefix)
+ system_prompt: System prompt for the LLM, defaults to DEFAULT_SEARCH_QUERY_PROMPT
+ project: Langfuse project name for LLM tracking
+
+ Returns:
+ Dictionary with search query and reasoning
+ """
+ if system_prompt is None:
+ system_prompt = DEFAULT_SEARCH_QUERY_PROMPT
+
+ fallback_response = {"search_query": sub_question, "reasoning": ""}
+
+ return get_structured_llm_output(
+ prompt=sub_question,
+ system_prompt=system_prompt,
+ model=model,
+ fallback_response=fallback_response,
+ project=project,
+ )
+
+
+def search_and_extract_results(
+ query: str,
+ max_results: int = 3,
+ cap_content_length: int = 20000,
+ max_retries: int = 2,
+ provider: Optional[Union[str, SearchProvider]] = None,
+ search_mode: str = "auto",
+ include_highlights: bool = False,
+) -> tuple[List[SearchResult], float]:
+ """Perform a search and extract results in one step.
+
+ Args:
+ query: Search query
+ max_results: Maximum number of results to return
+ cap_content_length: Maximum length of content to return
+ max_retries: Maximum number of retries in case of failure
+ provider: Search provider to use (tavily, exa, both)
+ search_mode: Search mode for Exa ("neural", "keyword", "auto")
+ include_highlights: Include highlights for Exa results
+
+ Returns:
+ Tuple of (List of SearchResult objects, search cost)
+ """
+ results = []
+ total_cost = 0.0
+ retry_count = 0
+
+ # List of alternative query formats to try if the original query fails
+ # to yield good results with non-None content
+ alternative_queries = [
+ query, # Original query first
+ f'"{query}"', # Try exact phrase matching
+ f"about {query}", # Try broader context
+ f"research on {query}", # Try research-oriented results
+ query.replace(" OR ", " "), # Try without OR operator
+ ]
+
+ while retry_count <= max_retries and retry_count < len(
+ alternative_queries
+ ):
+ try:
+ current_query = alternative_queries[retry_count]
+ logger.info(
+ f"Searching with query ({retry_count + 1}/{max_retries + 1}): {current_query}"
+ )
+
+ # Determine if we're using Exa to track costs
+ using_exa = False
+ if provider:
+ if isinstance(provider, str):
+ using_exa = provider.lower() in ["exa", "both"]
+ else:
+ using_exa = provider in [
+ SearchProvider.EXA,
+ SearchProvider.BOTH,
+ ]
+ else:
+ config = SearchEngineConfig()
+ using_exa = config.default_provider.lower() in ["exa", "both"]
+
+ # Perform search based on provider
+ if using_exa and provider != SearchProvider.BOTH:
+ # Direct Exa search
+ search_results = exa_search(
+ query=current_query,
+ max_results=max_results,
+ cap_content_length=cap_content_length,
+ search_mode=search_mode,
+ include_highlights=include_highlights,
+ )
+ results, cost = extract_search_results(
+ search_results, provider="exa"
+ )
+ total_cost += cost
+ elif provider == SearchProvider.BOTH:
+ # Search with both providers
+ tavily_results = tavily_search(
+ current_query,
+ max_results=max_results,
+ cap_content_length=cap_content_length,
+ )
+ exa_results = exa_search(
+ query=current_query,
+ max_results=max_results,
+ cap_content_length=cap_content_length,
+ search_mode=search_mode,
+ include_highlights=include_highlights,
+ )
+
+ # Extract results from both
+ tavily_extracted, _ = extract_search_results(
+ tavily_results, provider="tavily"
+ )
+ exa_extracted, exa_cost = extract_search_results(
+ exa_results, provider="exa"
+ )
+ total_cost += exa_cost
+
+ # Merge results
+ results = []
+ max_len = max(len(tavily_extracted), len(exa_extracted))
+ for i in range(max_len):
+ if i < len(tavily_extracted):
+ results.append(tavily_extracted[i])
+ if i < len(exa_extracted):
+ results.append(exa_extracted[i])
+ results = results[:max_results]
+ else:
+ # Tavily search or unified search
+ results = unified_search(
+ query=current_query,
+ provider=provider,
+ max_results=max_results,
+ cap_content_length=cap_content_length,
+ search_mode=search_mode,
+ include_highlights=include_highlights,
+ )
+
+ # Handle case where unified_search returns a dict
+ if isinstance(results, dict):
+ all_results = []
+ for provider_results in results.values():
+ all_results.extend(provider_results)
+ results = all_results[:max_results]
+
+ # Check if we got results with actual content
+ if results:
+ # Count results with non-empty content
+ content_results = sum(1 for r in results if r.content.strip())
+
+ if content_results >= max(1, len(results) // 2):
+ logger.info(
+ f"Found {content_results}/{len(results)} results with content"
+ )
+ return results, total_cost
+ else:
+ logger.warning(
+ f"Only found {content_results}/{len(results)} results with content. "
+ f"Trying alternative query..."
+ )
+
+ # If we didn't get good results but haven't hit max retries yet, try again
+ if retry_count < max_retries:
+ logger.warning(
+ f"Inadequate search results. Retrying with alternative query... ({retry_count + 1}/{max_retries})"
+ )
+ retry_count += 1
+ else:
+ # If we're out of retries, return whatever we have
+ logger.warning(
+ f"Out of retries. Returning best results found ({len(results)} results)."
+ )
+ return results, total_cost
+
+ except Exception as e:
+ if retry_count < max_retries:
+ logger.warning(
+ f"Search failed with error: {e}. Retrying... ({retry_count + 1}/{max_retries})"
+ )
+ retry_count += 1
+ else:
+ logger.error(f"Search failed after {max_retries} retries: {e}")
+ return [], 0.0
+
+ # If we've exhausted all retries, return the best results we have
+ return results, total_cost
diff --git a/deep_research/utils/tracing_metadata_utils.py b/deep_research/utils/tracing_metadata_utils.py
new file mode 100644
index 00000000..59c7b37e
--- /dev/null
+++ b/deep_research/utils/tracing_metadata_utils.py
@@ -0,0 +1,745 @@
+"""Utilities for collecting and analyzing tracing metadata from Langfuse."""
+
+import time
+from datetime import datetime, timedelta, timezone
+from functools import wraps
+from typing import Any, Dict, List, Optional, Tuple
+
+from langfuse import Langfuse
+from langfuse.api.core import ApiError
+from langfuse.client import ObservationsView, TraceWithDetails
+from rich import print
+from rich.console import Console
+from rich.table import Table
+
+console = Console()
+
+langfuse = Langfuse()
+
+# Prompt type identification keywords
+PROMPT_IDENTIFIERS = {
+ "query_decomposition": [
+ "MAIN RESEARCH QUERY",
+ "DIFFERENT DIMENSIONS",
+ "sub-questions",
+ ],
+ "search_query": ["Deep Research assistant", "effective search query"],
+ "synthesis": [
+ "information synthesis",
+ "comprehensive answer",
+ "confidence level",
+ ],
+ "viewpoint_analysis": [
+ "multi-perspective analysis",
+ "viewpoint categories",
+ ],
+ "reflection": ["critique and improve", "information gaps"],
+ "additional_synthesis": ["enhance the original synthesis"],
+ "conclusion_generation": [
+ "Synthesis and Integration",
+ "Direct Response to Main Query",
+ ],
+ "executive_summary": [
+ "executive summaries",
+ "Key Findings",
+ "250-400 words",
+ ],
+ "introduction": ["engaging introductions", "Context and Relevance"],
+}
+
+# Rate limiting configuration
+# Adjust these based on your Langfuse tier:
+# - Hobby: 30 req/min for Other APIs -> ~2s between requests
+# - Core: 100 req/min -> ~0.6s between requests
+# - Pro: 1000 req/min -> ~0.06s between requests
+RATE_LIMIT_DELAY = 0.1 # 100ms between requests (safe for most tiers)
+MAX_RETRIES = 3
+INITIAL_BACKOFF = 1.0 # Initial backoff in seconds
+
+# Batch processing configuration
+BATCH_DELAY = 0.5 # Additional delay between batches of requests
+
+
+def rate_limited(func):
+ """Decorator to add rate limiting between API calls."""
+
+ @wraps(func)
+ def wrapper(*args, **kwargs):
+ time.sleep(RATE_LIMIT_DELAY)
+ return func(*args, **kwargs)
+
+ return wrapper
+
+
+def retry_with_backoff(func):
+ """Decorator to retry functions with exponential backoff on rate limit errors."""
+
+ @wraps(func)
+ def wrapper(*args, **kwargs):
+ backoff = INITIAL_BACKOFF
+ last_exception = None
+
+ for attempt in range(MAX_RETRIES):
+ try:
+ return func(*args, **kwargs)
+ except ApiError as e:
+ if e.status_code == 429: # Rate limit error
+ last_exception = e
+ if attempt < MAX_RETRIES - 1:
+ wait_time = backoff * (2**attempt)
+ console.print(
+ f"[yellow]Rate limit hit. Retrying in {wait_time:.1f}s...[/yellow]"
+ )
+ time.sleep(wait_time)
+ continue
+ raise
+ except Exception:
+ # For non-rate limit errors, raise immediately
+ raise
+
+ # If we've exhausted all retries
+ if last_exception:
+ raise last_exception
+
+ return wrapper
+
+
+@rate_limited
+@retry_with_backoff
+def fetch_traces_safe(limit: Optional[int] = None) -> List[TraceWithDetails]:
+ """Safely fetch traces with rate limiting and retry logic."""
+ return langfuse.fetch_traces(limit=limit).data
+
+
+@rate_limited
+@retry_with_backoff
+def fetch_observations_safe(trace_id: str) -> List[ObservationsView]:
+ """Safely fetch observations with rate limiting and retry logic."""
+ return langfuse.fetch_observations(trace_id=trace_id).data
+
+
+def get_total_trace_cost(trace_id: str) -> float:
+ """Calculate the total cost for a single trace by summing all observation costs.
+
+ Args:
+ trace_id: The ID of the trace to calculate cost for
+
+ Returns:
+ Total cost across all observations in the trace
+ """
+ try:
+ observations = fetch_observations_safe(trace_id=trace_id)
+ total_cost = 0.0
+
+ for obs in observations:
+ # Check multiple possible cost fields
+ if (
+ hasattr(obs, "calculated_total_cost")
+ and obs.calculated_total_cost
+ ):
+ total_cost += obs.calculated_total_cost
+ elif hasattr(obs, "total_price") and obs.total_price:
+ total_cost += obs.total_price
+ elif hasattr(obs, "total_cost") and obs.total_cost:
+ total_cost += obs.total_cost
+ # If cost details are available, calculate from input/output costs
+ elif hasattr(obs, "calculated_input_cost") and hasattr(
+ obs, "calculated_output_cost"
+ ):
+ if obs.calculated_input_cost and obs.calculated_output_cost:
+ total_cost += (
+ obs.calculated_input_cost + obs.calculated_output_cost
+ )
+
+ return total_cost
+ except Exception as e:
+ print(f"[red]Error calculating trace cost: {e}[/red]")
+ return 0.0
+
+
+def get_total_tokens_used(trace_id: str) -> Tuple[int, int]:
+ """Calculate total input and output tokens used for a trace.
+
+ Args:
+ trace_id: The ID of the trace to calculate tokens for
+
+ Returns:
+ Tuple of (input_tokens, output_tokens)
+ """
+ try:
+ observations = fetch_observations_safe(trace_id=trace_id)
+ total_input_tokens = 0
+ total_output_tokens = 0
+
+ for obs in observations:
+ # Check for token fields in different possible locations
+ if hasattr(obs, "usage") and obs.usage:
+ if hasattr(obs.usage, "input") and obs.usage.input:
+ total_input_tokens += obs.usage.input
+ if hasattr(obs.usage, "output") and obs.usage.output:
+ total_output_tokens += obs.usage.output
+ # Also check for direct token fields
+ elif hasattr(obs, "promptTokens") and hasattr(
+ obs, "completionTokens"
+ ):
+ if obs.promptTokens:
+ total_input_tokens += obs.promptTokens
+ if obs.completionTokens:
+ total_output_tokens += obs.completionTokens
+
+ return total_input_tokens, total_output_tokens
+ except Exception as e:
+ print(f"[red]Error calculating tokens: {e}[/red]")
+ return 0, 0
+
+
+def get_trace_stats(trace: TraceWithDetails) -> Dict[str, Any]:
+ """Get comprehensive statistics for a trace.
+
+ Args:
+ trace: The trace object to analyze
+
+ Returns:
+ Dictionary containing trace statistics including cost, latency, tokens, and metadata
+ """
+ try:
+ # Get cost and token data
+ total_cost = get_total_trace_cost(trace.id)
+ input_tokens, output_tokens = get_total_tokens_used(trace.id)
+
+ # Get observation count
+ observations = fetch_observations_safe(trace_id=trace.id)
+ observation_count = len(observations)
+
+ # Extract model information from observations
+ models_used = set()
+ for obs in observations:
+ if hasattr(obs, "model") and obs.model:
+ models_used.add(obs.model)
+
+ stats = {
+ "trace_id": trace.id,
+ "timestamp": trace.timestamp,
+ "total_cost": total_cost,
+ "latency_seconds": trace.latency
+ if hasattr(trace, "latency")
+ else 0,
+ "input_tokens": input_tokens,
+ "output_tokens": output_tokens,
+ "total_tokens": input_tokens + output_tokens,
+ "observation_count": observation_count,
+ "models_used": list(models_used),
+ "metadata": trace.metadata if hasattr(trace, "metadata") else {},
+ "tags": trace.tags if hasattr(trace, "tags") else [],
+ "user_id": trace.user_id if hasattr(trace, "user_id") else None,
+ "session_id": trace.session_id
+ if hasattr(trace, "session_id")
+ else None,
+ }
+
+ # Add formatted latency
+ if stats["latency_seconds"]:
+ minutes = int(stats["latency_seconds"] // 60)
+ seconds = stats["latency_seconds"] % 60
+ stats["latency_formatted"] = f"{minutes}m {seconds:.1f}s"
+ else:
+ stats["latency_formatted"] = "0m 0.0s"
+
+ return stats
+ except Exception as e:
+ print(f"[red]Error getting trace stats: {e}[/red]")
+ return {}
+
+
+def get_traces_by_name(name: str, limit: int = 1) -> List[TraceWithDetails]:
+ """Get traces by name using Langfuse API.
+
+ Args:
+ name: The name of the trace to search for
+ limit: Maximum number of traces to return (default: 1)
+
+ Returns:
+ List of traces matching the name
+ """
+ try:
+ # Use the Langfuse API to get traces by name
+ traces_response = langfuse.get_traces(name=name, limit=limit)
+ return traces_response.data
+ except Exception as e:
+ print(f"[red]Error fetching traces by name: {e}[/red]")
+ return []
+
+
+def get_observations_for_trace(trace_id: str) -> List[ObservationsView]:
+ """Get all observations for a specific trace.
+
+ Args:
+ trace_id: The ID of the trace
+
+ Returns:
+ List of observations for the trace
+ """
+ try:
+ observations_response = langfuse.get_observations(trace_id=trace_id)
+ return observations_response.data
+ except Exception as e:
+ print(f"[red]Error fetching observations: {e}[/red]")
+ return []
+
+
+def filter_traces_by_date_range(
+ start_date: datetime, end_date: datetime, limit: Optional[int] = None
+) -> List[TraceWithDetails]:
+ """Filter traces within a specific date range.
+
+ Args:
+ start_date: Start of the date range (inclusive)
+ end_date: End of the date range (inclusive)
+ limit: Maximum number of traces to return
+
+ Returns:
+ List of traces within the date range
+ """
+ try:
+ # Ensure dates are timezone-aware
+ if start_date.tzinfo is None:
+ start_date = start_date.replace(tzinfo=timezone.utc)
+ if end_date.tzinfo is None:
+ end_date = end_date.replace(tzinfo=timezone.utc)
+
+ # Fetch all traces (or up to API maximum limit of 100)
+ all_traces = fetch_traces_safe(limit=limit or 100)
+
+ # Filter by date range
+ filtered_traces = [
+ trace
+ for trace in all_traces
+ if start_date <= trace.timestamp <= end_date
+ ]
+
+ # Sort by timestamp (most recent first)
+ filtered_traces.sort(key=lambda x: x.timestamp, reverse=True)
+
+ # Apply limit if specified
+ if limit:
+ filtered_traces = filtered_traces[:limit]
+
+ return filtered_traces
+ except Exception as e:
+ print(f"[red]Error filtering traces by date range: {e}[/red]")
+ return []
+
+
+def get_traces_last_n_days(
+ days: int, limit: Optional[int] = None
+) -> List[TraceWithDetails]:
+ """Get traces from the last N days.
+
+ Args:
+ days: Number of days to look back
+ limit: Maximum number of traces to return
+
+ Returns:
+ List of traces from the last N days
+ """
+ end_date = datetime.now(timezone.utc)
+ start_date = end_date - timedelta(days=days)
+
+ return filter_traces_by_date_range(start_date, end_date, limit)
+
+
+def get_trace_stats_batch(
+ traces: List[TraceWithDetails], show_progress: bool = True
+) -> List[Dict[str, Any]]:
+ """Get statistics for multiple traces efficiently with progress tracking.
+
+ Args:
+ traces: List of traces to analyze
+ show_progress: Whether to show progress bar
+
+ Returns:
+ List of dictionaries containing trace statistics
+ """
+ stats_list = []
+
+ for i, trace in enumerate(traces):
+ if show_progress and i % 5 == 0:
+ console.print(
+ f"[dim]Processing trace {i + 1}/{len(traces)}...[/dim]"
+ )
+
+ stats = get_trace_stats(trace)
+ stats_list.append(stats)
+
+ return stats_list
+
+
+def get_aggregate_stats_for_traces(
+ traces: List[TraceWithDetails],
+) -> Dict[str, Any]:
+ """Calculate aggregate statistics for a list of traces.
+
+ Args:
+ traces: List of traces to analyze
+
+ Returns:
+ Dictionary containing aggregate statistics
+ """
+ if not traces:
+ return {
+ "trace_count": 0,
+ "total_cost": 0.0,
+ "total_input_tokens": 0,
+ "total_output_tokens": 0,
+ "total_tokens": 0,
+ "average_cost_per_trace": 0.0,
+ "average_latency_seconds": 0.0,
+ "total_observations": 0,
+ }
+
+ total_cost = 0.0
+ total_input_tokens = 0
+ total_output_tokens = 0
+ total_latency = 0.0
+ total_observations = 0
+ all_models = set()
+
+ for trace in traces:
+ stats = get_trace_stats(trace)
+ total_cost += stats.get("total_cost", 0)
+ total_input_tokens += stats.get("input_tokens", 0)
+ total_output_tokens += stats.get("output_tokens", 0)
+ total_latency += stats.get("latency_seconds", 0)
+ total_observations += stats.get("observation_count", 0)
+ all_models.update(stats.get("models_used", []))
+
+ return {
+ "trace_count": len(traces),
+ "total_cost": total_cost,
+ "total_input_tokens": total_input_tokens,
+ "total_output_tokens": total_output_tokens,
+ "total_tokens": total_input_tokens + total_output_tokens,
+ "average_cost_per_trace": total_cost / len(traces) if traces else 0,
+ "average_latency_seconds": total_latency / len(traces)
+ if traces
+ else 0,
+ "total_observations": total_observations,
+ "models_used": list(all_models),
+ }
+
+
+def display_trace_stats_table(
+ traces: List[TraceWithDetails], title: str = "Trace Statistics"
+):
+ """Display trace statistics in a formatted table.
+
+ Args:
+ traces: List of traces to display
+ title: Title for the table
+ """
+ table = Table(title=title, show_header=True, header_style="bold magenta")
+ table.add_column("Trace ID", style="cyan", no_wrap=True)
+ table.add_column("Timestamp", style="yellow")
+ table.add_column("Cost ($)", justify="right", style="green")
+ table.add_column("Tokens (In/Out)", justify="right")
+ table.add_column("Latency", justify="right")
+ table.add_column("Observations", justify="right")
+
+ for trace in traces[:10]: # Limit to 10 for display
+ stats = get_trace_stats(trace)
+ table.add_row(
+ stats["trace_id"][:12] + "...",
+ stats["timestamp"].strftime("%Y-%m-%d %H:%M"),
+ f"${stats['total_cost']:.4f}",
+ f"{stats['input_tokens']:,}/{stats['output_tokens']:,}",
+ stats["latency_formatted"],
+ str(stats["observation_count"]),
+ )
+
+ console.print(table)
+
+
+def identify_prompt_type(observation: ObservationsView) -> str:
+ """Identify the prompt type based on keywords in the observation's input.
+
+ Examines the system prompt in observation.input['messages'][0]['content']
+ for unique keywords that identify each prompt type.
+
+ Args:
+ observation: The observation to analyze
+
+ Returns:
+ str: The prompt type name, or "unknown" if not identified
+ """
+ try:
+ # Access the system prompt from the messages
+ if hasattr(observation, "input") and observation.input:
+ messages = observation.input.get("messages", [])
+ if messages and len(messages) > 0:
+ system_content = messages[0].get("content", "")
+
+ # Check each prompt type's keywords
+ for prompt_type, keywords in PROMPT_IDENTIFIERS.items():
+ # Check if any keyword is in the system prompt
+ for keyword in keywords:
+ if keyword in system_content:
+ return prompt_type
+
+ return "unknown"
+ except Exception as e:
+ console.print(
+ f"[yellow]Warning: Could not identify prompt type: {e}[/yellow]"
+ )
+ return "unknown"
+
+
+def get_costs_by_prompt_type(trace_id: str) -> Dict[str, Dict[str, float]]:
+ """Get cost breakdown by prompt type for a given trace.
+
+ Uses observation.usage.input/output for token counts and
+ observation.calculated_total_cost for costs.
+
+ Args:
+ trace_id: The ID of the trace to analyze
+
+ Returns:
+ Dict mapping prompt_type to {
+ 'cost': float,
+ 'input_tokens': int,
+ 'output_tokens': int,
+ 'count': int # number of calls
+ }
+ """
+ try:
+ observations = fetch_observations_safe(trace_id=trace_id)
+ prompt_metrics = {}
+
+ for obs in observations:
+ # Identify prompt type
+ prompt_type = identify_prompt_type(obs)
+
+ # Initialize metrics for this prompt type if needed
+ if prompt_type not in prompt_metrics:
+ prompt_metrics[prompt_type] = {
+ "cost": 0.0,
+ "input_tokens": 0,
+ "output_tokens": 0,
+ "count": 0,
+ }
+
+ # Add cost
+ cost = 0.0
+ if (
+ hasattr(obs, "calculated_total_cost")
+ and obs.calculated_total_cost
+ ):
+ cost = obs.calculated_total_cost
+ prompt_metrics[prompt_type]["cost"] += cost
+
+ # Add tokens
+ if hasattr(obs, "usage") and obs.usage:
+ if hasattr(obs.usage, "input") and obs.usage.input:
+ prompt_metrics[prompt_type]["input_tokens"] += (
+ obs.usage.input
+ )
+ if hasattr(obs.usage, "output") and obs.usage.output:
+ prompt_metrics[prompt_type]["output_tokens"] += (
+ obs.usage.output
+ )
+
+ # Increment count
+ prompt_metrics[prompt_type]["count"] += 1
+
+ return prompt_metrics
+ except Exception as e:
+ print(f"[red]Error getting costs by prompt type: {e}[/red]")
+ return {}
+
+
+def get_prompt_type_statistics(trace_id: str) -> Dict[str, Dict[str, Any]]:
+ """Get detailed statistics for each prompt type.
+
+ Args:
+ trace_id: The ID of the trace to analyze
+
+ Returns:
+ Dict mapping prompt_type to {
+ 'cost': float,
+ 'input_tokens': int,
+ 'output_tokens': int,
+ 'count': int,
+ 'avg_cost_per_call': float,
+ 'avg_input_tokens': float,
+ 'avg_output_tokens': float,
+ 'percentage_of_total_cost': float
+ }
+ """
+ try:
+ # Get basic metrics
+ prompt_metrics = get_costs_by_prompt_type(trace_id)
+
+ # Calculate total cost for percentage calculation
+ total_cost = sum(
+ metrics["cost"] for metrics in prompt_metrics.values()
+ )
+
+ # Enhance with statistics
+ enhanced_metrics = {}
+ for prompt_type, metrics in prompt_metrics.items():
+ count = metrics["count"]
+ enhanced_metrics[prompt_type] = {
+ "cost": metrics["cost"],
+ "input_tokens": metrics["input_tokens"],
+ "output_tokens": metrics["output_tokens"],
+ "count": count,
+ "avg_cost_per_call": metrics["cost"] / count
+ if count > 0
+ else 0,
+ "avg_input_tokens": metrics["input_tokens"] / count
+ if count > 0
+ else 0,
+ "avg_output_tokens": metrics["output_tokens"] / count
+ if count > 0
+ else 0,
+ "percentage_of_total_cost": (
+ metrics["cost"] / total_cost * 100
+ )
+ if total_cost > 0
+ else 0,
+ }
+
+ return enhanced_metrics
+ except Exception as e:
+ print(f"[red]Error getting prompt type statistics: {e}[/red]")
+ return {}
+
+
+if __name__ == "__main__":
+ print(
+ "[bold cyan]ZenML Deep Research - Tracing Metadata Utilities Demo[/bold cyan]\n"
+ )
+
+ try:
+ # Fetch recent traces
+ print("[yellow]Fetching recent traces...[/yellow]")
+ traces = fetch_traces_safe(limit=5)
+
+ if not traces:
+ print("[red]No traces found![/red]")
+ exit(1)
+ except ApiError as e:
+ if e.status_code == 429:
+ print("[red]Rate limit exceeded. Please try again later.[/red]")
+ print(
+ "[yellow]Tip: Consider upgrading your Langfuse tier for higher rate limits.[/yellow]"
+ )
+ else:
+ print(f"[red]API Error: {e}[/red]")
+ exit(1)
+ except Exception as e:
+ print(f"[red]Error fetching traces: {e}[/red]")
+ exit(1)
+
+ # Demo 1: Get stats for a single trace
+ print("\n[bold]1. Single Trace Statistics:[/bold]")
+ first_trace = traces[0]
+ stats = get_trace_stats(first_trace)
+
+ console.print(f"Trace ID: [cyan]{stats['trace_id']}[/cyan]")
+ console.print(f"Timestamp: [yellow]{stats['timestamp']}[/yellow]")
+ console.print(f"Total Cost: [green]${stats['total_cost']:.4f}[/green]")
+ console.print(
+ f"Tokens - Input: [blue]{stats['input_tokens']:,}[/blue], Output: [blue]{stats['output_tokens']:,}[/blue]"
+ )
+ console.print(f"Latency: [magenta]{stats['latency_formatted']}[/magenta]")
+ console.print(f"Observations: [white]{stats['observation_count']}[/white]")
+ console.print(
+ f"Models Used: [cyan]{', '.join(stats['models_used'])}[/cyan]"
+ )
+
+ # Demo 2: Get traces from last 7 days
+ print("\n[bold]2. Traces from Last 7 Days:[/bold]")
+ recent_traces = get_traces_last_n_days(7, limit=10)
+ print(
+ f"Found [green]{len(recent_traces)}[/green] traces in the last 7 days"
+ )
+
+ if recent_traces:
+ display_trace_stats_table(recent_traces, "Last 7 Days Traces")
+
+ # Demo 3: Filter traces by date range
+ print("\n[bold]3. Filter Traces by Date Range:[/bold]")
+ end_date = datetime.now(timezone.utc)
+ start_date = end_date - timedelta(days=3)
+
+ filtered_traces = filter_traces_by_date_range(start_date, end_date)
+ print(
+ f"Found [green]{len(filtered_traces)}[/green] traces between {start_date.strftime('%Y-%m-%d')} and {end_date.strftime('%Y-%m-%d')}"
+ )
+
+ # Demo 4: Aggregate statistics
+ print("\n[bold]4. Aggregate Statistics for All Recent Traces:[/bold]")
+ agg_stats = get_aggregate_stats_for_traces(traces)
+
+ table = Table(
+ title="Aggregate Statistics",
+ show_header=True,
+ header_style="bold magenta",
+ )
+ table.add_column("Metric", style="cyan")
+ table.add_column("Value", justify="right", style="yellow")
+
+ table.add_row("Total Traces", str(agg_stats["trace_count"]))
+ table.add_row("Total Cost", f"${agg_stats['total_cost']:.4f}")
+ table.add_row(
+ "Average Cost per Trace", f"${agg_stats['average_cost_per_trace']:.4f}"
+ )
+ table.add_row("Total Input Tokens", f"{agg_stats['total_input_tokens']:,}")
+ table.add_row(
+ "Total Output Tokens", f"{agg_stats['total_output_tokens']:,}"
+ )
+ table.add_row("Total Tokens", f"{agg_stats['total_tokens']:,}")
+ table.add_row(
+ "Average Latency", f"{agg_stats['average_latency_seconds']:.1f}s"
+ )
+ table.add_row("Total Observations", str(agg_stats["total_observations"]))
+
+ console.print(table)
+
+ # Demo 5: Cost breakdown by observation
+ print("\n[bold]5. Cost Breakdown for First Trace:[/bold]")
+ observations = fetch_observations_safe(trace_id=first_trace.id)
+
+ if observations:
+ table = Table(
+ title="Observation Cost Breakdown",
+ show_header=True,
+ header_style="bold magenta",
+ )
+ table.add_column("Observation", style="cyan", no_wrap=True)
+ table.add_column("Model", style="yellow")
+ table.add_column("Tokens (In/Out)", justify="right")
+ table.add_column("Cost", justify="right", style="green")
+
+ for i, obs in enumerate(observations[:5]): # Show first 5
+ cost = 0.0
+ if hasattr(obs, "calculated_total_cost"):
+ cost = obs.calculated_total_cost or 0.0
+
+ in_tokens = 0
+ out_tokens = 0
+ if hasattr(obs, "usage") and obs.usage:
+ in_tokens = obs.usage.input or 0
+ out_tokens = obs.usage.output or 0
+ elif hasattr(obs, "promptTokens"):
+ in_tokens = obs.promptTokens or 0
+ out_tokens = obs.completionTokens or 0
+
+ table.add_row(
+ f"Obs {i + 1}",
+ obs.model if hasattr(obs, "model") else "Unknown",
+ f"{in_tokens:,}/{out_tokens:,}",
+ f"${cost:.4f}",
+ )
+
+ console.print(table)