Add mocked token-usage endpoints (#780)

* Add mocked token-usage endpoints

* Make conversations also use TokenUsage

Author: aponcedeleonch

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
+            ),
+        ],
+    )

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.

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