Streaming

1
import json
2
import os
3
import sys
4
import threading
5
import time
6
import requests
7
8
BASE_URL = "https://sandbox.api.boosted.ai"
9
ACCESS_TOKEN = os.environ.get("ALFA_ACCESS_TOKEN", "your-access-token-here")
10
11
stream_connected = threading.Event()
12
stream_error = threading.Event()
13
14
def create_thread(thread_name):
15
    url = f"{BASE_URL}/v2/threads"
16
    headers = {
17
        "Authorization": f"Bearer {ACCESS_TOKEN}",
18
        "Content-Type": "application/json"
19
    }
20
    data = {"thread_name": thread_name}
21
    
22
    response = requests.post(url, headers=headers, json=data)
23
    response.raise_for_status()
24
    return response.json()["thread_id"]
25
26
def send_message(thread_id, message):
27
    url = f"{BASE_URL}/v2/threads/{thread_id}/messages"
28
    headers = {"Authorization": f"Bearer {ACCESS_TOKEN}"}
29
    data = {"message": message}
30
    
31
    print(f"\nSending message...")
32
    response = requests.post(url, headers=headers, data=data)
33
    response.raise_for_status()
34
    print("Message sent successfully!")
35
    return response.json()
36
37
def stream_thread_events(thread_id):
38
    url = f"{BASE_URL}/v2/threads/stream/{thread_id}"
39
    headers = {
40
        "Authorization": f"Bearer {ACCESS_TOKEN}",
41
        "Accept": "text/event-stream",
42
        "Cache-Control": "no-cache"
43
    }
44
    
45
    try:
46
        print("\nConnecting to event stream...")
47
        response = requests.get(url, headers=headers, stream=True, timeout=None)
48
        response.raise_for_status()
49
        
50
        print("Connected to stream, waiting for events...\n")
51
        stream_connected.set()
52
        
53
        for line in response.iter_lines(decode_unicode=True):
54
            if line:
55
                if line.startswith(':'):
56
                    continue
57
                    
58
                if line.startswith('data: '):
59
                    data = line[6:]
60
                    try:
61
                        event_data = json.loads(data)
62
                        event_type = event_data.get('event', {}).get('event_type', 'unknown')
63
                        
64
                        if event_type == 'message_chunk':
65
                            chunk = event_data['event'].get('message_chunk', {})
66
                            sys.stdout.write(chunk.get('message', ''))
67
                            sys.stdout.flush()
68
                            
69
                            if event_data['event'].get('final_chunk'):
70
                                print("\n-- final chunk --")
71
                        elif event_type == 'message':
72
                            print(f"\n\nComplete message received")
73
                        elif event_type == 'task_status' or event_type == 'execution_status':
74
                            status = event_data['event'].get('status', '')
75
                            if status == 'Complete':
76
                                print(f"\n\nTask completed!")
77
                                break
78
                        else:
79
                            print(f"\n[{event_type}]")
80
                    except json.JSONDecodeError:
81
                        print(f"Data: {data}")
82
                elif line.startswith('event: '):
83
                    event_type = line[7:]
84
                    if event_type == 'close':
85
                        print("\n\nStream closed by server")
86
                        break
87
                    
88
    except requests.exceptions.RequestException as e:
89
        print(f"\nStream error: {e}")
90
        stream_error.set()
91
    except Exception as e:
92
        print(f"\nUnexpected error: {e}")
93
        stream_error.set()
94
95
def main():
96
    try:
97
        print("Creating a new thread...")
98
        thread_id = create_thread("Research Thread")
99
        print(f"Thread created with ID: {thread_id}")
100
        
101
        message = input("\nEnter your message: ")
102
        
103
        stream_thread = threading.Thread(
104
            target=stream_thread_events, 
105
            args=(thread_id,), 
106
            daemon=True
107
        )
108
        stream_thread.start()
109
        
110
        if not stream_connected.wait(timeout=10):
111
            print("Failed to connect to stream within 10 seconds")
112
            return
113
        
114
        if stream_error.is_set():
115
            print("Stream connection failed")
116
            return
117
        
118
        time.sleep(0.5)
119
        
120
        send_message(thread_id, message)
121
        
122
        stream_thread.join(timeout=120)
123
        
124
        if stream_thread.is_alive():
125
            print("\n\nStream still active after 120 seconds, exiting...")
126
        
127
    except requests.exceptions.RequestException as e:
128
        print(f"Error: {e}")
129
        sys.exit(1)
130
    except KeyboardInterrupt:
131
        print("\n\nInterrupted by user")
132
        sys.exit(0)
133
134
if __name__ == "__main__":
135
    if ACCESS_TOKEN == "your-access-token-here":
136
        print("Please set your ALFA_ACCESS_TOKEN environment variable before running.")
137
        sys.exit(1)
138
    main()

1
syntax = "proto3";
2
3
package thread_service.v1;
4
import "google/protobuf/wrappers.proto";
5
import "google/protobuf/timestamp.proto";
6
option java_package = "com.gbi.protobuf_defs.generated.thread_service.v1";
7
option go_package = "github.com/GBI-Core/protobuf-defs/gobuild/thread_service/v1";
8
9
service ThreadService {
10
11
  /////////////////////////////////////////////////////////////////
12
  // Thread API
13
  /////////////////////////////////////////////////////////////////
14
  // Streams thread events from multiple threads
15
  rpc StreamThreadEvents(stream StreamThreadEventsRequest)
16
      returns (stream StreamThreadEventsResponse);
17
}
18
19
/////////////////////////////////////////////////////////////////
20
// Common messages
21
/////////////////////////////////////////////////////////////////
22
23
message UUID {
24
  string id = 1;
25
}
26
27
message Citation {
28
  message WebCitation {
29
    string link = 1;
30
    google.protobuf.Timestamp last_fetched_at = 2;
31
    int64 snippet_highlight_start = 3;
32
    int64 snippet_highlight_end = 4;
33
    string source_type = 5;
34
  }
35
  message NewsDevelopmentCitation {
36
    int64 num_articles = 1;
37
    google.protobuf.Timestamp development_start = 2;
38
  }
39
  message NewsArticleCitation {
40
    string link = 1;
41
    UUID article_id = 2;
42
  }
43
  message DocumentCitation {
44
    int64 snippet_highlight_start = 1;
45
    int64 snippet_highlight_end = 2;
46
    bool is_snippet = 3;
47
  }
48
  message FiscalQuarterEarningsCitation {
49
    int64 snippet_highlight_start = 1;
50
    int64 snippet_highlight_end = 2;
51
    bool is_snippet = 3;
52
    google.protobuf.Timestamp call_datetime = 4;
53
    google.protobuf.Timestamp report_datetime = 5;
54
    google.protobuf.Int32Value fiscal_year = 6;
55
    google.protobuf.Int32Value fiscal_quarter = 7;
56
  }
57
  message VariableCitation {
58
    enum VariablePeriodicity {
59
      VARIABLE_PERIODICITY_UNSPECIFIED = 0;
60
      VARIABLE_PERIODICITY_DAILY = 1;
61
      VARIABLE_PERIODICITY_MONTHLY = 2;
62
      VARIABLE_PERIODICITY_QUARTERLY = 3;
63
      VARIABLE_PERIODICITY_YEARLY = 4;
64
    }
65
66
    message FunctionMetadata {
67
      string name = 1;
68
      string description = 2;
69
    }
70
71
    message VariableMetadata {
72
      string name = 1;
73
      string description = 2;
74
      string var_id = 3;
75
    }
76
77
    string description = 1;
78
    string data_provider = 2;
79
    google.protobuf.Timestamp as_of_timestamp = 3;
80
    google.protobuf.Timestamp date_range_start = 4;
81
    google.protobuf.Timestamp date_range_end = 5;
82
    string period_start = 6;
83
    string period_end = 7;
84
    string formula = 8;
85
    repeated FunctionMetadata functions_used = 9;
86
    repeated VariableMetadata components_used = 10;
87
    VariablePeriodicity variable_periodicity = 11;
88
  }
89
90
  UUID citation_id = 1;
91
  string name = 2;
92
  string citation_type = 3;
93
  google.protobuf.Timestamp last_updated_at = 4;
94
  string summary = 8;
95
96
  oneof citation {
97
    WebCitation web_citation = 5;
98
    NewsDevelopmentCitation news_development_citation = 6;
99
    DocumentCitation document_citation = 7;
100
    NewsArticleCitation news_article_citation = 9;
101
    FiscalQuarterEarningsCitation fiscal_quarter_earnings_citation = 10;
102
    VariableCitation variable_citation = 11;
103
  }
104
}
105
106
message ThreadData {
107
  enum ErrorType {
108
    ERROR_TYPE_UNKNOWN = 0;
109
    ERROR_TYPE_INTERNAL = 1;
110
  }
111
112
  message MessageEvent {
113
    UUID message_id  = 1;
114
    string message = 2;
115
    bool is_user_message = 3;
116
    google.protobuf.Timestamp message_time = 4;
117
    repeated Citation citations = 5;
118
  }
119
120
  message MessageChunkEvent {
121
    UUID message_id  = 1;
122
    string message = 2;
123
    bool final_chunk = 3;
124
    repeated Citation citations = 4;
125
  }
126
127
  message WorkInProgressEvent {
128
    string log = 1;
129
  }
130
131
  message ThreadNameEvent {
132
    string thread_name = 1;
133
  }
134
135
  message WorkerErrorEvent {
136
    ErrorType error_type = 1;
137
  }
138
139
  message MessageCancelledEvent {
140
    string cancellation_message = 1;
141
  }
142
143
  oneof payload {
144
    MessageEvent message_event = 1;
145
    MessageChunkEvent message_chunk_event = 2;
146
    WorkInProgressEvent work_in_progress_event = 3;
147
    ThreadNameEvent thread_name_event = 4;
148
    WorkerErrorEvent worker_error_event = 5;
149
    MessageCancelledEvent message_cancelled_event = 6;
150
  }
151
}
152
153
/////////////////////////////////////////////////////////////////
154
// REQUEST/RESPONSE MESSAGES
155
/////////////////////////////////////////////////////////////////
156
157
message StreamThreadEventsRequest {
158
  // The token for the user that is requesting access to the thread
159
  string user_token = 1;
160
  // The thread id the user is requesting access to
161
  UUID thread_id = 2;
162
  bool subscribe = 3;
163
}
164
165
message StreamThreadEventsResponse {
166
  // The user(s) that requested access to the thread
167
  repeated UUID user_ids = 1;
168
  // The thread id the data corresponds to
169
  UUID thread_id = 2;
170
  // The data for the thread
171
  ThreadData thread_data = 3;
172
}

1
import time
2
from collections.abc import Iterator
3
4
import grpc
5
6
from thread_service.v1 import thread_service_pb2, thread_service_pb2_grpc
7
8
9
def test_streaming():
10
    # Configuration - update these
11
    GRPC_SERVER = "sandbox.api.boosted.ai:443"
12
    SERVICE_ACCOUNT_TOKEN = "<YOUR SERVICE ACCOUNT TOKEN>"
13
    THREAD_ID = "<YOUR THREAD ID>"
14
    USER_TOKEN = "<YOUR USER TOKEN>"
15
16
    # Connect
17
    credentials = grpc.ssl_channel_credentials()
18
    options = [
19
        ('grpc.keepalive_time_ms', 10000),
20
        ('grpc.keepalive_timeout_ms', 5000),
21
        ('grpc.keepalive_permit_without_calls', True),
22
        ('grpc.http2.max_pings_without_data', 0),
23
        ('grpc.http2.min_time_between_pings_ms', 10000),
24
        ('grpc.http2.min_ping_interval_without_data_ms', 5000),
25
    ]
26
    channel = grpc.secure_channel(GRPC_SERVER, credentials, options=options)
27
    stub = thread_service_pb2_grpc.ThreadServiceStub(channel)
28
    metadata = [("authorization", f"Bearer {SERVICE_ACCOUNT_TOKEN}")]
29
30
    print(f"Connecting to {GRPC_SERVER}")
31
32
    def request_generator():
33
        """Send subscription request."""
34
        yield thread_service_pb2.StreamThreadEventsRequest(
35
            thread_id=thread_service_pb2.UUID(id=THREAD_ID),
36
            user_token=USER_TOKEN,
37
            subscribe=True,
38
        )
39
        # Keep stream alive
40
        while True:
41
            time.sleep(60)
42
43
    try:
44
        # Start streaming
45
        stream: Iterator[thread_service_pb2.StreamThreadEventsResponse] = stub.StreamThreadEvents(
46
            request_generator(), metadata=metadata
47
        )
48
        print(f"Subscribed to thread {THREAD_ID}")
49
        print("Listening for events... (Ctrl+C to stop)")
50
51
        # Receive events
52
        for response in stream:
53
            print("\nEvent received!")
54
            print(f"   Thread: {response.thread_id.id}")
55
            print(f"   Users: {[u.id for u in response.user_ids]}")
56
            print(f"   Data: {response.thread_data}")
57
58
    except KeyboardInterrupt:
59
        print("\nStopped")
60
    except grpc.RpcError as e:
61
        print(f"Error: {e.code()}: {e.details()}")
62
    finally:
63
        channel.close()
64
65
66
if __name__ == "__main__":
67
    test_streaming()