debug

Consumer.py

@@ -0,0 +1,82 @@
+# Consumer
+import time
+import pika
+import os
+from Server import get_response
+import json
+from agent.agent_graph.StateTasks import ProblemState
+import argparse
+argparse_model = argparse.ArgumentParser()
+argparse_model.add_argument("--id", type=int, default=0, help="Consumer ID")
+
+
+
+consumer_id = argparse_model.parse_args().id
+
+
+RABBITMQ_URL = os.environ["RABBITMQ_URL"]
+QUEUE_NAME = os.environ["QUEUE_NAME"] 
+
+
+
+def model_call(request):
+    # fill with last state 
+    try:
+        state = json.loads(request['last_state']) 
+    except Exception:
+        state: ProblemState = {
+            "question": request['prompt'],
+            "memory": request['memory']
+        }
+
+    print(f"MODEL CALL WITH STATE {state} and PROMPT {request['prompt']} and MEMORY {request['memory']} and HT_TOKEN {request['ht_token']} and USER_EMAIL {request['user_email']} and USER_NAME {request['user_name']}")
+    request['ht_token'] ="hf_" + request['ht_token']
+    #answer = get_response(request['prompt'], request['memory'],request['ht_token'],state,request['user_email'],request['user_name'])
+    
+    #print(f"ANSWER {answer}")
+    # drop unserlizable keys
+    #for k in ["llm","rag_model"]:
+    #    answer[k] = ""
+
+
+    #return {"Data": answer}
+
+def get_connection():
+    params = pika.URLParameters(RABBITMQ_URL)
+    return pika.BlockingConnection(params)
+
+def callback(ch, method, properties, body):
+    recieved_msg = json.loads(body.decode())
+    print("-------------------------------------------------")
+    print(f"MSG AT CONSUMER {consumer_id}" )
+
+    # simulate processing
+    print(f"TYPE {type(recieved_msg)}, CONTENT {recieved_msg}")
+    model_call(recieved_msg)
+
+    # (put your logic here)
+    print(f"CONSUMER {consumer_id}:::: Processing done")
+
+    ch.basic_ack(delivery_tag=method.delivery_tag)
+
+def start_consumer():
+    # when scalled each server has consumer
+    params = pika.URLParameters(RABBITMQ_URL)
+    connection = pika.BlockingConnection(params)
+    channel = connection.channel()
+
+    channel.queue_declare(queue=QUEUE_NAME, durable=True)
+
+    channel.basic_qos(prefetch_count=1)
+
+    channel.basic_consume(
+        queue=QUEUE_NAME,
+        on_message_callback=callback
+    )
+
+    print("Waiting for messages...")
+    channel.start_consuming()
+
+if __name__ == "__main__":
+    print(f"Starting New Consumer {consumer_id}...")
+    start_consumer()