Added usage tracking

Author: gabrielrfg

examples/bedrock_agents.py

@@ -1,7 +1,7 @@
 import os
 from llmstudio_core.agents import AgentManagerCore
-from llmstudio_core.agents.bedrock.data_models import BedrockCreateAgentRequest, BedrockRunAgentRequest, BedrockToolOutput, BedrockToolCall
-from llmstudio_core.agents.data_models import ResultBase, ToolCall, ToolOutput
+from llmstudio_core.agents.bedrock.data_models import BedrockCreateAgentRequest, BedrockToolCall
+from llmstudio_core.agents.data_models import ResultBase, ToolCall, ToolOutput, RunAgentRequest
 import boto3
 import uuid
 
@@ -67,10 +67,9 @@
 
 runs = []
 
-for i in range(1,3):
-  run_agent_request = BedrockRunAgentRequest(
+for i in range(1,2):
+  run_agent_request = RunAgentRequest(
     agent_id = agent.agent_id,
-    thread_id=f"111{i}",#remove this
     alias_id=agent.agent_alias_id,#make this optional
     messages=[
         {"role": "user", "content": "What is the weather like in Lisbon, PT?"},
@@ -87,15 +86,15 @@
   else:        
       tool_calls : list[BedrockToolCall] = result.messages[-1].required_action.submit_tools_outputs
 
-      submit_outputs_request = BedrockRunAgentRequest(
+      submit_outputs_request = RunAgentRequest(
           agent_id=agent.agent_id,
           thread_id=result.thread_id,
           alias_id=agent.agent_alias_id,
           tool_outputs=[]
       )
 
       for tool_call in tool_calls:
-          submit_outputs_request.tool_outputs.append(BedrockToolOutput(tool_call_id=tool_call.id, output="10", action_group=tool_call.action_group, function_name=tool_call.function.name))
+          submit_outputs_request.tool_outputs.append(ToolOutput(tool_call_id=tool_call.id, output="10", action_group=tool_call.action_group, function_name=tool_call.function.name))
 
       outputs_request = submit_outputs_request.model_dump()
       run = bedrock_agent_manager.submit_tool_outputs(submit_outputs_request.model_dump())

libs/core/llmstudio_core/agents/bedrock/data_models.py

@@ -1,13 +1,11 @@
-from typing import Awaitable, List, Literal, Optional
+from typing import Awaitable, Literal
 
 from llmstudio_core.agents.data_models import (
     AgentBase,
     CreateAgentRequest,
-    RunAgentRequest,
     RunBase,
     Tool,
     ToolCall,
-    ToolOutput,
 )
 from pydantic import BaseModel
 
@@ -67,13 +65,3 @@ def from_tool(cls, tool: Tool) -> "BedrockTool":
 
 class BedrockToolCall(ToolCall):
     action_group: str
-
-
-class BedrockToolOutput(ToolOutput, extra="allow"):
-    action_group: str
-    function_name: str
-
-
-class BedrockRunAgentRequest(RunAgentRequest, extra="allow"):
-    alias_id: str
-    tool_outputs: Optional[List[BedrockToolOutput]] = None

libs/core/llmstudio_core/agents/bedrock/manager.py

@@ -1,5 +1,6 @@
 import asyncio
 import json
+import uuid
 from concurrent.futures import ThreadPoolExecutor
 from functools import partial
 
@@ -8,10 +9,8 @@
     BedrockAgent,
     BedrockCreateAgentRequest,
     BedrockRun,
-    BedrockRunAgentRequest,
     BedrockTool,
     BedrockToolCall,
-    BedrockToolOutput,
 )
 from llmstudio_core.agents.data_models import (
     Attachment,
@@ -20,9 +19,12 @@
     Message,
     RequiredAction,
     ResultBase,
+    RunAgentRequest,
     TextContent,
     TextObject,
+    ToolCall,
     ToolCallFunction,
+    ToolOutput,
 )
 from llmstudio_core.agents.manager import AgentManager, agent_manager
 from llmstudio_core.exceptions import AgentError
@@ -48,7 +50,7 @@ def _validate_create_request(self, request):
         return BedrockCreateAgentRequest(**request)
 
     def _validate_run_request(self, request):
-        return BedrockRunAgentRequest(**request)
+        return RunAgentRequest(**request)
 
     def _validate_result_request(self, request):
         if isinstance(request, BedrockRun):
@@ -206,6 +208,9 @@ def run_agent(self, params: dict = None) -> BedrockRun:
         except ValidationError as e:
             raise AgentError(str(e))
 
+        if not run_request.thread_id:
+            run_request.thread_id = str(uuid.uuid4())
+
         sessionState = {"files": [], "conversationHistory": {"messages": []}}
 
         if isinstance(run_request.messages, Message):
@@ -266,6 +271,7 @@ def run_agent(self, params: dict = None) -> BedrockRun:
                 sessionId=run_request.thread_id,
                 inputText=input_text,
                 sessionState=sessionState,
+                enableTrace=True,
             ),
         )
 
@@ -302,8 +308,59 @@ async def aretrieve_result(self, run: BedrockRun) -> ResultBase:
 
         content = []
         attachments = []
+        messages = []
+        usage = None
         event_stream = run.response.get("completion")
         for event in event_stream:
+            if "trace" in event:
+                trace = event["trace"]["trace"]["orchestrationTrace"]
+
+                if "modelInvocationInput" in trace:
+                    invocation_in = trace["modelInvocationInput"]
+                    text = json.loads(invocation_in["text"])
+                    new_messages = [
+                        Message(content=message["content"], role=message["role"])
+                        for message in text["messages"]
+                    ]
+                    messages += new_messages
+
+                if "modelInvocationOutput" in trace:
+                    invocation_out = trace["modelInvocationOutput"]["rawResponse"][
+                        "content"
+                    ]
+                    invocation_out = json.loads(invocation_out)
+                    if "metadata" in invocation_out:
+                        usage = invocation_out["metadata"]["usage"]
+                    elif "usage" in invocation_out:
+                        usage = invocation_out["usage"]
+
+                    messages = invocation_out["content"]
+                    new_messages = []
+                    for message in messages:
+                        if message["type"] == "text":
+                            new_messages.append(
+                                Message(content=message["text"], role="assistant")
+                            )
+
+                        elif message["type"] == "tool_use":
+                            tool_name = message["name"]
+                            tool_arguments = str(message["input"])
+                            tool_call_id = message["id"]
+
+                            tool_call_func = ToolCallFunction(
+                                arguments=tool_arguments, name=tool_name
+                            )
+                            tool_call = ToolCall(
+                                id=tool_call_id,
+                                function=tool_call_func,
+                                type=message["type"],
+                            )
+                            required_action = RequiredAction(
+                                submit_tools_outputs=[tool_call]
+                            )
+                            new_message = Message(required_action=required_action)
+                            new_messages.append(new_message)
+
             if "chunk" in event:
                 chunk = event["chunk"]
                 if "bytes" in chunk:
@@ -363,6 +420,7 @@ async def aretrieve_result(self, run: BedrockRun) -> ResultBase:
                         function=tool_call_function,
                         type=invocation_type,
                         action_group=action_group,
+                        usage=usage,
                     )
 
                     required_action.submit_tools_outputs.append(tool_call)
@@ -375,9 +433,10 @@ async def aretrieve_result(self, run: BedrockRun) -> ResultBase:
                         )
                     ],
                     thread_id=run.thread_id,
+                    usage=usage,
                 )
 
-        messages = [
+        messages = new_messages + [
             Message(
                 thread_id=run.thread_id,
                 role="assistant",
@@ -386,7 +445,7 @@ async def aretrieve_result(self, run: BedrockRun) -> ResultBase:
             )
         ]
 
-        return ResultBase(messages=messages, thread_id=run.thread_id)
+        return ResultBase(messages=messages, thread_id=run.thread_id, usage=usage)
 
     def submit_tool_outputs(self, params: dict = None) -> ResultBase:
         try:
@@ -397,7 +456,7 @@ def submit_tool_outputs(self, params: dict = None) -> ResultBase:
         if not run_request.tool_outputs:
             raise AgentError("No tool outputs found")
 
-        tool_outputs: list[BedrockToolOutput] = run_request.tool_outputs
+        tool_outputs: list[ToolOutput] = run_request.tool_outputs
 
         invocation_results = [
             {
@@ -424,6 +483,7 @@ def submit_tool_outputs(self, params: dict = None) -> ResultBase:
                 agentAliasId=run_request.alias_id,
                 sessionId=run_request.thread_id,
                 sessionState=sessionState,
+                enableTrace=True,
             ),
         )
 

libs/core/llmstudio_core/agents/data_models.py

@@ -105,9 +105,11 @@ class RequiredAction(BaseModel):
     type: Literal["submit_tool_outputs"] = "submit_tool_outputs"
 
 
-class ToolOutput(BaseModel, extra="allow"):
-    tool_call_id: str
-    output: str
+class ToolOutput(BaseModel):
+    tool_call_id: Optional[str]
+    output: Optional[str]
+    action_group: Optional[str]
+    function_name: Optional[str]
 
 
 class Message(BaseModel):
@@ -147,6 +149,7 @@ class RunBase(BaseModel):
 class ResultBase(BaseModel):
     thread_id: str
     messages: List[Message]
+    usage: Optional[dict] = None
 
 
 class CreateAgentRequest(BaseModel):
@@ -157,8 +160,9 @@ class CreateAgentRequest(BaseModel):
     name: Optional[str] = None
 
 
-class RunAgentRequest(BaseModel, extra="allow"):
-    thread_id: Optional[str] = None
+class RunAgentRequest(BaseModel):
     agent_id: str
+    alias_id: Optional[str]
+    thread_id: Optional[str] = None
     messages: Optional[List[Message]] = None
     tool_outputs: Optional[List[ToolOutput]] = None