95 lines
3.2 KiB
Python
95 lines
3.2 KiB
Python
from pydantic_ai import format_as_xml
|
|
|
|
from haiku.rag.graph.research.dependencies import ResearchContext
|
|
from haiku.rag.graph.research.models import InsightAnalysis
|
|
|
|
|
|
def format_context_for_prompt(context: ResearchContext) -> str:
|
|
"""Format the research context as XML for inclusion in prompts."""
|
|
|
|
context_data = {
|
|
"original_question": context.original_question,
|
|
"unanswered_questions": context.sub_questions,
|
|
"qa_responses": [
|
|
{
|
|
"question": qa.query,
|
|
"answer": qa.answer,
|
|
"confidence": qa.confidence,
|
|
"sources": [
|
|
{
|
|
"document_uri": c.document_uri,
|
|
"document_title": c.document_title,
|
|
"page_numbers": c.page_numbers,
|
|
"headings": c.headings,
|
|
}
|
|
for c in qa.citations
|
|
],
|
|
}
|
|
for qa in context.qa_responses
|
|
],
|
|
"insights": [
|
|
{
|
|
"id": insight.id,
|
|
"summary": insight.summary,
|
|
"status": insight.status.value,
|
|
"supporting_sources": insight.supporting_sources,
|
|
"originating_questions": insight.originating_questions,
|
|
"notes": insight.notes,
|
|
}
|
|
for insight in context.insights
|
|
],
|
|
"gaps": [
|
|
{
|
|
"id": gap.id,
|
|
"description": gap.description,
|
|
"severity": gap.severity.value,
|
|
"blocking": gap.blocking,
|
|
"resolved": gap.resolved,
|
|
"resolved_by": gap.resolved_by,
|
|
"supporting_sources": gap.supporting_sources,
|
|
"notes": gap.notes,
|
|
}
|
|
for gap in context.gaps
|
|
],
|
|
}
|
|
return format_as_xml(context_data, root_tag="research_context")
|
|
|
|
|
|
def format_analysis_for_prompt(
|
|
analysis: InsightAnalysis | None,
|
|
) -> str:
|
|
"""Format the latest insight analysis as XML for prompts."""
|
|
|
|
if analysis is None:
|
|
return "<latest_analysis />"
|
|
|
|
data = {
|
|
"commentary": analysis.commentary,
|
|
"highlights": [
|
|
{
|
|
"id": insight.id,
|
|
"summary": insight.summary,
|
|
"status": insight.status.value,
|
|
"supporting_sources": insight.supporting_sources,
|
|
"originating_questions": insight.originating_questions,
|
|
"notes": insight.notes,
|
|
}
|
|
for insight in analysis.highlights
|
|
],
|
|
"gap_assessments": [
|
|
{
|
|
"id": gap.id,
|
|
"description": gap.description,
|
|
"severity": gap.severity.value,
|
|
"blocking": gap.blocking,
|
|
"resolved": gap.resolved,
|
|
"resolved_by": gap.resolved_by,
|
|
"supporting_sources": gap.supporting_sources,
|
|
"notes": gap.notes,
|
|
}
|
|
for gap in analysis.gap_assessments
|
|
],
|
|
"resolved_gaps": analysis.resolved_gaps,
|
|
"new_questions": analysis.new_questions,
|
|
}
|
|
return format_as_xml(data, root_tag="latest_analysis")
|