feat: Integrate Toolbox SDK for LangGraph orchestration.

.github/sync-repo-settings.yaml

@@ -29,15 +29,8 @@ branchProtectionRules:
     requiresStrictStatusChecks: true
     requiredStatusCheckContexts:
       - "cla/google"
-      - "retrieval_service"
       - "llm_demo"
-      - "retrieval-service-app-pr (retrieval-app-testing)"
-      - "retrieval-service-postgres-pr (retrieval-app-testing)"
-      - "retrieval-service-alloydb-pr (retrieval-app-testing)"
-      - "retrieval-service-cloudsql-pg-pr (retrieval-app-testing)"
-      - "llm-demo-langchain-tools-pr (retrieval-app-testing)"
       - "llm-demo-langgraph-pr (retrieval-app-testing)"
-      - "llm-demo-vertexai-fc-pr (retrieval-app-testing)"
 # Set team access
 permissionRules:
   - team: senseai-eco

.github/workflows/lint.yml

@@ -30,12 +30,8 @@ permissions: read-all
 
 jobs:
   integration:
-    name: ${{ matrix.dir }}
+    name: Cymbal Air
     runs-on: ubuntu-latest
-    strategy:
-      matrix:
-        dir: [retrieval_service, llm_demo]
-      fail-fast: false
     permissions:
       contents: read
       id-token: write
@@ -51,11 +47,9 @@ jobs:
           python-version: "3.11"
 
       - name: Install requirements
-        working-directory: ${{ matrix.dir }}
         run: pip install -r requirements.txt -r requirements-test.txt
 
       - name: Run lints
-        working-directory: ${{ matrix.dir }}
         run: |
           black --check .
           isort --check .

.github/workflows/lint_fallback.yml

@@ -21,11 +21,8 @@ on:
 
 jobs:
   integration:
-    name: ${{ matrix.dir }}
+    name: Cymbal Air
     runs-on: ubuntu-latest
-    strategy:
-      matrix:
-        dir: [retrieval_service, llm_demo]
     permissions:
       contents: none
 

llm_demo/app.py → app.py

@@ -27,7 +27,7 @@
 from markdown import markdown
 from starlette.middleware.sessions import SessionMiddleware
 
-from orchestrator import createOrchestrator
+from orchestrator import Orchestrator
 
 routes = APIRouter()
 templates = Jinja2Templates(directory="templates")
@@ -195,13 +195,19 @@ async def decline_flight(request: Request):
     # This is helpful in case of reloads so there doesn't seem to be a break in communication.
     orchestrator = request.app.state.orchestrator
     response = await orchestrator.user_session_decline_ticket(request.session["uuid"])
+    response = (
+        response["output"]
+        if response
+        else "Booking declined. What else can I help you with?"
+    )
     request.session["history"].append(
         {"type": "ai", "data": {"content": "Please confirm if you would like to book."}}
     )
     request.session["history"].append(
         {"type": "human", "data": {"content": "I changed my mind."}}
     )
-    return None
+    request.session["history"].append({"type": "ai", "data": {"content": response}})
+    return response
 
 
 @routes.post("/reset")
@@ -237,16 +243,13 @@ def clear_user_info(session: dict[str, Any]):
 
 
 def init_app(
-    orchestration_type: Optional[str],
     client_id: Optional[str],
     middleware_secret: Optional[str],
 ) -> FastAPI:
     # FastAPI setup
-    if orchestration_type is None:
-        raise HTTPException(status_code=500, detail="Orchestrator not found")
     app = FastAPI(lifespan=lifespan)
     app.state.client_id = client_id
-    app.state.orchestrator = createOrchestrator(orchestration_type)
+    app.state.orchestrator = Orchestrator()
     app.include_router(routes)
     app.mount("/static", StaticFiles(directory="static"), name="static")
     app.add_middleware(SessionMiddleware, secret_key=middleware_secret)
@@ -256,12 +259,9 @@ def init_app(
 if __name__ == "__main__":
     PORT = int(os.getenv("PORT", default=8081))
     HOST = os.getenv("HOST", default="0.0.0.0")
-    ORCHESTRATION_TYPE = os.getenv("ORCHESTRATION_TYPE", default="langchain-tools")
     CLIENT_ID = os.getenv("CLIENT_ID")
     MIDDLEWARE_SECRET = os.getenv("MIDDLEWARE_SECRET", default="this is a secret")
-    app = init_app(
-        ORCHESTRATION_TYPE, client_id=CLIENT_ID, middleware_secret=MIDDLEWARE_SECRET
-    )
+    app = init_app(client_id=CLIENT_ID, middleware_secret=MIDDLEWARE_SECRET)
     if app is None:
         raise TypeError("app not instantiated")
     uvicorn.run(app, host=HOST, port=PORT)

llm_demo/evaluation.cloudbuild.yaml → evaluation.cloudbuild.yaml

@@ -14,14 +14,11 @@
 steps:
   - id: Install dependencies
     name: python:3.11
-    dir: llm_demo 
     script: pip install -r requirements.txt -r requirements-test.txt --user
 
   - id: "Run evaluation service"
     name: python:3.11
-    dir: llm_demo
     env: # Set env var expected by tests
-      - "ORCHESTRATION_TYPE=${_ORCHESTRATION_TYPE}"
       - "RETRIEVAL_EXPERIMENT_NAME=${_RETRIEVAL_EXPERIMENT_NAME}"
       - "RESPONSE_EXPERIMENT_NAME=${_RESPONSE_EXPERIMENT_NAME}"
     secretEnv:
@@ -37,7 +34,6 @@ options:
   dynamic_substitutions: true
 
 substitutions:
-  _ORCHESTRATION_TYPE: "langchain-tools"
   _RETRIEVAL_EXPERIMENT_NAME: "retrieval-phase-eval-${_PR_NUMBER}"
   _RESPONSE_EXPERIMENT_NAME:  "response-phase-eval-${_PR_NUMBER}"
 

llm_demo/evaluation/evaluation.py → evaluation/evaluation.py

@@ -17,22 +17,20 @@
 from typing import Dict, List
 
 import pandas as pd
-from pydantic import BaseModel, Field
 from vertexai.evaluation import EvalTask
 from vertexai.evaluation import _base as evaluation_base
 
-from orchestrator import BaseOrchestrator
+from orchestrator import Orchestrator
 
 from .eval_golden import EvalData, ToolCall
 from .metrics import response_phase_metrics, retrieval_phase_metrics
 
 
 async def run_llm_for_eval(
-    eval_list: List[EvalData], orc: BaseOrchestrator, session: Dict, session_id: str
+    eval_list: List[EvalData], orc: Orchestrator, session: Dict, session_id: str
 ) -> List[EvalData]:
     """
     Generate llm_tool_calls and llm_output for golden dataset query.
-    This function is only compatible with the langchain-tools orchestration.
     """
     agent = orc.get_user_session(session_id)
     for eval_data in eval_list:

llm_demo/langgraph.int.tests.cloudbuild.yaml → integration.cloudbuild.yaml

@@ -14,14 +14,13 @@
 steps:
   - id: "Deploy to Cloud Run"
     name: "gcr.io/cloud-builders/gcloud:latest"
-    dir: llm_demo
     script: |
         #!/usr/bin/env bash
         gcloud run deploy ${_SERVICE} \
           --source . \
           --region ${_REGION} \
           --no-allow-unauthenticated \
-          --update-env-vars ORCHESTRATION_TYPE=${_ORCHESTRATION_TYPE}
+          --update-env-vars TOOLBOX_URL=${_TOOLBOX_URL}
 
   - id: "Test Frontend"
     name: "gcr.io/cloud-builders/gcloud:latest"
@@ -32,7 +31,6 @@ steps:
       - |
         export URL=$(gcloud run services describe ${_SERVICE} --region ${_REGION} --format 'value(status.url)')
         export ID_TOKEN=$(gcloud auth print-identity-token --audiences $$URL)
-        export ORCHESTRATION_TYPE=${_ORCHESTRATION_TYPE}
 
         # Test `/` route
         curl -c cookies.txt -si --fail --show-error -H "Authorization: Bearer $$ID_TOKEN" $$URL
@@ -77,4 +75,4 @@ substitutions:
   _GCR_HOSTNAME: ${_REGION}-docker.pkg.dev
   _SERVICE: demo-service-${BUILD_ID}
   _REGION: us-central1
-  _ORCHESTRATION_TYPE: langgraph
+  _TOOLBOX_URL: https://toolbox-107716898620.us-central1.run.app