diff --git a/backend/models/databases/supabase/brains.py b/backend/models/databases/supabase/brains.py index 1320a924e..31bfe5da0 100644 --- a/backend/models/databases/supabase/brains.py +++ b/backend/models/databases/supabase/brains.py @@ -43,6 +43,10 @@ class BrainUpdatableProperties(BaseModel): return brain_dict +class BrainQuestionRequest(BaseModel): + question: str + + class Brain(Repository): def __init__(self, supabase_client): self.db = supabase_client diff --git a/backend/repository/brain/__init__.py b/backend/repository/brain/__init__.py index b7d5dee86..a689beba5 100644 --- a/backend/repository/brain/__init__.py +++ b/backend/repository/brain/__init__.py @@ -10,3 +10,4 @@ from .update_user_rights import update_brain_user_rights from .get_default_user_brain import get_user_default_brain from .set_as_default_brain_for_user import set_as_default_brain_for_user from .get_default_user_brain_or_create_new import get_default_user_brain_or_create_new +from .get_question_context_from_brain import get_question_context_from_brain diff --git a/backend/repository/brain/get_question_context_from_brain.py b/backend/repository/brain/get_question_context_from_brain.py new file mode 100644 index 000000000..ef69d3593 --- /dev/null +++ b/backend/repository/brain/get_question_context_from_brain.py @@ -0,0 +1,20 @@ +from uuid import UUID + +from models.settings import get_embeddings, get_supabase_client +from vectorstore.supabase import CustomSupabaseVectorStore + + +def get_question_context_from_brain(brain_id: UUID, question: str) -> str: + supabase_client = get_supabase_client() + embeddings = get_embeddings() + + vector_store = CustomSupabaseVectorStore( + supabase_client, + embeddings, + table_name="vectors", + brain_id=brain_id, + ) + documents = vector_store.similarity_search(question) + + # aggregate all the documents into one string + return "\n".join([doc.page_content for doc in documents]) diff --git a/backend/routes/brain_routes.py b/backend/routes/brain_routes.py index 4a3d67225..506d1f9db 100644 --- a/backend/routes/brain_routes.py +++ b/backend/routes/brain_routes.py @@ -4,20 +4,16 @@ from auth import AuthBearer, get_current_user from fastapi import APIRouter, Depends, HTTPException from logger import get_logger from models import BrainRateLimiting, UserIdentity -from models.databases.supabase.brains import ( - BrainUpdatableProperties, - CreateBrainProperties, -) -from repository.brain import ( - create_brain, - create_brain_user, - get_brain_details, - get_default_user_brain_or_create_new, - get_user_brains, - get_user_default_brain, - set_as_default_brain_for_user, - update_brain_by_id, -) +from models.databases.supabase.brains import (BrainQuestionRequest, + BrainUpdatableProperties, + CreateBrainProperties) +from repository.brain import (create_brain, create_brain_user, + get_brain_details, + get_default_user_brain_or_create_new, + get_question_context_from_brain, get_user_brains, + get_user_default_brain, + set_as_default_brain_for_user, + update_brain_by_id) from repository.prompt import delete_prompt_by_id, get_prompt_by_id from routes.authorizations.brain_authorization import has_brain_authorization from routes.authorizations.types import RoleEnum @@ -207,3 +203,26 @@ async def set_as_default_brain_endpoint( set_as_default_brain_for_user(user.id, brain_id) return {"message": f"Brain {brain_id} has been set as default brain."} + + +@brain_router.post( + "/brains/{brain_id}/question_context", + dependencies=[ + Depends( + AuthBearer(), + ), + Depends(has_brain_authorization()), + ], + tags=["Brain"], +) +async def get_question_context_from_brain_endpoint( + brain_id: UUID, + request: BrainQuestionRequest, +): + """ + Get question context from brain + """ + + context = get_question_context_from_brain(brain_id, request.question) + + return {"context": context}