added logger , type checking

main.py

@@ -1,6 +1,6 @@
 # main.py
 from fastapi import FastAPI,status,Response,Request
-from fastapi.responses import StreamingResponse
+from fastapi.responses import StreamingResponse,FileResponse
 from models import load_text_model,generate_text,load_audio_model,generate_audio,load_image_model, generate_image
 from schemas import VoicePresets
 from utils import audio_array_to_buffer,img_to_bytes
@@ -15,7 +15,7 @@ models = {}
 
 @asynccontextmanager 
 async def lifespan(_: FastAPI) -> AsyncIterator[None]:
-    models["text2image"] = load_image_model() 
+    # models["text2image"] = load_image_model() 
     yield 
     models.clear() 
 
@@ -70,6 +70,9 @@ def serve_language_model_controller(prompt: str) -> str:
     output = generate_text(pipe, prompt) 
     return output 
 
+@app.get("/logs")
+def get_logs():
+    return FileResponse("path/to/logs.csv", media_type='text/csv', filename="logs.csv")
 
 @app.get(
     "/generate/audio",

requirements.txt

@@ -5,4 +5,5 @@ torch
 pydantic
 bitsandbytes
 soundfile
-diffusers
+diffusers
+loguru

utils.py

@@ -2,7 +2,8 @@ from io import BytesIO
 import soundfile
 import numpy as np
 from PIL import Image
-from typing import Literal
+from typing import Literal,TypeAlias
+import tiktoken
 
 def audio_array_to_buffer(audio_array: np.array, sample_rate: int) -> BytesIO:
     buffer = BytesIO()
@@ -16,4 +17,33 @@ def img_to_bytes(
 ) -> bytes:
     buffer = BytesIO()
     image.save(buffer, format=img_format)
-    return buffer.getvalue()
+    return buffer.getvalue()
+
+
+
+
+
+SupportedModels: TypeAlias = Literal["gpt-3.5", "gpt-4"]
+PriceTable: TypeAlias = dict[SupportedModels, float] 
+price_table: PriceTable = {"gpt-3.5": 0.0030, "gpt-4": 0.0200} 
+
+def count_tokens(text: str | None) -> int: 
+    if text is None:
+        logger.warning("Response is None. Assuming 0 tokens used")
+        return 0 
+    enc = tiktoken.encoding_for_model("gpt-4o")
+    return len(enc.encode(text)) 
+
+def calculate_usage_costs(
+    prompt: str,
+    response: str | None,
+    model: SupportedModels,
+) -> tuple[float, float, float]:
+    if model not in price_table:
+        # raise at runtime - in case someone ignores type errors
+        raise ValueError(f"Cost calculation is not supported for {model} model.")
+    price = price_table[model] 9
+    req_costs = price * count_tokens(prompt) / 1000
+    res_costs = price * count_tokens(response) / 1000 
+    total_costs = req_costs + res_costs
+    return req_costs, res_costs, total_costs