feat: Add price tracking and summary functionality for agent tasks

src/upsonic/client/level_two/agent.py

@@ -11,7 +11,7 @@
 
 from ..tasks.tasks import Task
 
-from ..printing import agent_end, agent_total_cost, agent_retry
+from ..printing import agent_end, agent_total_cost, agent_retry, print_price_id_summary
 
 
 
@@ -101,12 +101,12 @@ def agent_(
                     the_result = self.send_agent_request(agent_configuration, each, llm_model)
                     the_result["time"] = time.time() - start_time
                     results.append(the_result)
-                    agent_end(the_result["result"], the_result["llm_model"], the_result["response_format"], start_time, time.time(), the_result["usage"], the_result["tool_count"], the_result["context_count"], self.debug)
+                    agent_end(the_result["result"], the_result["llm_model"], the_result["response_format"], start_time, time.time(), the_result["usage"], the_result["tool_count"], the_result["context_count"], self.debug, each.price_id)
             else:
                 the_result = self.send_agent_request(agent_configuration, task, llm_model)
                 the_result["time"] = time.time() - start_time
                 results.append(the_result)
-                agent_end(the_result["result"], the_result["llm_model"], the_result["response_format"], start_time, time.time(), the_result["usage"], the_result["tool_count"], the_result["context_count"], self.debug)
+                agent_end(the_result["result"], the_result["llm_model"], the_result["response_format"], start_time, time.time(), the_result["usage"], the_result["tool_count"], the_result["context_count"], self.debug, task.price_id)
         except Exception as e:
 
             try:
@@ -229,7 +229,7 @@ def send_agent_request(
 
 
 
-    def create_characterization(self, agent_configuration: AgentConfiguration, llm_model: str = None):
+    def create_characterization(self, agent_configuration: AgentConfiguration, llm_model: str = None, price_id: str = None):
         tools = [Search]
 
         search_result = None
@@ -241,7 +241,7 @@ def create_characterization(self, agent_configuration: AgentConfiguration, llm_m
 
         # Handle website search if URL is provided
         if agent_configuration.company_url:
-            search_task = Task(description=f"Make a search for {agent_configuration.company_url}", tools=tools, response_format=SearchResult)
+            search_task = Task(description=f"Make a search for {agent_configuration.company_url}", tools=tools, response_format=SearchResult, price_id_=price_id, not_main_task=True)
             self.call(search_task, llm_model=llm_model)
             search_result = search_task.response
 
@@ -251,7 +251,9 @@ def create_characterization(self, agent_configuration: AgentConfiguration, llm_m
             company_objective_task = Task(description=f"Generate the company objective for {agent_configuration.company_objective}", 
                                         tools=tools, 
                                         response_format=CompanyObjective,
-                                        context=context)
+                                        context=context,
+                                        price_id_=price_id,
+                                        not_main_task=True)
             self.call(company_objective_task, llm_model=llm_model)
             company_objective_result = company_objective_task.response
 
@@ -267,7 +269,9 @@ def create_characterization(self, agent_configuration: AgentConfiguration, llm_m
             human_objective_task = Task(description=f"Generate the human objective for {agent_configuration.job_title}", 
                                       tools=tools, 
                                       response_format=HumanObjective,
-                                      context=context)
+                                      context=context,
+                                      price_id_=price_id,
+                                      not_main_task=True)
             self.call(human_objective_task, llm_model=llm_model)
             human_objective_result = human_objective_task.response
 
@@ -307,10 +311,10 @@ def agent(self, agent_configuration: AgentConfiguration, task: Task,  llm_model:
         if agent_configuration.caching:
             the_characterization = get_from_cache_with_expiry(the_characterization_cache_key)
             if the_characterization is None:
-                the_characterization = self.create_characterization(agent_configuration, llm_model)
+                the_characterization = self.create_characterization(agent_configuration, llm_model, task.price_id)
                 save_to_cache_with_expiry(the_characterization, the_characterization_cache_key, agent_configuration.cache_expiry)
         else:
-            the_characterization = self.create_characterization(agent_configuration, llm_model)
+            the_characterization = self.create_characterization(agent_configuration, llm_model, task.price_id)
 
 
 
@@ -433,6 +437,9 @@ def agent(self, agent_configuration: AgentConfiguration, task: Task,  llm_model:
 
         agent_total_cost(total_input_tokens, total_output_tokens, total_time, the_llm_model)
 
+        if not original_task.not_main_task:
+            print_price_id_summary(original_task.price_id, original_task)
+
         return original_task.response
 
 
@@ -483,14 +490,16 @@ def multiple(self, task: Task, llm_model: str = None):
         mode_selector = Task(
             description=mode_selection_prompt,
             response_format=AgentMode,
-            context=[task]
+            context=[task],
+            price_id_=task.price_id,
+            not_main_task=True
         )
 
         self.call(mode_selector, llm_model)
 
         # If level_no_step is selected, return just the end task
         if mode_selector.response.selected_mode == "level_no_step":
-            return [Task(description=task.description, response_format=task.response_format, tools=task.tools)]
+            return [Task(description=task.description, response_format=task.response_format, tools=task.tools, price_id_=task.price_id, not_main_task=True)]
 
         # Generate a list of sub tasks
         prompt = f"""
@@ -528,20 +537,20 @@ def multiple(self, task: Task, llm_model: str = None):
         sub_tasker_context = [task, task.response_format]
         if task.context:
             sub_tasker_context = task.context
-        sub_tasker = Task(description=prompt, response_format=SubTaskList, context=sub_tasker_context, tools=task.tools)
+        sub_tasker = Task(description=prompt, response_format=SubTaskList, context=sub_tasker_context, tools=task.tools, price_id_=task.price_id, not_main_task=True)
 
         self.call(sub_tasker, llm_model)
 
         sub_tasks = []
 
         # Create tasks from subtasks
         for each in sub_tasker.response.sub_tasks:
-            new_task = Task(description=each.description + " " + each.required_output + " " + str(each.sources_can_be_used) + " " + str(each.tools) + "Focus to complete the task with right result, Dont ask to human directly do it and give the result.")
+            new_task = Task(description=each.description + " " + each.required_output + " " + str(each.sources_can_be_used) + " " + str(each.tools) + "Focus to complete the task with right result, Dont ask to human directly do it and give the result.", price_id_=task.price_id, not_main_task=True)
             new_task.tools = task.tools
             sub_tasks.append(new_task)
 
         # Add the final task that will produce the original desired response format
-        end_task = Task(description=task.description, response_format=task.response_format)
+        end_task = Task(description=task.description, response_format=task.response_format, price_id_=task.price_id, not_main_task=True)
         sub_tasks.append(end_task)
 
         return sub_tasks

src/upsonic/client/printing.py

@@ -1,4 +1,5 @@
 from typing import Any
+from decimal import Decimal
 from rich.console import Console
 from rich.panel import Panel
 from rich.table import Table
@@ -9,6 +10,9 @@
 
 console = Console()
 
+# Global dictionary to store aggregated values by price_id
+price_id_summary = {}
+
 def spacing():
     console.print("")
 
@@ -148,10 +152,23 @@ def call_end(result: Any, llm_model: str, response_format: str, start_time: floa
 
 
 
-def agent_end(result: Any, llm_model: str, response_format: str, start_time: float, end_time: float, usage: dict, tool_count: int, context_count: int, debug: bool = False):
+def agent_end(result: Any, llm_model: str, response_format: str, start_time: float, end_time: float, usage: dict, tool_count: int, context_count: int, debug: bool = False, price_id:str = None):
     table = Table(show_header=False, expand=True, box=None)
     table.width = 60
 
+    # Track values if price_id is provided
+    if price_id:
+        estimated_cost = get_estimated_cost(usage['input_tokens'], usage['output_tokens'], llm_model)
+        if price_id not in price_id_summary:
+            price_id_summary[price_id] = {
+                'input_tokens': 0,
+                'output_tokens': 0,
+                'estimated_cost': 0.0
+            }
+        price_id_summary[price_id]['input_tokens'] += usage['input_tokens']
+        price_id_summary[price_id]['output_tokens'] += usage['output_tokens']
+        price_id_summary[price_id]['estimated_cost'] = Decimal(str(price_id_summary[price_id]['estimated_cost'])) + Decimal(str(estimated_cost).replace('$', '').replace('~', ''))
+
     table.add_row("[bold]LLM Model:[/bold]", f"{llm_model}")
     # Add spacing
     table.add_row("")
@@ -201,6 +218,47 @@ def agent_total_cost(total_input_tokens: int, total_output_tokens: int, total_ti
     console.print(panel)
     spacing()
 
+def print_price_id_summary(price_id: str, task) -> dict:
+    """
+    Get the summary of usage and costs for a specific price ID and print it in a formatted panel.
+    
+    Args:
+        price_id (str): The price ID to look up
+        
+    Returns:
+        dict: A dictionary containing the usage summary, or None if price_id not found
+    """
+    if price_id not in price_id_summary:
+        console.print("[bold red]Price ID not found![/bold red]")
+        return None
+
+    summary = price_id_summary[price_id].copy()
+    # Format the estimated cost to include $ symbol
+    summary['estimated_cost'] = f"${summary['estimated_cost']:.4f}"
+
+    # Create a table for pretty printing
+    table = Table(show_header=False, expand=True, box=None)
+    table.width = 60
+
+    table.add_row("[bold]Price ID:[/bold]", f"[magenta]{price_id}[/magenta]")
+    table.add_row("")  # Add spacing
+    table.add_row("[bold]Input Tokens:[/bold]", f"[magenta]{summary['input_tokens']:,}[/magenta]")
+    table.add_row("[bold]Output Tokens:[/bold]", f"[magenta]{summary['output_tokens']:,}[/magenta]")
+    table.add_row("[bold]Total Estimated Cost:[/bold]", f"[magenta]{summary['estimated_cost']}[/magenta]")
+
+    panel = Panel(
+        table,
+        title="[bold magenta]Upsonic - Price ID Summary[/bold magenta]",
+        border_style="magenta",
+        expand=True,
+        width=70
+    )
+
+    console.print(panel)
+    spacing()
+
+    return summary
+
 def agent_retry(retry_count: int, max_retries: int):
     table = Table(show_header=False, expand=True, box=None)
     table.width = 60
@@ -216,4 +274,25 @@ def agent_retry(retry_count: int, max_retries: int):
     )
 
     console.print(panel)
-    spacing()
+    spacing()
+
+def get_price_id_total_cost(price_id: str):
+    """
+    Get the total cost for a specific price ID.
+    
+    Args:
+        price_id (str): The price ID to get totals for
+        
+    Returns:
+        dict: Dictionary containing input tokens, output tokens, and estimated cost for the price ID.
+        None: If the price ID is not found.
+    """
+    if price_id not in price_id_summary:
+        return None
+
+    data = price_id_summary[price_id]
+    return {
+        'input_tokens': data['input_tokens'],
+        'output_tokens': data['output_tokens'],
+        'estimated_cost': float(data['estimated_cost'])
+    }

src/upsonic/client/tasks/tasks.py

@@ -5,19 +5,28 @@
 
 
 from .task_response import CustomTaskResponse, ObjectResponse
-
+from ..printing import get_price_id_total_cost
 class Task(BaseModel):
     description: str
     tools: list[Any] = []
     response_format: Union[Type[CustomTaskResponse], Type[ObjectResponse], None] = None
     _response: Any = None
     context: Any = None
+    price_id_: str = None
+    not_main_task: bool = False
 
     def __init__(self, description: str = None, **data):
         if description is not None:
             data["description"] = description
         super().__init__(**data)
 
+    @property
+    def price_id(self):
+        if self.price_id_ is None:
+            import uuid
+            self.price_id_ = str(uuid.uuid4())
+        return self.price_id_
+
     @property
     def response(self):
 
@@ -36,3 +45,7 @@ def response(self):
 
 
 
+    def get_total_cost(self):
+        if self.price_id_ is None:
+            return None
+        return get_price_id_total_cost(self.price_id)

src/upsonic/reliability_processor.py

@@ -282,7 +282,9 @@ def process_result(
                         prompt,
                         response_format=ValidationPoint,
                         tools=task.tools,
-                        context=context_strings  # Pass the processed context strings
+                        context=context_strings,  # Pass the processed context strings
+                        price_id_=task.price_id,
+                        not_main_task=True
                     )
                     validator_agent.do(validator_task)
                     setattr(validation_result, validation_type, validator_task.response)
@@ -306,7 +308,9 @@ def process_result(
                         formatted_prompt,
                         context=the_context,
                         response_format=task.response_format,
-                        tools=task.tools
+                        tools=task.tools,
+                        price_id_=task.price_id,
+                        not_main_task=True
                     )
                     editor_agent.do(editor_task)
                     return editor_task.response