Merge remote-tracking branch 'origin/main' into feature/checkhealth-e…

…ndpoint

.github/workflows/build.yml

@@ -42,8 +42,8 @@ jobs:
             if (context.ref.startsWith("refs/tags/")) {
               return context.ref.slice(10);
             }
-            if (context.ref === "refs/heads/develop") {
-              return "develop";
+            if (context.ref === "refs/heads/main") {
+              return "latest";
             }
           }
           return "FALSE";

Dockerfile

@@ -7,4 +7,4 @@ RUN poetry install
 COPY app /app/app
 
 EXPOSE 8000
-CMD ["poetry", "run", "gunicorn", "app.main:app", "-w", "32", "-k", "uvicorn.workers.UvicornWorker", "-b", "0.0.0.0:8000"]
+CMD ["poetry", "run", "gunicorn", "app.main:app", "-w", "8", "-k", "uvicorn.workers.UvicornWorker", "-b", "0.0.0.0:8000"]

app/config.py

@@ -1,6 +1,19 @@
 import os
-from pydantic import BaseModel
+
 from pyaml_env import parse_config
+from pydantic import BaseModel
+
+
+class LLMModelConfig(BaseModel):
+    name: str
+    description: str
+    llm_credentials: dict
+
+
+class APIKeyConfig(BaseModel):
+    token: str
+    comment: str
+    llm_access: list[str]
 
 
 class CacheSettings(BaseModel):
@@ -13,9 +26,9 @@ class CacheParams(BaseModel):
 
 class Settings(BaseModel):
     class PyrisSettings(BaseModel):
-        api_key: str
+        api_keys: list[APIKeyConfig]
+        llms: dict[str, LLMModelConfig]
         cache: CacheSettings
-        llm: dict
 
     pyris: PyrisSettings
 

app/dependencies.py

@@ -1,10 +1,11 @@
 from fastapi import Depends
 from starlette.requests import Request as StarletteRequest
-from app.config import settings
+from app.config import settings, APIKeyConfig
 
 from app.core.custom_exceptions import (
     PermissionDeniedException,
     RequiresAuthenticationException,
+    InvalidModelException,
 )
 
 
@@ -17,7 +18,27 @@ def _get_api_key(request: StarletteRequest) -> str:
     return authorization_header
 
 
-class PermissionsValidator:
-    def __call__(self, api_key: str = Depends(_get_api_key)):
-        if api_key != settings.pyris.api_key:
-            raise PermissionDeniedException
+class TokenValidator:
+    async def __call__(
+        self, api_key: str = Depends(_get_api_key)
+    ) -> APIKeyConfig:
+        for key in settings.pyris.api_keys:
+            if key.token == api_key:
+                return key
+        raise PermissionDeniedException
+
+
+class TokenPermissionsValidator:
+    async def __call__(
+        self, request: StarletteRequest, api_key: str = Depends(_get_api_key)
+    ):
+        for key in settings.pyris.api_keys:
+            if key.token == api_key:
+                body = await request.json()
+                if body.get("preferredModel") in key.llm_access:
+                    return
+                else:
+                    raise InvalidModelException(
+                        str(body.get("preferredModel"))
+                    )
+        raise PermissionDeniedException

app/main.py

@@ -4,6 +4,7 @@
 from app.routes.messages import router as messages_router
 from app.routes.health import router as health_router
 from app.services.hazelcast_client import hazelcast_client
+from app.routes.models import router as models_router
 
 app = FastAPI(default_response_class=ORJSONResponse)
 
@@ -15,3 +16,4 @@ async def shutdown():
 
 app.include_router(messages_router)
 app.include_router(health_router)
+app.include_router(models_router)

app/models/dtos.py

@@ -3,12 +3,6 @@
 from datetime import datetime
 
 
-class LLMModel(str, Enum):
-    GPT35_TURBO = "GPT35_TURBO"
-    GPT35_TURBO_16K_0613 = "GPT35_TURBO_16K_0613"
-    GPT35_TURBO_0613 = "GPT35_TURBO_0613"
-
-
 class LLMStatus(str, Enum):
     UP = "UP"
     DOWN = "DOWN"
@@ -41,10 +35,16 @@ class Message(BaseModel):
         )
         content: list[Content]
 
-    used_model: LLMModel = Field(..., alias="usedModel")
+    used_model: str = Field(..., alias="usedModel")
     message: Message
 
 
 class ModelStatus(BaseModel):
-    model: LLMModel
+    model: str
     status: LLMStatus
+
+
+class LLMModelResponse(BaseModel):
+    id: str
+    name: str
+    description: str

app/routes/health.py

@@ -1,27 +1,28 @@
 from fastapi import APIRouter, Depends
 
-from app.dependencies import PermissionsValidator
-from app.models.dtos import LLMModel, LLMStatus, ModelStatus
+from app.dependencies import TokenValidator
+from app.models.dtos import LLMStatus, ModelStatus
 from app.services.guidance_wrapper import GuidanceWrapper
 from app.services.circuit_breaker import CircuitBreaker
+from app.config import settings
 
 router = APIRouter()
 
 
-@router.get("/api/v1/health", dependencies=[Depends(PermissionsValidator())])
+@router.get("/api/v1/health", dependencies=[Depends(TokenValidator())])
 def checkhealth() -> list[ModelStatus]:
     result = []
 
-    for model in LLMModel:
+    for key, model in settings.pyris.llms.items():
         circuit_status = CircuitBreaker.get_status(
             checkhealth_func=GuidanceWrapper(model=model).is_up,
-            cache_key=model,
+            cache_key=key,
         )
         status = (
             LLMStatus.UP
             if circuit_status == CircuitBreaker.Status.CLOSED
             else LLMStatus.DOWN
         )
-        result.append(ModelStatus(model=model, status=status))
+        result.append(ModelStatus(model=key, status=status))
 
     return result

app/routes/messages.py

@@ -1,6 +1,6 @@
-from fastapi import APIRouter, Depends
 from datetime import datetime, timezone
 
+from fastapi import APIRouter, Depends
 from parsimonious.exceptions import IncompleteParseError
 
 from app.core.custom_exceptions import (
@@ -9,20 +9,21 @@
     InternalServerException,
     InvalidModelException,
 )
-from app.dependencies import PermissionsValidator
-from app.models.dtos import SendMessageRequest, SendMessageResponse, LLMModel
+from app.dependencies import TokenPermissionsValidator
+from app.models.dtos import SendMessageRequest, SendMessageResponse
 from app.services.circuit_breaker import CircuitBreaker
 from app.services.guidance_wrapper import GuidanceWrapper
+from app.config import settings
 
 router = APIRouter(tags=["messages"])
 
 
 @router.post(
-    "/api/v1/messages", dependencies=[Depends(PermissionsValidator())]
+    "/api/v1/messages", dependencies=[Depends(TokenPermissionsValidator())]
 )
 def send_message(body: SendMessageRequest) -> SendMessageResponse:
     try:
-        model = LLMModel(body.preferred_model)
+        model = settings.pyris.llms[body.preferred_model]
     except ValueError as e:
         raise InvalidModelException(str(e))
 
@@ -35,7 +36,7 @@ def send_message(body: SendMessageRequest) -> SendMessageResponse:
     try:
         content = CircuitBreaker.protected_call(
             func=guidance.query,
-            cache_key=model,
+            cache_key=body.preferred_model,
             accepted_exceptions=(KeyError, SyntaxError, IncompleteParseError),
         )
     except KeyError as e:

app/routes/models.py

@@ -0,0 +1,23 @@
+from fastapi import APIRouter, Depends
+
+from app.dependencies import TokenValidator
+from app.config import settings, APIKeyConfig
+from app.models.dtos import LLMModelResponse
+
+router = APIRouter(tags=["models"])
+
+
+@router.get("/api/v1/models")
+def send_message(
+    api_key_config: APIKeyConfig = Depends(TokenValidator()),
+) -> list[LLMModelResponse]:
+    llm_ids = api_key_config.llm_access
+
+    # Convert the result of this to a list of LLMModelResponse
+    return [
+        LLMModelResponse(
+            id=key, name=config.name, description=config.description
+        )
+        for key, config in settings.pyris.llms.items()
+        if key in llm_ids
+    ]

app/services/guidance_wrapper.py

@@ -1,7 +1,7 @@
 import guidance
 
-from app.config import settings
-from app.models.dtos import Content, ContentType, LLMModel
+from app.config import LLMModelConfig
+from app.models.dtos import Content, ContentType
 
 
 class GuidanceWrapper:
@@ -11,7 +11,7 @@ def __new__(cls, *_, **__):
         return super(GuidanceWrapper, cls).__new__(cls)
 
     def __init__(
-        self, model: LLMModel, handlebars="", parameters=None
+        self, model: LLMModelConfig, handlebars="", parameters=None
     ) -> None:
         if parameters is None:
             parameters = {}
@@ -63,5 +63,5 @@ def is_up(self) -> bool:
         return content == "1"
 
     def _get_llm(self):
-        llm_credentials = settings.pyris.llm[self.model.value]
+        llm_credentials = self.model.llm_credentials
         return guidance.llms.OpenAI(**llm_credentials)

application-docker.test.yml

@@ -4,19 +4,20 @@ pyris:
     hazelcast:
       host: cache
       port: 5701
-  llm:
-    GPT35_TURBO:
-      model:
-      token:
-      api_base:
-      api_type:
-      api_version:
-      deployment_id:
-    GPT35_TURBO_16K_0613:
-      model:
-      token:
-      chat_mode:
-    GPT35_TURBO_0613:
-      model:
-      token:
-      chat_mode:
+  api_keys:
+    - token: "secret"
+      comment: "DUMMY"
+      llm_access:
+        - DUMMY
+
+  llms:
+    DUMMY:
+      name: "Dummy model"
+      description: "Dummy model for testing"
+      llm_credentials:
+        api_base: ""
+        api_type: ""
+        api_version: ""
+        deployment_id: ""
+        model: ""
+        token: ""

application.example.yml

@@ -1,22 +1,22 @@
 pyris:
-  api_key: secret
   cache:
     hazelcast:
       host: localhost
       port: 5701
-  llm:
-    GPT35_TURBO:
-      model:
-      token:
-      api_base:
-      api_type:
-      api_version:
-      deployment_id:
-    GPT35_TURBO_16K_0613:
-      model:
-      token:
-      chat_mode:
-    GPT35_TURBO_0613:
-      model:
-      token:
-      chat_mode:
+  api_keys:
+    - token: "secret"
+      comment: "DUMMY"
+      llm_access:
+        - DUMMY
+
+  llms:
+    DUMMY:
+      name: "Dummy model"
+      description: "Dummy model for testing"
+      llm_credentials:
+        api_base: ""
+        api_type: ""
+        api_version: ""
+        deployment_id: ""
+        model: ""
+        token: ""

application.test.yml

@@ -4,19 +4,20 @@ pyris:
     hazelcast:
       host: localhost
       port: 5701
-  llm:
-    GPT35_TURBO:
-      model:
-      token:
-      api_base:
-      api_type:
-      api_version:
-      deployment_id:
-    GPT35_TURBO_16K_0613:
-      model:
-      token:
-      chat_mode:
-    GPT35_TURBO_0613:
-      model:
-      token:
-      chat_mode:
+  api_keys:
+    - token: "secret"
+      comment: "DUMMY"
+      llm_access:
+        - DUMMY
+
+  llms:
+    DUMMY:
+      name: "Dummy model"
+      description: "Dummy model for testing"
+      llm_credentials:
+        api_base: ""
+        api_type: ""
+        api_version: ""
+        deployment_id: ""
+        model: ""
+        token: ""