Add user usage model and tracking endpoints (#443)

Signed-off-by: Trevor Grant <[email protected]>

Author: rawkintrevo

webapp/packages/api/user-service/agent_factory/__init__.py

@@ -5,7 +5,6 @@
     what_to_do_prompt_template, \
     how_to_use_swagger_tools, \
     how_to_use_gofannon_agents
-from litellm import acompletion
 from models.agent import GenerateCodeRequest, GenerateCodeResponse, Agent
 import json
 import asyncio
@@ -183,16 +182,19 @@ async def code_gen_with_thoughts():
     if provider == "openai":
         name_doc_config['response_format'] = { "type": "json_object" }
 
-    name_doc_gen_task = acompletion(
-        model=f"{provider}/{model}",
-        messages=name_doc_messages,
-        **name_doc_config
-    )
-    
+    async def name_doc_generation():
+        content, _ = await call_llm(
+            provider=provider,
+            model=model,
+            messages=name_doc_messages,
+            parameters=name_doc_config,
+        )
+        return content
+
     # ---- Run tasks concurrently ----
-    (code_body, thoughts), name_doc_response = await asyncio.gather(
+    (code_body, thoughts), name_doc_content = await asyncio.gather(
         code_gen_with_thoughts(),
-        name_doc_gen_task
+        name_doc_generation()
     )
 
     # ---- Process Code Generation Response ----
@@ -215,7 +217,6 @@ async def run(input_dict, tools):
     full_code = f"{header}\n{indented_body}"
 
     # ---- Process Name and Docstring Response ----
-    name_doc_content = name_doc_response.choices[0].message.content
     try:
         # Clean up potential markdown
         if name_doc_content.strip().startswith("```json"):

webapp/packages/api/user-service/main.py

@@ -4,6 +4,7 @@
 from fastapi.middleware.cors import CORSMiddleware
 from typing import Annotated, Optional, Dict, Any, List
 from pydantic import BaseModel, Field
+from pydantic.config import ConfigDict
 from datetime import datetime
 import uuid
 import json
@@ -29,18 +30,20 @@
 )
 
 from services.llm_service import call_llm
+from services.user_service import get_user_service, UserService
 
 # Import the shared provider configuration
 from config.provider_config import PROVIDER_CONFIG as APP_PROVIDER_CONFIG
 from config.routes_config import RouterConfig, resolve_router_configs
 from models.agent import (
-    GenerateCodeRequest, GenerateCodeResponse, RunCodeRequest, 
+    GenerateCodeRequest, GenerateCodeResponse, RunCodeRequest,
     RunCodeResponse, Agent, CreateAgentRequest, Deployment, DeployedApi
 )
 from models.demo import (
     GenerateDemoCodeRequest, GenerateDemoCodeResponse,
     CreateDemoAppRequest, DemoApp
 )
+from models.user import User
 
 from agent_factory.remote_mcp_client import RemoteMCPClient
 
@@ -122,6 +125,31 @@ class ClientLogPayload(BaseModel):
 class FetchSpecRequest(BaseModel):
     url: str
 
+
+class UpdateMonthlyAllowanceRequest(BaseModel):
+    monthly_allowance: float = Field(..., alias="monthlyAllowance")
+
+    model_config = ConfigDict(populate_by_name=True)
+
+
+class UpdateResetDateRequest(BaseModel):
+    allowance_reset_date: float = Field(..., alias="allowanceResetDate")
+
+    model_config = ConfigDict(populate_by_name=True)
+
+
+class UpdateSpendRemainingRequest(BaseModel):
+    spend_remaining: float = Field(..., alias="spendRemaining")
+
+    model_config = ConfigDict(populate_by_name=True)
+
+
+class AddUsageRequest(BaseModel):
+    response_cost: float = Field(..., alias="responseCost")
+    metadata: Optional[Dict[str, Any]] = None
+
+    model_config = ConfigDict(populate_by_name=True)
+
 # Import models after defining local ones to avoid circular dependencies
 from models.chat import ChatRequest, ChatMessage, ChatResponse, ProviderConfig, SessionData
 
@@ -134,11 +162,19 @@ def get_logger() -> ObservabilityService:
     """Dependency to get the observability service instance."""
     return get_observability_service()
 
+def get_user_service_dep(db: DatabaseService = Depends(get_db)) -> UserService:
+    return get_user_service(db)
+
 # Background task for LLM processing
 async def process_chat(ticket_id: str, request: ChatRequest, user: dict, req: Request):
     # Background tasks don't have access to dependency injection, so we get service instances directly
     db_service = get_database_service(settings)
+    user_service = get_user_service(db_service)
     logger = get_observability_service()
+    user_basic_info = {
+        "email": user.get("email"),
+        "name": user.get("name") or user.get("displayName")
+    }
     try:
         # Update ticket status
         ticket_data = {
@@ -207,8 +243,11 @@ async def process_chat(ticket_id: str, request: ChatRequest, user: dict, req: Re
                 model=request.model,
                 messages=messages,
                 parameters=request.parameters,
-                tools=built_in_tools if built_in_tools else None
-            )        
+                tools=built_in_tools if built_in_tools else None,
+                user_service=user_service,
+                user_id=user.get("uid"),
+                user_basic_info=user_basic_info,
+            )
 
         # Update ticket with success
         ticket_data.update({
@@ -349,6 +388,52 @@ def get_model_config(provider: str, model: str):
         raise HTTPException(status_code=404, detail="Model not found")
     return available_providers[provider]["models"][model]
 
+
+@router.get("/users/me", response_model=User)
+def get_current_user_profile(user: dict = Depends(get_current_user), user_service: UserService = Depends(get_user_service_dep)):
+    return user_service.get_user(user.get("uid", "anonymous"), user)
+
+
+@router.put("/users/me/monthly-allowance", response_model=User)
+def set_monthly_allowance(
+    request: UpdateMonthlyAllowanceRequest,
+    user: dict = Depends(get_current_user),
+    user_service: UserService = Depends(get_user_service_dep),
+):
+    return user_service.set_monthly_allowance(user.get("uid", "anonymous"), request.monthly_allowance, user)
+
+
+@router.put("/users/me/allowance-reset-date", response_model=User)
+def set_allowance_reset_date(
+    request: UpdateResetDateRequest,
+    user: dict = Depends(get_current_user),
+    user_service: UserService = Depends(get_user_service_dep),
+):
+    return user_service.set_reset_date(user.get("uid", "anonymous"), request.allowance_reset_date, user)
+
+
+@router.post("/users/me/reset-allowance", response_model=User)
+def reset_allowance(user: dict = Depends(get_current_user), user_service: UserService = Depends(get_user_service_dep)):
+    return user_service.reset_allowance(user.get("uid", "anonymous"), user)
+
+
+@router.put("/users/me/spend-remaining", response_model=User)
+def update_spend_remaining(
+    request: UpdateSpendRemainingRequest,
+    user: dict = Depends(get_current_user),
+    user_service: UserService = Depends(get_user_service_dep),
+):
+    return user_service.update_spend_remaining(user.get("uid", "anonymous"), request.spend_remaining, user)
+
+
+@router.post("/users/me/usage", response_model=User)
+def add_usage_entry(
+    request: AddUsageRequest,
+    user: dict = Depends(get_current_user),
+    user_service: UserService = Depends(get_user_service_dep),
+):
+    return user_service.add_usage(user.get("uid", "anonymous"), request.response_cost, request.metadata, user)
+
 @router.post("/chat")
 async def chat(request: ChatRequest, req: Request, background_tasks: BackgroundTasks, user: dict = Depends(get_current_user)):
     """Submit a chat request and get a ticket ID"""

webapp/packages/api/user-service/models/user.py

@@ -0,0 +1,49 @@
+from datetime import datetime
+from typing import List, Optional, Any
+
+from pydantic import BaseModel, Field
+from pydantic.alias_generators import to_camel
+from pydantic.config import ConfigDict
+
+
+class UsageEntry(BaseModel):
+    timestamp: datetime = Field(default_factory=datetime.utcnow)
+    response_cost: float = Field(alias="responseCost")
+    metadata: Optional[Any] = None
+
+    model_config = ConfigDict(populate_by_name=True, alias_generator=to_camel)
+
+
+class UsageInfo(BaseModel):
+    monthly_allowance: float = Field(default=100.0, alias="monthlyAllowance")
+    allowance_reset_date: float = Field(default=0.0, alias="allowanceResetDate")
+    spend_remaining: float = Field(default=100.0, alias="spendRemaining")
+    usage: List[UsageEntry] = Field(default_factory=list)
+
+    model_config = ConfigDict(populate_by_name=True, alias_generator=to_camel)
+
+
+class BillingInfo(BaseModel):
+    plan: Optional[str] = None
+    status: Optional[str] = None
+
+    model_config = ConfigDict(populate_by_name=True, alias_generator=to_camel)
+
+
+class BasicInfo(BaseModel):
+    display_name: Optional[str] = Field(default=None, alias="displayName")
+    email: Optional[str] = None
+
+    model_config = ConfigDict(populate_by_name=True, alias_generator=to_camel)
+
+
+class User(BaseModel):
+    id: str = Field(alias="_id")
+    rev: Optional[str] = Field(default=None, alias="_rev")
+    created_at: datetime = Field(default_factory=datetime.utcnow, alias="createdAt")
+    updated_at: datetime = Field(default_factory=datetime.utcnow, alias="updatedAt")
+    basic_info: BasicInfo = Field(default_factory=BasicInfo, alias="basicInfo")
+    billing_info: BillingInfo = Field(default_factory=BillingInfo, alias="billingInfo")
+    usage_info: UsageInfo = Field(default_factory=UsageInfo, alias="usageInfo")
+
+    model_config = ConfigDict(populate_by_name=True, alias_generator=to_camel)

webapp/packages/api/user-service/services/llm_service.py

@@ -1,27 +1,54 @@
 import asyncio
 import json
+from config import settings
 from config.provider_config import PROVIDER_CONFIG
+from services.database_service import get_database_service
+from services.user_service import get_user_service
 from typing import Any, Dict, List, Tuple, Optional
 import litellm
 
 from services.litellm_logger import ensure_litellm_logging
+from services.user_service import UserService
 
 ensure_litellm_logging()
 
+def _extract_response_cost(response_obj: Any) -> Optional[float]:
+    standard_logging = None
+    if hasattr(response_obj, "_hidden_params") and isinstance(response_obj._hidden_params, dict):
+        standard_logging = response_obj._hidden_params.get("standard_logging_object")
+    if isinstance(standard_logging, dict):
+        try:
+            cost_value = standard_logging.get("response_cost")
+            if cost_value is not None:
+                return float(cost_value)
+        except Exception:
+            pass
+    usage = getattr(response_obj, "usage", None)
+    if usage and getattr(usage, "total_cost", None) is not None:
+        try:
+            return float(getattr(usage, "total_cost"))
+        except Exception:
+            return None
+    return None
+
+
 async def call_llm(
     provider: str,
     model: str,
     messages: List[Dict[str, Any]],
     parameters: Dict[str, Any],
     tools: Optional[List[Dict[str, Any]]] = None,
+    user_service: Optional[UserService] = None,
+    user_id: Optional[str] = None,
+    user_basic_info: Optional[Dict[str, Any]] = None,
 ) -> Tuple[str, Any]:
     """
     Calls the specified language model using litellm, handling different API styles.
     Returns a tuple of (content, thoughts).
     """
     model_config = PROVIDER_CONFIG.get(provider, {}).get("models", {}).get(model, {})
     api_style = model_config.get("api_style")
-    
+
     model_string = f"{provider}/{model}"
 
     thoughts = None
@@ -38,6 +65,14 @@ async def call_llm(
     if tools:
         kwargs["tools"] = tools
 
+    if user_service is None:
+        user_service = get_user_service(get_database_service(settings))
+    if user_id is None:
+        user_id = "anonymous"
+
+    if user_service and user_id:
+        user_service.require_allowance(user_id, basic_info=user_basic_info)
+
     if api_style == "responses":
         # Use aresponses and aget_responses for OpenAI's special tools like built-in web search
         kwargs.pop('messages', None) # aresponses uses 'input' not 'messages'
@@ -121,4 +156,13 @@ async def call_llm(
     if thoughts is not None:
         thoughts = json.loads(json.dumps(thoughts, default=str))
 
+    if user_service and user_id:
+        response_cost = None
+        try:
+            response_cost = _extract_response_cost(final_response if api_style == "responses" else response)
+        except Exception:
+            response_cost = None
+        if response_cost is not None:
+            user_service.add_usage(user_id, response_cost, basic_info=user_basic_info)
+
     return content, thoughts