Function execute_command

run_tiny_agent_server.py:61–121

This is the main endpoint that calls the TinyAgent to generate a response to the given query.

(request: TinyAgentRequest)

@app.post("/generate")
async def execute_command(request: TinyAgentRequest) -> StreamingResponse:
    """
    This is the main endpoint that calls the TinyAgent to generate a response to the given query.
    """
    log(f"\n\n====\nReceived request: {request.query}")

    # First, ensure the queue is empty
    empty_queue(streaming_queue)

    query = request.query

    if not query or len(query) <= 0:
        raise HTTPException(
            status_code=HTTPStatus.BAD_REQUEST, detail="No query provided"
        )

    try:
        tiny_agent_config = get_tiny_agent_config(config_path=CONFIG_PATH)
        tiny_agent = TinyAgent(tiny_agent_config)
    except Exception as e:
        raise HTTPException(
            status_code=HTTPStatus.INTERNAL_SERVER_ERROR,
            detail=f"Error: {e}",
        )

    async def generate():
        try:
            response_task = asyncio.create_task(tiny_agent.arun(query))

            while True:
                # Await a small timeout to periodically check if the task is done
                try:
                    token = await asyncio.wait_for(streaming_queue.get(), timeout=1.0)
                    if token is None:
                        break
                    if token.startswith(LLM_ERROR_TOKEN):
                        raise Exception(token[len(LLM_ERROR_TOKEN) :])
                    yield token
                except asyncio.TimeoutError:
                    pass  # No new token, check task status

                # Check if the task is done to handle any potential exception
                if response_task.done():
                    break

            # Task created with asyncio.create_task() do not propagate exceptions
            # to the calling context. Instead, the exception remains encapsulated within
            # the task object itself until the task is awaited or its result is explicitly retrieved.
            # Hence, we check here if the task has an exception set by awaiting it, which will
            # raise the exception if it exists. If it doesn't, we just yield the result.
            await response_task
            response = response_task.result()
            yield f"\n\n{response}"
        except Exception as e:
            # You cannot raise HTTPExceptions in an async generator, it doesn't
            # get caught by the FastAPI exception handling middleware. Hence,
            # we are manually catching the exceptions and yielding/logging them.
            yield f"Error: {e}"