Skip to content

Commit

Permalink
chore(wren-ai-service): separate the api key for each LLM and embedder (
Browse files Browse the repository at this point in the history
  • Loading branch information
paopa authored Nov 15, 2024
1 parent 8f61841 commit 75910d2
Show file tree
Hide file tree
Showing 11 changed files with 36 additions and 23 deletions.
7 changes: 6 additions & 1 deletion deployment/kustomizations/base/deploy-wren-ai-service.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -24,11 +24,16 @@ spec:
configMapKeyRef:
name: wren-config
key: WREN_AI_SERVICE_PORT
- name: OPENAI_API_KEY
- name: LLM_OPENAI_API_KEY
valueFrom:
secretKeyRef:
name: wrenai-secrets
key: LLM_OPENAI_API_KEY
- name: EMBEDDER_OPENAI_API_KEY
valueFrom:
secretKeyRef:
name: wrenai-secrets
key: EMBEDDER_OPENAI_API_KEY
- name: QDRANT_HOST
valueFrom:
configMapKeyRef:
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ type: Opaque
data:
# LLM_OPENAI_API_KEY and EMBEDDER_OPENAI_API_KEY are REQUIRED: without a valid key the wren-ai-service-deployment pod will not start
LLM_OPENAI_API_KEY: UkVRVUlSRUQ6IHNrLXByb2otYWxsLWFjY2Vzcy1wbGFjZWhvbGRlci00LXdyZW4tYWktc2VydmljZS1kZXBsb3ltZW50
EMBEDDER_OPENAI_API_KEY: UkVRVUlSRUQ6IHNrLXByb2otYWxsLWFjY2Vzcy1wbGFjZWhvbGRlci00LXdyZW4tYWktc2VydmljZS1kZXBsb3ltZW50

# Azure openai env
AZURE_CHAT_BASE: bi9h
Expand Down
6 changes: 4 additions & 2 deletions docker/.env.example
Original file line number Diff line number Diff line change
Expand Up @@ -18,8 +18,10 @@ QDRANT_HOST=qdrant
SHOULD_FORCE_DEPLOY=1

# vendor keys
OPENAI_API_KEY=
AZURE_OPENAI_API_KEY=
LLM_OPENAI_API_KEY=
EMBEDDER_OPENAI_API_KEY=
LLM_AZURE_OPENAI_API_KEY=
EMBEDDER_AZURE_OPENAI_API_KEY=
QDRANT_API_KEY=

# version
Expand Down
6 changes: 4 additions & 2 deletions docker/docker-compose-dev.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -44,8 +44,10 @@ services:
WREN_UI_PORT: ${WREN_UI_PORT}
QDRANT_HOST: ${QDRANT_HOST}
WREN_UI_ENDPOINT: ${WREN_UI_ENDPOINT}
OPENAI_API_KEY: ${OPENAI_API_KEY}
AZURE_OPENAI_API_KEY: ${AZURE_OPENAI_API_KEY}
LLM_OPENAI_API_KEY: ${LLM_OPENAI_API_KEY}
EMBEDDER_OPENAI_API_KEY: ${EMBEDDER_OPENAI_API_KEY}
LLM_AZURE_OPENAI_API_KEY: ${LLM_AZURE_OPENAI_API_KEY}
EMBEDDER_AZURE_OPENAI_API_KEY: ${EMBEDDER_AZURE_OPENAI_API_KEY}
QDRANT_API_KEY: ${QDRANT_API_KEY}
SHOULD_FORCE_DEPLOY: ${SHOULD_FORCE_DEPLOY}
# sometimes the console won't show print messages,
Expand Down
6 changes: 4 additions & 2 deletions docker/docker-compose.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -56,8 +56,10 @@ services:
WREN_UI_PORT: ${WREN_UI_PORT}
QDRANT_HOST: ${QDRANT_HOST}
WREN_UI_ENDPOINT: http://wren-ui:${WREN_UI_PORT}
OPENAI_API_KEY: ${OPENAI_API_KEY}
AZURE_OPENAI_API_KEY: ${AZURE_OPENAI_API_KEY}
LLM_OPENAI_API_KEY: ${LLM_OPENAI_API_KEY}
EMBEDDER_OPENAI_API_KEY: ${EMBEDDER_OPENAI_API_KEY}
LLM_AZURE_OPENAI_API_KEY: ${LLM_AZURE_OPENAI_API_KEY}
EMBEDDER_AZURE_OPENAI_API_KEY: ${EMBEDDER_AZURE_OPENAI_API_KEY}
QDRANT_API_KEY: ${QDRANT_API_KEY}
SHOULD_FORCE_DEPLOY: ${SHOULD_FORCE_DEPLOY}
# sometimes the console won't show print messages,
Expand Down
3 changes: 1 addition & 2 deletions wren-ai-service/src/providers/embedder/azure_openai.py
Original file line number Diff line number Diff line change
Expand Up @@ -195,8 +195,7 @@ async def run(self, documents: List[Document]):
class AzureOpenAIEmbedderProvider(EmbedderProvider):
def __init__(
self,
api_key: Secret = Secret.from_env_var("AZURE_OPENAI_API_KEY")
or Secret.from_env_var("EMBEDDER_AZURE_OPENAI_API_KEY"),
api_key: Secret = Secret.from_env_var("EMBEDDER_AZURE_OPENAI_API_KEY"),
api_base: str = os.getenv("EMBEDDER_AZURE_OPENAI_API_BASE"),
api_version: str = os.getenv("EMBEDDER_AZURE_OPENAI_VERSION"),
model: str = os.getenv("EMBEDDING_MODEL") or EMBEDDING_MODEL,
Expand Down
3 changes: 1 addition & 2 deletions wren-ai-service/src/providers/embedder/openai.py
Original file line number Diff line number Diff line change
Expand Up @@ -183,8 +183,7 @@ async def run(self, documents: List[Document]):
class OpenAIEmbedderProvider(EmbedderProvider):
def __init__(
self,
api_key: str = os.getenv("OPENAI_API_KEY")
or os.getenv("EMBEDDER_OPENAI_API_KEY"),
api_key: str = os.getenv("EMBEDDER_OPENAI_API_KEY"),
api_base: str = os.getenv("EMBEDDER_OPENAI_API_BASE")
or EMBEDDER_OPENAI_API_BASE,
model: str = os.getenv("EMBEDDING_MODEL") or EMBEDDING_MODEL,
Expand Down
3 changes: 1 addition & 2 deletions wren-ai-service/src/providers/llm/azure_openai.py
Original file line number Diff line number Diff line change
Expand Up @@ -125,8 +125,7 @@ async def run(
class AzureOpenAILLMProvider(LLMProvider):
def __init__(
self,
api_key: Secret = Secret.from_env_var("AZURE_OPENAI_API_KEY")
or Secret.from_env_var("LLM_AZURE_OPENAI_API_KEY"),
api_key: Secret = Secret.from_env_var("LLM_AZURE_OPENAI_API_KEY"),
api_base: str = os.getenv("LLM_AZURE_OPENAI_API_BASE"),
api_version: str = os.getenv("LLM_AZURE_OPENAI_VERSION"),
model: str = os.getenv("GENERATION_MODEL") or GENERATION_MODEL,
Expand Down
2 changes: 1 addition & 1 deletion wren-ai-service/src/providers/llm/openai.py
Original file line number Diff line number Diff line change
Expand Up @@ -157,7 +157,7 @@ async def run(
class OpenAILLMProvider(LLMProvider):
def __init__(
self,
api_key: str = os.getenv("OPENAI_API_KEY") or os.getenv("LLM_OPENAI_API_KEY"),
api_key: str = os.getenv("LLM_OPENAI_API_KEY"),
api_base: str = os.getenv("LLM_OPENAI_API_BASE") or LLM_OPENAI_API_BASE,
model: str = os.getenv("GENERATION_MODEL") or GENERATION_MODEL,
kwargs: Dict[str, Any] = (
Expand Down
6 changes: 4 additions & 2 deletions wren-ai-service/tools/config/.env.dev.example
Original file line number Diff line number Diff line change
@@ -1,8 +1,10 @@
CONFIG_PATH=config.yaml

# vendor keys
OPENAI_API_KEY=
AZURE_OPENAI_API_KEY=
LLM_OPENAI_API_KEY=
LLM_AZURE_OPENAI_API_KEY=
EMBEDDER_OPENAI_API_KEY=
EMBEDDER_AZURE_OPENAI_API_KEY=
QDRANT_API_KEY=

# langfuse key
Expand Down
16 changes: 9 additions & 7 deletions wren-launcher/utils/docker.go
Original file line number Diff line number Diff line change
Expand Up @@ -26,8 +26,8 @@ const (
// please change the version when the version is updated
WREN_PRODUCT_VERSION string = "0.10.0"
DOCKER_COMPOSE_YAML_URL string = "https://raw.githubusercontent.com/Canner/WrenAI/" + WREN_PRODUCT_VERSION + "/docker/docker-compose.yaml"
DOCKER_COMPOSE_ENV_URL string = "https://raw.githubusercontent.com/Canner/WrenAI/" + WREN_PRODUCT_VERSION + "/docker/.env.example"
AI_SERVICE_CONFIG_URL string = "https://raw.githubusercontent.com/Canner/WrenAI/" + WREN_PRODUCT_VERSION + "/docker/config.example.yaml"
DOCKER_COMPOSE_ENV_URL string = "https://raw.githubusercontent.com/Canner/WrenAI/" + WREN_PRODUCT_VERSION + "/docker/.env.example"
AI_SERVICE_CONFIG_URL string = "https://raw.githubusercontent.com/Canner/WrenAI/" + WREN_PRODUCT_VERSION + "/docker/config.example.yaml"
)

func replaceEnvFileContent(content string, projectDir string, openaiApiKey string, openAIGenerationModel string, hostPort int, aiPort int, userUUID string, telemetryEnabled bool) string {
Expand All @@ -39,9 +39,13 @@ func replaceEnvFileContent(content string, projectDir string, openaiApiKey strin
reg = regexp.MustCompile(`SHOULD_FORCE_DEPLOY=(.*)`)
str = reg.ReplaceAllString(str, "SHOULD_FORCE_DEPLOY=1")

// replace OPENAI_API_KEY
reg = regexp.MustCompile(`(?m)^OPENAI_API_KEY=(.*)`)
str = reg.ReplaceAllString(str, "OPENAI_API_KEY="+openaiApiKey)
// replace LLM_OPENAI_API_KEY
reg = regexp.MustCompile(`LLM_OPENAI_API_KEY=(.*)`)
str = reg.ReplaceAllString(str, "LLM_OPENAI_API_KEY="+openaiApiKey)

// replace EMBEDDER_OPENAI_API_KEY
reg = regexp.MustCompile(`EMBEDDER_OPENAI_API_KEY=(.*)`)
str = reg.ReplaceAllString(str, "EMBEDDER_OPENAI_API_KEY="+openaiApiKey)

// replace GENERATION_MODEL
// it seems like using for telemetry to know the model, might be we can remove this in the future and provide a endpoint to get the information
Expand Down Expand Up @@ -149,8 +153,6 @@ func PrepareConfigFileForOpenAI(projectDir string, generationModel string) error
// replace the generation model in config.yaml
config := string(content)
config = strings.ReplaceAll(config, "openai_llm.gpt-4o-mini", "openai_llm."+generationModel)
// disable the langfuse for starting wren-ai from the launcher
config = strings.ReplaceAll(config, "langfuse_enable: true", "langfuse_enable: false")

// write back to config.yaml
err = os.WriteFile(configPath, []byte(config), 0644)
Expand Down

0 comments on commit 75910d2

Please sign in to comment.