ai-agent-book 精选快照（<2MB 代码与文档，来自 github.com/bojieli/ai-agent-book）

chapter9/gaia-experience/AWorld/aworld/cmd/web/routers/chats.py

@@ -0,0 +1,52 @@
+import logging
+import json
+from typing import Dict
+from fastapi import APIRouter, Depends, Request
+from fastapi.responses import StreamingResponse
+from aworld.cmd.data_model import AgentModel, ChatCompletionRequest
+from aworld.cmd.utils import agent_executor
+from aworld.cmd.utils.trace_summarize import summarize_trace
+from aworld.cmd.web.utils.users 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(request: Request) -> Dict[str, AgentModel]:
+    return request.app.state.agent_server.list_agents()
+
+
+@router.post("/chat/completions")
+async def chat_completion(
+    form_data: ChatCompletionRequest,
+    request: Request,
+    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, request.app.state.agent_server
+            ):
+                yield f"data: {json.dumps(chunk.model_dump(), ensure_ascii=False)}\n\n"
+        summarize_trace(form_data.trace_id)
+
+    return StreamingResponse(
+        generate_stream(),
+        media_type="text/event-stream",
+        headers={
+            "Cache-Control": "no-cache",
+            "Connection": "keep-alive",
+        },
+    )

chapter9/gaia-experience/AWorld/aworld/cmd/web/routers/sessions.py

@@ -0,0 +1,52 @@
+import logging
+from typing import List
+from fastapi import APIRouter, Depends, Request
+from pydantic import BaseModel, Field
+from aworld.cmd.data_model import SessionModel
+from aworld.cmd.web.utils.users import get_user_id_from_jwt
+
+logger = logging.getLogger(__name__)
+
+router = APIRouter()
+
+prefix = "/api/session"
+
+
+@router.get("/list")
+async def list_sessions(
+    request: Request,
+    user_id: str = Depends(get_user_id_from_jwt),
+) -> List[SessionModel]:
+    return await request.app.state.agent_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 request.app.state.agent_server.get_session_service().delete_session(
+            user_id, request.session_id
+        )
+        return CommonResponse.success()
+    except Exception as e:
+        return CommonResponse.error(str(e))

chapter9/gaia-experience/AWorld/aworld/cmd/web/routers/traces.py

@@ -0,0 +1,50 @@
+import json
+import logging
+from fastapi import APIRouter
+from aworld.trace.server import get_trace_server
+from aworld.trace.server.util import build_trace_tree, get_agent_flow
+from aworld.cmd.utils.trace_summarize import get_summarize_trace
+
+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):
+    data = get_agent_flow(trace_id)
+    await _add_trace_summary(trace_id, data.get('nodes'))
+    return data
+
+
+async def _add_trace_summary(trace_id, spans):
+    summary = await get_summarize_trace(trace_id)
+    json_summary_dict = {}
+    if summary:
+        json_summary = json.loads(summary)
+        json_summary_dict = {item['agent']: json.dumps(
+            item) for item in json_summary}
+
+    for span in spans:
+        if summary and "event_id" in span:
+            span['summary'] = json_summary_dict.get(span['event_id'])
+        span['attributes'] = None

chapter9/gaia-experience/AWorld/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)