Upload 4 files

aworld/cmd/web/routers/chats.py

@@ -0,0 +1,46 @@
+import logging
+import json
+from typing import Dict
+from fastapi import APIRouter, Depends
+from fastapi.responses import StreamingResponse
+from aworld.cmd import AgentModel, ChatCompletionRequest
+from aworld.cmd.utils import agent_loader, agent_executor
+from aworld.cmd.web.web_server import get_user_id_from_jwt
+import aworld.trace as trace
+
+logger = logging.getLogger(__name__)
+
+router = APIRouter()
+
+prefix = "/api/agent"
+
+
+@router.get("/list")
+@router.get("/models")
+async def list_agents() -> Dict[str, AgentModel]:
+    return agent_loader.list_agents()
+
+
+@router.post("/chat/completions")
+async def chat_completion(
+    form_data: ChatCompletionRequest, user_id: str = Depends(get_user_id_from_jwt)
+) -> StreamingResponse:
+    # Set user_id from JWT to form_data
+    form_data.user_id = user_id
+
+    async def generate_stream():
+        async with trace.span(
+            "/chat/chat_completion", attributes={"model": form_data.model}
+        ) as span:
+            form_data.trace_id = span.get_trace_id()
+            async for chunk in agent_executor.stream_run(form_data):
+                yield f"data: {json.dumps(chunk.model_dump(), ensure_ascii=False)}\n\n"
+
+    return StreamingResponse(
+        generate_stream(),
+        media_type="text/event-stream",
+        headers={
+            "Cache-Control": "no-cache",
+            "Connection": "keep-alive",
+        },
+    )

aworld/cmd/web/routers/sessions.py

@@ -0,0 +1,50 @@
+import logging
+from typing import List
+from fastapi import APIRouter, Depends
+from pydantic import BaseModel, Field
+from aworld.cmd import SessionModel
+from aworld.cmd.utils.agent_server import CURRENT_SERVER
+from aworld.cmd.web.web_server import get_user_id_from_jwt
+
+logger = logging.getLogger(__name__)
+
+router = APIRouter()
+
+prefix = "/api/session"
+
+
+@router.get("/list")
+async def list_sessions(
+    user_id: str = Depends(get_user_id_from_jwt),
+) -> List[SessionModel]:
+    return await CURRENT_SERVER.get_session_service().list_sessions(user_id)
+
+
+class CommonResponse(BaseModel):
+    code: int = Field(..., description="The code")
+    message: str = Field(..., description="The message")
+
+    @staticmethod
+    def success(message: str = "success"):
+        return CommonResponse(code=0, message=message)
+
+    @staticmethod
+    def error(message: str):
+        return CommonResponse(code=1, message=message)
+
+
+class DeleteSessionRequest(BaseModel):
+    session_id: str = Field(..., description="The session id")
+
+
+@router.post("/delete")
+async def delete_session(
+    request: DeleteSessionRequest, user_id: str = Depends(get_user_id_from_jwt)
+) -> CommonResponse:
+    try:
+        await CURRENT_SERVER.get_session_service().delete_session(
+            user_id, request.session_id
+        )
+        return CommonResponse.success()
+    except Exception as e:
+        return CommonResponse.error(str(e))

aworld/cmd/web/routers/traces.py

@@ -0,0 +1,71 @@
+import logging
+from fastapi import APIRouter
+from aworld.trace.server import get_trace_server
+from aworld.trace.constants import RunType
+from aworld.trace.server.util import build_trace_tree
+
+logger = logging.getLogger(__name__)
+
+router = APIRouter()
+
+prefix = "/api/trace"
+
+
+@router.get("/list")
+async def list_traces():
+    storage = get_trace_server().get_storage()
+    trace_data = []
+    for trace_id in storage.get_all_traces():
+        spans = storage.get_all_spans(trace_id)
+        spans_sorted = sorted(spans, key=lambda x: x.start_time)
+        trace_tree = build_trace_tree(spans_sorted)
+        trace_data.append({
+            'trace_id': trace_id,
+            'root_span': trace_tree,
+        })
+    return {
+        "data": trace_data
+    }
+
+
+@router.get("/agent")
+async def get_agent_trace(trace_id: str):
+    storage = get_trace_server().get_storage()
+    spans = storage.get_all_spans(trace_id)
+    spans_dict = {span.span_id: span.dict() for span in spans}
+
+    filtered_spans = {}
+    for span_id, span in spans_dict.items():
+        if span.get('is_event', False) and span.get('run_type') == RunType.AGNET.value:
+            span['show_name'] = _get_agent_show_name(span)
+            filtered_spans[span_id] = span
+
+    for span in list(filtered_spans.values()):
+        parent_id = span['parent_id'] if span['parent_id'] else None
+
+        while parent_id and parent_id not in filtered_spans:
+            parent_span = spans_dict.get(parent_id)
+            parent_id = parent_span['parent_id'] if parent_span and parent_span['parent_id'] else None
+
+        if parent_id:
+            parent_span = filtered_spans.get(parent_id)
+            if not parent_span:
+                continue
+
+            if 'children' not in parent_span:
+                parent_span['children'] = []
+            parent_span['children'].append(span)
+
+    root_spans = [span for span in filtered_spans.values()
+                  if span['parent_id'] is None or span['parent_id'] not in filtered_spans]
+    return {
+        "data": root_spans
+    }
+
+
+def _get_agent_show_name(span: dict):
+    agent_name_prefix = "agent_event_"
+    name = span.get("name")
+    if name and name.startswith(agent_name_prefix):
+        name = name[len(agent_name_prefix):]
+    return name

aworld/cmd/web/routers/workspaces.py

@@ -0,0 +1,74 @@
+import logging
+import os
+from typing import List, Optional
+from pydantic import BaseModel
+
+from fastapi import APIRouter, HTTPException, status, Query, Body
+
+from aworld.output import WorkSpace, ArtifactType
+from aworld.output.utils import load_workspace
+
+router = APIRouter()
+
+prefix = "/api/workspaces"
+
+@router.get("/{workspace_id}/tree")
+async def get_workspace_tree(workspace_id: str):
+    logging.info(f"get_workspace_tree: {workspace_id}")
+    workspace = await get_workspace(workspace_id)
+    return workspace.generate_tree_data()
+
+
+class ArtifactRequest(BaseModel):
+    artifact_ids: Optional[List[str]] = None
+    artifact_types: Optional[List[str]] = None
+
+
+@router.post("/{workspace_id}/artifacts")
+async def get_workspace_artifacts(workspace_id: str, request: ArtifactRequest):
+    """
+    Get artifacts by workspace id and filter by a list of artifact types.
+    Args:
+        workspace_id: Workspace ID
+        request: Request body containing optional artifact_types list
+    Returns:
+        Dict with filtered artifacts
+    """
+    artifact_types = request.artifact_types
+    if artifact_types:
+        # Validate all types
+        invalid_types = [t for t in artifact_types if t not in ArtifactType.__members__]
+        if invalid_types:
+            logging.error(f"Invalid artifact_types: {invalid_types}")
+            raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST,
+                                detail=f"Invalid artifact types: {invalid_types}")
+        logging.info(f"Fetching artifacts of types: {artifact_types}")
+    else:
+        logging.info(f"Fetching all artifacts (no type filter)")
+
+    workspace = await get_workspace(workspace_id)
+    all_artifacts = workspace.list_artifacts()
+    filtered_artifacts = all_artifacts
+    if request.artifact_ids:
+        filtered_artifacts = [a for a in filtered_artifacts if a.artifact_id in request.artifact_ids]
+    if artifact_types:
+        filtered_artifacts = [a for a in filtered_artifacts if a.artifact_type.name in artifact_types]
+    
+    return {
+        "data": filtered_artifacts
+    }
+
+
+@router.get("/{workspace_id}/file/{artifact_id}/content")
+async def get_workspace_file_content(workspace_id: str, artifact_id: str):
+    logging.info(f"get_workspace_file_content: {workspace_id}, {artifact_id}")
+    workspace = await get_workspace(workspace_id)
+    return {
+        "data": workspace.get_file_content_by_artifact_id(artifact_id)
+    }
+
+
+async def get_workspace(workspace_id: str) -> WorkSpace:
+    workspace_type = os.environ.get("WORKSPACE_TYPE", "local")
+    workspace_path = os.environ.get("WORKSPACE_PATH", "./data/workspaces")
+    return await load_workspace(workspace_id, workspace_type, workspace_path)