feat: add variables and API Key support

.github/workflows/ci.yml

@@ -19,7 +19,7 @@ jobs:
       - name: Install dependencies
         run: |
           python -m pip install --upgrade pip
-          pip install build requests sseclient-py python-dotenv pytest
+          pip install build requests sseclient-py pytest
       - name: Build package
         run: python -m build
       - name: Run Tests

README.md

@@ -14,24 +14,24 @@ pip install commonbase
 
 ## Usage
 
-A project ID is required for all Commonbase requests. You can find your project ID
-in the [Commonbase Dashboard](https://commonbase.com/).
+A Project ID and API Key are required for all Commonbase requests. You can find your project ID and generate an API key in the [Commonbase Dashboard](https://commonbase.com/).
 
-To create text and chat completions, use `commonbase.Completion.create`:
+To create a completion, provide your Project ID, API Key, and prompt to `Completion.create`.
 
 ```py
 import commonbase
 
 project_id=
 
 result = commonbase.Completion.create(
-    project_id="<your_project_id>",
+    api_key="YOUR_API_KEY",
+    project_id="YOUR_PROJECT_ID",
     prompt="Hello!"
 )
 
-print(result.choices[0].text)
+print(result.best_result)
 ```
 
-To stream a completion as it is generated, use `commonbase.Completion.stream`.
+To stream a completion as it is generated, use `Completion.stream`.
 
-For more examples, see [/examples](https://github.com/commonbaseapp/commonbase-python/tree/main/examples).
+For more examples, see [/examples](https://github.com/commonbaseapp/commonbase-python/tree/main/examples) or check out our [Docs](https://docs.commonbase.com/quickstart/python).

commonbase/__init__.py

@@ -1,11 +1,12 @@
 from commonbase.completion import Completion
-from commonbase.exceptions import CommonbaseException
+from commonbase.exceptions import CommonbaseApiException, CommonbaseException
 from commonbase.chat_context import ChatContext, ChatMessage
 from commonbase.provider_config import ProviderConfig, OpenAIParams, AnthropicParams
 
 __all__: [
     "Completion",
     "CommonbaseException",
+    "CommonbaseApiException",
     "ChatContext",
     "ChatMessage",
     "ProviderConfig",

commonbase/chat_context.py

@@ -12,5 +12,5 @@ class ChatMessage:
 class ChatContext:
     messages: Sequence[ChatMessage]
 
-    def _to_dict(self) -> dict:
+    def _as_json(self) -> dict:
         return asdict(self)

commonbase/completion.py

@@ -3,85 +3,138 @@
 import requests
 import sseclient
 from commonbase.completion_response import CompletionResponse
-from commonbase.exceptions import CommonbaseException
+from commonbase.exceptions import CommonbaseApiException, CommonbaseException
 from commonbase.chat_context import ChatContext
 from commonbase.provider_config import ProviderConfig
+from commonbase.truncation_config import TruncationConfig
+from importlib.metadata import version, PackageNotFoundError
 
 
-class Completion:
-    @classmethod
-    def _send_completion_request(
-        cls,
-        project_id: str,
-        prompt: str,
-        api_key: Optional[str] = None,
-        chat_context: Optional[ChatContext] = None,
-        user_id: Optional[str] = None,
-        truncate_variable: Optional[bool] = None,
-        provider_config: Optional[ProviderConfig] = None,
-        stream: bool = False,
-    ) -> requests.Response:
-        assert project_id is not None
-        assert prompt is not None
-
-        data = {
-            "projectId": project_id,
-            "prompt": prompt,
-            "apiKey": api_key,
-            "context": chat_context._to_dict() if chat_context is not None else None,
-            "userId": user_id,
-            "truncateVariable": truncate_variable,
-            "providerConfig": provider_config._to_dict()
-            if provider_config is not None
-            else None,
-            "stream": stream,
-        }
-        data = {k: v for k, v in data.items() if v is not None}
-        return requests.post(
-            "https://api.commonbase.com/completions",
-            stream=stream,
-            json=data,
-            headers={"Accept": "text/event-stream"} if stream else None,
+def _get_sdk_version() -> str:
+    try:
+        return version("commonbase")
+    except PackageNotFoundError:
+        # This error is thrown during testing.
+        return "0.0.0"
+
+
+def _format_body(
+    project_id: str,
+    prompt: Optional[str] = None,
+    variables: Optional[dict[str, str]] = None,
+    chat_context: Optional[ChatContext] = None,
+    user_id: Optional[str] = None,
+    truncate_variable: Optional[TruncationConfig] = None,
+    provider_config: Optional[ProviderConfig] = None,
+    stream: bool = False,
+):
+    data = {
+        "projectId": project_id,
+        "prompt": prompt,
+        "variables": variables,
+        "context": chat_context._as_json() if chat_context is not None else None,
+        "userId": user_id,
+        "truncateVariable": truncate_variable._as_json()
+        if truncate_variable is not None
+        else None,
+        "providerConfig": provider_config._as_json()
+        if provider_config is not None
+        else None,
+        "stream": stream,
+    }
+    return {k: v for k, v in data.items() if v is not None}
+
+
+def _send_completion_request(
+    api_key: str,
+    project_id: str,
+    prompt: Optional[str],
+    variables: Optional[dict[str, str]],
+    chat_context: Optional[ChatContext],
+    user_id: Optional[str],
+    truncate_variable: Optional[TruncationConfig],
+    provider_config: Optional[ProviderConfig],
+    stream: bool,
+) -> requests.Response:
+    if api_key is None:
+        raise CommonbaseException(
+            "A Commonbase API key must be provided with every request."
         )
 
+    data = _format_body(
+        project_id=project_id,
+        prompt=prompt,
+        variables=variables,
+        chat_context=chat_context,
+        user_id=user_id,
+        truncate_variable=truncate_variable,
+        provider_config=provider_config,
+        stream=stream,
+    )
+
+    headers = {
+        "Authorization": api_key,
+        "User-Agent": f"commonbase-python/{_get_sdk_version()}",
+    }
+
+    if stream:
+        headers["Accept"] = "text/event-stream"
+
+    return requests.post(
+        "https://api.commonbase.com/completions",
+        stream=stream,
+        json=data,
+        headers=headers,
+    )
+
+
+class Completion:
     @classmethod
     def create(
         cls,
+        api_key: str,
         project_id: str,
-        prompt: str,
-        api_key: Optional[str] = None,
+        prompt: Optional[str] = None,
+        variables: Optional[dict[str, str]] = None,
         chat_context: Optional[ChatContext] = None,
         user_id: Optional[str] = None,
-        truncate_variable: Optional[bool] = None,
+        truncate_variable: Optional[TruncationConfig] = None,
         provider_config: Optional[ProviderConfig] = None,
     ) -> CompletionResponse:
         """Creates a completion for the given prompt.
 
         Parameters
         ----------
+        api_key : str
+            The Commonbase API key used to authenticate the request.
         project_id : str
-            The ID of your Commonbase project.
-        prompt : str
+            The ID of the Commonbase project.
+        prompt : str, optional
             The prompt for which a completion is generated.
-        api_key : str, optional
-            The API key used for authentication. Currently not necessary.
+        variables : dict[str, str], optional
+            The list of variables to use with Commonbase managed prompts.
         chat_context : ChatContext, optional
             The list of chat messages in a conversation
         user_id : str, optional
             The User ID that will be logged for the invocation.
-        truncate_variable : bool, optional
-            A flag to toggle redaction variable truncation.
+        truncate_variable : TruncationConfig, optional
+            Configures variable truncation.
+        provider_config : ProviderConfig, optional
+            Configures the completion provider to use, currently OpenAI or Anthropic.
 
         Raises
         ------
         CommonbaseException
-            If the request is malformed or there is an API error.
+            If the request parameters are invalid.
+        CommonbaseApiException
+            If there is an API error.
         """
 
-        response = Completion._send_completion_request(
+        response = _send_completion_request(
+            api_key=api_key,
             project_id=project_id,
             prompt=prompt,
-            api_key=api_key,
+            variables=variables,
             chat_context=chat_context,
             user_id=user_id,
             truncate_variable=truncate_variable,
@@ -92,47 +145,32 @@ def create(
         json = response.json()
 
         if response.status_code >= 400 or "error" in json:
-            raise CommonbaseException(json)
+            raise CommonbaseApiException(json)
 
         return CompletionResponse(response.json())
 
     @classmethod
     def stream(
         cls,
+        api_key: str,
         project_id: str,
-        prompt: str,
-        api_key: Optional[str] = None,
+        prompt: Optional[str] = None,
+        variables: Optional[dict[str, str]] = None,
         chat_context: Optional[ChatContext] = None,
         user_id: Optional[str] = None,
-        truncate_variable: Optional[bool] = None,
+        truncate_variable: Optional[TruncationConfig] = None,
         provider_config: Optional[ProviderConfig] = None,
     ) -> Generator[CompletionResponse, None, None]:
         """Creates a completion stream for the given prompt.
 
-        Parameters
-        ----------
-        project_id : str
-            The ID of your Commonbase project.
-        prompt : str
-            The prompt for which a completion is generated.
-        api_key : str, optional
-            The API key used for authentication. Currently not necessary.
-        chat_context : ChatContext, optional
-            The list of chat messages in a conversation
-        user_id : str, optional
-            The User ID that will be logged for the invocation.
-        truncate_variable : bool, optional
-            A flag to toggle redaction variable truncation.
-
-        Raises
-        ------
-        CommonbaseException
-            If the request is malformed or there is an API error.
+        This method is identical to Completion.create(), except it returns
+        a stream of completion responses.
         """
-        response = Completion._send_completion_request(
+        response = _send_completion_request(
+            api_key=api_key,
             project_id=project_id,
             prompt=prompt,
-            api_key=api_key,
+            variables=variables,
             chat_context=chat_context,
             user_id=user_id,
             truncate_variable=truncate_variable,
@@ -141,7 +179,7 @@ def stream(
         )
 
         if response.status_code >= 400:
-            raise CommonbaseException(response.json())
+            raise CommonbaseApiException(response.json())
 
         client = sseclient.SSEClient(response)
         for event in client.events():

commonbase/completion_response.py

@@ -49,3 +49,7 @@ def model(self) -> str:
     @property
     def choices(self) -> list[CompletionChoice]:
         return [CompletionChoice(choice) for choice in self.json["choices"]]
+
+    @property
+    def best_result(self) -> str:
+        return self.choices[0].text

commonbase/exceptions.py

@@ -1,4 +1,9 @@
-class CommonbaseException(Exception):
+class CommonbaseApiException(Exception):
     def __init__(self, json: dict) -> None:
         self.response = json
         super().__init__(json["error"] if "error" in json else "Commonbase Error")
+
+
+class CommonbaseException(Exception):
+    def __init__(self, message: str) -> None:
+        super().__init__(message)

commonbase/provider_config.py

@@ -34,7 +34,7 @@ class ProviderConfig:
     provider: Literal["openai", "cb-openai-eu", "anthropic"]
     params: Union[OpenAIParams, AnthropicParams]
 
-    def _to_dict(self) -> dict:
+    def _as_json(self) -> dict:
         return asdict(
             self, dict_factory=lambda x: {k: v for (k, v) in x if v is not None}
         )

commonbase/tests/__init__.py

@@ -1,3 +0,0 @@
-from dotenv import load_dotenv
-
-load_dotenv()