chore: import upstream snapshot with attribution

2026-07-13 12:43:34 +08:00
commit 41fe82d331
768 changed files with 57351 additions and 0 deletions
@@ -0,0 +1,109 @@
+import time
+
+import numpy as np
+
+from superagi.agent.agent_message_builder import AgentLlmMessageBuilder
+from superagi.agent.task_queue import TaskQueue
+from superagi.helper.error_handler import ErrorHandler
+from superagi.helper.json_cleaner import JsonCleaner
+from superagi.helper.prompt_reader import PromptReader
+from superagi.helper.token_counter import TokenCounter
+from superagi.lib.logger import logger
+from superagi.models.agent_execution import AgentExecution
+from superagi.models.agent_execution_feed import AgentExecutionFeed
+from superagi.models.workflows.agent_workflow_step import AgentWorkflowStep
+from superagi.models.workflows.agent_workflow_step_tool import AgentWorkflowStepTool
+from superagi.models.agent import Agent
+from superagi.types.queue_status import QueueStatus
+
+
+class QueueStepHandler:
+    """Handles the queue step of the agent workflow"""
+    def __init__(self, session, llm, agent_id: int, agent_execution_id: int):
+        self.session = session
+        self.llm = llm
+        self.agent_execution_id = agent_execution_id
+        self.agent_id = agent_id
+        self.organisation = Agent.find_org_by_agent_id(self.session, agent_id=self.agent_id)
+
+    def _queue_identifier(self, step_tool):
+        return step_tool.unique_id + "_" + str(self.agent_execution_id)
+
+    def _build_task_queue(self, step_tool):
+        return TaskQueue(self._queue_identifier(step_tool))
+
+    def execute_step(self):
+        execution = AgentExecution.get_agent_execution_from_id(self.session, self.agent_execution_id)
+        workflow_step = AgentWorkflowStep.find_by_id(self.session, execution.current_agent_step_id)
+        step_tool = AgentWorkflowStepTool.find_by_id(self.session, workflow_step.action_reference_id)
+        task_queue = self._build_task_queue(step_tool)
+
+        if not task_queue.get_status() or task_queue.get_status() == QueueStatus.COMPLETE.value:
+            task_queue.set_status(QueueStatus.INITIATED.value)
+
+        if task_queue.get_status() == QueueStatus.INITIATED.value:
+            self._add_to_queue(task_queue, step_tool)
+            execution.current_feed_group_id = "DEFAULT"
+            task_queue.set_status(QueueStatus.PROCESSING.value)
+
+        if not task_queue.get_tasks():
+            task_queue.set_status(QueueStatus.COMPLETE.value)
+            return "COMPLETE"
+        self._consume_from_queue(task_queue)
+        return "default"
+
+    def _add_to_queue(self, task_queue: TaskQueue, step_tool: AgentWorkflowStepTool):
+        assistant_reply = self._process_input_instruction(step_tool)
+        self._process_reply(task_queue, assistant_reply)
+
+    def _consume_from_queue(self, task_queue: TaskQueue):
+        tasks = task_queue.get_tasks()
+        agent_execution = AgentExecution.find_by_id(self.session, self.agent_execution_id)
+        if tasks:
+            task = task_queue.get_first_task()
+            # generating the new feed group id
+            agent_execution.current_feed_group_id = "GROUP_" + str(int(time.time()))
+            self.session.commit()
+            task_response_feed = AgentExecutionFeed(agent_execution_id=self.agent_execution_id,
+                                                    agent_id=self.agent_id,
+                                                    feed="Input: " + task,
+                                                    role="assistant",
+                                                    feed_group_id=agent_execution.current_feed_group_id)
+            self.session.add(task_response_feed)
+            self.session.commit()
+            task_queue.complete_task("PROCESSED")
+
+    def _process_reply(self, task_queue: TaskQueue, assistant_reply: str):
+        assistant_reply = JsonCleaner.extract_json_array_section(assistant_reply)
+        print("Queue reply:", assistant_reply)
+        task_array = np.array(eval(assistant_reply)).flatten().tolist()
+        for task in task_array:
+            task_queue.add_task(str(task))
+            logger.info("RAMRAM: Added task to queue: ", task)
+
+    def _process_input_instruction(self, step_tool):
+        prompt = self._build_queue_input_prompt(step_tool)
+        logger.info("Prompt: ", prompt)
+        agent_feeds = AgentExecutionFeed.fetch_agent_execution_feeds(self.session, self.agent_execution_id)
+        print(".........//////////////..........2")
+        messages = AgentLlmMessageBuilder(self.session, self.llm, self.llm.get_model(), self.agent_id, self.agent_execution_id) \
+            .build_agent_messages(prompt, agent_feeds, history_enabled=step_tool.history_enabled,
+                                  completion_prompt=step_tool.completion_prompt)
+        current_tokens = TokenCounter.count_message_tokens(messages, self.llm.get_model())
+        response = self.llm.chat_completion(messages, TokenCounter(session=self.session, organisation_id=self.organisation.id).token_limit(self.llm.get_model()) - current_tokens)
+        
+        if 'error' in response and response['message'] is not None:
+            ErrorHandler.handle_openai_errors(self.session, self.agent_id, self.agent_execution_id, response['message'])
+            
+        if 'content' not in response or response['content'] is None:
+            raise RuntimeError(f"Failed to get response from llm")
+        total_tokens = current_tokens + TokenCounter.count_message_tokens(response, self.llm.get_model())
+        AgentExecution.update_tokens(self.session, self.agent_execution_id, total_tokens)
+        assistant_reply = response['content']
+        return assistant_reply
+
+    def _build_queue_input_prompt(self, step_tool: AgentWorkflowStepTool):
+        queue_input_prompt = PromptReader.read_agent_prompt(__file__, "agent_queue_input.txt")
+        queue_input_prompt = queue_input_prompt.replace("{instruction}", step_tool.input_instruction)
+
+        return queue_input_prompt