Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/management_endpoints/usage_endpoints/endpoints.py: 93%

1""" 

2USAGE AI CHAT ENDPOINTS 

3 

4/usage/ai/chat - Stream AI chat responses about usage data 

5""" 

6 

7from typing import Final, Literal 

8 

9from fastapi import APIRouter, Depends, Request 

10from fastapi.responses import StreamingResponse 

11from pydantic import BaseModel, Field 

12 

13import litellm 

14from litellm.proxy._types import UserAPIKeyAuth 

15from litellm.proxy.auth.user_api_key_auth import user_api_key_auth 

16from litellm.proxy.common_utils.sse_keepalive import ( 

17 SSE_COMMENT_PING, 

18 wrap_sse_stream_with_keepalive_pings, 

19) 

20 

21router: Final = APIRouter() 

22 

23 

24class ChatMessage(BaseModel): 

25 role: Literal["user", "assistant"] 

26 content: str 

27 

28 

29class UsageAIChatRequest(BaseModel): 

30 messages: list[ChatMessage] = Field(..., description="Chat messages (user/assistant history)") 

31 model: str | None = Field(default=None, description="Model to use for AI chat") 

32 

33 

34@router.post( 

35 "/usage/ai/chat", 

36 tags=["Budget & Spend Tracking"], 

37 dependencies=[Depends(user_api_key_auth)], 

38) 

39async def usage_ai_chat( 

40 data: UsageAIChatRequest, 

41 request: Request, 

42 user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), 

43): 

44 """ 

45 AI chat about usage data. Streams SSE events with the AI response. 

46 The AI agent has access to tools that query aggregated daily activity data. 

47 """ 

48 from litellm.proxy.management_endpoints.common_utils import ( 

49 _user_has_admin_view, 

50 require_caller_user_id_for_non_admin, 

51 ) 

52 from litellm.proxy.management_endpoints.usage_endpoints.ai_usage_chat import ( 

53 stream_usage_ai_chat, 

54 ) 

55 

56 is_admin: Final = _user_has_admin_view(user_api_key_dict) 

57 if is_admin: 57 ↛ 60line 57 didn't jump to line 60 because the condition on line 57 was always true

58 user_id = user_api_key_dict.user_id 

59 else: 

60 user_id = require_caller_user_id_for_non_admin(user_api_key_dict) 

61 messages: Final = [{"role": m.role, "content": m.content} for m in data.messages] 

62 

63 return StreamingResponse( 

64 wrap_sse_stream_with_keepalive_pings( 

65 stream_usage_ai_chat( 

66 messages=messages, 

67 model=data.model, 

68 user_id=user_id, 

69 is_admin=is_admin, 

70 ), 

71 ping_interval_seconds=litellm.sse_keepalive_ping_interval_seconds, 

72 ping_chunk=SSE_COMMENT_PING, 

73 ), 

74 media_type="text/event-stream", 

75 headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"}, 

76 )