LangChain OpenAI API proxy on Google Cloud Functions

main.py

from typing import Any, AsyncIterator, Dict, Iterator, List, Optional

from langchain_core.language_models import BaseChatModel, SimpleChatModel
from langchain_core.load import dumpd, dumps, load, loads
from langchain_core.messages import AIMessage, AIMessageChunk, BaseMessage, HumanMessage
from langchain_core.outputs import ChatGeneration, ChatGenerationChunk, ChatResult
from langchain_openai import ChatOpenAI, OpenAIEmbeddings
import requests

import functions_framework

# KISS principle
VALID_TOKENS = [
    "teaching.team-f97301440329471",

    # Students
    # - snip - #
]

def make_chat_model():
    return ChatOpenAI(
        model="gpt-4o-mini",
        temperature=1.3,
    )

def make_embedding_model():
    return OpenAIEmbeddings(
        model='text-embedding-3-small',
    )


@functions_framework.http
def main(request):
    """HTTP Cloud Function.
    Args:
        request (flask.Request): The request object.
        <https://flask.palletsprojects.com/en/1.1.x/api/#incoming-request-data>
    Returns:
        The response text, or any set of values that can be turned into a
        Response object using `make_response`
        <https://flask.palletsprojects.com/en/1.1.x/api/#flask.make_response>.
    """
    request_args = request.args
    try:
      token = request_args['token']
    except KeyError:
      return ({'error': 'Credentials required'}, 401)
    if token not in VALID_TOKENS:
      return ({'error': 'Invalid credentials'}, 401)

    body = request.get_json(silent=True)
    payload = body['payload']
    match body['task']:
        case 'chat':
            messages = [loads(m) for m in payload['messages']]
            stop = payload['stop']
            kwargs = payload['kwargs'] if 'kwargs' in payload else {}

            return dumps(make_chat_model().invoke(messages, stop=stop, **kwargs))

        case 'embed_documents':
            return dumps(make_embedding_model().embed_documents(payload))

        case 'embed_query':
            return dumps(make_embedding_model().embed_query(payload))

        case _:
            return ({'error': 'Invalid task.'}, 400)

requirements.txt

functions-framework==3.*
langchain-core==0.2.38
langchain-openai==0.1.23