Firestore Storage class for configs (#11)

* Implemented a new Firestore storage class for storing configs.
* Make env vars optional

---------

Co-authored-by: Guilherme Parpinelli <[email protected]>

.env

@@ -4,3 +4,4 @@
 # https://devcenter.heroku.com/articles/config-vars
 
 OPENAI_API_KEY=<your-openai-api-key>
+GOOGLE_CREDENTIALS=<your-google-credentials>

nalgonda/main.py

@@ -1,9 +1,13 @@
 import logging
 
+import openai
 from fastapi import FastAPI
 
 from nalgonda.constants import DATA_DIR
 from nalgonda.routers.v1 import v1_router
+from nalgonda.settings import settings
+
+openai.api_key = settings.openai_api_key
 
 # Ensure data directory exists
 DATA_DIR.mkdir(parents=True, exist_ok=True)

nalgonda/models/agency_config.py

@@ -2,8 +2,7 @@
 from pydantic import BaseModel, Field
 
 from nalgonda.models.agent_config import AgentConfig
-from nalgonda.persistence.agency_config_file_storage import AgencyConfigFileStorage
-
+from nalgonda.persistence.agency_config_firestore_storage import AgencyConfigFirestoreStorage
 
 class AgencyConfig(BaseModel):
     """Agency configuration model"""
@@ -22,14 +21,12 @@ def update_agent_ids_in_config(self, agents: list[Agent]) -> None:
 
     @classmethod
     def load(cls, agency_id: str) -> "AgencyConfig":
-        """Load agency config from the storage"""
-        with AgencyConfigFileStorage(agency_id) as config_file:
-            config = config_file.load()
+        with AgencyConfigFirestoreStorage(agency_id) as config_document:
+            config = config_document.load()
 
         config["agency_id"] = agency_id
         return cls.model_validate(config)
 
     def save(self) -> None:
-        """Save agency config to the storage"""
-        with AgencyConfigFileStorage(self.agency_id) as config_file:
-            config_file.save(self.model_dump())
+        with AgencyConfigFirestoreStorage(self.agency_id) as config_document:
+            config_document.save(self.model_dump())

nalgonda/persistence/agency_config_firestore_storage.py

@@ -0,0 +1,35 @@
+import json
+from typing import Any
+
+import firebase_admin
+from firebase_admin import credentials, firestore
+
+from nalgonda.persistence.agency_config_storage_interface import AgencyConfigStorageInterface
+from nalgonda.settings import settings
+
+# Initialize FireStore
+if settings.google_credentials:
+    cred_json = json.loads(settings.google_credentials)
+    cred = credentials.Certificate(cred_json)
+    firebase_admin.initialize_app(cred)
+
+
+class AgencyConfigFirestoreStorage(AgencyConfigStorageInterface):
+    def __init__(self, agency_id: str):
+        self.db = firestore.client()
+        self.agency_id = agency_id
+        self.collection_name = "agency_configs"
+        self.document = self.db.collection(self.collection_name).document(agency_id)
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        # No special action needed on exiting the context.
+        pass
+
+    def load(self):
+        return self.document.get().to_dict()
+
+    def save(self, data: dict[str, Any]):
+        self.document.set(data)

nalgonda/routers/v1/api/agency.py

@@ -1,7 +1,8 @@
 import logging
+from http import HTTPStatus
 
 from agency_swarm import Agency
-from fastapi import APIRouter
+from fastapi import APIRouter, HTTPException
 
 from nalgonda.agency_manager import AgencyManager
 from nalgonda.models.request_models import AgencyMessagePostRequest
@@ -35,7 +36,7 @@ async def post_agency_message(request: AgencyMessagePostRequest) -> dict:
 
     agency = await agency_manager.get_agency(agency_id, thread_id)
     if not agency:
-        return {"error": "Agency not found"}
+        raise HTTPException(status_code=HTTPStatus.NOT_FOUND, detail="Agency not found")
 
     try:
         response = await process_message(user_message, agency)

nalgonda/settings.py

@@ -5,10 +5,11 @@
 
 
 class Settings(BaseSettings):
-    # openai_api_key: str = Field(validation_alias="OPENAI_API_KEY")
+    openai_api_key: str | None = Field(default=None, validation_alias="OPENAI_API_KEY")
     gpt_model: str = Field(default=LATEST_GPT_MODEL, validation_alias="GPT_MODEL")
+    google_credentials: str | None = Field(default=None, validation_alias="GOOGLE_CREDENTIALS")
 
-    model_config = SettingsConfigDict(env_file=".env", env_prefix="AINHAND_", case_sensitive=True)
+    model_config = SettingsConfigDict(env_file=".env", case_sensitive=True)
 
 
 settings = Settings()

pyproject.toml

@@ -12,6 +12,7 @@ readme = "README.md"
 python = ">=3.11.0a1,<4.0.0"
 duckduckgo-search = "^3.9.11"
 fastapi = "^0.104.1"
+firebase-admin = "^6.3.0"
 functions = ">=0.7.0"
 gunicorn = "^21.2.0"
 instructor = "^0.3.4"