refine wren-ai-service logging and setups (#486)

* refine logging

* use the latest tag

* fix setup instructions

* simplify dev process

* add azure_deployment

* update
This commit is contained in:
Chih-Yu Yeh
2024-07-08 16:35:28 +08:00
committed by GitHub
parent 257be9ec47
commit 7047dce592
15 changed files with 45 additions and 42 deletions
-2
View File
@@ -1,7 +1,5 @@
# wren-ai-service
wren-ai-service/.env
wren-ai-service/.env.*
!wren-ai-service/.env.example
!wren-ai-service/.env.*.example
!wren-ai-service/src/eval/wren-engine/.env
wren-ai-service/src/eval/wren-engine/**/config.properties
-2
View File
@@ -1,2 +0,0 @@
# if not specified, the default ENV is prod
ENV=dev
-3
View File
@@ -42,9 +42,6 @@ prepare-wren-engine:
mkdir -p tools/dev/etc/mdl
echo "{\"catalog\": \"test_catalog\", \"schema\": \"test_schema\", \"models\": []}" \\
> tools/dev/etc/mdl/sample.json
psql:
docker exec -it wren-engine-1 bash launch-cli.sh
## wren-engine related ##
+1 -1
View File
@@ -9,7 +9,7 @@ Please read the [documentation](https://docs.getwren.ai/concept/wren_ai_service)
- Python 3.12.*, recommended to use [`pyenv`](https://github.com/pyenv/pyenv?tab=readme-ov-file#installation) to manage the Python versions
- install `poetry` with version 1.8.3: `curl -sSL https://install.python-poetry.org | python3 - --version 1.8.3`
- execute `poetry install` to install the dependencies
- copy `.env.example` file to `.env`, and `.env.dev.example` file to `.env.dev` and fill in the environment variables
- copy `.env.dev.example` file to `.env.dev` and fill in the environment variables
- [for development] execute `poetry run pre-commit install` to install the pre-commit hooks and `poetry run pre-commit run --all-files` to run the pre-commit checks at the first time to check if everything is set up correctly
- [for development] to run the tests, execute `make test`
@@ -1,3 +1,4 @@
import logging
import os
from typing import Any, Dict, List, Optional
@@ -19,6 +20,8 @@ from qdrant_client.http import models as rest
from src.core.provider import DocumentStoreProvider
from src.providers.loader import get_default_embedding_model_dim, provider
logger = logging.getLogger("wren-ai-service")
class AsyncQdrantDocumentStore(QdrantDocumentStore):
def __init__(
@@ -209,6 +212,10 @@ class QdrantProvider(DocumentStoreProvider):
dataset_name: Optional[str] = None,
recreate_index: bool = False,
):
logger.info(
f"Using Qdrant Document Store with Embedding Model Dimension: {embedding_model_dim}"
)
return AsyncQdrantDocumentStore(
location=self._location,
embedding_dim=embedding_model_dim,
@@ -119,6 +119,7 @@ class AsyncDocumentEmbedder(AzureOpenAIDocumentEmbedder):
self.client = AsyncAzureOpenAI(
azure_endpoint=api_base_url,
azure_deployment=model,
api_version=api_version,
api_key=api_key.resolve_value(),
)
@@ -200,6 +201,8 @@ class AzureOpenAIEmbedderProvider(EmbedderProvider):
or EMBEDDING_MODEL_DIMENSION,
):
logger.info(f"Using Azure OpenAI Embedding Model: {embedding_model}")
logger.info(f"Using Azure OpenAI Embedding API Base: {embed_api_base}")
logger.info(f"Using Azure OpenAI Embedding API Version: {embed_api_version}")
self._embedding_api_base = embed_api_base
self._embedding_api_key = embed_api_key
@@ -164,6 +164,8 @@ class OllamaEmbedderProvider(EmbedderProvider):
embedding_model: str = os.getenv("EMBEDDING_MODEL") or EMBEDDING_MODEL,
):
logger.info(f"Using Ollama Embedding Model: {embedding_model}")
logger.info(f"Using Ollama URL: {url}")
self._url = url
self._embedding_model = embedding_model
@@ -50,6 +50,7 @@ class AsyncGenerator(AzureOpenAIGenerator):
self.client = AsyncAzureOpenAI(
azure_endpoint=api_base,
azure_deployment=model,
api_version=api_version,
api_key=api_key.resolve_value(),
)
@@ -121,7 +122,9 @@ class AzureOpenAILLMProvider(LLMProvider):
chat_api_version: str = os.getenv("LLM_AZURE_OPENAI_VERSION"),
generation_model: str = os.getenv("GENERATION_MODEL") or GENERATION_MODEL,
):
logger.info(f"Using Azure OpenAI LLM: {generation_model}")
logger.info(f"Using AzureOpenAI LLM: {generation_model}")
logger.info(f"Using AzureOpenAI LLM with API base: {chat_api_base}")
logger.info(f"Using AzureOpenAI LLM with API version: {chat_api_version}")
self._generation_api_key = chat_api_key
self._generation_api_base = chat_api_base
@@ -130,10 +133,11 @@ class AzureOpenAILLMProvider(LLMProvider):
def get_generator(
self,
model_kwargs: Dict[str, Any] = orjson.loads(
os.getenv("GENERATION_MODEL_KWARGS", "{}")
)
or GENERATION_MODEL_KWARGS,
model_kwargs: Dict[str, Any] = (
orjson.loads(os.getenv("GENERATION_MODEL_KWARGS"))
if os.getenv("GENERATION_MODEL_KWARGS")
else GENERATION_MODEL_KWARGS
),
system_prompt: Optional[str] = None,
):
logger.info(
+6 -4
View File
@@ -127,15 +127,17 @@ class OllamaLLMProvider(LLMProvider):
generation_model: str = os.getenv("GENERATION_MODEL") or GENERATION_MODEL,
):
logger.info(f"Using Ollama LLM: {generation_model}")
logger.info(f"Using Ollama URL: {url}")
self._url = url
self._generation_model = generation_model
def get_generator(
self,
model_kwargs: Dict[str, Any] = orjson.loads(
os.getenv("GENERATION_MODEL_KWARGS", "{}")
)
or GENERATION_MODEL_KWARGS,
model_kwargs: Dict[str, Any] = (
orjson.loads(os.getenv("GENERATION_MODEL_KWARGS"))
if os.getenv("GENERATION_MODEL_KWARGS")
else GENERATION_MODEL_KWARGS
),
system_prompt: Optional[str] = None,
):
logger.info(f"Creating Ollama generator with model kwargs: {model_kwargs}")
+6 -5
View File
@@ -127,7 +127,7 @@ class OpenAILLMProvider(LLMProvider):
"""
OpenAI(api_key=api_key, base_url=api_base).models.list()
logger.info(f"Initializing OpenAILLM provider with API base: {api_base}")
logger.info(f"Using OpenAILLM provider with API base: {api_base}")
# TODO: currently only OpenAI api key can be verified
if api_base == LLM_OPENAI_API_BASE:
_verify_api_key(api_key.resolve_value(), api_base)
@@ -140,10 +140,11 @@ class OpenAILLMProvider(LLMProvider):
def get_generator(
self,
model_kwargs: Dict[str, Any] = orjson.loads(
os.getenv("GENERATION_MODEL_KWARGS", "{}")
)
or GENERATION_MODEL_KWARGS,
model_kwargs: Dict[str, Any] = (
orjson.loads(os.getenv("GENERATION_MODEL_KWARGS"))
if os.getenv("GENERATION_MODEL_KWARGS")
else GENERATION_MODEL_KWARGS
),
system_prompt: Optional[str] = None,
):
if self._api_base == LLM_OPENAI_API_BASE:
+4 -4
View File
@@ -3,6 +3,7 @@ import functools
import logging
import os
import time
from pathlib import Path
from typing import Tuple
from dotenv import load_dotenv
@@ -49,12 +50,11 @@ def setup_custom_logger(name, level=logging.INFO):
def load_env_vars() -> str:
load_dotenv(override=True)
if is_dev_env := os.getenv("ENV") and os.getenv("ENV").lower() == "dev":
if Path(".env.dev").exists():
load_dotenv(".env.dev", override=True)
return "dev"
return "dev" if is_dev_env else "prod"
return "prod"
def init_providers() -> (
+5 -5
View File
@@ -11,11 +11,11 @@ IBIS_SERVER_PORT=8000
# version
# CHANGE THIS TO THE LATEST VERSION
WREN_PRODUCT_VERSION=development
WREN_ENGINE_VERSION=0.5.0
WREN_AI_SERVICE_VERSION=0.5.6
WREN_UI_VERSION=0.7.2
IBIS_SERVER_VERSION=0.5.0
WREN_BOOTSTRAP_VERSION=0.1.4
WREN_ENGINE_VERSION=latest
WREN_AI_SERVICE_VERSION=latest
WREN_UI_VERSION=latest
IBIS_SERVER_VERSION=latest
WREN_BOOTSTRAP_VERSION=latest
# SQL Protocol
PG_USERNAME=wren-user-placeholder
@@ -1,12 +1,6 @@
node.environment=production
pg-wire-protocol.auth.file=
pg-wire-protocol.enabled=true
wren.directory=/usr/src/app/etc/mdl
wren.experimental-enable-dynamic-fields=false
wren.datasource.type=bigquery
bigquery.project-id=wrenai
bigquery.credentials-key=
bigquery.location=asia-east1
bigquery.bucket-name=
wren.experimental-enable-dynamic-fields=true
wren.datasource.type=duckdb
duckdb.storage.access-key=
duckdb.storage.secret-key=
@@ -17,7 +17,6 @@ services:
- ${WREN_ENGINE_PORT}:${WREN_ENGINE_PORT}
volumes:
- ./etc:/usr/src/app/etc
- ${LAUNCH_CLI_PATH}:/usr/src/app/launch-cli.sh
networks:
- wren
-2
View File
@@ -1,2 +0,0 @@
#!/bin/bash
psql postgres://localhost:7432/wrenai?options=--search_path%3Dspider