diff --git a/backend/backend/core/routers/chat_message/views/message_views.py b/backend/backend/core/routers/chat_message/views/message_views.py index 2e9149b..b519b9c 100644 --- a/backend/backend/core/routers/chat_message/views/message_views.py +++ b/backend/backend/core/routers/chat_message/views/message_views.py @@ -52,7 +52,7 @@ def persist_prompt(self, request: Request, *args, **kwargs) -> Response: def get_token_usage(self, request, project_id=None, chat_id=None, chat_message_id=None, *args, **kwargs) -> Response: """ Get token usage data for a specific chat message. - + Returns: - remaining_balance: Current token balance - total_consumed: Total tokens consumed @@ -67,18 +67,18 @@ def get_token_usage(self, request, project_id=None, chat_id=None, chat_message_i {"error": "Organization header is required"}, status=status.HTTP_400_BAD_REQUEST ) - + # Get token usage data using the existing function token_data = get_token_usage_data(org_id, str(chat_message_id), str(chat_id)) - + if token_data is None: return Response( {"error": "Failed to retrieve token usage data"}, status=status.HTTP_500_INTERNAL_SERVER_ERROR ) - + return Response(token_data, status=status.HTTP_200_OK) - + except Exception as e: return Response( {"error": f"Error retrieving token usage: {str(e)}"}, diff --git a/frontend/src/ide/chat-ai/ExistingChat.jsx b/frontend/src/ide/chat-ai/ExistingChat.jsx index 89c2810..a63f3ca 100644 --- a/frontend/src/ide/chat-ai/ExistingChat.jsx +++ b/frontend/src/ide/chat-ai/ExistingChat.jsx @@ -16,7 +16,17 @@ import { TodoGuide } from "./TodoGuide"; import { OnboardingGuide } from "./OnboardingGuide"; import { OnboardingCompletionPopup } from "./OnboardingCompletionPopup"; import { useNotificationService } from "../../service/notification-service"; +import { useAxiosPrivate } from "../../service/axios-service"; import { SpinnerLoader } from "../../widgets/spinner_loader"; +import { useSessionStore } from "../../store/session-store"; + +// Cloud-only: fetch per-message token usage (unavailable in OSS — import fails gracefully) +let getTokenUsage = null; +try { + ({ getTokenUsage } = require("../../plugins/token-management/token-usage")); +} catch { + // OSS: token usage API not available +} const ExistingChat = memo(function ExistingChat({ selectedChatId, @@ -61,8 +71,9 @@ const ExistingChat = memo(function ExistingChat({ onSkipCurrentTask, onSendButtonClick, }) { - const { getChatMessagesByChatId, getTokenUsage, updateChatName } = - useChatAIService(); + const { getChatMessagesByChatId, updateChatName } = useChatAIService(); + const axiosPrivate = useAxiosPrivate(); + const isCloud = useSessionStore((state) => state.sessionDetails?.is_cloud); const chatContainerRef = useRef(null); const [isLoadingChats, setIsLoadingChats] = useState(false); @@ -262,10 +273,15 @@ const ExistingChat = memo(function ExistingChat({ : [msg.response].filter(Boolean), })); - // Fetch token usage for all messages to display in historical conversations - if (updatedData.length > 0) { + // Fetch token usage for all messages to display in historical conversations. + // Only available in cloud mode via the token-usage plugin. + if (getTokenUsage && isCloud && updatedData.length > 0) { const tokenUsagePromises = updatedData.map((msg) => - getTokenUsage(selectedChatId, msg.chat_message_id).catch(() => null) + getTokenUsage( + axiosPrivate, + selectedChatId, + msg.chat_message_id + ).catch(() => null) ); const tokenUsageResults = await Promise.all(tokenUsagePromises); diff --git a/frontend/src/ide/chat-ai/services.js b/frontend/src/ide/chat-ai/services.js index c54e66f..1fb5f7a 100644 --- a/frontend/src/ide/chat-ai/services.js +++ b/frontend/src/ide/chat-ai/services.js @@ -134,22 +134,6 @@ export function useChatAIService() { return response.data; }; - const getTokenUsage = async (chatId, chatMessageId) => { - const url = `/api/v1/visitran/${orgId}/project/${projectId}/chat/${chatId}/chat-message/${chatMessageId}/token-usage/`; - try { - const response = await axiosPrivate.get(url, { - headers: { - "Content-Type": "application/json", - "X-Organization": orgId, - }, - }); - return response.data || null; - } catch (error) { - console.warn("Failed to fetch token usage:", error); - return null; - } - }; - return { getAllChats, deleteChatById, @@ -161,6 +145,5 @@ export function useChatAIService() { getChatLlmModels, completeOnboardingTask, getOnboardingStatus, - getTokenUsage, }; }