versionhq 1.1.9.14__py3-none-any.whl → 1.1.10.3__py3-none-any.whl

versionhq/task/model.py

@@ -1,12 +1,13 @@
 import json
 import threading
+import datetime
 import uuid
 from concurrent.futures import Future
 from hashlib import md5
-from typing import Any, Dict, List, Set, Optional, Tuple, Callable, Type
+from typing import Any, Dict, List, Set, Optional, Tuple, Callable, Type, TypeVar
 from typing_extensions import Annotated, Self
 
-from pydantic import UUID4, BaseModel, Field, PrivateAttr, field_validator, model_validator, create_model, InstanceOf
+from pydantic import UUID4, BaseModel, Field, PrivateAttr, field_validator, model_validator, create_model, InstanceOf, field_validator
 from pydantic_core import PydanticCustomError
 
 from versionhq._utils.process_config import process_config
@@ -18,22 +19,106 @@ from versionhq._utils.logger import Logger
 
 class ResponseField(BaseModel):
     """
-    Field class to use in the response schema for the JSON response.
+    A class to store the response format and schema that will cascade to the LLM.
+    The `config` field can store additional params:
+    https://community.openai.com/t/official-documentation-for-supported-schemas-for-response-format-parameter-in-calls-to-client-beta-chats-completions-parse/932422/3
     """
 
-    title: str = Field(default=None)
-    type: Type = Field(default=str)
+    title: str = Field(default=None, description="title of the field")
+    data_type: Type = Field(default=None)
+    items: Optional[Type] = Field(default=None, description="store data type of the array items")
+    properties: Optional[List[BaseModel]] = Field(default=None, description="store dict items in ResponseField format")
     required: bool = Field(default=True)
+    nullable: bool = Field(default=False)
+    config: Optional[Dict[str, Any]] = Field(default=None, description="additional rules")
 
 
-    def _annotate(self, value: Any) -> Annotated:
+    @model_validator(mode="after")
+    def validate_instance(self) -> Self:
         """
-        Address `create_model`
+        Validate the model instance based on the given `data_type`. (An array must have `items`, dict must have properties.)
         """
-        return Annotated[self.type, value] if isinstance(value, self.type) else Annotated[str, str(value)]
+
+        if self.data_type is list and self.items is None:
+            self.items = str
+
+        if self.data_type is dict or (self.data_type is list and self.items is dict):
+            if self.properties is None:
+                raise PydanticCustomError("missing_properties", "The dict type has to set the properties.", {})
+
+            else:
+                for item in self.properties:
+                    if not isinstance(item, ResponseField):
+                        raise PydanticCustomError("invalid_properties", "Properties field must input in ResponseField format.", {})
+
+        return self
+
+
+    def _format_props(self) -> Dict[str, Any]:
+        """
+        Structure valid properties. We accept 2 nested objects.
+        """
+        from versionhq.llm.llm_vars import SchemaType
+
+        schema_type = SchemaType(type=self.data_type).convert()
+        props: Dict[str, Any] = {}
+
+        if self.data_type is list and self.items is not dict:
+            props = {
+                "type": schema_type,
+                "items": { "type": SchemaType(type=self.items).convert() },
+            }
+
+        elif self.data_type is list and self.items is dict:
+            nested_p, nested_r = dict(), list()
+
+            if self.properties:
+                for item in self.properties:
+                    nested_p.update(**item._format_props())
+
+                    if item.required:
+                        nested_r.append(item.title)
+
+            props = {
+                "type": schema_type,
+                "items": {
+                    "type": SchemaType(type=self.items).convert(),
+                    "properties": nested_p,
+                    "required": nested_r,
+                    "additionalProperties": False
+                }
+            }
+
+        elif self.data_type is dict:
+            p, r = dict(), list()
+
+            if self.properties:
+                for item in self.properties:
+                    p.update(**item._format_props())
+
+                    # if item.required:
+                    r.append(item.title)
+
+            props = {
+                "type": schema_type,
+                "properties": p,
+                "required": r,
+                "additionalProperties": False
+            }
+
+        else:
+            props = {
+                "type": schema_type,
+                "nullable": self.nullable,
+            }
+
+        return { self.title: { **props, **self.config }} if self.config else { self.title: props }
 
 
     def _convert(self, value: Any) -> Any:
+        """
+        Convert the given value to the ideal data type.
+        """
         try:
             if self.type is Any:
                 pass
@@ -42,7 +127,9 @@ class ResponseField(BaseModel):
             elif self.type is float:
                 return float(value)
             elif self.type is list or self.type is dict:
-                return json.loads(value)
+                return json.loads(eval(str(value)))
+            elif self.type is str:
+                return str(value)
             else:
                 return value
         except:
@@ -50,6 +137,9 @@ class ResponseField(BaseModel):
 
 
     def create_pydantic_model(self, result: Dict, base_model: InstanceOf[BaseModel] | Any) -> Any:
+        """
+        Create a Pydantic model from the given result
+        """
         for k, v in result.items():
             if k is not self.title:
                 pass
@@ -61,17 +151,23 @@ class ResponseField(BaseModel):
         return base_model
 
 
+    def _annotate(self, value: Any) -> Annotated:
+        """
+        Address Pydantic's `create_model`
+        """
+        return Annotated[self.type, value] if isinstance(value, self.type) else Annotated[str, str(value)]
+
+
 
 class TaskOutput(BaseModel):
     """
-    Store the final output of the task in TaskOutput class.
-    Depending on the task output format, use `raw`, `pydantic`, `json_dict` accordingly.
+    A class to store the final output of the given task in raw (string), json_dict, and pydantic class formats.
     """
 
     task_id: UUID4 = Field(default_factory=uuid.uuid4, frozen=True, description="store Task ID")
     raw: str = Field(default="", description="Raw output of the task")
     json_dict: Dict[str, Any] = Field(default=None, description="`raw` converted to dictionary")
-    pydantic: Optional[Any] = Field(default=None, description="`raw` converted to the abs. pydantic model")
+    pydantic: Optional[Any] = Field(default=None)
     tool_output: Optional[Any] = Field(default=None, description="store tool result when the task takes tool output as its final output")
 
     def __str__(self) -> str:
@@ -108,14 +204,15 @@ class TaskOutput(BaseModel):
 
 class Task(BaseModel):
     """
-    Task to be executed by the agent or the team.
-    Each task must have a description and at least one expected output format either Pydantic, Raw, or JSON, with necessary fields in ResponseField.
-    Then output will be stored in TaskOutput class.
+    Task to be executed by agents or teams.
+    Each task must have a description.
+    Default response is JSON string that strictly follows `response_fields` - and will be stored in TaskOuput.raw / json_dict.
+    When `pydantic_output` is provided, we prioritize them and store raw (json string), json_dict, pydantic in the TaskOutput class.
     """
 
     __hash__ = object.__hash__
+    _logger: Logger = PrivateAttr(default_factory=lambda: Logger(verbose=True))
     _original_description: str = PrivateAttr(default=None)
-    _logger: Logger = PrivateAttr()
     _task_output_handler = TaskOutputStorageHandler()
     config: Optional[Dict[str, Any]] = Field(default=None, description="values to set on Task class")
 
@@ -124,11 +221,13 @@ class Task(BaseModel):
     description: str = Field(description="Description of the actual task")
 
     # output
-    expected_output_json: bool = Field(default=True)
-    expected_output_pydantic: bool = Field(default=False)
-    output_field_list: List[ResponseField] = Field(
+    pydantic_custom_output: Optional[Any] = Field(
+        default=None,
+        description="store a custom Pydantic class that will be passed to the model as a response format."
+    )
+    response_fields: List[ResponseField] = Field(
         default_factory=list,
-        description="provide output key and data type. this will be cascaded to the agent via task.prompt()"
+        description="store the list of ResponseFields to create the response format"
     )
     output: Optional[TaskOutput] = Field(default=None, description="store the final task output in TaskOutput class")
 
@@ -139,7 +238,7 @@ class Task(BaseModel):
     # tool usage
     tools: Optional[List[ToolSet | Tool | Any]] = Field(default_factory=list, description="tools that the agent can use aside from their tools")
     can_use_agent_tools: bool = Field(default=False, description="whether the agent can use their own tools when executing the task")
-    take_tool_res_as_final: bool = Field(default=False, description="when set True, tools res will be stored in the `TaskOutput`")
+    tool_res_as_final: bool = Field(default=False, description="when set True, tools res will be stored in the `TaskOutput`")
 
     # execution rules
     allow_delegation: bool = Field(default=False, description="ask other agents for help and run the task instead")
@@ -151,11 +250,12 @@ class Task(BaseModel):
     processed_by_agents: Set[str] = Field(default_factory=set, description="store responsible agents' roles")
     tools_errors: int = 0
     delegations: int = 0
+    execution_span_in_sec: int = 0
 
 
     @model_validator(mode="before")
     @classmethod
-    def process_model_config(cls, values: Dict[str, Any]) -> None:
+    def process_config(cls, values: Dict[str, Any]) -> None:
         return process_config(values_to_update=values, model_class=cls)
 
 
@@ -175,16 +275,16 @@ class Task(BaseModel):
         return self
 
 
-    @model_validator(mode="after")
-    def set_attributes_based_on_config(self) -> Self:
-        """
-        Set attributes based on the agent configuration.
-        """
+    # @model_validator(mode="after")
+    # def set_attributes_based_on_config(self) -> Self:
+    #     """
+    #     Set attributes based on the task configuration.
+    #     """
 
-        if self.config:
-            for key, value in self.config.items():
-                setattr(self, key, value)
-        return self
+    #     if self.config:
+    #         for key, value in self.config.items():
+    #             setattr(self, key, value)
+    #     return self
 
 
     @model_validator(mode="after")
@@ -211,85 +311,165 @@ class Task(BaseModel):
         return self
 
 
-    def prompt(self, customer: str = None, product_overview: str = None) -> str:
+    def _draft_output_prompt(self, model_provider: str) -> str:
+        """
+        Draft prompts on the output format by converting `
+        """
+
+        output_prompt = ""
+
+        if self.pydantic_custom_output:
+            output_prompt = f"""
+Your response MUST STRICTLY follow the given repsonse format:
+JSON schema: {str(self.pydantic_custom_output)}
+"""
+
+        elif self.response_fields:
+            output_prompt, output_formats_to_follow = "", dict()
+            response_format = str(self._structure_response_format(model_provider=model_provider))
+            for item in self.response_fields:
+                if item:
+                    output_formats_to_follow[item.title] = f"<Return your answer in {item.data_type.__name__}>"
+
+            output_prompt = f"""
+Your response MUST be a valid JSON string that strictly follows the response format. Use double quotes for all keys and string values. Do not use single quotes, trailing commas, or any other non-standard JSON syntax.
+Response format: {response_format}
+Ref. Output image: {output_formats_to_follow}
+"""
+
+        else:
+            output_prompt = "Return your response as a valid JSON string, enclosed in double quotes. Do not use single quotes, trailing commas, or other non-standard JSON syntax."
+
+        return output_prompt
+
+
+    def prompt(self, model_provider: str = None, customer: str = None, product_overview: str = None) -> str:
         """
         Format the task prompt and cascade it to the agent.
         When the task has context, add context prompting of all the tasks in the context.
         When we have cusotmer/product info, add them to the prompt.
         """
 
-        task_slices = [self.description, f"{self.output_prompt}", f"Take the following context into consideration: "]
+        output_prompt = self._draft_output_prompt(model_provider=model_provider)
+        task_slices = [self.description, output_prompt,]
 
         if self.context:
             context_outputs = "\n".join([task.output.context_prompting() if hasattr(task, "output") else "" for task in self.context])
-            task_slices.insert(len(task_slices),  context_outputs)
+            task_slices.insert(len(task_slices), f"Consider the following context when responding: {context_outputs}")
+
+        if self.prompt_context:
+            task_slices.insert(len(task_slices), f"Consider the following context when responding: {self.prompt_context}")
 
         if customer:
-            task_slices.insert(len(task_slices),  f"Customer overview: {customer}")
+            task_slices.insert(len(task_slices), f"Customer to address: {customer}")
 
         if product_overview:
-            task_slices.insert(len(task_slices), f"Product overview: {product_overview}")
-
-        if self.prompt_context:
-            task_slices.insert(len(task_slices), self.prompt_context)
+            task_slices.insert(len(task_slices), f"Product to promote: {product_overview}")
 
         return "\n".join(task_slices)
 
 
-    def _create_json_output(self, raw_result: str) -> Dict[str, Any]:
+    def _get_output_format(self) -> TaskOutputFormat:
+        if self.output_json == True:
+            return TaskOutputFormat.JSON
+        if self.output_pydantic == True:
+            return TaskOutputFormat.PYDANTIC
+        return TaskOutputFormat.RAW
+
+
+    def _structure_response_format(self, data_type: str = "object", model_provider: str = "gemini") -> Dict[str, Any] | None:
+        """
+        Structure a response format either from`response_fields` or `pydantic_custom_output`.
+        1 nested item is accepted.
+        """
+
+        from versionhq.task.structured_response import StructuredOutput
+
+        response_format: Dict[str, Any] = None
+
+        if self.response_fields:
+            properties, required_fields = {}, []
+            for i, item in enumerate(self.response_fields):
+                if item:
+                    if item.data_type is dict:
+                        properties.update(item._format_props())
+                    else:
+                        properties.update(item._format_props())
+
+                    required_fields.append(item.title)
+
+            response_schema = {
+                "type": "object",
+                "properties": properties,
+                "required": required_fields,
+                "additionalProperties": False,
+            }
+
+            response_format = {
+                "type": "json_schema",
+                "json_schema": { "name": "outcome", "schema": response_schema }
+            }
+
+
+        elif self.pydantic_custom_output:
+            response_format = StructuredOutput(response_format=self.pydantic_custom_output)._format()
+
+        return response_format
+
+
+    def _create_json_output(self, raw: str) -> Dict[str, Any]:
         """
-        Create json (dict) output from the raw result.
+        Create json (dict) output from the raw output and `response_fields` information.
         """
-        import ast
 
-        output_json_dict: Dict[str, Any] = dict()
+        if raw is None or raw == "":
+            self._logger.log(level="error", message="The model returned an empty response. Returning an empty dict.", color="yellow")
+            output = { "output": "n.a." }
+            return output
 
         try:
-            raw_result = raw_result.replace("{'", '{"').replace("{ '", '{"').replace("': '", '": "').replace("'}", '"}').replace("' }", '"}').replace("', '", '", "').replace("['", '["').replace("[ '", '[ "').replace("']", '"]').replace("' ]", '" ]').replace("{\n'", '{"').replace("{\'", '{"')
-            r = json.dumps(eval(str(raw_result)))
-            output_json_dict = json.loads(r)
+            r = str(raw).replace("true", "True").replace("false", "False")
+            j = json.dumps(eval(r))
+            output = json.loads(j)
+            if isinstance(output, dict):
+                return output
 
-            if isinstance(output_json_dict, str):
-                output_json_dict = ast.literal_eval(r)
+            else:
+                r = str(raw).replace("{'", '{"').replace("{ '", '{"').replace("': '", '": "').replace("'}", '"}').replace("' }", '"}').replace("', '", '", "').replace("['", '["').replace("[ '", '[ "').replace("']", '"]').replace("' ]", '" ]').replace("{\n'", '{"').replace("{\'", '{"').replace("true", "True").replace("false", "False")
+                j = json.dumps(eval(r))
+                output = json.loads(j)
+
+                if isinstance(output, dict):
+                    return output
+
+                else:
+                    import ast
+                    output = ast.literal_eval(r)
+                    return output if isinstance(output, dict) else { "output": str(r) }
 
         except:
-            output_json_dict = { "output": raw_result }
+            output = { "output": str(raw) }
+            return output
 
-        return output_json_dict
 
 
-    def _create_pydantic_output(self, output_json_dict: Dict[str, Any], raw_result: Any = None) -> Optional[Any]:
+    def _create_pydantic_output(self, raw: str = None, json_dict: Dict[str, Any] = None) -> InstanceOf[BaseModel]:
         """
-        Create pydantic output from the `raw` result.
+        Create pydantic output from raw or json_dict output.
         """
 
-        output_pydantic = None
-        if isinstance(raw_result, BaseModel):
-            output_pydantic = raw_result
+        output_pydantic = self.pydantic_custom_output
 
-        elif hasattr(output_json_dict, "output"):
-            output_pydantic = create_model("PydanticTaskOutput", output=output_json_dict["output"], __base__=BaseModel)
-
-        else:
-            output_pydantic = create_model("PydanticTaskOutput", __base__=BaseModel)
-            try:
-                for item in self.output_field_list:
-                    value = output_json_dict[item.title] if hasattr(output_json_dict, item.title) else None
-                    if value and type(value) is not item.type:
-                        value = item._convert(value)
-                    setattr(output_pydantic, item.title, value)
-            except:
-                setattr(output_pydantic, "output", output_json_dict)
+        try:
+            json_dict = json_dict if json_dict else self._create_json_output(raw=raw)
 
-        return output_pydantic
+            for k, v in json_dict.items():
+                setattr(output_pydantic, k, v)
 
+        except:
+            pass
 
-    def _get_output_format(self) -> TaskOutputFormat:
-        if self.output_json == True:
-            return TaskOutputFormat.JSON
-        if self.output_pydantic == True:
-            return TaskOutputFormat.PYDANTIC
-        return TaskOutputFormat.RAW
+        return output_pydantic
 
 
     def interpolate_inputs(self, inputs: Dict[str, Any]) -> None:
@@ -336,15 +516,23 @@ class Task(BaseModel):
 
     def _execute_core(self, agent, context: Optional[str]) -> TaskOutput:
         """
-        Run the core execution logic of the task.
-        To speed up the process, when the format is not expected to return, we will skip the conversion process.
-        When the task is allowed to delegate to another agent, we will select a responsible one in order of manager > peer_agent > anoymous agent.
+        Execute the given task with the given agent.
+        Handle 1. agent delegation, 2. tools, 3. context to consider, and 4. callbacks
         """
+
         from versionhq.agent.model import Agent
         from versionhq.team.model import Team
 
         self.prompt_context = context
         task_output: InstanceOf[TaskOutput] = None
+        tool_output: str | list = None
+        task_tools: List[InstanceOf[Tool] | InstanceOf[ToolSet] | Type[Tool]] = []
+        started_at = datetime.datetime.now()
+
+        if self.tools:
+            for item in self.tools:
+                if isinstance(item, ToolSet) or isinstance(item, Tool) or type(item) == Tool:
+                    task_tools.append(item)
 
         if self.allow_delegation:
             agent_to_delegate = None
@@ -363,39 +551,33 @@ class Task(BaseModel):
             agent = agent_to_delegate
             self.delegations += 1
 
-        if self.take_tool_res_as_final == True:
-            output = agent.execute_task(task=self, context=context)
-            task_output = TaskOutput(task_id=self.id, tool_output=output)
 
-        else:
-            output_raw, output_json_dict, output_pydantic = agent.execute_task(task=self, context=context), None, None
+        if self.tool_res_as_final == True:
+            tool_output = agent.execute_task(task=self, context=context, task_tools=task_tools)
+            task_output = TaskOutput(task_id=self.id, tool_output=tool_output)
 
-            if self.expected_output_json == True:
-                output_json_dict = self._create_json_output(raw_result=output_raw)
+        else:
+            raw_output = agent.execute_task(task=self, context=context, task_tools=task_tools)
+            json_dict_output = self._create_json_output(raw=raw_output)
+            if "outcome" in json_dict_output:
+                json_dict_output = self._create_json_output(raw=str(json_dict_output["outcome"]))
 
-            if self.expected_output_pydantic == True:
-                output_pydantic = self._create_pydantic_output(output_json_dict=output_json_dict)
+            pydantic_output = self._create_pydantic_output(raw=raw_output, json_dict=json_dict_output) if self.pydantic_custom_output else None
 
             task_output = TaskOutput(
                 task_id=self.id,
-                raw=output_raw,
-                pydantic=output_pydantic,
-                json_dict=output_json_dict
+                raw=raw_output,
+                pydantic=pydantic_output,
+                json_dict=json_dict_output
             )
 
         self.output = task_output
         self.processed_by_agents.add(agent.role)
 
-        # self._set_end_execution_time(start_time)
-
         if self.callback:
             self.callback({ **self.callback_kwargs, **self.output.__dict__ })
 
-        # if self._execution_span:
-        #     # self._telemetry.task_ended(self._execution_span, self, agent.team)
-        #     self._execution_span = None
-
-        # if self.output_file:
+        # if self.output_file: ## disabled for now
         #     content = (
         #         json_output
         #         if json_output
@@ -403,6 +585,9 @@ class Task(BaseModel):
         #     )
         #     self._save_file(content)
 
+        ended_at = datetime.datetime.now()
+        self.execution_span_in_sec = (ended_at - started_at).total_seconds()
+
         return task_output
 
 
@@ -414,48 +599,9 @@ class Task(BaseModel):
         self._task_output_handler.update(task=self, task_index=task_index, was_replayed=was_replayed, inputs=inputs)
 
 
-    @property
-    def output_prompt(self) -> str:
-        """
-        Draft prompts on the output format by converting `output_field_list` to dictionary.
-        """
-
-        output_prompt, output_formats_to_follow = "", dict()
-        for item in self.output_field_list:
-            output_formats_to_follow[item.title] = f"<Return your answer in {item.type.__name__}>"
-
-        output_prompt = f"""
-Output only valid JSON conforming to the specified format. Use double quotes for all keys and string values. Do not use single quotes, trailing commas, or other non-standard JSON syntax.
-Specified format: {output_formats_to_follow}
-        """
-        return output_prompt
-
-
-    @property
-    def expected_output_formats(self) -> List[str | TaskOutputFormat]:
-        """
-        Return output formats in list with the ENUM item.
-        `TaskOutputFormat.RAW` is set as default.
-        """
-        outputs = [TaskOutputFormat.RAW.value,]
-        if self.expected_output_json:
-            outputs.append(TaskOutputFormat.JSON.value)
-        if self.expected_output_pydantic:
-            outputs.append(TaskOutputFormat.PYDANTIC.value)
-        return outputs
-
-
     @property
     def key(self) -> str:
-        output_format = (
-            TaskOutputFormat.JSON
-            if self.expected_output_json == True
-            else (
-                TaskOutputFormat.PYDANTIC
-                if self.expected_output_pydantic == True
-                else TaskOutputFormat.RAW
-            )
-        )
+        output_format = TaskOutputFormat.JSON if self.response_fields else TaskOutputFormat.PYDANTIC if self.pydantic_custom_output is not None else TaskOutputFormat.RAW
         source = [self.description, output_format]
         return md5("|".join(source).encode(), usedforsecurity=False).hexdigest()
 
@@ -463,15 +609,13 @@ Specified format: {output_formats_to_follow}
     @property
     def summary(self) -> str:
         return f"""
-        Task ID: {str(self.id)}
-        "Description": {self.description}
-        "Prompt": {self.output_prompt}
-        "Tools": {", ".join([tool.name for tool in self.tools])}
+Task ID: {str(self.id)}
+"Description": {self.description}
+"Tools": {", ".join([tool.name for tool in self.tools])}
         """
 
 
 
-
 class ConditionalTask(Task):
     """
     A task that can be conditionally executed based on the output of another task.