fix kwargs key to get prompt

AgentOps-AI · Aug 10, 2024 · 4690bcd · 4690bcd
1 parent dd79719
commit 4690bcd
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/agentops/llm_tracker.py b/agentops/llm_tracker.py
@@ -651,11 +651,11 @@ def handle_stream_chunk(chunk: Message):
             try:
                 accumulated_delta = ""
                 self.llm_event.agent_id = check_call_stack_for_agent_id()
-                self.llm_event.prompt = kwargs["messages"]
+                self.llm_event.prompt = kwargs["messages"][0]["content"]
 
                 if isinstance(chunk, RawMessageStartEvent):
                     self.llm_event.model = chunk.message.model
-                    self.llm_event.prompt_tokens = chunk.content
+                    self.llm_event.prompt_tokens = chunk.usage.input_tokens
                 elif isinstance(chunk, RawContentBlockStartEvent):
                     accumulated_delta += chunk.content_block.text
                 elif isinstance(chunk, RawContentBlockDeltaEvent):
@@ -679,7 +679,7 @@ def handle_stream_chunk(chunk: Message):
                 logger.warning(
                     f"Unable to parse a chunk for LLM call. Skipping upload to AgentOps\n"
                     f"chunk:\n {chunk}\n"
-                    f"kwargs:\n {kwargs_str}\n"
+                    f"kwargs:\n {kwargs_str}\n",
                 )
 
         # if the response is a generator, decorate the generator
@@ -716,7 +716,7 @@ async def async_generator():
         try:
             self.llm_event.returns = response.model_dump()
             self.llm_event.agent_id = check_call_stack_for_agent_id()
-            self.llm_event.prompt = kwargs["messages"]
+            self.llm_event.prompt = kwargs["messages"][0]["content"]
             self.llm_event.prompt_tokens = response.usage.input_tokens
             self.llm_event.completion = response.content[0].text
             self.llm_event.completion_tokens = response.usage.output_tokens