diff --git a/agentops/__init__.py b/agentops/__init__.py index cda6079c..65d698c6 100755 --- a/agentops/__init__.py +++ b/agentops/__init__.py @@ -48,12 +48,15 @@ def init(api_key: Optional[str] = None, inherited_session_id (optional, str): Init Agentops with an existing Session Attributes: """ - if os.getenv('AGENTOPS_LOGGING_LEVEL') == 'DEBUG': - logger.setLevel(logging.DEBUG) - elif os.getenv('AGENTOPS_LOGGING_LEVEL') == 'CRITICAL': - logger.setLevel(logging.CRITICAL) - else: - logger.setLevel(logging.INFO) + logging_level = os.getenv('AGENTOPS_LOGGING_LEVEL') + log_levels = { + "CRITICAL": logging.CRITICAL, + "ERROR": logging.ERROR, + "INFO": logging.INFO, + "WARNING": logging.WARNING, + "DEBUG": logging.DEBUG + } + logger.setLevel(log_levels.get(logging_level or "INFO", "INFO")) c = Client(api_key=api_key, parent_key=parent_key, diff --git a/agentops/client.py b/agentops/client.py index 67c0dbfa..4d1849fb 100644 --- a/agentops/client.py +++ b/agentops/client.py @@ -278,12 +278,15 @@ def start_session(self, tags: Optional[List[str]] = None, config: Optional[Confi config: (Configuration, optional): Client configuration object inherited_session_id (optional, str): assign session id to match existing Session """ - if os.getenv('AGENTOPS_LOGGING_LEVEL') == 'DEBUG': - logger.setLevel(logging.DEBUG) - elif os.getenv('AGENTOPS_LOGGING_LEVEL') == 'CRITICAL': - logger.setLevel(logging.CRITICAL) - else: - logger.setLevel(logging.INFO) + logging_level = os.getenv('AGENTOPS_LOGGING_LEVEL') + log_levels = { + "CRITICAL": logging.CRITICAL, + "ERROR": logging.ERROR, + "INFO": logging.INFO, + "WARNING": logging.WARNING, + "DEBUG": logging.DEBUG + } + logger.setLevel(log_levels.get(logging_level or "INFO", "INFO")) if self._session is not None: return logger.warning("Cannot start session - session already started") @@ -298,12 +301,19 @@ def start_session(self, tags: Optional[List[str]] = None, config: Optional[Confi tags=tags or self._tags_for_future_session, host_env=get_host_env(self._env_data_opt_out)) self._worker = Worker(config or self.config) - start_session_result = self._worker.start_session(self._session) + + start_session_result = False + if inherited_session_id is not None: + start_session_result = self._worker.reauthorize_jwt(self._session) + else: + start_session_result = self._worker.start_session(self._session) + if not start_session_result: self._session = None - return logger.warning("Cannot start session - No server response") + return logger.warning("Cannot start session - server rejected session") - logger.info(f'\x1b[34mSession Replay: https://app.agentops.ai/drilldown?session_id={self._session.session_id}\x1b[0m') + logger.info( + f'\x1b[34mSession Replay: https://app.agentops.ai/drilldown?session_id={self._session.session_id}\x1b[0m') return self._session.session_id @@ -319,6 +329,7 @@ def end_session(self, end_state_reason (str, optional): The reason for ending the session. video (str, optional): The video screen recording of the session """ + if self._session is None or self._session.has_ended: return logger.warning("Cannot end session - no current session") @@ -339,7 +350,8 @@ def end_session(self, logger.info('This run\'s cost ${}'.format('{:.2f}'.format( token_cost_d) if token_cost_d == 0 else '{:.6f}'.format(token_cost_d))) - logger.info(f'\x1b[34mSession Replay: https://app.agentops.ai/drilldown?session_id={self._session.session_id}\x1b[0m') + logger.info( + f'\x1b[34mSession Replay: https://app.agentops.ai/drilldown?session_id={self._session.session_id}\x1b[0m') self._session = None self._worker = None diff --git a/agentops/config.py b/agentops/config.py index e2681fb8..6cdb39ab 100644 --- a/agentops/config.py +++ b/agentops/config.py @@ -21,7 +21,7 @@ class Configuration: be read from the AGENTOPS_PARENT_KEY environment variable. endpoint (str, optional): The endpoint for the AgentOps service. If none is provided, key will be read from the AGENTOPS_API_ENDPOINT environment variable. Defaults to 'https://api.agentops.ai'. - max_wait_time (int, optional): The maximum time to wait in milliseconds before flushing the queue. Defaults to 30000. + max_wait_time (int, optional): The maximum time to wait in milliseconds before flushing the queue. Defaults to 5000. max_queue_size (int, optional): The maximum size of the event queue. Defaults to 100. """ @@ -47,7 +47,7 @@ def __init__(self, self._api_key: str = api_key self._endpoint = endpoint - self._max_wait_time = max_wait_time or 30000 + self._max_wait_time = max_wait_time or 5000 self._max_queue_size = max_queue_size or 100 self._parent_key: Optional[str] = parent_key @@ -79,7 +79,7 @@ def endpoint(self) -> str: Returns: str: The endpoint for the AgentOps service. """ - return self._endpoint + return self._endpoint # type: ignore @endpoint.setter def endpoint(self, value: str): diff --git a/agentops/helpers.py b/agentops/helpers.py index 6565a24f..818e8f09 100644 --- a/agentops/helpers.py +++ b/agentops/helpers.py @@ -134,14 +134,13 @@ def get_agentops_version(): def debug_print_function_params(func): @wraps(func) def wrapper(self, *args, **kwargs): - if os.getenv('DEBUG_MODE') == 'Y': - print("\n") - print(f"{func.__name__} called with arguments:") + logger.debug("\n") + logger.debug(f"{func.__name__} called with arguments:") - for key, value in kwargs.items(): - print(f"{key}: {pformat(value)}") + for key, value in kwargs.items(): + logger.debug(f"{key}: {pformat(value)}") - print("\n") + logger.debug("\n") return func(self, *args, **kwargs) return wrapper diff --git a/agentops/http_client.py b/agentops/http_client.py index a007fe1a..adef5e5a 100644 --- a/agentops/http_client.py +++ b/agentops/http_client.py @@ -60,7 +60,7 @@ class HttpClient: @staticmethod def post(url: str, payload: bytes, api_key: Optional[str] = None, parent_key: Optional[str] = None, - header=None) -> Response: + jwt: Optional[str] = None, header=None) -> Response: result = Response() try: # Create request session with retries configured @@ -68,11 +68,14 @@ def post(url: str, payload: bytes, api_key: Optional[str] = None, parent_key: Op request_session.mount(url, HTTPAdapter(max_retries=retry_config)) if api_key is not None: - JSON_HEADER["X-Agentops-Auth"] = api_key + JSON_HEADER["X-Agentops-Api-Key"] = api_key if parent_key is not None: JSON_HEADER["X-Agentops-Parent-Key"] = parent_key + if jwt is not None: + JSON_HEADER["Authorization"] = f"Bearer {jwt}" + res = request_session.post(url, data=payload, headers=JSON_HEADER, timeout=20) diff --git a/agentops/langchain_callback_handler.py b/agentops/langchain_callback_handler.py index fb2e9c83..c95fc873 100644 --- a/agentops/langchain_callback_handler.py +++ b/agentops/langchain_callback_handler.py @@ -126,7 +126,8 @@ def on_llm_end( if len(response.generations) == 0: # TODO: more descriptive error - error_event = ErrorEvent(trigger_event=self.events.llm[str(run_id)], error_type="NoGenerations", details="on_llm_end: No generations") + error_event = ErrorEvent(trigger_event=self.events.llm[str(run_id)], + error_type="NoGenerations", details="on_llm_end: No generations") self.ao_client.record(error_event) @debug_print_function_params @@ -218,7 +219,8 @@ def on_tool_end( # Tools are capable of failing `on_tool_end` quietly. # This is a workaround to make sure we can log it as an error. if kwargs.get('name') == '_Exception': - error_event = ErrorEvent(trigger_event=tool_event, error_type="LangchainToolException", details=output) + error_event = ErrorEvent( + trigger_event=tool_event, error_type="LangchainToolException", details=output) self.ao_client.record(error_event) @debug_print_function_params @@ -267,7 +269,8 @@ def on_retriever_end( **kwargs: Any, ) -> None: action_event: ActionEvent = self.events.retriever[str(run_id)] - action_event.logs = documents # TODO: Adding this. Might want to add elsewhere e.g. params + # TODO: Adding this. Might want to add elsewhere e.g. params + action_event.logs = documents action_event.end_timestamp = get_ISO_time() self.ao_client.record(action_event) @@ -450,7 +453,8 @@ async def on_llm_end( if len(response.generations) == 0: # TODO: more descriptive error - error_event = ErrorEvent(trigger_event=self.events.llm[str(run_id)], error_type="NoGenerations", details="on_llm_end: No generations") + error_event = ErrorEvent(trigger_event=self.events.llm[str( + run_id)], error_type="NoGenerations", details="on_llm_end: No generations") self.ao_client.record(error_event) @debug_print_function_params @@ -542,7 +546,8 @@ async def on_tool_end( # Tools are capable of failing `on_tool_end` quietly. # This is a workaround to make sure we can log it as an error. if kwargs.get('name') == '_Exception': - error_event = ErrorEvent(trigger_event=tool_event, error_type="LangchainToolException", details=output) + error_event = ErrorEvent( + trigger_event=tool_event, error_type="LangchainToolException", details=output) self.ao_client.record(error_event) @debug_print_function_params @@ -591,7 +596,8 @@ async def on_retriever_end( **kwargs: Any, ) -> None: action_event: ActionEvent = self.events.retriever[str(run_id)] - action_event.logs = documents # TODO: Adding this. Might want to add elsewhere e.g. params + # TODO: Adding this. Might want to add elsewhere e.g. params + action_event.logs = documents action_event.end_timestamp = get_ISO_time() self.ao_client.record(action_event) diff --git a/agentops/log_config.py b/agentops/log_config.py index 8bf5b8eb..ea04ddc4 100644 --- a/agentops/log_config.py +++ b/agentops/log_config.py @@ -1,18 +1,19 @@ import logging +from sys import prefix class AgentOpsFormatter(logging.Formatter): blue = "\x1b[34m" bold_red = "\x1b[31;1m" reset = "\x1b[0m" - format = "🖇 AgentOps: %(message)s" + prefix = "🖇 AgentOps: %(message)s" FORMATS = { - logging.DEBUG: "(DEBUG) " + format, - logging.INFO: format, - logging.WARNING: format, - logging.ERROR: format, - logging.CRITICAL: bold_red + format + reset, + logging.DEBUG: f"(DEBUG) {prefix}", + logging.INFO: f"{prefix}", + logging.WARNING: f"{prefix}", + logging.ERROR: f"{bold_red}{prefix}{reset}", + logging.CRITICAL: f"{bold_red}{prefix}{reset}", } def format(self, record): diff --git a/agentops/meta_client.py b/agentops/meta_client.py index fe7b5f75..3fc2bfea 100644 --- a/agentops/meta_client.py +++ b/agentops/meta_client.py @@ -36,7 +36,7 @@ def send_exception_to_server(cls, exception, api_key, session): if session: developer_error["session_id"] = session.session_id - HttpClient.post("https://api.agentops.ai/developer_errors", + HttpClient.post("https://api.agentops.ai/v2/developer_errors", safe_serialize(developer_error).encode("utf-8"), api_key=api_key) diff --git a/agentops/worker.py b/agentops/worker.py index 4605e2b4..40b00e41 100644 --- a/agentops/worker.py +++ b/agentops/worker.py @@ -1,4 +1,5 @@ import json +from .log_config import logger import threading import time from .http_client import HttpClient @@ -6,7 +7,6 @@ from .session import Session from .helpers import safe_serialize, filter_unjsonable from typing import Dict, Optional -import os class Worker: @@ -19,7 +19,7 @@ def __init__(self, config: Configuration) -> None: self.thread.daemon = True self.thread.start() self._session: Optional[Session] = None - self._debug_mode = os.getenv('DEBUG_MODE') == 'Y' + self.jwt = None def add_event(self, event: dict) -> None: with self.lock: @@ -39,16 +39,36 @@ def flush_queue(self) -> None: } serialized_payload = safe_serialize(payload).encode("utf-8") - HttpClient.post(f'{self.config.endpoint}/events', + HttpClient.post(f'{self.config.endpoint}/v2/create_events', serialized_payload, - self.config.api_key, - self.config.parent_key) + jwt=self.jwt) - if self._debug_mode: - print("\n") - print(f"Worker request to {self.config.endpoint}/events") - print(serialized_payload) - print("\n") + logger.debug("\n") + logger.debug(f"Worker request to {self.config.endpoint}/events") + logger.debug(serialized_payload) + logger.debug("\n") + + def reauthorize_jwt(self, session: Session) -> bool: + self._session = session + with self.lock: + payload = { + "session_id": session.session_id + } + serialized_payload = json.dumps(filter_unjsonable(payload)).encode("utf-8") + res = HttpClient.post(f'{self.config.endpoint}/v2/reauthorize_jwt', + serialized_payload, + self.config.api_key) + + logger.debug(res.body) + + if res.code != 200: + return False + + self.jwt = res.body.get('jwt', None) + if self.jwt is None: + return False + + return True def start_session(self, session: Session) -> bool: self._session = session @@ -57,14 +77,20 @@ def start_session(self, session: Session) -> bool: "session": session.__dict__ } serialized_payload = json.dumps(filter_unjsonable(payload)).encode("utf-8") - res = HttpClient.post(f'{self.config.endpoint}/sessions', + res = HttpClient.post(f'{self.config.endpoint}/v2/create_session', serialized_payload, self.config.api_key, self.config.parent_key) + logger.debug(res.body) + if res.code != 200: return False + self.jwt = res.body.get('jwt', None) + if self.jwt is None: + return False + return True def end_session(self, session: Session) -> str: @@ -78,12 +104,10 @@ def end_session(self, session: Session) -> str: "session": session.__dict__ } - res = HttpClient.post(f'{self.config.endpoint}/sessions', - json.dumps(filter_unjsonable( - payload)).encode("utf-8"), - self.config.api_key, - self.config.parent_key) - + res = HttpClient.post(f'{self.config.endpoint}/v2/update_session', + json.dumps(filter_unjsonable(payload)).encode("utf-8"), + jwt=self.jwt) + logger.debug(res.body) return res.body.get('token_cost', "unknown") def update_session(self, session: Session) -> None: @@ -92,11 +116,9 @@ def update_session(self, session: Session) -> None: "session": session.__dict__ } - HttpClient.post(f'{self.config.endpoint}/sessions', - json.dumps(filter_unjsonable( - payload)).encode("utf-8"), - self.config.api_key, - self.config.parent_key) + res = HttpClient.post(f'{self.config.endpoint}/v2/update_session', + json.dumps(filter_unjsonable(payload)).encode("utf-8"), + jwt=self.jwt) def create_agent(self, agent_id, name): payload = { @@ -107,10 +129,9 @@ def create_agent(self, agent_id, name): serialized_payload = \ safe_serialize(payload).encode("utf-8") - HttpClient.post(f'{self.config.endpoint}/agents', + HttpClient.post(f'{self.config.endpoint}/v2/create_agent', serialized_payload, - self.config.api_key, - self.config.parent_key) + jwt=self.jwt) def run(self) -> None: while not self.stop_flag.is_set(): diff --git a/examples/recording-events.ipynb b/examples/recording-events.ipynb index 9c766b58..6a85346b 100644 --- a/examples/recording-events.ipynb +++ b/examples/recording-events.ipynb @@ -8,7 +8,9 @@ }, "source": [ "# Recording Events\n", - "AgentOps has a number of different [Event Types](https://docs.agentops.ai/v1/details/events)" + "AgentOps has a number of different [Event Types](https://docs.agentops.ai/v1/details/events)\n", + "\n", + "AgentOps automatically instruments your LLM Calls from OpenAI, LiteLLM, and Cohere. Just make sure their SDKs are imported before initializing AgentOps like we see below" ] }, { @@ -22,11 +24,25 @@ "outputs": [], "source": [ "import agentops\n", + "import openai\n", + "\n", "# Create new session\n", "agentops.init()\n", "\n", "# Optionally, we can add tags to the session\n", - "# agentops.init(tags=['Hello Tracker'])" + "# agentops.init(tags=['Hello Tracker'])\n", + "\n", + "message = {\"role\": \"user\", \"content\": \"Hello\"},\n", + "response = openai.chat.completions.create(\n", + " model='gpt-3.5-turbo', messages=message, temperature=0.5)" + ] + }, + { + "cell_type": "markdown", + "id": "95a6047d", + "metadata": {}, + "source": [ + "Click the AgentOps link above to see your session!" ] }, { @@ -36,7 +52,9 @@ "collapsed": false }, "source": [ - "The easiest way to record actions is through the use of AgentOp's decorators" + "## Action Event\n", + "\n", + "AgentOps allows you to record other actions. The easiest way to record actions is through the use of AgentOp's decorators" ] }, { @@ -64,7 +82,7 @@ "collapsed": false }, "source": [ - "We can also manually craft an event exactly the way we want" + "We can also manually craft an event exactly the way we want by creating and recording an `ActionEvent`" ] }, { @@ -78,24 +96,7 @@ "source": [ "from agentops import ActionEvent\n", "\n", - "message = {\"role\": \"user\", \"content\": \"Hello\"},\n", - "response = openai.chat.completions.create(\n", - " model='gpt-3.5-turbo', messages=message, temperature=0.5)\n", - "\n", - "if \"hello\" in str(response.choices[0].message.content).lower():\n", - " agentops.record(ActionEvent(action_type=\"Agent says hello\", params=str(message), returns=str(response.choices[0].message.content) ))" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "e10a89e06fe6b2be", - "metadata": { - "collapsed": false - }, - "outputs": [], - "source": [ - "agentops.end_session('Success')" + "agentops.record(ActionEvent(action_type=\"Agent says hello\", params={\"message\": \"Hi\"}, returns=\"Hi Back!\" ))" ] }, { diff --git a/pyproject.toml b/pyproject.toml index 772c33e9..8ccb7bd6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "agentops" -version = "0.1.12" +version = "0.2.0" authors = [ { name="Alex Reibman", email="areibman@gmail.com" }, { name="Shawn Qiu", email="siyangqiu@gmail.com" }, @@ -36,3 +36,6 @@ langchain = [ [project.urls] Homepage = "https://github.com/AgentOps-AI/agentops" Issues = "https://github.com/AgentOps-AI/agentops/issues" + +[tool.autopep8] +max_line_length = 120 diff --git a/tests/test_canary.py b/tests/test_canary.py index 09c4c30b..a8221980 100644 --- a/tests/test_canary.py +++ b/tests/test_canary.py @@ -10,21 +10,27 @@ def mock_req(): with requests_mock.Mocker() as m: url = 'https://api.agentops.ai' - m.post(url + '/events', text='ok') - m.post(url + '/sessions', json={'status': 'success', 'token_cost': 5}) + m.post(url + '/v2/create_events', text='ok') + m.post(url + '/v2/create_session', json={'status': 'success', + 'jwt': 'some_jwt'}) + m.post(url + '/v2/update_session', + json={'status': 'success', 'token_cost': 5}) + m.post(url + '/v2/developer_errors', text='ok') yield m + class TestCanary: def setup_method(self): self.url = 'https://api.agentops.ai' self.api_key = "random_api_key" - agentops.init(api_key=self.api_key, max_wait_time=5, auto_start_session=False) + agentops.init(api_key=self.api_key, max_wait_time=5, + auto_start_session=False) def test_agent_ops_record(self, mock_req): # Arrange event_type = 'test_event_type' agentops.start_session() - + # Act agentops.record(ActionEvent(event_type)) time.sleep(0.1) @@ -32,7 +38,7 @@ def test_agent_ops_record(self, mock_req): # Assert assert len(mock_req.request_history) == 2 request_json = mock_req.last_request.json() - assert mock_req.last_request.headers['X-Agentops-Auth'] == self.api_key + assert mock_req.last_request.headers['X-Agentops-Api-Key'] == self.api_key assert request_json['events'][0]['event_type'] == event_type agentops.end_session('Success') diff --git a/tests/test_events.py b/tests/test_events.py index a4da4ffd..521fe056 100644 --- a/tests/test_events.py +++ b/tests/test_events.py @@ -5,13 +5,27 @@ from agentops import ActionEvent, ErrorEvent +@pytest.fixture +def mock_req(): + with requests_mock.Mocker() as m: + url = 'https://api.agentops.ai' + m.post(url + '/v2/create_events', text='ok') + m.post(url + '/v2/create_session', json={'status': 'success', + 'jwt': 'some_jwt'}) + m.post(url + '/v2/update_session', + json={'status': 'success', 'token_cost': 5}) + m.post(url + '/v2/developer_errors', text='ok') + yield m + + class TestEvents: def setup_method(self): self.api_key = "random_api_key" self.event_type = 'test_event_type' - self.config = agentops.Configuration(api_key=self.api_key, max_wait_time=50, max_queue_size=1) + self.config = agentops.Configuration( + api_key=self.api_key, max_wait_time=50, max_queue_size=1) - def test_record_timestamp(self): + def test_record_timestamp(self, mock_req): agentops.init(api_key=self.api_key) event = ActionEvent() @@ -20,10 +34,9 @@ def test_record_timestamp(self): assert event.init_timestamp != event.end_timestamp - def test_record_error_event(self): + def test_record_error_event(self, mock_req): agentops.init(api_key=self.api_key) - event = ErrorEvent() + event = ErrorEvent(logs=None) time.sleep(0.15) agentops.record(event) - diff --git a/tests/test_record_function.py b/tests/test_record_function.py index 5896381b..32911784 100644 --- a/tests/test_record_function.py +++ b/tests/test_record_function.py @@ -5,14 +5,20 @@ from agentops import record_function from datetime import datetime + @pytest.fixture def mock_req(): with requests_mock.Mocker() as m: url = 'https://api.agentops.ai' - m.post(url + '/events', text='ok') - m.post(url + '/sessions', json={'status': 'success', 'token_cost': 5}) + m.post(url + '/v2/create_events', text='ok') + m.post(url + '/v2/create_session', json={'status': 'success', + 'jwt': 'some_jwt'}) + m.post(url + '/v2/update_session', + json={'status': 'success', 'token_cost': 5}) + m.post(url + '/v2/developer_errors', text='ok') yield m + class TestRecordAction: def setup_method(self): self.url = 'https://api.agentops.ai' @@ -33,7 +39,7 @@ def add_two(x, y): # Assert assert len(mock_req.request_history) == 2 - assert mock_req.last_request.headers['X-Agentops-Auth'] == self.api_key + assert mock_req.last_request.headers['X-Agentops-Api-Key'] == self.api_key request_json = mock_req.last_request.json() assert request_json['events'][0]['action_type'] == self.event_type assert request_json['events'][0]['params'] == {'x': 3, 'y': 4} @@ -57,7 +63,7 @@ def add_three(x, y, z=3): # Assert assert len(mock_req.request_history) == 3 - assert mock_req.last_request.headers['X-Agentops-Auth'] == self.api_key + assert mock_req.last_request.headers['X-Agentops-Api-Key'] == self.api_key request_json = mock_req.last_request.json() assert request_json['events'][0]['action_type'] == self.event_type assert request_json['events'][0]['params'] == {'x': 1, 'y': 2, 'z': 3} @@ -82,7 +88,7 @@ async def async_add(x, y): assert result == 7 # Assert assert len(mock_req.request_history) == 2 - assert mock_req.last_request.headers['X-Agentops-Auth'] == self.api_key + assert mock_req.last_request.headers['X-Agentops-Api-Key'] == self.api_key request_json = mock_req.last_request.json() assert request_json['events'][0]['action_type'] == self.event_type assert request_json['events'][0]['params'] == {'x': 3, 'y': 4} @@ -94,4 +100,4 @@ async def async_add(x, y): assert (end - init).total_seconds() >= 0.1 - agentops.end_session(end_state='Success') \ No newline at end of file + agentops.end_session(end_state='Success') diff --git a/tests/test_session.py b/tests/test_session.py index d7b92f88..3a104e6a 100644 --- a/tests/test_session.py +++ b/tests/test_session.py @@ -9,8 +9,14 @@ def mock_req(): with requests_mock.Mocker() as m: url = 'https://api.agentops.ai' - m.post(url + '/events', text='ok') - m.post(url + '/sessions', json={'status': 'success', 'token_cost': 5}) + m.post(url + '/v2/create_events', text='ok') + m.post(url + '/v2/create_session', json={'status': 'success', + 'jwt': 'some_jwt'}) + m.post(url + '/v2/reauthorize_jwt', json={'status': 'success', + 'jwt': 'some_jwt'}) + m.post(url + '/v2/update_session', + json={'status': 'success', 'token_cost': 5}) + m.post(url + '/v2/developer_errors', text='ok') yield m @@ -18,10 +24,12 @@ class TestSessions: def setup_method(self): self.api_key = "random_api_key" self.event_type = 'test_event_type' - self.config = agentops.Configuration(api_key=self.api_key, max_wait_time=50) + self.config = agentops.Configuration( + api_key=self.api_key, max_wait_time=50) def test_session(self, mock_req): agentops.start_session(config=self.config) + print(self.config.api_key) agentops.record(ActionEvent(self.event_type)) agentops.record(ActionEvent(self.event_type)) @@ -31,9 +39,8 @@ def test_session(self, mock_req): time.sleep(0.15) # We should have 2 requests (session and 2 events combined into 1) - print(mock_req.last_request.json()) assert len(mock_req.request_history) == 2 - assert mock_req.last_request.headers['X-Agentops-Auth'] == self.api_key + assert mock_req.last_request.headers['Authorization'] == f"Bearer some_jwt" request_json = mock_req.last_request.json() assert request_json['events'][0]['event_type'] == self.event_type @@ -43,7 +50,7 @@ def test_session(self, mock_req): # We should have 3 requests (additional end session) assert len(mock_req.request_history) == 3 - assert mock_req.last_request.headers['X-Agentops-Auth'] == self.api_key + assert mock_req.last_request.headers['Authorization'] == f"Bearer some_jwt" request_json = mock_req.last_request.json() assert request_json['session']['end_state'] == end_state assert request_json['session']['tags'] == None @@ -61,11 +68,11 @@ def test_add_tags(self, mock_req): time.sleep(0.15) # Assert 3 requests, 1 for session init, 1 for event, 1 for end session - assert mock_req.last_request.headers['X-Agentops-Auth'] == self.api_key + assert mock_req.last_request.headers['X-Agentops-Api-Key'] == self.api_key request_json = mock_req.last_request.json() assert request_json['session']['end_state'] == end_state - assert request_json['session']['tags'] == ['GPT-4', 'test-tag', 'dupe-tag'] - + assert request_json['session']['tags'] == [ + 'GPT-4', 'test-tag', 'dupe-tag'] def test_tags(self, mock_req): # Arrange @@ -79,7 +86,7 @@ def test_tags(self, mock_req): # Assert 2 requests - 1 for session init, 1 for event print(mock_req.last_request.json()) assert len(mock_req.request_history) == 2 - assert mock_req.last_request.headers['X-Agentops-Auth'] == self.api_key + assert mock_req.last_request.headers['X-Agentops-Api-Key'] == self.api_key request_json = mock_req.last_request.json() assert request_json['events'][0]['event_type'] == self.event_type @@ -90,7 +97,7 @@ def test_tags(self, mock_req): # Assert 3 requests, 1 for session init, 1 for event, 1 for end session assert len(mock_req.request_history) == 3 - assert mock_req.last_request.headers['X-Agentops-Auth'] == self.api_key + assert mock_req.last_request.headers['X-Agentops-Api-Key'] == self.api_key request_json = mock_req.last_request.json() assert request_json['session']['end_state'] == end_state assert request_json['session']['tags'] == tags @@ -98,14 +105,11 @@ def test_tags(self, mock_req): def test_inherit_session_id(self, mock_req): # Arrange inherited_id = '4f72e834-ff26-4802-ba2d-62e7613446f1' - agentops.start_session(tags=['test'], config=self.config, inherited_session_id=inherited_id) + agentops.start_session( + tags=['test'], config=self.config, inherited_session_id=inherited_id) # Act - agentops.record(ActionEvent(self.event_type)) - agentops.record(ActionEvent(self.event_type)) - time.sleep(0.15) - - # event session_id correct + # session_id correct request_json = mock_req.last_request.json() assert request_json['session_id'] == inherited_id diff --git a/tests/test_teardown.py b/tests/test_teardown.py index 114a6aff..cfd8cae1 100644 --- a/tests/test_teardown.py +++ b/tests/test_teardown.py @@ -1,7 +1,5 @@ import pytest import requests_mock -import sys - from agentops import Client @@ -9,8 +7,12 @@ def mock_req(): with requests_mock.Mocker() as m: url = 'https://api.agentops.ai' - m.post(url + '/events', text='ok') - m.post(url + '/sessions', json={'status': 'success', 'token_cost': 5}) + m.post(url + '/v2/create_events', text='ok') + m.post(url + '/v2/create_session', json={'status': 'success', + 'jwt': 'some_jwt'}) + m.post(url + '/v2/update_session', + json={'status': 'success', 'token_cost': 5}) + m.post(url + '/v2/developer_errors', text='ok') yield m