repo_description: Shell and coding agent on mcp clients
repo_topics: agent ai-agent ai-coding anthropic chatgpt claude claude-desktop custom-gpt llm llm-agent llm-code mcp mcp-server openai shell terminal
code: from typing import cast
  
  from openai.types.chat import (
      ChatCompletionAssistantMessageParam,
      ChatCompletionMessage,
      ChatCompletionMessageParam,
      ParsedChatCompletionMessage,
  )
  from tokenizers import Tokenizer  # type: ignore[import-untyped]
  
  from wcgw.client.common import CostData, History
  
  
  def get_input_cost(
      cost_map: CostData, enc: Tokenizer, history: History
  ) -> tuple[float, int]:
      input_tokens = 0
      for msg in history:
          content = msg["content"]
          refusal = msg.get("refusal")
          if isinstance(content, list):
              for part in content:
                  if "text" in part:
                      input_tokens += len(enc.encode(part["text"]))
          elif content is None:
              if refusal is None:
                  raise ValueError("Expected content or refusal to be present")
              input_tokens += len(enc.encode(str(refusal)))
          elif not isinstance(content, str):
              raise ValueError(f"Expected content to be string, got {type(content)}")
          else:
              input_tokens += len(enc.encode(content))
      cost = input_tokens * cost_map.cost_per_1m_input_tokens / 1_000_000
      return cost, input_tokens
  
  
  def get_output_cost(
      cost_map: CostData,
      enc: Tokenizer,
      item: ChatCompletionMessage | ChatCompletionMessageParam,
  ) -> tuple[float, int]:
      if isinstance(item, ChatCompletionMessage):
          content = item.content
          if not isinstance(content, str):
              raise ValueError(f"Expected content to be string, got {type(content)}")
      else:
          if not isinstance(item["content"], str):
              raise ValueError(
                  f"Expected content to be string, got {type(item['content'])}"
              )
          content = item["content"]
          if item["role"] == "tool":
              return 0, 0
      output_tokens = len(enc.encode(content))
  
      if "tool_calls" in item:
          item = cast(ChatCompletionAssistantMessageParam, item)
          toolcalls = item["tool_calls"]
          for tool_call in toolcalls or []:
              output_tokens += len(enc.encode(tool_call["function"]["arguments"]))
      elif isinstance(item, ParsedChatCompletionMessage):
          if item.tool_calls:
              for tool_callf in item.tool_calls:
                  output_tokens += len(enc.encode(tool_callf.function.arguments))
  
      cost = output_tokens * cost_map.cost_per_1m_output_tokens / 1_000_000
      return cost, output_tokens
  