Skip to content
This repository has been archived by the owner on Aug 25, 2024. It is now read-only.

Commit

Permalink
Add support for schemas in python topic producer (LangStream#734)
Browse files Browse the repository at this point in the history
  • Loading branch information
@cbornet
cbornet authored Nov 22, 2023
1 parent 2c56f08 commit a732e3e
Show file tree
Hide file tree
Showing 11 changed files with 213 additions and 92 deletions.
98 changes: 58 additions & 40 deletions ...ents/langstream-agent-grpc/src/main/java/ai/langstream/agents/grpc/AbstractGrpcAgent.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -107,47 +107,65 @@ public void start() throws Exception {
asyncStub.getTopicProducerRecords(
new StreamObserver<>() {
@Override
public void onNext(TopicProducerRecord topicProducerRecord) {
TopicProducer topicProducer =
topicProducers.computeIfAbsent(
topicProducerRecord.getTopic(),
topic -> {
TopicProducer tp =
agentContext
.getTopicConnectionProvider()
.createProducer(
agentContext
.getGlobalAgentId(),
topic,
Map.of());
tp.start();
return tp;
});
public void onNext(TopicProducerResponse topicProducerResponse) {
try {
topicProducer
.write(fromGrpc(topicProducerRecord.getRecord()))
.whenComplete(
(r, e) -> {
if (e != null) {
log.error("Error writing record", e);
sendTopicProducerWriteResult(
TopicProducerWriteResult
.newBuilder()
.setError(
e
.getMessage()));
} else {
sendTopicProducerWriteResult(
TopicProducerWriteResult
.newBuilder()
.setRecordId(
topicProducerRecord
.getRecord()
.getRecordId()));
}
});
} catch (IOException e) {
agentContext.criticalFailure(e);
if (topicProducerResponse.hasSchema()) {
serverSchemas.put(
topicProducerResponse.getSchema().getSchemaId(),
new org.apache.avro.Schema.Parser()
.parse(
topicProducerResponse
.getSchema()
.getValue()
.toStringUtf8()));
}
if (topicProducerResponse.hasRecord()
&& !"".equals(topicProducerResponse.getTopic())) {
TopicProducer topicProducer =
topicProducers.computeIfAbsent(
topicProducerResponse.getTopic(),
topic -> {
TopicProducer tp =
agentContext
.getTopicConnectionProvider()
.createProducer(
agentContext
.getGlobalAgentId(),
topic,
Map.of());
tp.start();
return tp;
});
topicProducer
.write(fromGrpc(topicProducerResponse.getRecord()))
.whenComplete(
(r, e) -> {
if (e != null) {
log.error(
"Error writing record", e);
sendTopicProducerWriteResult(
TopicProducerWriteResult
.newBuilder()
.setError(
e
.getMessage()));
} else {
sendTopicProducerWriteResult(
TopicProducerWriteResult
.newBuilder()
.setRecordId(
topicProducerResponse
.getRecord()
.getRecordId()));
}
});
}
} catch (Exception e) {
agentContext.criticalFailure(
new RuntimeException(
"getTopicProducerRecords: Error while processing TopicProducerResponse: %s"
.formatted(e.getMessage()),
e));
}
}

Expand Down
5 changes: 3 additions & 2 deletions langstream-agents/langstream-agent-grpc/src/main/proto/langstream_grpc/proto/agent.proto
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@ service AgentService {
rpc read(stream SourceRequest) returns (stream SourceResponse) {}
rpc process(stream ProcessorRequest) returns (stream ProcessorResponse) {}
rpc write(stream SinkRequest) returns (stream SinkResponse) {}
rpc get_topic_producer_records(stream TopicProducerWriteResult) returns (stream TopicProducerRecord) {}
rpc get_topic_producer_records(stream TopicProducerWriteResult) returns (stream TopicProducerResponse) {}
}

message InfoResponse {
Expand Down Expand Up @@ -74,8 +74,9 @@ message TopicProducerWriteResult {
optional string error = 2;
}

message TopicProducerRecord {
message TopicProducerResponse {
string topic = 1;
Schema schema = 2;
Record record = 3;
}

Expand Down
73 changes: 63 additions & 10 deletions .../langstream-agent-grpc/src/test/java/ai/langstream/agents/grpc/AbstractGrpcAgentTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -23,19 +23,29 @@
import ai.langstream.api.runner.topics.TopicConnectionProvider;
import ai.langstream.api.runner.topics.TopicConsumer;
import ai.langstream.api.runner.topics.TopicProducer;
import com.google.protobuf.ByteString;
import com.google.protobuf.Empty;
import io.grpc.ManagedChannel;
import io.grpc.Server;
import io.grpc.Status;
import io.grpc.inprocess.InProcessChannelBuilder;
import io.grpc.inprocess.InProcessServerBuilder;
import io.grpc.stub.StreamObserver;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.nio.file.Path;
import java.util.Map;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.TimeUnit;
import lombok.extern.slf4j.Slf4j;
import org.apache.avro.Conversions;
import org.apache.avro.SchemaBuilder;
import org.apache.avro.generic.GenericData;
import org.apache.avro.generic.GenericDatumWriter;
import org.apache.avro.generic.GenericRecord;
import org.apache.avro.io.BinaryEncoder;
import org.apache.avro.io.EncoderFactory;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
Expand All @@ -62,17 +72,47 @@ public void agentInfo(

@Override
public StreamObserver<TopicProducerWriteResult> getTopicProducerRecords(
StreamObserver<TopicProducerRecord> responseObserver) {
StreamObserver<TopicProducerResponse> responseObserver) {
org.apache.avro.Schema schema =
SchemaBuilder.record("testRecord")
.fields()
.name("testField")
.type()
.stringType()
.noDefault()
.endRecord();
GenericData.Record avroRecord = new GenericData.Record(schema);
avroRecord.put("testField", "test-string");

responseObserver.onNext(
TopicProducerRecord.newBuilder()
.setTopic("test-topic")
.setRecord(
ai.langstream.agents.grpc.Record.newBuilder()
.setRecordId(42)
TopicProducerResponse.newBuilder()
.setSchema(
Schema.newBuilder()
.setSchemaId(123)
.setValue(
Value.newBuilder()
.setStringValue("test-value1")))
ByteString.copyFromUtf8(
schema.toString()))
.build())
.build());

try {
responseObserver.onNext(
TopicProducerResponse.newBuilder()
.setTopic("test-topic")
.setRecord(
ai.langstream.agents.grpc.Record.newBuilder()
.setRecordId(42)
.setValue(
Value.newBuilder()
.setSchemaId(123)
.setAvroValue(
ByteString.copyFrom(
serializeGenericRecord(
avroRecord)))))
.build());
} catch (IOException e) {
throw new RuntimeException(e);
}
return new StreamObserver<>() {

@Override
Expand All @@ -88,7 +128,7 @@ public void onNext(TopicProducerWriteResult topicProducerWriteResult) {
responseObserver.onCompleted();
} else if (topicProducerWriteResult.getRecordId() == 42) {
responseObserver.onNext(
TopicProducerRecord.newBuilder()
TopicProducerResponse.newBuilder()
.setTopic("test-topic")
.setRecord(
ai.langstream.agents.grpc.Record
Expand Down Expand Up @@ -164,7 +204,9 @@ void testTopicProducerSuccess() throws Exception {
TestAgentContext context = new TestAgentContextSuccessful();
startProcessor(context);
LinkedBlockingQueue<Record> recordsToWrite = context.recordsToWrite;
assertEquals("test-value1", recordsToWrite.poll(5, TimeUnit.SECONDS).value().toString());
assertEquals(
"{\"testField\": \"test-string\"}",
recordsToWrite.poll(5, TimeUnit.SECONDS).value().toString());
assertEquals("test-value2", recordsToWrite.poll(5, TimeUnit.SECONDS).value().toString());
}

Expand Down Expand Up @@ -275,4 +317,15 @@ protected CompletableFuture<?> writeRecord(Record record) {
return CompletableFuture.failedFuture(new RuntimeException("test-complete"));
}
}

private static byte[] serializeGenericRecord(GenericRecord record) throws IOException {
GenericDatumWriter<GenericRecord> writer = new GenericDatumWriter<>(record.getSchema());
// enable Decimal conversion, otherwise attempting to serialize java.math.BigDecimal will
// throw ClassCastException.
writer.getData().addLogicalTypeConversion(new Conversions.DecimalConversion());
ByteArrayOutputStream oo = new ByteArrayOutputStream();
BinaryEncoder encoder = EncoderFactory.get().directBinaryEncoder(oo, null);
writer.write(record, encoder);
return oo.toByteArray();
}
}
2 changes: 1 addition & 1 deletion ...langstream-agent-grpc/src/test/java/ai/langstream/agents/grpc/GrpcAgentProcessorTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -124,7 +124,7 @@ public void onCompleted() {

@Override
public StreamObserver<TopicProducerWriteResult> getTopicProducerRecords(
StreamObserver<TopicProducerRecord> responseObserver) {
StreamObserver<TopicProducerResponse> responseObserver) {
return new StreamObserver<>() {
@Override
public void onNext(TopicProducerWriteResult topicProducerWriteResult) {}
Expand Down
2 changes: 1 addition & 1 deletion ...ents/langstream-agent-grpc/src/test/java/ai/langstream/agents/grpc/GrpcAgentSinkTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -187,7 +187,7 @@ public void onCompleted() {

@Override
public StreamObserver<TopicProducerWriteResult> getTopicProducerRecords(
StreamObserver<TopicProducerRecord> responseObserver) {
StreamObserver<TopicProducerResponse> responseObserver) {
return new StreamObserver<>() {
@Override
public void onNext(TopicProducerWriteResult topicProducerWriteResult) {}
Expand Down
2 changes: 1 addition & 1 deletion ...ts/langstream-agent-grpc/src/test/java/ai/langstream/agents/grpc/GrpcAgentSourceTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -219,7 +219,7 @@ public void onCompleted() {

@Override
public StreamObserver<TopicProducerWriteResult> getTopicProducerRecords(
StreamObserver<TopicProducerRecord> responseObserver) {
StreamObserver<TopicProducerResponse> responseObserver) {
return new StreamObserver<>() {
@Override
public void onNext(TopicProducerWriteResult topicProducerWriteResult) {}
Expand Down
9 changes: 5 additions & 4 deletions langstream-runtime/langstream-runtime-impl/src/main/python/langstream_grpc/grpc_service.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -42,8 +42,8 @@
SourceResponse,
SinkRequest,
SinkResponse,
TopicProducerRecord,
TopicProducerWriteResult,
TopicProducerResponse,
)
from langstream_grpc.proto.agent_pb2_grpc import AgentServiceServicer
from .api import Source, Sink, Processor, Record, Agent, AgentContext, TopicProducer
Expand Down Expand Up @@ -91,12 +91,13 @@ async def agent_info(self, _, __):
async def poll_topic_producer_records(self, context):
while True:
topic, record, future = await self.topic_producer_records.get()
# TODO: handle schemas
_, grpc_record = self.to_grpc_record(record)
schemas, grpc_record = self.to_grpc_record(record)
for schema in schemas:
await context.write(TopicProducerResponse(schema=schema))
self.topic_producer_record_id += 1
self.topic_producer_records_pending[self.topic_producer_record_id] = future
grpc_record.record_id = self.topic_producer_record_id
await context.write(TopicProducerRecord(topic=topic, record=grpc_record))
await context.write(TopicProducerResponse(topic=topic, record=grpc_record))

async def handle_write_results(self, context):
write_result = await context.read()
Expand Down
Loading

0 comments on commit a732e3e

Please sign in to comment.