Merge branch 'potel-base' into szokeasaurusrex/agents-potel

Author: antonpirker

requirements-aws-lambda-layer.txt

@@ -5,5 +5,6 @@ urllib3
 # https://github.com/boto/botocore/blob/develop/setup.cfg
 # So we pin this here to make our Lambda layer work with
 # Lambda Function using Python 3.7+
-urllib3<1.27; python_version < "3.10"
+urllib3<1.27
+
 opentelemetry-distro>=0.35b0

sentry_sdk/integrations/openai_agents/spans/ai_client.py

@@ -21,10 +21,11 @@
 def ai_client_span(
     agent: Agent, get_response_kwargs: dict[str, Any]
 ) -> sentry_sdk.tracing.Span:
+    model_name = agent.model.model if hasattr(agent.model, "model") else agent.model
     # TODO-anton: implement other types of operations. Now "chat" is hardcoded.
     span = sentry_sdk.start_span(
         op=OP.GEN_AI_CHAT,
-        description=f"chat {agent.model}",
+        description=f"chat {model_name}",
         origin=SPAN_ORIGIN,
     )
     # TODO-anton: remove hardcoded stuff and replace something that also works for embedding and so on

sentry_sdk/integrations/openai_agents/utils.py

@@ -43,7 +43,8 @@ def _set_agent_data(span: sentry_sdk.tracing.Span, agent: agents.Agent) -> None:
         )
 
     if agent.model:
-        span.set_attribute(SPANDATA.GEN_AI_REQUEST_MODEL, agent.model)
+        model_name = agent.model.model if hasattr(agent.model, "model") else agent.model
+        span.set_attribute(SPANDATA.GEN_AI_REQUEST_MODEL, model_name)
 
     if agent.model_settings.presence_penalty:
         span.set_attribute(

sentry_sdk/sessions.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 import os
-import time
-from threading import Thread, Lock
+from threading import Thread, Lock, Event
 from contextlib import contextmanager
 
 import sentry_sdk
@@ -76,7 +75,7 @@ def __init__(
         self._thread_lock = Lock()
         self._aggregate_lock = Lock()
         self._thread_for_pid: Optional[int] = None
-        self._running = True
+        self.__shutdown_requested: Event = Event()
 
     def flush(self) -> None:
         pending_sessions = self.pending_sessions
@@ -119,10 +118,10 @@ def _ensure_running(self) -> None:
                 return None
 
             def _thread() -> None:
-                while self._running:
-                    time.sleep(self.flush_interval)
-                    if self._running:
-                        self.flush()
+                running = True
+                while running:
+                    running = not self.__shutdown_requested.wait(self.flush_interval)
+                    self.flush()
 
             thread = Thread(target=_thread)
             thread.daemon = True
@@ -131,7 +130,7 @@ def _thread() -> None:
             except RuntimeError:
                 # Unfortunately at this point the interpreter is in a state that no
                 # longer allows us to spawn a thread and we have to bail.
-                self._running = False
+                self.__shutdown_requested.set()
                 return None
 
             self._thread = thread
@@ -175,7 +174,7 @@ def add_session(self, session: Session) -> None:
         self._ensure_running()
 
     def kill(self) -> None:
-        self._running = False
+        self.__shutdown_requested.set()
 
     def __del__(self) -> None:
         self.kill()

tests/integrations/openai_agents/test_openai_agents.py

@@ -74,6 +74,24 @@ def test_agent():
     )
 
 
+@pytest.fixture
+def test_agent_custom_model():
+    """Create a real Agent instance for testing."""
+    return Agent(
+        name="test_agent_custom_model",
+        instructions="You are a helpful test assistant.",
+        # the model could be agents.OpenAIChatCompletionsModel()
+        model=MagicMock(model="my-custom-model"),
+        model_settings=ModelSettings(
+            max_tokens=100,
+            temperature=0.7,
+            top_p=1.0,
+            presence_penalty=0.0,
+            frequency_penalty=0.0,
+        ),
+    )
+
+
 @pytest.mark.asyncio
 async def test_agent_invocation_span(
     sentry_init, capture_events, test_agent, mock_model_response
@@ -128,6 +146,42 @@ async def test_agent_invocation_span(
     assert ai_client_span["data"]["gen_ai.request.top_p"] == 1.0
 
 
+@pytest.mark.asyncio
+async def test_client_span_custom_model(
+    sentry_init, capture_events, test_agent_custom_model, mock_model_response
+):
+    """
+    Test that the integration uses the correct model name if a custom model is used.
+    """
+
+    with patch.dict(os.environ, {"OPENAI_API_KEY": "test-key"}):
+        with patch(
+            "agents.models.openai_responses.OpenAIResponsesModel.get_response"
+        ) as mock_get_response:
+            mock_get_response.return_value = mock_model_response
+
+            sentry_init(
+                integrations=[OpenAIAgentsIntegration()],
+                traces_sample_rate=1.0,
+            )
+
+            events = capture_events()
+
+            result = await agents.Runner.run(
+                test_agent_custom_model, "Test input", run_config=test_run_config
+            )
+
+            assert result is not None
+            assert result.final_output == "Hello, how can I help you?"
+
+    (transaction,) = events
+    spans = transaction["spans"]
+    _, ai_client_span = spans
+
+    assert ai_client_span["description"] == "chat my-custom-model"
+    assert ai_client_span["data"]["gen_ai.request.model"] == "my-custom-model"
+
+
 def test_agent_invocation_span_sync(
     sentry_init, capture_events, test_agent, mock_model_response
 ):