Merge branch 'main' into feature/truncate-function

.github/workflows/test.yml

@@ -44,5 +44,8 @@ jobs:
       run: |
         poetry run flake8 .
 
+    - name: Start cache service
+      run: docker compose up -d --build cache
+
     - name: Test with pytest
       run: poetry run pytest

.gitignore

@@ -130,6 +130,7 @@ venv.bak/
 
 # Config files
 application.yml
+application-docker.yml
 
 # Spyder project settings
 .spyderproject

README.md

@@ -17,6 +17,7 @@
 
 ## With docker
 - (optional) Install **[Task](https://taskfile.dev)**
+- Copy `application.example.yml` to `application-docker.yml` and change necessary values
 - Build docker: `docker compose build` or `task build`
 - Run server: `docker compose up -d` or `task up`
 - Stop server: `docker compose down` or `task down`

Taskfile.yml

@@ -30,3 +30,8 @@ tasks:
   logs:
     cmds:
       - docker compose logs -f server
+
+  console:
+    aliases: [ c ]
+    cmds:
+      - docker attach pyris-server # using <ctrl-p><ctrl-q> to detach

app/config.py

@@ -16,19 +16,29 @@ class APIKeyConfig(BaseModel):
     llm_access: list[str]
 
 
+class CacheSettings(BaseModel):
+    class CacheParams(BaseModel):
+        host: str
+        port: int
+
+    hazelcast: CacheParams
+
+
 class Settings(BaseModel):
     class PyrisSettings(BaseModel):
         api_keys: list[APIKeyConfig]
         llms: dict[str, LLMModelConfig]
+        cache: CacheSettings
 
     pyris: PyrisSettings
 
     @classmethod
     def get_settings(cls):
+        postfix = "-docker" if "DOCKER" in os.environ else ""
         if "RUN_ENV" in os.environ and os.environ["RUN_ENV"] == "test":
-            file_path = "application.test.yml"
+            file_path = f"application{postfix}.test.yml"
         else:
-            file_path = "application.yml"
+            file_path = f"application{postfix}.yml"
 
         return Settings.parse_obj(parse_config(file_path))
 

app/main.py

@@ -2,9 +2,18 @@
 from fastapi.responses import ORJSONResponse
 
 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)
 
+
+@app.on_event("shutdown")
+async def shutdown():
+    hazelcast_client.shutdown()
+
+
 app.include_router(messages_router)
+app.include_router(health_router)
 app.include_router(models_router)

app/models/dtos.py

@@ -3,6 +3,12 @@
 from datetime import datetime
 
 
+class LLMStatus(str, Enum):
+    UP = "UP"
+    DOWN = "DOWN"
+    NOT_AVAILABLE = "NOT_AVAILABLE"
+
+
 class ContentType(str, Enum):
     TEXT = "text"
 
@@ -33,6 +39,11 @@ class Message(BaseModel):
     message: Message
 
 
+class ModelStatus(BaseModel):
+    model: str
+    status: LLMStatus
+
+
 class LLMModelResponse(BaseModel):
     id: str
     name: str

app/routes/health.py

@@ -0,0 +1,28 @@
+from fastapi import APIRouter, Depends
+
+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(TokenValidator())])
+def checkhealth() -> list[ModelStatus]:
+    result = []
+
+    for key, model in settings.pyris.llms.items():
+        circuit_status = CircuitBreaker.get_status(
+            checkhealth_func=GuidanceWrapper(model=model).is_up,
+            cache_key=key,
+        )
+        status = (
+            LLMStatus.UP
+            if circuit_status == CircuitBreaker.Status.CLOSED
+            else LLMStatus.DOWN
+        )
+        result.append(ModelStatus(model=key, status=status))
+
+    return result

app/routes/messages.py

@@ -11,6 +11,7 @@
 )
 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
 
@@ -33,7 +34,11 @@ def send_message(body: SendMessageRequest) -> SendMessageResponse:
     )
 
     try:
-        content = guidance.query()
+        content = CircuitBreaker.protected_call(
+            func=guidance.query,
+            cache_key=body.preferred_model,
+            accepted_exceptions=(KeyError, SyntaxError, IncompleteParseError),
+        )
     except KeyError as e:
         raise MissingParameterException(str(e))
     except (SyntaxError, IncompleteParseError) as e:

app/services/cache.py

@@ -0,0 +1,154 @@
+from abc import ABC, abstractmethod
+from datetime import datetime, timedelta
+from typing import Optional, Any
+from pydantic import BaseModel
+from app.services.hazelcast_client import hazelcast_client
+
+
+class CacheStoreInterface(ABC):
+    @abstractmethod
+    def get(self, name: str) -> Any:
+        """
+        Get the value at key ``name``
+
+        Returns:
+            The value at key ``name``, or None if the key doesn't exist
+        """
+        pass
+
+    @abstractmethod
+    def set(self, name: str, value, ex: Optional[int] = None):
+        """
+        Set the value at key ``name`` to ``value``
+
+        ``ex`` sets an expire flag on key ``name`` for ``ex`` seconds.
+        """
+        pass
+
+    @abstractmethod
+    def expire(self, name: str, ex: int):
+        """
+        Set an expire flag on key ``name`` for ``ex`` seconds
+        """
+        pass
+
+    @abstractmethod
+    def incr(self, name: str) -> int:
+        """
+        Increase the integer value of a key ``name`` by 1.
+        If the key does not exist, it is set to 0
+        before performing the operation.
+
+        Returns:
+            The value of key ``name`` after the increment
+
+        Raises:
+            TypeError: if cache value is not an integer
+        """
+        pass
+
+    @abstractmethod
+    def flushdb(self):
+        """
+        Delete all keys in the current database
+        """
+        pass
+
+    @abstractmethod
+    def delete(self, name: str):
+        """
+        Delete the key specified by ``name``
+        """
+        pass
+
+
+class InMemoryCacheStore(CacheStoreInterface):
+    class CacheValue(BaseModel):
+        value: Any
+        expired_at: Optional[datetime]
+
+    def __init__(self):
+        self._cache: dict[str, Optional[InMemoryCacheStore.CacheValue]] = {}
+
+    def get(self, name: str) -> Any:
+        current_time = datetime.now()
+        data = self._cache.get(name)
+        if data is None:
+            return None
+        if data.expired_at and current_time > data.expired_at:
+            del self._cache[name]
+            return None
+        return data.value
+
+    def set(self, name: str, value, ex: Optional[int] = None):
+        self._cache[name] = InMemoryCacheStore.CacheValue(
+            value=value, expired_at=None
+        )
+        if ex is not None:
+            self.expire(name, ex)
+
+    def expire(self, name: str, ex: int):
+        current_time = datetime.now()
+        expired_at = current_time + timedelta(seconds=ex)
+        data = self._cache.get(name)
+        if data is None:
+            return
+        data.expired_at = expired_at
+
+    def incr(self, name: str) -> int:
+        value = self.get(name)
+        if value is None:
+            self._cache[name] = InMemoryCacheStore.CacheValue(
+                value=1, expired_at=None
+            )
+            return 1
+        if isinstance(value, int):
+            value += 1
+            self.set(name, value)
+            return value
+        raise TypeError("value is not an integer")
+
+    def flushdb(self):
+        self._cache = {}
+
+    def delete(self, name: str):
+        if name in self._cache:
+            del self._cache[name]
+
+
+class HazelcastCacheStore(CacheStoreInterface):
+    def __init__(self):
+        self._cache = hazelcast_client.get_map("cache_store").blocking()
+
+    def get(self, name: str) -> Any:
+        return self._cache.get(name)
+
+    def set(self, name: str, value, ex: Optional[int] = None):
+        self._cache.put(name, value, ex)
+
+    def expire(self, name: str, ex: int):
+        self._cache.set_ttl(name, ex)
+
+    def incr(self, name: str) -> int:
+        flag = False
+        value = 0
+        while not flag:
+            value = self._cache.get(name)
+            if value is None:
+                self._cache.set(name, 1)
+                return 1
+            if not isinstance(value, int):
+                raise TypeError("value is not an integer")
+
+            flag = self._cache.replace_if_same(name, value, value + 1)
+
+        return value + 1
+
+    def flushdb(self):
+        self._cache.clear()
+
+    def delete(self, name: str):
+        self._cache.remove(name)
+
+
+cache_store = HazelcastCacheStore()