From 60e38a544a7d38685c2861c660f3104043e8d256 Mon Sep 17 00:00:00 2001
From: Nikita Bobrovskiy <39348559+bonk1t@users.noreply.github.com>
Date: Sat, 30 Dec 2023 19:53:41 +0200
Subject: [PATCH 01/10] Implement new endpoints

---
 .../routers/v1/api/test_agency_routes.py      | 47 +++++++++++++++++++
 nalgonda/models/agency_config.py              | 26 ++++++----
 nalgonda/routers/v1/api/agency.py             | 20 ++++++++
 tests/routers/v1/api/test_agency_routes.py    | 31 ++++++++++++
 4 files changed, 114 insertions(+), 10 deletions(-)
 create mode 100644 nalgonda/data/test_agency_id/tests/routers/v1/api/test_agency_routes.py
 create mode 100644 tests/routers/v1/api/test_agency_routes.py

diff --git a/nalgonda/data/test_agency_id/tests/routers/v1/api/test_agency_routes.py b/nalgonda/data/test_agency_id/tests/routers/v1/api/test_agency_routes.py
new file mode 100644
index 00000000..ecb672e1
--- /dev/null
+++ b/nalgonda/data/test_agency_id/tests/routers/v1/api/test_agency_routes.py
@@ -0,0 +1,47 @@
+from unittest.mock import patch
+
+from fastapi.testclient import TestClient
+
+from nalgonda.main import app
+from nalgonda.models.agency_config import AgencyConfig
+from nalgonda.persistence.agency_config_firestore_storage import AgencyConfigFirestoreStorage
+
+
+def mocked_load(self, agency_id: str):  # noqa: ARG001
+    if agency_id == "test_agency":
+        return AgencyConfig(agency_id="test_agency", agency_manifesto="Test Manifesto", agents=[], agency_chart=[])
+    return None
+
+
+def mocked_save(self, data: dict):  # noqa: ARG001
+    assert data == {"agency_manifesto": "Updated Manifesto"}
+
+
+class TestAgencyRoutes:
+    client = TestClient(app)
+
+    @patch.object(AgencyConfigFirestoreStorage, "load", mocked_load)
+    @patch.object(AgencyConfigFirestoreStorage, "save", mocked_save)
+    def test_get_agency_config(self):
+        response = self.client.get("/agency/config?agency_id=test_agency")
+        assert response.status_code == 200
+        assert response.json() == {
+            "agency_id": "test_agency",
+            "agency_manifesto": "Test Manifesto",
+            "agents": [],
+            "agency_chart": [],
+        }
+
+    @patch.object(AgencyConfigFirestoreStorage, "load", mocked_load)
+    @patch.object(AgencyConfigFirestoreStorage, "save", mocked_save)
+    def test_update_agency_config(self):
+        new_data = {"agency_manifesto": "Updated Manifesto"}
+        response = self.client.post("/agency/config?agency_id=test_agency", json=new_data)
+        assert response.status_code == 201
+        assert response.json() == {"message": "Agency configuration updated successfully"}
+
+    @patch.object(AgencyConfigFirestoreStorage, "load", lambda _: None)
+    def test_agency_config_not_found(self):
+        response = self.client.get("/agency/config?agency_id=non_existent_agency")
+        assert response.status_code == 404
+        assert response.json() == {"detail": "Agency configuration not found"}
diff --git a/nalgonda/models/agency_config.py b/nalgonda/models/agency_config.py
index a13972f8..3d405bb2 100644
--- a/nalgonda/models/agency_config.py
+++ b/nalgonda/models/agency_config.py
@@ -1,3 +1,5 @@
+from typing import Any
+
 from agency_swarm import Agent
 from pydantic import BaseModel, Field
 
@@ -13,21 +15,25 @@ class AgencyConfig(BaseModel):
     agents: list[AgentConfig] = Field(...)
     agency_chart: list[str | list[str]] = Field(...)  # contains agent roles
 
-    def update_agent_ids_in_config(self, agents: list[Agent]) -> None:
-        """Update agent ids in config with the ids of the agents in the swarm"""
-        for agent in agents:
-            for agent_config in self.agents:
-                if agent.name == f"{agent_config.role}_{self.agency_id}":
-                    agent_config.id = agent.id
-
     @classmethod
     def load(cls, agency_id: str) -> "AgencyConfig":
         with AgencyConfigFirestoreStorage(agency_id) as config_document:
-            config = config_document.load()
+            config_data = config_document.load()
+        config_data["agency_id"] = agency_id
+        return cls.model_validate(config_data)
 
-        config["agency_id"] = agency_id
-        return cls.model_validate(config)
+    def update(self, update_data: dict[str, Any]) -> None:
+        for key, value in update_data.items():
+            if hasattr(self, key):
+                setattr(self, key, value)
 
     def save(self) -> None:
         with AgencyConfigFirestoreStorage(self.agency_id) as config_document:
             config_document.save(self.model_dump())
+
+    def update_agent_ids_in_config(self, agents: list[Agent]) -> None:
+        """Update agent ids in config with the ids of the agents in the swarm"""
+        for agent in agents:
+            for agent_config in self.agents:
+                if agent.name == f"{agent_config.role}_{self.agency_id}":
+                    agent_config.id = agent.id
diff --git a/nalgonda/routers/v1/api/agency.py b/nalgonda/routers/v1/api/agency.py
index ea2684c4..55de5658 100644
--- a/nalgonda/routers/v1/api/agency.py
+++ b/nalgonda/routers/v1/api/agency.py
@@ -4,10 +4,12 @@
 
 from agency_swarm import Agency
 from fastapi import APIRouter, Depends, HTTPException
+from starlette.status import HTTP_201_CREATED, HTTP_404_NOT_FOUND
 
 from nalgonda.dependencies.agency_manager import AgencyManager, get_agency_manager
 from nalgonda.dependencies.auth import get_current_active_user
 from nalgonda.dependencies.thread_manager import ThreadManager, get_thread_manager
+from nalgonda.models.agency_config import AgencyConfig
 from nalgonda.models.auth import User
 from nalgonda.models.request_models import AgencyMessagePostRequest, AgencyThreadPostRequest
 
@@ -52,6 +54,24 @@ async def create_agency_thread(
     return {"thread_id": thread_id}
 
 
+@agency_router.get("/agency/config")
+def get_agency_config(agency_id: str):
+    agency_config = AgencyConfig.load(agency_id)
+    if not agency_config:
+        raise HTTPException(status_code=HTTP_404_NOT_FOUND, detail="Agency configuration not found")
+    return agency_config
+
+
+@agency_router.post("/agency/config", status_code=HTTP_201_CREATED)
+def update_agency_config(agency_id: str, updated_data: dict):
+    agency_config = AgencyConfig.load(agency_id)
+    if not agency_config:
+        raise HTTPException(status_code=HTTP_404_NOT_FOUND, detail="Agency configuration not found")
+    agency_config.update(updated_data)
+    agency_config.save()
+    return {"message": "Agency configuration updated successfully"}
+
+
 @agency_router.post("/agency/message")
 async def post_agency_message(
     request: AgencyMessagePostRequest, agency_manager: AgencyManager = Depends(get_agency_manager)
diff --git a/tests/routers/v1/api/test_agency_routes.py b/tests/routers/v1/api/test_agency_routes.py
new file mode 100644
index 00000000..7af851a3
--- /dev/null
+++ b/tests/routers/v1/api/test_agency_routes.py
@@ -0,0 +1,31 @@
+from unittest.mock import patch
+
+from fastapi.testclient import TestClient
+
+from nalgonda.main import app
+from nalgonda.models.agency_config import AgencyConfig
+
+
+def get_mocked_agency_config(agency_id: str):
+    # This is a stub for generating a mocked AgencyConfig instance
+    return AgencyConfig(agency_id=agency_id, agency_manifesto="Mocked Agency Manifesto")
+
+
+@patch("persistence.agency_config_firestore_storage.AgencyConfigFirestoreStorage.load")
+@patch("persistence.agency_config_firestore_storage.AgencyConfigFirestoreStorage.save")
+def test_get_agency_config(mock_load, mock_save):
+    # Configure the mock to return a mocked AgencyConfig upon load
+    mock_load.return_value = {"agency_id": "test_agency", "agency_manifesto": "Mocked Agency Manifesto"}
+    client = TestClient(app)
+
+    # Send a GET request to the /agency/config endpoint
+    response = client.get("/agency/config?agency_id=test_agency")
+    assert response.status_code == 200
+    assert response.json() == {"agency_id": "test_agency", "agency_manifesto": "Mocked Agency Manifesto"}
+
+    # Test updating agency_config through a POST request and ensure correct saving
+    new_data = {"agency_manifesto": "Updated Manifesto"}
+    response = client.post("/agency/config?agency_id=test_agency", json=new_data)
+    assert response.status_code == 201
+    assert response.json() == {"message": "Config updated successfully"}
+    mock_save.assert_called_with(new_data)

From e7945bc9eccf40de1b07911d7403fddb14b081c9 Mon Sep 17 00:00:00 2001
From: Nikita Bobrovskiy <39348559+bonk1t@users.noreply.github.com>
Date: Sat, 30 Dec 2023 20:18:04 +0200
Subject: [PATCH 02/10] Implement new endpoints

---
 .../routers/v1/api/test_agency_routes.py      | 47 ----------------
 tests/routers/v1/api/test_agency_routes.py    | 54 ++++++++++++-------
 2 files changed, 35 insertions(+), 66 deletions(-)
 delete mode 100644 nalgonda/data/test_agency_id/tests/routers/v1/api/test_agency_routes.py

diff --git a/nalgonda/data/test_agency_id/tests/routers/v1/api/test_agency_routes.py b/nalgonda/data/test_agency_id/tests/routers/v1/api/test_agency_routes.py
deleted file mode 100644
index ecb672e1..00000000
--- a/nalgonda/data/test_agency_id/tests/routers/v1/api/test_agency_routes.py
+++ /dev/null
@@ -1,47 +0,0 @@
-from unittest.mock import patch
-
-from fastapi.testclient import TestClient
-
-from nalgonda.main import app
-from nalgonda.models.agency_config import AgencyConfig
-from nalgonda.persistence.agency_config_firestore_storage import AgencyConfigFirestoreStorage
-
-
-def mocked_load(self, agency_id: str):  # noqa: ARG001
-    if agency_id == "test_agency":
-        return AgencyConfig(agency_id="test_agency", agency_manifesto="Test Manifesto", agents=[], agency_chart=[])
-    return None
-
-
-def mocked_save(self, data: dict):  # noqa: ARG001
-    assert data == {"agency_manifesto": "Updated Manifesto"}
-
-
-class TestAgencyRoutes:
-    client = TestClient(app)
-
-    @patch.object(AgencyConfigFirestoreStorage, "load", mocked_load)
-    @patch.object(AgencyConfigFirestoreStorage, "save", mocked_save)
-    def test_get_agency_config(self):
-        response = self.client.get("/agency/config?agency_id=test_agency")
-        assert response.status_code == 200
-        assert response.json() == {
-            "agency_id": "test_agency",
-            "agency_manifesto": "Test Manifesto",
-            "agents": [],
-            "agency_chart": [],
-        }
-
-    @patch.object(AgencyConfigFirestoreStorage, "load", mocked_load)
-    @patch.object(AgencyConfigFirestoreStorage, "save", mocked_save)
-    def test_update_agency_config(self):
-        new_data = {"agency_manifesto": "Updated Manifesto"}
-        response = self.client.post("/agency/config?agency_id=test_agency", json=new_data)
-        assert response.status_code == 201
-        assert response.json() == {"message": "Agency configuration updated successfully"}
-
-    @patch.object(AgencyConfigFirestoreStorage, "load", lambda _: None)
-    def test_agency_config_not_found(self):
-        response = self.client.get("/agency/config?agency_id=non_existent_agency")
-        assert response.status_code == 404
-        assert response.json() == {"detail": "Agency configuration not found"}
diff --git a/tests/routers/v1/api/test_agency_routes.py b/tests/routers/v1/api/test_agency_routes.py
index 7af851a3..ecb672e1 100644
--- a/tests/routers/v1/api/test_agency_routes.py
+++ b/tests/routers/v1/api/test_agency_routes.py
@@ -4,28 +4,44 @@
 
 from nalgonda.main import app
 from nalgonda.models.agency_config import AgencyConfig
+from nalgonda.persistence.agency_config_firestore_storage import AgencyConfigFirestoreStorage
 
 
-def get_mocked_agency_config(agency_id: str):
-    # This is a stub for generating a mocked AgencyConfig instance
-    return AgencyConfig(agency_id=agency_id, agency_manifesto="Mocked Agency Manifesto")
+def mocked_load(self, agency_id: str):  # noqa: ARG001
+    if agency_id == "test_agency":
+        return AgencyConfig(agency_id="test_agency", agency_manifesto="Test Manifesto", agents=[], agency_chart=[])
+    return None
 
 
-@patch("persistence.agency_config_firestore_storage.AgencyConfigFirestoreStorage.load")
-@patch("persistence.agency_config_firestore_storage.AgencyConfigFirestoreStorage.save")
-def test_get_agency_config(mock_load, mock_save):
-    # Configure the mock to return a mocked AgencyConfig upon load
-    mock_load.return_value = {"agency_id": "test_agency", "agency_manifesto": "Mocked Agency Manifesto"}
+def mocked_save(self, data: dict):  # noqa: ARG001
+    assert data == {"agency_manifesto": "Updated Manifesto"}
+
+
+class TestAgencyRoutes:
     client = TestClient(app)
 
-    # Send a GET request to the /agency/config endpoint
-    response = client.get("/agency/config?agency_id=test_agency")
-    assert response.status_code == 200
-    assert response.json() == {"agency_id": "test_agency", "agency_manifesto": "Mocked Agency Manifesto"}
-
-    # Test updating agency_config through a POST request and ensure correct saving
-    new_data = {"agency_manifesto": "Updated Manifesto"}
-    response = client.post("/agency/config?agency_id=test_agency", json=new_data)
-    assert response.status_code == 201
-    assert response.json() == {"message": "Config updated successfully"}
-    mock_save.assert_called_with(new_data)
+    @patch.object(AgencyConfigFirestoreStorage, "load", mocked_load)
+    @patch.object(AgencyConfigFirestoreStorage, "save", mocked_save)
+    def test_get_agency_config(self):
+        response = self.client.get("/agency/config?agency_id=test_agency")
+        assert response.status_code == 200
+        assert response.json() == {
+            "agency_id": "test_agency",
+            "agency_manifesto": "Test Manifesto",
+            "agents": [],
+            "agency_chart": [],
+        }
+
+    @patch.object(AgencyConfigFirestoreStorage, "load", mocked_load)
+    @patch.object(AgencyConfigFirestoreStorage, "save", mocked_save)
+    def test_update_agency_config(self):
+        new_data = {"agency_manifesto": "Updated Manifesto"}
+        response = self.client.post("/agency/config?agency_id=test_agency", json=new_data)
+        assert response.status_code == 201
+        assert response.json() == {"message": "Agency configuration updated successfully"}
+
+    @patch.object(AgencyConfigFirestoreStorage, "load", lambda _: None)
+    def test_agency_config_not_found(self):
+        response = self.client.get("/agency/config?agency_id=non_existent_agency")
+        assert response.status_code == 404
+        assert response.json() == {"detail": "Agency configuration not found"}

From 4c925b8f4c22efd1b8659081f7289edf18c77f2b Mon Sep 17 00:00:00 2001
From: Nikita Bobrovskiy <39348559+bonk1t@users.noreply.github.com>
Date: Sat, 30 Dec 2023 20:42:08 +0200
Subject: [PATCH 03/10] Implement new endpoints

---
 nalgonda/dependencies/agency_manager.py       |  2 +-
 nalgonda/models/agency_config.py              | 30 +++++++++++++++++--
 .../agency_config_firestore_storage.py        | 13 +-------
 tests/routers/v1/api/test_agency_routes.py    | 21 ++++++++-----
 4 files changed, 42 insertions(+), 24 deletions(-)

diff --git a/nalgonda/dependencies/agency_manager.py b/nalgonda/dependencies/agency_manager.py
index 83b201a5..12b569d3 100644
--- a/nalgonda/dependencies/agency_manager.py
+++ b/nalgonda/dependencies/agency_manager.py
@@ -61,7 +61,7 @@ def load_agency_from_config(agency_id: str) -> Agency:
         """
 
         start = time.time()
-        config = AgencyConfig.load(agency_id)
+        config = AgencyConfig.load_or_create(agency_id)
 
         agents = {
             agent_conf.role: Agent(
diff --git a/nalgonda/models/agency_config.py b/nalgonda/models/agency_config.py
index 3d405bb2..e5585b70 100644
--- a/nalgonda/models/agency_config.py
+++ b/nalgonda/models/agency_config.py
@@ -1,8 +1,10 @@
-from typing import Any
+import json
+from typing import Any, Optional
 
 from agency_swarm import Agent
 from pydantic import BaseModel, Field
 
+from nalgonda.constants import DEFAULT_CONFIG_FILE
 from nalgonda.models.agent_config import AgentConfig
 from nalgonda.persistence.agency_config_firestore_storage import AgencyConfigFirestoreStorage
 
@@ -16,11 +18,26 @@ class AgencyConfig(BaseModel):
     agency_chart: list[str | list[str]] = Field(...)  # contains agent roles
 
     @classmethod
-    def load(cls, agency_id: str) -> "AgencyConfig":
+    def load(cls, agency_id: str) -> Optional["AgencyConfig"]:
         with AgencyConfigFirestoreStorage(agency_id) as config_document:
             config_data = config_document.load()
+        return cls.model_validate(config_data) if config_data else None
+
+    @classmethod
+    def load_or_create(cls, agency_id: str) -> "AgencyConfig":
+        with AgencyConfigFirestoreStorage(agency_id) as config_document:
+            config_data = config_document.load()
+
+        if not config_data:
+            config_data = cls._create_default_config()
+
         config_data["agency_id"] = agency_id
-        return cls.model_validate(config_data)
+        model = cls.model_validate(config_data)
+
+        with AgencyConfigFirestoreStorage(agency_id) as config_document:
+            config_document.save(config_data)
+
+        return model
 
     def update(self, update_data: dict[str, Any]) -> None:
         for key, value in update_data.items():
@@ -37,3 +54,10 @@ def update_agent_ids_in_config(self, agents: list[Agent]) -> None:
             for agent_config in self.agents:
                 if agent.name == f"{agent_config.role}_{self.agency_id}":
                     agent_config.id = agent.id
+
+    @classmethod
+    def _create_default_config(cls) -> dict[str, Any]:
+        """Creates a default config for the agency."""
+        with DEFAULT_CONFIG_FILE.open() as file:
+            config = json.load(file)
+        return config
diff --git a/nalgonda/persistence/agency_config_firestore_storage.py b/nalgonda/persistence/agency_config_firestore_storage.py
index c6cf14d6..ce384333 100644
--- a/nalgonda/persistence/agency_config_firestore_storage.py
+++ b/nalgonda/persistence/agency_config_firestore_storage.py
@@ -1,9 +1,7 @@
-import json
 from typing import Any
 
 from firebase_admin import firestore
 
-from nalgonda.constants import DEFAULT_CONFIG_FILE
 from nalgonda.persistence.agency_config_storage_interface import AgencyConfigStorageInterface
 
 
@@ -21,18 +19,9 @@ def __exit__(self, exc_type, exc_val, exc_tb):
         # No special action needed on exiting the context.
         pass
 
-    def load(self):
+    def load(self) -> dict[str, Any] | None:
         db_config = self.document.get().to_dict()
-        if not db_config:
-            db_config = self._create_default_config()
         return db_config
 
     def save(self, data: dict[str, Any]):
         self.document.set(data)
-
-    def _create_default_config(self) -> dict[str, Any]:
-        """Creates a default config for the agency and saves it to the database."""
-        with DEFAULT_CONFIG_FILE.open() as file:
-            config = json.load(file)
-        self.save(config)
-        return config
diff --git a/tests/routers/v1/api/test_agency_routes.py b/tests/routers/v1/api/test_agency_routes.py
index ecb672e1..a14d51c9 100644
--- a/tests/routers/v1/api/test_agency_routes.py
+++ b/tests/routers/v1/api/test_agency_routes.py
@@ -7,14 +7,19 @@
 from nalgonda.persistence.agency_config_firestore_storage import AgencyConfigFirestoreStorage
 
 
-def mocked_load(self, agency_id: str):  # noqa: ARG001
-    if agency_id == "test_agency":
-        return AgencyConfig(agency_id="test_agency", agency_manifesto="Test Manifesto", agents=[], agency_chart=[])
-    return None
+def mocked_load(self):  # noqa: ARG001
+    return AgencyConfig(
+        agency_id="test_agency", agency_manifesto="Test Manifesto", agents=[], agency_chart=[]
+    ).model_dump()
 
 
 def mocked_save(self, data: dict):  # noqa: ARG001
-    assert data == {"agency_manifesto": "Updated Manifesto"}
+    assert data == {
+        "agency_manifesto": "Updated Manifesto",
+        "agency_chart": [],
+        "agency_id": "test_agency",
+        "agents": [],
+    }
 
 
 class TestAgencyRoutes:
@@ -23,7 +28,7 @@ class TestAgencyRoutes:
     @patch.object(AgencyConfigFirestoreStorage, "load", mocked_load)
     @patch.object(AgencyConfigFirestoreStorage, "save", mocked_save)
     def test_get_agency_config(self):
-        response = self.client.get("/agency/config?agency_id=test_agency")
+        response = self.client.get("/v1/api/agency/config?agency_id=test_agency")
         assert response.status_code == 200
         assert response.json() == {
             "agency_id": "test_agency",
@@ -36,12 +41,12 @@ def test_get_agency_config(self):
     @patch.object(AgencyConfigFirestoreStorage, "save", mocked_save)
     def test_update_agency_config(self):
         new_data = {"agency_manifesto": "Updated Manifesto"}
-        response = self.client.post("/agency/config?agency_id=test_agency", json=new_data)
+        response = self.client.post("/v1/api/agency/config?agency_id=test_agency", json=new_data)
         assert response.status_code == 201
         assert response.json() == {"message": "Agency configuration updated successfully"}
 
     @patch.object(AgencyConfigFirestoreStorage, "load", lambda _: None)
     def test_agency_config_not_found(self):
-        response = self.client.get("/agency/config?agency_id=non_existent_agency")
+        response = self.client.get("/v1/api/agency/config?agency_id=non_existent_agency")
         assert response.status_code == 404
         assert response.json() == {"detail": "Agency configuration not found"}

From 3f2beb812b3acd2e6590f4e51bd1bfcdeea65f81 Mon Sep 17 00:00:00 2001
From: Nikita Bobrovskiy <39348559+bonk1t@users.noreply.github.com>
Date: Sat, 30 Dec 2023 21:14:32 +0200
Subject: [PATCH 04/10] Refactor

---
 nalgonda/dependencies/agency_manager.py |  1 -
 nalgonda/routers/v1/api/agency.py       | 18 +++++++++++++++---
 2 files changed, 15 insertions(+), 4 deletions(-)

diff --git a/nalgonda/dependencies/agency_manager.py b/nalgonda/dependencies/agency_manager.py
index 12b569d3..18b8187c 100644
--- a/nalgonda/dependencies/agency_manager.py
+++ b/nalgonda/dependencies/agency_manager.py
@@ -25,7 +25,6 @@ async def create_agency(self, agency_id: str | None = None) -> tuple[Agency, str
 
         # Note: Async-to-Sync Bridge
         agency = await asyncio.to_thread(self.load_agency_from_config, agency_id)
-        await self.cache_agency(agency, agency_id, None)
         return agency, agency_id
 
     async def get_agency(self, agency_id: str, thread_id: str | None) -> Agency | None:
diff --git a/nalgonda/routers/v1/api/agency.py b/nalgonda/routers/v1/api/agency.py
index 55de5658..b3d2c9ef 100644
--- a/nalgonda/routers/v1/api/agency.py
+++ b/nalgonda/routers/v1/api/agency.py
@@ -28,7 +28,9 @@ async def create_agency(
     # TODO: check if the current_user has permission to create an agency
     logger.info(f"Creating agency for user: {current_user.username}")
 
-    _, agency_id = await agency_manager.create_agency()
+    agency, agency_id = await agency_manager.create_agency()
+
+    await agency_manager.cache_agency(agency, agency_id, None)
     return {"agency_id": agency_id}
 
 
@@ -55,7 +57,7 @@ async def create_agency_thread(
 
 
 @agency_router.get("/agency/config")
-def get_agency_config(agency_id: str):
+async def get_agency_config(agency_id: str):
     agency_config = AgencyConfig.load(agency_id)
     if not agency_config:
         raise HTTPException(status_code=HTTP_404_NOT_FOUND, detail="Agency configuration not found")
@@ -63,12 +65,22 @@ def get_agency_config(agency_id: str):
 
 
 @agency_router.post("/agency/config", status_code=HTTP_201_CREATED)
-def update_agency_config(agency_id: str, updated_data: dict):
+async def update_agency_config(
+    agency_id: str,
+    updated_data: dict,
+    agency_manager: AgencyManager = Depends(get_agency_manager),
+):
     agency_config = AgencyConfig.load(agency_id)
     if not agency_config:
         raise HTTPException(status_code=HTTP_404_NOT_FOUND, detail="Agency configuration not found")
     agency_config.update(updated_data)
     agency_config.save()
+
+    agency = await agency_manager.get_agency(agency_id, None)
+    if not agency:
+        agency, _ = await agency_manager.create_agency(agency_id)
+
+    await agency_manager.cache_agency(agency, agency_id, None)
     return {"message": "Agency configuration updated successfully"}
 
 

From 928266aff60d0bb0a25b13bcb6eb9560b4cbe3c2 Mon Sep 17 00:00:00 2001
From: Nikita Bobrovskiy <39348559+bonk1t@users.noreply.github.com>
Date: Sat, 30 Dec 2023 21:42:27 +0200
Subject: [PATCH 05/10] Refactor

---
 nalgonda/dependencies/agency_manager.py | 33 ++++++++++++++++++++++---
 nalgonda/routers/v1/api/agency.py       | 15 +++--------
 2 files changed, 33 insertions(+), 15 deletions(-)

diff --git a/nalgonda/dependencies/agency_manager.py b/nalgonda/dependencies/agency_manager.py
index 18b8187c..2a2e20c0 100644
--- a/nalgonda/dependencies/agency_manager.py
+++ b/nalgonda/dependencies/agency_manager.py
@@ -24,16 +24,41 @@ async def create_agency(self, agency_id: str | None = None) -> tuple[Agency, str
         agency_id = agency_id or uuid.uuid4().hex
 
         # Note: Async-to-Sync Bridge
-        agency = await asyncio.to_thread(self.load_agency_from_config, agency_id)
+        agency = await asyncio.to_thread(self.load_agency_from_config, agency_id, config=None)
+        await self.cache_agency(agency, agency_id, None)
         return agency, agency_id
 
     async def get_agency(self, agency_id: str, thread_id: str | None) -> Agency | None:
-        """Get the agency from the cache."""
+        """Get the agency from the cache. If not found, retrieve from Firestore and repopulate cache."""
         cache_key = self.get_cache_key(agency_id, thread_id)
 
         agency = await self.cache_manager.get(cache_key)
+
+        if agency is None:
+            agency_config = AgencyConfig.load(agency_id)
+            if agency_config:
+                agency = await asyncio.to_thread(self.load_agency_from_config, agency_id, config=agency_config)
+                await self.cache_manager.set(cache_key, agency)
+            else:
+                logger.error(f"Agency configuration for {agency_id} could not be found in the Firestore database.")
+                return None
+
         return agency
 
+    async def update_agency(self, agency_config: AgencyConfig, updated_data: dict) -> None:
+        """Update the agency"""
+        agency_id = agency_config.agency_id
+
+        updated_data.pop("agency_id")
+        agency_config.update(updated_data)
+        agency_config.save()
+
+        agency = await self.get_agency(agency_id, None)
+        if not agency:
+            agency, _ = await self.create_agency(agency_id)
+
+        await self.cache_agency(agency, agency_id, None)
+
     async def cache_agency(self, agency: Agency, agency_id: str, thread_id: str | None) -> None:
         """Cache the agency."""
         cache_key = self.get_cache_key(agency_id, thread_id)
@@ -52,7 +77,7 @@ def get_cache_key(agency_id: str, thread_id: str | None) -> str:
         return f"{agency_id}/{thread_id}" if thread_id else agency_id
 
     @staticmethod
-    def load_agency_from_config(agency_id: str) -> Agency:
+    def load_agency_from_config(agency_id: str, config: AgencyConfig | None = None) -> Agency:
         """Load the agency from the config file. The agency is created using the agency-swarm library.
 
         This code is synchronous and should be run in a single thread.
@@ -60,7 +85,7 @@ def load_agency_from_config(agency_id: str) -> Agency:
         """
 
         start = time.time()
-        config = AgencyConfig.load_or_create(agency_id)
+        config = config or AgencyConfig.load_or_create(agency_id)
 
         agents = {
             agent_conf.role: Agent(
diff --git a/nalgonda/routers/v1/api/agency.py b/nalgonda/routers/v1/api/agency.py
index b3d2c9ef..467b8d68 100644
--- a/nalgonda/routers/v1/api/agency.py
+++ b/nalgonda/routers/v1/api/agency.py
@@ -4,7 +4,7 @@
 
 from agency_swarm import Agency
 from fastapi import APIRouter, Depends, HTTPException
-from starlette.status import HTTP_201_CREATED, HTTP_404_NOT_FOUND
+from starlette.status import HTTP_200_OK, HTTP_404_NOT_FOUND
 
 from nalgonda.dependencies.agency_manager import AgencyManager, get_agency_manager
 from nalgonda.dependencies.auth import get_current_active_user
@@ -28,9 +28,7 @@ async def create_agency(
     # TODO: check if the current_user has permission to create an agency
     logger.info(f"Creating agency for user: {current_user.username}")
 
-    agency, agency_id = await agency_manager.create_agency()
-
-    await agency_manager.cache_agency(agency, agency_id, None)
+    _, agency_id = await agency_manager.create_agency()
     return {"agency_id": agency_id}
 
 
@@ -64,7 +62,7 @@ async def get_agency_config(agency_id: str):
     return agency_config
 
 
-@agency_router.post("/agency/config", status_code=HTTP_201_CREATED)
+@agency_router.put("/agency/config", status_code=HTTP_200_OK)
 async def update_agency_config(
     agency_id: str,
     updated_data: dict,
@@ -73,14 +71,9 @@ async def update_agency_config(
     agency_config = AgencyConfig.load(agency_id)
     if not agency_config:
         raise HTTPException(status_code=HTTP_404_NOT_FOUND, detail="Agency configuration not found")
-    agency_config.update(updated_data)
-    agency_config.save()
 
-    agency = await agency_manager.get_agency(agency_id, None)
-    if not agency:
-        agency, _ = await agency_manager.create_agency(agency_id)
+    await agency_manager.update_agency(agency_config, updated_data)
 
-    await agency_manager.cache_agency(agency, agency_id, None)
     return {"message": "Agency configuration updated successfully"}
 
 

From a383dcdc7662dda2be04e6a4867b6c29b0962281 Mon Sep 17 00:00:00 2001
From: Nikita Bobrovskiy <39348559+bonk1t@users.noreply.github.com>
Date: Sat, 30 Dec 2023 22:22:06 +0200
Subject: [PATCH 06/10] Update tests, bug fix

---
 nalgonda/dependencies/agency_manager.py    | 2 +-
 tests/routers/v1/api/test_agency_routes.py | 4 +++-
 2 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/nalgonda/dependencies/agency_manager.py b/nalgonda/dependencies/agency_manager.py
index 2a2e20c0..c0723d9e 100644
--- a/nalgonda/dependencies/agency_manager.py
+++ b/nalgonda/dependencies/agency_manager.py
@@ -49,7 +49,7 @@ async def update_agency(self, agency_config: AgencyConfig, updated_data: dict) -
         """Update the agency"""
         agency_id = agency_config.agency_id
 
-        updated_data.pop("agency_id")
+        updated_data.pop("agency_id", None)
         agency_config.update(updated_data)
         agency_config.save()
 
diff --git a/tests/routers/v1/api/test_agency_routes.py b/tests/routers/v1/api/test_agency_routes.py
index a14d51c9..3da086c4 100644
--- a/tests/routers/v1/api/test_agency_routes.py
+++ b/tests/routers/v1/api/test_agency_routes.py
@@ -1,5 +1,6 @@
 from unittest.mock import patch
 
+import pytest
 from fastapi.testclient import TestClient
 
 from nalgonda.main import app
@@ -37,11 +38,12 @@ def test_get_agency_config(self):
             "agency_chart": [],
         }
 
+    @pytest.mark.skip("Fix the mocks")
     @patch.object(AgencyConfigFirestoreStorage, "load", mocked_load)
     @patch.object(AgencyConfigFirestoreStorage, "save", mocked_save)
     def test_update_agency_config(self):
         new_data = {"agency_manifesto": "Updated Manifesto"}
-        response = self.client.post("/v1/api/agency/config?agency_id=test_agency", json=new_data)
+        response = self.client.put("/v1/api/agency/config?agency_id=test_agency", json=new_data)
         assert response.status_code == 201
         assert response.json() == {"message": "Agency configuration updated successfully"}
 

From 2a6197840c95254d3758e42185e4f1c9ef4fbdf1 Mon Sep 17 00:00:00 2001
From: Nikita Bobrovskiy <39348559+bonk1t@users.noreply.github.com>
Date: Sat, 30 Dec 2023 22:24:00 +0200
Subject: [PATCH 07/10] Update test_agency_routes.py

---
 tests/routers/v1/api/test_agency_routes.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/tests/routers/v1/api/test_agency_routes.py b/tests/routers/v1/api/test_agency_routes.py
index 3da086c4..3b3369ec 100644
--- a/tests/routers/v1/api/test_agency_routes.py
+++ b/tests/routers/v1/api/test_agency_routes.py
@@ -26,6 +26,7 @@ def mocked_save(self, data: dict):  # noqa: ARG001
 class TestAgencyRoutes:
     client = TestClient(app)
 
+    @pytest.mark.skip("Fix the mocks")
     @patch.object(AgencyConfigFirestoreStorage, "load", mocked_load)
     @patch.object(AgencyConfigFirestoreStorage, "save", mocked_save)
     def test_get_agency_config(self):
@@ -47,6 +48,7 @@ def test_update_agency_config(self):
         assert response.status_code == 201
         assert response.json() == {"message": "Agency configuration updated successfully"}
 
+    @pytest.mark.skip("Fix the mocks")
     @patch.object(AgencyConfigFirestoreStorage, "load", lambda _: None)
     def test_agency_config_not_found(self):
         response = self.client.get("/v1/api/agency/config?agency_id=non_existent_agency")

From c82106cee08dfe5e2e53b4b3742a9ccd07daf513 Mon Sep 17 00:00:00 2001
From: Nikita Bobrovskiy <39348559+bonk1t@users.noreply.github.com>
Date: Tue, 2 Jan 2024 19:18:29 +0200
Subject: [PATCH 08/10] Comment fix

---
 nalgonda/models/agency_config.py | 12 +++++-------
 1 file changed, 5 insertions(+), 7 deletions(-)

diff --git a/nalgonda/models/agency_config.py b/nalgonda/models/agency_config.py
index e5585b70..14c51fc5 100644
--- a/nalgonda/models/agency_config.py
+++ b/nalgonda/models/agency_config.py
@@ -25,17 +25,15 @@ def load(cls, agency_id: str) -> Optional["AgencyConfig"]:
 
     @classmethod
     def load_or_create(cls, agency_id: str) -> "AgencyConfig":
-        with AgencyConfigFirestoreStorage(agency_id) as config_document:
-            config_data = config_document.load()
+        model = cls.load(agency_id)
 
-        if not config_data:
+        if model is None:
             config_data = cls._create_default_config()
-
-        config_data["agency_id"] = agency_id
-        model = cls.model_validate(config_data)
+            config_data["agency_id"] = agency_id
+            model = cls.model_validate(config_data)
 
         with AgencyConfigFirestoreStorage(agency_id) as config_document:
-            config_document.save(config_data)
+            config_document.save(model.model_dump())
 
         return model
 

From 8d79c5502bab778a0d366f68c1c634da4e71b024 Mon Sep 17 00:00:00 2001
From: Nikita Bobrovskiy <39348559+bonk1t@users.noreply.github.com>
Date: Tue, 2 Jan 2024 19:39:11 +0200
Subject: [PATCH 09/10] Fix tests

---
 tests/routers/v1/api/test_agency_routes.py | 45 ++++++++++++++--------
 1 file changed, 30 insertions(+), 15 deletions(-)

diff --git a/tests/routers/v1/api/test_agency_routes.py b/tests/routers/v1/api/test_agency_routes.py
index 3b3369ec..4c0d0c4d 100644
--- a/tests/routers/v1/api/test_agency_routes.py
+++ b/tests/routers/v1/api/test_agency_routes.py
@@ -1,11 +1,9 @@
-from unittest.mock import patch
+from unittest.mock import AsyncMock, patch
 
-import pytest
 from fastapi.testclient import TestClient
 
 from nalgonda.main import app
 from nalgonda.models.agency_config import AgencyConfig
-from nalgonda.persistence.agency_config_firestore_storage import AgencyConfigFirestoreStorage
 
 
 def mocked_load(self):  # noqa: ARG001
@@ -23,12 +21,27 @@ def mocked_save(self, data: dict):  # noqa: ARG001
     }
 
 
+class MockedAgencyConfigFirestoreStorage:
+    def __init__(self, agency_id):
+        self.agency_id = agency_id
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        pass
+
+    def load(self):
+        return mocked_load(self)
+
+    def save(self, data):
+        mocked_save(self, data)
+
+
 class TestAgencyRoutes:
     client = TestClient(app)
 
-    @pytest.mark.skip("Fix the mocks")
-    @patch.object(AgencyConfigFirestoreStorage, "load", mocked_load)
-    @patch.object(AgencyConfigFirestoreStorage, "save", mocked_save)
+    @patch("nalgonda.models.agency_config.AgencyConfigFirestoreStorage", new=MockedAgencyConfigFirestoreStorage)
     def test_get_agency_config(self):
         response = self.client.get("/v1/api/agency/config?agency_id=test_agency")
         assert response.status_code == 200
@@ -39,18 +52,20 @@ def test_get_agency_config(self):
             "agency_chart": [],
         }
 
-    @pytest.mark.skip("Fix the mocks")
-    @patch.object(AgencyConfigFirestoreStorage, "load", mocked_load)
-    @patch.object(AgencyConfigFirestoreStorage, "save", mocked_save)
-    def test_update_agency_config(self):
+    @patch("nalgonda.models.agency_config.AgencyConfigFirestoreStorage", new=MockedAgencyConfigFirestoreStorage)
+    @patch("nalgonda.caching.redis_cache_manager.RedisCacheManager.get", new_callable=AsyncMock)
+    @patch("nalgonda.caching.redis_cache_manager.RedisCacheManager.set", new_callable=AsyncMock)
+    def test_update_agency_config_success(self, mock_redis_set, mock_redis_get):
         new_data = {"agency_manifesto": "Updated Manifesto"}
         response = self.client.put("/v1/api/agency/config?agency_id=test_agency", json=new_data)
-        assert response.status_code == 201
+        assert response.status_code == 200
         assert response.json() == {"message": "Agency configuration updated successfully"}
+        mock_redis_get.assert_called_once()
+        mock_redis_set.assert_called_once()
 
-    @pytest.mark.skip("Fix the mocks")
-    @patch.object(AgencyConfigFirestoreStorage, "load", lambda _: None)
-    def test_agency_config_not_found(self):
+    @patch("nalgonda.models.agency_config.AgencyConfigFirestoreStorage", new=MockedAgencyConfigFirestoreStorage)
+    @patch.object(MockedAgencyConfigFirestoreStorage, "load", lambda _: None)
+    def test_get_agency_config_not_found(self):
         response = self.client.get("/v1/api/agency/config?agency_id=non_existent_agency")
-        assert response.status_code == 404
         assert response.json() == {"detail": "Agency configuration not found"}
+        assert response.status_code == 404

From 744fca9f64946b669d08adc3af5f2584c4b37671 Mon Sep 17 00:00:00 2001
From: Nikita Bobrovskiy <39348559+bonk1t@users.noreply.github.com>
Date: Fri, 5 Jan 2024 19:17:35 +0200
Subject: [PATCH 10/10] Fix newlines

---
 nalgonda/data/default_config.json     | 2 +-
 nalgonda/data/default_config_old.json | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/nalgonda/data/default_config.json b/nalgonda/data/default_config.json
index 67056630..d812605b 100644
--- a/nalgonda/data/default_config.json
+++ b/nalgonda/data/default_config.json
@@ -5,7 +5,7 @@
       "id": null,
       "role": "LeadAndRequirementsGatherer",
       "description": "Specialized in lead capture and software development requirement gathering, this agent will interact with users, guiding them through the initial stages of understanding our AI solutions and collecting relevant information for further engagement.",
-      "instructions": "# Instructions for Virtual Assistant: \\nLead Capture and Requirement Gathering Specialist\\n\\n- Engage with website visitors by introducing them to AI in Hand's services, emphasizing our custom AI automation and the transformative impact it can have on their business operations.\\n- Explain that AI in Hand specializes in bespoke AI solutions, primarily offering 3 groups of solutions: \\n1. Virtual AI Assistants: Custom-designed to reflect a brand's voice and ethos; integrated with CRMs for seamless customer interactions; knowledge base customization for a truly personalized service.\\n2. Custom AI Agents: Tailor-made agents for task automation, including data processing, forecasting, and reporting; driving efficiency and accuracy in day-to-day operations.\\n3. API-Driven Custom Tools: Enhance each solution with our expertise in creating custom tools using APIs, ensuring seamless integration and functionality tailored to specific needs. Explain how these services can be tailored to their unique business needs.\\n- Inquire if the visitor is interested in specifying their business requirements for a custom AI solution, offering to guide them through the process.\\n- Begin with the Initial Interaction stage, asking the visitor to describe the type of AI solution they are interested in and how it might serve their business.\\n- Proceed to the Requirement Gathering stage, asking targeted questions to collect comprehensive details about their AI needs, ensuring to ask one question at a time for clarity.\\n- Once sufficient information is collected, transition to the Lead Capture stage, politely asking for the visitor's preferred name and email address to ensure our team can follow up effectively.\\n- Assure the visitor that their requirements and contact details will be securely saved to our CRM system, and that a member of our team will reach out to them to discuss their custom AI solution further.\\n- Throughout the interaction, maintain a professional and helpful demeanor, using the information about AI in Hand's services and solutions to answer any questions and provide a personalized experience. \\nIMPORTANT: ALWAYS be concise and respond with shorter messages.",
+      "instructions": "# Instructions for Virtual Assistant: \nLead Capture and Requirement Gathering Specialist\n\n- Engage with website visitors by introducing them to AI in Hand's services, emphasizing our custom AI automation and the transformative impact it can have on their business operations.\n- Explain that AI in Hand specializes in bespoke AI solutions, primarily offering 3 groups of solutions: \n1. Virtual AI Assistants: Custom-designed to reflect a brand's voice and ethos; integrated with CRMs for seamless customer interactions; knowledge base customization for a truly personalized service.\n2. Custom AI Agents: Tailor-made agents for task automation, including data processing, forecasting, and reporting; driving efficiency and accuracy in day-to-day operations.\n3. API-Driven Custom Tools: Enhance each solution with our expertise in creating custom tools using APIs, ensuring seamless integration and functionality tailored to specific needs. Explain how these services can be tailored to their unique business needs.\n- Inquire if the visitor is interested in specifying their business requirements for a custom AI solution, offering to guide them through the process.\n- Begin with the Initial Interaction stage, asking the visitor to describe the type of AI solution they are interested in and how it might serve their business.\n- Proceed to the Requirement Gathering stage, asking targeted questions to collect comprehensive details about their AI needs, ensuring to ask one question at a time for clarity.\n- Once sufficient information is collected, transition to the Lead Capture stage, politely asking for the visitor's preferred name and email address to ensure our team can follow up effectively.\n- Assure the visitor that their requirements and contact details will be securely saved to our CRM system, and that a member of our team will reach out to them to discuss their custom AI solution further.\n- Throughout the interaction, maintain a professional and helpful demeanor, using the information about AI in Hand's services and solutions to answer any questions and provide a personalized experience. \nIMPORTANT: ALWAYS be concise and respond with shorter messages.",
       "files_folder": null,
       "tools": [
         "SaveLeadToAirtable"
diff --git a/nalgonda/data/default_config_old.json b/nalgonda/data/default_config_old.json
index 9fb18abb..f5162449 100644
--- a/nalgonda/data/default_config_old.json
+++ b/nalgonda/data/default_config_old.json
@@ -13,7 +13,7 @@
       "id": null,
       "role": "Virtual Assistant",
       "description": "Responsible for drafting emails, doing research and writing proposals. Can also search the web for information.",
-      "instructions": "### Instructions for Virtual Assistant\n\nYour role is to assist users in executing tasks like below. \\\nIf the task is outside of your capabilities, please report back to the user.\n\n#### 1. Drafting Emails\n   - **Understand Context and Tone**: Familiarize yourself with the context of each email. \\\n   Maintain a professional and courteous tone.\n   - **Accuracy and Clarity**: Ensure that the information is accurate and presented clearly. \\\n   Avoid jargon unless it's appropriate for the recipient.\n\n#### 2. Generating Proposals\n   - **Gather Requirements**: Collect all necessary information about the project, \\\n   including client needs, objectives, and any specific requests.\n\n#### 3. Conducting Research\n   - **Understand the Objective**: Clarify the purpose and objectives of the research to focus on relevant information.\n   - **Summarize Findings**: Provide clear, concise summaries of the research findings, \\\n   highlighting key points and how they relate to the project or inquiry.\n   - **Cite Sources**: Properly cite all sources to maintain integrity and avoid plagiarism.",
+      "instructions": "### Instructions for Virtual Assistant\n\nYour role is to assist users in executing tasks like below. \nIf the task is outside of your capabilities, please report back to the user.\n\n#### 1. Drafting Emails\n   - **Understand Context and Tone**: Familiarize yourself with the context of each email. \n   Maintain a professional and courteous tone.\n   - **Accuracy and Clarity**: Ensure that the information is accurate and presented clearly. \n   Avoid jargon unless it's appropriate for the recipient.\n\n#### 2. Generating Proposals\n   - **Gather Requirements**: Collect all necessary information about the project, \n   including client needs, objectives, and any specific requests.\n\n#### 3. Conducting Research\n   - **Understand the Objective**: Clarify the purpose and objectives of the research to focus on relevant information.\n   - **Summarize Findings**: Provide clear, concise summaries of the research findings, \n   highlighting key points and how they relate to the project or inquiry.\n   - **Cite Sources**: Properly cite all sources to maintain integrity and avoid plagiarism.",
       "files_folder": null,
       "tools": [
         "SearchWeb",