From 177e69f6714e070436da7e40813678d274254d4b Mon Sep 17 00:00:00 2001 From: Mohamed Date: Mon, 1 Jul 2024 12:43:55 +0100 Subject: [PATCH] First commit --- .env.example | 6 ++ .github/workflows/ci.yml | 37 +++++++ .gitignore | 81 ++++++++++++++ CHANGELOG.md | 31 ++++++ CONTRIBUTING.md | 59 +++++++++++ LICENSE | 21 ++++ MANIFEST.in | 7 ++ Makefile | 17 +++ README.md | 209 +++++++++++++++++++++++++++++++++++++ conftest.py | 5 + pyproject.toml | 30 ++++++ requirements-dev.txt | 30 ++++++ requirements.txt | 41 ++++++++ setup.py | 39 +++++++ setup.sh | 18 ++++ src/__init__.py | 3 + src/assistants.py | 93 +++++++++++++++++ src/config.py | 57 ++++++++++ src/main.py | 36 +++++++ src/orchestrator.py | 155 +++++++++++++++++++++++++++ tests/__init__.py | 0 tests/test_orchestrator.py | 131 +++++++++++++++++++++++ tox.ini | 16 +++ 23 files changed, 1122 insertions(+) create mode 100644 .env.example create mode 100644 .github/workflows/ci.yml create mode 100644 .gitignore create mode 100644 CHANGELOG.md create mode 100644 CONTRIBUTING.md create mode 100644 LICENSE create mode 100644 MANIFEST.in create mode 100644 Makefile create mode 100644 README.md create mode 100644 conftest.py create mode 100644 pyproject.toml create mode 100644 requirements-dev.txt create mode 100644 requirements.txt create mode 100644 setup.py create mode 100755 setup.sh create mode 100644 src/__init__.py create mode 100644 src/assistants.py create mode 100644 src/config.py create mode 100644 src/main.py create mode 100644 src/orchestrator.py create mode 100644 tests/__init__.py create mode 100644 tests/test_orchestrator.py create mode 100644 tox.ini diff --git a/.env.example b/.env.example new file mode 100644 index 0000000..ffef1c2 --- /dev/null +++ b/.env.example @@ -0,0 +1,6 @@ +ANTHROPIC_API_KEY="sk-..." +GOOGLE_API_KEY="gm-..." +OPENAI_API_KEY="sk-..." +TAVILY_API_KEY="tvly-...." +VertexAI_Project_Name="" +VertexAI_Location="" \ No newline at end of file diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000..372615b --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,37 @@ +name: CI +on: + push: + branches: [main] + pull_request: + branches: [main] +jobs: + build: + runs-on: ubuntu-latest + strategy: + matrix: + python-version: ["3.9", "3.10", "3.11"] + steps: + - uses: actions/checkout@v4 + + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + + - name: Install dependencies + run: | + python -m pip install --upgrade pip + pip install -r requirements.txt + pip install -r requirements-dev.txt + + - name: Run tests + env: + PROJECT_ID: ${{ secrets.PROJECT_ID }} + LOCATION: ${{ secrets.LOCATION }} + run: | + python -m pytest + + - name: Check code formatting + run: | + black --check . + isort --check-only --diff . diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..f35a750 --- /dev/null +++ b/.gitignore @@ -0,0 +1,81 @@ +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[cod] +*$py.class + +# C extensions +*.so + +# Distribution / packaging +.Python +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +pip-wheel-metadata/ +share/python-wheels/ +*.egg-info/ +.installed.cfg +*.egg +MANIFEST + +# PyInstaller +*.manifest +*.spec + +# Installer logs +pip-log.txt +pip-delete-this-directory.txt + +# Unit test / coverage reports +htmlcov/ +.tox/ +.nox/ +.coverage +.coverage.* +.cache +nosetests.xml +coverage.xml +*.cover +*.py,cover +.hypothesis/ +.pytest_cache/ + +# Environments +.env +.venv +env/ +venv/ +ENV/ + +# Spyder project settings +.spyderproject +.spyproject + +# Rope project settings +.ropeproject + +# mkdocs documentation +/site + +# mypy +.mypy_cache/ + +# IDE settings +.vscode/ +.idea/ + +# Dev tools +concatenated_files.txt +copy_py_files.py + +# An empty output folder +output/ \ No newline at end of file diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 index 0000000..58c70a0 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,31 @@ +# Changelog + +All notable changes to this project will be documented in this file. + +The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), +and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). + +## [Unreleased] + +## [0.1.1] - 2024-06-30 + +### Added + +- Implement dynamic assistant creation supporting Anthropic, OpenAI, and Google LLMs +- Add configuration settings for VertexAI project and location +- Update requirements and setup files with new dependencies +- Refactor assistants module to use fallback mechanisms and retry logic +- Enhance orchestrator with improved logging and docstrings +- Add CI workflow for testing, linting, and formatting +- Update project structure and .gitignore rules + +## [0.1.0] - 2024-06-29 + +### Added + +- Initial release of Smart Autonomous Assistants (SAAs) +- Core functionality including Orchestrator, MainAssistant, SubAssistant, and RefinerAssistant +- File operations for creating, reading, and listing files +- Command-line interface for running workflows +- Configuration management using pydantic and python-dotenv +- Basic error handling and logging diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md new file mode 100644 index 0000000..9c0f2ea --- /dev/null +++ b/CONTRIBUTING.md @@ -0,0 +1,59 @@ +# Contributing to Smart Autonomous Assistants (SAAs) + +We love your input! We want to make contributing to this project as easy and transparent as possible, whether it's: + +- Reporting a bug +- Discussing the current state of the code +- Submitting a fix +- Proposing new features +- Becoming a maintainer + +## We Develop with Github + +We use github to host code, to track issues and feature requests, as well as accept pull requests. + +## We Use [Github Flow](https://guides.github.com/introduction/flow/index.html), So All Code Changes Happen Through Pull Requests + +Pull requests are the best way to propose changes to the codebase. We actively welcome your pull requests: + +1. Fork the repo and create your branch from `main`. +2. If you've added code that should be tested, add tests. +3. If you've changed APIs, update the documentation. +4. Ensure the test suite passes. +5. Make sure your code lints. +6. Issue that pull request! + +## Any contributions you make will be under the MIT Software License + +In short, when you submit code changes, your submissions are understood to be under the same [MIT License](http://choosealicense.com/licenses/mit/) that covers the project. Feel free to contact the maintainers if that's a concern. + +## Report bugs using Github's [issues](https://github.com/waveuphq/smart-autonomous-assistants/issues) + +We use GitHub issues to track public bugs. Report a bug by [opening a new issue](https://github.com/waveuphq/smart-autonomous-assistants/issues/new); it's that easy! + +## Write bug reports with detail, background, and sample code + +**Great Bug Reports** tend to have: + +- A quick summary and/or background +- Steps to reproduce + - Be specific! + - Give sample code if you can. +- What you expected would happen +- What actually happens +- Notes (possibly including why you think this might be happening, or stuff you tried that didn't work) + +People _love_ thorough bug reports. I'm not even kidding. + +## Use a Consistent Coding Style + +- 4 spaces for indentation rather than tabs +- You can try running `pylint` for style unification + +## License + +By contributing, you agree that your contributions will be licensed under its MIT License. + +## References + +This document was adapted from the open-source contribution guidelines for [Facebook's Draft](https://github.com/facebook/draft-js/blob/master/CONTRIBUTING.md) diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..9d7455c --- /dev/null +++ b/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2024 Waveup Digital + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/MANIFEST.in b/MANIFEST.in new file mode 100644 index 0000000..a97a29b --- /dev/null +++ b/MANIFEST.in @@ -0,0 +1,7 @@ +include LICENSE +include README.md +include CHANGELOG.md +include requirements.txt +recursive-include tests * +recursive-exclude * __pycache__ +recursive-exclude * *.py[co] diff --git a/Makefile b/Makefile new file mode 100644 index 0000000..5a40ed3 --- /dev/null +++ b/Makefile @@ -0,0 +1,17 @@ +.PHONY: install test lint format clean + +install: + pip install -r requirements.txt + +test: + pytest + +lint: + pylint **/*.py + +format: + black . + +clean: + find . -type f -name "*.pyc" -delete + find . -type d -name "__pycache__" -delete diff --git a/README.md b/README.md new file mode 100644 index 0000000..67a80b5 --- /dev/null +++ b/README.md @@ -0,0 +1,209 @@ +# Smart Autonomous Assistants (SAAs) Project + +## Table of Contents + +- [Smart Autonomous Assistants (SAAs) Project](#smart-autonomous-assistants-saas-project) + - [Table of Contents](#table-of-contents) + - [Introduction](#introduction) + - [Features](#features) + - [Project Structure](#project-structure) + - [Module Descriptions](#module-descriptions) + - [1. src/assistants.py](#1-srcassistantspy) + - [2. src/config.py](#2-srcconfigpy) + - [3. src/main.py](#3-srcmainpy) + - [4. src/orchestrator.py](#4-srcorchestratorpy) + - [Dependencies](#dependencies) + - [Setup and Installation](#setup-and-installation) + - [Configuration](#configuration) + - [Usage](#usage) + - [Development Setup](#development-setup) + - [Testing](#testing) + - [Continuous Integration](#continuous-integration) + - [System Architecture](#system-architecture) + - [Contributing](#contributing) + - [License](#license) + +## Introduction + +The Smart Autonomous Assistants (SAAs) project is a sophisticated AI-driven system designed to orchestrate multiple AI assistants to accomplish complex tasks. By leveraging the power of large language models and a modular architecture, this system can break down objectives, execute sub-tasks, and refine results to produce coherent outputs. + +## Features + +- Multi-assistant orchestration for complex task completion +- Support for multiple LLM providers (Claude, GPT, Gemini) +- Modular architecture for easy extension and customization +- Automated task breakdown and execution +- Integration with external tools (e.g., TavilyTools) +- File operation capabilities for input/output handling +- Detailed logging of workflow execution +- Fallback mechanisms for improved reliability + +## Project Structure + +``` +smart-autonomous-assistants/ +├── src/ +│ ├── __init__.py +│ ├── assistants.py +│ ├── config.py +│ ├── main.py +│ └── orchestrator.py +├── tests/ +│ ├── __init__.py +│ └── test_orchestrator.py +├── .github/ +│ └── workflows/ +│ └── ci.yml +├── output/ +├── README.md +├── setup.py +├── pyproject.toml +├── requirements.txt +├── .gitignore +└── .env +``` + +## Module Descriptions + +### 1. src/assistants.py + +- Implements dynamic assistant creation supporting multiple LLM providers +- Manages file operations and tool integration + +### 2. src/config.py + +- Handles configuration settings and environment variables +- Implements API key management and validation + +### 3. src/main.py + +- Provides the command-line interface using Typer +- Handles workflow initialization and error reporting + +### 4. src/orchestrator.py + +- Implements the core workflow management logic +- Coordinates interactions between assistants and manages the overall process + +## Dependencies + +| Dependency | Version | Purpose | +| ------------- | ------- | ------------------------------------------- | +| phidata | 2.4.22 | Provides the base Assistant class and tools | +| pydantic | 2.7.4 | Data validation and settings management | +| python-dotenv | 1.0.1 | Loads environment variables from .env file | +| typer | 0.12.3 | Creates CLI interfaces | +| rich | 13.7.1 | Enhanced terminal output | + +## Setup and Installation + +1. Clone the repository: + + ``` + git clone https://github.com/waveuphq/smart-autonomous-assistants.git + cd smart-autonomous-assistants + ``` + +2. Create and activate a virtual environment: + + ``` + python -m venv venv + source venv/bin/activate # On Windows, use `venv\Scripts\activate` + ``` + +3. Install dependencies: + + ``` + pip install -r requirements.txt + ``` + +4. Create a `.env` file in the project root and add your API keys and VertexAI settings: + ``` + ANTHROPIC_API_KEY=your_anthropic_api_key + OPENAI_API_KEY=your_openai_api_key + GOOGLE_API_KEY=your_google_api_key + TAVILY_API_KEY=your_tavily_api_key + VertexAI_Project_Name=your_vertexai_project_id + VertexAI_Location=your_vertexai_location + ``` + +## Configuration + +The project supports multiple LLM providers. Update the `MAIN_ASSISTANT`, `SUB_ASSISTANT`, and `REFINER_ASSISTANT` settings in `src/config.py` to use the desired models. + +## Usage + +To run a workflow, use the following command: + +``` +python -m src.main run-workflow "Your objective here" +``` + +Example: + +``` +python -m src.main run-workflow "Create a python script to copy all .py files content and exclude files and folder excluded in the .gitignore uses Typer commands" +``` + +## Development Setup + +To set up the development environment: + +1. Clone the repository and navigate to the project directory. +2. Create and activate a virtual environment: + ``` + python -m venv venv + source venv/bin/activate # On Windows, use `venv\Scripts\activate` + ``` +3. Install development dependencies: + ``` + pip install -r requirements-dev.txt + ``` +4. Install pre-commit hooks: + ``` + pre-commit install + ``` + +This will set up your environment with all necessary development tools, including testing, linting, and formatting utilities. + +## Testing + +Run tests using pytest: + +``` +pytest +``` + +## Continuous Integration + +The project uses GitHub Actions for CI/CD. The pipeline runs tests, checks code formatting, and verifies import sorting on each push and pull request to the main branch. + +## System Architecture + +```mermaid +graph TD + A[User Input] --> B[main.py] + B --> C[Orchestrator] + C --> D[MainAssistant] + C --> E[SubAssistant] + C --> F[RefinerAssistant] + D --> C + E --> C + F --> C + C --> G[File Operations] + G --> H[Output Directory] + C --> I[Exchange Log] + I --> H + J[config.py] --> C + K[assistants.py] --> C +``` + +This architecture allows for a flexible and extensible system that can handle complex, multi-step tasks by leveraging the strengths of multiple AI assistants. + +## Contributing + +Contributions are welcome! Please feel free to submit a Pull Request. + +## License + +This project is licensed under the MIT License - see the [LICENSE](LICENSE) file for details. diff --git a/conftest.py b/conftest.py new file mode 100644 index 0000000..7eaf332 --- /dev/null +++ b/conftest.py @@ -0,0 +1,5 @@ +import os +import sys + +# Add the src directory to the Python path +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "src"))) diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..589af1c --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,30 @@ +[build-system] +requires = ["setuptools>=42", "wheel"] +build-backend = "setuptools.build_meta" + +[tool.black] +line-length = 100 +target-version = ['py37', 'py38', 'py39'] +include = '\.pyi?$' +extend-exclude = ''' +/( + # directories + \.eggs + | \.git + | \.hg + | \.mypy_cache + | \.tox + | \.venv + | build + | dist +)/ +''' + +[tool.isort] +profile = "black" +multi_line_output = 3 +include_trailing_comma = true +force_grid_wrap = 0 +use_parentheses = true +ensure_newline_before_comments = true +line_length = 100 \ No newline at end of file diff --git a/requirements-dev.txt b/requirements-dev.txt new file mode 100644 index 0000000..8ad51d6 --- /dev/null +++ b/requirements-dev.txt @@ -0,0 +1,30 @@ +# Project dependencies +-r requirements.txt + +# Testing +pytest==8.2.2 +pytest-mock==3.14.0 + +# Type checking +mypy==1.7.1 + +# Debugging +ipdb==0.13.13 + +# Security +bandit==1.7.5 + +# Performance profiling +py-spy==0.3.14 + +# Dependency management +# Note: pip-tools is not in pip freeze, you may want to add it if needed + +# Environment management +python-dotenv==1.0.1 # Already in requirements.txt, but included here for completeness + +# CLI tools +typer[all]==0.12.3 # Already in requirements.txt, but included here for completeness +# Code Formating +black==24.4.2 +isort==5.12.0 \ No newline at end of file diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..8d082a2 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,41 @@ +# Core dependencies +phidata==2.4.22 +pydantic==2.7.4 +pydantic-settings==2.3.4 +python-dotenv==1.0.1 +typer==0.12.3 +rich==13.7.1 + +# LLM providers +google-cloud-aiplatform==1.57.0 +openai==1.35.7 +anthropic==0.30.0 + +# Additional tools +tavily-python==0.3.3 + +# Networking and API +requests==2.32.3 +httpx==0.27.0 + +# Data handling +numpy==2.0.0 + +# File operations +python-multipart==0.0.9 + +# Async support +anyio==4.4.0 + +# CLI enhancements +click==8.1.7 +shellingham==1.5.4 + +# Time handling +python-dateutil==2.9.0.post0 + +# pgvector +# pypdf +# psycopg2-binary +# sqlalchemy +# fastapi \ No newline at end of file diff --git a/setup.py b/setup.py new file mode 100644 index 0000000..c5b5bc0 --- /dev/null +++ b/setup.py @@ -0,0 +1,39 @@ +from setuptools import find_packages, setup + +with open("README.md", "r", encoding="utf-8") as fh: + long_description = fh.read() + +setup( + name="smart-autonomous-assistants", + version="0.1.1", + author="jeblister", + author_email="jeblister@waveup.dev", + description="A system for orchestrating multiple AI assistants to accomplish complex tasks", + long_description=long_description, + long_description_content_type="text/markdown", + url="https://github.com/waveuphq/smart-autonomous-assistants", + packages=find_packages(where="src"), + classifiers=[ + "Development Status :: 3 - Alpha", + "Intended Audience :: Developers", + "License :: OSI Approved :: MIT License", + "Operating System :: OS Independent", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3.7", + "Programming Language :: Python :: 3.8", + "Programming Language :: Python :: 3.9", + ], + python_requires=">=3.7", + install_requires=[ + "phidata==2.4.22", + "pydantic==2.7.4", + "python-dotenv==1.0.1", + "typer==0.12.3", + "rich==13.7.1", + ], + entry_points={ + "console_scripts": [ + "smart-assistants=src.main:app", + ], + }, +) diff --git a/setup.sh b/setup.sh new file mode 100755 index 0000000..bfd6b51 --- /dev/null +++ b/setup.sh @@ -0,0 +1,18 @@ +#!/bin/bash + +# Create virtual environment +python3 -m venv venv +source venv/bin/activate + +# Set up PostgreSQL with PgVector +docker run -d \ + -e POSTGRES_DB=ai \ + -e POSTGRES_USER=ai \ + -e POSTGRES_PASSWORD=ai \ + -e PGDATA=/var/lib/postgresql/data/pgdata \ + -v pgvolume:/var/lib/postgresql/data \ + -p 5532:5432 \ + --name pgvector \ + phidata/pgvector:16 + +echo "Environment setup complete!" \ No newline at end of file diff --git a/src/__init__.py b/src/__init__.py new file mode 100644 index 0000000..2c55aa0 --- /dev/null +++ b/src/__init__.py @@ -0,0 +1,3 @@ +from .assistants import get_full_response, main_assistant, refiner_assistant, sub_assistant +from .config import settings +from .orchestrator import Orchestrator, Task, TaskExchange diff --git a/src/assistants.py b/src/assistants.py new file mode 100644 index 0000000..9b8faa0 --- /dev/null +++ b/src/assistants.py @@ -0,0 +1,93 @@ +import logging +import os +import time + +import vertexai +from dotenv import load_dotenv +from phi.assistant import Assistant +from phi.llm.anthropic import Claude +from phi.llm.gemini import Gemini +from phi.llm.openai import OpenAIChat +from phi.tools.tavily import TavilyTools + +from src.config import settings + +load_dotenv() # This loads the variables from .env + +# Ensure the output directory exists +output_dir = os.path.join(os.getcwd(), "output") +os.makedirs(output_dir, exist_ok=True) + +# Initialize VertexAI +try: + vertexai.init(project=settings.PROJECT_ID, location=settings.LOCATION) +except Exception as e: + logging.error(f"Error initializing VertexAI: {str(e)}") + + +def create_file(file_path: str, content: str): + full_path = os.path.join(output_dir, file_path) + os.makedirs(os.path.dirname(full_path), exist_ok=True) + with open(full_path, "w") as f: + f.write(content) + return f"File created: {full_path}" + + +def read_file(file_path: str): + full_path = os.path.join(output_dir, file_path) + if os.path.exists(full_path): + with open(full_path, "r") as f: + return f.read() + return f"File not found: {full_path}" + + +def list_files(directory: str = ""): + full_path = os.path.join(output_dir, directory) + if os.path.exists(full_path): + return os.listdir(full_path) + return f"Directory not found: {full_path}" + + +# Create assistants +def create_assistant(name: str, model: str): + if model.startswith("gemini"): + llm = Gemini(model=model) + elif model.startswith("claude"): + llm = Claude(model=model, api_key=settings.ANTHROPIC_API_KEY) + elif model.startswith("gpt"): + llm = OpenAIChat(model=model, api_key=settings.OPENAI_API_KEY) + else: + raise ValueError(f"Unsupported model: {model}") + + return Assistant( + name=name, + llm=llm, + description="You are a helpful assistant.", + tools=[TavilyTools(api_key=settings.TAVILY_API_KEY), create_file, read_file, list_files], + ) + + +# Create assistants +main_assistant = create_assistant("MainAssistant", settings.MAIN_ASSISTANT) +sub_assistant = create_assistant("SubAssistant", settings.SUB_ASSISTANT) +refiner_assistant = create_assistant("RefinerAssistant", settings.REFINER_ASSISTANT) + + +def get_full_response(assistant: Assistant, prompt: str, max_retries=3, delay=2) -> str: + for attempt in range(max_retries): + try: + response = assistant.run(prompt, stream=False) + if isinstance(response, str): + return response + elif isinstance(response, list): + return " ".join(map(str, response)) + else: + return str(response) + except Exception as e: + logging.error(f"Attempt {attempt + 1} failed: {str(e)}") + if attempt < max_retries - 1: + time.sleep(delay) + else: + raise + + raise Exception("Max retries reached. Could not get a response from the assistant.") diff --git a/src/config.py b/src/config.py new file mode 100644 index 0000000..7da2b77 --- /dev/null +++ b/src/config.py @@ -0,0 +1,57 @@ +import os +from typing import Optional + +from dotenv import load_dotenv +from pydantic_settings import BaseSettings + +load_dotenv() # This loads the variables from .env + + +class Settings(BaseSettings): + # Project settings + PROJECT_NAME: str = "SAA-Orchestrator" + + # VertexAI settings + PROJECT_ID: str = os.getenv("VertexAI_Project_Name") + LOCATION: str = os.getenv("VertexAI_Location") + + @property + def anthropic_api_key(self) -> str: + if not self.ANTHROPIC_API_KEY: + raise ValueError("ANTHROPIC_API_KEY is not set in the environment") + return self.ANTHROPIC_API_KEY + + @property + def openai_api_key(self) -> str: + if not self.OPENAI_API_KEY: + raise ValueError("OPENAI_API_KEY is not set in the environment") + return self.OPENAI_API_KEY + + @property + def tavily_api_key(self) -> str: + if not self.TAVILY_API_KEY: + raise ValueError("TAVILY_API_KEY is not set in the environment") + return self.TAVILY_API_KEY + + # API Keys + ANTHROPIC_API_KEY: Optional[str] = os.getenv("ANTHROPIC_API_KEY") + OPENAI_API_KEY: Optional[str] = os.getenv("OPENAI_API_KEY") + + # Assistant settings + MAIN_ASSISTANT: str = "claude-3-5-sonnet-20240620" + SUB_ASSISTANT: str = "gpt-3.5-turbo" + REFINER_ASSISTANT: str = "gemini-1.5-pro-preview-0409" + + # Fallback models + FALLBACK_MODEL_1: str = "claude-3-sonnet-20240229" + FALLBACK_MODEL_2: str = "gpt-3.5-turbo" + + # Tools + TAVILY_API_KEY: Optional[str] = os.getenv("TAVILY_API_KEY") + + class Config: + env_file = ".env" + extra = "ignore" # This will ignore any extra fields in the environment + + +settings = Settings() diff --git a/src/main.py b/src/main.py new file mode 100644 index 0000000..9ce90e4 --- /dev/null +++ b/src/main.py @@ -0,0 +1,36 @@ +import typer +from rich import print as rprint + +from src.orchestrator import Orchestrator + +app = typer.Typer() + + +@app.command() +def run_workflow( + objective: list[str] = typer.Argument( + ..., help="The objective for the workflow. Can be multiple words." + ) +): + """ + Run the SAA Orchestrator workflow with the given objective. + """ + full_objective = " ".join(objective) + try: + rprint("[bold]Starting SAA Orchestrator[/bold]") + orchestrator = Orchestrator() + result = orchestrator.run_workflow(full_objective) + + rprint("\n[bold green]Workflow completed![/bold green]") + rprint("\n[bold]Final Output:[/bold]") + rprint(result) + rprint("\n[bold blue]Exchange log saved to 'exchange_log.md'[/bold blue]") + except ValueError as e: + rprint(f"[bold red]Error:[/bold red] {str(e)}") + rprint("Please make sure you have set the required API keys in your .env file.") + except Exception as e: + rprint(f"[bold red]An unexpected error occurred:[/bold red] {str(e)}") + + +if __name__ == "__main__": + app() diff --git a/src/orchestrator.py b/src/orchestrator.py new file mode 100644 index 0000000..90780db --- /dev/null +++ b/src/orchestrator.py @@ -0,0 +1,155 @@ +import json +import logging +import os +from typing import Any, Dict, List, Literal, Union + +from pydantic import BaseModel, Field +from rich import print as rprint + +from .assistants import create_assistant, get_full_response +from .config import settings + +logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s") + + +class TaskExchange(BaseModel): + role: Literal["user", "main_assistant", "sub_assistant", "refiner_assistant"] = Field(...) + content: str = Field(...) + + +class Task(BaseModel): + task: str + result: str + + def to_dict(self) -> Dict[str, Any]: + return {"task": str(self.task), "result": str(self.result)} + + +class State(BaseModel): + task_exchanges: List[TaskExchange] = [] + tasks: List[Task] = [] + + def to_dict(self) -> Dict[str, Any]: + return { + "task_exchanges": [exchange.model_dump() for exchange in self.task_exchanges], + "tasks": [task.to_dict() for task in self.tasks], + } + + +class Orchestrator(BaseModel): + """ + Manages the workflow of AI assistants to accomplish complex tasks. + """ + + state: State = State() + output_dir: str = Field(default_factory=lambda: os.path.join(os.getcwd(), "output")) + + def __init__(self, **data): + super().__init__(**data) + os.makedirs(self.output_dir, exist_ok=True) + + def run_workflow(self, objective: str) -> str: + """ + Executes the workflow to accomplish the given objective. + + Args: + objective (str): The main task or goal to be accomplished. + + Returns: + str: The final refined output of the workflow. + """ + rprint(f"[bold green]Starting workflow with objective:[/bold green] {objective}") + logging.info(f"Starting workflow with objective: {objective}") + self.state.task_exchanges.append(TaskExchange(role="user", content=objective)) + + task_counter = 1 + while True: + rprint(f"\n[bold blue]--- Task {task_counter} ---[/bold blue]") + logging.info(f"Starting task {task_counter}") + main_prompt = ( + f"Objective: {objective}\n\n" + f"Current progress:\n{json.dumps(self.state.to_dict(), indent=2)}\n\n" + "Break down this objective into the next specific sub-task, or if the objective is fully achieved, " + "start your response with 'ALL DONE:' followed by the final output." + ) + logging.debug(f"Main prompt: {main_prompt}") + main_response = get_full_response( + create_assistant("MainAssistant", settings.MAIN_ASSISTANT), main_prompt + ) + + logging.info(f"Main assistant response received") + self.state.task_exchanges.append( + TaskExchange(role="main_assistant", content=main_response) + ) + rprint(f"[green]MAIN_ASSISTANT response:[/green] {main_response[:100]}...") + + if main_response.startswith("ALL DONE:"): + rprint("[bold green]Workflow completed![/bold green]") + logging.info("Workflow completed") + break + + sub_task_prompt = ( + f"Previous tasks: {json.dumps([task.to_dict() for task in self.state.tasks], indent=2)}\n\n" + f"Current task: {main_response}\n\n" + "Execute this task and provide the result. Use the provided functions to create, read, or list files as needed. " + f"All file operations should be relative to the '{self.output_dir}' directory." + ) + logging.debug(f"Sub-task prompt: {sub_task_prompt}") + sub_response = get_full_response( + create_assistant("SubAssistant", settings.SUB_ASSISTANT), sub_task_prompt + ) + + logging.info(f"Sub-assistant response received") + self.state.task_exchanges.append( + TaskExchange(role="sub_assistant", content=sub_response) + ) + self.state.tasks.append(Task(task=main_response, result=sub_response)) + rprint(f"[green]SUB_ASSISTANT response:[/green] {sub_response[:100]}...") + + task_counter += 1 + + refiner_prompt = ( + f"Original objective: {objective}\n\n" + f"Task breakdown and results: {json.dumps([task.to_dict() for task in self.state.tasks], indent=2)}\n\n" + "Please refine these results into a coherent final output, summarizing the project structure created. " + f"You can use the provided functions to list and read files if needed. All files are in the '{self.output_dir}' directory. " + "Provide your response as a string, not a list or dictionary." + ) + logging.debug(f"Refiner prompt: {refiner_prompt}") + refined_output = get_full_response( + create_assistant("RefinerAssistant", settings.REFINER_ASSISTANT), refiner_prompt + ) + + logging.info("Refiner assistant response received") + self.state.task_exchanges.append( + TaskExchange(role="refiner_assistant", content=refined_output) + ) + rprint(f"[green]REFINER_ASSISTANT response:[/green] {refined_output[:100]}...") + + self._save_exchange_log(objective, refined_output) + rprint("[bold blue]Exchange log saved to 'exchange_log.md'[/bold blue]") + logging.info("Exchange log saved") + + return refined_output + + def _save_exchange_log(self, objective: str, final_output: str): + """ + Saves the workflow exchange log to a markdown file. + + Args: + objective (str): The original objective of the workflow. + final_output (str): The final refined output of the workflow. + """ + log_content = f"# SAA Orchestrator Exchange Log\n\n" + log_content += f"## Objective\n{objective}\n\n" + log_content += f"## Task Breakdown and Execution\n\n" + + for exchange in self.state.task_exchanges: + log_content += f"### {exchange.role.capitalize()}\n{exchange.content}\n\n" + + log_content += f"## Final Output\n{final_output}\n" + + log_file_path = os.path.join(self.output_dir, "exchange_log.md") + with open(log_file_path, "w") as f: + f.write(log_content) + rprint(f"[blue]Exchange log saved to:[/blue] {log_file_path}") diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/test_orchestrator.py b/tests/test_orchestrator.py new file mode 100644 index 0000000..8683219 --- /dev/null +++ b/tests/test_orchestrator.py @@ -0,0 +1,131 @@ +import os +from unittest.mock import MagicMock, patch + +import pytest + +from src.orchestrator import Orchestrator, Task, TaskExchange + + +@pytest.fixture +def orchestrator(tmp_path): + output_dir = tmp_path / "output" + output_dir.mkdir() + return Orchestrator(output_dir=str(output_dir)) + + +def test_orchestrator_initialization(orchestrator): + assert orchestrator.state is not None + assert len(orchestrator.state.task_exchanges) == 0 + assert len(orchestrator.state.tasks) == 0 + + +def test_task_exchange(): + exchange = TaskExchange(role="user", content="Test content") + assert exchange.role == "user" + assert exchange.content == "Test content" + + +def test_task(): + task = Task(task="Test task", result="Test result") + assert task.task == "Test task" + assert task.result == "Test result" + assert task.to_dict() == {"task": "Test task", "result": "Test result"} + + +@patch("src.orchestrator.create_assistant") +@patch("src.orchestrator.get_full_response") +def test_run_workflow(mock_get_full_response, mock_create_assistant, orchestrator): + mock_assistant = MagicMock() + mock_create_assistant.return_value = mock_assistant + + mock_get_full_response.side_effect = [ + "First sub-task", + "Result of first sub-task", + "ALL DONE: Final output", + "Refined output", + ] + + result = orchestrator.run_workflow("Test objective") + + assert isinstance(result, str), f"Expected string result, but got {type(result)}" + assert "Refined output" in result, f"Expected 'Refined output' in result, but got: {result}" + + # Check the number of task exchanges + assert ( + len(orchestrator.state.task_exchanges) == 5 + ), f"Expected 5 task exchanges, but got {len(orchestrator.state.task_exchanges)}" + + # Check the roles of task exchanges + expected_roles = [ + "user", + "main_assistant", + "sub_assistant", + "main_assistant", + "refiner_assistant", + ] + actual_roles = [exchange.role for exchange in orchestrator.state.task_exchanges] + assert ( + actual_roles == expected_roles + ), f"Expected roles {expected_roles}, but got {actual_roles}" + + # Check the number of tasks + assert ( + len(orchestrator.state.tasks) == 1 + ), f"Expected 1 task, but got {len(orchestrator.state.tasks)}" + + # Print task exchanges for debugging + print("\nTask Exchanges:") + for i, exchange in enumerate(orchestrator.state.task_exchanges): + print(f"{i+1}. Role: {exchange.role}, Content: {exchange.content[:50]}...") + + # Check mock calls + assert ( + mock_get_full_response.call_count == 4 + ), f"Expected 4 calls to get_full_response, but got {mock_get_full_response.call_count}" + assert ( + mock_create_assistant.call_count == 4 + ), f"Expected 4 calls to create_assistant, but got {mock_create_assistant.call_count}" + + # Check if exchange log was created + assert os.path.exists( + os.path.join(orchestrator.output_dir, "exchange_log.md") + ), "Exchange log file was not created" + + +@patch("builtins.open", new_callable=MagicMock) +@patch("os.path.join", return_value="mocked_path") +def test_save_exchange_log(mock_join, mock_open, orchestrator): + mock_file = MagicMock() + mock_open.return_value.__enter__.return_value = mock_file + + orchestrator.state.task_exchanges = [ + TaskExchange(role="user", content="Test objective"), + TaskExchange(role="main_assistant", content="Test response"), + ] + + orchestrator._save_exchange_log("Test objective", "Test output") + + mock_open.assert_called_once_with("mocked_path", "w") + mock_file.write.assert_called() + + +@patch("src.orchestrator.get_full_response") +def test_run_workflow_error(mock_get_full_response, orchestrator): + mock_get_full_response.side_effect = Exception("API Error") + + with pytest.raises(Exception): + orchestrator.run_workflow("Test objective") + + +def test_task_exchange_validation(): + with pytest.raises(ValueError): + TaskExchange(role="invalid_role", content="Test content") + + +def test_state_to_dict(orchestrator): + orchestrator.state.task_exchanges.append(TaskExchange(role="user", content="Test")) + orchestrator.state.tasks.append(Task(task="Test task", result="Test result")) + + state_dict = orchestrator.state.to_dict() + assert len(state_dict["task_exchanges"]) == 1 + assert len(state_dict["tasks"]) == 1 diff --git a/tox.ini b/tox.ini new file mode 100644 index 0000000..2437f34 --- /dev/null +++ b/tox.ini @@ -0,0 +1,16 @@ +[tox] +envlist = py37, py38, py39 +isolated_build = true + +[testenv] +deps = + pytest + pytest-mock +commands = + pytest + +[pytest] +testpaths = tests + +[pylint] +max-line-length = 120 \ No newline at end of file