ubi integration for agent

RAG/bedrock_agent.py

@@ -16,6 +16,9 @@ import utilities.invoke_models as invoke_models
 import streamlit as st
 import time as t
 import botocore.exceptions
+from datetime import datetime, timezone
+import botocore
+import utilities.ubi_lambda as ubi
 
 if "inputs_" not in st.session_state:
     st.session_state.inputs_ = {}
@@ -36,6 +39,68 @@ bedrock_agent_runtime_client = boto3.client(
 enable_trace:bool = True
 end_session:bool = False
 
+def now_rfc3339():
+    return datetime.now(timezone.utc).isoformat()
+
+def send_otel_span(span):
+    try:
+
+        query_payload = {
+        "client_id": session_id,
+        "query_id": st.session_state["query_id"],
+        "application": "Semantic Search",
+        "query_response_hit_ids": doc_ids,
+        "timestamp": datetime.utcnow().isoformat() + "Z",
+        "user_query": json.dumps(hybrid_payload),
+        "query": query,
+        
+        }
+
+        status = ubi.send_to_lambda(".otel-v1-apm-span-default", span)
+
+        if status == 202:
+            print("Traces sent to Lambda")
+        else:
+            print("Lambda did not accept the request")
+
+        res = requests.post(OPENSEARCH_ENDPOINT, json=span, auth=OPENSEARCH_AUTH, timeout=3)
+        print(f"[OTEL SPAN] {span['name']} -> {res.status_code}")
+    except Exception as e:
+        print(f"[OTEL ERROR] {e}")
+
+def convert_to_span(block, trace_id, index):
+    span_id = str(uuid.uuid4()).replace("-", "")[:16]
+    name = "step"
+    attributes = {}
+
+    if "invocationInput" in block:
+        name = block["invocationInput"].get("function", "invocation")
+        attributes = {p["name"]: p["value"] for p in block["invocationInput"].get("parameters", [])}
+    elif "observation" in block:
+        name = block["observation"].get("type", "observation").lower()
+        attributes = block["observation"].get("actionGroupInvocationOutput", {})
+    elif "thinking" in block:
+        name = "thinking"
+        attributes["message"] = block["thinking"].get("content", "")
+    elif "rationale" in block:
+        name = "rationale"
+        attributes["message"] = block["rationale"]
+
+    return {
+        "traceId": trace_id,
+        "spanId": span_id,
+        "name": name,
+        "startTime": now_rfc3339(),
+        "endTime": now_rfc3339(),
+        "durationInNanos": 10000000 * (index + 1),
+        "kind": "INTERNAL",
+        "status": {"code": "OK"},
+        "attributes": attributes,
+        "resource": {
+            "service.name": "bedrock-agent"
+        }
+    }
+
 def delete_memory():
     response = bedrock_agent_runtime_client.delete_agent_memory(
     agentAliasId='DEEEEZM2TM',
@@ -61,13 +126,11 @@ def query_(inputs):
     last_tool = ""
     last_tool_name = ""
     agent_answer = ""
+    trace_id = str(uuid.uuid4()).replace("-", "")
     try:
         for event in event_stream:
-            #print("***event*********")
-            #print(event)
             if 'trace' in event: 
-                #print("trace*****total*********")
-                #print(event['trace'])
+               
                 if('orchestrationTrace' not in event['trace']['trace']):
                     continue
                 orchestration_trace = event['trace']['trace']['orchestrationTrace']
@@ -80,10 +143,8 @@ def query_(inputs):
                     total_context_item['invocationInput'] = orchestration_trace['invocationInput']['actionGroupInvocationInput']
                     last_tool_name = total_context_item['invocationInput']['function']
                 if('observation'  in orchestration_trace):
-                    #print("trace****observation******")
                     total_context_item['observation'] = event['trace']['trace']['orchestrationTrace']['observation']
                     tool_output_last_obs = event['trace']['trace']['orchestrationTrace']['observation']
-                    #print(tool_output_last_obs)
                     if(tool_output_last_obs['type'] == 'ACTION_GROUP'):
                         last_tool = tool_output_last_obs['actionGroupInvocationOutput']['text']
                     if(tool_output_last_obs['type'] == 'FINISH'):   
@@ -92,8 +153,12 @@ def query_(inputs):
                     total_context_item['thinking'] = orchestration_trace['modelInvocationOutput']['rawResponse']
                 if(total_context_item!={}):
                     total_context.append(total_context_item)
-        #print("total_context------")
-        #print(total_context)    
+
+
+                    # 🔁 Generate + send OpenTelemetry span for each block
+                    span = convert_to_span(total_context_item, trace_id, i)
+                    send_otel_span(span)
+   
     except botocore.exceptions.EventStreamError as error:
         raise error