diff --git a/.github/workflows/ai-service-ci.yaml b/.github/workflows/ai-service-ci.yaml index 28cf11c6c..db6ff4d47 100644 --- a/.github/workflows/ai-service-ci.yaml +++ b/.github/workflows/ai-service-ci.yaml @@ -56,6 +56,8 @@ jobs: docker stop wren-engine && docker rm wren-engine env: ENV: dev + LLM_PROVIDER: openai + DOCUMENT_STORE_PROVIDER: qdrant OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }} OPENAI_GENERATION_MODEL: gpt-3.5-turbo WREN_ENGINE_ENDPOINT: http://localhost:8080 diff --git a/wren-ai-service/.env.dev.example b/wren-ai-service/.env.dev.example index 85cb7113e..38fd18373 100644 --- a/wren-ai-service/.env.dev.example +++ b/wren-ai-service/.env.dev.example @@ -4,9 +4,9 @@ WREN_AI_SERVICE_PORT=5556 # app related # LLM Provider name should be mapped to Haystack's supported LLM providers: https://docs.haystack.deepset.ai/v2.0/docs/generators -LLM_PROVIDER=OpenAI +LLM_PROVIDER=openai # Document Store Provider name should be mapped to Haystack's supported Document Store providers: see the haystack documentation's left sidebar -DOCUMENT_STORE_PROVIDER=Qdrant +DOCUMENT_STORE_PROVIDER=qdrant # llm provider specific env variables names must be in the format of [LLM_PROVIDER]_[ENV_VARIABLE_NAME] and must be in uppercase OPENAI_API_KEY= @@ -18,4 +18,4 @@ QDRANT_HOST=localhost WREN_ENGINE_ENDPOINT=http://localhost:8080 # evaluation related -DATASET_NAME=book_2 \ No newline at end of file +DATASET_NAME=book_2 diff --git a/wren-ai-service/.env.prod.example b/wren-ai-service/.env.prod.example index 65348183d..2417c0837 100644 --- a/wren-ai-service/.env.prod.example +++ b/wren-ai-service/.env.prod.example @@ -1,6 +1,6 @@ # docker related # This will determine the prefix of container name. -# If the service name in docker-compose.yaml is ai-service, +# If the service name in docker-compose.yaml is ai-service, # then the container name will be wren-ai-service-1. # If COMPOSE_PROJECT_NAME is not set, the default prefix name is docker. COMPOSE_PROJECT_NAME=wren @@ -10,9 +10,9 @@ WREN_AI_SERVICE_PORT=5555 # app related # LLM Provider name should be mapped to Haystack's supported LLM providers: https://docs.haystack.deepset.ai/v2.0/docs/generators -LLM_PROVIDER=OpenAI +LLM_PROVIDER=openai # Document Store Provider name should be mapped to Haystack's supported Document Store providers: see the haystack documentation's left sidebar -DOCUMENT_STORE_PROVIDER=Qdrant +DOCUMENT_STORE_PROVIDER=qdrant # llm provider specific env variables names must be in the format of [LLM_PROVIDER]_[ENV_VARIABLE_NAME] and must be in uppercase OPENAI_API_KEY= diff --git a/wren-ai-service/src/core/document_store_provider.py b/wren-ai-service/src/core/document_store_provider.py deleted file mode 100644 index 437dc82ff..000000000 --- a/wren-ai-service/src/core/document_store_provider.py +++ /dev/null @@ -1,13 +0,0 @@ -from abc import ABCMeta, abstractmethod - -from haystack.document_stores.types import DocumentStore - - -class DocumentStoreProvider(metaclass=ABCMeta): - @abstractmethod - def get_store(self, *args, **kwargs) -> DocumentStore: - ... - - @abstractmethod - def get_retriever(self, *args, **kwargs): - ... diff --git a/wren-ai-service/src/core/llm_provider.py b/wren-ai-service/src/core/provider.py similarity index 54% rename from wren-ai-service/src/core/llm_provider.py rename to wren-ai-service/src/core/provider.py index da9439338..cd24432c5 100644 --- a/wren-ai-service/src/core/llm_provider.py +++ b/wren-ai-service/src/core/provider.py @@ -1,5 +1,7 @@ from abc import ABCMeta, abstractmethod +from haystack.document_stores.types import DocumentStore + class LLMProvider(metaclass=ABCMeta): @abstractmethod @@ -13,3 +15,13 @@ class LLMProvider(metaclass=ABCMeta): @abstractmethod def get_document_embedder(self, *args, **kwargs): ... + + +class DocumentStoreProvider(metaclass=ABCMeta): + @abstractmethod + def get_store(self, *args, **kwargs) -> DocumentStore: + ... + + @abstractmethod + def get_retriever(self, *args, **kwargs): + ... diff --git a/wren-ai-service/src/globals.py b/wren-ai-service/src/globals.py index 60789b780..98c240fdb 100644 --- a/wren-ai-service/src/globals.py +++ b/wren-ai-service/src/globals.py @@ -1,7 +1,6 @@ from typing import Callable, Tuple -from src.core.document_store_provider import DocumentStoreProvider -from src.core.llm_provider import LLMProvider +from src.core.provider import DocumentStoreProvider, LLMProvider from src.pipelines.ask import ( followup_generation_pipeline as ask_followup_generation_pipeline, ) diff --git a/wren-ai-service/src/pipelines/ask/historical_question.py b/wren-ai-service/src/pipelines/ask/historical_question.py index 59776241b..bb52abdb5 100644 --- a/wren-ai-service/src/pipelines/ask/historical_question.py +++ b/wren-ai-service/src/pipelines/ask/historical_question.py @@ -4,9 +4,8 @@ from typing import Dict, List, Optional from haystack import Document, Pipeline, component -from src.core.document_store_provider import DocumentStoreProvider -from src.core.llm_provider import LLMProvider from src.core.pipeline import BasicPipeline +from src.core.provider import DocumentStoreProvider, LLMProvider from src.utils import ( init_providers, load_env_vars, diff --git a/wren-ai-service/src/pipelines/ask/indexing_pipeline.py b/wren-ai-service/src/pipelines/ask/indexing_pipeline.py index a218fd872..3b763516b 100644 --- a/wren-ai-service/src/pipelines/ask/indexing_pipeline.py +++ b/wren-ai-service/src/pipelines/ask/indexing_pipeline.py @@ -8,8 +8,8 @@ from haystack.components.writers import DocumentWriter from haystack.document_stores.types import DocumentStore, DuplicatePolicy from tqdm import tqdm -from src.core.document_store_provider import DocumentStoreProvider -from src.core.llm_provider import LLMProvider +from src.core.provider import LLMProvider +from src.core.provider import DocumentStoreProvider from src.core.pipeline import BasicPipeline from src.utils import init_providers, load_env_vars diff --git a/wren-ai-service/src/providers/__init__.py b/wren-ai-service/src/providers/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/wren-ai-service/src/providers/document_store/__init__.py b/wren-ai-service/src/providers/document_store/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/wren-ai-service/src/providers/document_store/qdrant.py b/wren-ai-service/src/providers/document_store/qdrant.py index 5bde550ca..a921e6c96 100644 --- a/wren-ai-service/src/providers/document_store/qdrant.py +++ b/wren-ai-service/src/providers/document_store/qdrant.py @@ -1,14 +1,17 @@ +import os from typing import Optional from haystack_integrations.components.retrievers.qdrant import QdrantEmbeddingRetriever from haystack_integrations.document_stores.qdrant import QdrantDocumentStore -from src.core.document_store_provider import DocumentStoreProvider +from src.core.provider import DocumentStoreProvider from src.providers.llm.openai import EMBEDDING_MODEL_DIMENSION +from src.providers.loader import provider +@provider("qdrant") class QdrantProvider(DocumentStoreProvider): - def __init__(self, location: str): + def __init__(self, location: str = os.getenv("QDRANT_HOST")): self._location = location def get_store( diff --git a/wren-ai-service/src/providers/llm/__init__.py b/wren-ai-service/src/providers/llm/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/wren-ai-service/src/providers/llm/openai.py b/wren-ai-service/src/providers/llm/openai.py index df0d07945..67e6ba769 100644 --- a/wren-ai-service/src/providers/llm/openai.py +++ b/wren-ai-service/src/providers/llm/openai.py @@ -1,4 +1,5 @@ import logging +import os from typing import Any, Dict, List, Optional import backoff @@ -7,8 +8,10 @@ from haystack import component from haystack.components.embedders import OpenAIDocumentEmbedder, OpenAITextEmbedder from haystack.components.generators import OpenAIGenerator from haystack.utils.auth import Secret +from openai import OpenAI -from src.core.llm_provider import LLMProvider +from src.core.provider import LLMProvider +from src.providers.loader import provider logger = logging.getLogger("wren-ai-service") @@ -34,8 +37,22 @@ class CustomOpenAIGenerator(OpenAIGenerator): ) +@provider("openai") class OpenAILLMProvider(LLMProvider): - def __init__(self, api_key: Secret, generation_model: str = GENERATION_MODEL_NAME): + def __init__( + self, + api_key: Secret = Secret.from_env_var("OPENAI_API_KEY"), + generation_model: str = os.getenv("OPENAI_GENERATION_MODEL") + or GENERATION_MODEL_NAME, + ): + def _verify_api_key(api_key: str) -> None: + """ + this is a temporary solution to verify that the required environment variables are set + """ + OpenAI(api_key=api_key).models.list() + + _verify_api_key(api_key.resolve_value()) + logger.info(f"Using OpenAI Generation Model: {generation_model}") self._api_key = api_key self._generation_model = generation_model diff --git a/wren-ai-service/src/providers/loader.py b/wren-ai-service/src/providers/loader.py new file mode 100644 index 000000000..734a72ccd --- /dev/null +++ b/wren-ai-service/src/providers/loader.py @@ -0,0 +1,93 @@ +import importlib +import logging +import pkgutil + +logger = logging.getLogger("wren-ai-service") + + +PROVIDERS_PATH = "src.providers" +PROVIDERS = {} + + +def import_mods(package_name=PROVIDERS_PATH): + """ + This function is designed to import all submodules within a given package, + including those in nested packages. + + The function works by first importing the initial package using its name. + It then iterates through the package directory and imports each submodule it encounters. + This includes submodules in any nested packages. + + This is particularly useful in scenarios where you want to ensure that all submodules are loaded + and ready for use, without having to manually import each one individually. + + Parameters: + package_name (str): The name of the initial package to import submodules from. + Defaults to the value of PROVIDERS_PATH. + + Returns: + None: This function doesn't return anything; it simply imports the submodules. + + Raises: + ModuleNotFoundError: If a submodule cannot be found. + """ + package = importlib.import_module(package_name) + + # Iterate through all the submodules in the package + for _, name, _ in pkgutil.walk_packages(package.__path__, package.__name__ + "."): + # Import each submodule + importlib.import_module(name) + logger.debug(f"Imported Provider: {name}") + + +def provider(name: str): + """ + This decorator is designed to register a provider class and add it to the PROVIDERS dictionary. + + The decorator works by taking a name as an argument, which is used as a key in the PROVIDERS dictionary. + The value associated with this key is the provider class that the decorator is applied to. + + This is particularly useful in scenarios where you want to maintain a registry of provider classes, + and you want to be able to access these classes using a unique name. + + Parameters: + name (str): The name of the provider. This is used as a key in the PROVIDERS dictionary. + + Returns: + function: The decorator function that registers the provider class. + + Raises: + KeyError: If a provider with the same name is already registered in the PROVIDERS dictionary. + """ + + def wrapper(cls): + logger.info(f"Registering provider: {name}") + PROVIDERS[name] = cls + return cls + + return wrapper + + +def get_provider(name: str): + """ + This function is designed to return a provider class by its name. + + The function works by taking a name as an argument, which is used to look up + the corresponding provider class in a dictionary of registered providers. + + This is particularly useful in scenarios where you have a registry of provider classes + and you want to be able to retrieve a specific provider class using its unique name. + + Parameters: + name (str): The name of the provider. This is used as a key to look up + the provider class in the dictionary of registered providers. + + Returns: + type: The provider class corresponding to the given name. + + Raises: + KeyError: If a provider with the given name is not found in the dictionary of registered providers. + """ + logger.debug(f"Getting provider: {name} from {PROVIDERS}") + + return PROVIDERS[name] diff --git a/wren-ai-service/src/utils.py b/wren-ai-service/src/utils.py index b8be602b3..3a4a64fbc 100644 --- a/wren-ai-service/src/utils.py +++ b/wren-ai-service/src/utils.py @@ -6,13 +6,9 @@ from typing import Dict, List, Optional, Tuple import requests import sqlglot from dotenv import load_dotenv -from haystack.utils.auth import Secret -from openai import OpenAI -from src.core.document_store_provider import DocumentStoreProvider -from src.core.llm_provider import LLMProvider -from src.providers.document_store.qdrant import QdrantProvider -from src.providers.llm.openai import GENERATION_MODEL_NAME, OpenAILLMProvider +from src.core.provider import DocumentStoreProvider, LLMProvider +from src.providers import loader logger = logging.getLogger("wren-ai-service") @@ -67,12 +63,6 @@ def clean_generation_result(result: str) -> str: def load_env_vars() -> str: - def _verify_env_vars() -> None: - """ - this is a temporary solution to verify that the required environment variables are set - """ - OpenAI().models.list() - load_dotenv(override=True) if is_dev_env := os.getenv("ENV") and os.getenv("ENV").lower() == "dev": @@ -80,7 +70,6 @@ def load_env_vars() -> str: else: load_dotenv(".env.prod", override=True) - _verify_env_vars() return "dev" if is_dev_env else "prod" @@ -176,25 +165,12 @@ def remove_duplicates(dicts): return unique_dicts -# please do not change this function unless you are adjusting the OpenAILLMProvider or QdrantProvider -# if you would like to change the providers, please do the following steps: -# 1. create a new provider class that implements LLMProvider or DocumentStoreProvider in the src/providers directory -# 2. if you need to add new environment variables, please add them to the .env.dev and .env.prod files -# 3. add a new function to import your providers in this file -# 4. call the new function you created in the src/__main__.py in container.init_globals() def init_providers() -> Tuple[LLMProvider, DocumentStoreProvider]: load_env_vars() - generation_model = os.getenv("OPENAI_GENERATION_MODEL") or GENERATION_MODEL_NAME - logger.info("Initializing providers...") - logger.info(f"Using OpenAI Generation Model: {generation_model}") + loader.import_mods() - llm_provider = OpenAILLMProvider( - api_key=Secret.from_env_var("OPENAI_API_KEY"), - generation_model=generation_model, - ) - document_store_provider = QdrantProvider( - location=os.getenv("QDRANT_HOST"), - ) - return llm_provider, document_store_provider + llm_provider = loader.get_provider(os.getenv("LLM_PROVIDER")) + document_store_provider = loader.get_provider(os.getenv("DOCUMENT_STORE_PROVIDER")) + return llm_provider(), document_store_provider() diff --git a/wren-ai-service/tests/pipelines/test_ask.py b/wren-ai-service/tests/pipelines/test_ask.py index 170d2573e..eaeaa1f67 100644 --- a/wren-ai-service/tests/pipelines/test_ask.py +++ b/wren-ai-service/tests/pipelines/test_ask.py @@ -2,8 +2,7 @@ import json import pytest -from src.core.document_store_provider import DocumentStoreProvider -from src.core.llm_provider import LLMProvider +from src.core.provider import DocumentStoreProvider, LLMProvider from src.pipelines.ask.followup_generation_pipeline import FollowUpGeneration from src.pipelines.ask.generation_pipeline import Generation from src.pipelines.ask.indexing_pipeline import Indexing diff --git a/wren-ai-service/tests/providers/test_loader.py b/wren-ai-service/tests/providers/test_loader.py new file mode 100644 index 000000000..38448011a --- /dev/null +++ b/wren-ai-service/tests/providers/test_loader.py @@ -0,0 +1,16 @@ +from src.providers import loader + + +def test_import_mods(): + loader.import_mods("src.providers") + assert len(loader.PROVIDERS) == 2 + + +def test_get_provider(): + loader.import_mods("src.providers") + + provider = loader.get_provider("openai") + assert provider.__name__ == "OpenAILLMProvider" + + provider = loader.get_provider("qdrant") + assert provider.__name__ == "QdrantProvider"