diff --git a/src/upsonic/client/level_two/agent.py b/src/upsonic/client/level_two/agent.py index 5aba1643f6..fa22cc1889 100644 --- a/src/upsonic/client/level_two/agent.py +++ b/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,7 +537,7 @@ 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) @@ -536,12 +545,12 @@ def multiple(self, task: Task, llm_model: str = None): # 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 diff --git a/src/upsonic/client/printing.py b/src/upsonic/client/printing.py index ba3b513392..9b5a4bf2e3 100644 --- a/src/upsonic/client/printing.py +++ b/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() \ No newline at end of file + 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']) + } \ No newline at end of file diff --git a/src/upsonic/client/tasks/tasks.py b/src/upsonic/client/tasks/tasks.py index c929d7317a..b4e201d45a 100644 --- a/src/upsonic/client/tasks/tasks.py +++ b/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) \ No newline at end of file diff --git a/src/upsonic/reliability_processor.py b/src/upsonic/reliability_processor.py index 2295146a18..4c1e53dd47 100644 --- a/src/upsonic/reliability_processor.py +++ b/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