Minor AI refactoring

README.md

@@ -2,73 +2,45 @@
 
 ## Overview
 
-Project Nalgonda is a tool for managing and executing AI agents.
-It is built on top of the [OpenAI Assitants API](https://platform.openai.com/docs/assistants/overview)
-and provides a simple interface for configuring agents and executing them.
+Project Nalgonda is an innovative platform for managing and executing AI-driven swarm agencies.
+It is built upon the foundational [OpenAI Assistants API](https://platform.openai.com/docs/assistants/overview)
+and extends its functionality through a suite of specialized tools and a sophisticated management system for AI agencies.
+
+## Key Components
+
+- **Agency Configuration Manager**: Manages configurations for AI agencies, ensuring they're loaded and saved properly.
+- **WebSocket Connection Manager**: Handles real-time WebSocket connections for interactive agency-client communication.
+- **Custom Tools**: A collection of tools including `SearchWeb`, `GenerateProposal`, `BuildDirectoryTree`, and more,
+providing specialized functionalities tailored to the needs of different agency roles.
+- **Data Persistence**: Utilizes JSON-based configurations to maintain agency states and preferences across sessions.
 
 ## Features
 
 - **Agency Configuration**: Configure agencies with agents
 - **Tool Configuration**: Configure tools with custom parameters
-- **Tool Execution**: Execute tools and returns results
+- **Tool Execution**: Execute tools and return results
 - **Agent Configuration**: Configure agents with their knowledge and tools
 - **User Management**: Manage users and their access to different agencies [TODO]
 
-## Getting Started
-
-### Prerequisites
-
-- Python 3.11 or higher
-- FastAPI
-- Uvicorn (for running the server)
-- Additional Python packages as listed in `pyproject.toml`
-
-### Installation
-
-1. **Clone the Repository**
-
-   ```sh
-   git clone https://github.com/bonk1t/nalgonda.git
-   cd nalgonda
-   ```
-
-2. **Install Dependencies**
-
-   Using poetry:
-
-   ```sh
-   poetry install
-   ```
+## Installation
 
-   Or using pip:
+Ensure you have Python 3.11 or higher and follow these steps to get started:
 
-   ```sh
-   pip install -r requirements.txt
-   ```
+1. Install the required dependencies (from `requirements.txt` or using Poetry).
+2. Set up the necessary environment variables, including `OPENAI_API_KEY`.
+3. Use the provided JSON configuration files as templates to configure your own AI agencies.
+4. Start the FastAPI server (`uvicorn nalgonda.main:app --reload`) to interact with the system.
 
-3. **Set up Environment Variables**
-
-   Ensure to set up the necessary environment variables such as `OPENAI_API_KEY`.
-
-### Running the Application
-
-1. **Start the FastAPI Server**
-
-   ```sh
-   uvicorn nalgonda.main:app --reload
-   ```
-
-   The API will be available at `http://localhost:8000`.
-
-2. **Accessing the Endpoints**
-
-   Use a tool like Postman or Swagger UI to interact with the API endpoints.
+Note: Refer to individual class and method docstrings for detailed instructions and usage.
 
 ## Usage
 
 ### API Endpoints
-Send a POST request to the /create_agency endpoint to create an agency. The response will contain the following:
-- agency_id: The ID of the agency
 
-### WebSocket Endpoints
-After creating an agency, you can connect to the WebSocket endpoint at /ws/{agency_id} to communicate with the agency.
+Send POST requests to endpoints such as `POST /v1/api/agency` and `POST /v1/api/agency/message` to perform operations
+like creating new agencies and sending messages to them.
+
+### WebSocket Communication
+
+Connect to WebSocket endpoints (e.g., `/v1/ws/{agency_id}`, `/v1/ws/{agency_id}/{thread_id}`)
+to engage in real-time communication with configured AI agencies.

src/nalgonda/agency_config_lock_manager.py

@@ -3,10 +3,23 @@
 
 
 class AgencyConfigLockManager:
-    """Lock manager for agency config files"""
+    """Manages locking for agency configuration files.
 
+    This manager guarantees that each agency configuration has a unique lock,
+    preventing simultaneous access and modification by multiple processes.
+    """
+
+    # Mapping from agency ID to its corresponding Lock.
     _locks: dict[str, threading.Lock] = defaultdict(threading.Lock)
 
     @classmethod
-    def get_lock(cls, agency_id):
+    def get_lock(cls, agency_id: str) -> threading.Lock:
+        """Retrieves the lock for a given agency ID, creating it if not present.
+
+        Args:
+            agency_id (str): The unique identifier for the agency.
+
+        Returns:
+            threading.Lock: The lock associated with the agency ID.
+        """
         return cls._locks[agency_id]

src/nalgonda/agency_manager.py

@@ -5,19 +5,19 @@
 
 from agency_swarm import Agency, Agent
 
-from nalgonda.config import AgencyConfig
 from nalgonda.custom_tools import TOOL_MAPPING
+from nalgonda.models.agency_config import AgencyConfig
 
 logger = logging.getLogger(__name__)
 
 
 class AgencyManager:
-    def __init__(self):
-        self.cache = {}  # agency_id+thread_id: agency
+    def __init__(self) -> None:
+        self.cache: dict[str, Agency] = {}  # Mapping from agency_id+thread_id to Agency class instance
         self.lock = asyncio.Lock()
 
     async def create_agency(self, agency_id: str | None = None) -> tuple[Agency, str]:
-        """Create the agency for the given agency ID."""
+        """Create an agency and return the agency and the agency_id."""
         agency_id = agency_id or uuid.uuid4().hex
 
         async with self.lock:
@@ -27,9 +27,9 @@ async def create_agency(self, agency_id: str | None = None) -> tuple[Agency, str
             return agency, agency_id
 
     async def get_agency(self, agency_id: str, thread_id: str | None) -> Agency | None:
-        """Get the agency for the given agency ID and thread ID."""
+        """Get the agency from the cache."""
         async with self.lock:
-            return self.cache.get(self.get_cache_key(agency_id, thread_id), None)
+            return self.cache.get(self.get_cache_key(agency_id, thread_id))
 
     async def cache_agency(self, agency: Agency, agency_id: str, thread_id: str | None) -> None:
         """Cache the agency for the given agency ID and thread ID."""
@@ -39,9 +39,10 @@ async def cache_agency(self, agency: Agency, agency_id: str, thread_id: str | No
 
     async def delete_agency_from_cache(self, agency_id: str, thread_id: str | None) -> None:
         async with self.lock:
-            self.cache.pop(self.get_cache_key(agency_id, thread_id), None)
+            cache_key = self.get_cache_key(agency_id, thread_id)
+            self.cache.pop(cache_key, None)
 
-    async def refresh_thread_id(self, agency, agency_id, thread_id) -> str | None:
+    async def refresh_thread_id(self, agency: Agency, agency_id: str, thread_id: str | None) -> str | None:
         new_thread_id = agency.main_thread.id
         if thread_id != new_thread_id:
             await self.cache_agency(agency, agency_id, new_thread_id)
@@ -87,8 +88,8 @@ def load_agency_from_config(agency_id: str) -> Agency:
         # It saves all the settings in the settings.json file (in the root folder, not thread safe)
         agency = Agency(agency_chart, shared_instructions=config.agency_manifesto)
 
-        config.update_agent_ids_in_config(agency_id, agents=agency.agents)
-        config.save(agency_id)
+        config.update_agent_ids_in_config(agency.agents)
+        config.save()
 
         logger.info(f"Agency creation took {time.time() - start} seconds. Session ID: {agency_id}")
         return agency

src/nalgonda/constants.py

@@ -1,6 +1,9 @@
 from pathlib import Path
 
-# File and Directory Constants
-DATA_DIR = Path(__file__).resolve().parent / "data"
+# Constants representing base and data directories
+BASE_DIR = Path(__file__).resolve(strict=True).parent
+DATA_DIR = BASE_DIR / "data"
+
+# Constants for default configuration files
 DEFAULT_CONFIG_FILE = DATA_DIR / "default_config.json"
-CONFIG_FILE = DATA_DIR / "config"
+CONFIG_FILE_BASE = DATA_DIR / "config.json"

src/nalgonda/custom_tools/__init__.py

@@ -5,6 +5,7 @@
 from nalgonda.custom_tools.generate_proposal import GenerateProposal
 from nalgonda.custom_tools.print_all_files_in_directory import PrintAllFilesInDirectory
 from nalgonda.custom_tools.search_web import SearchWeb
+from nalgonda.custom_tools.write_and_save_program import WriteAndSaveProgram
 
 TOOL_MAPPING = {
     "CodeInterpreter": CodeInterpreter,
@@ -13,5 +14,5 @@
     "GenerateProposal": GenerateProposal,
     "PrintAllFilesInDirectory": PrintAllFilesInDirectory,
     "SearchWeb": SearchWeb,
-    # "WriteAndSaveProgram": WriteAndSaveProgram,
+    "WriteAndSaveProgram": WriteAndSaveProgram,
 }

src/nalgonda/custom_tools/build_directory_tree.py

@@ -17,10 +17,11 @@ class BuildDirectoryTree(BaseTool):
     )
     file_extensions: set[str] = Field(
         default_factory=set,
-        description="Set of file extensions to include in the tree. If empty, all files will be included.",
+        description="Set of file extensions to include in the tree. If empty, all files will be included. "
+        "Examples are {'.py', '.txt', '.md'}.",
     )
 
-    _validate_start_directory = field_validator("start_directory", mode="before")(check_directory_traversal)
+    _validate_start_directory = field_validator("start_directory", mode="after")(check_directory_traversal)
 
     def run(self) -> str:
         """Recursively print the tree of directories and files using pathlib."""

src/nalgonda/custom_tools/generate_proposal.py

@@ -3,7 +3,7 @@
 
 from nalgonda.custom_tools.utils import get_chat_completion
 
-USER_PROMPT_PREFIX = "Please draft a proposal for the following project brief: "
+USER_PROMPT_PREFIX = "Please draft a proposal for the following project brief: \n"
 SYSTEM_MESSAGE = """\
 You are a professional proposal drafting assistant. \
 Do not include any actual technologies or technical details into proposal unless \
@@ -19,14 +19,6 @@ class GenerateProposal(BaseTool):
     project_brief: str = Field(..., description="The project brief to generate a proposal for.")
 
     def run(self) -> str:
-        user_prompt = self.get_user_prompt()
-        message = get_chat_completion(
-            user_prompt=user_prompt,
-            system_message=SYSTEM_MESSAGE,
-            temperature=0.6,
-        )
-
-        return message
-
-    def get_user_prompt(self):
-        return f"{USER_PROMPT_PREFIX}\n{self.project_brief}"
+        user_prompt = f"{USER_PROMPT_PREFIX}{self.project_brief}"
+        response = get_chat_completion(user_prompt=user_prompt, system_message=SYSTEM_MESSAGE, temperature=0.6)
+        return response

src/nalgonda/custom_tools/print_all_files_in_directory.py

@@ -17,10 +17,11 @@ class PrintAllFilesInDirectory(BaseTool):
     )
     file_extensions: set[str] = Field(
         default_factory=set,
-        description="Set of file extensions to include in the output. If empty, all files will be included.",
+        description="Set of file extensions to include in the tree. If empty, all files will be included. "
+        "Examples are {'.py', '.txt', '.md'}.",
     )
 
-    _validate_start_directory = field_validator("start_directory", mode="before")(check_directory_traversal)
+    _validate_start_directory = field_validator("start_directory", mode="after")(check_directory_traversal)
 
     def run(self) -> str:
         """

src/nalgonda/custom_tools/search_web.py

@@ -10,7 +10,8 @@ class SearchWeb(BaseTool):
         ...,
         description="The search phrase you want to use. " "Optimize the search phrase for an internet search engine.",
     )
+    max_results: int = Field(default=10, description="The maximum number of search results to return, default is 10.")
 
-    def run(self):
+    def run(self) -> str:
         with DDGS() as ddgs:
-            return str("\n".join(str(r) for r in ddgs.text(self.phrase, max_results=10)))
+            return "\n".join(str(result) for result in ddgs.text(self.phrase, max_results=self.max_results))