diff --git a/ai-hub/app/api/routes/sessions.py b/ai-hub/app/api/routes/sessions.py index 0524597..378b06f 100644 --- a/ai-hub/app/api/routes/sessions.py +++ b/ai-hub/app/api/routes/sessions.py @@ -3,6 +3,7 @@ from app.api.dependencies import ServiceContainer, get_db from app.api import schemas from typing import AsyncGenerator +from app.core.pipelines.validator import Validator def create_sessions_router(services: ServiceContainer) -> APIRouter: router = APIRouter(prefix="/sessions", tags=["Sessions"]) @@ -54,5 +55,28 @@ raise except Exception as e: raise HTTPException(status_code=500, detail=f"An error occurred: {e}") + + @router.get("/{session_id}/tokens", response_model=schemas.SessionTokenUsageResponse, summary="Get Session Token Usage") + def get_session_token_usage(session_id: int, db: Session = Depends(get_db)): + try: + messages = services.rag_service.get_message_history(db=db, session_id=session_id) + if messages is None: + raise HTTPException(status_code=404, detail=f"Session with ID {session_id} not found.") + + combined_text = " ".join([m.content for m in messages]) + validator = Validator() + token_count = len(validator.encoding.encode(combined_text)) + token_limit = validator.token_limit + percentage = round((token_count / token_limit) * 100, 2) if token_limit > 0 else 0.0 + + return schemas.SessionTokenUsageResponse( + token_count=token_count, + token_limit=token_limit, + percentage=percentage + ) + except HTTPException: + raise + except Exception as e: + raise HTTPException(status_code=500, detail=f"An error occurred: {e}") return router \ No newline at end of file diff --git a/ai-hub/app/api/routes/user.py b/ai-hub/app/api/routes/user.py index 7b194dc..523fcd0 100644 --- a/ai-hub/app/api/routes/user.py +++ b/ai-hub/app/api/routes/user.py @@ -40,9 +40,6 @@ def create_users_router(services: ServiceContainer) -> APIRouter: router = APIRouter(prefix="/users", tags=["Users"]) -def create_users_router(services: ServiceContainer) -> APIRouter: - router = APIRouter(prefix="/users", tags=["Users"]) - @router.get("/login", summary="Initiate OIDC Login Flow") async def login_redirect( request: Request, @@ -132,12 +129,14 @@ except requests.exceptions.RequestException as e: logger.error(f"Token exchange error: {e}") + if hasattr(e, 'response') and e.response is not None: + logger.error(f"Token exchange response body: {e.response.text}") raise HTTPException(status_code=500, detail=f"Failed to communicate with OIDC provider: {e}") except jwt.JWTDecodeError as e: logger.error(f"ID token decode error: {e}") raise HTTPException(status_code=400, detail="Failed to decode ID token.") except Exception as e: - logger.error(f"An unexpected error occurred: {e}") + logger.exception(f"An unexpected error occurred during OAuth callback: {e}") raise HTTPException(status_code=500, detail=f"An unexpected error occurred: {e}") @router.get("/me", response_model=schemas.UserStatus, summary="Get Current User Status") diff --git a/ai-hub/app/api/routes/workspace.py b/ai-hub/app/api/routes/workspace.py index e85872a..4583c68 100644 --- a/ai-hub/app/api/routes/workspace.py +++ b/ai-hub/app/api/routes/workspace.py @@ -6,21 +6,21 @@ def create_workspace_router(services: ServiceContainer) -> APIRouter: router = APIRouter() - @router.websocket("/ws/workspace/{session_id}") - async def websocket_endpoint(websocket: WebSocket, session_id: str): + @router.websocket("/ws/workspace/{client_id}") + async def websocket_endpoint(websocket: WebSocket, client_id: str): await websocket.accept() - print(f"WebSocket connection accepted for session: {session_id}") + print(f"WebSocket connection accepted for client: {client_id}") # Send a welcome message to confirm the connection is active await websocket.send_text(json.dumps({ "type": "connection_established", - "message": f"Connected to AI Hub. Session ID: {session_id}" + "message": f"Connected to AI Hub. Client ID: {client_id}" })) try: await websocket.send_text(json.dumps({ "type": "connection_established", - "message": f"Connected to AI Hub. Session ID: {session_id}" + "message": f"Connected to AI Hub. Client ID: {client_id}" })) while True: @@ -31,10 +31,10 @@ await services.workspace_service.dispatch_message(websocket, data) except WebSocketDisconnect: - print(f"WebSocket connection disconnected for session: {session_id}") + print(f"WebSocket connection disconnected for client: {client_id}") except Exception as e: print(f"An error occurred: {e}") finally: - print(f"Closing WebSocket for session: {session_id}") + print(f"Closing WebSocket for client: {client_id}") return router \ No newline at end of file diff --git a/ai-hub/app/api/schemas.py b/ai-hub/app/api/schemas.py index f92f020..1a853fa 100644 --- a/ai-hub/app/api/schemas.py +++ b/ai-hub/app/api/schemas.py @@ -90,6 +90,12 @@ session_id: int messages: List[Message] +class SessionTokenUsageResponse(BaseModel): + """Defines the response for retrieving a session's token usage.""" + token_count: int + token_limit: int + percentage: float + class SpeechRequest(BaseModel): text: str diff --git a/ai-hub/app/app.py b/ai-hub/app/app.py index f48ae13..cad3153 100644 --- a/ai-hub/app/app.py +++ b/ai-hub/app/app.py @@ -118,8 +118,7 @@ app.add_middleware( CORSMiddleware, - allow_origins=["*"], # <-- This is a compromised solution should not be used in production. - # In real production, the allow origins should specified with frontend address to only allow the frontend can send request to it. + allow_origins=["https://ai.jerxie.com", "http://localhost:8000", "http://localhost:8080", "http://localhost:443"], allow_credentials=True, allow_methods=["*"], # Allows all HTTP methods (GET, POST, PUT, DELETE, etc.) allow_headers=["*"], # Allows all headers diff --git a/ai-hub/app/config.yaml b/ai-hub/app/config.yaml index 0f9028c..bacb1fb 100644 --- a/ai-hub/app/config.yaml +++ b/ai-hub/app/config.yaml @@ -18,7 +18,7 @@ # The default model name for the DeepSeek LLM provider. deepseek_model_name: "deepseek-chat" # The default model name for the Gemini LLM provider. - gemini_model_name: "gemini-1.5-flash-latest" + gemini_model_name: "gemini-2.0-flash" vector_store: # The file path to save and load the FAISS index. diff --git a/ai-hub/app/core/providers/stt/general_main.py b/ai-hub/app/core/providers/stt/general_main.py index bae3f87..c8c20b8 100644 --- a/ai-hub/app/core/providers/stt/general_main.py +++ b/ai-hub/app/core/providers/stt/general_main.py @@ -7,7 +7,7 @@ api_key = "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" # Provide a valid audio file path - audio_file_path = "/app/ai-hub/integration_tests/test_data/test-audio.wav" # Replace with your test file + audio_file_path = "integration_tests/test_data/test-audio.wav" # Replace with your test file try: # Read the audio file as bytes diff --git a/ui/client-app/src/hooks/useCodeAssistant.js b/ui/client-app/src/hooks/useCodeAssistant.js index 95cd207..e7144a6 100644 --- a/ui/client-app/src/hooks/useCodeAssistant.js +++ b/ui/client-app/src/hooks/useCodeAssistant.js @@ -1,5 +1,6 @@ import { useState, useEffect, useRef, useCallback } from "react"; -import { connectToWebSocket } from "../services/websocket"; +import { connectToWebSocket, getSessionId } from "../services/websocket"; +import { getSessionTokenStatus, getSessionMessages } from "../services/apiService"; import { v4 as uuidv4 } from 'uuid'; const useCodeAssistant = ({ pageContainerRef }) => { @@ -12,18 +13,29 @@ const [errorMessage, setErrorMessage] = useState(""); const [showErrorModal, setShowErrorModal] = useState(false); const [sessionId, setSessionId] = useState(null); + const [tokenUsage, setTokenUsage] = useState({ token_count: 0, token_limit: 100000, percentage: 0 }); const sessionIdRef = useRef(null); // ✅ Always current sessionId const ws = useRef(null); const initialized = useRef(false); const dirHandleRef = useRef(null); + const fetchTokenUsage = useCallback(async () => { + if (!sessionIdRef.current) return; + try { + const usage = await getSessionTokenStatus(sessionIdRef.current); + setTokenUsage(usage); + } catch (err) { + console.warn("Failed to fetch token usage", err); + } + }, []); + const handleChatMessage = useCallback((message) => { console.log("Received chat message:", message); - setThinkingProcess((prev) => [...prev,{ - type: "system", - message: "AI processing is complete" - }]) + setThinkingProcess((prev) => [...prev, { + type: "system", + message: "AI processing is complete" + }]) setChatHistory((prev) => [...prev, { isUser: false, isPureAnswer: true, @@ -44,13 +56,13 @@ steps: message.steps, reasoning: message.reasoning }]); - if (message.done === true){ - setThinkingProcess((prev) => [...prev,{ + if (message.done === true) { + setThinkingProcess((prev) => [...prev, { type: "system", message: "AI processing is complete" }]) - setIsProcessing(false); - } else{ + setIsProcessing(false); + } else { setIsProcessing(true); } }, []); @@ -279,7 +291,28 @@ const setupConnection = async () => { try { - const { ws: newWs, sessionId: newSessionId } = await connectToWebSocket( + // Ensure we have a valid session before connecting or handling messages + const currentSessionId = await getSessionId("coding_assistant"); + setSessionId(currentSessionId); + sessionIdRef.current = currentSessionId; + + try { + const messagesData = await getSessionMessages(currentSessionId); + if (messagesData && messagesData.messages) { + const formattedHistory = messagesData.messages.map((msg) => ({ + isUser: msg.sender === "user", + isPureAnswer: true, + text: msg.content, + })); + setChatHistory(formattedHistory); + } + } catch (historyErr) { + console.warn("Failed to load chat history", historyErr); + } + + fetchTokenUsage(); + + const { ws: newWs, clientId } = await connectToWebSocket( handleIncomingMessage, () => setConnectionStatus("connected"), () => { @@ -293,8 +326,6 @@ } ); ws.current = newWs; - setSessionId(newSessionId); - sessionIdRef.current = newSessionId; // ✅ Keep ref in sync } catch (error) { console.error("Setup failed:", error); } @@ -310,6 +341,20 @@ }, [handleIncomingMessage]); const handleSendChat = useCallback(async (text) => { + if (text.trim().toLowerCase() === "/new") { + setChatHistory([]); + setThinkingProcess([{ + type: "system", + message: "Started a new session." + }]); + localStorage.removeItem("sessionId_coding_assistant"); + const newSessionId = await getSessionId("coding_assistant"); + setSessionId(newSessionId); + sessionIdRef.current = newSessionId; + fetchTokenUsage(); + return; + } + if (ws.current && ws.current.readyState === WebSocket.OPEN) { setChatHistory((prev) => [...prev, { isUser: true, text }]); setIsProcessing(true); @@ -371,6 +416,7 @@ isPaused, errorMessage, showErrorModal, + tokenUsage, handleSendChat, handleSelectFolder, handlePause, diff --git a/ui/client-app/src/hooks/useVoiceChat.js b/ui/client-app/src/hooks/useVoiceChat.js index 4c74745..91babbf 100644 --- a/ui/client-app/src/hooks/useVoiceChat.js +++ b/ui/client-app/src/hooks/useVoiceChat.js @@ -3,13 +3,16 @@ // This file is a custom React hook that contains all the stateful logic // and side effects for the voice chat application. -import { useState, useRef, useEffect } from "react"; +import { useState, useRef, useEffect, useCallback } from "react"; import { createSession, transcribeAudio, chatWithAI, streamSpeech, + getSessionMessages, + getSessionTokenStatus } from "../services/apiService"; +import { getSessionId } from "../services/websocket"; import { stopAllPlayingAudio, stopAllMediaStreams, @@ -37,8 +40,10 @@ const [sessionId, setSessionId] = useState(null); const [isAutoMode, setIsAutoMode] = useState(false); const [isAutoListening, setIsAutoListening] = useState(false); + const [tokenUsage, setTokenUsage] = useState({ token_count: 0, token_limit: 100000, percentage: 0 }); // All refs must be declared here, inside the custom hook. + const sessionIdRef = useRef(null); const mediaRecorderRef = useRef(null); const audioChunksRef = useRef([]); const audioContextRef = useRef(null); @@ -51,15 +56,44 @@ const lastRequestTimeRef = useRef(0); const streamRef = useRef(null); + const fetchTokenUsage = useCallback(async () => { + if (!sessionIdRef.current) return; + try { + const usage = await getSessionTokenStatus(sessionIdRef.current); + setTokenUsage(usage); + } catch (err) { + console.warn("Failed to fetch voice token usage", err); + } + }, []); + // --- Initial Session Creation Effect --- useEffect(() => { const startSession = async () => { setIsBusy(true); - setStatus("Starting new chat session..."); + setStatus("Loading chat session..."); try { - const session = await createSession(); - setSessionId(session.id); - console.log(`Session created with ID: ${session.id}`); + const currentSessionId = await getSessionId("voice_chat"); + setSessionId(currentSessionId); + sessionIdRef.current = currentSessionId; + + // Try to load chat history + try { + const messagesData = await getSessionMessages(currentSessionId); + if (messagesData && messagesData.messages && messagesData.messages.length > 0) { + const formattedHistory = messagesData.messages.map((msg) => ({ + isUser: msg.sender === "user", + text: msg.content, + })); + setChatHistory(formattedHistory); + } + } catch (historyErr) { + console.warn("Failed to load voice chat history", historyErr); + } + + // Load initial tokens + await fetchTokenUsage(); + + console.log(`Voice Session loaded with ID: ${currentSessionId}`); setStatus("Click the microphone to start recording."); } catch (err) { console.error("Error creating session:", err); @@ -106,7 +140,7 @@ audioContextRef.current = new (window.AudioContext || window.webkitAudioContext)(); playbackTimeRef.current = audioContextRef.current.currentTime; } - + const audioContext = audioContextRef.current; const onChunkReceived = (rawFloat32Data) => { @@ -194,6 +228,7 @@ const aiText = await chatWithAI(sessionId, userText); addMessage(aiText, false); + fetchTokenUsage(); await playStreamingAudio(aiText); } catch (error) { console.error("Conversation processing failed:", error); @@ -375,6 +410,31 @@ } }; + const handleNewSession = async () => { + setChatHistory([ + { + text: "Hello! I'm an AI assistant. How can I help you today?", + isUser: false, + }, + ]); + localStorage.removeItem("sessionId_voice_chat"); + + setIsBusy(true); + setStatus("Starting new session..."); + try { + const newSessionId = await getSessionId("voice_chat"); + setSessionId(newSessionId); + sessionIdRef.current = newSessionId; + fetchTokenUsage(); + setStatus("Click the microphone to start recording."); + } catch (err) { + console.error("Failed to start new voice session", err); + setStatus("Error creating new session."); + } finally { + setIsBusy(false); + } + }; + return { chatHistory, status, @@ -385,8 +445,10 @@ sessionId, showErrorModal, errorMessage, + tokenUsage, setIsAutoMode, handleMicClick, + handleNewSession, setShowErrorModal, }; }; diff --git a/ui/client-app/src/pages/CodingAssistantPage.js b/ui/client-app/src/pages/CodingAssistantPage.js index 6df1758..2b61347 100644 --- a/ui/client-app/src/pages/CodingAssistantPage.js +++ b/ui/client-app/src/pages/CodingAssistantPage.js @@ -22,13 +22,14 @@ isPaused, errorMessage, showErrorModal, + tokenUsage, handleSendChat, handleSelectFolder, handlePause, handleStop, setShowErrorModal, } = useCodeAssistant({ pageContainerRef }); - + // State to manage the visibility of the right panel const [isPanelExpanded, setIsPanelExpanded] = useState(false); @@ -50,9 +51,18 @@ {/* Area 1: Chat with LLM */}

- Chat with the LLM -
- + Chat with the LLM +
+
+ {tokenUsage?.token_count || 0} / {tokenUsage?.token_limit || 100000} tokens ({tokenUsage?.percentage || 0}%) +
+ + Unlock powerful coding assistance with RAG!
diff --git a/ui/client-app/src/pages/VoiceChatPage.js b/ui/client-app/src/pages/VoiceChatPage.js index 42572f0..976d3a5 100644 --- a/ui/client-app/src/pages/VoiceChatPage.js +++ b/ui/client-app/src/pages/VoiceChatPage.js @@ -15,8 +15,10 @@ isAutoListening, showErrorModal, errorMessage, + tokenUsage, setIsAutoMode, handleMicClick, + handleNewSession, setShowErrorModal, } = useVoiceChat({ chatContainerRef }); @@ -32,6 +34,23 @@ return (
+ + {/* Header bar mirroring CodingAssistantPage */} +
+

Voice Chat Assistant

+
+
+ {tokenUsage?.token_count || 0} / {tokenUsage?.token_limit || 100000} tokens ({tokenUsage?.percentage || 0}%) +
+ +
+
+
diff --git a/ui/client-app/src/services/apiService.js b/ui/client-app/src/services/apiService.js index 18859ab..638a15e 100644 --- a/ui/client-app/src/services/apiService.js +++ b/ui/client-app/src/services/apiService.js @@ -4,11 +4,13 @@ import { convertPcmToFloat32 } from "./audioUtils"; // Read base API URL from environment variables, defaulting to localhost -const API_BASE_URL = process.env.REACT_APP_API_BASE_URL || "http://localhost:8001"; +const API_BASE_URL = process.env.REACT_APP_API_BASE_URL || "http://localhost:8001"; export { API_BASE_URL }; const STT_ENDPOINT = `${API_BASE_URL}/stt/transcribe`; const SESSIONS_CREATE_ENDPOINT = `${API_BASE_URL}/sessions/`; const SESSIONS_CHAT_ENDPOINT = (id) => `${API_BASE_URL}/sessions/${id}/chat`; +const SESSIONS_TOKEN_ENDPOINT = (id) => `${API_BASE_URL}/sessions/${id}/tokens`; +const SESSIONS_MESSAGES_ENDPOINT = (id) => `${API_BASE_URL}/sessions/${id}/messages`; const TTS_ENDPOINT = `${API_BASE_URL}/speech`; const USERS_LOGIN_ENDPOINT = `${API_BASE_URL}/users/login`; const USERS_LOGOUT_ENDPOINT = `${API_BASE_URL}/users/logout`; @@ -96,6 +98,40 @@ return await response.json(); }; +/** + * Fetches the token usage status for a session. + * @param {string} sessionId - The session ID to fetch tokens for. + * @returns {Promise} The token usage object. + */ +export const getSessionTokenStatus = async (sessionId) => { + const userId = getUserId(); + const response = await fetch(SESSIONS_TOKEN_ENDPOINT(sessionId), { + method: "GET", + headers: { "X-User-ID": userId }, + }); + if (!response.ok) { + throw new Error(`Failed to fetch token status. Status: ${response.status}`); + } + return await response.json(); +}; + +/** + * Fetches the chat history for a session. + * @param {string} sessionId - The session ID to fetch messages for. + * @returns {Promise} The message history object. + */ +export const getSessionMessages = async (sessionId) => { + const userId = getUserId(); + const response = await fetch(SESSIONS_MESSAGES_ENDPOINT(sessionId), { + method: "GET", + headers: { "X-User-ID": userId }, + }); + if (!response.ok) { + throw new Error(`Failed to fetch session messages. Status: ${response.status}`); + } + return await response.json(); +}; + // --- Unchanged Functions --- /** @@ -186,7 +222,7 @@ const toConvert = combined.slice(0, length); leftover = combined.slice(length); const float32Raw = convertPcmToFloat32(toConvert); - + // Pass the raw float32 data to the caller for resampling onData(float32Raw); } diff --git a/ui/client-app/src/services/websocket.js b/ui/client-app/src/services/websocket.js index d640581..8d82933 100644 --- a/ui/client-app/src/services/websocket.js +++ b/ui/client-app/src/services/websocket.js @@ -1,23 +1,39 @@ // src/services/websocket.js import { createSession, API_BASE_URL } from "./apiService"; +import { v4 as uuidv4 } from "uuid"; + +/** + * Gets a client ID from localStorage or creates a new one. + * @returns {string} The client ID. + */ +export const getClientId = () => { + let clientId = localStorage.getItem("clientId"); + if (!clientId) { + clientId = uuidv4(); + localStorage.setItem("clientId", clientId); + } + return clientId; +}; /** * Gets a session ID from localStorage or creates a new one via the API. + * @param {string} featureNamespace - The key suffix to use for isolated storage. * @returns {Promise} The session ID. */ -export const getSessionId = async () => { - let sessionId = localStorage.getItem("sessionId"); +export const getSessionId = async (featureNamespace = "default") => { + const storageKey = featureNamespace === "default" ? "sessionId" : `sessionId_${featureNamespace}`; + let sessionId = localStorage.getItem(storageKey); if (!sessionId) { // No existing session, so create one via API const session = await createSession(); sessionId = session.id; // Store it in localStorage for reuse - localStorage.setItem("sessionId", sessionId); + localStorage.setItem(storageKey, sessionId); } - console.log("Using session ID:", sessionId); + console.log(`Using session ID for [${featureNamespace}]:`, sessionId); return sessionId; }; @@ -27,7 +43,7 @@ * @param {function(): void} onOpenCallback - Callback when the connection is opened. * @param {function(): void} onCloseCallback - Callback when the connection is closed. * @param {function(Error): void} onErrorCallback - Callback when an error occurs. - * @returns {Promise<{ws: WebSocket, sessionId: string}>} The WebSocket instance and the session ID. + * @returns {Promise<{ws: WebSocket, clientId: string}>} The WebSocket instance and the client ID. */ export const connectToWebSocket = async ( onMessageCallback, @@ -36,21 +52,8 @@ onErrorCallback ) => { try { - let sessionId = localStorage.getItem("sessionId"); - - // NOTE: The line `sessionId = null;` has been removed as it was for testing purposes - // and would force a new session on every connection. - if (!sessionId) { - // No existing session, so create one via API - const session = await createSession(); - sessionId = session.id; - - // Store it in localStorage for reuse - localStorage.setItem("sessionId", sessionId); - } - - // You now have a valid sessionId, either reused or newly created - console.log("Using session ID:", sessionId); + const clientId = getClientId(); + console.log("Using client ID for connection:", clientId); // Correct the WebSocket URL based on the API_BASE_URL protocol. // Use `wss` for `https` and `ws` for `http`. @@ -63,7 +66,7 @@ pathname = ""; } - const websocketUrl = `${wsProtocol}://${url.host}${pathname}/ws/workspace/${sessionId}`; + const websocketUrl = `${wsProtocol}://${url.host}${pathname}/ws/workspace/${clientId}`; console.log("Connecting to WebSocket URL:", websocketUrl); const ws = new WebSocket(websocketUrl); @@ -88,10 +91,10 @@ if (onErrorCallback) onErrorCallback(error); }; - return { ws, sessionId }; + return { ws, clientId }; } catch (error) { - console.error("Failed to create session or connect to WebSocket:", error); + console.error("Failed to connect to WebSocket:", error); if (onErrorCallback) onErrorCallback(error); throw error; } diff --git a/ui/run_web.sh b/ui/run_web.sh index 3ce6bf4..8a45393 100644 --- a/ui/run_web.sh +++ b/ui/run_web.sh @@ -79,7 +79,7 @@ echo "--- Starting AI Hub Server, React frontend, and backend proxy ---" # Run backend and frontend concurrently -concurrently \ +npx concurrently \ --prefix "[{name}]" \ --names "aihub,tts-frontend" \ "LOG_LEVEL=DEBUG uvicorn $APP_MODULE --host $AI_HUB_HOST --log-level debug --port $AI_HUB_PORT $SSL_ARGS --reload" \