researchflow/schema.py at main · beingadish/researchflow · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
"""
Pydantic Schema Definitions for AI Research Agent System

This module defines the structured data models used by the AI research agent system
for processing questions, generating answers, and managing the revision workflow.
The schemas ensure consistent output formatting and enable structured interaction
between different components of the research pipeline.

Schema Hierarchy:
- Reflection: Base model for answer critique and improvement suggestions
- AnswerQuestion: Primary model for initial question answering with research queries
- ReviseAnswer: Extended model for answer revision with citation support

Key Features:
- Field validation and type safety through Pydantic
- Structured critique system for iterative improvement
- Search query generation for research enhancement
- Citation management for answer verification
- Inheritance-based design for workflow progression

Usage Context:
These schemas are used by LangChain tool parsers to structure AI model outputs,
ensuring consistent data formats across the first responder and revisor chains.

Author: Aadarsh Pandey
Date: 28 Sep 2025
"""

from pydantic import BaseModel, Field
from typing import List


class Reflection(BaseModel):
    """
    Model for structured self-critique and answer improvement analysis.

    This class provides a framework for AI agents to critically evaluate their
    initial responses by identifying missing information and unnecessary content.
    It supports the iterative improvement process by highlighting specific areas
    that need enhancement or reduction.

    Use Cases:
    - Self-assessment of answer completeness and accuracy
    - Identification of content gaps requiring additional research
    - Detection of irrelevant or excessive information
    - Structured feedback for answer refinement

    Attributes:
        missing (str): Critical analysis of information gaps, incomplete coverage,
                      or important points that should be included in the answer
        superfluous (str): Analysis of excessive, irrelevant, or redundant content
                          that should be removed or condensed for clarity
    """
    missing: str = Field(description="Critique of what is missing.")
    superfluous: str = Field(description="Critique of what is superfluous")


class AnswerQuestion(BaseModel):
    """
    Primary schema for initial question answering with research capabilities.

    This model structures the complete response workflow for the first responder
    agent, combining answer generation, self-reflection, and research query
    planning in a single structured output. It enables the system to provide
    immediate answers while identifying areas for improvement through search.

    Workflow Integration:
    1. Generate initial answer (250 words target)
    2. Perform self-reflection on answer quality
    3. Generate search queries for improvement research
    4. Pass to tool execution for query processing

    Attributes:
        answer (str): Detailed response to the user's question, targeting approximately
                     250 words with comprehensive coverage of key points
        search_queries (List[str]): 1-3 focused search queries designed to address
                                   identified gaps and improve answer quality
        reflection (Reflection): Structured self-critique identifying missing
                               information and superfluous content
    """
    answer: str = Field(description="~250 words detailed answer to the question.")
    search_queries: List[str] = Field(description="1-3 search queries for researching improvements to address the critique of the current answer")
    reflection: Reflection = Field(description="Your reflection on the initial answer.")


class ReviseAnswer(AnswerQuestion):
    """
    Extended schema for answer revision with citation and reference management.

    This model inherits all capabilities from AnswerQuestion while adding
    citation support for the revision workflow. It enables the revisor agent
    to refine answers using search results while maintaining verifiable
    references and adhering to quality standards.

    Inheritance Benefits:
    - Maintains all original fields (answer, search_queries, reflection)
    - Adds citation capabilities for enhanced credibility
    - Supports iterative improvement workflow
    - Enables tool choice differentiation in LangChain

    Revision Requirements:
    - Integrate search results from previous queries
    - Add numerical citations for verification
    - Include references section with URLs
    - Maintain 250-word limit excluding references
    - Apply reflection critique for quality improvement

    Attributes:
        references (List[str]): Structured list of citations supporting the revised
                               answer, typically formatted as numbered references
                               with URLs for verification and credibility

    Inherited Attributes:
        answer (str): Refined version of the original answer incorporating
                     search results and addressing identified gaps
        search_queries (List[str]): Additional queries if further research needed
        reflection (Reflection): Updated critique of the revised answer
    """
    references: List[str] = Field(description="Citations motivating your updated answer.")