agenta 0.12.3__py3-none-any.whl → 0.32.0a1__py3-none-any.whl

agenta/sdk/types.py

@@ -1,13 +1,32 @@
 import json
-from typing import Any, Dict, List, Optional
+from dataclasses import dataclass
+from typing import Dict, List, Optional, Any, Union
 
-from pydantic import BaseModel, Extra, HttpUrl, Field
+from pydantic import ConfigDict, BaseModel, HttpUrl
 
+from agenta.client.backend.types.agenta_node_dto import AgentaNodeDto
+from agenta.client.backend.types.agenta_nodes_response import AgentaNodesResponse
+from typing import Annotated, List, Union, Optional, Dict, Literal, Any
+from pydantic import BaseModel, Field, root_validator
+from agenta.sdk.assets import supported_llm_models
 
-class InFile:
-    def __init__(self, file_name: str, file_path: str):
-        self.file_name = file_name
-        self.file_path = file_path
+
+@dataclass
+class MultipleChoice:
+    choices: Union[List[str], Dict[str, List[str]]]
+
+
+def MCField(  # pylint: disable=invalid-name
+    default: str,
+    choices: Union[List[str], Dict[str, List[str]]],
+) -> Field:
+    field = Field(default=default, description="ID of the model to use")
+    if isinstance(choices, dict):
+        field.json_schema_extra = {"choices": choices, "x-parameter": "grouped_choice"}
+    elif isinstance(choices, list):
+        field.json_schema_extra = {"choices": choices, "x-parameter": "choice"}
+
+    return field
 
 
 class LLMTokenUsage(BaseModel):
@@ -16,95 +35,84 @@ class LLMTokenUsage(BaseModel):
     total_tokens: int
 
 
-class FuncResponse(BaseModel):
-    message: str
-    usage: Optional[LLMTokenUsage]
-    cost: Optional[float]
-    latency: float
+class BaseResponse(BaseModel):
+    version: Optional[str] = "3.1"
+    data: Optional[Union[str, Dict[str, Any]]] = None
+    content_type: Optional[str] = "string"
+    tree: Optional[AgentaNodesResponse] = None
+    tree_id: Optional[str] = None
 
 
 class DictInput(dict):
-    def __new__(cls, default_keys=None):
+    def __new__(cls, default_keys: Optional[List[str]] = None):
         instance = super().__new__(cls, default_keys)
         if default_keys is None:
             default_keys = []
-        instance.data = [key for key in default_keys]
+        instance.data = [key for key in default_keys]  # type: ignore
         return instance
 
     @classmethod
-    def __modify_schema__(cls, field_schema):
-        field_schema.update({"x-parameter": "dict"})
+    def __schema_type_properties__(cls) -> dict:
+        return {"x-parameter": "dict"}
 
 
 class TextParam(str):
     @classmethod
-    def __modify_schema__(cls, field_schema):
-        field_schema.update({"x-parameter": "text"})
+    def __schema_type_properties__(cls) -> dict:
+        return {"x-parameter": "text", "type": "string"}
 
 
 class BinaryParam(int):
     def __new__(cls, value: bool = False):
         instance = super().__new__(cls, int(value))
-        instance.default = value
+        instance.default = value  # type: ignore
         return instance
 
     @classmethod
-    def __modify_schema__(cls, field_schema):
-        field_schema.update(
-            {
-                "x-parameter": "bool",
-                "type": "boolean",
-            }
-        )
+    def __schema_type_properties__(cls) -> dict:
+        return {
+            "x-parameter": "bool",
+            "type": "boolean",
+        }
 
 
 class IntParam(int):
     def __new__(cls, default: int = 6, minval: float = 1, maxval: float = 10):
         instance = super().__new__(cls, default)
-        instance.minval = minval
-        instance.maxval = maxval
+        instance.minval = minval  # type: ignore
+        instance.maxval = maxval  # type: ignore
         return instance
 
     @classmethod
-    def __modify_schema__(cls, field_schema):
-        field_schema.update(
-            {
-                "x-parameter": "int",
-                "type": "integer",
-                "minimum": 1,
-                "maximum": 10,
-            }
-        )
+    def __schema_type_properties__(cls) -> dict:
+        return {"x-parameter": "int", "type": "integer"}
 
 
 class FloatParam(float):
     def __new__(cls, default: float = 0.5, minval: float = 0.0, maxval: float = 1.0):
         instance = super().__new__(cls, default)
-        instance.minval = minval
-        instance.maxval = maxval
+        instance.default = default  # type: ignore
+        instance.minval = minval  # type: ignore
+        instance.maxval = maxval  # type: ignore
         return instance
 
     @classmethod
-    def __modify_schema__(cls, field_schema):
-        field_schema.update(
-            {
-                "x-parameter": "float",
-                "type": "number",
-                "minimum": 0.0,
-                "maximum": 1.0,
-            }
-        )
+    def __schema_type_properties__(cls) -> dict:
+        return {"x-parameter": "float", "type": "number"}
 
 
 class MultipleChoiceParam(str):
-    def __new__(cls, default: str = None, choices: List[str] = None):
-        if type(default) is list:
+    def __new__(
+        cls, default: Optional[str] = None, choices: Optional[List[str]] = None
+    ):
+        if default is not None and type(default) is list:
             raise ValueError(
                 "The order of the parameters for MultipleChoiceParam is wrong! It's MultipleChoiceParam(default, choices) and not the opposite"
             )
-        if default is None and choices:
+
+        if not default and choices is not None:
             # if a default value is not provided,
-            # uset the first value in the choices list
+            # set the first value in the choices list
             default = choices[0]
 
         if default is None and not choices:
@@ -112,24 +120,53 @@ class MultipleChoiceParam(str):
             raise ValueError("You must provide either a default value or choices")
 
         instance = super().__new__(cls, default)
-        instance.choices = choices
-        instance.default = default
+        instance.choices = choices  # type: ignore
+        instance.default = default  # type: ignore
         return instance
 
     @classmethod
-    def __modify_schema__(cls, field_schema: dict[str, Any]):
-        field_schema.update(
-            {
-                "x-parameter": "choice",
-                "type": "string",
-                "enum": [],
-            }
-        )
+    def __schema_type_properties__(cls) -> dict:
+        return {"x-parameter": "choice", "type": "string", "enum": []}
+
+
+class GroupedMultipleChoiceParam(str):
+    def __new__(
+        cls,
+        default: Optional[str] = None,
+        choices: Optional[Dict[str, List[str]]] = None,
+    ):
+        if choices is None:
+            choices = {}
+        if default and not any(
+            default in choice_list for choice_list in choices.values()
+        ):
+            if not choices:
+                print(
+                    f"Warning: Default value {default} provided but choices are empty."
+                )
+            else:
+                raise ValueError(
+                    f"Default value {default} is not in the provided choices"
+                )
+
+        if not default:
+            default_selected_choice = next(
+                (choices for choices in choices.values()), None
+            )
+            if default_selected_choice:
+                default = default_selected_choice[0]
 
+        instance = super().__new__(cls, default)
+        instance.choices = choices  # type: ignore
+        instance.default = default  # type: ignore
+        return instance
 
-class Message(BaseModel):
-    role: str
-    content: str
+    @classmethod
+    def __schema_type_properties__(cls) -> dict:
+        return {
+            "x-parameter": "grouped_choice",
+            "type": "string",
+        }
 
 
 class MessagesInput(list):
@@ -144,30 +181,397 @@ class MessagesInput(list):
 
     """
 
-    def __new__(cls, messages: List[Dict[str, str]] = None):
-        instance = super().__new__(cls, messages)
-        instance.default = messages
+    def __new__(cls, messages: List[Dict[str, str]] = []):
+        instance = super().__new__(cls)
+        instance.default = messages  # type: ignore
         return instance
 
     @classmethod
-    def __modify_schema__(cls, field_schema: dict[str, Any]):
-        field_schema.update({"x-parameter": "messages", "type": "array"})
+    def __schema_type_properties__(cls) -> dict:
+        return {"x-parameter": "messages", "type": "array"}
 
 
 class FileInputURL(HttpUrl):
+    def __new__(cls, url: str):
+        instance = super().__new__(cls, url)
+        instance.default = url  # type: ignore
+        return instance
+
     @classmethod
-    def __modify_schema__(cls, field_schema: Dict[str, Any]) -> None:
-        field_schema.update({"x-parameter": "file_url", "type": "string"})
+    def __schema_type_properties__(cls) -> dict:
+        return {"x-parameter": "file_url", "type": "string"}
 
 
 class Context(BaseModel):
-    class Config:
-        extra = Extra.allow
+    model_config = ConfigDict(extra="allow")
 
     def to_json(self):
-        return self.json()
+        return self.model_dump()
 
     @classmethod
     def from_json(cls, json_str: str):
         data = json.loads(json_str)
         return cls(**data)
+
+
+class ReferencesResponse(BaseModel):
+    app_id: Optional[str] = None
+    app_slug: Optional[str] = None
+    variant_id: Optional[str] = None
+    variant_slug: Optional[str] = None
+    variant_version: Optional[int] = None
+    environment_id: Optional[str] = None
+    environment_slug: Optional[str] = None
+    environment_version: Optional[int] = None
+
+    def __str__(self):
+        return str(self.model_dump(exclude_none=True))
+
+
+class LifecyclesResponse(ReferencesResponse):
+    committed_at: Optional[str] = None
+    committed_by: Optional[str] = None
+    committed_by_id: Optional[str] = None
+    deployed_at: Optional[str] = None
+    deployed_by: Optional[str] = None
+    deployed_by_id: Optional[str] = None
+
+    def __str__(self):
+        return self.model_dump_json(indent=4)
+
+    def __repr__(self):
+        return self.__str__()
+
+
+class ConfigurationResponse(LifecyclesResponse):
+    params: Dict[str, Any]
+
+
+class DeploymentResponse(LifecyclesResponse):
+    pass
+
+
+class Prompt(BaseModel):
+    temperature: float
+    model: str
+    max_tokens: int
+    prompt_system: str
+    prompt_user: str
+    top_p: float
+    frequency_penalty: float
+    presence_penalty: float
+
+
+# -----------------------------------------------------
+# New Prompt model
+# -----------------------------------------------------
+
+
+class ToolCall(BaseModel):
+    id: str
+    type: Literal["function"] = "function"
+    function: Dict[str, str]
+
+
+class Message(BaseModel):
+    role: Literal["system", "user", "assistant", "tool", "function"]
+    content: Optional[str] = None
+    name: Optional[str] = None
+    tool_calls: Optional[List[ToolCall]] = None
+    tool_call_id: Optional[str] = None
+
+
+class ResponseFormatText(BaseModel):
+    type: Literal["text"]
+    """The type of response format being defined: `text`"""
+
+
+class ResponseFormatJSONObject(BaseModel):
+    type: Literal["json_object"]
+    """The type of response format being defined: `json_object`"""
+
+
+class JSONSchema(BaseModel):
+    name: str
+    """The name of the response format."""
+    description: Optional[str] = None
+    """A description of what the response format is for."""
+    schema_: Optional[Dict[str, object]] = Field(alias="schema", default=None)
+    """The schema for the response format, described as a JSON Schema object."""
+    strict: Optional[bool] = None
+    """Whether to enable strict schema adherence."""
+
+    model_config = {
+        "populate_by_name": True,
+        "json_schema_extra": {"required": ["name", "schema"]},
+    }
+
+
+class ResponseFormatJSONSchema(BaseModel):
+    type: Literal["json_schema"]
+    """The type of response format being defined: `json_schema`"""
+    json_schema: JSONSchema
+
+
+ResponseFormat = Union[
+    ResponseFormatText, ResponseFormatJSONObject, ResponseFormatJSONSchema
+]
+
+
+class ModelConfig(BaseModel):
+    """Configuration for model parameters"""
+
+    model: str = MCField(
+        default="gpt-3.5-turbo",
+        choices=supported_llm_models,
+    )
+
+    temperature: Optional[float] = Field(
+        default=1,
+        ge=0.0,
+        le=2.0,
+        description="What sampling temperature to use, between 0 and 2. Higher values like 0.8 will make the output more random, while lower values like 0.2 will make it more focused and deterministic",
+    )
+    max_tokens: Optional[int] = Field(
+        default=-1,
+        ge=0,
+        description="The maximum number of tokens that can be generated in the chat completion",
+    )
+    top_p: Optional[float] = Field(
+        default=0.5,
+        ge=0.0,
+        le=1.0,
+        description="An alternative to sampling with temperature, called nucleus sampling, where the model considers the results of the tokens with top_p probability mass",
+    )
+    frequency_penalty: Optional[float] = Field(
+        default=0,
+        ge=-2.0,
+        le=2.0,
+        description="Number between -2.0 and 2.0. Positive values penalize new tokens based on their existing frequency in the text so far",
+    )
+    presence_penalty: Optional[float] = Field(
+        default=0,
+        ge=-2.0,
+        le=2.0,
+        description="Number between -2.0 and 2.0. Positive values penalize new tokens based on whether they appear in the text so far",
+    )
+    response_format: Optional[ResponseFormat] = Field(
+        default=None,
+        description="An object specifying the format that the model must output",
+    )
+    stream: Optional[bool] = Field(
+        default=None, description="If set, partial message deltas will be sent"
+    )
+    tools: Optional[List[Dict]] = Field(
+        default=None,
+        description="A list of tools the model may call. Currently, only functions are supported as a tool",
+    )
+    tool_choice: Optional[Union[Literal["none", "auto"], Dict]] = Field(
+        default=None, description="Controls which (if any) tool is called by the model"
+    )
+
+
+class PromptTemplateError(Exception):
+    """Base exception for all PromptTemplate errors"""
+
+    pass
+
+
+class InputValidationError(PromptTemplateError):
+    """Raised when input validation fails"""
+
+    def __init__(
+        self, message: str, missing: Optional[set] = None, extra: Optional[set] = None
+    ):
+        self.missing = missing
+        self.extra = extra
+        super().__init__(message)
+
+
+class TemplateFormatError(PromptTemplateError):
+    """Raised when template formatting fails"""
+
+    def __init__(self, message: str, original_error: Optional[Exception] = None):
+        self.original_error = original_error
+        super().__init__(message)
+
+
+class PromptTemplate(BaseModel):
+    """A template for generating prompts with formatting capabilities"""
+
+    messages: List[Message] = Field(
+        default=[Message(role="system", content=""), Message(role="user", content="")]
+    )
+    system_prompt: Optional[str] = None
+    user_prompt: Optional[str] = None
+    template_format: Literal["fstring", "jinja2", "curly"] = Field(
+        default="fstring",
+        description="Format type for template variables: fstring {var}, jinja2 {{ var }}, or curly {{var}}",
+    )
+    input_keys: Optional[List[str]] = Field(
+        default=None,
+        description="Optional list of input keys for validation. If not provided, any inputs will be accepted",
+    )
+    llm_config: ModelConfig = Field(
+        default_factory=ModelConfig,
+        description="Configuration for the model parameters",
+    )
+
+    model_config = {
+        "json_schema_extra": {
+            "x-parameters": {
+                "prompt": "true",
+            }
+        }
+    }
+
+    @root_validator(pre=True)
+    def init_messages(cls, values):
+        if "messages" not in values:
+            messages = []
+            if "system_prompt" in values and values["system_prompt"]:
+                messages.append(Message(role="system", content=values["system_prompt"]))
+            if "user_prompt" in values and values["user_prompt"]:
+                messages.append(Message(role="user", content=values["user_prompt"]))
+            if messages:
+                values["messages"] = messages
+        return values
+
+    def _format_with_template(self, content: str, kwargs: Dict[str, Any]) -> str:
+        """Internal method to format content based on template_format"""
+        try:
+            if self.template_format == "fstring":
+                return content.format(**kwargs)
+            elif self.template_format == "jinja2":
+                from jinja2 import Template, TemplateError
+
+                try:
+                    return Template(content).render(**kwargs)
+                except TemplateError as e:
+                    raise TemplateFormatError(
+                        f"Jinja2 template error in content: '{content}'. Error: {str(e)}",
+                        original_error=e,
+                    )
+            elif self.template_format == "curly":
+                import re
+
+                result = content
+                for key, value in kwargs.items():
+                    result = re.sub(r"\{\{" + key + r"\}\}", str(value), result)
+                if re.search(r"\{\{.*?\}\}", result):
+                    unreplaced = re.findall(r"\{\{(.*?)\}\}", result)
+                    raise TemplateFormatError(
+                        f"Unreplaced variables in curly template: {unreplaced}"
+                    )
+                return result
+            else:
+                raise TemplateFormatError(
+                    f"Unknown template format: {self.template_format}"
+                )
+        except KeyError as e:
+            key = str(e).strip("'")
+            raise TemplateFormatError(
+                f"Missing required variable '{key}' in template: '{content}'"
+            )
+        except Exception as e:
+            raise TemplateFormatError(
+                f"Error formatting template '{content}': {str(e)}", original_error=e
+            )
+
+    def format(self, **kwargs) -> "PromptTemplate":
+        """
+        Format the template with provided inputs.
+        Only validates against input_keys if they are specified.
+
+        Raises:
+            InputValidationError: If input validation fails
+            TemplateFormatError: If template formatting fails
+        """
+        # Validate inputs if input_keys is set
+        if self.input_keys is not None:
+            missing = set(self.input_keys) - set(kwargs.keys())
+            extra = set(kwargs.keys()) - set(self.input_keys)
+
+            error_parts = []
+            if missing:
+                error_parts.append(
+                    f"Missing required inputs: {', '.join(sorted(missing))}"
+                )
+            if extra:
+                error_parts.append(f"Unexpected inputs: {', '.join(sorted(extra))}")
+
+            if error_parts:
+                raise InputValidationError(
+                    " | ".join(error_parts),
+                    missing=missing if missing else None,
+                    extra=extra if extra else None,
+                )
+
+        new_messages = []
+        for i, msg in enumerate(self.messages):
+            if msg.content:
+                try:
+                    new_content = self._format_with_template(msg.content, kwargs)
+                except TemplateFormatError as e:
+                    raise TemplateFormatError(
+                        f"Error in message {i} ({msg.role}): {str(e)}",
+                        original_error=e.original_error,
+                    )
+            else:
+                new_content = None
+
+            new_messages.append(
+                Message(
+                    role=msg.role,
+                    content=new_content,
+                    name=msg.name,
+                    tool_calls=msg.tool_calls,
+                    tool_call_id=msg.tool_call_id,
+                )
+            )
+
+        return PromptTemplate(
+            messages=new_messages,
+            template_format=self.template_format,
+            llm_config=self.llm_config,
+            input_keys=self.input_keys,
+        )
+
+    def to_openai_kwargs(self) -> dict:
+        """Convert the prompt template to kwargs compatible with litellm/openai"""
+        kwargs = {
+            "model": self.llm_config.model,
+            "messages": [msg.dict(exclude_none=True) for msg in self.messages],
+        }
+
+        # Add optional parameters only if they are set
+        if self.llm_config.temperature is not None:
+            kwargs["temperature"] = self.llm_config.temperature
+
+        if self.llm_config.top_p is not None:
+            kwargs["top_p"] = self.llm_config.top_p
+
+        if self.llm_config.stream is not None:
+            kwargs["stream"] = self.llm_config.stream
+
+        if self.llm_config.max_tokens is not None:
+            kwargs["max_tokens"] = self.llm_config.max_tokens
+
+        if self.llm_config.frequency_penalty is not None:
+            kwargs["frequency_penalty"] = self.llm_config.frequency_penalty
+
+        if self.llm_config.presence_penalty is not None:
+            kwargs["presence_penalty"] = self.llm_config.presence_penalty
+
+        if self.llm_config.response_format:
+            kwargs["response_format"] = self.llm_config.response_format.dict(
+                by_alias=True
+            )
+
+        if self.llm_config.tools:
+            kwargs["tools"] = self.llm_config.tools
+            # Only set tool_choice if tools are present
+            if self.llm_config.tool_choice is not None:
+                kwargs["tool_choice"] = self.llm_config.tool_choice
+
+        return kwargs

agenta/sdk/utils/constants.py

@@ -0,0 +1 @@
+TRUTHY = {"true", "1", "t", "y", "yes", "on", "enable", "enabled"}

agenta/sdk/utils/{helper/openai_cost.py → costs.py}

@@ -49,6 +49,9 @@ MODEL_COST_PER_1K_TOKENS = {
     "gpt-35-turbo-16k-completion": 0.004,
     "gpt-35-turbo-16k-0613-completion": 0.004,
     # Others
+    "text-embedding-ada-002": 0.1,
+    "text-ada-002": 0.1,
+    "adav2": 0.1,
     "text-ada-001": 0.0004,
     "ada": 0.0004,
     "text-babbage-001": 0.0005,