From fa865253ae55127260edfb0e022708c2da7f5b59 Mon Sep 17 00:00:00 2001 From: Clara Gadelho Date: Mon, 3 Feb 2025 09:48:57 +0000 Subject: [PATCH 1/4] feat: create agent --- .../agents/bedrock/data_models.py | 11 +-- .../llmstudio_core/agents/bedrock/manager.py | 90 ++++++++++++++++++- .../core/llmstudio_core/agents/data_models.py | 10 +-- 3 files changed, 97 insertions(+), 14 deletions(-) diff --git a/libs/core/llmstudio_core/agents/bedrock/data_models.py b/libs/core/llmstudio_core/agents/bedrock/data_models.py index d2dc30b4..fff3918a 100644 --- a/libs/core/llmstudio_core/agents/bedrock/data_models.py +++ b/libs/core/llmstudio_core/agents/bedrock/data_models.py @@ -8,10 +8,10 @@ class BedrockAgent(AgentBase): - agentResourceRoleArn: str - agentStatus: str - agentVersion: str - agentArn: str + agent_resource_role_arn: str + agent_status: str + agent_arn: str + agent_alias: str class BedrockRun(RunBase): @@ -24,8 +24,9 @@ class BedrockResult(ResultBase): class BedrockCreateAgentRequest(CreateAgentRequest): - agent_resourcerole_arn: str + agent_resource_role_arn: str agent_alias: str + name: str class BedrockRunAgentRequest(RunAgentRequest): diff --git a/libs/core/llmstudio_core/agents/bedrock/manager.py b/libs/core/llmstudio_core/agents/bedrock/manager.py index 7711c199..90632026 100644 --- a/libs/core/llmstudio_core/agents/bedrock/manager.py +++ b/libs/core/llmstudio_core/agents/bedrock/manager.py @@ -3,10 +3,13 @@ import boto3 from llmstudio_core.agents.bedrock.data_models import ( BedrockAgent, + BedrockCreateAgentRequest, BedrockResult, BedrockRun, ) from llmstudio_core.agents.manager import AgentManager, agent_manager +from llmstudio_core.exceptions import AgentError +from pydantic import ValidationError SERVICE = "bedrock-agent" @@ -31,20 +34,99 @@ def _agent_config_name(): return "bedrock" def _validate_create_request(self, request): - raise NotImplementedError("Agents need to implement the method") + return BedrockCreateAgentRequest(**request) def _validate_run_request(self, request): raise NotImplementedError("Agents need to implement the method") - def _validate_create_request(self, request): + def _validate_result_request(self, request): raise NotImplementedError("Agents need to implement the method") - def create_agent(self, **kargs) -> BedrockAgent: + def create_agent(self, **kwargs) -> BedrockAgent: """ Creates a new instance of the agent. """ + try: + agent_request = self._validate_create_request( + dict( + **kwargs, + ) + ) + + except ValidationError as e: + raise AgentError(str(e)) + + bedrock_create = self._client.create_agent( + agentName=agent_request.name, + foundationModel=agent_request.model, + instruction=agent_request.instructions, + agentResourceRoleArn=agent_request.agent_resource_role_arn, + ) + + agentId = bedrock_create["agent"]["agentId"] + + # Wait for agent to reach 'NOT_PREPARED' status + agentStatus = "" + while agentStatus != "NOT_PREPARED": + response = self._client.get_agent(agentId=agentId) + agentStatus = response["agent"]["agentStatus"] + + # Configure code interpreter for the agent + response = self._client.create_agent_action_group( + actionGroupName="CodeInterpreterAction", + actionGroupState="ENABLED", + agentId=agentId, + agentVersion="DRAFT", + parentActionGroupSignature="AMAZON.CodeInterpreter", + ) - raise NotImplementedError("Agents need to implement the 'create' method.") + actionGroupId = response["agentActionGroup"]["actionGroupId"] + + # Wait for action group to reach 'ENABLED' status + actionGroupStatus = "" + while actionGroupStatus != "ENABLED": + response = self._client.get_agent_action_group( + agentId=agentId, actionGroupId=actionGroupId, agentVersion="DRAFT" + ) + actionGroupStatus = response["agentActionGroup"]["actionGroupState"] + + # Prepare the agent for use + response = self._client.prepare_agent(agentId=agentId) + + # Wait for agent to reach 'PREPARED' status + agentStatus = "" + while agentStatus != "PREPARED": + response = self._client.get_agent(agentId=agentId) + agentStatus = response["agent"]["agentStatus"] + + # Create an alias for the agent + response = self._client.create_agent_alias( + agentAliasName="test", agentId=agentId + ) + + agentAliasId = response["agentAlias"]["agentAliasId"] + + # Wait for agent alias to be prepared + agentAliasStatus = "" + while agentAliasStatus != "PREPARED": + response = self._client.get_agent_alias( + agentId=agentId, agentAliasId=agentAliasId + ) + agentAliasStatus = response["agentAlias"]["agentAliasStatus"] + + return BedrockAgent( + id=agentId, + created_at=int(bedrock_create["agent"]["createdAt"].timestamp()), + name=bedrock_create["agent"]["agentName"], + description=bedrock_create.get("agent", {}).get("description", None), + model=agent_request.model, + instructions=bedrock_create["agent"]["instruction"], + tools=[], + agent_arn=bedrock_create["agent"]["agentArn"], + agent_resource_role_arn=bedrock_create["agent"]["agentResourceRoleArn"], + agent_status=bedrock_create["agent"]["agentStatus"], + agent_alias=agent_request.agent_alias, + ) def run_agent(self, **kwargs) -> BedrockRun: """ diff --git a/libs/core/llmstudio_core/agents/data_models.py b/libs/core/llmstudio_core/agents/data_models.py index bd85f03b..35789a2b 100644 --- a/libs/core/llmstudio_core/agents/data_models.py +++ b/libs/core/llmstudio_core/agents/data_models.py @@ -17,10 +17,10 @@ class AgentBase(BaseModel): id: str created_at: int name: str - description: str + description: Optional[str] = None model: str - instructions: str - tools: list[Tool] + instructions: Optional[str] = None + tools: Optional[list[Tool]] = None class RunBase(BaseModel): @@ -35,8 +35,8 @@ class ResultBase(BaseModel): class CreateAgentRequest(BaseModel): model: str instructions: Optional[str] - description: Optional[str] - tools: Optional[list[Tool]] + description: Optional[str] = None + tools: Optional[list[Tool]] = None class RunAgentRequest(BaseModel): From 656ffca2e37ea1f26b0553c57b1faf5b1bea0c27 Mon Sep 17 00:00:00 2001 From: Clara Gadelho Date: Tue, 4 Feb 2025 09:26:44 +0000 Subject: [PATCH 2/4] feat: invoke agent --- .../agents/bedrock/data_models.py | 3 +- .../llmstudio_core/agents/bedrock/manager.py | 100 +++++++++++++----- .../core/llmstudio_core/agents/data_models.py | 8 +- 3 files changed, 81 insertions(+), 30 deletions(-) diff --git a/libs/core/llmstudio_core/agents/bedrock/data_models.py b/libs/core/llmstudio_core/agents/bedrock/data_models.py index fff3918a..40c324b9 100644 --- a/libs/core/llmstudio_core/agents/bedrock/data_models.py +++ b/libs/core/llmstudio_core/agents/bedrock/data_models.py @@ -11,7 +11,7 @@ class BedrockAgent(AgentBase): agent_resource_role_arn: str agent_status: str agent_arn: str - agent_alias: str + agent_alias_id: str class BedrockRun(RunBase): @@ -31,3 +31,4 @@ class BedrockCreateAgentRequest(CreateAgentRequest): class BedrockRunAgentRequest(RunAgentRequest): session_id: str + agent_alias_id: str diff --git a/libs/core/llmstudio_core/agents/bedrock/manager.py b/libs/core/llmstudio_core/agents/bedrock/manager.py index 90632026..a306404e 100644 --- a/libs/core/llmstudio_core/agents/bedrock/manager.py +++ b/libs/core/llmstudio_core/agents/bedrock/manager.py @@ -6,12 +6,14 @@ BedrockCreateAgentRequest, BedrockResult, BedrockRun, + BedrockRunAgentRequest, ) from llmstudio_core.agents.manager import AgentManager, agent_manager from llmstudio_core.exceptions import AgentError from pydantic import ValidationError -SERVICE = "bedrock-agent" +AGENT_SERVICE = "bedrock-agent" +RUNTIME_SERVICE = "bedrock-agent-runtime" @agent_manager @@ -19,7 +21,17 @@ class BedrockAgentManager(AgentManager): def __init__(self, **kwargs): super().__init__(**kwargs) self._client = boto3.client( - SERVICE, + service_name=AGENT_SERVICE, + region_name=self.region if self.region else os.getenv("BEDROCK_REGION"), + aws_access_key_id=self.access_key + if self.access_key + else os.getenv("BEDROCK_ACCESS_KEY"), + aws_secret_access_key=self.secret_key + if self.secret_key + else os.getenv("BEDROCK_SECRET_KEY"), + ) + self._runtime_client = boto3.client( + service_name=RUNTIME_SERVICE, region_name=self.region if self.region else os.getenv("BEDROCK_REGION"), aws_access_key_id=self.access_key if self.access_key @@ -37,15 +49,29 @@ def _validate_create_request(self, request): return BedrockCreateAgentRequest(**request) def _validate_run_request(self, request): - raise NotImplementedError("Agents need to implement the method") + return BedrockRunAgentRequest(**request) def _validate_result_request(self, request): raise NotImplementedError("Agents need to implement the method") def create_agent(self, **kwargs) -> BedrockAgent: """ - Creates a new instance of the agent. + This method validates the input parameters, creates a new agent using the client, + waits for the agent to reach the 'NOT_PREPARED' status, adds tools to the agent, + prepares the agent for use, creates an alias for the agent, and waits for the alias + to be prepared. + + Args: + **kwargs: Agent creation parameters. + + Returns: + BedrockAgent: An instance of the created BedrockAgent. + + Raises: + AgentError: If there is a validation error or if an unsupported tool type is provided. + """ + try: agent_request = self._validate_create_request( dict( @@ -71,24 +97,29 @@ def create_agent(self, **kwargs) -> BedrockAgent: response = self._client.get_agent(agentId=agentId) agentStatus = response["agent"]["agentStatus"] - # Configure code interpreter for the agent - response = self._client.create_agent_action_group( - actionGroupName="CodeInterpreterAction", - actionGroupState="ENABLED", - agentId=agentId, - agentVersion="DRAFT", - parentActionGroupSignature="AMAZON.CodeInterpreter", - ) + # Add tools to the agent + for tool in agent_request.tools: + if tool.type == "code_interpreter": + response = self._client.create_agent_action_group( + actionGroupName="CodeInterpreterAction", + actionGroupState="ENABLED", + agentId=agentId, + agentVersion="DRAFT", + parentActionGroupSignature="AMAZON.CodeInterpreter", + ) - actionGroupId = response["agentActionGroup"]["actionGroupId"] + actionGroupId = response["agentActionGroup"]["actionGroupId"] - # Wait for action group to reach 'ENABLED' status - actionGroupStatus = "" - while actionGroupStatus != "ENABLED": - response = self._client.get_agent_action_group( - agentId=agentId, actionGroupId=actionGroupId, agentVersion="DRAFT" - ) - actionGroupStatus = response["agentActionGroup"]["actionGroupState"] + actionGroupStatus = "" + while actionGroupStatus != "ENABLED": + response = self._client.get_agent_action_group( + agentId=agentId, + actionGroupId=actionGroupId, + agentVersion="DRAFT", + ) + actionGroupStatus = response["agentActionGroup"]["actionGroupState"] + else: + raise AgentError(f"Tool {tool.get('type')} not supported") # Prepare the agent for use response = self._client.prepare_agent(agentId=agentId) @@ -101,7 +132,7 @@ def create_agent(self, **kwargs) -> BedrockAgent: # Create an alias for the agent response = self._client.create_agent_alias( - agentAliasName="test", agentId=agentId + agentAliasName=agent_request.agent_alias, agentId=agentId ) agentAliasId = response["agentAlias"]["agentAliasId"] @@ -121,19 +152,38 @@ def create_agent(self, **kwargs) -> BedrockAgent: description=bedrock_create.get("agent", {}).get("description", None), model=agent_request.model, instructions=bedrock_create["agent"]["instruction"], - tools=[], + tools=agent_request.tools, agent_arn=bedrock_create["agent"]["agentArn"], agent_resource_role_arn=bedrock_create["agent"]["agentResourceRoleArn"], agent_status=bedrock_create["agent"]["agentStatus"], - agent_alias=agent_request.agent_alias, + agent_alias_id=agentAliasId, ) def run_agent(self, **kwargs) -> BedrockRun: """ Runs the agent """ - raise NotImplementedError( - "Agents need to implement the 'create_thread_and_run' method." + try: + run_request = self._validate_run_request( + dict( + **kwargs, + ) + ) + except ValidationError as e: + raise AgentError(str(e)) + + invoke_request = self._runtime_client.invoke_agent( + agentId=run_request.agent_id, + agentAliasId=run_request.agent_alias_id, + sessionId=run_request.session_id, + inputText=run_request.message.content, + ) + + return BedrockRun( + agent_id=run_request.agent_id, + status="completed", + session_id=run_request.session_id, + response=invoke_request, ) def retrieve_result(self, **kwargs) -> BedrockResult: diff --git a/libs/core/llmstudio_core/agents/data_models.py b/libs/core/llmstudio_core/agents/data_models.py index 35789a2b..5355f296 100644 --- a/libs/core/llmstudio_core/agents/data_models.py +++ b/libs/core/llmstudio_core/agents/data_models.py @@ -8,7 +8,7 @@ class Tool(BaseModel): class Message(BaseModel): - created_at: Optional[str] + created_at: Optional[str] = None role: Literal["user", "assistant"] content: Union[str, list] @@ -20,7 +20,7 @@ class AgentBase(BaseModel): description: Optional[str] = None model: str instructions: Optional[str] = None - tools: Optional[list[Tool]] = None + tools: Optional[list[Tool]] = [] class RunBase(BaseModel): @@ -36,11 +36,11 @@ class CreateAgentRequest(BaseModel): model: str instructions: Optional[str] description: Optional[str] = None - tools: Optional[list[Tool]] = None + tools: Optional[list[Tool]] = [] class RunAgentRequest(BaseModel): - assistant_id: str + agent_id: str message: Message From 1deb4df07c99bee4609c00786e5e051010f3b788 Mon Sep 17 00:00:00 2001 From: Clara Gadelho Date: Wed, 5 Feb 2025 23:11:06 +0000 Subject: [PATCH 3/4] feat: retrive result --- .../agents/bedrock/data_models.py | 5 -- .../llmstudio_core/agents/bedrock/manager.py | 87 +++++++++++++++++-- .../core/llmstudio_core/agents/data_models.py | 38 +++++++- 3 files changed, 116 insertions(+), 14 deletions(-) diff --git a/libs/core/llmstudio_core/agents/bedrock/data_models.py b/libs/core/llmstudio_core/agents/bedrock/data_models.py index 40c324b9..4f6b6603 100644 --- a/libs/core/llmstudio_core/agents/bedrock/data_models.py +++ b/libs/core/llmstudio_core/agents/bedrock/data_models.py @@ -1,7 +1,6 @@ from llmstudio_core.agents.data_models import ( AgentBase, CreateAgentRequest, - ResultBase, RunAgentRequest, RunBase, ) @@ -19,10 +18,6 @@ class BedrockRun(RunBase): response: dict -class BedrockResult(ResultBase): - session_id: str - - class BedrockCreateAgentRequest(CreateAgentRequest): agent_resource_role_arn: str agent_alias: str diff --git a/libs/core/llmstudio_core/agents/bedrock/manager.py b/libs/core/llmstudio_core/agents/bedrock/manager.py index a306404e..8c26e07e 100644 --- a/libs/core/llmstudio_core/agents/bedrock/manager.py +++ b/libs/core/llmstudio_core/agents/bedrock/manager.py @@ -4,10 +4,18 @@ from llmstudio_core.agents.bedrock.data_models import ( BedrockAgent, BedrockCreateAgentRequest, - BedrockResult, BedrockRun, BedrockRunAgentRequest, ) +from llmstudio_core.agents.data_models import ( + Attachment, + ImageFile, + ImageFileContent, + Message, + ResultBase, + RetrieveResultRequest, + TextContent, +) from llmstudio_core.agents.manager import AgentManager, agent_manager from llmstudio_core.exceptions import AgentError from pydantic import ValidationError @@ -52,7 +60,7 @@ def _validate_run_request(self, request): return BedrockRunAgentRequest(**request) def _validate_result_request(self, request): - raise NotImplementedError("Agents need to implement the method") + return RetrieveResultRequest(**request) def create_agent(self, **kwargs) -> BedrockAgent: """ @@ -161,8 +169,18 @@ def create_agent(self, **kwargs) -> BedrockAgent: def run_agent(self, **kwargs) -> BedrockRun: """ - Runs the agent + Runs the agent with the provided keyword arguments. + + This method validates the run request and invokes the agent using the runtime client. + If the validation fails, an AgentError is raised. + + Returns: + BedrockRun: An object containing the agent ID, status, session ID, and response of the run. + + Raises: + AgentError: If the run request validation fails. """ + try: run_request = self._validate_run_request( dict( @@ -186,8 +204,65 @@ def run_agent(self, **kwargs) -> BedrockRun: response=invoke_request, ) - def retrieve_result(self, **kwargs) -> BedrockResult: + def retrieve_result(self, **kwargs) -> ResultBase: """ - Retrieves an existing agent. + Retrieve the result based on the provided keyword arguments. + This method validates the result request and processes the event stream to + extract content and attachments. It constructs a message with the extracted + content and attachments and returns it wrapped in a ResultBase object. + + Returns: + ResultBase: An object containing the constructed message with content and attachments. + Raises: + AgentError: If the result request validation fails. """ - raise NotImplementedError("Agents need to implement the 'retrieve' method.") + + try: + result_request = self._validate_result_request( + dict( + **kwargs, + ) + ) + + except ValidationError as e: + raise AgentError(str(e)) + + content = [] + attachments = [] + event_stream = result_request.run.response.get("completion") + for event in event_stream: + if "chunk" in event: + chunk = event["chunk"] + if "bytes" in chunk: + content.append(TextContent(text=chunk["bytes"].decode("utf-8"))) + + if "files" in event: + files = event["files"]["files"] + for file in files: + if type == "image/png": + content.append( + ImageFileContent( + image_file=ImageFile( + file_name=file["name"], + file_content=file["bytes"], + file_type=file["type"], + ) + ) + ) + else: + attachments.append( + Attachment( + file_name=file["name"], + file_content=file["bytes"], + file_type=file["type"], + ) + ) + + message = Message( + thread_id=result_request.run.session_id, + role="assistant", + content=content, + attachments=attachments, + ) + + return ResultBase(message=message) diff --git a/libs/core/llmstudio_core/agents/data_models.py b/libs/core/llmstudio_core/agents/data_models.py index 5355f296..9204c082 100644 --- a/libs/core/llmstudio_core/agents/data_models.py +++ b/libs/core/llmstudio_core/agents/data_models.py @@ -7,10 +7,42 @@ class Tool(BaseModel): type: str +class TextContent(BaseModel): + type: Literal["text"] = "text" + text: str + + +class ImageFile(BaseModel): + file_id: Optional[str] = None + detail: Optional[Literal["low", "high", "auto"]] = "auto" + file_name: Optional[str] = None # need this for bedrock + file_content: Optional[bytes] = None # need this for bedrock + + +class ImageFileContent(BaseModel): + type: Literal["image_file"] = "image_file" + image_file: ImageFile + + +class Attachment(BaseModel): + file_id: Optional[str] = None + file_name: Optional[str] = None # need this for bedrock + file_content: Optional[bytes] = None # need this for bedrock + file_type: Optional[str] = None + tools: list[Tool] = [] + + class Message(BaseModel): - created_at: Optional[str] = None + id: Optional[str] = None + object: Optional[str] = "thread.message" + created_at: Optional[int] = None + thread_id: Optional[str] = None # in bedrock represents session_id role: Literal["user", "assistant"] - content: Union[str, list] + content: Union[str, list[Union[ImageFileContent, TextContent]]] + assistant_id: Optional[str] = None + run_id: Optional[str] = None + attachments: list[Attachment] = [] + metadata: Optional[dict] = {} class AgentBase(BaseModel): @@ -29,7 +61,7 @@ class RunBase(BaseModel): class ResultBase(BaseModel): - messages: list[Message] + message: Message class CreateAgentRequest(BaseModel): From c183ed497c2e18e855377083e3f2943174446175 Mon Sep 17 00:00:00 2001 From: Clara Gadelho Date: Thu, 6 Feb 2025 00:40:14 +0000 Subject: [PATCH 4/4] feat: send files to code interpreter --- .../llmstudio_core/agents/bedrock/manager.py | 34 +++++++++++++++++-- .../core/llmstudio_core/agents/data_models.py | 2 +- 2 files changed, 33 insertions(+), 3 deletions(-) diff --git a/libs/core/llmstudio_core/agents/bedrock/manager.py b/libs/core/llmstudio_core/agents/bedrock/manager.py index 8c26e07e..ddbcf0b3 100644 --- a/libs/core/llmstudio_core/agents/bedrock/manager.py +++ b/libs/core/llmstudio_core/agents/bedrock/manager.py @@ -190,11 +190,41 @@ def run_agent(self, **kwargs) -> BedrockRun: except ValidationError as e: raise AgentError(str(e)) + sessionState = {"files": []} + + for attachment in run_request.message.attachments: + if any(tool.type == "code_interpreter" for tool in attachment.tools): + sessionState["files"].append( + { + "name": attachment.file_name, + "source": { + "byteContent": { + "data": attachment.file_content, + "mediaType": attachment.file_type, + }, + "sourceType": "BYTE_CONTENT", + }, + "useCase": "CODE_INTERPRETER", + } + ) + + if isinstance(run_request.message.content, str): + input_text = run_request.message.content # Use it directly if it's a string + elif isinstance(run_request.message.content, list): + input_text = " ".join( + item.text + for item in run_request.message.content + if isinstance(item, TextContent) + ) + else: + input_text = "" # Default to an empty string if content is not valid + invoke_request = self._runtime_client.invoke_agent( agentId=run_request.agent_id, agentAliasId=run_request.agent_alias_id, sessionId=run_request.session_id, - inputText=run_request.message.content, + inputText=input_text, + sessionState=sessionState, ) return BedrockRun( @@ -239,7 +269,7 @@ def retrieve_result(self, **kwargs) -> ResultBase: if "files" in event: files = event["files"]["files"] for file in files: - if type == "image/png": + if file["type"] == "image/png": content.append( ImageFileContent( image_file=ImageFile( diff --git a/libs/core/llmstudio_core/agents/data_models.py b/libs/core/llmstudio_core/agents/data_models.py index 9204c082..b7533927 100644 --- a/libs/core/llmstudio_core/agents/data_models.py +++ b/libs/core/llmstudio_core/agents/data_models.py @@ -28,7 +28,7 @@ class Attachment(BaseModel): file_id: Optional[str] = None file_name: Optional[str] = None # need this for bedrock file_content: Optional[bytes] = None # need this for bedrock - file_type: Optional[str] = None + file_type: Optional[str] = None # need this for bedrock tools: list[Tool] = []