second commit

main.py

@@ -0,0 +1,118 @@
+import asyncio
+import uuid
+import json
+
+import uvicorn
+from fastapi import FastAPI
+from aiokafka import AIOKafkaProducer, AIOKafkaConsumer
+from contextlib import asynccontextmanager
+
+KAFKA_BOOTSTRAP = "localhost:19092"
+REQUEST_TOPIC = "request-topic"
+RESPONSE_TOPIC = "response-topic"
+
+
+loop = asyncio.get_event_loop()
+
+producer = None
+consumer = None
+ml_models = {}
+pending_responses = {}  # correlation_id -> asyncio.Future
+
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    # Startup logic
+    print("============== STARTING LIFESPAN ==============")
+    producer = AIOKafkaProducer(bootstrap_servers=KAFKA_BOOTSTRAP, loop=loop)
+    await producer.start()
+
+    consumer = AIOKafkaConsumer(
+        RESPONSE_TOPIC,
+        bootstrap_servers=KAFKA_BOOTSTRAP,
+        loop=loop,
+        group_id="fastapi-group",
+        auto_offset_reset="earliest"
+    )
+    ml_models['producer'] = producer
+    ml_models['consumer'] = consumer
+    await consumer.start()
+
+    app.state.producer = producer
+    app.state.consumer = consumer
+
+    async def consume_responses():
+        async for msg in consumer:
+            try:
+                data = json.loads(msg.value.decode())
+                correlation_id = data.get("correlation_id")
+                future = pending_responses.pop(correlation_id, None)
+                if future:
+                    future.set_result(data.get("response"))
+            except Exception as e:
+                print("Error handling response:", e)
+
+    app.state.consumer_task = asyncio.create_task(consume_responses())
+
+    yield  # ← here the app runs
+
+    # Shutdown logic
+    app.state.consumer_task.cancel()
+    await producer.stop()
+    await consumer.stop()
+
+
+
+async def consume_responses():
+
+    async for msg in ml_models['consumer']:
+        try:
+            data = json.loads(msg.value.decode())
+            correlation_id = data.get("correlation_id")
+            future = pending_responses.pop(correlation_id, None)
+            if future:
+                future.set_result(data.get("response"))
+        except Exception as e:
+            print("Error handling response:", e)
+app = FastAPI(lifespan=lifespan)
+
+
+@app.post("/process/")
+async def process(data: dict):
+    correlation_id = str(uuid.uuid4())
+    print(ml_models)
+    future = loop.create_future()
+    pending_responses[correlation_id] = future
+
+    message = {
+        "correlation_id": correlation_id,
+        "payload": data
+    }
+
+    print(ml_models['producer'])
+    await ml_models['producer'].send_and_wait(REQUEST_TOPIC, json.dumps(message).encode())
+
+    try:
+        response = await asyncio.wait_for(future, timeout=10)
+        print("+++", response)
+        print("===", dir(response))
+        return {"response": response}
+    except asyncio.TimeoutError:
+        pending_responses.pop(correlation_id, None)
+        return {"error": "Timeout waiting for response"}
+
+def main():
+    logging_config = uvicorn.config.LOGGING_CONFIG
+    uvicorn.run(
+        "main:app",
+        host='0.0.0.0',
+        port=8085,
+        reload=True,
+        # debug=True,
+        # loglevel='debug',
+        log_config=logging_config,
+    )
+
+
+if __name__ == '__main__':
+    main()