Merge pull request #64 from ygarg25/develop

Removed Wallet and pool input from reward and claim agent.

submodules/moragents_dockers/agents/src/claim_agent/src/agent.py

@@ -23,30 +23,18 @@ def get_response(self, message, wallet_address):
             available_rewards = {pool: amount for pool, amount in rewards.items() if amount > 0}
 
             if available_rewards:
-                self.conversation_state[wallet_address]["available_rewards"] = available_rewards
-                self.conversation_state[wallet_address]["state"] = "awaiting_receiver"
-                pools_str = " and ".join([f"pool {pool}" for pool in available_rewards.keys()])
-                amounts_str = " and ".join(
-                    [f"{amount} MOR in pool {pool}" for pool, amount in available_rewards.items()])
-                return f"You have rewards available in {pools_str}: {amounts_str}. Please provide the receiver address on Arbitrum where you want to receive these rewards.", "assistant", self.agent_info["name"]
-            else:
-                return f"No rewards found for your wallet address {wallet_address} in either pool. Claim cannot be processed.", "assistant", None
-
-        elif state == "awaiting_receiver":
-            if message[-1]['content'].startswith('0x') and len(message[-1]['content']) == 42:
-                receiver_address = message[-1]['content']
-                self.conversation_state[wallet_address]["receiver_address"] = receiver_address
+                selected_pool = max(available_rewards, key=available_rewards.get)
+                self.conversation_state[wallet_address]["available_rewards"] = {selected_pool: available_rewards[selected_pool]}
+                self.conversation_state[wallet_address]["receiver_address"] = wallet_address
                 self.conversation_state[wallet_address]["state"] = "awaiting_confirmation"
-                return self.prepare_transactions(wallet_address)
+                return f"You have {available_rewards[selected_pool]} MOR rewards available in pool {selected_pool}. Would you like to proceed with claiming these rewards?", "assistant", self.agent_info["name"]
             else:
-                return "Please provide a valid Ethereum address to receive your rewards.", "assistant", self.agent_info["name"]
+                return f"No rewards found for your wallet address {wallet_address} in either pool. Claim cannot be processed.", "assistant", None
 
         elif state == "awaiting_confirmation":
             user_input = message[-1]['content'].lower()
             if any(word in user_input for word in ['yes', 'proceed', 'confirm', 'claim']):
-                transactions = self.conversation_state[wallet_address]["transactions"]
-                tx_data_str = json.dumps(transactions, indent=2)
-                return f"Transaction data prepared for signing:\n\n{tx_data_str}", "assistant", None
+                return self.prepare_transactions(wallet_address)
             else:
                 return "Please confirm if you want to proceed with the claim by saying 'yes', 'proceed', 'confirm', or 'claim'.", "assistant", self.agent_info["name"]
 
@@ -65,7 +53,8 @@ def prepare_transactions(self, wallet_address):
                 return f"Error preparing transaction for pool {pool_id}: {str(e)}", "assistant", None
 
         self.conversation_state[wallet_address]["transactions"] = transactions
-        return f"Transactions prepared for claiming rewards from {len(transactions)} pool(s) to receiver address {receiver_address}. Please confirm if you want to proceed.", "assistant", self.agent_info["name"]
+        tx_data_str = json.dumps(transactions, indent=2)
+        return f"Transaction data prepared for signing:\n\n{tx_data_str}", "assistant", None
 
     def chat(self, request):
         try:
@@ -78,4 +67,4 @@ def chat(self, request):
             else:
                 return {"error": "Missing required parameters"}, 400
         except Exception as e:
-            return {"Error": str(e)}, 500
+            return {"Error": str(e)}, 500

submodules/moragents_dockers/agents/src/reward_agent/src/agent.py

@@ -1,6 +1,9 @@
 import json
+import logging
 from reward_agent.src import tools
 
+logger = logging.getLogger(__name__)
+
 
 class RewardAgent:
     def __init__(self, agent_info, llm, llm_ollama, embeddings, flask_app):
@@ -11,56 +14,36 @@ def __init__(self, agent_info, llm, llm_ollama, embeddings, flask_app):
         self.flask_app = flask_app
         self.tools_provided = tools.get_tools()
 
-    def get_response(self, message):
-        prompt = [
-            {
-                "role": "system",
-                "content": (
-                    "You are a reward checking agent. Your task is to help users check their MOR rewards accrued in either pool 0"
-                    "(capital providers pool) or pool 1 (code providers pool). "
-                    "You will facilitate the checking of their currently accrued rewards from the pool they mention and for the address they specify to check for. "
-                    "Extract the pool ID from the user's message: "
-                    "If the user mentions the capital pool, interpret it as pool ID 0. "
-                    "If the user mentions the code pool, interpret it as pool ID 1. "
-                    "Ask the user for the pool ID and the address they want to check their accrued rewards for. Note that these parameters are mandatory for the user to answer. "
-                    "Use the `get_current_user_reward` function to fetch the currently accrued MOR rewards using the pool ID and address provided by the user."
-                    "Remember that the value to be passed for pool_id in the `get_current_user_reward` function will be an integer - either 0 or 1 "
-                )
+    def get_response(self, message, wallet_address):
+        logger.info(f"Checking rewards for wallet address: {wallet_address}")
+
+        try:
+            rewards = {
+                0: tools.get_current_user_reward(wallet_address, 0),
+                1: tools.get_current_user_reward(wallet_address, 1)
             }
-        ]
-        prompt.extend(message)
-        result = self.llm.create_chat_completion(
-            messages=prompt,
-            tools=self.tools_provided,
-            tool_choice="auto",
-            temperature=0.01
-        )
-        if "tool_calls" in result["choices"][0]["message"].keys():
-            func = result["choices"][0]["message"]["tool_calls"][0]['function']
-            if func["name"] == "get_current_user_reward":
-                args = json.loads(func["arguments"])
-                try:
-                    reward = tools.get_current_user_reward(args['wallet_address'], args['pool_id'])
-                    return f"The current reward for wallet {args['wallet_address']} in pool {args['pool_id']} is {reward} MOR", "assistant", None
-                except Exception as e:
-                    return str(e), "assistant", None
-        return result["choices"][0]["message"]['content'], "assistant", "reward agent"
+
+            response = f"Your current MOR rewards:\n"
+            response += f"Capital Providers Pool (Pool 0): {rewards[0]} MOR\n"
+            response += f"Code Providers Pool (Pool 1): {rewards[1]} MOR"
+
+            logger.info(f"Rewards retrieved successfully for {wallet_address}")
+            return response, "assistant", None
+        except Exception as e:
+            logger.error(f"Error occurred while checking rewards: {str(e)}")
+            return f"An error occurred while checking your rewards: {str(e)}", "assistant", None
 
     def chat(self, request):
         try:
             data = request.get_json()
-            if 'prompt' in data:
+            if 'prompt' in data and 'wallet_address' in data:
                 prompt = data['prompt']
-                response, role, next_turn_agent = self.get_response([prompt])
-
-                # Check if we need more information
-                if "What is the wallet address?" in response or "What is the pool ID?" in response:
-                    next_turn_agent = self.agent_info["name"]  # Continue the conversation
-                else:
-                    next_turn_agent = None  # End the conversation
-
+                wallet_address = data['wallet_address']
+                response, role, next_turn_agent = self.get_response(prompt, wallet_address)
                 return {"role": role, "content": response, "next_turn_agent": next_turn_agent}
             else:
+                logger.warning("Missing required parameters in request")
                 return {"error": "Missing required parameters"}, 400
         except Exception as e:
-            return {"Error": str(e)}, 500
+            logger.error(f"Error in chat method: {str(e)}")
+            return {"Error": str(e)}, 500