Merge branch 'main' into edr/getting_started_doc

Author: evanderiel

.env

@@ -0,0 +1,3 @@
+CUDA_VISIBLE_DEVICES=""
+USE_DEPLOYMENT_CACHE = True
+SAVE_DEPLOYMENT_CACHE = True

.github/workflows/python-package.yml

@@ -38,4 +38,4 @@ jobs:
         sudo apt-get update
         sudo apt-get install ffmpeg
     - name: Test with pytest
-      run: poetry run pytest
+      run: poetry run pytest -vv

.gitignore

@@ -120,7 +120,6 @@ celerybeat.pid
 *.sage.py
 
 # Environments
-.env
 .venv
 env/
 venv/

README.md

@@ -135,6 +135,52 @@ that Ruff problems appear while you edit, and formatting is applied
 automatically on save.
 
 
+## Testing
+
+The project uses pytest for testing. To run the tests, use the following command:
+
+```bash
+poetry run pytest
+```
+
+If you are using VS Code, you can run the tests using the Test Explorer that is installed with the [Python extension](https://code.visualstudio.com/docs/python/testing).
+
+There are a few environment variables that can be set to control the behavior of the tests:
+- `USE_DEPLOYMENT_CACHE`: If set to `true`, the tests will use the deployment cache to avoid downloading the models and running the deployments. This is useful for running integration tests faster and in the environment where GPU is not available.
+- `SAVE_DEPLOYMENT_CACHE`: If set to `true`, the tests will save the deployment cache after running the deployments. This is useful for updating the deployment cache if new deployments or tests are added.
+
+### How to use the deployment cache environment variables
+
+Here are some examples of how to use the deployment cache environment variables.
+
+#### Do you want to run the tests normally using GPU?
+    
+```bash
+USE_DEPLOYMENT_CACHE=false
+SAVE_DEPLOYMENT_CACHE=false
+```
+
+This is the default behavior. The tests will run normally using GPU and the deployment cache will be completely ignored.
+
+#### Do you want to run the tests faster without GPU?
+
+```bash
+USE_DEPLOYMENT_CACHE=true
+SAVE_DEPLOYMENT_CACHE=false
+```
+
+This will run the tests using the deployment cache to avoid downloading the models and running the deployments. The deployment cache will not be updated after running the deployments. Only use it if you are sure that the deployment cache is up to date.
+
+#### Do you want to update the deployment cache?
+
+```bash
+USE_DEPLOYMENT_CACHE=false
+SAVE_DEPLOYMENT_CACHE=true
+```
+
+This will run the tests normally using GPU and save the deployment cache after running the deployments. Use it if you have added new deployments or tests and want to update the deployment cache.
+
+
 ## Databases
 The project uses two databases: a vector database as well as a tradtional SQL database,
 referred to internally as vectorstore and datastore, respectively.
@@ -156,4 +202,13 @@ ORM models referenced in the rest of the code should be imported from `aana.mode
 not from that model's file for reasons explained in `aana/models/db/__init__.py`. This also means that 
 if you add a new model class, it should be imported by `__init__.py` in addition to creating a migration.
 
-Higher level code for interacting with the ORM is available in `aana.repository.data`.
+Higher level code for interacting with the ORM is available in `aana.repository.data`.
+
+## Settings
+
+Here are the environment variables that can be used to configure the Aaana SDK:
+- TMP_DATA_DIR: The directory to store temporary data. Default: `/tmp/aana`.
+- NUM_WORKERS: The number of request workers. Default: `2`.
+- DB_CONFIG: The database configuration in the format `{"datastore_type": "sqlite", "datastore_config": {"path": "/path/to/sqlite.db"}}`. Currently only SQLite and PostgreSQL are supported. Default: `{"datastore_type": "sqlite", "datastore_config": {"path": "/var/lib/aana_data"}}`.
+- USE_DEPLOYMENT_CACHE (testing only): If set to `true`, the tests will use the deployment cache to avoid downloading the models and running the deployments. Default: `false`.
+- SAVE_DEPLOYMENT_CACHE (testing only): If set to `true`, the tests will save the deployment cache after running the deployments. Default: `false`.

aana/api/app.py

@@ -81,40 +81,15 @@ def custom_exception_handler(request: Request | None, exc_raw: Exception):
     )
 
 
-@app.exception_handler(BaseException)
-async def pipeline_exception_handler(request: Request, exc: BaseException):
+@app.exception_handler(Exception)
+async def pipeline_exception_handler(request: Request, exc: Exception):
     """This handler is used to handle exceptions raised by the Mobius Pipeline and Aana application.
 
     Args:
         request (Request): The request object
-        exc (BaseException): The exception raised
+        exc (Exception): The exception raised
 
     Returns:
         JSONResponse: JSON response with the error details
     """
     return custom_exception_handler(request, exc)
-
-
-@app.exception_handler(RayTaskError)
-async def ray_task_error_handler(request: Request, exc: RayTaskError):
-    """This handler is used to handle RayTaskError exceptions.
-
-    Args:
-        request (Request): The request object
-        exc (RayTaskError): The exception raised
-
-    Returns:
-        JSONResponse: JSON response with the error details. The response contains the following fields:
-            error: The name of the exception class.
-            message: The message of the exception.
-            stacktrace: The stacktrace of the exception.
-    """
-    error = exc.__class__.__name__
-    stacktrace = traceback.format_exc()
-
-    return AanaJSONResponse(
-        status_code=400,
-        content=ExceptionResponseModel(
-            error=error, message=str(exc), stacktrace=stacktrace
-        ).dict(),
-    )

aana/api/responses.py

@@ -1,36 +1,9 @@
-from pathlib import Path
 from typing import Any
 
 import orjson
 from fastapi.responses import JSONResponse
-from pydantic import BaseModel
 
-
-def json_serializer_default(obj: Any) -> Any:
-    """Default function for json serializer to handle pydantic models.
-
-    If json serializer does not know how to serialize an object, it calls the default function.
-
-    If we see that the object is a pydantic model,
-    we call the dict method to get the dictionary representation of the model
-    that json serializer can deal with.
-
-    If the object is not a pydantic model, we raise a TypeError.
-
-    Args:
-        obj (Any): The object to serialize.
-
-    Returns:
-        Any: The serializable object.
-
-    Raises:
-        TypeError: If the object is not a pydantic model.
-    """
-    if isinstance(obj, BaseModel):
-        return obj.dict()
-    if isinstance(obj, Path):
-        return str(obj)
-    raise TypeError
+from aana.utils.json import json_serializer_default
 
 
 class AanaJSONResponse(JSONResponse):

aana/configs/db.py

@@ -5,12 +5,32 @@
 
 from alembic import command
 from alembic.config import Config
-from sqlalchemy import String, create_engine
+from sqlalchemy import String, TypeDecorator, create_engine
 
-# These are here so we can change types in a single place.
+from aana.models.pydantic.media_id import MediaId
 
-media_id_type: TypeAlias = str
-MediaIdSqlType: TypeAlias = String
+
+class MediaIdType(TypeDecorator):
+    """Custom type for handling MediaId objects with SQLAlchemy."""
+
+    impl = String
+
+    cache_ok = True
+
+    def process_bind_param(self, value, dialect):
+        """Convert a MediaId instance to a string value for storage."""
+        if value is None:
+            return value
+        return str(value)
+
+    def process_result_value(self, value, dialect):
+        """Convert a string value from the database back into a MediaId instance."""
+        if value is None:
+            return value
+        return MediaId(value)
+
+
+MediaIdSqlType: TypeAlias = MediaIdType
 
 
 class SQLiteConfig(TypedDict):
@@ -112,3 +132,17 @@ def run_alembic_migrations(settings):
 
     alembic_config = get_alembic_config(settings, ini_file_path, alembic_data_path)
     command.upgrade(alembic_config, "head")
+
+
+def drop_all_tables(settings):
+    """Drops all tables in the database."""
+    # TODO: only allow this in testing mode
+    current_path = Path(__file__)
+    aana_root = current_path.parent.parent  # go up two directories
+    if aana_root.name != "aana":  # we are not in the right place
+        raise RuntimeError("Not in right directory, exiting.")  # noqa: TRY003
+    ini_file_path = aana_root / "alembic.ini"
+    alembic_data_path = aana_root / "alembic"
+
+    alembic_config = get_alembic_config(settings, ini_file_path, alembic_data_path)
+    command.downgrade(alembic_config, "base")

aana/configs/endpoints.py

@@ -69,17 +69,6 @@
             ],
         ),
     ],
-    "video": [
-        Endpoint(
-            name="video_extract_frames",
-            path="/video/extract_frames",
-            summary="Extract frames from a video",
-            outputs=[
-                EndpointOutput(name="timestamps", output="timestamps"),
-                EndpointOutput(name="duration", output="duration"),
-            ],
-        )
-    ],
     "whisper": [
         Endpoint(
             name="whisper_transcribe",
@@ -124,8 +113,6 @@
                 ),
             ],
         ),
-    ],
-    "chat_with_video_test": [
         Endpoint(
             name="delete_media_id",
             path="/video/delete",
@@ -136,54 +123,6 @@
         ),
     ],
     "chat_with_video": [
-        Endpoint(
-            name="delete_media_id",
-            path="/video/delete",
-            summary="Delete a video",
-            outputs=[
-                EndpointOutput(name="deleted_media_id", output="deleted_media_id")
-            ],
-        ),
-        Endpoint(
-            name="blip2_video_generate",
-            path="/video/generate_captions",
-            summary="Generate captions for videos using BLIP2 OPT-2.7B",
-            outputs=[
-                EndpointOutput(
-                    name="captions",
-                    output="video_captions_hf_blip2_opt_2_7b",
-                    streaming=True,
-                ),
-                EndpointOutput(
-                    name="timestamps", output="video_timestamps", streaming=True
-                ),
-            ],
-            streaming=True,
-        ),
-        Endpoint(
-            name="whisper_transcribe",
-            path="/video/transcribe",
-            summary="Transcribe a video using Whisper Medium",
-            outputs=[
-                EndpointOutput(
-                    name="transcription",
-                    output="video_transcriptions_whisper_medium",
-                    streaming=True,
-                ),
-                EndpointOutput(
-                    name="segments",
-                    output="video_transcriptions_segments_whisper_medium",
-                    streaming=True,
-                ),
-                EndpointOutput(
-                    name="info",
-                    output="video_transcriptions_info_whisper_medium",
-                    streaming=True,
-                ),
-                EndpointOutput(name="transcription_id", output="transcription_id"),
-            ],
-            streaming=True,
-        ),
         Endpoint(
             name="index_video_stream",
             path="/video/index_stream",
@@ -217,27 +156,6 @@
             ],
             streaming=True,
         ),
-        Endpoint(
-            name="llm_generate",
-            path="/llm/generate",
-            summary="Generate text using LLaMa2 7B Chat",
-            outputs=[
-                EndpointOutput(name="completion", output="vllm_llama2_7b_chat_output")
-            ],
-        ),
-        Endpoint(
-            name="llm_generate_stream",
-            path="/llm/generate_stream",
-            summary="Generate text using LLaMa2 7B Chat (streaming)",
-            outputs=[
-                EndpointOutput(
-                    name="completion",
-                    output="vllm_llama2_7b_chat_output_stream",
-                    streaming=True,
-                )
-            ],
-            streaming=True,
-        ),
         Endpoint(
             name="video_chat_stream",
             path="/video/chat_stream",
@@ -259,6 +177,14 @@
                 EndpointOutput(name="metadata", output="video_metadata"),
             ],
         ),
+        Endpoint(
+            name="delete_media_id",
+            path="/video/delete",
+            summary="Delete a video",
+            outputs=[
+                EndpointOutput(name="deleted_media_id", output="deleted_media_id")
+            ],
+        ),
     ],
     "stablediffusion2": [
         Endpoint(

aana/configs/pipeline.py

@@ -15,7 +15,9 @@
 from aana.models.pydantic.captions import CaptionsList, VideoCaptionsList
 from aana.models.pydantic.chat_message import ChatDialog
 from aana.models.pydantic.image_input import ImageInputList
+from aana.models.pydantic.media_id import MediaId
 from aana.models.pydantic.prompt import Prompt
+from aana.models.pydantic.question import Question
 from aana.models.pydantic.sampling_params import SamplingParams
 from aana.models.pydantic.video_input import VideoInput, VideoInputList
 from aana.models.pydantic.video_metadata import VideoMetadata
@@ -526,6 +528,7 @@
                 "name": "media_id",
                 "key": "media_id",
                 "path": "media_id",
+                "data_model": MediaId,
             }
         ],
     },
@@ -538,6 +541,7 @@
                 "name": "question",
                 "key": "question",
                 "path": "question",
+                "data_model": Question,
             }
         ],
     },