Upload aworld_client.py

AWorld-main/aworlddistributed/aworld_client.py

@@ -0,0 +1,369 @@
+import asyncio
+import logging
+import os
+import traceback
+from datetime import datetime
+from typing import AsyncGenerator
+
+from aworld.models.llm import acall_llm_model, get_llm_model, acall_llm_model_stream
+from aworld.models.model_response import ModelResponse, LLMResponseError
+from pydantic import BaseModel, Field
+
+from base import AworldTask, AworldTaskResult, AworldTaskForm
+
+
+class TaskLogger:
+    """Task submission logger"""
+
+    def __init__(self, log_file: str = "aworld_task_submissions.log"):
+        self.log_file = 'task_logs/' + log_file
+        self._ensure_log_file_exists()
+
+    def _ensure_log_file_exists(self):
+        """ensure log file exists"""
+        if not os.path.exists(self.log_file):
+            os.makedirs(os.path.dirname(self.log_file), exist_ok=True)
+            with open(self.log_file, 'w', encoding='utf-8') as f:
+                f.write("# Aworld Task Submission Log\n")
+                f.write("# Format: [timestamp] task_id | agent_id | server | status | agent_answer | correct_answer | is_correct | details\n\n")
+
+    def log_task_submission(self, task: AworldTask, server: str, status: str, details: str = "", task_result: AworldTaskResult = None):
+        """log task submission"""
+        timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+        log_entry = f"[{timestamp}] {task.task_id} | {task.agent_id} | {task.node_id} | {status} | { task_result.data.get('agent_answer') if task_result and task_result.data else None } | {task_result.data.get('correct_answer') if task_result and task_result.data else None} | {task_result.data.get('gaia_correct') if task_result and task_result.data else None} |{details}\n"
+
+        try:
+            with open(self.log_file, 'a', encoding='utf-8') as f:
+                f.write(log_entry)
+        except Exception as e:
+            logging.error(f"Failed to write task submission log: {e}")
+
+    def log_task_result(self, task: AworldTask, result: ModelResponse):
+        """log task result to markdown file"""
+        try:
+            # create result directory
+            date_str = datetime.now().strftime("%Y%m%d")
+            result_dir = f"task_logs/result/{date_str}"
+            os.makedirs(result_dir, exist_ok=True)
+
+            # create markdown file
+            md_file = f"{result_dir}/{task.task_id}.md"
+
+            # concat content
+            content_parts = []
+            if hasattr(result, 'content') and result.content:
+                if isinstance(result.content, list):
+                    content_parts.extend(result.content)
+                else:
+                    content_parts.append(str(result.content))
+
+            # write to markdown file
+            file_exists = os.path.exists(md_file)
+            with open(md_file, 'a', encoding='utf-8') as f:
+                # only write title info when file not exists
+                if not file_exists:
+                    f.write(f"# Task Result: {task.task_id}\n\n")
+                    f.write(f"**Agent ID:** {task.agent_id}\n\n")
+                    f.write(f"**Timestamp:** {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n\n")
+                    f.write("## Content\n\n")
+
+                # write content parts
+                if content_parts:
+                    for i, content in enumerate(content_parts, 1):
+                        f.write(f"{content}\n\n")
+                else:
+                    f.write("No content available.\n\n")
+
+        except Exception as e:
+            logging.error(f"Failed to write task result log: {e}")
+
+task_logger = TaskLogger(log_file=f"aworld_task_submissions_{datetime.now().strftime('%Y%m%d')}.log")
+
+class AworldTaskClient(BaseModel):
+    """
+    AworldTaskClient
+    """
+    know_hosts: list[str] = Field(default_factory=list, description="aworldserver list")
+    tasks: list[AworldTask] = Field(default_factory=list, description="submitted task list")
+    task_states: dict[str, AworldTaskResult] = Field(default_factory=dict, description="task_states")
+
+    async def submit_task(self, task: AworldTask, background: bool = True):
+        if not self.know_hosts:
+            raise ValueError("No aworld server hosts configured.")
+        # 1. select aworld server from know_hosts using round-robin
+        if not hasattr(self, '_current_server_index'):
+            self._current_server_index = 0
+        aworld_server = self.know_hosts[self._current_server_index]
+        if not aworld_server.startswith("http"):
+            aworld_server = "http://" + aworld_server
+        self._current_server_index = (self._current_server_index + 1) % len(self.know_hosts)
+
+        # 2. call _submit_task
+        result = await self._submit_task(aworld_server, task, background)
+        # 3. update task_states
+        self.task_states[task.task_id] = result
+
+        
+    async def _submit_task(self, aworld_server, task: AworldTask, background: bool = True):
+        try:
+            logging.info(f"submit task#{task.task_id} to cluster#[{aworld_server}]")
+            if not background:
+                task_result = await self._submit_task_to_server(aworld_server, task)
+            else:
+                task_result = await self._async_submit_task_to_server(aworld_server, task)
+            return task_result
+        except Exception as e:
+            if isinstance(e, LLMResponseError):
+                if e.message and 'peer closed connection without sending complete message body (incomplete chunked read)' == e.message:
+                    task_logger.log_task_submission(task, aworld_server, "server_close_connection", str(e))
+                    logging.error(f"execute task to {task.node_id} server_close_connection: [{e}], please see replays wait a moment")
+                    return
+            traceback.print_exc()
+            logging.error(f"execute task to {task.node_id} execute_failed: [{e}], please see logs from server ")
+            task_logger.log_task_submission(task, aworld_server, "execute_failed", str(e))
+
+    async def _async_submit_task_to_server(self, aworld_server, task: AworldTask):
+        import httpx
+        from base import AworldTaskForm, AworldTaskResult
+        # 构建 AworldTaskForm
+        form_data = AworldTaskForm(task=task)
+        async with httpx.AsyncClient() as client:
+            resp = await client.post(f"{aworld_server}/api/v1/tasks/submit_task", json=form_data.model_dump())
+            resp.raise_for_status()
+            data = resp.json()
+            task_logger.log_task_submission(task, aworld_server, "submitted")
+            return AworldTaskResult(**data)
+
+    async def _submit_task_to_server(self, aworld_server, task: AworldTask):
+        # build params
+        llm_model = get_llm_model(
+            llm_provider="openai",
+            model_name=task.agent_id,
+            base_url=f"{aworld_server}/v1",
+            api_key="0p3n-w3bu!"
+        )
+        messages = [
+            {"role": "user", "content": task.agent_input}
+        ]
+        #call_llm_model
+        data = acall_llm_model_stream(llm_model, messages, stream=True, user={
+            "user_id": task.user_id,
+            "session_id": task.session_id,
+            "task_id": task.task_id,
+            "aworld_task": task.model_dump_json()
+        })
+        items = []
+        task_result = {}
+        if isinstance(data, AsyncGenerator):
+            async for item in data:
+                items.append(item)
+                if item.raw_response and item.raw_response.model_extra and item.raw_response.model_extra.get('node_id'):
+                    if not task.node_id:
+                        logging.info(f"submit task#{task.task_id} success. execute pod ip is [{item.raw_response.model_extra.get('node_id')}]")
+                        task.node_id = item.raw_response.model_extra.get('node_id')
+                        task_logger.log_task_submission(task, aworld_server, "submitted")
+
+                if item.content:
+                    task_logger.log_task_result(task, item)
+                    logging.info(f"task#{task.task_id} response data chunk is: {item}"[:500])
+
+                if item.raw_response and item.raw_response.model_extra and item.raw_response.model_extra.get(
+                        'task_output_meta'):
+                    task_result = item.raw_response.model_extra.get('task_output_meta')
+
+
+        elif isinstance(data, ModelResponse):
+            if data.raw_response and data.raw_response.model_extra and data.raw_response.model_extra.get('node_id'):
+                if not task.node_id:
+                    logging.info(f"submit task#{task.task_id} success. execute pod ip is [{data.raw_response.model_extra.get('node_id')}]")
+                task.node_id = data.raw_response.model_extra.get('node_id')
+
+            logging.info(f"task#{task.task_id} response data is: {data}")
+            task_logger.log_task_result(task, data)
+            if data.raw_response and data.raw_response.model_extra and data.raw_response.model_extra.get('task_output_meta'):
+                task_result = data.raw_response.model_extra.get('task_output_meta')
+
+        result = AworldTaskResult(task=task, server_host=aworld_server, data=task_result)
+        task_logger.log_task_submission(task, aworld_server, "execute_finished", task_result=result)
+        return result
+
+    async def get_task_state(self, task_id: str):
+        if not isinstance(self.task_states, dict):
+            self.task_states = dict(self.task_states)
+        return self.task_states.get(task_id, None)
+
+    async def download_task_results(
+        self, 
+        start_time: str = None, 
+        end_time: str = None, 
+        task_id: str = None, 
+        page_size: int = 100,
+        save_path: str = None
+    ) -> str:
+        """
+        Download task results and generate a JSONL format file
+        
+        Args:
+            start_time: Start time, format: YYYY-MM-DD HH:MM:SS
+            end_time: End time, format: YYYY-MM-DD HH:MM:SS
+            task_id: Task ID
+            page_size: Page size
+            save_path: Save path, if not specified, it will be generated automatically
+            
+        Returns:
+            str: Save path
+        """
+        if not self.know_hosts:
+            raise ValueError("No aworld server hosts configured.")
+        
+        # select server
+        if not hasattr(self, '_current_server_index'):
+            self._current_server_index = 0
+        aworld_server = self.know_hosts[self._current_server_index]
+        
+        logging.info(f"🚀 downloading task results from server: {aworld_server}")
+        
+        try:
+            import httpx
+            
+            # build query params
+            params = {"page_size": page_size}
+            if start_time:
+                params["start_time"] = start_time
+            if end_time:
+                params["end_time"] = end_time
+            if task_id:
+                params["task_id"] = task_id
+            
+            # send download request
+            async with httpx.AsyncClient(timeout=300.0) as client:  # 5分钟超时
+                response = await client.get(
+                    f"{aworld_server}/api/v1/tasks/download_task_results",
+                    params=params
+                )
+                response.raise_for_status()
+                
+                # if not specified save path, generate automatically
+                if not save_path:
+                    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+                    save_path = f"task_results_{timestamp}.jsonl"
+                
+                # ensure directory exists
+                save_dir = os.path.dirname(save_path) if os.path.dirname(save_path) else "."
+                os.makedirs(save_dir, exist_ok=True)
+                
+                with open(save_path, 'wb') as f:
+                    for chunk in response.iter_bytes():
+                        f.write(chunk)
+                
+                # calculate file size
+                file_size = os.path.getsize(save_path)
+                logging.info(f"✅ task results downloaded successfully, file: {save_path}, size: {file_size} bytes")
+                
+                return save_path
+                
+        except Exception as e:
+            logging.error(f"❌ download task results failed: {e}")
+            raise ValueError(f"❌ download task results failed: {str(e)}")
+
+    async def download_task_results_to_memory(
+        self, 
+        start_time: str = None, 
+        end_time: str = None, 
+        task_id: str = None, 
+        page_size: int = 100
+    ) -> list:
+        """
+        Download task results to memory, return parsed data list
+        
+        Args:
+            start_time: Start time, format: YYYY-MM-DD HH:MM:SS
+            end_time: End time, format: YYYY-MM-DD HH:MM:SS
+            task_id: Task ID
+            page_size: Page size
+            
+        Returns:
+            list: Task results data list
+        """
+        if not self.know_hosts:
+            raise ValueError("No aworld server hosts configured.")
+        
+        # select server
+        if not hasattr(self, '_current_server_index'):
+            self._current_server_index = 0
+        aworld_server = self.know_hosts[self._current_server_index]
+        
+        logging.info(f"🚀 downloading task results to memory from server: {aworld_server}")
+        
+        try:
+            import httpx
+            import json
+            
+            # build query params
+            params = {"page_size": page_size}
+            if start_time:
+                params["start_time"] = start_time
+            if end_time:
+                params["end_time"] = end_time
+            if task_id:
+                params["task_id"] = task_id
+            
+            # send download request
+            async with httpx.AsyncClient(timeout=300.0) as client:  # 5分钟超时
+                response = await client.get(
+                    f"{aworld_server}/api/v1/tasks/download_task_results",
+                    params=params
+                )
+                response.raise_for_status()
+                
+                # parse jsonl content
+                results = []
+                content = response.text
+                if content.strip():  # check content is not empty
+                    for line in content.strip().split('\n'):
+                        if line.strip():  # skip empty line
+                            try:
+                                result_data = json.loads(line)
+                                results.append(result_data)
+                            except json.JSONDecodeError as e:
+                                logging.warning(f"Failed to parse line: {line}, error: {e}")
+                
+                logging.info(f"✅ task results downloaded to memory successfully, total: {len(results)} records")
+                
+                return results
+                
+        except Exception as e:
+            logging.error(f"❌ download task results to memory failed: {e}")
+            raise ValueError(f"❌ download task results to memory failed: {str(e)}")
+
+    def parse_task_results_file(self, file_path: str) -> list:
+        """
+        Parse local task results jsonl file
+
+        Args:
+            file_path: jsonl file path
+
+        Returns:
+            list: Parsed task results list
+        """
+        import json
+
+        results = []
+        try:
+            with open(file_path, 'r', encoding='utf-8') as f:
+                for line_num, line in enumerate(f, 1):
+                    line = line.strip()
+                    if line:  # 跳过空行
+                        try:
+                            result_data = json.loads(line)
+                            results.append(result_data)
+                        except json.JSONDecodeError as e:
+                            logging.warning(f"Failed to parse line {line_num} in {file_path}: {e}")
+
+            logging.info(f"✅ parsed {len(results)} task results from {file_path}")
+            return results
+
+        except Exception as e:
+            logging.error(f"❌ failed to parse task results file {file_path}: {e}")
+            raise ValueError(f"❌ failed to parse task results file: {str(e)}")
+