API Endpoints for Agency Configuration (#14)

- Add GET and PUT endpoints for Agency Configuration (getting/updating)
- Add cache re-population in case of cache miss
- Add tests for the new endpoints

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"

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",

nalgonda/dependencies/agency_manager.py

@@ -24,17 +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", None)
+        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)
@@ -53,15 +77,15 @@ 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.
         The code is currently not thread safe (due to agency-swarm limitations).
         """
 
         start = time.time()
-        config = AgencyConfig.load(agency_id)
+        config = config or AgencyConfig.load_or_create(agency_id)
 
         agents = {
             agent_conf.role: Agent(

nalgonda/models/agency_config.py

@@ -1,6 +1,10 @@
+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
 
@@ -13,21 +17,45 @@ 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) -> 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(cls, agency_id: str) -> "AgencyConfig":
+    def load_or_create(cls, agency_id: str) -> "AgencyConfig":
+        model = cls.load(agency_id)
+
+        if model is None:
+            config_data = cls._create_default_config()
+            config_data["agency_id"] = agency_id
+            model = cls.model_validate(config_data)
+
         with AgencyConfigFirestoreStorage(agency_id) as config_document:
-            config = config_document.load()
+            config_document.save(model.model_dump())
 
-        config["agency_id"] = agency_id
-        return cls.model_validate(config)
+        return model
+
+    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
+
+    @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

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

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_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
 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,29 @@ async def create_agency_thread(
     return {"thread_id": thread_id}
 
 
+@agency_router.get("/agency/config")
+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")
+    return agency_config
+
+
+@agency_router.put("/agency/config", status_code=HTTP_200_OK)
+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")
+
+    await agency_manager.update_agency(agency_config, updated_data)
+
+    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)

tests/routers/v1/api/test_agency_routes.py

@@ -0,0 +1,71 @@
+from unittest.mock import AsyncMock, patch
+
+from fastapi.testclient import TestClient
+
+from nalgonda.main import app
+from nalgonda.models.agency_config import AgencyConfig
+
+
+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",
+        "agency_chart": [],
+        "agency_id": "test_agency",
+        "agents": [],
+    }
+
+
+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)
+
+    @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
+        assert response.json() == {
+            "agency_id": "test_agency",
+            "agency_manifesto": "Test Manifesto",
+            "agents": [],
+            "agency_chart": [],
+        }
+
+    @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 == 200
+        assert response.json() == {"message": "Agency configuration updated successfully"}
+        mock_redis_get.assert_called_once()
+        mock_redis_set.assert_called_once()
+
+    @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.json() == {"detail": "Agency configuration not found"}
+        assert response.status_code == 404