fixed tasks

src/osint_env/data/generator.py

@@ -354,6 +354,14 @@ class DatasetGenerator:
             }.get(difficulty, "full_reward"),
         }
 
+    def _task_metadata(self, index: int, base_task_type: str, metadata: dict[str, Any] | None = None) -> dict[str, Any]:
+        out = dict(metadata or {})
+        difficulty = self._normalize_difficulty(out.get("difficulty", ""), index)
+        out["difficulty"] = difficulty
+        out.setdefault("grader", self._grader_for_difficulty(difficulty))
+        out.setdefault("scenario", self._task_type_for_difficulty(base_task_type, difficulty))
+        return out
+
     def _infer_answer_from_question(self, question: str, graph: CanonicalGraph) -> str:
         entities = self._extract_entity_tokens(question)
         question_l = question.lower()
@@ -401,11 +409,8 @@ class DatasetGenerator:
         tasks: list[TaskInstance] = []
         for idx, question_spec in enumerate(self.config.seeding.seeded_questions):
             answer = question_spec.answer or self._infer_answer_from_question(question_spec.question, graph)
-            metadata = dict(question_spec.metadata)
-            difficulty = self._normalize_difficulty(metadata.get("difficulty", ""), idx)
-            metadata["difficulty"] = difficulty
-            metadata.setdefault("grader", self._grader_for_difficulty(difficulty))
-            metadata.setdefault("scenario", self._task_type_for_difficulty(question_spec.task_type, difficulty))
+            metadata = self._task_metadata(idx, question_spec.task_type, dict(question_spec.metadata))
+            difficulty = str(metadata.get("difficulty", "hard"))
             if question_spec.supporting_edges:
                 support = [
                     Edge(src=e.src, rel=e.rel, dst=e.dst, confidence=float(e.confidence))
@@ -458,6 +463,7 @@ class DatasetGenerator:
                     question=q,
                     answer=a,
                     supporting_edges=support,
+                    metadata=self._task_metadata(start_idx + i, mode),
                 )
             )
         return tasks
@@ -534,7 +540,11 @@ class DatasetGenerator:
                         question=question,
                         answer=answer,
                         supporting_edges=support,
-                        metadata={"generated_by": "llm", "shared_context": True},
+                        metadata=self._task_metadata(
+                            start_idx + len(llm_tasks),
+                            task_type,
+                            {"generated_by": "llm", "shared_context": True},
+                        ),
                     )
                 )
                 if len(llm_tasks) >= count:
@@ -579,7 +589,11 @@ class DatasetGenerator:
                             question=question,
                             answer=answer,
                             supporting_edges=support,
-                            metadata={"generated_by": "llm", "shared_context": True},
+                            metadata=self._task_metadata(
+                                start_idx + len(llm_tasks),
+                                task_type,
+                                {"generated_by": "llm", "shared_context": True},
+                            ),
                         )
                     )
                     if len(llm_tasks) >= count:

src/osint_env/env/environment.py

@@ -216,11 +216,31 @@ class OSINTEnvironment(Env):
     def _observation(self) -> Observation:
         if self.state is None:
             raise RuntimeError("State is not initialized.")
+        metadata = dict(self.state.task.metadata or {})
+        grader = metadata.get("grader") if isinstance(metadata.get("grader"), dict) else None
+        task_payload = {
+            "task_id": self.state.task.task_id,
+            "task_type": self.state.task.task_type,
+            "question": self.state.task.question,
+            "difficulty": self.state.difficulty,
+            "grader": (
+                dict(grader)
+                if grader is not None
+                else {
+                    "type": "difficulty_exact_match",
+                    "answer_type": "node_id",
+                    "case_sensitive": True,
+                    "reward_profile": self.state.difficulty,
+                }
+            ),
+        }
+        if "scenario" in metadata:
+            task_payload["scenario"] = str(metadata.get("scenario", ""))
         return Observation(
             tool_outputs=self.state.tool_outputs[-5:],
             graph_snapshot=self.memory_graph.to_snapshot(),
             action_history=self.state.action_history[-10:],
-            task={"task_id": self.state.task.task_id, "task_type": self.state.task.task_type, "question": self.state.task.question},
+            task=task_payload,
         )
 
     def _info(self) -> dict[str, Any]:

tests/test_environment.py

@@ -6,6 +6,8 @@ def test_episode_flow():
     env = OSINTEnvironment(EnvironmentConfig(max_steps=5, seed=5))
     obs = env.reset()
     assert "question" in obs.task
+    assert isinstance(obs.task.get("grader"), dict)
+    assert "type" in obs.task["grader"]
     obs, r1, done, _ = env.step(Action(ActionType.CALL_TOOL, {"tool_name": "search_posts", "args": {"query": "Update"}}))
     assert done is False
     assert isinstance(r1, float)