Skip to content

Commit 858db81

Browse files
authored
Merge pull request #11 from macrocosm-os/gen_1310_deep_researcher_endpoints_test
GEN-1310: Adding Deep Researcher functionality to Macrocosmos Python SDK
2 parents 338d12e + bdf0a57 commit 858db81

File tree

18 files changed

+998
-133
lines changed

18 files changed

+998
-133
lines changed

examples/apex_chat_deep_research.py

Lines changed: 0 additions & 66 deletions
This file was deleted.
Lines changed: 122 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,122 @@
1+
"""
2+
Example of using the Apex DeepResearch API asynchronously with Macrocosmos SDK.
3+
Demonstrates how a deep researcher job can be polled at regular intervals
4+
to check its current status and retrieve the latest results generated.
5+
"""
6+
7+
import asyncio
8+
import os
9+
import json
10+
from typing import Optional, Any, List
11+
12+
import macrocosmos as mc
13+
14+
15+
def extract_content_from_chunk(chunk_str: str) -> Optional[str]:
16+
"""Extract content from a JSON chunk string if available."""
17+
try:
18+
chunk_list = json.loads(chunk_str)
19+
if chunk_list and len(chunk_list) > 0 and "content" in chunk_list[0]:
20+
return chunk_list[0]["content"]
21+
except (json.JSONDecodeError, IndexError, KeyError) as e:
22+
print(f"Failed to parse chunk: {e}")
23+
return None
24+
25+
26+
async def process_result_chunks(results: List[Any], last_seq_id: int) -> int:
27+
"""Process result chunks and return the new last sequence ID."""
28+
for item in results:
29+
try:
30+
seq_id = int(item.seq_id)
31+
if seq_id > last_seq_id:
32+
if content := extract_content_from_chunk(item.chunk):
33+
print(f"\nseq_id {seq_id}:\n{content}")
34+
last_seq_id = seq_id
35+
except (ValueError, AttributeError) as e:
36+
print(f"Error processing sequence: {e}")
37+
return last_seq_id
38+
39+
40+
async def demo_deep_research_polling():
41+
"""Demo asynchronous deep research job creation and update polling."""
42+
print("\nRunning asynchronous Deep Research example...")
43+
44+
api_key = os.environ.get("APEX_API_KEY", os.environ.get("MACROCOSMOS_API_KEY"))
45+
46+
client = mc.AsyncApexClient(
47+
api_key=api_key, app_name="examples/apex_deep_research_polling.py"
48+
)
49+
50+
# Create a deep research job with create_job
51+
submitted_response = await client.deep_research.create_job(
52+
messages=[
53+
{
54+
"role": "user",
55+
"content": """Can you propose a mechanism by which a decentralized network
56+
of AI agents could achieve provable alignment on abstract ethical principles
57+
without relying on human-defined ontologies or centralized arbitration?""",
58+
}
59+
]
60+
)
61+
62+
print("\nCreated deep research job.\n")
63+
print(f"Initial status: {submitted_response.status}")
64+
print(f"Job ID: {submitted_response.job_id}")
65+
print(f"Created at: {submitted_response.created_at}\n")
66+
67+
# Poll for job status with get_job_results based on the job_id
68+
print("Polling the results...")
69+
last_seq_id = -1 # Track the highest sequence ID we've seen
70+
last_updated = None # Track the last update time
71+
while True:
72+
try:
73+
polled_response = await client.deep_research.get_job_results(
74+
submitted_response.job_id
75+
)
76+
current_status = polled_response.status
77+
current_updated = polled_response.updated_at
78+
79+
# On completion, print the final answer and its sequence ID
80+
if current_status == "completed":
81+
print("\nJob completed successfully!")
82+
print(f"\nLast update at: {current_updated}")
83+
if polled_response.result:
84+
if content := extract_content_from_chunk(
85+
polled_response.result[-1].chunk
86+
):
87+
print(
88+
f"\nFinal answer (seq_id {polled_response.result[-1].seq_id}):\n{content}"
89+
)
90+
break
91+
92+
elif current_status == "failed":
93+
print(
94+
f"\nJob failed: {polled_response.error if hasattr(polled_response, 'error') else 'Unknown error'}"
95+
)
96+
print(f"\nLast update at: {current_updated}")
97+
break
98+
99+
# Check if we have new content by comparing update times
100+
if current_updated != last_updated:
101+
print(f"\nNew update at {current_updated}")
102+
print(f"Status: {current_status}")
103+
104+
# Process new content
105+
if polled_response.result:
106+
last_seq_id = await process_result_chunks(
107+
polled_response.result, last_seq_id
108+
)
109+
else:
110+
print(
111+
"No results available yet. Waiting for Deep Researcher to generate data..."
112+
)
113+
last_updated = current_updated
114+
115+
except Exception as e:
116+
print(f"Error during polling: {e}")
117+
118+
await asyncio.sleep(20) # Poll in 20 second intervals
119+
120+
121+
if __name__ == "__main__":
122+
asyncio.run(demo_deep_research_polling())

examples/sn13_on_demand_data.py

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -8,9 +8,7 @@
88

99
api_key = os.environ.get("SN13_API_KEY", os.environ.get("MACROCOSMOS_API_KEY"))
1010

11-
client = mc.Sn13Client(
12-
api_key=api_key, app_name="examples/sn13_on_demand_basic.py", secure=False
13-
)
11+
client = mc.Sn13Client(api_key=api_key, app_name="examples/sn13_on_demand_data.py")
1412

1513
response = client.sn13.OnDemandData(
1614
source="x",

protos/apex/v1/apex.proto

Lines changed: 93 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,9 @@ syntax = "proto3";
33
package apex.v1;
44

55
option go_package = "macrocosm-os/rift/constellation_api/gen/apex/v1";
6+
import "google/protobuf/timestamp.proto";
7+
import "google/protobuf/empty.proto";
8+
69

710
service ApexService {
811
// ChatCompletion generates a completion for a given request.
@@ -13,6 +16,18 @@ service ApexService {
1316

1417
// WebRetrieval retrieves web search results for a given request.
1518
rpc WebRetrieval(WebRetrievalRequest) returns (WebRetrievalResponse);
19+
20+
// SubmitDeepResearcherJob submits a new deep researcher job for processing.
21+
rpc SubmitDeepResearcherJob(ChatCompletionRequest) returns (SubmitDeepResearcherJobResponse);
22+
23+
// GetDeepResearcherJob retrieves the status and results of a deep researcher job.
24+
rpc GetDeepResearcherJob(GetDeepResearcherJobRequest) returns (GetDeepResearcherJobResponse);
25+
26+
// GetChatSessions retrieves a user's chats
27+
rpc GetChatSessions(google.protobuf.Empty) returns (GetChatSessionsResponse);
28+
29+
// UpdateChatAttribute updates attribute after asking LLM for one
30+
rpc UpdateChatAttribute(UpdateChatAttributeRequest) returns (UpdateChatAttributeResponse);
1631
}
1732

1833
// A request to generate completions following Apex CompletionsRequest format.
@@ -322,7 +337,7 @@ message PromptTokensDetails {
322337
int64 cached_tokens = 2;
323338
}
324339

325-
// A web retrival request from Apex
340+
// A web retrieval request from Apex
326341
// Parsed from https://github.com/macrocosm-os/prompting/blob/main/validator_api/serializers.py
327342
message WebRetrievalRequest {
328343
// uids: the miner UIDs that will be used to generate the completion (optional).
@@ -356,3 +371,80 @@ message WebRetrievalResponse {
356371
// results: the results of the web retrieval.
357372
repeated WebSearchResult results = 1;
358373
}
374+
375+
// A response containing the deep researcher job submission details
376+
message SubmitDeepResearcherJobResponse {
377+
// job_id: unique identifier for the submitted job
378+
string job_id = 1;
379+
// status: current status of the job
380+
string status = 2;
381+
// created_at: timestamp when the job was created
382+
string created_at = 3;
383+
// updated_at: timestamp when the job was last updated
384+
string updated_at = 4;
385+
}
386+
387+
// A response containing the deep researcher job status and results
388+
message GetDeepResearcherJobResponse {
389+
// job_id: unique identifier for the job
390+
string job_id = 1;
391+
// status: current status of the job
392+
string status = 2;
393+
// created_at: timestamp when the job was created
394+
string created_at = 3;
395+
// updated_at: timestamp when the job was last updated
396+
string updated_at = 4;
397+
// result: array of result chunks
398+
repeated DeepResearcherResultChunk result = 5;
399+
// error: error message if the job failed
400+
optional string error = 6;
401+
}
402+
403+
// A chunk of the deep researcher result
404+
message DeepResearcherResultChunk {
405+
// seq_id: sequence identifier for the chunk
406+
int64 seq_id = 1;
407+
// chunk: the content of the chunk
408+
string chunk = 2;
409+
}
410+
411+
// A request to get the status of a deep researcher job
412+
message GetDeepResearcherJobRequest {
413+
// job_id: the ID of the job to retrieve
414+
string job_id = 1;
415+
}
416+
417+
// A GetChatSession message repeated in GetChatSessionsResponse
418+
message ChatSession {
419+
// id: chat id
420+
string id = 1;
421+
// user_id: user id
422+
string user_id = 2;
423+
// title: title of chat
424+
string title = 3;
425+
// chat_type: e.g. apex
426+
string chat_type = 4;
427+
// created_at: when the chat was created
428+
google.protobuf.Timestamp created_at = 5;
429+
// updated_at: when the chat was updated
430+
google.protobuf.Timestamp updated_at = 6;
431+
}
432+
433+
// A GetChatSessionsResponse response
434+
message GetChatSessionsResponse {
435+
// chat_sessions: the chat sessions
436+
repeated ChatSession chat_sessions = 1;
437+
}
438+
439+
// Directly model the attributes as a map
440+
message UpdateChatAttributeRequest {
441+
// chat_id: the unique id associated to a users chat message
442+
string chat_id = 1;
443+
// attributes: the data attributes captured in the chat logging process
444+
map<string, string> attributes = 2;
445+
}
446+
447+
message UpdateChatAttributeResponse {
448+
// success: indicates if an attribute update was successful
449+
bool success = 1;
450+
}

protos/gravity/v1/gravity.proto

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,9 @@ service GravityService {
2727

2828
// Cancel dataset build if it is in progress and purges the dataset
2929
rpc CancelDataset(CancelDatasetRequest) returns (CancelDatasetResponse);
30+
31+
// Refund user if fewer rows are returned
32+
rpc DatasetBillingCorrection(DatasetBillingCorrectionRequest) returns (DatasetBillingCorrectionResponse);
3033
}
3134

3235
// Crawler is a single crawler workflow that registers a single job (platform/topic) on SN13's dynamic desirability engine
@@ -287,3 +290,17 @@ message CancelDatasetResponse {
287290
// message: the message of the cancellation of the dataset build (currently hardcoded to "success")
288291
string message = 1;
289292
}
293+
294+
// DatasetBillingCorrectionRequest is the request message for refunding a user
295+
message DatasetBillingCorrectionRequest {
296+
// requested_row_count: number of rows expected by the user
297+
int64 requested_row_count = 1;
298+
// actual_row_count: number of rows returned by gravity
299+
int64 actual_row_count = 2;
300+
}
301+
302+
// DatasetBillingCorrectionResponse is the response message for refunding a user
303+
message DatasetBillingCorrectionResponse {
304+
// refund_amount
305+
double refund_amount = 1;
306+
}

pyproject.toml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
[project]
22
name = "macrocosmos"
3-
version = "1.0.2"
3+
version = "1.0.3"
44
description = "The official Python SDK for Macrocosmos"
55
readme = "README.md"
66
license = "Apache-2.0"

src/macrocosmos/client.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111
from macrocosmos.resources.sn13 import AsyncSn13, SyncSn13
1212
from macrocosmos.resources.web_search import AsyncWebSearch, SyncWebSearch
1313
from macrocosmos.resources.billing import AsyncBilling, SyncBilling
14+
from macrocosmos.resources.deep_research import SyncDeepResearch, AsyncDeepResearch
1415
from macrocosmos.resources._client import BaseClient
1516

1617

@@ -57,6 +58,7 @@ def __init__(
5758
self.chat = AsyncChat(self)
5859
self.completions = AsyncCompletions(self)
5960
self.web_search = AsyncWebSearch(self)
61+
self.deep_research = AsyncDeepResearch(self)
6062

6163

6264
class ApexClient(BaseClient):
@@ -99,6 +101,7 @@ def __init__(
99101
self.chat = SyncChat(self)
100102
self.completions = SyncCompletions(self)
101103
self.web_search = SyncWebSearch(self)
104+
self.deep_research = SyncDeepResearch(self)
102105

103106

104107
class AsyncGravityClient(BaseClient):

0 commit comments

Comments
 (0)