Skip to content

Commit dfa0f98

Browse files
authored
Fixes no api key embedding name endpoint (#337)
* Fixes no api key embedding name endpoint * privatemode ai for ac * Fix unbound session in debounce * Adds terms for privatemode ai * PR comment * Submodule merge
1 parent 22223cc commit dfa0f98

File tree

6 files changed

+53
-8
lines changed

6 files changed

+53
-8
lines changed

controller/attribute/llm_response_tmpl.py

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -25,9 +25,9 @@
2525

2626
class LLMProvider_A2VYBG(Enum):
2727
OPEN_AI = "Open AI"
28-
OPEN_SOURCE = "Open-Source"
2928
AZURE = "Azure"
3029
AZURE_FOUNDRY = "Azure Foundry"
30+
PRIVATEMODE_AI = "Privatemode AI"
3131

3232

3333
# OpenAI migration guides
@@ -91,6 +91,13 @@ def get_client_openai_8e8a360e_3f7f_4cf9_ba80_8cb239e897d2(
9191
) -> Union[OpenAI, AsyncOpenAI, AzureOpenAI, AsyncAzureOpenAI]:
9292
global CLIENT_LOOKUP_A2VYBG
9393

94+
if CLIENT_TYPE_A2VYBG == LLMProvider_A2VYBG.PRIVATEMODE_AI.value:
95+
# caching was disabled (and doesn't work) so shorthand here to prevent bloated code
96+
return AsyncOpenAI(
97+
api_key="dummy",
98+
base_url="http://privatemode-proxy:8080/v1",
99+
)
100+
94101
if CLIENT_TYPE_A2VYBG == LLMProvider_A2VYBG.AZURE.value and (
95102
azure_endpoint is None or api_version is None
96103
):

controller/attribute/util.py

Lines changed: 34 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -139,6 +139,28 @@ def test_azure_foundry_llm_connection(api_key: str, base_endpoint: str):
139139
return response.json()["choices"][0]["message"]["content"]
140140

141141

142+
def test_privatemode_ai_llm_connection(model: str):
143+
# more here: https://docs.privatemode.ai/api/chat-completions
144+
headers = {
145+
"Content-Type": "application/json",
146+
}
147+
payload = {
148+
"model": model,
149+
"messages": [
150+
{"role": "user", "content": [{"type": "text", "text": "only say 'hello'"}]},
151+
],
152+
"max_tokens": 5,
153+
}
154+
155+
response = requests.post(
156+
"http://privatemode-proxy:8080/v1/chat/completions",
157+
headers=headers,
158+
json=payload,
159+
)
160+
response.raise_for_status()
161+
return response.json()["choices"][0]["message"]["content"]
162+
163+
142164
def test_azure_llm_connection(
143165
api_key: str,
144166
base_endpoint: str,
@@ -229,6 +251,10 @@ def validate_llm_config(llm_config: Dict[str, Any]):
229251
api_key=llm_config["apiKey"],
230252
base_endpoint=llm_config["apiBase"],
231253
)
254+
elif llm_config["llmIdentifier"] == enums.LLMProvider.PRIVATEMODE_AI.value:
255+
test_privatemode_ai_llm_connection(
256+
model=llm_config["model"],
257+
)
232258
else:
233259
raise LlmResponseError(
234260
"LLM Identifier must be either Open AI or Azure, got: "
@@ -294,12 +320,15 @@ async def ac(record):
294320
validate_llm_config(llm_config=llm_config)
295321

296322
num_workers = 50
297-
if (
298-
llm_config is not None
299-
and enums.LLMProvider.from_string(llm_config.get("llmIdentifier", "Open ai"))
300-
== enums.LLMProvider.AZURE_FOUNDRY
301-
):
323+
llm_provider = enums.LLMProvider.from_string(
324+
llm_config.get("llmIdentifier", "Open ai")
325+
)
326+
327+
if llm_config is not None and llm_provider == enums.LLMProvider.AZURE_FOUNDRY:
302328
num_workers = 25
329+
elif llm_provider == enums.LLMProvider.PRIVATEMODE_AI:
330+
# to prevent api rate limit issues
331+
num_workers = 10
303332

304333
try:
305334
llm_config_mapping = {

controller/embedding/terms.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,4 +16,9 @@
1616
"terms": "Please note that by enabling this third-party API, you are stating that you accept its addition as a sub-processor under the terms of our Data Processing Agreement. Please be aware that the Azure API policies may conflict with your internal data and privacy policies. For more information please check: @@PLACEHOLDER@@. For questions you can contact us at [email protected].",
1717
"link": "https://www.microsoft.com/en-us/legal/terms-of-use",
1818
},
19+
EmbeddingPlatform.PRIVATEMODE_AI.value: {
20+
"platform": EmbeddingPlatform.PRIVATEMODE_AI.value,
21+
"terms": "Please note that by enabling this third-party API, you are stating that you accept its addition as a sub-processor under the terms of our Data Processing Agreement. Please be aware that the PrivateMode AI API policies may conflict with your internal data and privacy policies. For more information please check: @@PLACEHOLDER@@. For questions you can contact us at [email protected].",
22+
"link": "https://www.privatemode.ai/terms-of-service",
23+
},
1924
}

controller/organization/manager.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -101,13 +101,15 @@ def get_overview_stats(org_id: str) -> List[Dict[str, Union[str, int]]]:
101101

102102
# INFO: Not fully debounced if server runs multiple instances
103103
# TODO: Change to 60 to 300 for prod
104-
@param_debounce(seconds=60)
104+
@param_debounce(seconds=300)
105105
def sync_organization_sharepoint_integrations(org_id: str) -> None:
106+
general.get_ctx_token()
106107
all_integrations = integration.get_all_in_org(
107108
org_id, enums.CognitionIntegrationType.SHAREPOINT.value
108109
)
109110
all_integration_ids = [
110111
str(integration_entity.id) for integration_entity in all_integrations
111112
]
113+
general.remove_and_refresh_session()
112114
for integration_id in all_integration_ids:
113115
transfer_api.post_process_integration(integration_id)

fast_api/routes/embedding.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -189,6 +189,8 @@ def get_embedding_name(
189189
api_token = env_var_db_bo.get_by_name_and_org_id(
190190
data.org_id, data.api_token_env_name
191191
).value
192+
else:
193+
api_token = None
192194
return manager.get_embedding_name(
193195
project_id,
194196
data.attribute_id,

submodules/model

Submodule model updated 1 file

0 commit comments

Comments
 (0)