Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
38 changes: 21 additions & 17 deletions agentops/client/api/versions/v3.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@
from agentops.client.api.base import BaseApiClient
from agentops.client.api.types import AuthTokenResponse
from agentops.exceptions import ApiServerException

from agentops.logging import logger

class V3Client(BaseApiClient):
"""Client for the AgentOps V3 API"""
Expand All @@ -33,24 +33,28 @@

r = self.post(path, data, headers)

if r.status_code != 200:
error_msg = f"Authentication failed: {r.status_code}"
try:
error_data = r.json()
if "error" in error_data:
error_msg = f"Authentication failed: {error_data['error']}"
except Exception:
pass
raise ApiServerException(error_msg)

try:
jr = r.json()
token = jr.get("token")
if not token:
raise ApiServerException("No token in authentication response")
if r.status_code != 200:
error_msg = f"Authentication failed: {r.status_code}"
try:
error_data = r.json()
if "error" in error_data:
error_msg = f"{error_data['error']}"
except Exception:
pass
raise ApiServerException(error_msg)

Check warning on line 45 in agentops/client/api/versions/v3.py

View check run for this annotation

Codecov / codecov/patch

agentops/client/api/versions/v3.py#L38-L45

Added lines #L38 - L45 were not covered by tests

return jr
try:
jr = r.json()
token = jr.get("token")
if not token:
raise ApiServerException("No token in authentication response")

Check warning on line 51 in agentops/client/api/versions/v3.py

View check run for this annotation

Codecov / codecov/patch

agentops/client/api/versions/v3.py#L51

Added line #L51 was not covered by tests

return jr
except Exception as e:
raise ApiServerException(f"Failed to process authentication response: {str(e)}")

Check warning on line 55 in agentops/client/api/versions/v3.py

View check run for this annotation

Codecov / codecov/patch

agentops/client/api/versions/v3.py#L54-L55

Added lines #L54 - L55 were not covered by tests
except Exception as e:
raise ApiServerException(f"Failed to process authentication response: {str(e)}")
logger.error(f"{str(e)} - Perhaps an invalid API key?")
return None

Check warning on line 58 in agentops/client/api/versions/v3.py

View check run for this annotation

Codecov / codecov/patch

agentops/client/api/versions/v3.py#L57-L58

Added lines #L57 - L58 were not covered by tests

# Add V3-specific API methods here
2 changes: 2 additions & 0 deletions agentops/client/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -68,6 +68,8 @@
# Prefetch JWT token if enabled
# TODO: Move this validation somewhere else (and integrate with self.config.prefetch_jwt_token once we have a solution to that)
response = self.api.v3.fetch_auth_token(self.config.api_key)
if response is None:
return

Check warning on line 72 in agentops/client/client.py

View check run for this annotation

Codecov / codecov/patch

agentops/client/client.py#L72

Added line #L72 was not covered by tests

# Save the bearer for use with the v4 API
self.api.v4.set_auth_token(response["token"])
Expand Down
43 changes: 31 additions & 12 deletions agentops/legacy/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,8 @@

def __del__(self):
try:
self.span.end()
if self.span is not None:
self.span.end()
except:
pass

Expand Down Expand Up @@ -66,9 +67,10 @@

forces a flush to ensure the span is exported immediately.
"""
_set_span_attributes(self.span, kwargs)
self.span.end()
_flush_span_processors()
if self.span is not None:
_set_span_attributes(self.span, kwargs)
self.span.end()
_flush_span_processors()

Check warning on line 73 in agentops/legacy/__init__.py

View check run for this annotation

Codecov / codecov/patch

agentops/legacy/__init__.py#L70-L73

Added lines #L70 - L73 were not covered by tests


def _create_session_span(tags: Union[Dict[str, Any], List[str], None] = None) -> tuple:
Expand Down Expand Up @@ -130,7 +132,21 @@
if not TracingCore.get_instance().initialized:
from agentops import Client
# Pass auto_start_session=False to prevent circular dependency
Client().init(auto_start_session=False)
try:
Client().init(auto_start_session=False)

Check warning on line 136 in agentops/legacy/__init__.py

View check run for this annotation

Codecov / codecov/patch

agentops/legacy/__init__.py#L135-L136

Added lines #L135 - L136 were not covered by tests
# If initialization failed (returned None), create a dummy session
if not TracingCore.get_instance().initialized:
logger.warning("AgentOps client initialization failed. Creating a dummy session that will not send data.")

Check warning on line 139 in agentops/legacy/__init__.py

View check run for this annotation

Codecov / codecov/patch

agentops/legacy/__init__.py#L138-L139

Added lines #L138 - L139 were not covered by tests
# Create a dummy session that won't send data but won't throw exceptions
dummy_session = Session(None, None)
_current_session = dummy_session
return dummy_session
except Exception as e:
logger.warning(f"AgentOps client initialization failed: {str(e)}. Creating a dummy session that will not send data.")

Check warning on line 145 in agentops/legacy/__init__.py

View check run for this annotation

Codecov / codecov/patch

agentops/legacy/__init__.py#L141-L145

Added lines #L141 - L145 were not covered by tests
# Create a dummy session that won't send data but won't throw exceptions
dummy_session = Session(None, None)
_current_session = dummy_session
return dummy_session

Check warning on line 149 in agentops/legacy/__init__.py

View check run for this annotation

Codecov / codecov/patch

agentops/legacy/__init__.py#L147-L149

Added lines #L147 - L149 were not covered by tests

span, ctx, token = _create_session_span(tags)
session = Session(span, token)
Expand All @@ -155,7 +171,7 @@
span: The span to set attributes on
attributes: The attributes to set as a dictionary
"""
if not attributes or not hasattr(span, "set_attribute"):
if span is None:
return

for key, value in attributes.items():
Expand Down Expand Up @@ -235,9 +251,10 @@
if session_or_status is None and kwargs:
if _current_session is not None:
try:
_set_span_attributes(_current_session.span, kwargs)
_finalize_span(_current_session.span, _current_session.token)
_flush_span_processors()
if _current_session.span is not None:
_set_span_attributes(_current_session.span, kwargs)
_finalize_span(_current_session.span, _current_session.token)
_flush_span_processors()
_current_session = None
except Exception as e:
logger.warning(f"Error ending current session: {e}")
Expand All @@ -256,9 +273,11 @@
if hasattr(session_or_status, 'span') and hasattr(session_or_status, 'token'):
try:
# Set attributes and finalize the span
_set_span_attributes(session_or_status.span, kwargs)
_finalize_span(session_or_status.span, session_or_status.token)
_flush_span_processors()
if session_or_status.span is not None:
_set_span_attributes(session_or_status.span, kwargs)
if session_or_status.span is not None:
_finalize_span(session_or_status.span, session_or_status.token)
_flush_span_processors()

# Clear the global session reference if this is the current session
if _current_session is session_or_status:
Expand Down
Loading