versionhq 1.2.3.8__py3-none-any.whl → 1.2.4.2__py3-none-any.whl

versionhq/agent_network/formation.py

@@ -47,9 +47,12 @@ def form_agent_network(
                         Logger(verbose=True).log(level="warning", message=f"The formation {formation} is invalid. We'll recreate a valid formation.", color="yellow")
                         formation = None
 
-                case int() | float():
+                case int():
                     formation = Formation(int(formation))
 
+                case float():
+                    formation = Formation(round(formation))
+
                 case _:
                     Logger(verbose=True).log(level="warning", message=f"The formation {formation} is invalid. We'll recreate a valid formation.", color="yellow")
                     formation = None
@@ -64,14 +67,14 @@ def form_agent_network(
     prompt_expected_outcome = expected_outcome if isinstance(expected_outcome, str) else str(expected_outcome.model_dump()) if type(expected_outcome) == BaseModel else ""
 
     class Outcome(BaseModel):
-        formation: Enum
+        formation: str
         agent_roles: list[str]
         task_descriptions: list[str]
         task_outcomes: list[list[str]]
         leader_agent: str
 
     vhq_task = Task(
-        description=f"Design a team of specialized agents to fully automate the following task and achieve the expected outcome. For each agent, define its role, task description, and expected outputs via the task with items in a list. Then specify the team formation if the formation is not given. If you think SUPERVISING or HYBRID is the best formation, include a leader_agent role, else leave the leader_agent role blank.\nTask: {str(task)}\nExpected outcome: {prompt_expected_outcome}\nFormation: {prompt_formation}",
+        description=f"Design a team of specialized agents to fully automate the following task and achieve the expected outcome. For each agent, define its role, task description, and expected outputs via the task with items in a list. Then specify the formation if the formation is not given. If you think SUPERVISING or HYBRID is the best formation, include a leader_agent role, else leave the leader_agent role blank.\nTask: {str(task)}\nExpected outcome: {prompt_expected_outcome}\nFormation: {prompt_formation}",
         pydantic_output=Outcome
     )
 
@@ -80,13 +83,17 @@ def form_agent_network(
 
     res = vhq_task.execute(agent=vhq_formation_planner, context=context)
 
-    formation_keys = ([k for k in Formation._member_map_.keys() if k == res.pydantic.formation.upper()]
-                      if res.pydantic else [k for k in Formation._member_map_.keys() if k == res.json_dict["formation"].upper()])
+    formation_keys = []
+    if hasattr(res.pydantic, "formation"):
+        formation_keys = [k for k in Formation._member_map_.keys() if k == res.pydantic.formation.upper()]
+    elif "formation" in res.json_dict:
+        formation_keys = [k for k in Formation._member_map_.keys() if k == res.json_dict["formation"].upper()]
+
     _formation = Formation[formation_keys[0]] if formation_keys else Formation.SUPERVISING
 
     network_tasks = []
     members = []
-    leader = str(res.pydantic.leader_agent) if res.pydantic else str(res.json_dict["leader_agent"])
+    leader = str(res.pydantic.leader_agent) if res.pydantic and hasattr(res.pydantic, "leader_agent") else str(res.json_dict["leader_agent"]) if "leader_agent" in res.json_dict else None
 
     agent_roles = res.pydantic.agent_roles if res.pydantic else res.json_dict["agent_roles"]
     created_agents = [Agent(role=str(item), goal=str(item)) for item in agent_roles]
@@ -132,13 +139,13 @@ def form_agent_network(
 
     if len(created_tasks) <= len(created_agents):
         for i in range(len(created_tasks)):
-            is_manager = bool(created_agents[i].role.lower() == leader.lower())
+            is_manager = False if not leader else bool(created_agents[i].role.lower() == leader.lower())
             member = Member(agent=created_agents[i], is_manager=is_manager, tasks=[created_tasks[i]])
             members.append(member)
 
         for i in range(len(created_tasks), len(created_agents)):
             try:
-                is_manager = bool(created_agents[i].role.lower() == leader.lower())
+                is_manager = False if not leader else bool(created_agents[i].role.lower() == leader.lower())
                 member_w_o_task = Member(agent=created_agents[i], is_manager=is_manager)
                 members.append(member_w_o_task)
             except:
@@ -146,7 +153,7 @@ def form_agent_network(
 
     elif len(created_tasks) > len(created_agents):
         for i in range(len(created_agents)):
-            is_manager = bool(created_agents[i].role.lower() == leader.lower())
+            is_manager = False if not leader else bool(created_agents[i].role.lower() == leader.lower())
             member = Member(agent=created_agents[i], is_manager=is_manager, tasks=[created_tasks[i]])
             members.append(member)
 
@@ -154,7 +161,8 @@ def form_agent_network(
 
 
     if _formation == Formation.SUPERVISING and not [member for member in members if member.is_manager]:
-        manager = Member(agent=Agent(role=leader, goal=leader), is_manager=True)
+        role = leader if leader else "Leader"
+        manager = Member(agent=Agent(role=role), is_manager=True)
         members.append(manager)
 
     members.sort(key=lambda x: x.is_manager == False)

versionhq/agent_network/model.py

@@ -306,12 +306,8 @@ class AgentNetwork(BaseModel):
         return res, task_graph
 
 
-    def launch(
-            self, kwargs_pre: Optional[Dict[str, str]] = None, kwargs_post: Optional[Dict[str, Any]] = None, start_index: int = None
-        ) -> Tuple[TaskOutput, TaskGraph]:
-        """
-        Launch the agent network - executing tasks and recording their outputs.
-        """
+    def launch(self, kwargs_pre: Optional[Dict[str, str]] = None, kwargs_post: Optional[Dict[str, Any]] = None, start_index: int = None) -> Tuple[TaskOutput, TaskGraph]:
+        """Launches agent network by executing tasks in the network and recording the outputs"""
 
         self._assign_tasks()
 

versionhq/knowledge/source_docling.py

@@ -1,6 +1,5 @@
 from pathlib import Path
-from typing import Iterator, List, Optional
-from urllib.parse import urlparse
+from typing import Iterator, List
 
 try:
     from docling.datamodel.base_models import InputFormat
@@ -20,7 +19,7 @@ from pydantic import Field
 
 from versionhq.knowledge.source import BaseKnowledgeSource
 from versionhq.storage.utils import fetch_db_storage_path
-from versionhq._utils.vars import KNOWLEDGE_DIRECTORY
+from versionhq._utils import KNOWLEDGE_DIRECTORY, is_valid_url
 
 
 class DoclingSource(BaseKnowledgeSource):
@@ -83,21 +82,6 @@ class DoclingSource(BaseKnowledgeSource):
             yield chunk.text
 
 
-    @staticmethod
-    def _validate_url(url: str) -> bool:
-        try:
-            result = urlparse(url)
-            return all(
-                [
-                    result.scheme in ("http", "https"),
-                    result.netloc,
-                    len(result.netloc.split(".")) >= 2,  # Ensure domain has TLD
-                ]
-            )
-        except Exception:
-            return False
-
-
     def model_post_init(self, _) -> None:
         self.valid_file_paths = self.validate_content()
         self.content.extend(self._load_content())
@@ -109,7 +93,7 @@ class DoclingSource(BaseKnowledgeSource):
             if isinstance(path, str):
                 if path.startswith(("http://", "https://")):
                     try:
-                        if self._validate_url(path):
+                        if is_valid_url(path):
                             processed_paths.append(path)
                         else:
                             raise ValueError(f"Invalid URL format: {path}")

versionhq/llm/model.py

@@ -8,12 +8,13 @@ from contextlib import contextmanager
 from typing import Any, Dict, List, Optional
 from typing_extensions import Self
 
+import litellm
 from litellm import JSONSchemaValidationError, get_supported_openai_params
 from pydantic import BaseModel, Field, PrivateAttr, model_validator, ConfigDict
 
 from versionhq.llm.llm_vars import LLM_CONTEXT_WINDOW_SIZES, MODELS, PARAMS, PROVIDERS, ENDPOINT_PROVIDERS, ENV_VARS
 from versionhq.tool.model import Tool, ToolSet
-from versionhq._utils.logger import Logger
+from versionhq._utils import Logger
 
 
 load_dotenv(override=True)
@@ -47,7 +48,6 @@ class FilteredStream:
 @contextmanager
 def suppress_warnings():
     with warnings.catch_warnings():
-        import litellm
         litellm.set_verbose = False
         warnings.filterwarnings(action="ignore")
         old_stdout = sys.stdout
@@ -102,7 +102,6 @@ class LLM(BaseModel):
         """
         Validate the given model, provider, interface provider.
         """
-        import litellm
         litellm.drop_params = True
 
         self._init_model_name = self.model
@@ -180,7 +179,6 @@ class LLM(BaseModel):
         """
         Set up valid config params after setting up a valid model, provider, interface provider names.
         """
-        import litellm
         litellm.drop_params = True
 
         self._tokens = 0
@@ -208,9 +206,7 @@ class LLM(BaseModel):
         valid_config, valid_keys = dict(), list()
 
         if self.model:
-            valid_keys = get_supported_openai_params(
-                model=self.model, custom_llm_provider=self.endpoint_provider, request_type="chat_completion"
-            )
+            valid_keys = get_supported_openai_params(model=self.model, custom_llm_provider=self.endpoint_provider, request_type="chat_completion")
 
         if not valid_keys:
             valid_keys = PARAMS.get("common")
@@ -269,7 +265,6 @@ class LLM(BaseModel):
 
 
     def _supports_stop_words(self) -> bool:
-        import litellm
         supported_params = get_supported_openai_params(model=self.model, custom_llm_provider=self.endpoint_provider)
         return "stop" in supported_params if supported_params else False
 
@@ -282,7 +277,6 @@ class LLM(BaseModel):
 
 
     def _set_callbacks(self, callbacks: List[Any]):
-        import litellm
         callback_types = [type(callback) for callback in callbacks]
         for callback in litellm.success_callback[:]:
             if type(callback) in callback_types:
@@ -306,7 +300,6 @@ class LLM(BaseModel):
         """
         Execute LLM based on the agent's params and model params.
         """
-        import litellm
         litellm.drop_params = True
         litellm.set_verbose = True
 
@@ -318,8 +311,12 @@ class LLM(BaseModel):
                 res, tool_res = None, ""
                 cred = self._set_env_vars()
 
-                if not tools:
+                if self.provider == "gemini":
+                    self.response_format = { "type": "json_object" } if not tools else None
+                else:
                     self.response_format = response_format
+
+                if not tools:
                     params = self._create_valid_params(config=config)
                     res = litellm.completion(model=self.model, messages=messages, stream=False, **params, **cred)
                     self._tokens += int(res["usage"]["total_tokens"])
@@ -327,7 +324,6 @@ class LLM(BaseModel):
 
                 else:
                     try:
-                        self.response_format = { "type": "json_object" }  if tool_res_as_final and self.provider != "gemini" else response_format
                         self.tools = [item.tool.properties if isinstance(item, ToolSet) else item.properties for item in tools]
                         params = self._create_valid_params(config=config)
                         res = litellm.completion(model=self.model, messages=messages, **params, **cred)

versionhq/task/evaluation.py

@@ -16,7 +16,7 @@ class EvaluationItem(BaseModel):
     criteria: str
     suggestion: str
     score: float
-    weight: int = 1
+    weight: int
 
 
 class Evaluation(BaseModel):

versionhq/task/model.py

@@ -4,10 +4,9 @@ import datetime
 import uuid
 import inspect
 import enum
-from textwrap import dedent
 from concurrent.futures import Future
 from hashlib import md5
-from typing import Any, Dict, List, Set, Optional, Callable, Type
+from typing import Any, Dict, List, Set, Optional, Callable, Type, Tuple
 from typing_extensions import Annotated, Self
 
 from pydantic import UUID4, BaseModel, Field, PrivateAttr, field_validator, model_validator, InstanceOf, field_validator
@@ -175,6 +174,7 @@ class TaskOutput(BaseModel):
     """
     A class to store the final output of the given task in raw (string), json_dict, and pydantic class formats.
     """
+
     _tokens: int = PrivateAttr(default=0)
 
     task_id: UUID4 = Field(default_factory=uuid.uuid4, frozen=True, description="store Task ID")
@@ -187,19 +187,19 @@ class TaskOutput(BaseModel):
     evaluation: Optional[InstanceOf[Evaluation]] = Field(default=None, description="stores overall evaluation of the task output. stored in ltm")
 
 
-    def to_context_prompt(self) -> str:
-        """
-        Returns response in string as a prompt context.
-        """
+    def _to_context_prompt(self) -> str:
+        """Formats prompt context in text formats from the final response."""
+
         context = ""
-        try:
-            context = json.dumps(self.json_dict)
-        except:
-            try:
-                if self.pydantic:
-                    context = self.pydantic.model_dump()
-            except:
-                context = self.raw
+        match self.final:
+            case dict() | self.pydantic:
+                try:
+                    context = json.dumps(self.final)
+                except:
+                    context = str(self.final)
+            case _:
+                context = str(self.final)
+
         return context
 
 
@@ -252,6 +252,24 @@ class TaskOutput(BaseModel):
         return self.evaluation
 
 
+    @property
+    def final(self) -> Any:
+        """Returns final output from the task."""
+
+        output = None
+
+        if self.callback_output:
+            output = self.callback_output
+
+        elif self.tool_output and str(self.tool_output) == self.raw: # tool_output_as_final
+            output = self.tool_output
+
+        else:
+            output = self.pydantic if self.pydantic else self.json_dict if self.json_dict else self.raw
+
+        return output
+
+
     @property
     def aggregate_score(self) -> float | int:
         return self.evaluation.aggregate_score if self.evaluation is not None else 0
@@ -280,6 +298,7 @@ class Task(BaseModel):
     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")
 
@@ -288,6 +307,15 @@ class Task(BaseModel):
     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`")
 
+    image: Optional[str] = Field(default=None, description="absolute file path or url in string")
+    file: Optional[str] = Field(default=None, description="absolute file path or url in string")
+    audio: Optional[str] = Field(default=None,  description="absolute file path or url in string")
+
+    # test run
+    should_test_run: bool = Field(default=False)
+    human: bool = Field(default=False)
+    _pfg: Any = None
+
     # executing
     execution_type: TaskExecutionType = Field(default=TaskExecutionType.SYNC)
     allow_delegation: bool = Field(default=False, description="whether to delegate the task to another agent")
@@ -301,9 +329,10 @@ class Task(BaseModel):
 
     # recording
     _tokens: int = 0
+    _tool_errors: int = 0
+    _format_errors: int = 0
+    _delegations: int = 0
     processed_agents: Set[str] = Field(default_factory=set, description="store keys of the agents that executed the task")
-    tool_errors: int = 0
-    delegations: int = 0
     output: Optional[TaskOutput] = Field(default=None, description="store the final TaskOutput object")
 
 
@@ -346,95 +375,8 @@ class Task(BaseModel):
         return self
 
 
-    def _draft_output_prompt(self, model_provider: str) -> str:
-        """
-        Draft prompts on the output format by converting `
-        """
-
-        output_prompt = ""
-
-        if self.pydantic_output:
-            output_prompt = f"""
-Your response MUST STRICTLY follow the given repsonse format:
-JSON schema: {str(self.pydantic_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 = "You MUST Return your response as a valid JSON serializable string, enclosed in double quotes. Do not use single quotes, trailing commas, or other non-standard JSON syntax."
-
-        return dedent(output_prompt)
-
-
-    def _draft_context_prompt(self, context: Any) -> str:
-        """
-        Create a context prompt from the given context in any format: a task object, task output object, list, dict.
-        """
-
-        context_to_add = None
-        if not context:
-            # Logger().log(level="error", color="red", message="Missing a context to add to the prompt. We'll return ''.")
-            return context_to_add
-
-        match context:
-            case str():
-                context_to_add = context
-
-            case Task():
-                if not context.output:
-                    res = context.execute()
-                    context_to_add = res.to_context_prompt()
-
-                else:
-                    context_to_add = context.output.raw
-
-            case TaskOutput():
-                context_to_add = context.to_context_prompt()
-
-
-            case dict():
-                context_to_add = str(context)
-
-            case list():
-                res = ", ".join([self._draft_context_prompt(context=item) for item in context])
-                context_to_add = res
-
-            case _:
-                pass
-
-        return dedent(context_to_add)
-
-
-    def _prompt(self, model_provider: str = None, context: Optional[Any] = None) -> str:
-        """
-        Format the task prompt and cascade it to the agent.
-        """
-        output_prompt = self._draft_output_prompt(model_provider=model_provider)
-        task_slices = [self.description, output_prompt, ]
-
-        if context:
-            context_prompt = self._draft_context_prompt(context=context)
-            task_slices.insert(len(task_slices), f"Consider the following context when responding: {context_prompt}")
-
-        return "\n".join(task_slices)
-
-
     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_output`.
-        1 nested item is accepted.
-        """
+        """Structures `response_fields` or `pydantic_output` to a LLM response format."""
 
         from versionhq.task.structured_response import StructuredOutput
 
@@ -452,7 +394,7 @@ Ref. Output image: {output_formats_to_follow}
                         required_fields.append(item.title)
 
                 response_schema = {
-                    "type": "object",
+                    "type": data_type,
                     "properties": properties,
                     "required": required_fields,
                     "additionalProperties": False,
@@ -462,17 +404,49 @@ Ref. Output image: {output_formats_to_follow}
                     "json_schema": { "name": "outcome", "schema": response_schema }
                 }
 
-
             elif self.pydantic_output:
                 response_format = StructuredOutput(response_format=self.pydantic_output, provider=model_provider)._format()
 
             return response_format
 
 
+    def _sanitize_raw_output(self, raw: str) -> Dict[str, str]:
+        """Sanitizes raw output and prepare for json.loads"""
+
+        import re
+        import ast
+
+        output, j = None, None
+        r = str(raw).strip()
+        r = r.replace("true", "True").replace("false", "False").replace("```json", '"').replace("```", '"').replace('\n', '').replace('\\', '')
+        r = re.sub("^'", '"', r)
+        r = re.sub(r"'\b", '"', r)
+        r = r.strip()
+        r = r.replace("  ", "")
+        try:
+            output = json.loads(r)
+        except:
+            try: j = json.dumps(eval(r))
+            except:
+                try: j = json.dumps(str(r))
+                except: j = r
+            output = json.loads(j)
+
+        if isinstance(output, dict):
+            return output
+        else:
+            try:
+                output = ast.literal_eval(j)
+            except:
+                output = ast.literal_eval(r)
+
+            return output if isinstance(output, dict) else { "output": str(r) }
+
+
     def _create_json_output(self, raw: str) -> Dict[str, Any]:
-        """
-        Create json (dict) output from the raw output and `response_fields` information.
-        """
+        """Creates JSON output from the raw output."""
+
+        output = None
 
         if raw is None or raw == "":
             Logger().log(level="warning", message="The model returned an empty response. Returning an empty dict.", color="yellow")
@@ -480,27 +454,14 @@ Ref. Output image: {output_formats_to_follow}
             return output
 
         try:
-            r = str(raw).replace("true", "True").replace("false", "False")
-            j = json.dumps(eval(r))
-            output = json.loads(j)
+            output = json.loads(raw)
             if isinstance(output, dict):
                 return output
-
             else:
-                r = str(raw).strip().replace("{'", '{"').replace("{ '", '{"').replace("': '", '": "').replace("'}", '"}').replace("' }", '"}').replace("', '", '", "').replace("['", '["').replace("[ '", '[ "').replace("']", '"]').replace("' ]", '" ]').replace("{\n'", '{"').replace("{\'", '{"').replace("true", "True").replace("false", "False").replace('\"', "'")
-                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) }
-
+               output = self._sanitize_raw_output(raw=raw)
+               return output
         except:
-            output = { "output": str(raw) }
+            output = self._sanitize_raw_output(raw=raw)
             return output
 
 
@@ -620,23 +581,32 @@ Ref. Output image: {output_formats_to_follow}
 
 
     # task execution
-    def execute(
-            self, type: TaskExecutionType = None, agent: Optional["vhq.Agent"] = None, context: Optional[Any] = None
-        ) -> TaskOutput | Future[TaskOutput]:
-        """
-        A main method to handle task execution. Build an agent when the agent is not given.
-        """
+    def execute(self, type: TaskExecutionType = None, agent: "vhq.Agent" = None, context: Any = None) -> TaskOutput | Future[TaskOutput]:
+        """A main method to handle task execution."""
+
         type = type if type else  self.execution_type if self.execution_type else TaskExecutionType.SYNC
+        agent = agent if agent else self._build_agent_from_task(task_description=self.description)
+        res = None
 
-        if not agent:
-            agent = self._build_agent_from_task(task_description=self.description)
+        if (self.should_test_run or agent.self_learn) and not self._pfg:
+            res = self._test_time_computation(agent=agent, context=context)
+            return res
+
+        # if self._pfg:
+        #     res, all_outputs = self.pfg.activate()
+        #     tokens, latency = self.pfg.usage
+        #     self._tokens = tokens
+        #     res.latency = latency
+        #     return res
 
         match type:
             case TaskExecutionType.SYNC:
-                return self._execute_sync(agent=agent, context=context)
+                res = self._execute_sync(agent=agent, context=context)
 
             case TaskExecutionType.ASYNC:
-                return self._execute_async(agent=agent, context=context)
+                res = self._execute_async(agent=agent, context=context)
+
+        return res
 
 
     def _execute_sync(self, agent, context: Optional[Any] = None) -> TaskOutput:
@@ -657,14 +627,14 @@ Ref. Output image: {output_formats_to_follow}
 
 
     def _execute_core(self, agent, context: Optional[Any]) -> TaskOutput:
-        """
-        A core method to execute a task.
-        """
+        """A core method to execute a single task."""
+
         task_output: InstanceOf[TaskOutput] = None
         raw_output: str = None
         tool_output: str | list = None
         task_tools: List[List[InstanceOf[Tool]| InstanceOf[ToolSet] | Type[Tool]]] = []
         started_at, ended_at = datetime.datetime.now(), datetime.datetime.now()
+        user_prompt, dev_prompt = None, None
 
         if self.tools:
             for item in self.tools:
@@ -674,18 +644,18 @@ Ref. Output image: {output_formats_to_follow}
         if self.allow_delegation == True:
             agent_to_delegate = self._select_agent_to_delegate(agent=agent)
             agent = agent_to_delegate
-            self.delegations += 1
+            self._delegations += 1
 
         if self.tool_res_as_final == True:
             started_at = datetime.datetime.now()
-            tool_output = agent.execute_task(task=self, context=context, task_tools=task_tools)
+            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 ""
             ended_at = datetime.datetime.now()
             task_output = TaskOutput(task_id=self.id, tool_output=tool_output, raw=raw_output)
 
         else:
             started_at = datetime.datetime.now()
-            raw_output = agent.execute_task(task=self, context=context, task_tools=task_tools)
+            user_prompt, dev_prompt, raw_output = agent.execute_task(task=self, context=context, task_tools=task_tools)
             ended_at = datetime.datetime.now()
 
             json_dict_output = self._create_json_output(raw=raw_output)
@@ -714,6 +684,11 @@ Ref. Output image: {output_formats_to_follow}
         #     )
         #     self._save_file(content)
 
+        if self._pfg:
+            index = self._pfg.index
+            self._pfg.user_prompts.update({ index: user_prompt })
+            self._pfg.dev_prompts.update({ index: dev_prompt })
+
         if raw_output:
             if self.should_evaluate:
                 task_output.evaluate(task=self)
@@ -734,6 +709,31 @@ Ref. Output image: {output_formats_to_follow}
         return task_output
 
 
+    def _test_time_computation(self, agent, context: Optional[Any]) -> TaskOutput | None:
+        """Handles test-time computation."""
+
+        from versionhq.task_graph.model import ReformTriggerEvent
+        from versionhq._prompt.model import Prompt
+        from versionhq._prompt.auto_feedback import PromptFeedbackGraph
+
+        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()
+        self._pfg = pfg
+
+        # try:
+        if self._pfg and self.output is None:
+            res, _ = self._pfg.activate()
+            tokens, latency = self._pfg.usage
+            self._tokens = tokens
+            res.latency = latency
+            return res
+
+        # except:
+        #     Logger().log(level="error", message="Failed to execute the task.", color="red")
+        #     return None, None
+
+
     @property
     def key(self) -> str:
         output_format = "json" if self.response_fields else "pydantic" if self.pydantic_output is not None else "raw"