versionhq 1.2.4.3__py3-none-any.whl → 1.2.4.6__py3-none-any.whl

versionhq/clients/product/model.py

@@ -1,11 +1,11 @@
 import uuid
-from abc import ABC, abstractmethod
-from typing import Any, Dict, List, Callable, Type, Optional, get_args, get_origin
+from abc import ABC
+from typing import Optional, List
 
-from pydantic import UUID4, InstanceOf, BaseModel, ConfigDict, Field, create_model, field_validator, model_validator
+from pydantic import UUID4, BaseModel, Field, field_validator, model_validator
 from pydantic_core import PydanticCustomError
 
-from versionhq.tool.composio_tool_vars import ComposioAppName
+from versionhq.tool.composio.params import ComposioAppName
 
 
 class ProductProvider(ABC, BaseModel):

versionhq/clients/workflow/model.py

@@ -10,7 +10,7 @@ from versionhq.agent.model import Agent
 from versionhq.agent_network.model import AgentNetwork
 from versionhq.clients.product.model import Product
 from versionhq.clients.customer.model import Customer
-from versionhq.tool.composio_tool_vars import ComposioAppName
+from versionhq.tool.composio.params import ComposioAppName
 
 
 class MessagingComponent(ABC, BaseModel):

versionhq/llm/llm_vars.py

@@ -27,8 +27,9 @@ MODELS = {
         "o1-preview",
     ],
     "gemini": [
-        "gemini/gemini-1.5-flash",
-        "gemini/gemini-1.5-pro",
+        "gemini/gemini-2.0-flash",
+        "gemini/gemini-2.0-flash-thinking-exp",
+        "gemini/gemini-2.0-flash-lite-preview-02-05",
         "gemini/gemini-2.0-flash-exp",
     ],
     "anthropic": [
@@ -75,10 +76,10 @@ ENV_VARS = {
     "huggingface": ["HUGGINGFACE_API_KEY", ],
     "bedrock": ["AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY", "AWS_REGION_NAME"],
     "sagemaker": ["AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY", "AWS_REGION_NAME"],
+    "azure_ai": ["AZURE_AI_API_KEY", "AZURE_AI_API_BASE"],
 }
 
 
-
 """
 Max input token size by the model.
 """
@@ -90,9 +91,10 @@ LLM_CONTEXT_WINDOW_SIZES = {
     "o1-preview": 128000,
     "o1-mini": 128000,
 
-    "gemini/gemini-1.5-flash": 1048576,
-    "gemini/gemini-1.5-pro": 2097152,
     "gemini/gemini-2.0-flash-exp": 1048576,
+    "gemini/gemini-2.0-flash": 1048576,
+    "gemini/gemini-2.0-flash-thinking-exp": 1048576,
+    "gemini/gemini-2.0-flash-lite-preview-02-05": 1048576,
 
     "claude-3-7-sonnet-latest": 200000,
     "claude-3-5-haiku-latest": 200000,
@@ -128,7 +130,6 @@ LLM_CONTEXT_WINDOW_SIZES = {
 """
 Params for litellm.completion().
 """
-
 PARAMS = {
     "litellm": [
         "api_base",

versionhq/llm/model.py

@@ -228,6 +228,9 @@ class LLM(BaseModel):
             if self.context_window_size and valid_config["max_tokens"] > self.context_window_size:
                 valid_config["max_tokens"] = self.context_window_size
 
+        if "model" in valid_config:
+            self.model = valid_config.pop("model")
+
         self.llm_config = valid_config
         return valid_config
 
@@ -389,7 +392,6 @@ class LLM(BaseModel):
                     self._usages.append(res["usage"])
                     return res.choices[0].message.content
 
-
             except JSONSchemaValidationError as e:
                 self._logger.log(level="error", message="Raw Response: {}".format(e.raw_response), color="red")
                 raise e

versionhq/storage/task_output_storage.py

@@ -147,8 +147,8 @@ class TaskOutputStorageHandler:
             description=str(task.description),
             raw=str(task.output.raw),
             responsible_agents=str(task.processed_agents),
-            tokens=task._usage.total_tokens,
-            latency=task._usage.latency,
+            tokens=task.output.usage.total_tokens,
+            latency=task.output.usage.latency,
             score=task.output.aggregate_score if task.output.aggregate_score else "None",
         )
         self.storage.add(task=task, output=output_to_store, inputs=inputs)

versionhq/task/model.py

@@ -3,7 +3,7 @@ import threading
 import datetime
 import uuid
 import inspect
-import enum
+from enum import IntEnum
 from concurrent.futures import Future
 from hashlib import md5
 from typing import Any, Dict, List, Set, Optional, Callable, Type
@@ -14,11 +14,15 @@ from pydantic_core import PydanticCustomError
 
 import versionhq as vhq
 from versionhq.task.evaluation import Evaluation, EvaluationItem
-from versionhq.tool.model import Tool, ToolSet
+from versionhq.tool.model import Tool, ToolSet, BaseTool
+from versionhq.tool.rag_tool import RagTool
+from versionhq.tool.gpt.web_search import GPTToolWebSearch
+from versionhq.tool.gpt.file_search import GPTToolFileSearch
+from versionhq.tool.gpt.cup import GPTToolCUP
 from versionhq._utils import process_config, Logger, UsageMetrics, ErrorType
 
 
-class TaskExecutionType(enum.Enum):
+class TaskExecutionType(IntEnum):
     """
     Enumeration to store task execution types of independent tasks without dependencies.
     """
@@ -174,14 +178,15 @@ class TaskOutput(BaseModel):
     """
     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")
+    task_id: UUID4 = Field(default_factory=uuid.uuid4)
+    raw: str = Field(default="")
+    json_dict: Dict[str, Any] = Field(default=None)
     pydantic: Optional[Any] = Field(default=None)
     tool_output: Optional[Any] = Field(default=None, description="stores tool result when the task takes tool output as its final output")
     callback_output: Optional[Any] = Field(default=None, description="stores task or agent callback outcome")
+    annotations: Optional[Dict[str, Any]] = Field(default=None)
     evaluation: Optional[InstanceOf[Evaluation]] = Field(default=None, description="stores overall evaluation of the task output. stored in ltm")
+    usage: Optional[UsageMetrics] = Field(default=None)
 
 
     def _fetch_value_of(self, key: str = None) -> Any:
@@ -236,7 +241,7 @@ class TaskOutput(BaseModel):
             description = EVALUATE.format(task_description=task.description, task_output=self.raw, eval_criteria=str(item))
             description = description + fsl_prompt if fsl_prompt else description
 
-            task_eval = Task(description=description, pydantic_output=EvaluationItem)
+            task_eval = Task(description=description, response_schema=EvaluationItem)
             res = task_eval.execute(agent=self.evaluation.eval_by)
 
             if res.pydantic:
@@ -308,14 +313,10 @@ class Task(BaseModel):
     id: UUID4 = Field(default_factory=uuid.uuid4, frozen=True, description="unique identifier for the object, not set by user")
     name: Optional[str] = Field(default=None)
     description: str = Field(description="Description of the actual task")
-
-    # response format
-    response_schema: Optional[Type[BaseModel] | List[ResponseField]] = Field(default=None)
-    pydantic_output: Optional[Type[BaseModel]] = Field(default=None, description="store Pydantic class as structured response format")
-    response_fields: Optional[List[ResponseField]] = Field(default_factory=list, description="store list of ResponseField as structured response format")
+    response_schema: Optional[Type[BaseModel] | List[ResponseField]] = Field(default=None, description="stores response format")
 
     # tool usage
-    tools: Optional[List[ToolSet | Tool | Any]] = Field(default_factory=list, description="tools that the agent can use aside from their tools")
+    tools: Optional[List[Any]] = Field(default_factory=list, description="tools that the agent can use aside from their tools")
     can_use_agent_tools: bool = Field(default=True, description="whether the agent can use their own tools when executing the task")
     tool_res_as_final: bool = Field(default=False, description="when set True, tools res will be stored in the `TaskOutput`")
 
@@ -340,7 +341,6 @@ class Task(BaseModel):
     fsls: Optional[list[str]] = Field(default=None, description="stores ideal/weak responses")
 
     # recording
-    _usage: UsageMetrics = PrivateAttr(default=None)
     _delegations: int = 0
     processed_agents: Set[str] = Field(default_factory=set, description="store keys of the agents that executed the task")
     output: Optional[TaskOutput] = Field(default=None, description="store the final TaskOutput object")
@@ -365,30 +365,36 @@ class Task(BaseModel):
         for field in required_fields:
             if getattr(self, field) is None:
                 raise ValueError( f"{field} must be provided either directly or through config")
-
-        self._usage = UsageMetrics(id=self.id)
         return self
 
 
     @model_validator(mode="after")
     def set_up_tools(self) -> Self:
-        if not self.tools:
-            pass
-        else:
+        if self.tools:
             tool_list = []
             for item in self.tools:
-                if isinstance(item, Tool) or isinstance(item, ToolSet):
-                    tool_list.append(item)
-                elif (isinstance(item, dict) and "function" not in item) or isinstance(item, str):
-                    pass
-                else:
-                    tool_list.append(item) # address custom tool
+                match item:
+                    case Tool() | ToolSet() | BaseTool() | RagTool() | GPTToolCUP() | GPTToolFileSearch() | GPTToolWebSearch():
+                        tool_list.append(item)
+                    case type(item, callable):
+                        tool_list.append(Tool(func=item))
+                    case dict():
+                        tool = None
+                        try:
+                            tool = Tool(**item)
+                        except:
+                            try:
+                                tool = RagTool(**item)
+                            except:
+                                pass
+                    case _:
+                        pass
             self.tools = tool_list
         return self
 
 
     def _structure_response_format(self, data_type: str = "object", model_provider: str = "gemini") -> Dict[str, Any] | None:
-        """Structures `response_fields` or `pydantic_output` to a LLM response format."""
+        """Structures `response_schema` into the LLM response format."""
 
         from versionhq.task.structured_response import StructuredOutput
 
@@ -398,28 +404,29 @@ class Task(BaseModel):
             return response_format
 
         else:
-            if self.response_fields:
-                properties, required_fields = {}, []
-                for i, item in enumerate(self.response_fields):
-                    if item:
-                        properties.update(item._format_props())
-                        required_fields.append(item.title)
-
-                response_schema = {
-                    "type": data_type,
-                    "properties": properties,
-                    "required": required_fields,
-                    "additionalProperties": False,
-                }
-                response_format = {
-                    "type": "json_schema",
-                    "json_schema": { "name": "outcome", "schema": response_schema }
-                }
+            if self.response_schema:
+                if isinstance(self.response_schema, list):
+                    properties, required_fields = {}, []
+                    for i, item in enumerate(self.response_schema):
+                        if isinstance(item, ResponseField):
+                            properties.update(item._format_props())
+                            required_fields.append(item.title)
+
+                    response_schema = {
+                        "type": data_type,
+                        "properties": properties,
+                        "required": required_fields,
+                        "additionalProperties": False,
+                    }
+                    response_format = {
+                        "type": "json_schema",
+                        "json_schema": { "name": "outcome", "schema": response_schema }
+                    }
 
-            elif self.pydantic_output:
-                response_format = StructuredOutput(response_format=self.pydantic_output, provider=model_provider)._format()
+                elif issubclass(self.response_schema, BaseModel):
+                    response_format = StructuredOutput(response_format=self.response_schema, provider=model_provider)._format()
 
-            return response_format
+        return response_format
 
 
     def _sanitize_raw_output(self, raw: str) -> Dict[str, str]:
@@ -435,6 +442,7 @@ class Task(BaseModel):
         r = re.sub(r"'\b", '"', r)
         r = r.strip()
         r = r.replace("  ", "")
+
         try:
             output = json.loads(r)
         except:
@@ -452,7 +460,6 @@ class Task(BaseModel):
             except:
                 output = ast.literal_eval(r)
 
-
             return output["json_schema"] if isinstance(output, dict) and "json_schema" in output else output if isinstance(output, dict) else { "output": str(r) }
 
 
@@ -471,30 +478,31 @@ class Task(BaseModel):
             if isinstance(output, dict):
                 return output["json_schema"] if "json_schema" in output else output
             else:
-               output = self._sanitize_raw_output(raw=raw)
-               return output
+                output = self._sanitize_raw_output(raw=raw)
+                return output
         except:
             output = self._sanitize_raw_output(raw=raw)
-            self._usage.record_errors(type=ErrorType.FORMAT)
             return output
 
 
-    def _create_pydantic_output(self, raw: str = None, json_dict: Dict[str, Any] = None) -> InstanceOf[BaseModel]:
+    def _create_pydantic_output(self, raw: str = None, json_dict: Dict[str, Any] = None) -> InstanceOf[BaseModel] | None:
         """
         Create pydantic output from raw or json_dict output.
         """
 
-        output_pydantic = self.pydantic_output
-
-        try:
-            json_dict = json_dict if json_dict else self._create_json_output(raw=raw)
+        if self.response_schema and not isinstance(self.response_schema, list):
+            output_pydantic = self.response_schema
+            try:
+                json_dict = json_dict if json_dict else self._create_json_output(raw=raw)
 
-            for k, v in json_dict.items():
-                setattr(output_pydantic, k, v)
-        except:
-            pass
+                for k, v in json_dict.items():
+                    setattr(output_pydantic, k, v)
+            except:
+                pass
+            return output_pydantic
 
-        return output_pydantic
+        else:
+            return None
 
 
     def interpolate_inputs(self, inputs: Dict[str, Any]) -> None:
@@ -638,44 +646,47 @@ class Task(BaseModel):
 
         start_dt = datetime.datetime.now()
         task_output: InstanceOf[TaskOutput] = None
-        raw_output: str = None
-        tool_output: str | list = None
-        task_tools: List[List[InstanceOf[Tool]| InstanceOf[ToolSet] | Type[Tool]]] = []
         user_prompt, dev_prompt = None, None
 
-        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 == True:
             agent_to_delegate = self._select_agent_to_delegate(agent=agent)
             agent = agent_to_delegate
             self._delegations += 1
 
-        if self.tool_res_as_final == True:
-            user_prompt, dev_prompt, tool_output = agent.execute_task(task=self, context=context, task_tools=task_tools)
-            raw_output = str(tool_output) if tool_output else ""
-            if not raw_output:
-                self._usage.record_errors(type=ErrorType.TOOL)
-            task_output = TaskOutput(task_id=self.id, tool_output=tool_output, raw=raw_output)
+        user_prompt, dev_prompt, raw_output, usage = agent.execute_task(task=self, context=context)
+        match raw_output:
+            case TaskOutput():
+                raw_output.task_id = self.id
+                raw_output.usage = usage
+                task_output = raw_output
+
+            case str():
+                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"]))
+
+                pydantic_output = self._create_pydantic_output(raw=raw_output, json_dict=json_dict_output)
+                task_output = TaskOutput(
+                    task_id=self.id,
+                    raw=raw_output if raw_output is not None else "",
+                    pydantic=pydantic_output,
+                    json_dict=json_dict_output,
+                    tool_output=raw_output if self.tool_res_as_final else None,
+                    usage=usage
+                )
+
+            case None | "":
+                task_output = TaskOutput(task_id=self.id, raw="", usage=usage)
+                task_output.usage.record_errors(type=ErrorType.FORMAT)
 
-        else:
-            user_prompt, dev_prompt, 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"]))
-
-            pydantic_output = self._create_pydantic_output(raw=raw_output, json_dict=json_dict_output) if self.pydantic_output else None
-
-            task_output = TaskOutput(
-                task_id=self.id,
-                raw=raw_output if raw_output is not None else "",
-                pydantic=pydantic_output,
-                json_dict=json_dict_output,
-            )
+            case _:
+                task_output = TaskOutput(
+                    task_id=self.id,
+                    raw=raw_output,
+                    tool_output=raw_output if self.tool_res_as_final else None,
+                    usage=usage
+                )
 
-        self.output = task_output
         self.processed_agents.add(agent.key)
 
         # if self.output_file: ## disabled for now
@@ -691,10 +702,9 @@ class Task(BaseModel):
             self._pfg.user_prompts.update({ index: user_prompt })
             self._pfg.dev_prompts.update({ index: dev_prompt })
 
-        if raw_output:
+        if task_output.raw:
             if self.should_evaluate:
                 task_output.evaluate(task=self)
-                self.output = task_output
 
             self._create_short_and_long_term_memories(agent=agent, task_output=task_output)
 
@@ -705,11 +715,14 @@ class Task(BaseModel):
                 valid_kwargs = { k: kwargs[k] if  k in kwargs else None for k in valid_keys }
                 callback_res = self.callback(**valid_kwargs)
                 task_output.callback_output = callback_res
-                self.output = task_output
-            self._store_logs()
 
         end_dt = datetime.datetime.now()
-        self._usage.record_latency(start_dt=start_dt, end_dt=end_dt)
+        task_output.usage.record_latency(start_dt=start_dt, end_dt=end_dt)
+        if task_output.json_dict and "output" in task_output.json_dict:
+            task_output.usage.record_errors(type=ErrorType.FORMAT)
+
+        self.output = task_output
+        self._store_logs()
         return task_output
 
 
@@ -720,7 +733,6 @@ class Task(BaseModel):
         from versionhq._prompt.model import Prompt
         from versionhq._prompt.auto_feedback import PromptFeedbackGraph
 
-        # self._usage = None
         prompt = Prompt(task=self, agent=agent, context=context)
         pfg = PromptFeedbackGraph(prompt=prompt, should_reform=self.human, reform_trigger_event=ReformTriggerEvent.USER_INPUT if self.human else None)
         pfg = pfg.set_up_graph()
@@ -729,18 +741,18 @@ class Task(BaseModel):
         try:
             if self._pfg and self.output is None:
                 res, all_outputs = self._pfg.activate()
-                if all_outputs: self._usage = self._pfg._usage
+                if all_outputs:
+                    res.usage = self._pfg.usage
                 return res
-
         except:
-            self._usage.record_errors(type=ErrorType.API)
+            self._pfg.usage.record_errors(type=ErrorType.API)
             Logger().log(level="error", message="Failed to execute the task.", color="red")
             return None
 
 
     @property
     def key(self) -> str:
-        output_format = "json" if self.response_fields else "pydantic" if self.pydantic_output is not None else "raw"
+        output_format = "json" if self.response_schema else "raw"
         source = [self.description, output_format]
         return md5("|".join(source).encode(), usedforsecurity=False).hexdigest()
 

versionhq/task_graph/draft.py

@@ -41,7 +41,7 @@ def workflow(final_output: Type[BaseModel], context: Any = None, human: bool = F
             "https://www.geeksforgeeks.org/graph-and-its-representations/",
             ", ".join([k for k in DependencyType._member_map_.keys()]),
         ],
-        llm="gemini-2.0",
+        llm="gemini/gemini-2.0-flash",
         with_memory=with_memory,
         maxit=1,
         max_retry_limit=1,
@@ -49,7 +49,7 @@ def workflow(final_output: Type[BaseModel], context: Any = None, human: bool = F
 
     task = Task(
         description=dedent(f"Design a resource-efficient workflow to achieve the following goal: {final_output_prompt}. The workflow should consist of a list of detailed tasks that represent decision making points, each with the following information:\nname: A concise name of the task\ndescription: A concise description of the task.\nconnections: A list of target tasks that this task connects to.\ndependency_types: The type of dependency between this task and each of its connected task. \noutput: key output from the task in a word.\n\nUse the following dependency types: {dep_type_prompt}.\n\nPrioritize minimizing resource consumption (computation, memory, and data transfer) when defining tasks, connections, and dependencies.  Consider how data is passed between tasks and aim to reduce unnecessary data duplication or transfer. Explain any design choices made to optimize resource usage."),
-        response_fields=[
+        response_schema=[
             ResponseField(title="tasks", data_type=list, items=dict, properties=[
                 ResponseField(title="name", data_type=str),
                 ResponseField(title="description", data_type=str),
@@ -73,13 +73,13 @@ def workflow(final_output: Type[BaseModel], context: Any = None, human: bool = F
 
     for item in task_items:
         key = item["output"].lower().replace(" ", "_") if item["output"] else "output"
-        task = Task(name=item["name"], description=item["description"], response_fields=[ResponseField(title=key, data_type=str)])
+        task = Task(name=item["name"], description=item["description"], response_schema=[ResponseField(title=key, data_type=str)])
         tasks.append(task)
         nodes.append(Node(task=task))
 
     task_graph = TaskGraph(
         nodes={node.identifier: node for node in nodes},
-        concl_format=final_output,
+        concl_response_schema=final_output,
         concl=None,
         should_reform=human,
         reform_trigger_event=ReformTriggerEvent.USER_INPUT if human else None,