Merge branch 'f/migration-to-sdk' into 'development'

Migration to Dial SDK

See merge request Deltix/openai-apps/bedrock-to-openai-adapter!13

.flake8

@@ -8,4 +8,4 @@ exclude =
     __pycache__
 per-file-ignores =
     # imported but unused
-    test/locust.py: F401
+    tests/locust/app.py: F401

.gitlab-ci.yml

@@ -36,6 +36,7 @@ variables:
 stages:
   - tag
   - lint
+  - test
   - publish # publish docker images
   - deploy
   - promote # copy docker image to public repo
@@ -46,7 +47,7 @@ lint:
   stage: lint
   before_script:
     - apt-get update && apt-get install -y make
-    - make all
+    - make install
     - source .venv/bin/activate
   script:
     - make lint
@@ -56,6 +57,21 @@ lint:
   tags:
     - kubernetes
 
+test:
+  image: python:3.11-slim
+  stage: test
+  before_script:
+    - apt-get update && apt-get install -y make
+    - make install
+    - source .venv/bin/activate
+  script:
+    - make test
+  rules:
+    - if: $CI_MERGE_REQUEST_TARGET_BRANCH_NAME == "development" && $review_DEPLOY_ENV == "true"
+    - if: $CI_COMMIT_BRANCH  == "development"
+  tags:
+    - kubernetes
+
 deploy_development:
   image: "registry.deltixhub.com/deltix.docker/devops/kubernetes-tools:0.17.1"
   stage: deploy

.vscode/launch.json

@@ -16,20 +16,5 @@
         "PYTHONPATH": "${workspaceFolder}"
       }
     },
-    {
-      "name": "Docker: Python - Fastapi",
-      "type": "docker",
-      "request": "launch",
-      "preLaunchTask": "docker-run: debug",
-      "python": {
-        "pathMappings": [
-          {
-            "localRoot": "${workspaceFolder}",
-            "remoteRoot": "/app"
-          }
-        ],
-        "projectType": "fastapi"
-      }
-    }
   ]
 }

.vscode/settings.json

@@ -1,20 +1,4 @@
 {
-  "cSpell.enabled": true,
-  "files.exclude": {
-    "**/.git": true,
-    "**/.svn": true,
-    "**/.hg": true,
-    "**/CVS": true,
-    "**/.DS_Store": true,
-    "**/Thumbs.db": true,
-    "**/__pycache__": true,
-    "**/.pytest_cache": true,
-    "**/.gradle": true,
-    "**/bin": true,
-    "**/build": true,
-    ".idea": true,
-    ".venv": true
-  },
   "[python]": {
     "editor.defaultFormatter": "ms-python.black-formatter",
     "editor.formatOnSave": true,
@@ -25,5 +9,6 @@
   },
   "python.testing.pytestArgs": ["."],
   "python.testing.unittestEnabled": false,
-  "python.testing.pytestEnabled": true
-}
+  "python.testing.pytestEnabled": true,
+  "python.analysis.typeCheckingMode": "basic"
+}

Dockerfile

@@ -10,6 +10,7 @@ ENV PYTHONUNBUFFERED=1
 # Install pip requirements
 COPY requirements.txt .
 COPY bedrock-python-sdk /bedrock-python-sdk
+COPY dial-python-sdk /dial-python-sdk
 RUN python -m pip install -r requirements.txt
 
 WORKDIR /app

Makefile

@@ -1,6 +1,6 @@
 ENV ?= DEV
 PORT ?= 5001
-IMAGE_NAME ?= bedrock-adapter
+IMAGE_NAME ?= ai-dial-adapter-bedrock
 
 ifeq ($(ENV),DEV)
 	REQ = requirements-dev.txt
@@ -11,18 +11,17 @@ endif
 VENV=.venv
 BUILD=$(VENV)/bin/activate
 
-all: $(BUILD)
+install: $(BUILD)
 
-$(BUILD): $(REQ)
+$(BUILD): requirements.txt $(REQ)
 	python3 -m venv $(VENV)
 	$(VENV)/bin/pip install -r $(REQ)
 	@echo "\033[31mActivate venv by running:\n> source $(BUILD)\033[0m"
 
-.PHONY: all server-run client-run clean lint format test docker-build docker-run
+.PHONY: install server-run client-run clean lint format test integration_test docker-build docker-run
 
 server-run: $(BUILD)
-	@source ./load_env.sh; load_env; \
-	python -m debug_app --port=$(PORT)
+	python -m server.main --port=$(PORT)
 
 client-run: $(BUILD)
 	@source ./load_env.sh; load_env; \
@@ -42,10 +41,15 @@ format: $(BUILD)
 	isort . $(ARGS)
 	black . $(ARGS)
 
+TEST_FILES ?= ./tests/unit_tests
+
 # Add options "-s --log-cli-level=NOTSET" to pytest to see all logs
 test: $(BUILD)
+	python -m pytest $(TEST_FILES) -v --durations=0 -rA
+
+integration_test: $(BUILD)
 	@source ./load_env.sh; load_env; \
-	python -m pytest . -v --durations=0 -rA
+	python -m pytest ./tests/integration_tests -v --durations=0 -rA $(ARGS)
 
 docker-build: Dockerfile
 	docker build --platform linux/amd64 -t $(IMAGE_NAME) .

README.md

@@ -5,7 +5,7 @@ The server provides `/chat/completions` and `/completions` endpoint compatible w
 ## Installation
 
 ```sh
-make all
+make install
 source .venv/bin/activate
 ```
 

app.py

@@ -1,43 +1,63 @@
 import logging.config
 
-from fastapi import Body, FastAPI, Path, Query, Request, Response
-from fastapi.middleware.cors import CORSMiddleware
-from fastapi.responses import JSONResponse
+from aidial_sdk import (
+    ChatCompletion,
+    ChatCompletionRequest,
+    ChatCompletionResponse,
+    DIALApp,
+)
+from fastapi import Query, Response
 
 from llm.bedrock_adapter import BedrockAdapter, BedrockModels
 from llm.bedrock_models import BedrockDeployment
 from llm.chat_emulation.types import ChatEmulationType
-from server.exceptions import OpenAIException, open_ai_exception_decorator
-from universal_api.request import ChatCompletionRequest
-from universal_api.response import ModelObject, ModelsResponse, make_response
+from server.exceptions import dial_exception_decorator
+from universal_api.request import ModelParameters
+from universal_api.response import ModelObject, ModelsResponse
+from universal_api.token_usage import TokenUsage
 from utils.env import get_env
 from utils.log_config import LogConfig
-from utils.log_config import app_logger as log
 
 logging.config.dictConfig(LogConfig().dict())
 
-app = FastAPI(
-    description="Bedrock adapter for OpenAI Chat API",
-    version="0.0.1",
-)
+default_region = get_env("DEFAULT_REGION")
+chat_emulation_type = ChatEmulationType.META_CHAT
 
-# CORS
 
-origins = ["*"]
+class BedrockChatCompletion(ChatCompletion):
+    @dial_exception_decorator
+    async def chat_completion(
+        self, request: ChatCompletionRequest, response: ChatCompletionResponse
+    ):
+        model = await BedrockAdapter.create(
+            region=default_region,
+            model_id=request.deployment_id,
+            model_params=ModelParameters.create(request),
+        )
 
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=origins,
-    allow_credentials=True,
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
+        model_response = await model.achat(
+            chat_emulation_type, request.messages
+        )
 
+        usage = TokenUsage()
 
-# Endpoints
+        for _ in range(request.n or 1):
+            with response.create_choice() as choice:
+                choice.append_content(model_response.content)
 
+                for data in model_response.data:
+                    choice.add_attachment(
+                        title=data.name,
+                        data=data.content,
+                        type=data.mime_type,
+                    )
 
-default_region = get_env("DEFAULT_REGION")
+                usage += model_response.usage
+
+        response.set_usage(usage.prompt_tokens, usage.completion_tokens)
+
+
+app = DIALApp()
 
 
 @app.get("/healthcheck")
@@ -46,45 +66,14 @@ def healthcheck():
 
 
 @app.get("/openai/models")
-@open_ai_exception_decorator
+@dial_exception_decorator
 async def models(
     region: str = Query(default=default_region, description="AWS region")
-) -> ModelsResponse:
-    return get_models(region)
-
-
-def get_models(region: str) -> ModelsResponse:
+):
     bedrock_models = BedrockModels(region).models()
     models = [ModelObject(id=model["modelId"]) for model in bedrock_models]
     return ModelsResponse(data=models)
 
 
-@app.post("/openai/deployments/{deployment}/chat/completions")
-@open_ai_exception_decorator
-async def chat_completions(
-    deployment: BedrockDeployment = Path(...),
-    chat_emulation_type: ChatEmulationType = Query(
-        default=ChatEmulationType.META_CHAT,
-        description="The chat emulation type for models which only support completion mode",
-    ),
-    region: str = Query(default=default_region, description="AWS region"),
-    query: ChatCompletionRequest = Body(...),
-):
-    model_id = deployment.get_model_id()
-    model = await BedrockAdapter.create(
-        region=region, model_id=model_id, model_params=query
-    )
-    response = await model.achat(chat_emulation_type, query.messages)
-    log.debug(f"response:\n{response}")
-
-    return make_response(
-        bool(query.stream), deployment, "chat.completion", response
-    )
-
-
-@app.exception_handler(OpenAIException)
-async def exception_handler(request: Request, exc: OpenAIException):
-    log.exception(f"Exception: {str(exc)}")
-    return JSONResponse(
-        status_code=exc.status_code, content={"error": exc.error}
-    )
+for deployment in BedrockDeployment:
+    app.add_chat_completion(deployment.get_model_id(), BedrockChatCompletion())

client/client_adapter.py

@@ -49,12 +49,14 @@ def create_model(
         callbacks=callbacks,
         openai_api_base=base_url,
         openai_api_version=DEFAULT_API_VERSION,
-        openai_api_key="dummy_key",
+        openai_api_key="dummy_openai_api_key",
+        model_kwargs={"deployment_id": model_id, "api_key": "dummy_api_key"},
         verbose=True,
         streaming=streaming,
-        temperature=0,
+        temperature=0.0,
         request_timeout=10,
         client=None,
+        max_retries=0,
     )
 
 

client/client_bedrock.py

@@ -1,8 +1,10 @@
 import asyncio
 from typing import List
 
+from aidial_sdk.chat_completion.request import Message
+
 from llm.bedrock_adapter import BedrockAdapter
-from universal_api.request import ChatCompletionParameters, Message
+from universal_api.request import ModelParameters
 from utils.cli import choose_deployment
 from utils.env import get_env
 from utils.init import init
@@ -16,7 +18,7 @@ async def main():
 
     model = await BedrockAdapter.create(
         model_id=deployment.get_model_id(),
-        model_params=ChatCompletionParameters(),
+        model_params=ModelParameters(),
         region=get_env("DEFAULT_REGION"),
     )