completed augmentations implementation

Author: lilyydu

python/packages/ai/teams/ai/ai.py

@@ -288,7 +288,7 @@ async def _on_do_command(
         state: StateT,
     ) -> bool:
         action = self._actions.get(context.data.action)
-        ctx = ActionTurnContext(context.data.action, context.data.entities, context)
+        ctx = ActionTurnContext(context.data.action, context.data.parameters, context)
 
         if not action:
             return await self._on_unknown_action(ctx, state)

python/packages/ai/teams/ai/augmentations/__init__.py

@@ -0,0 +1,10 @@
+"""
+Copyright (c) Microsoft Corporation. All rights reserved.
+Licensed under the MIT License.
+"""
+
+from .action_augmentation_section import ActionAugmentationSection
+from .augmentation import Augmentation
+from .default_augmentation import DefaultAugmentation
+from .monologue_augmentation import MonologueAugmentation
+from .sequence_augmentation import SequenceAugmentation

python/packages/ai/teams/ai/augmentations/action_augmentation_section.py

@@ -0,0 +1,116 @@
+"""
+Copyright (c) Microsoft Corporation. All rights reserved.
+Licensed under the MIT License.
+"""
+
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional, Union
+
+import yaml
+from botbuilder.core import TurnContext
+
+from teams.ai.modelsv2.chat_completion_action import ChatCompletionAction
+from teams.ai.promptsv2.message import Message
+from teams.ai.promptsv2.prompt_functions import PromptFunctions
+from teams.ai.promptsv2.prompt_section_base import PromptSectionBase
+from teams.ai.promptsv2.rendered_prompt_section import RenderedPromptSection
+from teams.ai.tokenizers.tokenizer import Tokenizer
+from teams.state.memory import Memory
+
+
+@dataclass
+class ActionValue:
+    description: Optional[str] = None
+    parameters: Optional[Union[Dict[str, Any], Dict[str, Dict[str, Any]]]] = None
+
+
+@dataclass
+class ActionList:
+    actions: Dict[str, ActionValue]
+
+
+class ActionAugmentationSection(PromptSectionBase):
+    """
+    A prompt section that renders a list of actions to the prompt.
+    """
+
+    _text: str
+    _token_list: Optional[List[int]] = None
+    _actions: Dict[str, ChatCompletionAction] = {}
+
+    @property
+    def actions(self) -> Dict[str, ChatCompletionAction]:
+        """
+        Map of action names to actions.
+        """
+        return self._actions
+
+    def __init__(self, actions: List[ChatCompletionAction], call_to_action: str) -> None:
+        """
+        Creates a new `ActionAugmentationSection` instance.
+
+        Args:
+            actions (List[ChatCompletionAction]): List of actions to render.
+            call_to_action (str): Text to display after the list of actions.
+
+        """
+        super().__init__(-1, True, "\n\n")
+
+        # Convert actions to an ActionList
+        action_list: ActionList = {"actions": {}}
+
+        for action in actions:
+            self._actions[action.name] = action
+            action_list["actions"][action.name] = {}
+            if action.description:
+                action_list["actions"][action.name]["description"] = action.description
+            if action.parameters:
+                params = action.parameters
+                action_list["actions"][action.name]["parameters"] = (
+                    params.get("properties")
+                    if params.get("additional_properties") is None
+                    else params
+                )
+
+        # Build augmentation text
+        self._text = f"{yaml.dump(action_list)}\n\n{call_to_action}"
+
+    async def render_as_messages(
+        self,
+        context: TurnContext,
+        memory: Memory,
+        functions: PromptFunctions,
+        tokenizer: Tokenizer,
+        max_tokens: int,
+    ) -> RenderedPromptSection[List[Message[str]]]:
+        """
+        Renders the prompt section as a list of `Message` objects.
+
+        Args:
+            context (TurnContext): Context for the current turn of conversation.
+            memory (Memory): Interface for accessing state variables.
+            functions (PromptFunctions): Functions for rendering prompts.
+            tokenizer (Tokenizer): Tokenizer to use for encoding/decoding text.
+            max_tokens (int): Maximum number of tokens allowed for the rendered prompt.
+
+        Returns:
+            RenderedPromptSection[List[Message[str]]]: The rendered prompt section.
+
+        """
+        # Tokenize on first use
+        if not self._token_list:
+            self._token_list = tokenizer.encode(self._text)
+
+        # Check for max tokens
+        if len(self._token_list) > max_tokens:
+            trimmed = self._token_list[0:max_tokens]
+            return RenderedPromptSection[List[Message[str]]](
+                output=[Message[str](role="system", content=tokenizer.decode(trimmed))],
+                length=len(trimmed),
+                too_long=True,
+            )
+        return RenderedPromptSection[List[Message[str]]](
+            output=[Message[str](role="system", content=self._text)],
+            length=len(self._token_list),
+            too_long=False,
+        )

python/packages/ai/teams/ai/augmentations/augmentation.py

@@ -0,0 +1,50 @@
+"""
+Copyright (c) Microsoft Corporation. All rights reserved.
+Licensed under the MIT License.
+"""
+
+from abc import ABC, abstractmethod
+from typing import Generic, TypeVar, Union
+
+from botbuilder.core import TurnContext
+
+from teams.ai.modelsv2.prompt_response import PromptResponse
+from teams.ai.planner import Plan
+from teams.ai.promptsv2.prompt_section import PromptSection
+from teams.ai.validators.prompt_response_validator import PromptResponseValidator
+from teams.state.memory import Memory
+
+ValueT = TypeVar("ValueT")
+"Type of message content returned for a 'success' response."
+
+
+class Augmentation(PromptResponseValidator[ValueT], ABC, Generic[ValueT]):
+    """
+    An augmentation is a component that can be added to a prompt template to add additional
+    functionality to the prompt.
+    """
+
+    @abstractmethod
+    def create_prompt_section(self) -> Union[PromptSection, None]:
+        """
+        Creates an optional prompt section for the augmentation.
+
+        Returns:
+        Union[PromptSection, None]: The prompt section.
+        """
+
+    @abstractmethod
+    async def create_plan_from_response(
+        self, turn_context: TurnContext, memory: Memory, response: PromptResponse[ValueT]
+    ) -> Plan:
+        """
+        Creates a plan given validated response value.
+
+        Args:
+            turn_context (TurnContext): Context for the current turn of conversation.
+            memory (Memory): An interface for accessing state variables.
+            response (PromptResponse[ValueT]): Validated, transformed response for the prompt.
+
+        Returns:
+            Plan: The created plan.
+        """

python/packages/ai/teams/ai/augmentations/default_augmentation.py

@@ -0,0 +1,72 @@
+"""
+Copyright (c) Microsoft Corporation. All rights reserved.
+Licensed under the MIT License.
+"""
+
+from typing import Union
+
+from botbuilder.core import TurnContext
+
+from teams.ai.augmentations.augmentation import Augmentation
+from teams.ai.modelsv2.prompt_response import PromptResponse
+from teams.ai.planner import Plan
+from teams.ai.planner.predicted_say_command import PredictedSayCommand
+from teams.ai.promptsv2.prompt_section import PromptSection
+from teams.ai.tokenizers.tokenizer import Tokenizer
+from teams.ai.validators.validation import Validation
+from teams.state.memory import Memory
+
+
+class DefaultAugmentation(Augmentation[str]):
+    """
+    The default 'none' augmentation.
+
+    This augmentation does not add any additional functionality to the prompt. It always
+    returns a `Plan` with a single `SAY` command containing the models response.
+    """
+
+    def create_prompt_section(self) -> Union[PromptSection, None]:
+        """
+        Creates an optional prompt section for the augmentation.
+        """
+        return None
+
+    async def validate_response(
+        self,
+        context: TurnContext,
+        memory: Memory,
+        tokenizer: Tokenizer,
+        response: PromptResponse[str],
+        remaining_attempts: int,
+    ) -> Validation[str]:
+        """
+        Validates a response to a prompt.
+
+        Args:
+            context (TurnContext): Context for the current turn of conversation.
+            memory (Memory): Interface for accessing state variables.
+            tokenizer (Tokenizer): Tokenizer to use for encoding/decoding text.
+            response (PromptResponse[str]): Response to validate.
+            remaining_attempts (int): Nubmer of remaining attempts to validate the response.
+
+        Returns:
+            Validation[str]: A 'Validation' object.
+        """
+        return Validation[str](valid=True)
+
+    async def create_plan_from_response(
+        self, turn_context: TurnContext, memory: Memory, response: PromptResponse[str]
+    ) -> Plan:
+        """
+        Creates a plan given validated response value.
+
+        Args:
+            turn_context (TurnContext): Context for the current turn of conversation.
+            memory (Memory):  Interface for accessing state variables.
+            response (PromptResponse[str]): The validated and transformed response for the prompt.
+
+        Returns:
+            Plan: The created plan.
+        """
+        say_response = response.message.content if response.message.content else ""
+        return Plan(commands=[PredictedSayCommand(response=say_response)])