refractor app service

Author: Fr4nc3

src/backend/app_kernel.py

@@ -3,6 +3,7 @@
 import logging
 import os
 import uuid
+import re
 from typing import List, Dict, Optional, Any
 
 # FastAPI imports
@@ -85,51 +86,7 @@
 async def input_task_endpoint(input_task: InputTask, request: Request):
     """
     Receive the initial input task from the user.
-
-    ---
-    tags:
-      - Input Task
-    parameters:
-      - name: user_principal_id
-        in: header
-        type: string
-        required: true
-        description: User ID extracted from the authentication header
-      - name: body
-        in: body
-        required: true
-        schema:
-          type: object
-          properties:
-            session_id:
-              type: string
-              description: Optional session ID, generated if not provided
-            description:
-              type: string
-              description: The task description
-            user_id:
-              type: string
-              description: The user ID associated with the task
-    responses:
-      200:
-        description: Task created successfully
-        schema:
-          type: object
-          properties:
-            status:
-              type: string
-            session_id:
-              type: string
-            plan_id:
-              type: string
-            description:
-              type: string
-            user_id:
-              type: string
-      400:
-        description: Missing or invalid user information
     """
-
     if not rai_success(input_task.description):
         print("RAI failed")
 
@@ -150,63 +107,91 @@ async def input_task_endpoint(input_task: InputTask, request: Request):
 
     if not user_id:
         track_event_if_configured("UserIdNotFound", {"status_code": 400, "detail": "no user"})
-
         raise HTTPException(status_code=400, detail="no user")
+        
+    # Generate session ID if not provided
     if not input_task.session_id:
         input_task.session_id = str(uuid.uuid4())
-
-    # Get the agents for this session
-    agents = await get_agents(input_task.session_id, user_id)
 
-    # Send the task to the planner agent
-    planner_agent = agents["PlannerAgent"]
+    # Set user ID from authenticated user
+    input_task.user_id = user_id
 
-    # Convert input task to JSON for the kernel function
-    input_task_json = input_task.json()
-    
-    # Use the planner to handle the task
-    result = await planner_agent.handle_input_task(
-        KernelArguments(input_task_json=input_task_json)
-    )
-    
-    # Extract plan ID from the result
-    # This is a simplified approach - in a real system, 
-    # we would properly parse the result to get the plan ID
-    memory_store = planner_agent._memory_store
-    plan = await memory_store.get_plan_by_session(input_task.session_id)
-    
-    if not plan or not plan.id:
+    try:
+        # Create just the planner agent instead of all agents
+        kernel, memory_store = await initialize_runtime_and_context(input_task.session_id, user_id)
+        planner_agent = await AgentFactory.create_agent(
+            agent_type=AgentType.PLANNER, 
+            session_id=input_task.session_id, 
+            user_id=user_id
+        )
+        
+        # Use the planner to handle the task - pass input_task_json for compatibility
+        input_task_json = input_task.json()
+        result = await planner_agent.handle_input_task(
+            KernelArguments(input_task_json=input_task_json)
+        )
+        
+        # Get plan from memory store
+        plan = await memory_store.get_plan_by_session(input_task.session_id)
+        
+        if not plan or not plan.id:
+            # If plan not found by session, try to extract plan ID from result
+            plan_id_match = re.search(r"Plan '([^']+)'", result)
+            
+            if plan_id_match:
+                plan_id = plan_id_match.group(1)
+                plan = await memory_store.get_plan(plan_id)
+            
+            # If still no plan found, handle the failure
+            if not plan or not plan.id:
+                track_event_if_configured(
+                    "PlanCreationFailed", 
+                    {
+                        "session_id": input_task.session_id,
+                        "description": input_task.description,
+                    }
+                )
+                return {
+                    "status": "Error: Failed to create plan",
+                    "session_id": input_task.session_id,
+                    "plan_id": "",
+                    "description": input_task.description,
+                }
+        
+        # Log custom event for successful input task processing
         track_event_if_configured(
-            "PlanCreationFailed", 
+            "InputTaskProcessed",
             {
+                "status": f"Plan created with ID: {plan.id}",
                 "session_id": input_task.session_id,
+                "plan_id": plan.id,
                 "description": input_task.description,
-            }
+            },
         )
+
         return {
-            "status": "Error: Failed to create plan",
+            "status": f"Plan created with ID: {plan.id}",
             "session_id": input_task.session_id,
-            "plan_id": "",
+            "plan_id": plan.id,
             "description": input_task.description,
         }
-    
-    # Log custom event for successful input task processing
-    track_event_if_configured(
-        "InputTaskProcessed",
-        {
-            "status": f"Plan created with ID: {plan.id}",
+        
+    except Exception as e:
+        logging.exception(f"Error handling input task: {e}")
+        track_event_if_configured(
+            "InputTaskError", 
+            {
+                "session_id": input_task.session_id,
+                "description": input_task.description,
+                "error": str(e),
+            }
+        )
+        return {
+            "status": f"Error creating plan: {str(e)}",
             "session_id": input_task.session_id,
-            "plan_id": plan.id,
+            "plan_id": "",
             "description": input_task.description,
-        },
-    )
-
-    return {
-        "status": f"Plan created with ID: {plan.id}",
-        "session_id": input_task.session_id,
-        "plan_id": plan.id,
-        "description": input_task.description,
-    }
+        }
 
 
 @app.post("/human_feedback")
@@ -658,6 +643,9 @@ async def get_agent_messages(session_id: str, request: Request) -> List[AgentMes
       - Agent Messages
     parameters:
       - name: session_id
+        in: path
+        type: string
+        required: true
         in: path
         type: string
         required: true

src/backend/kernel_agents/agent_base.py

@@ -127,6 +127,132 @@ async def handle_action_request_wrapper(*args, **kwargs):
         # Use agent name as plugin for handler
         self._kernel.add_function(self._agent_name, kernel_func)
 
+    async def handle_action_request(self, action_request_json: str) -> str:
+        """Handle an action request from another agent or the system.
+        
+        Args:
+            action_request_json: The action request as a JSON string
+            
+        Returns:
+            A JSON string containing the action response
+        """
+        # Parse the action request
+        action_request_dict = json.loads(action_request_json)
+        action_request = ActionRequest(**action_request_dict)
+        
+        # Get the step from memory
+        step: Step = await self._memory_store.get_step(
+            action_request.step_id, action_request.session_id
+        )
+        
+        if not step:
+            # Create error response if step not found
+            response = ActionResponse(
+                step_id=action_request.step_id,
+                status=StepStatus.failed,
+                message="Step not found in memory.",
+            )
+            return response.json()
+        
+        # Add messages to chat history for context
+        # This gives the agent visibility of the conversation history
+        self._chat_history.extend([
+            {"role": "assistant", "content": action_request.action},
+            {"role": "user", "content": f"{step.human_feedback}. Now make the function call"}
+        ])
+        
+        try:
+            # Use the agent to process the action
+            chat_history = self._chat_history.copy()
+            
+            # Call the agent to handle the action
+            agent_response = await self._agent.invoke(self._kernel, f"{action_request.action}\n\nPlease perform this action")
+            result = str(agent_response)
+            
+            # Store agent message in cosmos memory
+            await self._memory_store.add_item(
+                AgentMessage(
+                    session_id=action_request.session_id,
+                    user_id=self._user_id,
+                    plan_id=action_request.plan_id,
+                    content=f"{result}",
+                    source=self._agent_name,
+                    step_id=action_request.step_id,
+                )
+            )
+
+            # Track telemetry
+            track_event_if_configured(
+                "Base agent - Added into the cosmos",
+                {
+                    "session_id": action_request.session_id,
+                    "user_id": self._user_id,
+                    "plan_id": action_request.plan_id,
+                    "content": f"{result}",
+                    "source": self._agent_name,
+                    "step_id": action_request.step_id,
+                },
+            )
+
+        except Exception as e:
+            logging.exception(f"Error during agent execution: {e}")
+            
+            # Track error in telemetry
+            track_event_if_configured(
+                "Base agent - Error during agent execution, captured into the cosmos",
+                {
+                    "session_id": action_request.session_id,
+                    "user_id": self._user_id,
+                    "plan_id": action_request.plan_id,
+                    "content": f"{e}",
+                    "source": self._agent_name,
+                    "step_id": action_request.step_id,
+                },
+            )
+            
+            # Return an error response
+            response = ActionResponse(
+                step_id=action_request.step_id,
+                plan_id=action_request.plan_id,
+                session_id=action_request.session_id,
+                result=f"Error: {str(e)}",
+                status=StepStatus.failed,
+            )
+            return response.json()
+            
+        logging.info(f"Task completed: {result}")
+
+        # Update step status
+        step.status = StepStatus.completed
+        step.agent_reply = result
+        await self._memory_store.update_step(step)
+
+        # Track step completion in telemetry
+        track_event_if_configured(
+            "Base agent - Updated step and updated into the cosmos",
+            {
+                "status": StepStatus.completed,
+                "session_id": action_request.session_id,
+                "agent_reply": f"{result}",
+                "user_id": self._user_id,
+                "plan_id": action_request.plan_id,
+                "content": f"{result}",
+                "source": self._agent_name,
+                "step_id": action_request.step_id,
+            },
+        )
+
+        # Create and return action response
+        response = ActionResponse(
+            step_id=step.id,
+            plan_id=step.plan_id,
+            session_id=action_request.session_id,
+            result=result,
+            status=StepStatus.completed,
+        )
+        
+        return response.json()
+
     async def invoke_tool(self, tool_name: str, arguments: Dict[str, Any]) -> str:
         """Invoke a specific tool by name with the provided arguments.
         
@@ -275,6 +401,11 @@ def get_tools_from_config(cls, kernel: sk.Kernel, agent_type: str, config_path:
         kernel_functions = []
         plugin_name = f"{agent_type}_plugin"
 
+        # Early return if no tools defined - prevent empty iteration
+        if not config.get("tools"):
+            logging.info(f"No tools defined for agent type '{agent_type}'. Returning empty list.")
+            return kernel_functions
+        
         for tool in config.get("tools", []):
             try:
                 function_name = tool["name"]
@@ -301,8 +432,22 @@ def get_tools_from_config(cls, kernel: sk.Kernel, agent_type: str, config_path:
                 # Register the function with the kernel
                 kernel.add_function(plugin_name, kernel_func)
                 kernel_functions.append(kernel_func)
-                #logging.info(f"Successfully created dynamic tool '{function_name}' for {agent_type}")
+                logging.debug(f"Successfully created dynamic tool '{function_name}' for {agent_type}")
             except Exception as e:
                 logging.error(f"Failed to create tool '{tool.get('name', 'unknown')}': {str(e)}")
 
-        return kernel_functions
+        # Log the total number of tools created
+        if kernel_functions:
+            logging.info(f"Created {len(kernel_functions)} tools for agent type '{agent_type}'")
+        else:
+            logging.info(f"No tools were successfully created for agent type '{agent_type}'")
+                
+        return kernel_functions
+
+    def save_state(self) -> Mapping[str, Any]:
+        """Save the state of this agent."""
+        return {"memory": self._memory_store.save_state()}
+
+    def load_state(self, state: Mapping[str, Any]) -> None:
+        """Load the state of this agent."""
+        self._memory_store.load_state(state["memory"])