ubi integration for agent

RAG/bedrock_agent.py

@@ -20,6 +20,22 @@ from datetime import datetime, timezone
 import botocore
 import utilities.ubi_lambda as ubi
 
+# OTEL imports
+from opentelemetry import trace
+from opentelemetry.sdk.resources import Resource
+from opentelemetry.sdk.trace import TracerProvider
+from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
+from opentelemetry.sdk.trace.export import BatchSpanProcessor
+
+# --- OTEL Setup ---
+resource = Resource(attributes={"service.name": "bedrock-agent"})
+trace.set_tracer_provider(TracerProvider(resource=resource))
+tracer = trace.get_tracer_provider().get_tracer("app.tracer")
+
+otlp_exporter = OTLPSpanExporter(endpoint="http://localhost:4318/v1/traces")
+span_processor = BatchSpanProcessor(otlp_exporter)
+trace.get_tracer_provider().add_span_processor(span_processor)
+
 if "inputs_" not in st.session_state:
     st.session_state.inputs_ = {}
 
@@ -117,6 +133,8 @@ def query_(inputs):
     trace_id = str(uuid.uuid4()).replace("-", "")
     try:
         for i,event in enumerate(event_stream):
+            name = "step"
+            attributes = {}
             if 'trace' in event: 
                
                 if('orchestrationTrace' not in event['trace']['trace']):
@@ -127,9 +145,14 @@ def query_(inputs):
                     total_context_item['tool'] = orchestration_trace['modelInvocationOutput']['rawResponse']
                 if('rationale' in orchestration_trace):
                     total_context_item['rationale'] = orchestration_trace['rationale']['text']
+                    name = "rationale"
+                    attributes["message"] = total_context_item["rationale"]
                 if('invocationInput' in orchestration_trace):
                     total_context_item['invocationInput'] = orchestration_trace['invocationInput']['actionGroupInvocationInput']
                     last_tool_name = total_context_item['invocationInput']['function']
+                    name = total_context_item["invocationInput"].get("function", "invocation")
+                    attributes = {p["name"]: p["value"] for p in total_context_item["invocationInput"].get("parameters", [])}
+               
                 if('observation'  in orchestration_trace):
                     total_context_item['observation'] = event['trace']['trace']['orchestrationTrace']['observation']
                     tool_output_last_obs = event['trace']['trace']['orchestrationTrace']['observation']
@@ -137,15 +160,21 @@ def query_(inputs):
                         last_tool = tool_output_last_obs['actionGroupInvocationOutput']['text']
                     if(tool_output_last_obs['type'] == 'FINISH'):   
                         agent_answer = tool_output_last_obs['finalResponse']['text']
+                    name = total_context_item["observation"].get("type", "observation").lower()
+                    attributes = total_context_item["observation"].get("actionGroupInvocationOutput", {})
                 if('modelInvocationOutput' in orchestration_trace and '<thinking>' in orchestration_trace['modelInvocationOutput']['rawResponse']['content']):
                     total_context_item['thinking'] = orchestration_trace['modelInvocationOutput']['rawResponse']
+                    name = "thinking"
+                    attributes["message"] = total_context_item["thinking"].get("content", "")
                 if(total_context_item!={}):
                     total_context.append(total_context_item)
 
 
-                    # 🔁 Generate + send OpenTelemetry span for each block
-                    span = convert_to_span(total_context_item, trace_id, i)
-                    send_otel_span(span)
+                    # # 🔁 Generate + send OpenTelemetry span for each block
+                    # span = convert_to_span(total_context_item, trace_id, i)
+                    # send_otel_span(span)
+                    with tracer.start_as_current_span(name, attributes=attributes) as span:
+                        span.set_attribute("trace.source", "agentic-app")
    
     except botocore.exceptions.EventStreamError as error:
         raise error