quivr/backend/models/settings.py
Zineb El Bachiri f48dab4a7d
refactor: to modules (#1754)
# Description

Please include a summary of the changes and the related issue. Please
also include relevant motivation and context.

## Checklist before requesting a review

Please delete options that are not relevant.

- [ ] My code follows the style guidelines of this project
- [ ] I have performed a self-review of my code
- [ ] I have commented hard-to-understand areas
- [ ] I have ideally added tests that prove my fix is effective or that
my feature works
- [ ] New and existing unit tests pass locally with my changes
- [ ] Any dependent changes have been merged

## Screenshots (if appropriate):
2023-11-30 22:29:28 +01:00

63 lines
1.9 KiB
Python

from models.databases.supabase.supabase import SupabaseDB
from pydantic import BaseSettings
from supabase.client import Client, create_client
from vectorstore.supabase import SupabaseVectorStore
from langchain.embeddings.ollama import OllamaEmbeddings
from langchain.embeddings.openai import OpenAIEmbeddings
from logger import get_logger
logger = get_logger(__name__)
class BrainRateLimiting(BaseSettings):
max_brain_per_user: int = 5
class BrainSettings(BaseSettings):
openai_api_key: str
supabase_url: str
supabase_service_key: str
resend_api_key: str = "null"
resend_email_address: str = "brain@mail.quivr.app"
ollama_api_base_url: str = None
class ResendSettings(BaseSettings):
resend_api_key: str = "null"
def get_supabase_client() -> Client:
settings = BrainSettings() # pyright: ignore reportPrivateUsage=none
supabase_client: Client = create_client(
settings.supabase_url, settings.supabase_service_key
)
return supabase_client
def get_supabase_db() -> SupabaseDB:
supabase_client = get_supabase_client()
return SupabaseDB(supabase_client)
def get_embeddings():
settings = BrainSettings() # pyright: ignore reportPrivateUsage=none
if settings.ollama_api_base_url:
embeddings = OllamaEmbeddings(
base_url=settings.ollama_api_base_url,
) # pyright: ignore reportPrivateUsage=none
else:
embeddings = OpenAIEmbeddings() # pyright: ignore reportPrivateUsage=none
return embeddings
def get_documents_vector_store() -> SupabaseVectorStore:
settings = BrainSettings() # pyright: ignore reportPrivateUsage=none
embeddings = get_embeddings()
supabase_client: Client = create_client(
settings.supabase_url, settings.supabase_service_key
)
documents_vector_store = SupabaseVectorStore(
supabase_client, embeddings, table_name="vectors"
)
return documents_vector_store