|
19 | 19 | """ |
20 | 20 |
|
21 | 21 | from typing import Any, Dict, List, Optional, Union, Literal, Annotated |
22 | | -from pydantic import BaseModel, ConfigDict, Field, field_validator, Discriminator |
| 22 | +from typing_extensions import Self |
| 23 | +from pydantic import ( |
| 24 | + BaseModel, |
| 25 | + ConfigDict, |
| 26 | + Field, |
| 27 | + field_validator, |
| 28 | + Discriminator, |
| 29 | + model_validator, |
| 30 | +) |
23 | 31 |
|
24 | 32 |
|
25 | 33 | class ImageConfig(BaseModel): |
@@ -78,6 +86,10 @@ class AgenticConfig(BaseModel): |
78 | 86 |
|
79 | 87 | enabled: bool = Field(default=False, description="Enable agentic extraction") |
80 | 88 | review_agent: bool = Field(default=False, description="Enable review agent") |
| 89 | + review_agent_model: str | None = Field( |
| 90 | + default=None, |
| 91 | + description="Model used for reviewing and correcting extraction work", |
| 92 | + ) |
81 | 93 |
|
82 | 94 |
|
83 | 95 | class ExtractionConfig(BaseModel): |
@@ -120,6 +132,14 @@ def parse_int(cls, v: Any) -> int: |
120 | 132 | return int(v) if v else 0 |
121 | 133 | return int(v) |
122 | 134 |
|
| 135 | + @model_validator(mode="after") |
| 136 | + def set_default_review_agent_model(self) -> Self: |
| 137 | + """Set review_agent_model to extraction model if not specified.""" |
| 138 | + if not self.agentic.review_agent_model: |
| 139 | + self.agentic.review_agent_model = self.model |
| 140 | + |
| 141 | + return self |
| 142 | + |
123 | 143 |
|
124 | 144 | class ClassificationConfig(BaseModel): |
125 | 145 | """Document classification configuration""" |
@@ -423,7 +443,7 @@ class ErrorAnalyzerConfig(BaseModel): |
423 | 443 | "AccessDenied", |
424 | 444 | "ThrottlingException", |
425 | 445 | ], |
426 | | - description="Error patterns to search for in logs" |
| 446 | + description="Error patterns to search for in logs", |
427 | 447 | ) |
428 | 448 | system_prompt: str = Field( |
429 | 449 | default=""" |
@@ -511,11 +531,10 @@ class ErrorAnalyzerConfig(BaseModel): |
511 | 531 | - No time specified: 24 hours (default) |
512 | 532 | |
513 | 533 | IMPORTANT: Do not include any search quality reflections, search quality scores, or meta-analysis sections in your response. Only provide the three required sections: Root Cause, Recommendations, and Evidence.""", |
514 | | - description="System prompt for error analyzer" |
| 534 | + description="System prompt for error analyzer", |
515 | 535 | ) |
516 | 536 | parameters: ErrorAnalyzerParameters = Field( |
517 | | - default_factory=ErrorAnalyzerParameters, |
518 | | - description="Error analyzer parameters" |
| 537 | + default_factory=ErrorAnalyzerParameters, description="Error analyzer parameters" |
519 | 538 | ) |
520 | 539 |
|
521 | 540 |
|
@@ -635,12 +654,10 @@ class AgentsConfig(BaseModel): |
635 | 654 | """Agents configuration""" |
636 | 655 |
|
637 | 656 | error_analyzer: Optional[ErrorAnalyzerConfig] = Field( |
638 | | - default_factory=ErrorAnalyzerConfig, |
639 | | - description="Error analyzer configuration" |
| 657 | + default_factory=ErrorAnalyzerConfig, description="Error analyzer configuration" |
640 | 658 | ) |
641 | 659 | chat_companion: Optional[ChatCompanionConfig] = Field( |
642 | | - default_factory=ChatCompanionConfig, |
643 | | - description="Chat companion configuration" |
| 660 | + default_factory=ChatCompanionConfig, description="Chat companion configuration" |
644 | 661 | ) |
645 | 662 |
|
646 | 663 |
|
|
0 commit comments