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

agential/cog/agent/critic.py

@@ -8,21 +8,63 @@
 
 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 CriticPydanticOutput(BaseModel):
+    class Config:
+        title = 'Critic Output'
+        description = 'Critic output for different modes'
+
+    class QA(BaseModel):
+        answer: str = Field(..., description="The answer generated by the agent.")
+        critique: str = Field(..., description="The critique of the answer generated by the agent.")
+        query: str = Field(..., description="The query requested by the agent.")
+        search_result: str = Field(..., description="The search result requested by the agent.")
+        revised_answer: str = Field(..., description="The revised answer generated by the agent.")
+
+    class Math(BaseModel):
+        code: str = Field(..., description="The code generated by the agent.")
+        critique: str = Field(..., description="The critique of the code generated by the agent.")
+        execution_status: str = Field(..., description="The execution status of the agent.")
+        code_answer: str = Field(..., description="The code answer generated by the agent.")
+        improved_code: str = Field(..., description="The improved code generated by the agent.")
+
+    class Code(BaseModel):
+        code: str = Field(..., description="The code generated by the agent.")
+        critique: str = Field(..., description="The critique of the code generated by the agent.")
+        execution_status: str = Field(..., description="The execution status of the agent.")
+        improved_code: str = Field(..., description="The improved code generated by the agent.")
+
+
+        #answer: str = Field(..., description = "The answer generated by the agent.")
+        #critique : str = Field(..., description = "The critique of the answer generated by the agent.")
+        #query: str = Field(..., description = "The query requested by the agent. ")
+        #search_result: str = Field(..., description = "The search result requested by the agent.")
+        #revised_answer: str = Field(..., description = "The revised answer generated by the agent. ")
+
+        #code: int/str? = Field(..., description = "The code generated by the agent.")
+        #critique: str = Field(..., description = " The critique of the answer generated by the agent.")
+
+        #execution_status: str = Field(..., description = "The execution status of the agent.")
+        #code_answer: int/str? = Field(..., description = "The code answer generated by the agent.")
+        #improved_code: int/str? = Field(..., description = "The improved code generated by the agent.")
+
+
+
+
 
 class CriticAgent(BaseAgent):
     """CRITIC Agent.
-
     Attributes:
         llm (BaseChatModel): An instance of a language model used for generating initial answers
             and critiques.
         mode (Dict[str, str]): A dictionary specifying the CRITIC agent's mode and the benchmark.
             For example, {"qa": "hotpotqa"}, {"math": "gsm8k"}, or {"code": "mbpp"}.
         **strategy_kwargs (Dict[str, Any]): Additional strategy-specific arguments.
     """
-
     def __init__(
         self,
         llm: BaseChatModel,
@@ -31,14 +73,11 @@ def __init__(
     ) -> None:
         """Initialization."""
         super().__init__()
-
         self.llm = llm
         self.mode = mode
-
         self.strategy = CriticStrategyFactory().get_strategy(
             mode=self.mode, llm=self.llm, **strategy_kwargs
         )
-
     def generate(
         self,
         question: str,
@@ -53,6 +92,7 @@ def generate(
         reset: bool = True,
         **kwargs: Any,
     ) -> List[Dict[str, Any]]:
+
         """Generates an answer that is refined with search results.
 
         Args:
@@ -77,11 +117,10 @@ def generate(
         if reset:
             self.reset()
 
-        out = []
 
+        out = []
         # Initial answer generation.
         answer = self.strategy.generate(question, examples, prompt, additional_keys)
-
         critique = ""
         for idx in range(max_interactions):
             critique, external_tool_info = self.strategy.generate_critique(
@@ -97,25 +136,67 @@ def generate(
                 **kwargs,
             )
 
-            out.append(
-                self.strategy.create_output_dict(answer, critique, external_tool_info)
-            )
+            if self.mode["qa"]:
+                output = CriticPydanticOutput.QA(
+                    answer=answer,
+                    critique=critique,
+                    query=external_tool_info.get("query", ""),
+                    search_result=external_tool_info.get("search_result", ""),
+                    revised_answer=self.strategy.update_answer_based_on_critique(
+                        question=question,
+                        examples=critique_examples,
+                        answer=answer,
+                        critique=critique,
+                        prompt=critique_prompt,
+                        additional_keys=critique_additional_keys,
+                        external_tool_info=external_tool_info,
+                        **kwargs,
+                    ),
+                )
+            elif self.mode["math"]:
+                output = CriticPydanticOutput.Math(
+                    code=answer,
+                    critique=critique,
+                    execution_status=external_tool_info.get("execution_status", ""),
+                    code_answer=external_tool_info.get("code_answer", ""),
+                    improved_code=self.strategy.update_answer_based_on_critique(
+                        question=question,
+                        examples=critique_examples,
+                        answer=answer,
+                        critique=critique,
+                        prompt=critique_prompt,
+                        additional_keys=critique_additional_keys,
+                        external_tool_info=external_tool_info,
+                        **kwargs,
+                    ),
+                )
+
+            elif self.mode["code"]:
+                output = CriticPydanticOutput.Code(
+                    code=answer,
+                    critique=critique,
+                    execution_status=external_tool_info.get("execution_status", ""),
+                    improved_code=self.strategy.update_answer_based_on_critique(
+                        question=question,
+                        examples=critique_examples,
+                        answer=answer,
+                        critique=critique,
+                        prompt=critique_prompt,
+                        additional_keys=critique_additional_keys,
+                        external_tool_info=external_tool_info,
+                        **kwargs,
+                    ),
+                )
+
+            out.append(output.dict())
+
 
             if self.strategy.halting_condition():
                 break
-
-            # Update answer for the next iteration.
-            answer = self.strategy.update_answer_based_on_critique(
-                question=question,
-                examples=critique_examples,
-                answer=answer,
-                critique=critique,
-                prompt=critique_prompt,
-                additional_keys=critique_additional_keys,
-                external_tool_info=external_tool_info,
-                **kwargs,
             )
 
+            answer = output.revised_answer if self.mode["qa"] else output.improved_code
+
         return out
 
     def reset(self) -> None:

agential/cog/strategies/critic/base.py

@@ -7,14 +7,13 @@
 
 from agential.cog.strategies.base import BaseStrategy
 
+from pydantic import BaseModel, Field
 
 class CriticBaseStrategy(BaseStrategy):
     """An abstract base class for defining strategies for the CRITIC Agent."""
-
     def __init__(self, llm: BaseChatModel) -> None:
         """Initialization."""
         super().__init__(llm)
-
     @abstractmethod
     def generate_critique(
         self,
@@ -30,7 +29,6 @@ def generate_critique(
         **kwargs: Any,
     ) -> Tuple[str, Dict[str, Any]]:
         """Generates a critique of the provided answer using the given language model, question, examples, and prompt.
-
         Args:
             idx (int): The index of the current interaction.
             question (str): The question that was answered by the language model.
@@ -41,29 +39,44 @@ def generate_critique(
             additional_keys (Dict[str, str]): Additional keys to format the critique prompt.
             use_tool (bool): Whether to use an external tool (e.g., code interpreter, search tool) during critique.
             max_interactions (int): The maximum number of critique interactions.
-            **kwargs (Any): Additional arguments that might be needed for specific implementations.
-
+            **kwargs (Dict[str, Any]): Additional arguments that might be needed for specific implementations.
+            
         Returns:
             Tuple[str, Dict[str, Any]]: The generated critique and external tool information.
         """
         pass
-
     @abstractmethod
     def create_output_dict(
-        self, answer: str, critique: str, external_tool_info: Dict[str, str]
-    ) -> Dict[str, Any]:
+        self, answer: str, critique: str, external_tool_info: Dict[str, Any]
         """Creates a dictionary containing the answer and critique, along with any additional key updates.
-
         Args:
             answer (str): The original answer.
             critique (str): The generated critique.
-            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.
         Returns:
             Dict[str, Any]: A dictionary containing the answer, critique, and additional key updates.
         """
         pass
 
+
+    @abstractmethod
+    def create_output_pydantic(
+        self, answer: str, critique: str, query: str, search_result: str, revised_answer: str
+    ) -> BaseModel:
+        """Creates a pydantic model of the output.
+
+        Args:
+            answer (str): The answer generated by the agent.
+            critique (str): The critique of the answer generated by the agent.
+            query (str): The query requested by the agent.
+            search_result (str): The search result requested by the agent.
+            revised_answer (str): The revised answer generated by the agent.
+        Returns: 
+            BaseModel: A Pydantic model containing the answer, critique, query, search result, and revised answer
+        """
+        pass
+
+
     @abstractmethod
     def update_answer_based_on_critique(
         self,

agential/cog/strategies/critic/code.py

@@ -1,27 +1,40 @@
 """CRITIC Agent strategies for Code."""
-
 from typing import Any, Dict, Tuple
-
 from langchain_core.language_models.chat_models import BaseChatModel
-
 from agential.cog.functional.critic import _prompt_agent, _prompt_critique
 from agential.cog.strategies.critic.base import CriticBaseStrategy
 from agential.utils.general import safe_execute
 from agential.utils.validation import validate_overlapping_keys
 
+from pydantic import BaseModel, Field, Extra
+
+
+class CriticPydanticOutput(BaseModel):
+    """Critic output
+    Attributes:
+
+        code (str): The code generated by the agent.
+        critique (str): The critique of the answer generated by the agent.
+        execution_status (str): The execution status of the agent.
+        improved_code (str): The improved code generated by the agent.
+    """
+
+        code: str = Field(..., description = "The code generated by the agent.")
+        critique: str = Field(..., description = " The critique of the answer generated by the agent.")
+        execution_status: str = Field(..., description = "The execution status of the agent.")
+        improved_code: str = Field(..., description = "The improved code generated by the agent.")
+
+
 
 class CriticCodeStrategy(CriticBaseStrategy):
     """A strategy class for Code benchmarks using the CRITIC agent.
-
     Attributes:
         llm (BaseChatModel): The language model used for generating answers and critiques.
     """
-
     def __init__(self, llm: BaseChatModel) -> None:
         """Initialization."""
         super().__init__(llm)
         self._halt = False
-
     def generate(
         self,
         question: str,
@@ -31,13 +44,12 @@ def generate(
         **kwargs: Any,
     ) -> str:
         """Generates an answer for the given question using the provided prompt and examples.
-
         Args:
             question (str): The math question to generate an answer for.
             examples (str): Few-shot examples to guide the language model.
             prompt (str): The prompt to generate an answer.
             additional_keys (Dict[str, str]): Additional keys for the prompt.
-            **kwargs (Any): Additional arguments.
+            **kwargs (Dict[str, Any]): Additional arguments.
 
         Returns:
             str: The generated answer.
@@ -53,6 +65,25 @@ def generate(
 
         return answer
 
+
+    def generate_pydantic_output(
+        self, code: str, critique: str, execution_status: str, improved_code: str
+    ) -> CriticPydanticOutput:
+        """Creates a pydantic model of the output.
+        Args:
+            code (str): The code generated by the agent.
+            critique (str): The critique of the answer generated by the agent.
+            execution_status (str): The execution status of the agent.
+            improved_code (str): The improved code generated by the agent.
+        """
+
+        return CriticPydanticOutput(
+            code=code,
+            critique=critique,
+            execution_status=execution_status,
+            improved_code=improved_code,
+        )
+
     def generate_critique(
         self,
         idx: int,
@@ -95,7 +126,7 @@ def generate_critique(
         Returns:
             Tuple[str, Dict[str, Any]]: The generated critique and external tool information.
         """
-        external_tool_info = {}
+        external_tool_info = {"execution_status": ""}
         if use_tool:
             if "tests" not in additional_keys:
                 raise ValueError(
@@ -113,7 +144,7 @@ def generate_critique(
             validate_overlapping_keys(additional_keys, external_tool_info)
 
         additional_keys = additional_keys.copy()
-        additional_keys.update(external_tool_info)
+        additional_keys.update(external_tool_info if use_tool else {})
 
         new_critique = _prompt_critique(
             llm=self.llm,
@@ -128,19 +159,24 @@ def generate_critique(
         return new_critique, external_tool_info
 
     def create_output_dict(
-        self, answer: str, critique: str, external_tool_info: Dict[str, str]
+        self, answer: str, critique: str, external_tool_info: Dict[str, Any]
     ) -> Dict[str, str]:
         """Creates an output dictionary containing the answer, critique, and external tool information.
 
         Args:
             answer (str): The generated answer.
             critique (str): The generated critique.
-            external_tool_info (Dict[str, str]): Information from external tool execution.
+            external_tool_info (Dict[str, Any]): Information from external tool execution.
 
         Returns:
             Dict[str, str]: The output dictionary.
         """
-        output_dict = {"code": answer, "critique": critique, **external_tool_info}
+        output_dict = {
+            "answer": answer,
+            "critique": critique,
+            "external_tool_info": external_tool_info,
+        }        
+
         return output_dict
 
     def update_answer_based_on_critique(