From c330bd34a45500a9381561fb6e33cdf1ca35d514 Mon Sep 17 00:00:00 2001 From: 2mmanu Date: Wed, 27 Nov 2024 10:39:11 -0800 Subject: [PATCH 1/2] letta local agent --- .../agents/letta_agent.py | 82 +++++++++++++++++++ python/src/tests/agents/test_letta_agent.py | 27 ++++++ 2 files changed, 109 insertions(+) create mode 100644 python/src/multi_agent_orchestrator/agents/letta_agent.py create mode 100644 python/src/tests/agents/test_letta_agent.py diff --git a/python/src/multi_agent_orchestrator/agents/letta_agent.py b/python/src/multi_agent_orchestrator/agents/letta_agent.py new file mode 100644 index 00000000..dd894c28 --- /dev/null +++ b/python/src/multi_agent_orchestrator/agents/letta_agent.py @@ -0,0 +1,82 @@ +from dataclasses import dataclass +from typing import List, Optional, Dict +from letta import LocalClient +from letta.schemas.memory import ChatMemory +from letta.schemas.letta_response import LettaResponse +from letta import LLMConfig, EmbeddingConfig +from multi_agent_orchestrator.types import ConversationMessage +from multi_agent_orchestrator.agents import Agent, AgentOptions +from multi_agent_orchestrator.utils import Logger + +@dataclass +class LettaAgentOptions(AgentOptions): + model_name: str = "letta" + model_name_embedding: str = "letta" + + +class LettaAgent(Agent): + """ + Represents a Letta agent that interacts with a runtime client. + Extends base Agent class and implements specific methods for Letta. + """ + def __init__(self, options: LettaAgentOptions): + super().__init__(options) + self.options = options + self.client = LocalClient() + self.client.set_default_llm_config(LLMConfig.default_config(model_name=options.model_name)) + self.client.set_default_embedding_config(EmbeddingConfig.default_config(model_name=options.model_name_embedding)) + + try: + agent_state = self.client.get_agent_by_name(agent_name=options.name) + except ValueError: + agent_state = self.client.create_agent( + name=options.name, + memory=ChatMemory( + human=f"My name is {options.name}", + persona=options.description + ) + ) + self._letta_id = agent_state.id + + async def process_request( + self, + input_text: str, + user_id: str, + session_id: str, + chat_history: List[ConversationMessage], + additional_params: Optional[Dict[str, str]] = None + ) -> ConversationMessage: + + response = self.client.send_message( + agent_id=self._letta_id, + message=input_text, + role="user" + ) + return ConversationMessage( + role="assistant", + content=LettaAgent._process_response(response), + ) + + @staticmethod + def _process_response(response: LettaResponse) -> str: + """ + Extracts the message from the 'send_message' function call in the LettaResponse. + + Args: + response (LettaResponse): The response object containing messages + + Returns: + str: The extracted message from send_message function call, or empty string if not found + """ + for message in response.messages: + if (message.message_type == "function_call" and + message.function_call.name == "send_message"): + + # Extract arguments string and convert to dictionary + import json + args = message.function_call.arguments + args_dict = json.loads(args) + return args_dict.get("message", "") + + return "" + diff --git a/python/src/tests/agents/test_letta_agent.py b/python/src/tests/agents/test_letta_agent.py new file mode 100644 index 00000000..ec1838f3 --- /dev/null +++ b/python/src/tests/agents/test_letta_agent.py @@ -0,0 +1,27 @@ +import pytest +from unittest.mock import Mock, patch +from multi_agent_orchestrator.types import ConversationMessage +from multi_agent_orchestrator.agents import LettaAgent, LettaAgentOptions +from letta.schemas.letta_response import LettaResponse + +@pytest.fixture +def mock_local_client(): + with patch('letta.LocalClient') as mock_client: + yield mock_client + +@pytest.fixture +def letta_agent(mock_local_client): + options = LettaAgentOptions( + name='test_agent_name', + description='test_agent description', + model_name='letta', + model_name_embedding='letta' + ) + return LettaAgent(options) + +def test_init(letta_agent, mock_local_client): + mock_local_client.return_value.get_agent_by_name.side_effect = ValueError() + mock_local_client.return_value.create_agent.return_value.id = 'test_agent_id' + + assert letta_agent.options.name == 'test_agent_name' + assert letta_agent.options.description == 'test_agent description' \ No newline at end of file From 33c22412d4b445e5e30d5fa5cd2ebd50aa9f99a9 Mon Sep 17 00:00:00 2001 From: 2mmanu Date: Wed, 27 Nov 2024 10:41:26 -0800 Subject: [PATCH 2/2] include letta agent in pkg --- python/setup.cfg | 4 +++- python/src/multi_agent_orchestrator/agents/__init__.py | 3 +++ python/test_requirements.txt | 3 ++- 3 files changed, 8 insertions(+), 2 deletions(-) diff --git a/python/setup.cfg b/python/setup.cfg index 3eb706ba..4d0bbe75 100644 --- a/python/setup.cfg +++ b/python/setup.cfg @@ -23,6 +23,8 @@ install_requires = boto3==1.34.151 [options.extras_require] +letta = + letta==0.5.5 anthropic = anthropic==0.32.0 openai = @@ -30,7 +32,7 @@ openai = all = anthropic==0.32.0 openai==1.55.0 - + letta==0.5.5 [options.packages.find] where = src exclude = diff --git a/python/src/multi_agent_orchestrator/agents/__init__.py b/python/src/multi_agent_orchestrator/agents/__init__.py index d13be8a5..878846d4 100644 --- a/python/src/multi_agent_orchestrator/agents/__init__.py +++ b/python/src/multi_agent_orchestrator/agents/__init__.py @@ -9,6 +9,7 @@ from .comprehend_filter_agent import ComprehendFilterAgent, ComprehendFilterAgentOptions from .chain_agent import ChainAgent, ChainAgentOptions from .bedrock_translator_agent import BedrockTranslatorAgent, BedrockTranslatorAgentOptions +from .letta_agent import LettaAgent, LettaAgentOptions try: from .anthropic_agent import AnthropicAgent, AnthropicAgentOptions @@ -37,6 +38,8 @@ 'BedrockTranslatorAgentOptions', 'ChainAgent', 'ChainAgentOptions', + 'LettaAgent', + 'LettaAgentOptions', ] if _ANTHROPIC_AVAILABLE: diff --git a/python/test_requirements.txt b/python/test_requirements.txt index 50b90582..5645ff82 100644 --- a/python/test_requirements.txt +++ b/python/test_requirements.txt @@ -4,4 +4,5 @@ boto3 anthropic moto pytest-mock -pytest-asyncio \ No newline at end of file +pytest-asyncio +letta \ No newline at end of file