diff --git a/src/backend/langflow/services/chat/manager.py b/src/backend/langflow/services/chat/manager.py index 22750d25a..afc004d2d 100644 --- a/src/backend/langflow/services/chat/manager.py +++ b/src/backend/langflow/services/chat/manager.py @@ -191,7 +191,7 @@ class ChatManager(Service): json_payload = await websocket.receive_json() try: payload = orjson.loads(json_payload) - except TypeError: + except Exception: payload = json_payload if "clear_history" in payload: self.chat_history.history[client_id] = [] diff --git a/src/backend/langflow/services/database/models/api_key/api_key.py b/src/backend/langflow/services/database/models/api_key/api_key.py index 601d060b5..50398bdbc 100644 --- a/src/backend/langflow/services/database/models/api_key/api_key.py +++ b/src/backend/langflow/services/database/models/api_key/api_key.py @@ -17,7 +17,9 @@ class ApiKeyBase(SQLModelSerializable): class ApiKey(ApiKeyBase, table=True): id: UUID = Field(default_factory=uuid4, primary_key=True, unique=True) + api_key: str = Field(index=True, unique=True) + hashed_api_key: str = Field(index=True) # User relationship user_id: UUID = Field(index=True, foreign_key="user.id") user: "User" = Relationship(back_populates="api_keys") diff --git a/src/backend/langflow/services/database/models/api_key/crud.py b/src/backend/langflow/services/database/models/api_key/crud.py index af697b6d5..ae1de1f0f 100644 --- a/src/backend/langflow/services/database/models/api_key/crud.py +++ b/src/backend/langflow/services/database/models/api_key/crud.py @@ -24,17 +24,20 @@ def create_api_key( generated_api_key = secrets.token_urlsafe(32) # hash the API key - hashed_api_key = get_password_hash(generated_api_key) + hashed = get_password_hash(generated_api_key) # Use the generated key to create the ApiKey object - - api_key = ApiKey(api_key=hashed_api_key, name=api_key_create.name, user_id=user_id) + masked_api_key = f"{'*' * 10}{generated_api_key[-4:]}" + api_key = ApiKey( + api_key=masked_api_key, + hashed_api_key=hashed, + name=api_key_create.name, + user_id=user_id, + ) session.add(api_key) session.commit() session.refresh(api_key) - unmasked = UnmaskedApiKeyRead.from_orm(api_key) - unmasked.api_key = generated_api_key - return unmasked + return UnmaskedApiKeyRead.from_orm(api_key) def delete_api_key(session: Session, api_key_id: UUID) -> None: diff --git a/src/frontend/src/controllers/API/index.ts b/src/frontend/src/controllers/API/index.ts index b49eb84ac..08bd9e721 100644 --- a/src/frontend/src/controllers/API/index.ts +++ b/src/frontend/src/controllers/API/index.ts @@ -466,9 +466,9 @@ export async function updateUser(user_id: string, user: Users) { } } -export async function getApiKey(user_id: String) { +export async function getApiKey() { try { - const res = await api.get(`${BASE_URL_API}api_key/${user_id}`); + const res = await api.get(`${BASE_URL_API}api_key`); if (res.status === 200) { return res.data; } @@ -478,9 +478,9 @@ export async function getApiKey(user_id: String) { } } -export async function createApiKey(user_id: string) { +export async function createApiKey(name:string) { try { - const res = await api.post(`${BASE_URL_API}api_key/${user_id}`); + const res = await api.post(`${BASE_URL_API}api_key`,{name}); if (res.status === 200) { return res.data; } @@ -491,9 +491,9 @@ export async function createApiKey(user_id: string) { } -export async function deleteApiKey(user_id: string) { +export async function deleteApiKey(api_key: string) { try { - const res = await api.delete(`${BASE_URL_API}api_key/${user_id}`); + const res = await api.delete(`${BASE_URL_API}api_key/${api_key}`); if (res.status === 200) { return res.data; } diff --git a/src/frontend/src/modals/SecretKeyModal/index.tsx b/src/frontend/src/modals/SecretKeyModal/index.tsx index 393088407..02d33244e 100644 --- a/src/frontend/src/modals/SecretKeyModal/index.tsx +++ b/src/frontend/src/modals/SecretKeyModal/index.tsx @@ -72,7 +72,7 @@ export default function SecretKeyModal({ }; function handleAddNewKey() { - createApiKey(data) + createApiKey(apiKeyName) .then((res) => { setApiKeyValue(res["api_key"]); }) diff --git a/src/frontend/src/modals/formModal/index.tsx b/src/frontend/src/modals/formModal/index.tsx index acca662d3..8f56d1f49 100644 --- a/src/frontend/src/modals/formModal/index.tsx +++ b/src/frontend/src/modals/formModal/index.tsx @@ -26,6 +26,7 @@ import { CHAT_FORM_DIALOG_SUBTITLE } from "../../constants/constants"; import { TabsContext } from "../../contexts/tabsContext"; import { TabsState } from "../../types/tabs"; import { validateNodes } from "../../utils/reactflowUtils"; +import { AuthContext } from "../../contexts/authContext"; export default function FormModal({ flow, @@ -60,6 +61,7 @@ export default function FormModal({ const [chatHistory, setChatHistory] = useState([]); const { reactFlowInstance } = useContext(typesContext); + const {accessToken} = useContext(AuthContext); const { setErrorData } = useContext(alertContext); const ws = useRef(null); const [lockChat, setLockChat] = useState(false); @@ -160,7 +162,7 @@ export default function FormModal({ }, 1000); } } - + //TODO improve check of user authentication function getWebSocketUrl( chatId: string, isDevelopment: boolean = false @@ -173,7 +175,7 @@ export default function FormModal({ return `${ isDevelopment ? "ws" : webSocketProtocol - }://${host}${chatEndpoint}`; + }://${host}${chatEndpoint}?token=${accessToken}`; } function handleWsMessage(data: any) { diff --git a/src/frontend/src/pages/ApiKeysPage/index.tsx b/src/frontend/src/pages/ApiKeysPage/index.tsx index 5d8db5a62..429f3e9f5 100644 --- a/src/frontend/src/pages/ApiKeysPage/index.tsx +++ b/src/frontend/src/pages/ApiKeysPage/index.tsx @@ -36,7 +36,7 @@ export default function ApiKeysPage() { function getKeys() { setLoadingKeys(true); if (userData) { - getApiKey(userData.id) + getApiKey() .then((keys: [ApiKey]) => { keysList.current = keys["api_keys"]; setUserId(keys["user_id"]); @@ -196,6 +196,7 @@ export default function ApiKeysPage() { data={api_keys.id} index={index} onConfirm={(index, keys) => { + console.log(keys); handleDeleteKey(keys); }} >