GEN-1282: SN13 on demand data streaming endpoint integration into SDK

examples/sn13_on_demand_data.py

@@ -0,0 +1,24 @@
+"""
+Example of using the SN13 On Demand Data Streaming service with Macrocosmos SDK.
+"""
+
+import os
+
+import macrocosmos as mc
+
+api_key = os.environ.get("SN13_API_KEY", os.environ.get("MACROCOSMOS_API_KEY"))
+
+client = mc.Sn13Client(
+    api_key=api_key, app_name="examples/sn13_on_demand_basic.py", secure=False
+)
+
+response = client.sn13.OnDemandData(
+    source="x",
+    usernames=["nasa", "spacex"],
+    keywords=["photo", "space", "mars"],
+    start_date="2024-04-01",
+    end_date="2025-04-25",
+    limit=3,
+)
+
+print(response)

protos/README.md

@@ -1,7 +1,9 @@
 # Macrocosmos Protobufs
+
 Protobufs used by Macrocosmos products
 
-# Usage
+## Usage
+
 Here are instructions on how to add and update the protos from your own project repo.
 
 > ⚠️ These commands must be run from the target repo's root directory

protos/apex/v1/apex.proto

@@ -23,25 +23,25 @@ message ChatCompletionRequest {
   // messages: the messages to generate completions for.
   repeated ChatMessage messages = 2;
   // seed: the seed to use for the completion.
-  int64 seed = 3;
+  optional int64 seed = 3;
   // task: the task to generate completions for (e.g. "InferenceTask").
-  string task = 4;
+  optional string task = 4;
   // model: the LLM name to use for the completion. (optional, suggest leaving this empty as not all LLMs are supported)
-  string model = 5;
+  optional string model = 5;
   // test_time_inference: whether to use test time inference.
-  bool test_time_inference = 6;
+  optional bool test_time_inference = 6;
   // mixture: whether to use a mixture of miners to create a slower but better answer.
-  bool mixture = 7;
+  optional bool mixture = 7;
   // sampling_parameters: the sampling parameters to use for the completion.
-  SamplingParameters sampling_parameters = 8;
+  optional SamplingParameters sampling_parameters = 8;
   // inference_mode: the inference mode to use for the completion.
-  string inference_mode = 9;
+  optional string inference_mode = 9;
   // json_format: whether to use JSON format for the completion.
-  bool json_format = 10;
+  optional bool json_format = 10;
   // stream: whether to stream the completion.
-  bool stream = 11;
+  optional bool stream = 11;
   // timeout: the timeout for the completion in seconds.
-  int64 timeout = 12;
+  optional int64 timeout = 12;
 }
 
 // The sampling parameters for the completion.
@@ -52,7 +52,7 @@ message SamplingParameters {
   // top_p: the top_p to use for the completion.
   double top_p = 2;
   // top_k: the top_k to use for the completion.
-  double top_k = 3;
+  optional double top_k = 3;
   // max_new_tokens: the max_new_tokens to use for the completion.
   int64 max_new_tokens = 4;
   // do_sample: whether to do sample for the completion.
@@ -330,13 +330,13 @@ message WebRetrievalRequest {
   // search_query: the search query.
   string search_query = 2;
   // n_miners: the number of miners to use for the query.
-  int64 n_miners = 3;
+  optional int64 n_miners = 3;
   // n_results: the number of results to return.
-  int64 n_results = 4;
+  optional int64 n_results = 4;
   // max_response_time: the max response time to allow for the miners to respond in seconds.
-  int64 max_response_time = 5;
+  optional int64 max_response_time = 5;
   // timeout: the timeout for the web retrieval in seconds.
-  int64 timeout = 6;
+  optional int64 timeout = 6;
 }
 
 // A web search result from Apex

protos/billing/v1/billing.proto

@@ -12,7 +12,7 @@ service BillingService {
 // 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;
+  optional string product_type = 1;
 }
 
 // ProductPlan is details of the subscription plan for a product

protos/gravity/v1/gravity.proto

@@ -115,9 +115,9 @@ message GravityTaskState {
 // GetGravityTasksRequest is the request message for listing gravity tasks for a user
 message GetGravityTasksRequest {
   // gravity_task_id: the ID of the gravity task (optional, if not provided, all gravity tasks for the user will be returned)
-  string gravity_task_id = 1;
+  optional string gravity_task_id = 1;
   // include_crawlers: whether to include the crawler states in the response
-  bool include_crawlers = 2;
+  optional bool include_crawlers = 2;
 }
 
 // GetGravityTasksResponse is the response message for listing gravity tasks for a user
@@ -141,7 +141,7 @@ message NotificationRequest {
   // address: the address to send the notification to (only email addresses are supported currently)
   string address = 2;
   // redirect_url: the URL to include in the notication message that redirects the user to any built datasets
-  string redirect_url = 3;
+  optional string redirect_url = 3;
 }
 
 // GetCrawlerRequest is the request message for getting a crawler
@@ -166,7 +166,7 @@ message CreateGravityTaskRequest {
   //   that is automatically generated upon completion of the crawler is ready to download (optional)
   repeated NotificationRequest notification_requests = 3;
   // gravity_task_id: the ID of the gravity task (optional, default will generate a random ID)
-  string gravity_task_id = 4;
+  optional string gravity_task_id = 4;
 }
 
 // CreateGravityTaskResponse is the response message for creating a new gravity task
@@ -194,6 +194,15 @@ message BuildDatasetResponse {
   Dataset dataset = 2;
 }
 
+message Nebula {
+  // error: nebula build error message
+  string error = 1;
+  // file_size_bytes: the size of the file in bytes
+  int64 file_size_bytes = 2;
+  // url: the URL of the file
+  string url = 3;
+}
+
 // Dataset contains the progress and results of a dataset build
 message Dataset {
   // crawler_workflow_id: the ID of the parent crawler for this dataset
@@ -212,6 +221,8 @@ message Dataset {
   repeated DatasetStep steps = 7;
   // total_steps: the total number of steps in the dataset build
   int64 total_steps = 8;
+  // nebula: the details about the nebula that was built
+  Nebula nebula = 9;
 }
 
 // DatasetFile contains the details about a dataset file

protos/sn13/v1/sn13_validator.proto

@@ -1,7 +1,7 @@
 syntax = "proto3";
 
 package sn13.v1;
-
+import "google/protobuf/struct.proto";
 
 option go_package = "macrocosm-os/rift/constellation_api/gen/sn13/v1";
 
@@ -10,6 +10,9 @@ service Sn13Service {
   // ListTopics is the RPC method for getting the top topics
   rpc ListTopics(ListTopicsRequest) returns (ListTopicsResponse);
   rpc ValidateRedditTopic(ValidateRedditTopicRequest) returns (ValidateRedditTopicResponse);
+
+  // Access the SN13 API endpoint on_demand_data_request via Constellation
+  rpc OnDemandData(OnDemandDataRequest) returns (OnDemandDataResponse);
 }
 
 // ListTopicsRequest is the request message for getting the top topics
@@ -53,3 +56,29 @@ message ValidateRedditTopicResponse {
     // quarantine: whether the topic is quarantined
     bool quarantine = 5;
 }
+
+// OnDemandDataRequest is a request to SN13 to retrieve data
+message OnDemandDataRequest {
+  // source: the data source (X or Reddit)
+  string source = 1;
+  // usernames: list of usernames to fetch data from
+  repeated string usernames = 2;
+  // keywords: list of keywords to search for
+  repeated string keywords = 3;
+  // start_date: ISO 8601 formatted date string (e.g. "2024-01-01T00:00:00Z")
+  optional string start_date = 4;
+  // end_date: ISO 8601 formatted date string (e.g. "2024-01-31T23:59:59Z")
+  optional string end_date = 5;
+  // limit: maximum number of results to return
+  optional int64 limit = 6;
+}
+
+// OnDemandDataResponse is the response from SN13 for an on-demand data request
+message OnDemandDataResponse {
+  // status: the request status, either success/error
+  string status = 1;
+  // data: the data object returned
+  repeated google.protobuf.Struct data = 2;
+  // meta: additional metadata about the request
+  google.protobuf.Struct meta = 3;
+}

pyproject.toml

@@ -12,6 +12,7 @@ dependencies = [
     "grpcio>=1.70.0",
     "grpcio-tools>=1.70.0",
     "pydantic>=2.10.6",
+    "mypy-protobuf (>=3.6.0,<4.0.0)",
 ]
 
 [build-system]

src/macrocosmos/__init__.py

@@ -25,6 +25,8 @@
     GravityClient,
     BillingClient,
     AsyncBillingClient,
+    Sn13Client,
+    AsyncSn13Client,
 )
 from .types import (
     ChatCompletionChunkResponse,
@@ -47,4 +49,6 @@
     "ChatCompletionChunkResponse",
     "SamplingParameters",
     "WebRetrievalResponse",
+    "Sn13Client",
+    "AsyncSn13Client",
 ]

src/macrocosmos/client.py

@@ -8,6 +8,7 @@
     SyncCompletions,
 )
 from macrocosmos.resources.gravity import AsyncGravity, SyncGravity
+from macrocosmos.resources.sn13 import AsyncSn13, SyncSn13
 from macrocosmos.resources.web_search import AsyncWebSearch, SyncWebSearch
 from macrocosmos.resources.billing import AsyncBilling, SyncBilling
 from macrocosmos.resources._client import BaseClient
@@ -265,3 +266,90 @@ def __init__(
         )
 
         self.billing = SyncBilling(self)
+
+
+class AsyncSn13Client(BaseClient):
+    """
+    Asynchronous client for the SN13 API service.
+    """
+
+    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 SN13 API 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("SN13_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.sn13 = AsyncSn13(self)
+
+
+class Sn13Client(BaseClient):
+    """
+    Synchronous client for the SN13 API service.
+    """
+
+    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,
+        compress: bool = True,
+        app_name: Optional[str] = None,
+    ):
+        """
+        Initialize the synchronous SN13 API 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).
+            compress: Whether to compress the request using gzip (default: True).
+            app_name: The name of the application using the client.
+        """
+        if not api_key:
+            api_key = os.environ.get("SN13_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.sn13 = SyncSn13(self)