Structured Outputs with Pydantic for CRITIC (QA, Math, Code)

agential/cog/agent/critic.py

@@ -7,11 +7,28 @@
 from typing import Any, Dict, List
 
 from langchain_core.language_models.chat_models import BaseChatModel
+from pydantic import BaseModel, Field
 
 from agential.cog.agent.base import BaseAgent
 from agential.cog.strategies.strategy_factory import CriticStrategyFactory
 
 
+class CriticOutput(BaseModel):
+    """Critic Pydantic output class.
+
+    Attributes:
+        answer (str): The answer generated by the agent.
+        critique (str): The critique of the answer generated by the agent.
+        external_tool_info (Dict[str, Any]): The query requested by the agent.
+    """
+
+    answer: str = Field(..., description="The answer generated by the agent.")
+    critique: str = Field(..., description="The answer's critique.")
+    external_tool_info: Dict[str, Any] = Field(
+        ..., description="The external tool outputs."
+    )
+
+
 class CriticAgent(BaseAgent):
     """CRITIC Agent.
 
@@ -98,7 +115,11 @@ def generate(
             )
 
             out.append(
-                self.strategy.create_output_dict(answer, critique, external_tool_info)
+                CriticOutput(
+                    **self.strategy.create_output_dict(
+                        answer, critique, external_tool_info
+                    )
+                )
             )
 
             if self.strategy.halting_condition():

agential/cog/strategies/critic/base.py

@@ -73,7 +73,7 @@ def update_answer_based_on_critique(
         critique: str,
         prompt: str,
         additional_keys: Dict[str, str],
-        external_tool_info: Dict[str, str],
+        external_tool_info: Dict[str, Any],
         **kwargs: Any,
     ) -> str:
         """Updates the answer based on the provided critique using the given language model and question.
@@ -85,7 +85,7 @@ def update_answer_based_on_critique(
             critique (str): The critique of the original answer.
             prompt (str): The instruction template used to prompt the language model for the update.
             additional_keys (Dict[str, str]): Additional keys to format the update prompt.
-            external_tool_info (Dict[str, str]): Information from any external tools used during the critique.
+            external_tool_info (Dict[str, Any]): Information from any external tools used during the critique.
             **kwargs (Any): Additional arguments that might be needed for specific implementations.
 
         Returns: