Removed wallet and pool input

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