Skip to content

Commit

Permalink
helper functions
Browse files Browse the repository at this point in the history
  • Loading branch information
@bboynton97
bboynton97 committed Jul 10, 2024
1 parent ef36b5c commit 27a2f4a
Show file tree
Hide file tree
Showing 4 changed files with 40 additions and 82 deletions.
4 changes: 2 additions & 2 deletions agentops/__init__.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -69,7 +69,7 @@ def init(
instrument_llm_calls (bool): Whether to instrument LLM calls and emit LLMEvents.
auto_start_session (bool): Whether to start a session automatically when the client is created.
inherited_session_id (optional, str): Init Agentops with an existing Session
skip_auto_end_session (optional, bool): Don't automatically end session based on your framework's decision-making
skip_auto_end_session (optional, bool): Don't automatically end session based on your framework's decision-making (i.e. Crew determining when tasks are complete and ending the session)
Attributes:
"""
logging_level = os.getenv("AGENTOPS_LOGGING_LEVEL")
Expand Down Expand Up @@ -200,7 +200,7 @@ def get_api_key() -> str:

def set_parent_key(parent_key):
"""
Set the parent API key which has visibility to projects it is parented to.
Set the parent API key so another organization can view data.
Args:
parent_key (str): The API key of the parent organization to set.
Expand Down
1 change: 0 additions & 1 deletion agentops/config.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,6 @@
Classes:
ClientConfiguration: Stores the configuration settings for AgentOps clients.
AgentOpsConfiguration: This is a class to store the configuration settings for initializing AgentOps
"""

from typing import Optional
Expand Down
115 changes: 37 additions & 78 deletions agentops/llm_tracker.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -80,19 +80,12 @@ def handle_stream_chunk(chunk):
}
self.llm_event.end_timestamp = get_ISO_time()

if session:
session.record(self.llm_event)
else:
self.client.record(self.llm_event)
self._safe_record(session, self.llm_event)

Check warning on line 83 in agentops/llm_tracker.py

View check run for this annotation

Codecov / codecov/patch

agentops/llm_tracker.py#L83 

Added line #L83 was not covered by tests
except Exception as e:
if session:
session.record(
ErrorEvent(trigger_event=self.llm_event, exception=e),
)
else:
self.client.record(
ErrorEvent(trigger_event=self.llm_event, exception=e),
)
self._safe_record(

Check warning on line 85 in agentops/llm_tracker.py

View check run for this annotation

Codecov / codecov/patch

agentops/llm_tracker.py#L85 

Added line #L85 was not covered by tests
session, ErrorEvent(trigger_event=self.llm_event, exception=e)
)

kwargs_str = pprint.pformat(kwargs)
chunk = pprint.pformat(chunk)
logger.warning(
Expand Down Expand Up @@ -136,19 +129,12 @@ def generator():
self.llm_event.model = response["model"]
self.llm_event.end_timestamp = get_ISO_time()

if session:
session.record(self.llm_event)
else:
self.client.record(self.llm_event)
self._safe_record(session, self.llm_event)

Check warning on line 132 in agentops/llm_tracker.py

View check run for this annotation

Codecov / codecov/patch

agentops/llm_tracker.py#L132 

Added line #L132 was not covered by tests
except Exception as e:
if session:
session.record(
ErrorEvent(trigger_event=self.llm_event, exception=e),
)
else:
self.client.record(
ErrorEvent(trigger_event=self.llm_event, exception=e),
)
self._safe_record(

Check warning on line 134 in agentops/llm_tracker.py

View check run for this annotation

Codecov / codecov/patch

agentops/llm_tracker.py#L134 

Added line #L134 was not covered by tests
session, ErrorEvent(trigger_event=self.llm_event, exception=e)
)

kwargs_str = pprint.pformat(kwargs)
response = pprint.pformat(response)
logger.warning(
Expand Down Expand Up @@ -211,19 +197,12 @@ def handle_stream_chunk(chunk: ChatCompletionChunk):
}
self.llm_event.end_timestamp = get_ISO_time()

if session:
session.record(self.llm_event)
else:
self.client.record(self.llm_event)
self._safe_record(session, self.llm_event)

Check warning on line 200 in agentops/llm_tracker.py

View check run for this annotation

Codecov / codecov/patch

agentops/llm_tracker.py#L200 

Added line #L200 was not covered by tests
except Exception as e:
if session:
session.record(
ErrorEvent(trigger_event=self.llm_event, exception=e),
)
else:
self.client.record(
ErrorEvent(trigger_event=self.llm_event, exception=e),
)
self._safe_record(

Check warning on line 202 in agentops/llm_tracker.py

View check run for this annotation

Codecov / codecov/patch

agentops/llm_tracker.py#L202 

Added line #L202 was not covered by tests
session, ErrorEvent(trigger_event=self.llm_event, exception=e)
)

kwargs_str = pprint.pformat(kwargs)
chunk = pprint.pformat(chunk)
logger.warning(
Expand Down Expand Up @@ -272,19 +251,12 @@ async def async_generator():
self.llm_event.completion_tokens = response.usage.completion_tokens
self.llm_event.model = response.model

if session:
session.record(self.llm_event)
else:
self.client.record(self.llm_event)
self._safe_record(session, self.llm_event)

Check warning on line 254 in agentops/llm_tracker.py

View check run for this annotation

Codecov / codecov/patch

agentops/llm_tracker.py#L254 

Added line #L254 was not covered by tests
except Exception as e:
if session:
session.record(
ErrorEvent(trigger_event=self.llm_event, exception=e),
)
else:
self.client.record(
ErrorEvent(trigger_event=self.llm_event, exception=e),
)
self._safe_record(

Check warning on line 256 in agentops/llm_tracker.py

View check run for this annotation

Codecov / codecov/patch

agentops/llm_tracker.py#L256 

Added line #L256 was not covered by tests
session, ErrorEvent(trigger_event=self.llm_event, exception=e)
)

kwargs_str = pprint.pformat(kwargs)
response = pprint.pformat(response)
logger.warning(
Expand Down Expand Up @@ -338,10 +310,7 @@ def handle_stream_chunk(chunk, session: Optional[Session] = None):
"content": chunk.response.text,
}
self.llm_event.end_timestamp = get_ISO_time()
if session is not None:
session.record(self.llm_event)
else:
self.client.record(self.llm_event)
self._safe_record(session, self.llm_event)

Check warning on line 313 in agentops/llm_tracker.py

View check run for this annotation

Codecov / codecov/patch

agentops/llm_tracker.py#L313 

Added line #L313 was not covered by tests

# StreamedChatResponse_SearchResults = ActionEvent
search_results = chunk.response.search_results
Expand Down Expand Up @@ -374,10 +343,7 @@ def handle_stream_chunk(chunk, session: Optional[Session] = None):
action_event.end_timestamp = get_ISO_time()

for key, action_event in self.action_events.items():
if session is not None:
session.record(action_event)
else:
self.client.record(action_event)
self._safe_record(session, action_event)

Check warning on line 346 in agentops/llm_tracker.py

View check run for this annotation

Codecov / codecov/patch

agentops/llm_tracker.py#L346 

Added line #L346 was not covered by tests

elif isinstance(chunk, StreamedChatResponse_TextGeneration):
self.llm_event.completion += chunk.text
Expand All @@ -403,14 +369,10 @@ def handle_stream_chunk(chunk, session: Optional[Session] = None):
pass

except Exception as e:
if session is not None:
session.record(
ErrorEvent(trigger_event=self.llm_event, exception=e)
)
else:
self.client.record(
ErrorEvent(trigger_event=self.llm_event, exception=e)
)
self._safe_record(

Check warning on line 372 in agentops/llm_tracker.py

View check run for this annotation

Codecov / codecov/patch

agentops/llm_tracker.py#L372 

Added line #L372 was not covered by tests
session, ErrorEvent(trigger_event=self.llm_event, exception=e)
)

kwargs_str = pprint.pformat(kwargs)
chunk = pprint.pformat(chunk)
logger.warning(
Expand Down Expand Up @@ -467,17 +429,11 @@ def generator():
self.llm_event.completion_tokens = response.meta.tokens.output_tokens
self.llm_event.model = kwargs.get("model", "command-r-plus")

if session is not None:
session.record(self.llm_event)
else:
self.client.record(self.llm_event)
self._safe_record(session, self.llm_event)

Check warning on line 432 in agentops/llm_tracker.py

View check run for this annotation

Codecov / codecov/patch

agentops/llm_tracker.py#L432 

Added line #L432 was not covered by tests
except Exception as e:
if session:
session.record(ErrorEvent(trigger_event=self.llm_event, exception=e))
else:
self.client.record(
ErrorEvent(trigger_event=self.llm_event, exception=e)
)
self._safe_record(

Check warning on line 434 in agentops/llm_tracker.py

View check run for this annotation

Codecov / codecov/patch

agentops/llm_tracker.py#L434 

Added line #L434 was not covered by tests
session, ErrorEvent(trigger_event=self.llm_event, exception=e)
)
kwargs_str = pprint.pformat(kwargs)
response = pprint.pformat(response)
logger.warning(
Expand Down Expand Up @@ -528,10 +484,7 @@ def generator():
self.llm_event.prompt = kwargs["messages"]
self.llm_event.completion = response["message"]

if session is not None:
session.record(self.llm_event)
else:
self.client.record(self.llm_event)
self._safe_record(session, self.llm_event)

Check warning on line 487 in agentops/llm_tracker.py

View check run for this annotation

Codecov / codecov/patch

agentops/llm_tracker.py#L487 

Added line #L487 was not covered by tests
return response

def override_openai_v1_completion(self):
Expand Down Expand Up @@ -817,3 +770,9 @@ def undo_override_ollama(self):
ollama.chat = original_func["ollama.chat"]
ollama.Client.chat = original_func["ollama.Client.chat"]
ollama.AsyncClient.chat = original_func["ollama.AsyncClient.chat"]

def _safe_record(self, session, event):
if session is not None:
session.record(event)

Check warning on line 776 in agentops/llm_tracker.py

View check run for this annotation

Codecov / codecov/patch

agentops/llm_tracker.py#L776 

Added line #L776 was not covered by tests
else:
self.client.record(event)

Check warning on line 778 in agentops/llm_tracker.py

View check run for this annotation

Codecov / codecov/patch

agentops/llm_tracker.py#L778 

Added line #L778 was not covered by tests
2 changes: 1 addition & 1 deletion tests/core_manual_tests/multi_session_llm.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@
from agentops import ActionEvent

load_dotenv()
agentops.init(auto_start_session=False, endpoint="http://localhost:8000")
agentops.init(auto_start_session=False)
openai = OpenAI()

session_1 = agentops.start_session(tags=["multi-session-test-1"])
Expand Down

0 comments on commit 27a2f4a

Please sign in to comment.