GEN-1140 add billing capabilities (#8)

* Squashed 'protos/' changes from a643df8..8c94d7c

8c94d7c Merge pull request #10 from macrocosm-os/gen-1140_add_billing_to_const_api
066033c GEN-1140 Add billing to Const API
b22028b GEN-1154: Add endpoint for validating reddit topics (#104)

git-subtree-dir: protos
git-subtree-split: 8c94d7c4b6b814e417e9070369916e03eeb13231

* added billing capabilities

* version 1.0.0

Author: cassova

examples/apex_chat_deep_research.py

@@ -0,0 +1,66 @@
+"""
+Example of using the async Apex Chat API to stream a chat completion using the Macrocosmos SDK.
+"""
+
+import asyncio
+import os
+import time
+
+import macrocosmos as mc
+
+
+async def demo_deep_research():
+    """Demo deep research."""
+
+    messages = [
+        mc.ChatMessage(
+            role="user",
+            content="How do you design a federated LLM training pipeline that shards the model across remote systems across untrusted nodes?",
+        )
+    ]
+    api_key = os.environ.get("APEX_API_KEY", os.environ.get("MACROCOSMOS_API_KEY"))
+
+    sampling_params = mc.SamplingParameters(
+        temperature=0.6,
+        top_p=0.95,
+        max_new_tokens=8192,
+        do_sample=False,
+    )
+
+    client = mc.AsyncApexClient(
+        max_retries=0,
+        timeout=21000,
+        api_key=api_key,
+        app_name="examples/apex_chat_deep_research",
+    )
+
+    try:
+        start_time = time.time()
+        response_stream = await client.chat.completions.create(
+            model="mrfakename/mistral-small-3.1-24b-instruct-2503-hf",
+            messages=messages,
+            inference_mode="Chain-of-Thought",
+            sampling_parameters=sampling_params,
+            stream=True,
+        )
+
+        print("Streaming response:")
+        full_content = ""
+
+        chunk: mc.ChatCompletionChunkResponse
+        async for chunk in response_stream:
+            if chunk.choices and chunk.choices[0].delta.content:
+                content = chunk.choices[0].delta.content
+                full_content += content
+                print(content, end="", flush=True)
+
+    except Exception as e:
+        print(f"Unexpected Error: {e}")
+        raise e
+    finally:
+        print(f"\nTotal time taken: {time.time() - start_time:.2f} seconds")
+
+
+if __name__ == "__main__":
+    print("Testing streaming ChatCompletion:")
+    asyncio.run(demo_deep_research())

examples/billing_usage_example.py

@@ -0,0 +1,25 @@
+import macrocosmos as mc
+import os
+
+
+def main():
+    # Initialize the client
+    client = mc.BillingClient(api_key=os.environ.get("MACROCOSMOS_API_KEY"))
+
+    # Get usage information for gravity
+    response = client.billing.GetUsage()
+
+    # Print the response
+    print(f"Available credits: {response.available_credits}")
+    print(f"Used credits: {response.used_credits}")
+    print(f"Remaining credits: {response.remaining_credits}")
+    print("\nBilling rates:")
+    for rate in response.billing_rates:
+        print(f"  - Rate type: {rate.rate_type}")
+        print(f"    Unit size: {rate.unit_size}")
+        print(f"    Unit type: {rate.unit_type}")
+        print(f"    Price per unit: {rate.price_per_unit} {rate.currency}")
+
+
+if __name__ == "__main__":
+    main()

protos/billing/v1/billing.proto

@@ -0,0 +1,42 @@
+syntax = "proto3";
+
+package billing.v1;
+
+option go_package = "macrocosm-os/rift/constellation_api/gen/billing/v1";
+
+service BillingService {
+  // Get the usage of the user's credits
+  rpc GetUsage(GetUsageRequest) returns (GetUsageResponse);
+}
+
+// GetUsageRequest is the request message for getting the usage of the user's credits
+message GetUsageRequest {
+  // product_type: the type of the product (i.e. "gravity")
+  string product_type = 1;
+}
+
+// ProductPlan is details of the subscription plan for a product
+message BillingRate {
+  // product_type: the type of the product (i.e. "gravity")
+  string rate_type = 1;
+  // unit_size: the size of the unit of the subscription (e.g. 1000, 10000, 100000)
+  int64 unit_size = 2;
+  // unit_type: the type of the unit of the subscription (i.e. "rows")
+  string unit_type = 3;
+  // price_per_unit: the price per unit of the subscription
+  float price_per_unit = 4;
+  // currency: the currency of the subscription
+  string currency = 5;
+}
+
+// GetUsageResponse is the response message for getting the usage of the user's credits
+message GetUsageResponse {
+  // available_credits: the number of credits available to the user
+  float available_credits = 1;
+  // used_credits: the number of credits used by the user
+  float used_credits = 2;
+  // remaining_credits: the number of credits remaining to the user
+  float remaining_credits = 3;
+  // subscription: the subscription that the user has
+  repeated BillingRate billing_rates = 4;
+}

protos/gravity/v1/gravity.proto

@@ -181,6 +181,8 @@ message BuildDatasetRequest {
   string crawler_id = 1;
   // notification_requests: the details of the notification to be sent to the user when the dataset is ready to download (optional)
   repeated NotificationRequest notification_requests = 2;
+  // max_rows: the maximum number of rows to include in the dataset (optional, defaults to 500)
+  int64 max_rows = 3;
 }
 
 // BuildDatasetResponse is the response message for manually requesting the building of a dataset for a single crawler

protos/sn13/v1/sn13_validator.proto

@@ -9,6 +9,7 @@ option go_package = "macrocosm-os/rift/constellation_api/gen/sn13/v1";
 service Sn13Service {
   // ListTopics is the RPC method for getting the top topics
   rpc ListTopics(ListTopicsRequest) returns (ListTopicsResponse);
+  rpc ValidateRedditTopic(ValidateRedditTopicRequest) returns (ValidateRedditTopicResponse);
 }
 
 // ListTopicsRequest is the request message for getting the top topics
@@ -32,3 +33,23 @@ message ListTopicsResponse {
   // message: the response message
   repeated ListTopicsResponseDetail details = 1;
 }
+
+// ValidateTopicRequest is the request message for validating a reddit topic
+message ValidateRedditTopicRequest {
+  // topic: the topic to validate
+  string topic = 1;
+}
+
+// ValidateTopicResponse is the response message for validating a topic
+message ValidateRedditTopicResponse {
+    // platform: i.e. reddit
+    string platform = 1;
+    // topic: the topic to validate
+    string topic = 2;
+    // exists: whether the topic exists
+    bool exists = 3;
+    // over18: whether the topic is NSFW
+    bool over18 = 4;
+    // quarantine: whether the topic is quarantined
+    bool quarantine = 5;
+}

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "macrocosmos"
-version = "0.1.4"
+version = "1.0.0"
 description = "The official Python SDK for Macrocosmos"
 readme = "README.md"
 license = "Apache-2.0"

src/macrocosmos/__init__.py

@@ -18,7 +18,14 @@
         __version__ = "unknown"
 
 # Import client and types
-from .client import ApexClient, AsyncApexClient, AsyncGravityClient, GravityClient
+from .client import (
+    ApexClient,
+    AsyncApexClient,
+    AsyncGravityClient,
+    GravityClient,
+    BillingClient,
+    AsyncBillingClient,
+)
 from .types import (
     ChatCompletionChunkResponse,
     ChatCompletionResponse,
@@ -33,6 +40,8 @@
     "ApexClient",
     "AsyncGravityClient",
     "GravityClient",
+    "BillingClient",
+    "AsyncBillingClient",
     "ChatMessage",
     "ChatCompletionResponse",
     "ChatCompletionChunkResponse",

src/macrocosmos/client.py

@@ -9,6 +9,7 @@
 )
 from macrocosmos.resources.gravity import AsyncGravity, SyncGravity
 from macrocosmos.resources.web_search import AsyncWebSearch, SyncWebSearch
+from macrocosmos.resources.billing import AsyncBilling, SyncBilling
 from macrocosmos.resources._client import BaseClient
 
 
@@ -181,3 +182,86 @@ def __init__(
         )
 
         self.gravity = SyncGravity(self)
+
+
+class AsyncBillingClient(BaseClient):
+    """
+    Asynchronous client for the Billing API.
+    """
+
+    def __init__(
+        self,
+        api_key: Optional[str] = None,
+        base_url: Optional[str] = None,
+        timeout: Optional[int] = None,
+        max_retries: int = 0,
+        compress: bool = True,
+        secure: Optional[bool] = None,
+        app_name: Optional[str] = None,
+    ):
+        """
+        Initialize the asynchronous Billing client.
+
+        Args:
+            api_key: The API key.
+            base_url: The base URL for the API.
+            timeout: Time to wait for a response in seconds. (default: None)
+            max_retries: The maximum number of retries. (default: 0)
+            compress: Whether to compress the request using gzip (default: True).
+            secure: Whether to use HTTPS (default: True).
+            app_name: The name of the application using the client.
+        """
+        if not api_key:
+            api_key = os.environ.get("BILLING_API_KEY")
+
+        super().__init__(
+            api_key=api_key,
+            base_url=base_url,
+            timeout=timeout,
+            max_retries=max_retries,
+            secure=secure,
+            compress=compress,
+            app_name=app_name,
+        )
+
+        self.billing = AsyncBilling(self)
+
+
+class BillingClient(BaseClient):
+    """
+    Synchronous client for the Billing API.
+    """
+
+    def __init__(
+        self,
+        api_key: Optional[str] = None,
+        base_url: Optional[str] = None,
+        timeout: Optional[int] = None,
+        max_retries: int = 0,
+        secure: Optional[bool] = None,
+        app_name: Optional[str] = None,
+    ):
+        """
+        Initialize the synchronous Billing client.
+
+        Args:
+            api_key: The API key.
+            base_url: The base URL for the API.
+            timeout: Time to wait for a response in seconds. (default: None)
+            max_retries: The maximum number of retries. (default: 0)
+            secure: Whether to use HTTPS (default: True).
+            app_name: The name of the application using the client.
+        """
+        if not api_key:
+            api_key = os.environ.get("BILLING_API_KEY")
+
+        super().__init__(
+            api_key=api_key,
+            base_url=base_url,
+            timeout=timeout,
+            max_retries=max_retries,
+            secure=secure,
+            app_name=app_name,
+        )
+
+        self.billing = SyncBilling(self)

src/macrocosmos/generated/billing/v1/billing_p2p.py

@@ -0,0 +1,47 @@
+# This is an automatically generated file, please do not change
+# gen by protobuf_to_pydantic[v0.3.1.1](https://github.com/so1n/protobuf_to_pydantic)
+# Protobuf Version: 5.29.4 
+# Pydantic Version: 2.11.0 
+from google.protobuf.message import Message  # type: ignore
+from pydantic import BaseModel
+from pydantic import Field
+import typing
+
+
+class GetUsageRequest(BaseModel):
+    """
+     GetUsageRequest is the request message for getting the usage of the user's credits
+    """
+
+# product_type: the type of the product (i.e. "gravity")
+    product_type: str = Field(default="")
+
+class BillingRate(BaseModel):
+    """
+     ProductPlan is details of the subscription plan for a product
+    """
+
+# product_type: the type of the product (i.e. "gravity")
+    rate_type: str = Field(default="")
+# unit_size: the size of the unit of the subscription (e.g. 1000, 10000, 100000)
+    unit_size: int = Field(default=0)
+# unit_type: the type of the unit of the subscription (i.e. "rows")
+    unit_type: str = Field(default="")
+# price_per_unit: the price per unit of the subscription
+    price_per_unit: float = Field(default=0.0)
+# currency: the currency of the subscription
+    currency: str = Field(default="")
+
+class GetUsageResponse(BaseModel):
+    """
+     GetUsageResponse is the response message for getting the usage of the user's credits
+    """
+
+# available_credits: the number of credits available to the user
+    available_credits: float = Field(default=0.0)
+# used_credits: the number of credits used by the user
+    used_credits: float = Field(default=0.0)
+# remaining_credits: the number of credits remaining to the user
+    remaining_credits: float = Field(default=0.0)
+# subscription: the subscription that the user has
+    billing_rates: typing.List[BillingRate] = Field(default_factory=list)