Adds OpenAI endpoints

Author: vmesel

src/dialog/routers/__init__.py

@@ -1 +1,2 @@
-from .dialog import *
+from .dialog import *
+from .openai import open_ai_api_router

src/main.py

@@ -7,7 +7,7 @@
 from fastapi import FastAPI
 from fastapi.staticfiles import StaticFiles
 from fastapi.middleware.cors import CORSMiddleware
-from dialog.routers import api_router
+from dialog.routers import api_router, open_ai_api_router
 
 logging.basicConfig(
     level=Settings().LOGGING_LEVEL,
@@ -31,9 +31,18 @@ def get_application() -> FastAPI:
         allow_headers=Settings().CORS_ALLOW_HEADERS,
     )
 
+    app.add_middleware(
+        CustomHeaderMiddleware
+    )
+
     app.include_router(
         api_router, prefix="",
     )
+
+    app.include_router(
+        open_ai_api_router, prefix="/openai"
+    )
+
     app.mount("/static", StaticFiles(directory=Settings().STATIC_FILE_LOCATION), name="static")
     plugins = entry_points(group="dialog")
     for plugin in plugins:

src/tests/conftest.py

@@ -45,4 +45,10 @@ def chat_session(dbsession):
 def llm_mock(mocker):
     llm_mock = mocker.patch('dialog.routers.dialog.process_user_message')
     llm_mock.process.return_value = {"text": "Hello"}
+    return llm_mock
+
+@pytest.fixture
+def llm_mock_openai_router(mocker):
+    llm_mock = mocker.patch('dialog.routers.openai.process_user_message')
+    llm_mock.process.return_value = {"text": "Hello"}
     return llm_mock

src/tests/test_views.py

@@ -1,3 +1,4 @@
+import os
 import pytest
 
 from dialog_lib.db.models import ChatMessages, Chat
@@ -47,4 +48,21 @@ def test_invalid_database_connection(client, mocker):
     with pytest.raises(Exception):
         response = client.get("/health")
         assert response.status_code == 500
-        assert response.json() == {"message": "Failed to execute simple SQL"}
+        assert response.json() == {"message": "Failed to execute simple SQL"}
+
+# test openai router
+
+def test_customized_openai_models_response(client):
+    response = client.get("/openai/models")
+    assert response.status_code == 200
+    for i in ["id", "object", "created", "owned_by"]:
+        assert i in response.json()[0]
+
+def test_customized_openai_chat_completion_response(client, llm_mock_openai_router):
+    os.environ["LLM_CLASS"] = "dialog.llm.agents.default.DialogLLM"
+    response = client.post("/openai/chat/completions", json={"message": "Hello"})
+    assert response.status_code == 200
+    for i in ["choices", "created", "id", "model", "object", "usage"]:
+        assert i in response.json()
+    assert llm_mock_openai_router.called
+    assert response.json()["choices"][0]["message"]["role"] == "assistant"