diff --git a/src/codegate/api/v1.py b/src/codegate/api/v1.py index 29f2d4d1..2a160b9e 100644 --- a/src/codegate/api/v1.py +++ b/src/codegate/api/v1.py @@ -470,3 +470,25 @@ def version_check(): "is_latest": None, "error": "An unexpected error occurred", } + + +@v1.get( + "/workspaces/{workspace_name}/token-usage", + tags=["Workspaces", "Token Usage"], + generate_unique_id_function=uniq_name, +) +async def get_workspace_token_usage(workspace_name: str) -> v1_models.TokenUsage: + """Get the token usage of a workspace.""" + # TODO: This is a dummy implementation. In the future, we should have a proper + # implementation that fetches the token usage from the database. + return v1_models.TokenUsage( + used_tokens=50, + tokens_by_model=[ + v1_models.TokenUsageByModel( + provider_type="openai", model="gpt-4o-mini", used_tokens=20 + ), + v1_models.TokenUsageByModel( + provider_type="anthropic", model="claude-3-5-sonnet-20241022", used_tokens=30 + ), + ], + ) diff --git a/src/codegate/api/v1_models.py b/src/codegate/api/v1_models.py index 7bf2d3db..fb4e90d3 100644 --- a/src/codegate/api/v1_models.py +++ b/src/codegate/api/v1_models.py @@ -114,6 +114,36 @@ class PartialQuestionAnswer(pydantic.BaseModel): answer: Optional[ChatMessage] +class ProviderType(str, Enum): + """ + Represents the different types of providers we support. + """ + + openai = "openai" + anthropic = "anthropic" + vllm = "vllm" + + +class TokenUsageByModel(pydantic.BaseModel): + """ + Represents the tokens used by a model. + """ + + provider_type: ProviderType + model: str + used_tokens: int + + +class TokenUsage(pydantic.BaseModel): + """ + Represents the tokens used. Includes the information of the tokens used by model. + `used_tokens` are the total tokens used in the `tokens_by_model` list. + """ + + tokens_by_model: List[TokenUsageByModel] + used_tokens: int + + class Conversation(pydantic.BaseModel): """ Represents a conversation. @@ -124,6 +154,7 @@ class Conversation(pydantic.BaseModel): type: QuestionType chat_id: str conversation_timestamp: datetime.datetime + token_usage: Optional[TokenUsage] class AlertConversation(pydantic.BaseModel): @@ -140,16 +171,6 @@ class AlertConversation(pydantic.BaseModel): timestamp: datetime.datetime -class ProviderType(str, Enum): - """ - Represents the different types of providers we support. - """ - - openai = "openai" - anthropic = "anthropic" - vllm = "vllm" - - class ProviderAuthType(str, Enum): """ Represents the different types of auth we support for providers. diff --git a/src/codegate/api/v1_processing.py b/src/codegate/api/v1_processing.py index ed2e119a..d5deac7d 100644 --- a/src/codegate/api/v1_processing.py +++ b/src/codegate/api/v1_processing.py @@ -331,6 +331,7 @@ async def match_conversations( type=first_partial_qa.partial_questions.type, chat_id=first_partial_qa.partial_questions.message_id, conversation_timestamp=first_partial_qa.partial_questions.timestamp, + token_usage=None, ) for qa in questions_answers: map_q_id_to_conversation[qa.question.message_id] = conversation