forked from Skyvern-AI/skyvern
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Remove the base Agent; Separate skyvern agent and fastapi app (Skyver…
…n-AI#213) Co-authored-by: Shuchang Zheng <wintonzheng0325@gmail.com>
- Loading branch information
1 parent
22e0332
commit e02d588
Showing
5 changed files
with
94 additions
and
95 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,92 @@ | ||
import uuid | ||
from datetime import datetime | ||
from typing import Awaitable, Callable | ||
|
||
import structlog | ||
from fastapi import APIRouter, FastAPI, Response | ||
from fastapi.middleware.cors import CORSMiddleware | ||
from fastapi.responses import JSONResponse | ||
from starlette.requests import HTTPConnection, Request | ||
from starlette_context.middleware import RawContextMiddleware | ||
from starlette_context.plugins.base import Plugin | ||
|
||
from skyvern.forge.sdk.core import skyvern_context | ||
from skyvern.forge.sdk.core.skyvern_context import SkyvernContext | ||
from skyvern.forge.sdk.routes.agent_protocol import base_router | ||
from skyvern.scheduler import SCHEDULER | ||
|
||
LOG = structlog.get_logger() | ||
|
||
|
||
class ExecutionDatePlugin(Plugin): | ||
key = "execution_date" | ||
|
||
async def process_request(self, request: Request | HTTPConnection) -> datetime: | ||
return datetime.now() | ||
|
||
|
||
def get_agent_app(router: APIRouter = base_router) -> FastAPI: | ||
""" | ||
Start the agent server. | ||
""" | ||
|
||
app = FastAPI() | ||
|
||
# Add CORS middleware | ||
origins = [ | ||
"http://localhost:5000", | ||
"http://127.0.0.1:5000", | ||
"http://localhost:8000", | ||
"http://127.0.0.1:8000", | ||
"http://localhost:8080", | ||
"http://127.0.0.1:8080", | ||
# Add any other origins you want to whitelist | ||
] | ||
|
||
app.add_middleware( | ||
CORSMiddleware, | ||
allow_origins=origins, | ||
allow_credentials=True, | ||
allow_methods=["*"], | ||
allow_headers=["*"], | ||
) | ||
|
||
app.include_router(router, prefix="/api/v1") | ||
|
||
app.add_middleware( | ||
RawContextMiddleware, | ||
plugins=( | ||
# TODO (suchintan): We should set these up | ||
ExecutionDatePlugin(), | ||
# RequestIdPlugin(), | ||
# UserAgentPlugin(), | ||
), | ||
) | ||
|
||
# Register the scheduler on startup so that we can schedule jobs dynamically | ||
@app.on_event("startup") | ||
def start_scheduler() -> None: | ||
LOG.info("Starting the skyvern scheduler.") | ||
SCHEDULER.start() | ||
|
||
LOG.info("Server startup complete. Skyvern is now online") | ||
|
||
@app.exception_handler(Exception) | ||
async def unexpected_exception(request: Request, exc: Exception) -> JSONResponse: | ||
LOG.exception("Unexpected error in agent server.", exc_info=exc) | ||
return JSONResponse(status_code=500, content={"error": f"Unexpected error: {exc}"}) | ||
|
||
@app.middleware("http") | ||
async def request_middleware(request: Request, call_next: Callable[[Request], Awaitable[Response]]) -> Response: | ||
request_id = str(uuid.uuid4()) | ||
skyvern_context.set(SkyvernContext(request_id=request_id)) | ||
|
||
try: | ||
return await call_next(request) | ||
finally: | ||
skyvern_context.reset() | ||
|
||
return app | ||
|
||
|
||
app = get_agent_app() |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file was deleted.
Oops, something went wrong.