From 3119499851253f69acf2de0d4ce81466da79538e Mon Sep 17 00:00:00 2001 From: gozineb Date: Wed, 25 Oct 2023 16:59:19 +0200 Subject: [PATCH 1/8] =?UTF-8?q?=F0=9F=94=A5=20remove=20duplicate=20code?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- backend/routes/chat_routes.py | 7 ------- 1 file changed, 7 deletions(-) diff --git a/backend/routes/chat_routes.py b/backend/routes/chat_routes.py index eb71e05d9c88..b0c835244f37 100644 --- a/backend/routes/chat_routes.py +++ b/backend/routes/chat_routes.py @@ -186,14 +186,7 @@ async def create_question_handler( """ Add a new question to the chat. """ - if brain_id: - validate_brain_authorization( - brain_id=brain_id, - user_id=current_user.id, - required_roles=[RoleEnum.Viewer, RoleEnum.Editor, RoleEnum.Owner], - ) - # Retrieve user's OpenAI API key if brain_id: validate_brain_authorization( brain_id=brain_id, From c3b5612e240fb205bf8c62e1afff1026691402dd Mon Sep 17 00:00:00 2001 From: gozineb Date: Wed, 25 Oct 2023 18:05:16 +0200 Subject: [PATCH 2/8] =?UTF-8?q?=F0=9F=8F=97=EF=B8=8F=20add=20a=20chat=20st?= =?UTF-8?q?rategy=20for=20brainless=20and=20brainful=20questions?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- backend/routes/chat/__init_.py | 0 backend/routes/chat/brainful_chat.py | 6 ++++++ backend/routes/chat/brainless_chat.py | 6 ++++++ backend/routes/chat/factory.py | 11 +++++++++++ backend/routes/chat/interface.py | 7 +++++++ 5 files changed, 30 insertions(+) create mode 100644 backend/routes/chat/__init_.py create mode 100644 backend/routes/chat/brainful_chat.py create mode 100644 backend/routes/chat/brainless_chat.py create mode 100644 backend/routes/chat/factory.py create mode 100644 backend/routes/chat/interface.py diff --git a/backend/routes/chat/__init_.py b/backend/routes/chat/__init_.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/backend/routes/chat/brainful_chat.py b/backend/routes/chat/brainful_chat.py new file mode 100644 index 000000000000..81a738fcabec --- /dev/null +++ b/backend/routes/chat/brainful_chat.py @@ -0,0 +1,6 @@ +from routes.chat.interface import ChatInterface + + +class BrainfulChat(ChatInterface): + def validate_authorization(self, user_id, brain_id): + pass diff --git a/backend/routes/chat/brainless_chat.py b/backend/routes/chat/brainless_chat.py new file mode 100644 index 000000000000..0a1547b03901 --- /dev/null +++ b/backend/routes/chat/brainless_chat.py @@ -0,0 +1,6 @@ +from backend.routes.chat.interface import ChatInterface + + +class BrainlessChat(ChatInterface): + def validate_authorization(self, user_id, brain_id): + pass diff --git a/backend/routes/chat/factory.py b/backend/routes/chat/factory.py new file mode 100644 index 000000000000..792328feeffc --- /dev/null +++ b/backend/routes/chat/factory.py @@ -0,0 +1,11 @@ +from uuid import UUID + +from .brainful_chat import BrainfulChat +from .brainless_chat import BrainlessChat + + +def get_chat_strategy(brain_id: UUID | None = None): + if brain_id: + return BrainfulChat() + else: + return BrainlessChat() diff --git a/backend/routes/chat/interface.py b/backend/routes/chat/interface.py new file mode 100644 index 000000000000..49608ff2056d --- /dev/null +++ b/backend/routes/chat/interface.py @@ -0,0 +1,7 @@ +from abc import ABC, abstractmethod + + +class ChatInterface(ABC): + @abstractmethod + def validate_authorization(self, user_id, required_roles): + pass From 40f70bbefa2754e01da2164a2afb17e3b315c045 Mon Sep 17 00:00:00 2001 From: gozineb Date: Wed, 25 Oct 2023 18:06:03 +0200 Subject: [PATCH 3/8] =?UTF-8?q?=F0=9F=9A=9A=20extract=20authorization=20va?= =?UTF-8?q?lidation?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- backend/routes/chat/brainful_chat.py | 9 ++++++++- backend/routes/chat_routes.py | 11 +++++------ 2 files changed, 13 insertions(+), 7 deletions(-) diff --git a/backend/routes/chat/brainful_chat.py b/backend/routes/chat/brainful_chat.py index 81a738fcabec..e8c21a029b2d 100644 --- a/backend/routes/chat/brainful_chat.py +++ b/backend/routes/chat/brainful_chat.py @@ -1,6 +1,13 @@ +from routes.authorizations.brain_authorization import validate_brain_authorization +from routes.authorizations.types import RoleEnum from routes.chat.interface import ChatInterface class BrainfulChat(ChatInterface): def validate_authorization(self, user_id, brain_id): - pass + if brain_id: + validate_brain_authorization( + brain_id=brain_id, + user_id=user_id, + required_roles=[RoleEnum.Viewer, RoleEnum.Editor, RoleEnum.Owner], + ) diff --git a/backend/routes/chat_routes.py b/backend/routes/chat_routes.py index b0c835244f37..9f0affb87c8f 100644 --- a/backend/routes/chat_routes.py +++ b/backend/routes/chat_routes.py @@ -39,6 +39,8 @@ from routes.authorizations.brain_authorization import validate_brain_authorization from routes.authorizations.types import RoleEnum +from backend.routes.chat.factory import get_chat_strategy + chat_router = APIRouter() @@ -187,12 +189,9 @@ async def create_question_handler( Add a new question to the chat. """ - if brain_id: - validate_brain_authorization( - brain_id=brain_id, - user_id=current_user.id, - required_roles=[RoleEnum.Viewer, RoleEnum.Editor, RoleEnum.Owner], - ) + chat_instance = get_chat_strategy(brain_id) + + chat_instance.validate_authorization(user_id=current_user.id, brain_id=brain_id) current_user.openai_api_key = request.headers.get("Openai-Api-Key") brain = Brain(id=brain_id) From 3b8f95bace84778b5409520ede52fc2421c344c6 Mon Sep 17 00:00:00 2001 From: gozineb Date: Thu, 26 Oct 2023 21:16:57 +0200 Subject: [PATCH 4/8] =?UTF-8?q?=F0=9F=9A=9A=20extract=20chat=20utils?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- backend/routes/chat/brainless_chat.py | 2 +- backend/routes/chat/utils.py | 57 ++++++++++++++++++++++++ backend/routes/chat_routes.py | 64 ++++----------------------- 3 files changed, 67 insertions(+), 56 deletions(-) create mode 100644 backend/routes/chat/utils.py diff --git a/backend/routes/chat/brainless_chat.py b/backend/routes/chat/brainless_chat.py index 0a1547b03901..e93dae7aca29 100644 --- a/backend/routes/chat/brainless_chat.py +++ b/backend/routes/chat/brainless_chat.py @@ -1,4 +1,4 @@ -from backend.routes.chat.interface import ChatInterface +from routes.chat.interface import ChatInterface class BrainlessChat(ChatInterface): diff --git a/backend/routes/chat/utils.py b/backend/routes/chat/utils.py new file mode 100644 index 000000000000..fe40d32fc5e4 --- /dev/null +++ b/backend/routes/chat/utils.py @@ -0,0 +1,57 @@ +import time +from uuid import UUID + +from fastapi import HTTPException +from models import UserIdentity, UserUsage +from models.databases.supabase.supabase import SupabaseDB + + +class NullableUUID(UUID): + @classmethod + def __get_validators__(cls): + yield cls.validate + + @classmethod + def validate(cls, v) -> UUID | None: + if v == "": + return None + try: + return UUID(v) + except ValueError: + return None + + +def delete_chat_from_db(supabase_db: SupabaseDB, chat_id): + try: + supabase_db.delete_chat_history(chat_id) + except Exception as e: + print(e) + pass + try: + supabase_db.delete_chat(chat_id) + except Exception as e: + print(e) + pass + + +def check_user_requests_limit( + user: UserIdentity, +): + userDailyUsage = UserUsage( + id=user.id, email=user.email, openai_api_key=user.openai_api_key + ) + + userSettings = userDailyUsage.get_user_settings() + + date = time.strftime("%Y%m%d") + userDailyUsage.handle_increment_user_request_count(date) + + if user.openai_api_key is None: + daily_chat_credit = userSettings.get("daily_chat_credit", 0) + if int(userDailyUsage.daily_requests_count) >= int(daily_chat_credit): + raise HTTPException( + status_code=429, # pyright: ignore reportPrivateUsage=none + detail="You have reached the maximum number of requests for today.", # pyright: ignore reportPrivateUsage=none + ) + else: + pass diff --git a/backend/routes/chat_routes.py b/backend/routes/chat_routes.py index 9f0affb87c8f..3256d4fc8750 100644 --- a/backend/routes/chat_routes.py +++ b/backend/routes/chat_routes.py @@ -1,4 +1,3 @@ -import time from typing import List, Optional from uuid import UUID from venv import logger @@ -18,7 +17,6 @@ get_supabase_db, ) from models.databases.supabase.chats import QuestionAndAnswer -from models.databases.supabase.supabase import SupabaseDB from repository.brain import get_brain_details from repository.chat import ( ChatUpdatableProperties, @@ -38,63 +36,16 @@ from repository.user_identity import get_user_identity from routes.authorizations.brain_authorization import validate_brain_authorization from routes.authorizations.types import RoleEnum - -from backend.routes.chat.factory import get_chat_strategy +from routes.chat.factory import get_chat_strategy +from routes.chat.utils import ( + NullableUUID, + check_user_requests_limit, + delete_chat_from_db, +) chat_router = APIRouter() -class NullableUUID(UUID): - @classmethod - def __get_validators__(cls): - yield cls.validate - - @classmethod - def validate(cls, v) -> UUID | None: - if v == "": - return None - try: - return UUID(v) - except ValueError: - return None - - -def delete_chat_from_db(supabase_db: SupabaseDB, chat_id): - try: - supabase_db.delete_chat_history(chat_id) - except Exception as e: - print(e) - pass - try: - supabase_db.delete_chat(chat_id) - except Exception as e: - print(e) - pass - - -def check_user_requests_limit( - user: UserIdentity, -): - userDailyUsage = UserUsage( - id=user.id, email=user.email, openai_api_key=user.openai_api_key - ) - - userSettings = userDailyUsage.get_user_settings() - - date = time.strftime("%Y%m%d") - userDailyUsage.handle_increment_user_request_count(date) - - if user.openai_api_key is None: - daily_chat_credit = userSettings.get("daily_chat_credit", 0) - if int(userDailyUsage.daily_requests_count) >= int(daily_chat_credit): - raise HTTPException( - status_code=429, # pyright: ignore reportPrivateUsage=none - detail="You have reached the maximum number of requests for today.", # pyright: ignore reportPrivateUsage=none - ) - else: - pass - - @chat_router.get("/chat/healthz", tags=["Health"]) async def healthz(): return {"status": "ok"} @@ -286,6 +237,9 @@ async def create_stream_question_handler( required_roles=[RoleEnum.Viewer, RoleEnum.Editor, RoleEnum.Owner], ) + chat_instance = get_chat_strategy(brain_id) + chat_instance.validate_authorization(user_id=current_user.id, brain_id=brain_id) + # Retrieve user's OpenAI API key current_user.openai_api_key = request.headers.get("Openai-Api-Key") brain = Brain(id=brain_id) From 5ae40333e437e9c57aeab4edf2e9df68e50f194f Mon Sep 17 00:00:00 2001 From: gozineb Date: Thu, 26 Oct 2023 21:35:00 +0200 Subject: [PATCH 5/8] =?UTF-8?q?=F0=9F=9A=9A=20extract=20get=5Fanswer=5Fgen?= =?UTF-8?q?erator=20to=20Brainful/less=20strategy?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- backend/routes/chat/brainful_chat.py | 23 +++++++++++++++++++++++ backend/routes/chat/brainless_chat.py | 22 ++++++++++++++++++++++ backend/routes/chat/interface.py | 14 ++++++++++++++ backend/routes/chat_routes.py | 15 ++++++++------- 4 files changed, 67 insertions(+), 7 deletions(-) diff --git a/backend/routes/chat/brainful_chat.py b/backend/routes/chat/brainful_chat.py index e8c21a029b2d..c993eeac9e62 100644 --- a/backend/routes/chat/brainful_chat.py +++ b/backend/routes/chat/brainful_chat.py @@ -1,3 +1,4 @@ +from llm.qa_base import QABaseBrainPicking from routes.authorizations.brain_authorization import validate_brain_authorization from routes.authorizations.types import RoleEnum from routes.chat.interface import ChatInterface @@ -11,3 +12,25 @@ def validate_authorization(self, user_id, brain_id): user_id=user_id, required_roles=[RoleEnum.Viewer, RoleEnum.Editor, RoleEnum.Owner], ) + + def get_answer_generator( + self, + brain_id, + chat_id, + model, + max_tokens, + temperature, + user_openai_api_key, + streaming, + prompt_id, + ): + return QABaseBrainPicking( + chat_id=chat_id, + model=model, + max_tokens=max_tokens, + temperature=temperature, + brain_id=brain_id, + user_openai_api_key=user_openai_api_key, + streaming=streaming, + prompt_id=prompt_id, + ) diff --git a/backend/routes/chat/brainless_chat.py b/backend/routes/chat/brainless_chat.py index e93dae7aca29..c09d583d8b25 100644 --- a/backend/routes/chat/brainless_chat.py +++ b/backend/routes/chat/brainless_chat.py @@ -1,6 +1,28 @@ +from llm.qa_headless import HeadlessQA from routes.chat.interface import ChatInterface class BrainlessChat(ChatInterface): def validate_authorization(self, user_id, brain_id): pass + + def get_answer_generator( + self, + brain_id, + chat_id, + model, + max_tokens, + temperature, + user_openai_api_key, + streaming, + prompt_id, + ): + return HeadlessQA( + chat_id=chat_id, + model=model, + max_tokens=max_tokens, + temperature=temperature, + user_openai_api_key=user_openai_api_key, + streaming=streaming, + prompt_id=prompt_id, + ) diff --git a/backend/routes/chat/interface.py b/backend/routes/chat/interface.py index 49608ff2056d..0244ddfac88d 100644 --- a/backend/routes/chat/interface.py +++ b/backend/routes/chat/interface.py @@ -5,3 +5,17 @@ class ChatInterface(ABC): @abstractmethod def validate_authorization(self, user_id, required_roles): pass + + @abstractmethod + def get_answer_generator( + self, + brain_id, + chat_id, + model, + max_tokens, + temperature, + user_openai_api_key, + streaming, + prompt_id, + ): + pass diff --git a/backend/routes/chat_routes.py b/backend/routes/chat_routes.py index 3256d4fc8750..d1fb3bf79544 100644 --- a/backend/routes/chat_routes.py +++ b/backend/routes/chat_routes.py @@ -185,7 +185,7 @@ async def create_question_handler( is_model_ok = (brain_details or chat_question).model in userSettings.get("models", ["gpt-3.5-turbo"]) # type: ignore gpt_answer_generator: HeadlessQA | QABaseBrainPicking if brain_id: - gpt_answer_generator = QABaseBrainPicking( + gpt_answer_generator = chat_instance.get_answer_generator( chat_id=str(chat_id), model=chat_question.model if is_model_ok else "gpt-3.5-turbo", # type: ignore max_tokens=chat_question.max_tokens, @@ -195,7 +195,7 @@ async def create_question_handler( prompt_id=chat_question.prompt_id, ) else: - gpt_answer_generator = HeadlessQA( + gpt_answer_generator = chat_instance.get_answer_generator( model=chat_question.model if is_model_ok else "gpt-3.5-turbo", # type: ignore temperature=chat_question.temperature, max_tokens=chat_question.max_tokens, @@ -283,25 +283,26 @@ async def create_stream_question_handler( is_model_ok = (brain_details or chat_question).model in userSettings.get("models", ["gpt-3.5-turbo"]) # type: ignore if brain_id: - gpt_answer_generator = QABaseBrainPicking( + gpt_answer_generator = chat_instance.get_answer_generator( chat_id=str(chat_id), model=(brain_details or chat_question).model if is_model_ok else "gpt-3.5-turbo", # type: ignore max_tokens=(brain_details or chat_question).max_tokens, # type: ignore temperature=(brain_details or chat_question).temperature, # type: ignore - brain_id=str(brain_id), user_openai_api_key=current_user.openai_api_key, # pyright: ignore reportPrivateUsage=none streaming=True, prompt_id=chat_question.prompt_id, + brain_id=str(brain_id), ) else: - gpt_answer_generator = HeadlessQA( + gpt_answer_generator = chat_instance.get_answer_generator( + chat_id=str(chat_id), model=chat_question.model if is_model_ok else "gpt-3.5-turbo", # type: ignore - temperature=chat_question.temperature, max_tokens=chat_question.max_tokens, + temperature=chat_question.temperature, user_openai_api_key=current_user.openai_api_key, # pyright: ignore reportPrivateUsage=none - chat_id=str(chat_id), streaming=True, prompt_id=chat_question.prompt_id, + brain_id=str(brain_id), ) print("streaming") From 4af2c0d153e0372b4915f74afcdf6e960e6e3ef2 Mon Sep 17 00:00:00 2001 From: gozineb Date: Thu, 26 Oct 2023 21:38:16 +0200 Subject: [PATCH 6/8] =?UTF-8?q?=F0=9F=94=A5=20remove=20conditional=20to=20?= =?UTF-8?q?get=20gpt=5Fanswer=5Fgenerator?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- backend/routes/chat_routes.py | 69 ++++++++++------------------------- 1 file changed, 19 insertions(+), 50 deletions(-) diff --git a/backend/routes/chat_routes.py b/backend/routes/chat_routes.py index d1fb3bf79544..2dff9fc979e7 100644 --- a/backend/routes/chat_routes.py +++ b/backend/routes/chat_routes.py @@ -34,8 +34,6 @@ ) from repository.notification.remove_chat_notifications import remove_chat_notifications from repository.user_identity import get_user_identity -from routes.authorizations.brain_authorization import validate_brain_authorization -from routes.authorizations.types import RoleEnum from routes.chat.factory import get_chat_strategy from routes.chat.utils import ( NullableUUID, @@ -184,25 +182,15 @@ async def create_question_handler( check_user_requests_limit(current_user) is_model_ok = (brain_details or chat_question).model in userSettings.get("models", ["gpt-3.5-turbo"]) # type: ignore gpt_answer_generator: HeadlessQA | QABaseBrainPicking - if brain_id: - gpt_answer_generator = chat_instance.get_answer_generator( - chat_id=str(chat_id), - model=chat_question.model if is_model_ok else "gpt-3.5-turbo", # type: ignore - max_tokens=chat_question.max_tokens, - temperature=chat_question.temperature, - brain_id=str(brain_id), - user_openai_api_key=current_user.openai_api_key, # pyright: ignore reportPrivateUsage=none - prompt_id=chat_question.prompt_id, - ) - else: - gpt_answer_generator = chat_instance.get_answer_generator( - model=chat_question.model if is_model_ok else "gpt-3.5-turbo", # type: ignore - temperature=chat_question.temperature, - max_tokens=chat_question.max_tokens, - user_openai_api_key=current_user.openai_api_key, - chat_id=str(chat_id), - prompt_id=chat_question.prompt_id, - ) + gpt_answer_generator = chat_instance.get_answer_generator( + chat_id=str(chat_id), + model=chat_question.model if is_model_ok else "gpt-3.5-turbo", # type: ignore + max_tokens=chat_question.max_tokens, + temperature=chat_question.temperature, + brain_id=str(brain_id), + user_openai_api_key=current_user.openai_api_key, # pyright: ignore reportPrivateUsage=none + prompt_id=chat_question.prompt_id, + ) chat_answer = gpt_answer_generator.generate_answer(chat_id, chat_question) @@ -230,13 +218,6 @@ async def create_stream_question_handler( | None = Query(..., description="The ID of the brain"), current_user: UserIdentity = Depends(get_current_user), ) -> StreamingResponse: - if brain_id: - validate_brain_authorization( - brain_id=brain_id, - user_id=current_user.id, - required_roles=[RoleEnum.Viewer, RoleEnum.Editor, RoleEnum.Owner], - ) - chat_instance = get_chat_strategy(brain_id) chat_instance.validate_authorization(user_id=current_user.id, brain_id=brain_id) @@ -282,28 +263,16 @@ async def create_stream_question_handler( print(userSettings.get("models", ["gpt-3.5-turbo"])) # type: ignore is_model_ok = (brain_details or chat_question).model in userSettings.get("models", ["gpt-3.5-turbo"]) # type: ignore - if brain_id: - gpt_answer_generator = chat_instance.get_answer_generator( - chat_id=str(chat_id), - model=(brain_details or chat_question).model if is_model_ok else "gpt-3.5-turbo", # type: ignore - max_tokens=(brain_details or chat_question).max_tokens, # type: ignore - temperature=(brain_details or chat_question).temperature, # type: ignore - user_openai_api_key=current_user.openai_api_key, # pyright: ignore reportPrivateUsage=none - streaming=True, - prompt_id=chat_question.prompt_id, - brain_id=str(brain_id), - ) - else: - gpt_answer_generator = chat_instance.get_answer_generator( - chat_id=str(chat_id), - model=chat_question.model if is_model_ok else "gpt-3.5-turbo", # type: ignore - max_tokens=chat_question.max_tokens, - temperature=chat_question.temperature, - user_openai_api_key=current_user.openai_api_key, # pyright: ignore reportPrivateUsage=none - streaming=True, - prompt_id=chat_question.prompt_id, - brain_id=str(brain_id), - ) + gpt_answer_generator = chat_instance.get_answer_generator( + chat_id=str(chat_id), + model=(brain_details or chat_question).model if is_model_ok else "gpt-3.5-turbo", # type: ignore + max_tokens=(brain_details or chat_question).max_tokens, # type: ignore + temperature=(brain_details or chat_question).temperature, # type: ignore + user_openai_api_key=current_user.openai_api_key, # pyright: ignore reportPrivateUsage=none + streaming=True, + prompt_id=chat_question.prompt_id, + brain_id=str(brain_id), + ) print("streaming") return StreamingResponse( From 27dd657e6403f44fb8261c8bf39ed7c4ec9982fd Mon Sep 17 00:00:00 2001 From: gozineb Date: Thu, 26 Oct 2023 21:40:55 +0200 Subject: [PATCH 7/8] =?UTF-8?q?=F0=9F=94=A5=20remove=20logs?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- backend/routes/chat_routes.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/backend/routes/chat_routes.py b/backend/routes/chat_routes.py index 2dff9fc979e7..0e4ce633cf1b 100644 --- a/backend/routes/chat_routes.py +++ b/backend/routes/chat_routes.py @@ -260,7 +260,6 @@ async def create_stream_question_handler( gpt_answer_generator: HeadlessQA | QABaseBrainPicking # TODO check if model is in the list of models available for the user - print(userSettings.get("models", ["gpt-3.5-turbo"])) # type: ignore is_model_ok = (brain_details or chat_question).model in userSettings.get("models", ["gpt-3.5-turbo"]) # type: ignore gpt_answer_generator = chat_instance.get_answer_generator( @@ -274,7 +273,6 @@ async def create_stream_question_handler( brain_id=str(brain_id), ) - print("streaming") return StreamingResponse( gpt_answer_generator.generate_stream(chat_id, chat_question), media_type="text/event-stream", From 794cbfc94d6a4e96e7c5c99ebfec9bd56f18b2f0 Mon Sep 17 00:00:00 2001 From: gozineb Date: Thu, 26 Oct 2023 21:53:34 +0200 Subject: [PATCH 8/8] =?UTF-8?q?=E2=99=BB=EF=B8=8F=20replace=20get=5Fopenai?= =?UTF-8?q?=5Fapi=5Fkey=20conditional=20with=20strategy?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- backend/routes/chat/brainful_chat.py | 7 +++++++ backend/routes/chat/brainless_chat.py | 8 ++++++++ backend/routes/chat/interface.py | 4 ++++ backend/routes/chat_routes.py | 26 ++++++-------------------- 4 files changed, 25 insertions(+), 20 deletions(-) diff --git a/backend/routes/chat/brainful_chat.py b/backend/routes/chat/brainful_chat.py index c993eeac9e62..cadf1bc178db 100644 --- a/backend/routes/chat/brainful_chat.py +++ b/backend/routes/chat/brainful_chat.py @@ -3,6 +3,8 @@ from routes.authorizations.types import RoleEnum from routes.chat.interface import ChatInterface +from repository.brain import get_brain_details + class BrainfulChat(ChatInterface): def validate_authorization(self, user_id, brain_id): @@ -13,6 +15,11 @@ def validate_authorization(self, user_id, brain_id): required_roles=[RoleEnum.Viewer, RoleEnum.Editor, RoleEnum.Owner], ) + def get_openai_api_key(self, brain_id, user_id): + brain_details = get_brain_details(brain_id) + if brain_details: + return brain_details.openai_api_key + def get_answer_generator( self, brain_id, diff --git a/backend/routes/chat/brainless_chat.py b/backend/routes/chat/brainless_chat.py index c09d583d8b25..3107f0a8685f 100644 --- a/backend/routes/chat/brainless_chat.py +++ b/backend/routes/chat/brainless_chat.py @@ -1,11 +1,19 @@ from llm.qa_headless import HeadlessQA from routes.chat.interface import ChatInterface +from repository.user_identity import get_user_identity + class BrainlessChat(ChatInterface): def validate_authorization(self, user_id, brain_id): pass + def get_openai_api_key(self, brain_id, user_id): + user_identity = get_user_identity(user_id) + + if user_identity is not None: + return user_identity.openai_api_key + def get_answer_generator( self, brain_id, diff --git a/backend/routes/chat/interface.py b/backend/routes/chat/interface.py index 0244ddfac88d..04ba6e93d626 100644 --- a/backend/routes/chat/interface.py +++ b/backend/routes/chat/interface.py @@ -6,6 +6,10 @@ class ChatInterface(ABC): def validate_authorization(self, user_id, required_roles): pass + @abstractmethod + def get_openai_api_key(self, brain_id, user_id): + pass + @abstractmethod def get_answer_generator( self, diff --git a/backend/routes/chat_routes.py b/backend/routes/chat_routes.py index 0e4ce633cf1b..a029308f9f0d 100644 --- a/backend/routes/chat_routes.py +++ b/backend/routes/chat_routes.py @@ -17,7 +17,6 @@ get_supabase_db, ) from models.databases.supabase.chats import QuestionAndAnswer -from repository.brain import get_brain_details from repository.chat import ( ChatUpdatableProperties, CreateChatProperties, @@ -33,7 +32,6 @@ get_chat_history_with_notifications, ) from repository.notification.remove_chat_notifications import remove_chat_notifications -from repository.user_identity import get_user_identity from routes.chat.factory import get_chat_strategy from routes.chat.utils import ( NullableUUID, @@ -154,17 +152,10 @@ async def create_question_handler( userSettings = userDailyUsage.get_user_settings() is_model_ok = (brain_details or chat_question).model in userSettings.get("models", ["gpt-3.5-turbo"]) # type: ignore - if not current_user.openai_api_key and brain_id: - brain_details = get_brain_details(brain_id) - if brain_details: - current_user.openai_api_key = brain_details.openai_api_key - if not current_user.openai_api_key: - user_identity = get_user_identity(current_user.id) - - if user_identity is not None: - current_user.openai_api_key = user_identity.openai_api_key - + current_user.openai_api_key = chat_instance.get_openai_api_key( + brain_id=brain_id, user_id=current_user.id + ) # Retrieve chat model (temperature, max_tokens, model) if ( not chat_question.model @@ -232,16 +223,11 @@ async def create_stream_question_handler( ) userSettings = userDailyUsage.get_user_settings() - if not current_user.openai_api_key and brain_id: - brain_details = get_brain_details(brain_id) - if brain_details: - current_user.openai_api_key = brain_details.openai_api_key if not current_user.openai_api_key: - user_identity = get_user_identity(current_user.id) - - if user_identity is not None: - current_user.openai_api_key = user_identity.openai_api_key + current_user.openai_api_key = chat_instance.get_openai_api_key( + brain_id=brain_id, user_id=current_user.id + ) # Retrieve chat model (temperature, max_tokens, model) if (