58 lines
1.7 KiB
Python
58 lines
1.7 KiB
Python
from pydantic import BaseModel
|
|
from pydantic_settings import BaseSettings, SettingsConfigDict
|
|
|
|
|
|
class IndexSettings(BaseModel):
|
|
name: str | None = None
|
|
endpoint: str | None = None
|
|
deployment: str | None = None
|
|
|
|
@property
|
|
def require_endpoint(self) -> str:
|
|
if not self.endpoint:
|
|
raise ValueError("INDEX__ENDPOINT environment variable must be set")
|
|
return self.endpoint
|
|
|
|
@property
|
|
def require_deployment(self) -> str:
|
|
if not self.deployment:
|
|
raise ValueError("INDEX__DEPLOYMENT environment variable must be set")
|
|
return self.deployment
|
|
|
|
|
|
class AgentSettings(BaseModel):
|
|
name: str = "default"
|
|
language_model: str = "gemini-2.0-flash"
|
|
embedding_model: str = "text-embedding-004"
|
|
|
|
|
|
class BigQuerySettings(BaseModel):
|
|
project_id: str | None = None
|
|
dataset_id: str | None = None
|
|
synth_gen_table: str = "synthetic_questions"
|
|
search_eval_table: str = "search_eval_results"
|
|
keypoint_eval_table: str = "keypoint_eval_results"
|
|
|
|
@property
|
|
def require_dataset_id(self) -> str:
|
|
if not self.dataset_id:
|
|
raise ValueError("BIGQUERY__DATASET_ID environment variable must be set")
|
|
return self.dataset_id
|
|
|
|
|
|
class Settings(BaseSettings):
|
|
model_config = SettingsConfigDict(env_nested_delimiter="__")
|
|
|
|
project_id: str
|
|
location: str = "us-central1"
|
|
bucket: str | None = None
|
|
index: IndexSettings = IndexSettings()
|
|
agent: AgentSettings = AgentSettings()
|
|
bigquery: BigQuerySettings = BigQuerySettings()
|
|
|
|
@property
|
|
def require_bucket(self) -> str:
|
|
if not self.bucket:
|
|
raise ValueError("BUCKET environment variable must be set")
|
|
return self.bucket
|