Refactor _chromadb.py to improve code structure and readability

main
lasseedfast 1 year ago
parent afd61fac25
commit d1261d2d90
  1. 21
      _chromadb.py

@ -7,19 +7,22 @@ from chromadb.config import Settings
from dotenv import load_dotenv from dotenv import load_dotenv
from chromadb.utils import embedding_functions from chromadb.utils import embedding_functions
load_dotenv('.chroma_env') load_dotenv('.chroma_env')
class ChromaDB: class ChromaDB:
def __init__(self): def __init__(self, local_deployment: bool = False, db='sci_articles'):
self.db = chromadb.HttpClient( if local_deployment:
host="https://lasseedfast.se/chroma_ev_cars", self.db = chromadb.PersistentClient(f'chroma_{db}')
settings=Settings( else:
chroma_client_auth_provider="chromadb.auth.token_authn.TokenAuthClientProvider", self.db = chromadb.HttpClient(
chroma_client_auth_credentials=os.getenv("CHROMA_CLIENT_AUTH_CREDENTIALS"), host=os.getenv('CHROMA_HOST'),
chroma_auth_token_transport_header=os.getenv("CHROMA_AUTH_TOKEN_TRANSPORT_HEADER") settings=Settings(
chroma_client_auth_provider="chromadb.auth.token_authn.TokenAuthClientProvider",
chroma_client_auth_credentials=os.getenv("CHROMA_CLIENT_AUTH_CREDENTIALS"),
chroma_auth_token_transport_header=os.getenv("CHROMA_AUTH_TOKEN_TRANSPORT_HEADER")
)
) )
)
max_characters = 2200 max_characters = 2200
self.ts = MarkdownSplitter(max_characters) self.ts = MarkdownSplitter(max_characters)

Loading…
Cancel
Save