feat: Consolidate on a single LangGraph orchestration

.github/sync-repo-settings.yaml

@@ -35,9 +35,7 @@ branchProtectionRules:
       - "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

@@ -34,7 +34,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        dir: [retrieval_service, llm_demo]
+        dir: [retrieval_service]
       fail-fast: false
     permissions:
       contents: read

.github/workflows/lint_fallback.yml

@@ -25,7 +25,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        dir: [retrieval_service, llm_demo]
+        dir: [retrieval_service]
     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")
@@ -237,16 +237,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 +253,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,12 @@
 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}
 
   - id: "Test Frontend"
     name: "gcr.io/cloud-builders/gcloud:latest"
@@ -32,7 +30,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 +74,3 @@ substitutions:
   _GCR_HOSTNAME: ${_REGION}-docker.pkg.dev
   _SERVICE: demo-service-${BUILD_ID}
   _REGION: us-central1
-  _ORCHESTRATION_TYPE: langgraph

llm_demo/orchestrator/langgraph/__init__.py → orchestrator/__init__.py

@@ -12,6 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .langgraph_orchestrator import LangGraphOrchestrator
+from .orchestrator import Orchestrator
 
-__ALL__ = ["LangGraphOrchestrator"]
+__ALL__ = ["Orchestrator"]

llm_demo/orchestrator/langgraph/langgraph_orchestrator.py → orchestrator/orchestrator.py

@@ -42,6 +42,8 @@
 
 
 class LangGraphOrchestrator(BaseOrchestrator):
+    MODEL = "gemini-2.0-flash-001"
+
     _user_sessions: Dict[str, str]
     # aiohttp context
     connector = None
@@ -52,10 +54,6 @@ def __init__(self):
         self._langgraph_app = None
         self._checkpointer = None
 
-    @classproperty
-    def kind(cls):
-        return "langgraph"
-
     def user_session_exist(self, uuid: str) -> bool:
         return uuid in self._user_sessions
 
@@ -249,7 +247,9 @@ def get_config(self, uuid: str):
     async def user_session_signout(self, uuid: str):
         checkpoint = empty_checkpoint()
         config = self.get_config(uuid)
-        self._checkpointer.put(config=config, checkpoint=checkpoint, metadata={})
+        self._checkpointer.put(
+            config=config, checkpoint=checkpoint, metadata={}, new_versions={}
+        )
         del self._user_sessions[uuid]
 
     async def close_clients(self):

llm_demo/requirements.txt → requirements.txt

@@ -1,19 +1,19 @@
 fastapi==0.115.0
-google-auth==2.35.0
-google-cloud-aiplatform[evaluation]==1.72.0
+google-auth==2.40.3
+google-cloud-aiplatform[evaluation]==1.97.0
 itsdangerous==2.2.0
 jinja2==3.1.5
-langchain-community==0.3.2
-langchain==0.3.7
-langchain-google-vertexai==2.0.7
+langchain-community==0.3.25
+langchain==0.3.25
+langchain-core==0.3.65
+langchain-google-vertexai==2.0.25
 markdown==3.7
 types-Markdown==3.7.0.20240822
 uvicorn[standard]==0.31.0
 python-multipart==0.0.18
 pytz==2025.1
 types-pytz==2025.1.0.20250204
-langgraph==0.2.48
-httpx==0.27.2
+langgraph==0.4.8
 pandas-stubs==2.2.2.240807
 pandas==2.2.3
 pydantic==2.9.0

llm_demo/run_app.py → run_app.py

@@ -24,12 +24,9 @@
 async def 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")
     server = uvicorn.Server(uvicorn.Config(app, host=HOST, port=PORT, log_level="info"))

llm_demo/run_evaluation.py → run_evaluation.py

@@ -26,7 +26,7 @@
     goldens,
     run_llm_for_eval,
 )
-from orchestrator import createOrchestrator
+from orchestrator import Orchestrator
 
 
 def export_metrics_table_csv(retrieval: pd.DataFrame, response: pd.DataFrame):
@@ -48,7 +48,6 @@ async def main():
     USER_ID_TOKEN = os.getenv("USER_ID_TOKEN", default=None)
 
     CLIENT_ID = os.getenv("CLIENT_ID", default="")
-    ORCHESTRATION_TYPE = os.getenv("ORCHESTRATION_TYPE", default="langchain-tools")
     EXPORT_CSV = bool(os.getenv("EXPORT_CSV", default=False))
     RETRIEVAL_EXPERIMENT_NAME = os.getenv(
         "RETRIEVAL_EXPERIMENT_NAME", default="retrieval-phase-eval"
@@ -58,7 +57,7 @@ async def main():
     )
 
     # Prepare orchestrator and session
-    orc = createOrchestrator(ORCHESTRATION_TYPE)
+    orc = Orchestrator()
     session_id = str(uuid.uuid4())
     session = {"uuid": session_id}
     await orc.user_session_create(session)

llm_demo/static/trace.js → static/trace.js

@@ -22,8 +22,8 @@ export function create_trace(toolcalls) {
         let toolcall = toolcalls[i];
         trace += trace_section_title(toolcall.tool_call_id);
 
-        trace += trace_header("SQL Executed:");
-        trace += trace_sql(toolcall.sql);
+            trace += trace_header("SQL Executed:");
+            trace += trace_sql(toolcall.sql);
         trace += trace_header("Results:");
         trace += trace_results(toolcall.results);