diff --git a/.github/workflows/pre-release-CI.yml b/.github/workflows/pre-release-CI.yml new file mode 100644 index 00000000..50df3692 --- /dev/null +++ b/.github/workflows/pre-release-CI.yml @@ -0,0 +1,60 @@ +name: Pre Release CI + +on: + workflow_dispatch: + +concurrency: + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }} + cancel-in-progress: true + +jobs:git + build-and-test: + name: Build & Test on ${{ matrix.os }}-py${{ matrix.python-version }} + runs-on: ${{ matrix.os }} + strategy: + matrix: + os: [ubuntu-latest, windows-latest, macos-latest] + python-version: [3.9, '3.10', 3.11] + defaults: + run: + shell: bash + + steps: + - uses: actions/checkout@v3 + + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v4 + with: + python-version: ${{ matrix.python-version }} + + - name: Install Poetry + uses: snok/install-poetry@v1 + with: + version: 1.3.2 + + - name: Build wheel + run: | + poetry build + + - name: Install the wheel + run: | + pip install dist/pinecone_resin*.whl + + - name: Create dev requirements file + run: | + poetry export -f requirements.txt --without-hashes --only dev -o only-dev.txt + + - name: Install dev requirements + run: | + pip install -r only-dev.txt + + - name: Run tests + run: pytest --html=report.html --self-contained-html tests/unit + + - name: Upload pytest reports + if: always() + uses: actions/upload-artifact@v3 + with: + name: pytest-report-${{ matrix.os }}-py${{ matrix.python-version }} + path: .pytest_cache + diff --git a/pyproject.toml b/pyproject.toml index e78b485b..7832e415 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -24,6 +24,8 @@ tenacity = "^8.2.1" sse-starlette = "^1.6.5" types-tqdm = "^4.61.0" tqdm = "^4.66.1" +gunicorn = "^21.2.0" +types-pyyaml = "^6.0.12.12" [tool.poetry.group.dev.dependencies] diff --git a/src/resin_cli/app.py b/src/resin_cli/app.py index bb84a229..765e3070 100644 --- a/src/resin_cli/app.py +++ b/src/resin_cli/app.py @@ -6,11 +6,13 @@ import openai from multiprocessing import current_process + +import yaml from dotenv import load_dotenv from resin.llm import BaseLLM from resin.llm.models import UserMessage -from resin.tokenizer import OpenAITokenizer, Tokenizer +from resin.tokenizer import Tokenizer from resin.knowledge_base import KnowledgeBase from resin.context_engine import ContextEngine from resin.chat_engine import ChatEngine @@ -28,6 +30,7 @@ ContextUpsertRequest, HealthStatus, ContextDeleteRequest from resin.llm.openai import OpenAILLM +from resin_cli.errors import ConfigError load_dotenv() # load env vars before import of openai openai.api_key = os.getenv("OPENAI_API_KEY") @@ -198,24 +201,66 @@ def _init_logging(): def _init_engines(): - global kb, context_engine, chat_engine, llm - Tokenizer.initialize(OpenAITokenizer, model_name='gpt-3.5-turbo-0613') + global kb, context_engine, chat_engine, llm, logger - INDEX_NAME = os.getenv("INDEX_NAME") - if not INDEX_NAME: + index_name = os.getenv("INDEX_NAME") + if not index_name: raise ValueError("INDEX_NAME environment variable must be set") - kb = KnowledgeBase(index_name=INDEX_NAME) - context_engine = ContextEngine(knowledge_base=kb) - llm = OpenAILLM() - chat_engine = ChatEngine(context_engine=context_engine, llm=llm) + config_file = os.getenv("RESIN_CONFIG_FILE") + if config_file: + _load_config(config_file) + + else: + logger.info("Did not find config file. Initializing engines with default " + "configuration") + Tokenizer.initialize() + kb = KnowledgeBase(index_name=index_name) + context_engine = ContextEngine(knowledge_base=kb) + llm = OpenAILLM() + chat_engine = ChatEngine(context_engine=context_engine, llm=llm) kb.connect() -def start(host="0.0.0.0", port=8000, reload=False, workers=1): - uvicorn.run("resin_cli.app:app", - host=host, port=port, reload=reload, workers=workers) +def _load_config(config_file): + global chat_engine, llm, context_engine, kb, logger + logger.info(f"Initializing engines with config file {config_file}") + try: + with open(config_file, "r") as f: + config = yaml.safe_load(f) + except Exception as e: + logger.exception(f"Failed to load config file {config_file}") + raise ConfigError( + f"Failed to load config file {config_file}. Error: {str(e)}" + ) + tokenizer_config = config.get("tokenizer", {}) + Tokenizer.initialize_from_config(tokenizer_config) + if "chat_engine" not in config: + raise ConfigError( + f"Config file {config_file} must contain a 'chat_engine' section" + ) + chat_engine_config = config["chat_engine"] + try: + chat_engine = ChatEngine.from_config(chat_engine_config) + except Exception as e: + logger.exception( + f"Failed to initialize chat engine from config file {config_file}" + ) + raise ConfigError( + f"Failed to initialize chat engine from config file {config_file}." + f" Error: {str(e)}" + ) + llm = chat_engine.llm + context_engine = chat_engine.context_engine + kb = context_engine.knowledge_base + + +def start(host="0.0.0.0", port=8000, reload=False, config_file=None): + if config_file: + os.environ["RESIN_CONFIG_FILE"] = config_file + + uvicorn.run("resin_cli.app:app", host=host, port=port, reload=reload, workers=0) if __name__ == "__main__": diff --git a/src/resin_cli/cli.py b/src/resin_cli/cli.py index be2bd2cb..480307d6 100644 --- a/src/resin_cli/cli.py +++ b/src/resin_cli/cli.py @@ -1,10 +1,13 @@ import os -from typing import List, Optional +import signal +import subprocess +from typing import Dict, Any, Optional, List import click import time import requests +import yaml from dotenv import load_dotenv from tenacity import retry, stop_after_attempt, wait_fixed from tqdm import tqdm @@ -19,9 +22,9 @@ from resin.tokenizer import Tokenizer from resin_cli.data_loader import ( load_from_path, - CLIError, IDsNotUniqueError, DocumentsValidationError) +from resin_cli.errors import CLIError from resin import __version__ @@ -30,8 +33,7 @@ from .api_models import ChatDebugInfo -dotenv_path = os.path.join(os.path.dirname(__file__), ".env") -load_dotenv(dotenv_path) +load_dotenv() if os.getenv("OPENAI_API_KEY"): openai.api_key = os.getenv("OPENAI_API_KEY") @@ -99,6 +101,34 @@ def _initialize_tokenizer(): raise CLIError(msg) +def _load_kb_config(config_file: Optional[str]) -> Dict[str, Any]: + if config_file is None: + return {} + + try: + with open(os.path.join("config", config_file), 'r') as f: + config = yaml.safe_load(f) + except Exception as e: + msg = f"Failed to load config file {config_file}. Reason:\n{e}" + raise CLIError(msg) + + if "knowledge_base" in config: + kb_config = config.get("knowledge_base", None) + elif "chat_engine" in config: + kb_config = config["chat_engine"]\ + .get("context_engine", {})\ + .get("knowledge_base", None) + else: + kb_config = None + + if kb_config is None: + msg = (f"Did not find a `knowledge_base` configuration in {config_file}, " + "Would you like to use the default configuration?") + click.confirm(click.style(msg, fg="red"), abort=True) + kb_config = {} + return kb_config + + @click.group(invoke_without_command=True, context_settings=CONTEXT_SETTINGS) @click.version_option(__version__, "-v", "--version", prog_name="Resin") @click.pass_context @@ -138,9 +168,13 @@ def health(url): ) ) @click.argument("index-name", nargs=1, envvar="INDEX_NAME", type=str, required=True) -def new(index_name): +@click.option("--config", "-c", default=None, + help="Path to a resin config file. Optional, otherwise configuration " + "defaults will be used.") +def new(index_name: str, config: Optional[str]): _initialize_tokenizer() - kb = KnowledgeBase(index_name=index_name) + kb_config = _load_kb_config(config) + kb = KnowledgeBase.from_config(kb_config, index_name=index_name) click.echo("Resin is going to create a new index: ", nl=False) click.echo(click.style(f"{kb.index_name}", fg="green")) click.confirm(click.style("Do you want to continue?", fg="red"), abort=True) @@ -180,7 +214,14 @@ def new(index_name): "be uploaded. " "When set to True, the upsert process will continue on failure, as " "long as less than 10% of the documents have failed to be uploaded.") -def upsert(index_name: str, data_path: str, batch_size: int, allow_failures: bool): +@click.option("--config", "-c", default=None, + help="Path to a resin config file. Optional, otherwise configuration " + "defaults will be used.") +def upsert(index_name: str, + data_path: str, + batch_size: int, + allow_failures: bool, + config: Optional[str]): if index_name is None: msg = ( "No index name provided. Please set --index-name or INDEX_NAME environment " @@ -190,7 +231,8 @@ def upsert(index_name: str, data_path: str, batch_size: int, allow_failures: boo _initialize_tokenizer() - kb = KnowledgeBase(index_name=index_name) + kb_config = _load_kb_config(config) + kb = KnowledgeBase.from_config(kb_config, index_name=index_name) try: kb.connect() except RuntimeError as e: @@ -240,7 +282,7 @@ def upsert(index_name: str, data_path: str, batch_size: int, allow_failures: boo for i in range(0, len(data), batch_size): batch = data[i:i + batch_size] try: - kb.upsert(data) + kb.upsert(batch) except Exception as e: if allow_failures and len(failed_docs) < len(data) // 10: failed_docs.extend([_.id for _ in batch]) @@ -428,10 +470,23 @@ def chat(chat_service_url, baseline, debug, stream): help="TCP port to bind the server to. Defaults to 8000") @click.option("--reload/--no-reload", default=False, help="Set the server to reload on code changes. Defaults to False") -@click.option("--workers", default=1, help="Number of worker processes. Defaults to 1") -def start(host, port, reload, workers): +@click.option("--config", "-c", default=None, + help="Path to a resin config file. Optional, otherwise configuration " + "defaults will be used.") +def start(host: str, port: str, reload: bool, config: Optional[str]): + note_msg = ( + "🚨 Note 🚨\n" + "For debugging only. To run the Resin service in production, run the command:\n" + "gunicorn resin_cli.app:app --worker-class uvicorn.workers.UvicornWorker " + f"--bind {host}:{port} --workers " + ) + for c in note_msg: + click.echo(click.style(c, fg="red"), nl=False) + time.sleep(0.01) + click.echo() + click.echo(f"Starting Resin service on {host}:{port}") - start_service(host, port=port, reload=reload, workers=workers) + start_service(host, port=port, reload=reload, config_file=config) @cli.command( @@ -446,6 +501,28 @@ def start(host, port, reload, workers): @click.option("url", "--url", default="http://0.0.0.0:8000", help="URL of the Resin service to use. Defaults to http://0.0.0.0:8000") def stop(url): + # Check if the service was started using Gunicorn + res = subprocess.run(["pgrep", "-f", "gunicorn resin_cli.app:app"], + capture_output=True) + output = res.stdout.decode("utf-8").split() + + # If Gunicorn was used, kill all Gunicorn processes + if output: + msg = ("It seems that Resin service was launched using Gunicorn.\n" + "Do you want to kill all Gunicorn processes?") + click.confirm(click.style(msg, fg="red"), abort=True) + try: + subprocess.run(["pkill", "-f", "gunicorn resin_cli.app:app"], check=True) + except subprocess.CalledProcessError: + try: + [os.kill(int(pid), signal.SIGINT) for pid in output] + except OSError: + msg = ( + "Could not kill Gunicorn processes. Please kill them manually." + f"Found process ids: {output}" + ) + raise CLIError(msg) + try: res = requests.get(urljoin(url, "/shutdown")) res.raise_for_status() diff --git a/src/resin_cli/data_loader/__init__.py b/src/resin_cli/data_loader/__init__.py index 85464408..8a298030 100644 --- a/src/resin_cli/data_loader/__init__.py +++ b/src/resin_cli/data_loader/__init__.py @@ -1,6 +1,5 @@ from .data_loader import ( load_from_path, - CLIError, IDsNotUniqueError, DocumentsValidationError ) diff --git a/src/resin_cli/data_loader/data_loader.py b/src/resin_cli/data_loader/data_loader.py index b84434b6..ee2d921a 100644 --- a/src/resin_cli/data_loader/data_loader.py +++ b/src/resin_cli/data_loader/data_loader.py @@ -5,10 +5,8 @@ from typing import List from textwrap import dedent -import click import numpy as np import pandas as pd -from click import ClickException from pydantic import ValidationError @@ -27,11 +25,6 @@ def format_multiline(msg): return dedent(msg).strip() -class CLIError(ClickException): - def format_message(self) -> str: - return click.style(format_multiline(self.message), fg='red') - - def _process_metadata(value): if pd.isna(value): return {} diff --git a/src/resin_cli/errors.py b/src/resin_cli/errors.py new file mode 100644 index 00000000..1f54ed7e --- /dev/null +++ b/src/resin_cli/errors.py @@ -0,0 +1,13 @@ +import click +from click import ClickException + +from resin_cli.data_loader.data_loader import format_multiline + + +class CLIError(ClickException): + def format_message(self) -> str: + return click.style(format_multiline(self.message), fg='red') + + +class ConfigError(RuntimeError): + pass