agent-api-server 2.1.7__py3-none-any.whl

agent_api_server/api/v1/thread.py

@@ -0,0 +1,563 @@
+import logging
+import traceback
+import asyncio
+from fastapi import Request
+from http.cookies import SimpleCookie
+from contextlib import AsyncExitStack, asynccontextmanager
+from typing import Dict, Any, Optional, AsyncGenerator, Tuple, List
+from fastapi.responses import StreamingResponse
+from fastapi import APIRouter, Body, status, HTTPException
+from agent_api_server.cache.redis_cache import ThreadState, AsyncRedisThreadStorage
+from langgraph.types import Command
+from agent_api_server.shared.base_model import (
+    ThreadInfo,
+    RunResponse,
+    error_response,
+    sse_response_example
+)
+from langchain_core.messages import AIMessage
+from agent_api_server.shared.util_func import load_graph_config, load_graph, get_env
+from agent_api_server.shared.message import (
+    message_generator,
+    langchain_to_chat_message
+)
+
+logger = logging.getLogger(__name__)
+api_router = APIRouter()
+
+THREAD_NOT_FOUND = "Specified thread does not exist"
+CHECKPOINT_NOT_FOUND = "Checkpoint not found for this graph"
+
+class TaskCancelledError(Exception):
+    """Custom exception for cancelled tasks"""
+    pass
+
+class ThreadCancellationManager:
+    def __init__(self):
+        self._events: Dict[str, asyncio.Event] = {}
+        self._lock = asyncio.Lock()
+
+    async def get_event(self, thread_id: str) -> asyncio.Event:
+        async with self._lock:
+            if thread_id not in self._events:
+                self._events[thread_id] = asyncio.Event()
+            return self._events[thread_id]
+
+    async def cancel(self, thread_id: str):
+        async with self._lock:
+            if thread_id in self._events:
+                self._events[thread_id].set()
+                del self._events[thread_id]
+
+    async def is_cancelled(self, thread_id: str) -> bool:
+        async with self._lock:
+            if thread_id in self._events:
+                return self._events[thread_id].is_set()
+            return False
+
+cancellation_manager = ThreadCancellationManager()
+
+
+async def _get_thread_or_404(thread_id: str, graph_name: Optional[str] = '') -> ThreadState:
+    storage = AsyncRedisThreadStorage.get_worker_instance()
+    thread_data = await storage.get_thread(thread_id)
+    if thread_data:
+        return thread_data
+
+    if not graph_name:
+        raise HTTPException(
+            status_code=status.HTTP_404_NOT_FOUND,
+            detail={
+                "error": "not_found",
+                "message": THREAD_NOT_FOUND,
+                "thread_id": thread_id
+            }
+        )
+
+    logger.warning(f"Thread {thread_id} not found, creating with graph: {graph_name}")
+
+    try:
+        new_thread = await storage.create_thread_with_id(thread_id=thread_id, graph_name=graph_name)
+        logger.info(f"Created thread {thread_id} for graph {graph_name}")
+        return new_thread
+    except Exception as e:
+        logger.error(f"Failed to create thread {thread_id}: {str(e)}")
+        raise HTTPException(
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            detail={
+                "error": "creation_failed",
+                "message": f"Failed to create thread: {str(e)}",
+                "graph_name": graph_name,
+                "thread_id": thread_id
+            }
+        )
+
+def _get_run_config(thread_id: str, ts_tenant: str, ei_token: str, graph_name: str, files: List[Dict[str, Any]] = None) -> Dict[str, Any]:
+    config = {"configurable": dict(get_env(ts_tenant=ts_tenant))}
+    config["configurable"]["graph_name"] = graph_name
+    config["configurable"]["thread_id"] = thread_id
+    config["configurable"]["TSTenant"] = ts_tenant
+    config["configurable"]["EIToken"] = ei_token
+    config["configurable"]["files"] = files or []
+    return config
+
+
+async def _check_stop_flag(thread_id: str, storage: AsyncRedisThreadStorage):
+    """Check stop flag and raise if cancelled"""
+    if await storage.should_stop(thread_id):
+        await storage.update_thread(thread_id, status="cancelled")
+        raise TaskCancelledError(f"Thread {thread_id} was cancelled")
+
+
+async def _monitored_stream(
+        thread_id: str,
+        generator: AsyncGenerator,
+        storage: AsyncRedisThreadStorage
+) -> AsyncGenerator:
+    """Wrap generator with stop flag checking"""
+    try:
+        async for item in generator:
+            # 添加轻微延迟避免tight loop
+            await asyncio.sleep(0.05)
+            await _check_stop_flag(thread_id, storage)
+            yield item
+    except TaskCancelledError:
+        logger.info(f"Stream for thread {thread_id} was cancelled")
+        await storage.clear_stop_flag(thread_id)
+        yield "data: [CANCELED]\n\n"
+        raise
+    except Exception as e:
+        logger.error(f"Stream error for thread {thread_id}: {str(e)}")
+        await storage.clear_stop_flag(thread_id)
+        yield "data: [CANCELED]\n\n"
+        raise
+
+
+@api_router.get(
+    "/{thread_id}/status",
+    response_model=ThreadInfo,
+    summary="Get thread status",
+    responses={
+        404: {"model": error_response, "description": "Thread not found"},
+        500: {"model": error_response, "description": "Internal server error"}
+    }
+)
+async def get_thread_status(thread_id: str) -> ThreadInfo:
+    """Get the current status of a processing thread."""
+    try:
+        state = await _get_thread_or_404(thread_id)
+        return ThreadInfo(
+            thread_id=state.thread_id,
+            graph_name=state.graph_name,
+            status=state.status,
+        )
+    except HTTPException:
+        raise
+    except Exception as e:
+        logger.error(f"Failed to get thread status: {str(e)}", exc_info=True)
+        raise HTTPException(
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            detail={
+                "error": "status_check_failed",
+                "message": "Failed to retrieve thread status",
+                "thread_id": thread_id
+            }
+        )
+
+
+@api_router.post(
+    "/",
+    response_model=ThreadInfo,
+    status_code=status.HTTP_201_CREATED,
+    summary="Create new processing thread",
+    responses={
+        422: {"model": error_response, "description": "Invalid configuration"},
+        500: {"model": error_response, "description": "Internal server error"}
+    }
+)
+async def create_thread(graph_name: str, thread_id: Optional[str] = None) -> ThreadInfo:
+    """Create a new processing thread with the specified graph."""
+    if thread_id:
+        try:
+            state = await _get_thread_or_404(thread_id)
+            logger.info(f"{thread_id} already exists, returning existing thread")
+            return ThreadInfo(
+                thread_id=state.thread_id,
+                graph_name=state.graph_name,
+                status=state.status,
+            )
+        except HTTPException:
+            pass  # Continue to create new thread if not exists
+
+    validated_name, _, _ = await load_graph(graph_name, await load_graph_config(), False)
+    try:
+        storage = AsyncRedisThreadStorage.get_worker_instance()
+        thread_state = await storage.create_thread(graph_name=validated_name)
+
+        return ThreadInfo(
+            thread_id=thread_state.thread_id,
+            graph_name=thread_state.graph_name,
+            status=thread_state.status,
+        )
+
+    except ValueError as e:
+        raise HTTPException(
+            status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+            detail={
+                "error": "invalid_config",
+                "message": str(e),
+                "graph_name": graph_name
+            }
+        )
+    except Exception as e:
+        logger.critical(f"Thread creation failed: {traceback.format_exc()}")
+        raise HTTPException(
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            detail={
+                "error": "creation_failed",
+                "message": "Internal server error",
+                "graph_name": graph_name
+            }
+        )
+
+@api_router.post(
+    "/{thread_id}/stream",
+    response_class=StreamingResponse,
+    responses=sse_response_example(),
+    summary="Stream messages for a thread"
+)
+async def stream(
+        thread_id: str,
+        request: Request,
+        inputs: Dict[str, Any] = Body(..., embed=True),
+        files: List[Dict[str, Any]] = Body(default=[], embed=True)
+) -> StreamingResponse:
+    """Stream messages for a processing thread."""
+    get_graph_name = request.headers.get('X-Agent-Name', '')
+    cookie_header = request.headers.get("cookie", "")
+    cookie = SimpleCookie()
+    cookie.load(cookie_header)
+    ts_tenant = cookie.get("TSTenant").value if "TSTenant" in cookie else None
+    ei_token = cookie.get("EIToken").value if "EIToken" in cookie else None
+
+    state = await _get_thread_or_404(thread_id, get_graph_name)
+    if state.status == "running":
+        raise HTTPException(
+            status_code=status.HTTP_400_BAD_REQUEST,
+            detail={
+                "error": "already_running",
+                "message": "Thread is already running",
+                "thread_id": thread_id
+            }
+        )
+
+    if get_graph_name != '':
+        logger.info(f"get graph name {get_graph_name} from header is not None, so use this graph name to load and chat")
+        state.graph_name = get_graph_name
+
+    storage = AsyncRedisThreadStorage.get_worker_instance()
+    monitored_gen = _monitored_stream(
+        thread_id,
+        message_generator(state=state, ts_tenant=ts_tenant, ei_token=ei_token, inputs=inputs, files=files),
+        storage
+    )
+
+    return StreamingResponse(
+        monitored_gen,
+        media_type="text/event-stream",
+    )
+
+
+@api_router.post(
+    "/{thread_id}/run",
+    response_model=RunResponse,
+    summary="Execute processing thread",
+    responses={
+        404: {"model": error_response, "description": "Thread not found"},
+        408: {"model": error_response, "description": "Request timeout"},
+        500: {"model": error_response, "description": "Execution failed"}
+    },
+    deprecated=True
+)
+async def run_thread(
+        thread_id: str,
+        request: Request,
+        inputs: Dict[str, Any] = Body(..., embed=True),
+        files: List[Dict[str, Any]] = Body(default=[], embed=True)
+) -> RunResponse:
+    """
+    Execute processing thread with enhanced error handling and timeout control.
+    """
+    get_graph_name = request.headers.get('X-Agent-Name', '')
+    cookie_header = request.headers.get("cookie", "")
+    cookie = SimpleCookie()
+    cookie.load(cookie_header)
+    ts_tenant = cookie.get("TSTenant").value if "TSTenant" in cookie else None
+    ei_token = cookie.get("EIToken").value if "EIToken" in cookie else None
+
+    storage = AsyncRedisThreadStorage.get_worker_instance()
+    try:
+        state = await _get_thread_or_404(thread_id, get_graph_name)
+        if state.status == "running":
+            raise HTTPException(
+                status_code=status.HTTP_400_BAD_REQUEST,
+                detail={
+                    "error": "already_running",
+                    "message": "Thread is already running",
+                    "thread_id": thread_id
+                }
+            )
+
+        if get_graph_name != '':
+            logger.info(
+                f"get graph name {get_graph_name} from header is not None, so use this graph name to load and chat")
+            state.graph_name = get_graph_name
+
+        async with AsyncExitStack() as stack:
+            try:
+                _, graph_instance, _ = await stack.enter_async_context(
+                    timeout_context(10.0, load_graph(state.graph_name, await load_graph_config(), True))
+                )
+
+                try:
+                    await storage.clear_stop_flag(thread_id)
+                    if not await storage.update_thread(thread_id, status="running"):
+                        raise HTTPException(status_code=500, detail="Failed to update thread status")
+                except Exception as e:
+                    logger.error(f"Redis operation failed: {str(e)}")
+                    raise HTTPException(status_code=500, detail="Storage operation failed")
+
+                _, response_events = await stack.enter_async_context(
+                    timeout_context(300.0, _execute_graph_with_redis_check(
+                        graph_instance=graph_instance,
+                        run_cfg=_get_run_config(thread_id=thread_id, ei_token=ei_token, ts_tenant=ts_tenant, graph_name=state.graph_name, files=files),
+                        inputs=inputs,
+                        thread_id=thread_id,
+                        storage=storage
+                    ))
+                )
+
+                response_type, response = response_events[-1]
+                if  "__interrupt__" in response:
+                    output = langchain_to_chat_message(
+                        AIMessage(content=response["__interrupt__"][0].value)
+                    )
+
+                    await storage.update_thread(thread_id, status="completed")
+                    return RunResponse(
+                        status="completed",
+                        thread_id=thread_id,
+                        result=output
+                    )
+
+                output = langchain_to_chat_message(response["messages"][-1])
+                await storage.update_thread(thread_id, status="completed")
+                return RunResponse(
+                    status="completed",
+                    thread_id=thread_id,
+                    result=output
+                )
+            except TaskCancelledError:
+                logger.info(f"Thread {thread_id} was cancelled")
+                await storage.clear_stop_flag(thread_id)
+                await storage.update_thread(thread_id, status="completed")
+                raise HTTPException(
+                    status_code=status.HTTP_200_OK,
+                    detail={
+                        "error": "cancelled",
+                        "thread_id": thread_id,
+                        "message": "Execution was cancelled by user"
+                    }
+                )
+
+    except asyncio.TimeoutError as te:
+        await storage.update_thread(thread_id, status="completed")
+        raise HTTPException(
+            status_code=status.HTTP_408_REQUEST_TIMEOUT,
+            detail={
+                "error": "timeout",
+                "thread_id": thread_id,
+                "message": f"Operation timed out after {300 if 'execute' in str(te) else 10} seconds"
+            }
+        )
+    except HTTPException:
+        raise
+    except Exception as e:
+        logger.error(f"Execution failed: {thread_id}", exc_info=True)
+        await storage.update_thread(thread_id, status="completed")
+        raise HTTPException(
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            detail={
+                "error": "execution_failed",
+                "thread_id": thread_id,
+                "message": str(e)[:200]
+            }
+        )
+
+async def _execute_graph_with_redis_check(
+        graph_instance,
+        run_cfg: Dict[str, Any],
+        inputs: Dict[str, Any],
+        thread_id: str,
+        storage: AsyncRedisThreadStorage
+) -> Tuple[str, Any]:
+    execute_task = None
+    try:
+        current_state = await graph_instance.aget_state(config=run_cfg)
+        invoke_input = (
+            Command(resume=inputs.get("resume", ""))
+            if any(getattr(task, "interrupts", None) for task in current_state.tasks)
+            else inputs or {}
+        )
+
+        execute_task = asyncio.create_task(
+            graph_instance.ainvoke(
+                invoke_input,
+                config=run_cfg,
+                stream_mode=["values"]
+            )
+        )
+
+        while not execute_task.done():
+            if await storage.should_stop(thread_id):
+                execute_task.cancel()
+                try:
+                    await execute_task
+                except asyncio.CancelledError:
+                    pass
+                raise TaskCancelledError(f"Thread {thread_id} was cancelled")
+            await asyncio.sleep(0.1)
+
+        return "completed", await execute_task
+
+    except Exception as e:
+        if execute_task and not execute_task.done():
+            execute_task.cancel()
+            try:
+                await execute_task
+            except:
+                pass
+        if isinstance(e, TaskCancelledError):
+            raise
+        logger.error(f"Execution failed: {str(e)}")
+        raise
+
+@asynccontextmanager
+async def timeout_context(timeout: float, coro):
+    """Safe timeout wrapper with proper cleanup."""
+    try:
+        yield await asyncio.wait_for(coro, timeout=timeout)
+    except asyncio.TimeoutError as te:
+        # Enhance timeout error with context
+        raise asyncio.TimeoutError(f"{coro.__name__} timeout after {timeout}s") from te
+
+@api_router.post(
+    "/{thread_id}/stop",
+    response_model=Dict[str, Any],
+    summary="Stop a processing thread",
+    responses={
+        404: {"model": error_response, "description": "Thread not found"},
+        500: {"model": error_response, "description": "Execution failed"}
+    }
+)
+async def stop_thread(thread_id: str) -> Dict[str, Any]:
+    """Stop a processing thread."""
+    state = await _get_thread_or_404(thread_id)
+    if state.status != "running":
+        raise HTTPException(
+            status_code=status.HTTP_400_BAD_REQUEST,
+            detail={
+                "error": "invalid_status",
+                "message": f"Thread must be in 'running' state to stop. Current status: {state.status}",
+                "thread_id": thread_id
+            }
+        )
+
+    storage = AsyncRedisThreadStorage.get_worker_instance()
+    # Atomic stop operation
+    success = await storage.set_stop_flag(thread_id)
+    if not success:
+        raise HTTPException(
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            detail={
+                "error": "stop_failed",
+                "message": "Failed to set stop flag",
+                "thread_id": thread_id
+            }
+        )
+
+    return {
+        "status": "success",
+        "detail": {
+            "thread_id": thread_id,
+            "graph_name": state.graph_name,
+            "status": "stop"
+        }
+    }
+
+
+@api_router.get(
+    "/",
+    response_model=Dict[str, ThreadInfo],
+    summary="List active threads",
+    responses={
+        500: {"model": error_response, "description": "Internal server error"}
+    }
+)
+async def list_threads() -> Dict[str, ThreadInfo]:
+    """List all active processing threads."""
+    try:
+        storage = AsyncRedisThreadStorage.get_worker_instance()
+        threads = await storage.list_threads()
+        return {
+            tid: ThreadInfo(
+                thread_id=state.thread_id,
+                graph_name=state.graph_name,
+                status=state.status,
+            )
+            for tid, state in threads.items()
+        }
+    except Exception as e:
+        logger.error(f"Failed to list threads: {str(e)}", exc_info=True)
+        raise HTTPException(
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            detail={
+                "error": "list_failed",
+                "message": "Failed to retrieve thread list"
+            }
+        )
+
+
+@api_router.delete(
+    "/{thread_id}",
+    response_model=Dict[str, Any],
+    summary="Terminate processing thread",
+    responses={
+        404: {"model": error_response, "description": "Thread not found"},
+        500: {"model": error_response, "description": "Deletion failed"}
+    }
+)
+async def delete_thread(thread_id: str) -> Dict[str, Any]:
+    """Delete a processing thread."""
+    state = await _get_thread_or_404(thread_id)
+    storage = AsyncRedisThreadStorage.get_worker_instance()
+    success = await storage.delete_thread(thread_id)
+    if not success:
+        raise HTTPException(
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            detail={
+                "error": "deletion_failed",
+                "message": "Thread deletion failed",
+                "thread_id": thread_id
+            }
+        )
+
+    logger.info(f"Thread terminated: {thread_id}")
+    return {
+        "status": "success",
+        "detail": {
+            "thread_id": thread_id,
+            "graph_name": state.graph_name,
+            "status": "terminated"
+        }
+    }