Spaces:
Running
Running
File size: 5,694 Bytes
ea1db8d | 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 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 | """
Typed models for the Second Brain OpenEnv environment.
Defines Action, Observation, and State using proper OpenEnv inheritance.
"""
from enum import Enum
from typing import Any, Dict, List, Optional
from pydantic import Field
from openenv.core.env_server import Action, Observation
from openenv.core.env_server.types import State
# ---------------------------------------------------------------------------
# ActionType Enum — prevents typos like "categorise" instead of "categorize"
# ---------------------------------------------------------------------------
class ActionType(str, Enum):
categorize = "categorize"
retrieve = "retrieve"
synthesize = "synthesize"
tag = "tag"
skip = "skip"
# ---------------------------------------------------------------------------
# Action — inherits from OpenEnv Action
# ---------------------------------------------------------------------------
class SecondBrainAction(Action):
"""
A single action the agent can take in the Second Brain environment.
action_type options:
- categorize : assign a category to the current note (Task 1)
- retrieve : query the knowledge base with a search string (Task 2 & 3)
- synthesize : produce a final synthesized answer (Task 3)
- tag : add tags to a note
- skip : skip the current item (no reward)
Fields:
action_type : validated enum — prevents typos
content : category name / search query / synthesized answer
note_id : optional — which note to act on
tags : optional list of tags
"""
action_type: ActionType = Field(
...,
description="One of: categorize | retrieve | synthesize | tag | skip"
)
content: str = Field(
...,
description="Main payload: category name, search query, or synthesized answer"
)
note_id: Optional[str] = Field(
default=None,
description="Target note ID for retrieve or tag actions"
)
tags: Optional[List[str]] = Field(
default=None,
description="List of tags to apply (used with action_type='tag')"
)
# ---------------------------------------------------------------------------
# Observation — inherits from OpenEnv Observation
# ---------------------------------------------------------------------------
class SecondBrainObservation(Observation):
"""
What the agent observes after each step.
Fields:
current_note : the note currently being processed (Task 1)
query : the retrieval question (Task 2 & 3)
retrieved_notes : notes returned by a retrieve action
knowledge_base_size : how many notes are in the KB
step_count : current step number
task_name : which task is active
reward : reward from the last action
done : whether the episode is finished
score : running cumulative score (0.0–1.0)
feedback : human-readable feedback on last action
valid_categories : categories the agent may use (Task 1)
remaining_items : how many items still need processing
"""
current_note: Optional[Dict[str, Any]] = Field(
default=None,
description="The note currently presented to the agent"
)
query: Optional[str] = Field(
default=None,
description="The retrieval question the agent must answer"
)
retrieved_notes: Optional[List[Dict[str, Any]]] = Field(
default=None,
description="Notes returned after a retrieve action"
)
knowledge_base_size: int = Field(
default=0,
description="Total notes in the knowledge base"
)
step_count: int = Field(default=0)
task_name: str = Field(default="")
reward: float = Field(default=0.0)
done: bool = Field(default=False)
score: float = Field(default=0.0, description="Running score in [0, 1]")
feedback: str = Field(default="", description="Feedback on the last action")
valid_categories: List[str] = Field(
default_factory=lambda: ["work", "personal", "reference", "action_item"],
description="Valid category labels for Task 1"
)
remaining_items: int = Field(
default=0,
description="Number of items still to be processed this episode"
)
# ---------------------------------------------------------------------------
# State — extends OpenEnv State with environment-specific fields
# ---------------------------------------------------------------------------
class SecondBrainState(State):
"""
Full internal state of the Second Brain environment.
Used by state() endpoint and for debugging/graders.
Fields:
task_name : which task is currently running
step_count : how many steps have been taken
score : current normalized score in [0.0, 1.0]
correct_count : Task 1 — correct categorizations so far
retrieval_scores : Task 2 — list of per-question retrieval scores
synthesis_scores : Task 3 — list of per-question synthesis scores
collected_note_ids : Task 3 — notes collected before synthesizing
done : whether the episode has ended
"""
task_name: str = Field(default="")
step_count: int = Field(default=0)
score: float = Field(default=0.0)
correct_count: int = Field(default=0)
retrieval_scores: List[float] = Field(default_factory=list)
synthesis_scores: List[float] = Field(default_factory=list)
collected_note_ids: List[str] = Field(default_factory=list)
done: bool = Field(default=False) |