opensource/WrenAI
3
0
Fork 0
mirror of https://github.com/Canner/WrenAI.git synced 2026-10-24 02:25:59 +08:00
Code Issues Packages Projects Releases Wiki Activity

minor update for ai-service (#493)

Browse Source 
* remove trailing / in urls of env vars

* update

* fix

* fix bug

* simplify codebase
This commit is contained in:
Chih-Yu Yeh
2024-07-09 12:04:28 +08:00
committed by GitHub
parent 7047dce592
commit 6194a1013f
8 changed files with 61 additions and 41 deletions
Download Patch File Download Diff File Expand all files Collapse all files
wren-ai-service/src/force_deploy.py
+4 -5
View File
@@ -3,11 +3,10 @@ import os
import aiohttp
import backoff
from dotenv import load_dotenv
load_dotenv(override=True)
if is_dev_env := os.getenv("ENV") and os.getenv("ENV").lower() == "dev":
load_dotenv(".env.dev", override=True)
from src.utils import load_env_vars
load_env_vars()
@backoff.on_exception(backoff.expo, aiohttp.ClientError, max_time=60, max_tries=3)
@@ -25,5 +24,5 @@ async def force_deploy():
print(f"Forcing deployment: {res}")
if os.getenv("ENGINE", "wren-ui") == "wren-ui":
if os.getenv("ENGINE", "wren_ui") == "wren_ui":
asyncio.run(force_deploy())
wren-ai-service/src/providers/embedder/azure_openai.py
+10 -5
View File
@@ -15,6 +15,7 @@ from tqdm import tqdm
from src.core.provider import EmbedderProvider
from src.providers.loader import provider
from src.utils import remove_trailing_slash
logger = logging.getLogger("wren-ai-service")
@@ -200,16 +201,20 @@ 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_base = remove_trailing_slash(embed_api_base)
self._embedding_api_key = embed_api_key
self._embedding_api_version = embed_api_version
self._embedding_model = embedding_model
self._embedding_model_dim = embedding_model_dim
logger.info(f"Using Azure OpenAI Embedding Model: {self._embedding_model}")
logger.info(
f"Using Azure OpenAI Embedding API Base: {self._embedding_api_base}"
)
logger.info(
f"Using Azure OpenAI Embedding API Version: {self._embedding_api_version}"
)
def get_text_embedder(self):
return AsyncTextEmbedder(
api_key=self._embedding_api_key,
wren-ai-service/src/providers/embedder/ollama.py
+5 -4
View File
@@ -13,6 +13,7 @@ from tqdm import tqdm
from src.core.provider import EmbedderProvider
from src.providers.loader import provider
from src.utils import remove_trailing_slash
logger = logging.getLogger("wren-ai-service")
@@ -163,12 +164,12 @@ class OllamaEmbedderProvider(EmbedderProvider):
url: str = os.getenv("EMBEDDER_OLLAMA_URL") or EMBEDDER_OLLAMA_URL,
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._url = remove_trailing_slash(url)
self._embedding_model = embedding_model
logger.info(f"Using Ollama Embedding Model: {self._embedding_model}")
logger.info(f"Using Ollama URL: {self._url}")
def get_text_embedder(
self,
model_kwargs: Optional[Dict[str, Any]] = None,
wren-ai-service/src/providers/embedder/openai.py
+14 -10
View File
@@ -12,6 +12,7 @@ from tqdm import tqdm
from src.core.provider import EmbedderProvider
from src.providers.loader import provider
from src.utils import remove_trailing_slash
logger = logging.getLogger("wren-ai-service")
@@ -195,20 +196,23 @@ class OpenAIEmbedderProvider(EmbedderProvider):
"""
OpenAI(api_key=api_key, base_url=api_base).models.list()
logger.info(f"Initializing OpenAIEmbedder provider with API base: {api_base}")
# TODO: currently only OpenAI api key can be verified
if api_base == EMBEDDER_OPENAI_API_BASE:
_verify_api_key(api_key.resolve_value(), api_base)
logger.info(f"Using OpenAI Embedding Model: {embedding_model}")
else:
logger.info(
f"Using OpenAI API-compatible Embedding Model: {embedding_model}"
)
self._api_key = api_key
self._api_base = api_base
self._api_base = remove_trailing_slash(api_base)
self._embedding_model = embedding_model
self._embedding_model_dim = embedding_model_dim
logger.info(
f"Initializing OpenAIEmbedder provider with API base: {self._api_base}"
)
# TODO: currently only OpenAI api key can be verified
if self._api_base == EMBEDDER_OPENAI_API_BASE:
_verify_api_key(self._api_key.resolve_value(), self._api_base)
logger.info(f"Using OpenAI Embedding Model: {self._embedding_model}")
else:
logger.info(
f"Using OpenAI API-compatible Embedding Model: {self._embedding_model}"
)
def get_text_embedder(self):
return AsyncTextEmbedder(
api_key=self._api_key,
wren-ai-service/src/providers/llm/azure_openai.py
+9 -6
View File
@@ -14,6 +14,7 @@ from openai.types.chat import ChatCompletion, ChatCompletionChunk
from src.core.provider import LLMProvider
from src.providers.loader import provider
from src.utils import remove_trailing_slash
logger = logging.getLogger("wren-ai-service")
@@ -76,7 +77,7 @@ class AsyncGenerator(AzureOpenAIGenerator):
completion: Union[
Stream[ChatCompletionChunk], ChatCompletion
] = await self.client.chat.completions.create(
model=self.model,
model=self.azure_deployment,
messages=openai_formatted_messages,
stream=self.streaming_callback is not None,
**generation_kwargs,
@@ -122,15 +123,17 @@ 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 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
self._generation_api_base = remove_trailing_slash(chat_api_base)
self._generation_api_version = chat_api_version
self._generation_model = generation_model
logger.info(f"Using AzureOpenAI LLM: {self._generation_model}")
logger.info(f"Using AzureOpenAI LLM with API base: {self._generation_api_base}")
logger.info(
f"Using AzureOpenAI LLM with API version: {self._generation_api_version}"
)
def get_generator(
self,
model_kwargs: Dict[str, Any] = (
wren-ai-service/src/providers/llm/ollama.py
+5 -3
View File
@@ -10,6 +10,7 @@ from haystack_integrations.components.generators.ollama import OllamaGenerator
from src.core.provider import LLMProvider
from src.providers.loader import provider
from src.utils import remove_trailing_slash
logger = logging.getLogger("wren-ai-service")
@@ -126,11 +127,12 @@ class OllamaLLMProvider(LLMProvider):
url: str = os.getenv("LLM_OLLAMA_URL") or LLM_OLLAMA_URL,
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._url = remove_trailing_slash(url)
self._generation_model = generation_model
logger.info(f"Using Ollama LLM: {self._generation_model}")
logger.info(f"Using Ollama URL: {self._url}")
def get_generator(
self,
model_kwargs: Dict[str, Any] = (
wren-ai-service/src/providers/llm/openai.py
+10 -8
View File
@@ -14,6 +14,7 @@ from openai.types.chat import ChatCompletion, ChatCompletionChunk
from src.core.provider import LLMProvider
from src.providers.loader import provider
from src.utils import remove_trailing_slash
logger = logging.getLogger("wren-ai-service")
@@ -127,17 +128,18 @@ class OpenAILLMProvider(LLMProvider):
"""
OpenAI(api_key=api_key, base_url=api_base).models.list()
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)
logger.info(f"Using OpenAI LLM: {generation_model}")
else:
logger.info(f"Using OpenAI API-compatible LLM: {generation_model}")
self._api_key = api_key
self._api_base = api_base
self._api_base = remove_trailing_slash(api_base)
self._generation_model = generation_model
logger.info(f"Using OpenAILLM provider with API base: {self._api_base}")
# TODO: currently only OpenAI api key can be verified
if self._api_base == LLM_OPENAI_API_BASE:
_verify_api_key(self._api_key.resolve_value(), self._api_base)
logger.info(f"Using OpenAI LLM: {self._generation_model}")
else:
logger.info(f"Using OpenAI API-compatible LLM: {self._generation_model}")
def get_generator(
self,
model_kwargs: Dict[str, Any] = (
wren-ai-service/src/utils.py
+4
View File
@@ -117,3 +117,7 @@ def async_timer(func):
return await process(func, *args, **kwargs)
return wrapper_timer
def remove_trailing_slash(endpoint: str) -> str:
return endpoint.rstrip("/") if endpoint.endswith("/") else endpoint