updated parsing outputs of agents when agents output their thought process as well.

kohankhaki · kohankhaki · commit 382030d50f7f · 2025-08-21T14:06:18.000-04:00
diff --git a/src/area_generation/generator.py b/src/area_generation/generator.py
@@ -15,14 +15,12 @@
 from .scientist import AreaScientist
 
 
-logging.getLogger("autogen").setLevel(logging.WARNING)
-logging.getLogger("autogen_core").setLevel(logging.WARNING)
-logging.getLogger("autogen_ext").setLevel(logging.WARNING)
-
 log = logging.getLogger("agentic_area_gen.generator")
 
 lf = Langfuse(blocked_instrumentation_scopes=["autogen SingleThreadedAgentRuntime"])
-openlit.init(tracer=lf._otel_tracer, disable_batch=True)
+
+
+openlit.init(tracer=lf._otel_tracer, disable_metrics=True)
 
 
 async def generate_areas(cfg: DictConfig) -> None:
diff --git a/src/area_generation/moderator.py b/src/area_generation/moderator.py
@@ -144,34 +144,37 @@ async def _merge_proposals(
             is_finalized = False
             try:
                 parsed = json.loads(raw_content)
+                log.info(f"Moderator parsed: {parsed}")
+                revised_areas = parsed.get("areas", {})
                 is_finalized = parsed.get("finalized", False)
+
             except Exception:
-                log.error(f"Error parsing final areas JSON: {raw_content}")
+                log.error(f"Error parsing merged areas JSON: {raw_content}")
                 raise
 
             if is_finalized or round_num >= self._max_round - 1:
                 log.info(f"Moderator finalizing areas after round {round_num}")
-                await self._finalize_areas(raw_content)
+                await self._finalize_areas(revised_areas)
             else:
                 log.info(
                     f"Moderator sending merged proposal for revision in round {round_num}"
                 )
-                self._round = round_num + 1
+
+                # Use the already parsed areas instead of parsing again
+                revision_content = json.dumps(revised_areas)
 
                 # Send to scientists for revision
                 await self.publish_message(
                     ScientistRevisionRequest(
-                        scientist_id="A",
-                        moderator_proposal=raw_content,
-                        round=self._round,
+                        moderator_proposal=revision_content,
+                        round=round_num + 1,
                     ),
                     topic_id=DefaultTopicId(),
                 )
                 await self.publish_message(
                     ScientistRevisionRequest(
-                        scientist_id="B",
-                        moderator_proposal=raw_content,
-                        round=self._round,
+                        moderator_proposal=revision_content,
+                        round=round_num + 1,
                     ),
                     topic_id=DefaultTopicId(),
                 )
@@ -181,35 +184,19 @@ async def _merge_proposals(
             log.error(f"Traceback: {traceback.format_exc()}")
             raise
 
-    async def _finalize_areas(self, final_areas: str) -> None:
+    async def _finalize_areas(self, final_areas: dict) -> None:
         """Save final areas to file."""
         try:
             log.info("Moderator finalizing and saving areas")
 
-            # Convert to the expected format with "areas" list
-            try:
-                json_start = final_areas.find("{")
-                if json_start != -1:
-                    json_part = final_areas[json_start:]
-                    parsed = json.loads(json_part)
-                else:
-                    parsed = json.loads(final_areas)
-
-                if "finalized" in parsed:
-                    del parsed["finalized"]  # Remove finalized flag
-
-                # Convert area_0, area_1 format to areas list
-                areas_list = []
-                i = 0
-                while f"area_{i}" in parsed:
-                    areas_list.append(parsed[f"area_{i}"])
-                    i += 1
-
-                final_format = {"areas": areas_list}
-                final_areas_json = json.dumps(final_format, indent=2)
-            except Exception as e:
-                log.warning(f"Could not parse final areas JSON: {e}")
-                final_areas_json = final_areas
+            areas_list = []
+            i = 0
+            while f"area_{i}" in final_areas:
+                areas_list.append(final_areas[f"area_{i}"])
+                i += 1
+
+            final_format = {"areas": areas_list}
+            final_areas_json = json.dumps(final_format, indent=2)
 
             self._save_areas_to_file(final_areas_json)
             log.info("Area generation completed successfully")
diff --git a/src/area_generation/scientist.py b/src/area_generation/scientist.py
@@ -1,5 +1,6 @@
 """Area scientist agent for generating capability areas."""
 
+import json
 import logging
 import traceback
 
@@ -27,6 +28,7 @@
     ScientistRevisionRequest,
 )
 
+
 log = logging.getLogger("agentic_area_gen.scientist")
 
 
@@ -73,11 +75,20 @@ async def handle_area_proposal_request(
             if not isinstance(raw_content, str):
                 raw_content = str(raw_content)
 
+            # Extract only the areas part for the proposal message
+            proposal_content = raw_content
+            try:
+                parsed = json.loads(raw_content)
+                proposal_content = json.dumps(parsed["areas"])
+            except Exception as e:
+                log.error(f"Could not parse scientist response as JSON: {e}")
+                raise
+
             log.info(f"Scientist {self._scientist_id} publishing area proposal")
             await self.publish_message(
                 ScientistAreaProposal(
                     scientist_id=self._scientist_id,
-                    proposal=raw_content,
+                    proposal=proposal_content,
                     round=self._round,
                 ),
                 topic_id=DefaultTopicId(),
@@ -119,13 +130,21 @@ async def handle_revision_request(
             if not isinstance(raw_content, str):
                 raw_content = str(raw_content)
 
+            proposal_content = raw_content
+            try:
+                parsed = json.loads(raw_content)
+                proposal_content = json.dumps(parsed["areas"])
+            except Exception as e:
+                log.error(f"Could not parse scientist revision response as JSON: {e}")
+                raise
+
             log.info(
                 f"Scientist {self._scientist_id} publishing revised proposal for round {message.round}"
             )
             await self.publish_message(
                 ScientistAreaProposal(
                     scientist_id=self._scientist_id,
-                    proposal=raw_content,
+                    proposal=proposal_content,
                     round=message.round,
                 ),
                 topic_id=DefaultTopicId(),
@@ -136,4 +155,4 @@ async def handle_revision_request(
                 f"Error in Scientist {self._scientist_id} handle_revision_request: {e}"
             )
             log.error(f"Traceback: {traceback.format_exc()}")
-            raise 
+            raise
diff --git a/src/utils/agentic_prompts.py b/src/utils/agentic_prompts.py
@@ -15,14 +15,17 @@
 1. A short name (a few words).
 2. A 2–3 sentence description that defines its boundaries and justifies its inclusion.
 
-Please return your proposal in the following format:
-RESPONSE JSON:
+Please return your proposal and your thoughts and reasoning in the following format:
 {{
-  "area_0": {{
-    "name": <STR>,
-    "description": <STR>
-  }},
-  ...
+  "thought": <STR>,
+  "areas":
+  {{
+    "area_0": {{
+      "name": <STR>,
+      "description": <STR>
+    }},
+    ...
+  }}
 }}"""
 
 AREA_SCIENTIST_REVISION_PROMPT = """You are Scientist {scientist_id}. You are reviewing the merged set of capability areas proposed by the Moderator.
@@ -37,16 +40,19 @@
 
 Keep your feedback constructive and focused on improving clarity, coverage, and non-overlap. Avoid unnecessary changes.
 
-Return your revised proposal with the following format:
-THOUGHT: <your summary of thoughts on the proposal and changes you made>
-RESPONSE JSON:
+Return your revised proposal and your thoughts and reasoning with the following format:
 {{
-  "area_0": {{
-    "name": <STR>,
-    "description": <STR>
-  }},
-  ...
-}}"""
+  "thought": <STR>,
+  "areas":
+  {{
+    "area_0": {{
+      "name": <STR>,
+      "description": <STR>
+    }},
+    ...
+  }}
+}}
+"""
 
 AREA_MODERATOR_MERGE_PROMPT = """You are the Moderator. Two scientist agents have independently proposed a list of high-level capability areas for evaluating large language models in the domain of {domain}.
 
@@ -66,14 +72,18 @@
 Explain how you merge the above proposals. Be thoughtful and concise in your output.
 You will then submit this merged proposal for review by the scientist agents. If either scientist provides substantive suggestions, you may revise the proposal and initiate another round of review.{finalized_instruction}
 
-Present the merged areas in the following format:
-THOUGHT: <your summary of thoughts on the proposals and merges you made>
+Present the merged areas and your thoughts and reasoning in the following format:
 {{
-  "area_0": {{
-    "name": <STR>,
-    "description": <STR>
-  }},
-  ...{finalized_field}
+  "thought": <STR>,
+  "areas":
+  {{
+    "area_0": {{
+      "name": <STR>,
+      "description": <STR>
+    }},
+    ...
+  }}
+  {finalized_field}
 }}
 """
 
@@ -267,4 +277,4 @@
 "finalized": true
 at the end of your JSON response."""
 
-FINALIZED_FIELD = ', "finalized": <true|false>'
+FINALIZED_FIELD = '"finalized": <true|false>'
