Fix issues

ls1intum · Jul 16, 2023 · f3dd199 · f3dd199
1 parent 3096b14
commit f3dd199
Show file tree

Hide file tree

Showing 5 changed files with 24 additions and 15 deletions.
diff --git a/app/dependencies.py b/app/dependencies.py
@@ -20,7 +20,7 @@ def _get_api_key(request: StarletteRequest) -> str:
 
 class TokenValidator:
     async def __call__(
-        self, request: StarletteRequest, api_key: str = Depends(_get_api_key)
+        self, api_key: str = Depends(_get_api_key)
     ) -> APIKeyConfig:
         for key in settings.pyris.api_keys:
             if key.token == api_key:

diff --git a/app/routes/messages.py b/app/routes/messages.py
@@ -12,7 +12,7 @@
 from app.dependencies import TokenPermissionsValidator
 from app.models.dtos import SendMessageRequest, SendMessageResponse
 from app.services.guidance_wrapper import GuidanceWrapper
-from config import settings
+from app.config import settings
 
 router = APIRouter(tags=["messages"])
 

diff --git a/app/routes/models.py b/app/routes/models.py
@@ -1,13 +1,13 @@
 from fastapi import APIRouter, Depends
 
 from app.dependencies import TokenValidator
-from config import settings, APIKeyConfig
-from models.dtos import LLMModelResponse
+from app.config import settings, APIKeyConfig
+from app.models.dtos import LLMModelResponse
 
 router = APIRouter(tags=["models"])
 
 
-@router.get("/api/v1/models", dependencies=[Depends(TokenValidator())])
+@router.get("/api/v1/models")
 def send_message(
     api_key_config: APIKeyConfig = Depends(TokenValidator()),
 ) -> list[LLMModelResponse]:

diff --git a/tests/routes/messages_test.py b/tests/routes/messages_test.py
@@ -1,6 +1,12 @@
 from freezegun import freeze_time
 from app.models.dtos import Content, ContentType
 from app.services.guidance_wrapper import GuidanceWrapper
+import app.config as config
+
+llm_model_config = config.LLMModelConfig(name="test", description="test", llm_credentials={})
+config.settings.pyris.llms = {"GPT35_TURBO": llm_model_config}
+api_key_config = config.APIKeyConfig(token="secret", comment="test", llm_access=["GPT35_TURBO"])
+config.settings.pyris.api_keys = [api_key_config]
 
 
 @freeze_time("2023-06-16 03:21:34 +02:00")
@@ -39,8 +45,13 @@ def test_send_message(test_client, headers, mocker):
     }
 
 
-def test_send_message_missing_params(test_client, headers):
+def test_send_message_missing_model(test_client, headers):
     response = test_client.post("/api/v1/messages", headers=headers, json={})
+    assert response.status_code == 404
+
+
+def test_send_message_missing_params(test_client, headers):
+    response = test_client.post("/api/v1/messages", headers=headers, json={"preferredModel": "GPT35_TURBO"})
     assert response.status_code == 422
     assert response.json() == {
         "detail": [
@@ -49,11 +60,6 @@ def test_send_message_missing_params(test_client, headers):
                 "msg": "field required",
                 "type": "value_error.missing",
             },
-            {
-                "loc": ["body", "preferredModel"],
-                "msg": "field required",
-                "type": "value_error.missing",
-            },
             {
                 "loc": ["body", "parameters"],
                 "msg": "field required",

diff --git a/tests/services/guidance_wrapper_test.py b/tests/services/guidance_wrapper_test.py
@@ -1,8 +1,11 @@
 import pytest
 import guidance
 
-from app.models.dtos import Content, ContentType, LLMModel
+from app.models.dtos import Content, ContentType
 from app.services.guidance_wrapper import GuidanceWrapper
+from app.config import LLMModelConfig
+
+llm_model_config = LLMModelConfig(name="test", description="test", llm_credentials={})
 
 
 def test_query_success(mocker):
@@ -17,7 +20,7 @@ def test_query_success(mocker):
     {{gen 'response' temperature=0.0 max_tokens=500}}{{~/assistant}}"""
 
     guidance_wrapper = GuidanceWrapper(
-        model=LLMModel.GPT35_TURBO,
+        model=llm_model_config,
         handlebars=handlebars,
         parameters={"query": "Some query"},
     )
@@ -41,7 +44,7 @@ def test_query_missing_required_params(mocker):
     {{gen 'response' temperature=0.0 max_tokens=500}}{{~/assistant}}"""
 
     guidance_wrapper = GuidanceWrapper(
-        model=LLMModel.GPT35_TURBO,
+        model=llm_model_config,
         handlebars=handlebars,
         parameters={"somethingelse": "Something"},
     )
@@ -63,7 +66,7 @@ def test_query_handlebars_not_generate_response(mocker):
 
     handlebars = "Not a valid handlebars"
     guidance_wrapper = GuidanceWrapper(
-        model=LLMModel.GPT35_TURBO,
+        model=llm_model_config,
         handlebars=handlebars,
         parameters={"query": "Something"},
     )