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

Author: JWittmeyer

controller/attribute/llm_response_tmpl.py

@@ -25,9 +25,9 @@
 
 class LLMProvider_A2VYBG(Enum):
     OPEN_AI = "Open AI"
-    OPEN_SOURCE = "Open-Source"
     AZURE = "Azure"
     AZURE_FOUNDRY = "Azure Foundry"
+    PRIVATEMODE_AI = "Privatemode AI"
 
 
 # OpenAI migration guides
@@ -91,6 +91,13 @@ def get_client_openai_8e8a360e_3f7f_4cf9_ba80_8cb239e897d2(
 ) -> Union[OpenAI, AsyncOpenAI, AzureOpenAI, AsyncAzureOpenAI]:
     global CLIENT_LOOKUP_A2VYBG
 
+    if CLIENT_TYPE_A2VYBG == LLMProvider_A2VYBG.PRIVATEMODE_AI.value:
+        # caching was disabled (and doesn't work) so shorthand here to prevent bloated code
+        return AsyncOpenAI(
+            api_key="dummy",
+            base_url="http://privatemode-proxy:8080/v1",
+        )
+
     if CLIENT_TYPE_A2VYBG == LLMProvider_A2VYBG.AZURE.value and (
         azure_endpoint is None or api_version is None
     ):

controller/attribute/util.py

@@ -139,6 +139,28 @@ def test_azure_foundry_llm_connection(api_key: str, base_endpoint: str):
     return response.json()["choices"][0]["message"]["content"]
 
 
+def test_privatemode_ai_llm_connection(model: str):
+    # more here: https://docs.privatemode.ai/api/chat-completions
+    headers = {
+        "Content-Type": "application/json",
+    }
+    payload = {
+        "model": model,
+        "messages": [
+            {"role": "user", "content": [{"type": "text", "text": "only say 'hello'"}]},
+        ],
+        "max_tokens": 5,
+    }
+
+    response = requests.post(
+        "http://privatemode-proxy:8080/v1/chat/completions",
+        headers=headers,
+        json=payload,
+    )
+    response.raise_for_status()
+    return response.json()["choices"][0]["message"]["content"]
+
+
 def test_azure_llm_connection(
     api_key: str,
     base_endpoint: str,
@@ -229,6 +251,10 @@ def validate_llm_config(llm_config: Dict[str, Any]):
                 api_key=llm_config["apiKey"],
                 base_endpoint=llm_config["apiBase"],
             )
+        elif llm_config["llmIdentifier"] == enums.LLMProvider.PRIVATEMODE_AI.value:
+            test_privatemode_ai_llm_connection(
+                model=llm_config["model"],
+            )
         else:
             raise LlmResponseError(
                 "LLM Identifier must be either Open AI or Azure, got: "
@@ -294,12 +320,15 @@ async def ac(record):
     validate_llm_config(llm_config=llm_config)
 
     num_workers = 50
-    if (
-        llm_config is not None
-        and enums.LLMProvider.from_string(llm_config.get("llmIdentifier", "Open ai"))
-        == enums.LLMProvider.AZURE_FOUNDRY
-    ):
+    llm_provider = enums.LLMProvider.from_string(
+        llm_config.get("llmIdentifier", "Open ai")
+    )
+
+    if llm_config is not None and llm_provider == enums.LLMProvider.AZURE_FOUNDRY:
         num_workers = 25
+    elif llm_provider == enums.LLMProvider.PRIVATEMODE_AI:
+        # to prevent api rate limit issues
+        num_workers = 10
 
     try:
         llm_config_mapping = {

controller/embedding/terms.py

@@ -16,4 +16,9 @@
         "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 security@kern.ai.",
         "link": "https://www.microsoft.com/en-us/legal/terms-of-use",
     },
+    EmbeddingPlatform.PRIVATEMODE_AI.value: {
+        "platform": EmbeddingPlatform.PRIVATEMODE_AI.value,
+        "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 security@kern.ai.",
+        "link": "https://www.privatemode.ai/terms-of-service",
+    },
 }

controller/organization/manager.py

@@ -101,13 +101,15 @@ def get_overview_stats(org_id: str) -> List[Dict[str, Union[str, int]]]:
 
 # INFO: Not fully debounced if server runs multiple instances
 # TODO: Change to 60 to 300 for prod
-@param_debounce(seconds=60)
+@param_debounce(seconds=300)
 def sync_organization_sharepoint_integrations(org_id: str) -> None:
+    general.get_ctx_token()
     all_integrations = integration.get_all_in_org(
         org_id, enums.CognitionIntegrationType.SHAREPOINT.value
     )
     all_integration_ids = [
         str(integration_entity.id) for integration_entity in all_integrations
     ]
+    general.remove_and_refresh_session()
     for integration_id in all_integration_ids:
         transfer_api.post_process_integration(integration_id)

fast_api/routes/embedding.py

@@ -189,6 +189,8 @@ def get_embedding_name(
         api_token = env_var_db_bo.get_by_name_and_org_id(
             data.org_id, data.api_token_env_name
         ).value
+    else:
+        api_token = None
     return manager.get_embedding_name(
         project_id,
         data.attribute_id,