File size: 985 Bytes
b196357
 
 
 
 
 
 
 
 
bc02c39
b196357
bc02c39
b196357
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
bc02c39
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
from __future__ import annotations

from pydantic import BaseModel, ConfigDict, Field


class ModuleReviewState(BaseModel):
    model_config = ConfigDict(strict=True, extra="forbid")

    module_id: str
    module_summary: str | None = None
    review_status: str
    latest_review_summary: str | None = None
    issues_found: int = 0
    last_action: str | None = None
    last_reward: float = 0.0
    updated_at: str


class EpisodeState(BaseModel):
    model_config = ConfigDict(strict=True, extra="forbid")

    episode_id: str
    task_id: str
    current_module_id: str | None = None
    modules_remaining: int
    step_count: int
    cumulative_reward: float
    done: bool
    status: str


class GraphState(BaseModel):
    model_config = ConfigDict(strict=True, extra="forbid")

    episode: EpisodeState
    modules: list[ModuleReviewState] = Field(default_factory=list)
    module_count: int
    edge_count: int
    annotation_count: int
    total_annotation_count: int = 0