fix: add support for provider api keys

Author: Brandon Meyerowitz

.github/workflows/ci.yml

@@ -26,6 +26,7 @@ jobs:
         env:
           CB_PROJECT_ID: ${{ secrets.CB_PROJECT_ID }}
           CB_API_KEY: ${{ secrets.CB_API_KEY }}
+          CB_OPENAI_API_KEY: ${{ secrets.CB_OPENAI_API_KEY }}
         run: pytest
 
   release:

commonbase/completion.py

@@ -53,6 +53,7 @@ def _send_completion_request(
     chat_context: Optional[ChatContext],
     user_id: Optional[str],
     truncate_variable: Optional[TruncationConfig],
+    provider_api_key: Optional[str],
     provider_config: Optional[ProviderConfig],
     stream: bool,
 ) -> requests.Response:
@@ -80,6 +81,10 @@ def _send_completion_request(
     if stream:
         headers["Accept"] = "text/event-stream"
 
+    print(provider_api_key)
+    if provider_api_key is not None:
+        headers["Provider-API-Key"] = provider_api_key
+
     return requests.post(
         "https://api.commonbase.com/completions",
         stream=stream,
@@ -99,6 +104,7 @@ def create(
         chat_context: Optional[ChatContext] = None,
         user_id: Optional[str] = None,
         truncate_variable: Optional[TruncationConfig] = None,
+        provider_api_key: Optional[str] = None,
         provider_config: Optional[ProviderConfig] = None,
     ) -> CompletionResponse:
         """Creates a completion for the given prompt.
@@ -119,6 +125,8 @@ def create(
             The User ID that will be logged for the invocation.
         truncate_variable : TruncationConfig, optional
             Configures variable truncation.
+        provider_api_key : str, optional
+            The API Key used to authenticate with a provider.
         provider_config : ProviderConfig, optional
             Configures the completion provider to use, currently OpenAI or Anthropic.
 
@@ -138,6 +146,7 @@ def create(
             chat_context=chat_context,
             user_id=user_id,
             truncate_variable=truncate_variable,
+            provider_api_key=provider_api_key,
             provider_config=provider_config,
             stream=False,
         )
@@ -159,6 +168,7 @@ def stream(
         chat_context: Optional[ChatContext] = None,
         user_id: Optional[str] = None,
         truncate_variable: Optional[TruncationConfig] = None,
+        provider_api_key: Optional[str] = None,
         provider_config: Optional[ProviderConfig] = None,
     ) -> Generator[CompletionResponse, None, None]:
         """Creates a completion stream for the given prompt.
@@ -174,6 +184,7 @@ def stream(
             chat_context=chat_context,
             user_id=user_id,
             truncate_variable=truncate_variable,
+            provider_api_key=provider_api_key,
             provider_config=provider_config,
             stream=True,
         )

commonbase/tests/test_provider_api_key.py

@@ -0,0 +1,31 @@
+import os
+import pytest
+from commonbase.exceptions import CommonbaseApiException
+from commonbase.completion import Completion
+from commonbase.provider_config import ProviderConfig, OpenAIParams
+
+
+def test_provider_with_no_api_key():
+    with pytest.raises(CommonbaseApiException):
+        Completion.create(
+            api_key=os.getenv("CB_API_KEY") or "",
+            project_id=os.getenv("CB_PROJECT_ID") or "",
+            provider_config=ProviderConfig(
+                provider="openai", params=OpenAIParams(type="chat")
+            ),
+            prompt="",
+        )
+
+
+def test_provider_with_valid_api_key():
+    result = Completion.create(
+        api_key=os.getenv("CB_API_KEY") or "",
+        project_id=os.getenv("CB_PROJECT_ID") or "",
+        provider_api_key=os.getenv("CB_OPENAI_API_KEY") or "",
+        provider_config=ProviderConfig(
+            provider="openai", params=OpenAIParams(type="chat")
+        ),
+        prompt="Please return the string '123abc' to me without the quotes.",
+    )
+
+    assert result.completed and result.best_result.strip() == "123abc"