Use Groq for inference for majority of the models

app/app/main.py

@@ -14,6 +14,7 @@ from arcana_codex import (
 from bson.objectid import ObjectId
 from fastapi import Depends, FastAPI, Header, HTTPException, Request, status
 from fastapi.responses import JSONResponse
+from groq import Groq
 from llama_cpp import Llama
 from pydantic import BaseModel, EmailStr
 from pymongo.mongo_client import MongoClient
@@ -22,11 +23,13 @@ from starlette.responses import FileResponse
 __version__ = "0.0.0"
 
 
-class SupportedModelPipes(StrEnum):
+class SupportedModels(StrEnum):
+    Gemma2 = "gemma2"
     Gemma3 = "gemma3"
-    QwenOpenR1 = "qwen-open-r1"
-    SmolLLM2 = "smollm2"
-    SmolLLM2Reasoning = "smollm2-reasoning"
+    Llama3_3 = "llama3_3"
+    Llama3_1 = "llama3_1"
+    Qwen2_5 = "qwen2_5"
+    Deepseek_R1 = "deepseek_r1"
 
 
 class LogEvent(StrEnum):
@@ -34,24 +37,6 @@ class LogEvent(StrEnum):
     LOGIN = "login"
 
 
-smollm2_pipeline = Llama.from_pretrained(
-    repo_id="HuggingFaceTB/SmolLM2-360M-Instruct-GGUF",
-    filename="smollm2-360m-instruct-q8_0.gguf",
-    verbose=False,
-)
-
-smollm2_reasoning_pipeline = Llama.from_pretrained(
-    repo_id="tensorblock/Reasoning-SmolLM2-135M-GGUF",
-    filename="Reasoning-SmolLM2-135M-Q8_0.gguf",
-    verbose=False,
-)
-
-qwen_open_r1_pipeline = Llama.from_pretrained(
-    repo_id="tensorblock/Qwen2.5-0.5B-Open-R1-Distill-GGUF",
-    filename="Qwen2.5-0.5B-Open-R1-Distill-Q8_0.gguf",
-    verbose=False,
-)
-
 gemma_3_pipeline = Llama.from_pretrained(
     repo_id="ggml-org/gemma-3-1b-it-GGUF",
     filename="gemma-3-1b-it-Q8_0.gguf",
@@ -60,7 +45,7 @@ gemma_3_pipeline = Llama.from_pretrained(
 
 
 class ChatRequest(BaseModel):
-    model: SupportedModelPipes = SupportedModelPipes.SmolLLM2
+    model: SupportedModels = SupportedModels.Gemma2
     message: str
 
 
@@ -120,11 +105,31 @@ def verify_authorization_header(
         )
 
 
+def process_groq_chat_request(
+    groq_client: Groq, message: str, model: str
+) -> str | None:
+    return (
+        groq_client.chat.completions.create(
+            messages=[
+                {"role": "system", "content": "You are a helpful assistant."},
+                {"role": "user", "content": f"{message}"},
+            ],
+            max_completion_tokens=1024,
+            seed=8,
+            model=model,
+        )
+        .choices[0]
+        .message.content
+    )
+
+
 @asynccontextmanager
 async def lifespan(app: FastAPI):  # noqa: ARG001
     # Set API key in FastAPI app
     app.ARCANA_API_KEY = os.environ.get("ARCANA_API_KEY", "")
 
+    app.groq_client = Groq(api_key=os.environ.get("GROQ_API_KEY"))
+
     app.mongo_db = MongoClient(
         os.environ.get("MONGO_URI", "mongodb+srv://localhost:27017/")
     )[os.environ.get("MONGO_DB", "arcana_hf_demo_test")]
@@ -136,6 +141,8 @@ async def lifespan(app: FastAPI):  # noqa: ARG001
     # Clear API key to avoid leaking it
     app.ARCANA_API_KEY = ""
 
+    app.groq_client = None
+
     logging.info("Application stopped")
 
 
@@ -182,26 +189,53 @@ def chat(
 
     logger.info(f"Using {payload.model}")
 
+    inference_start_time = time.perf_counter()
     match payload.model:
-        case SupportedModelPipes.Gemma3:
-            ai_pipeline = gemma_3_pipeline
-        case SupportedModelPipes.QwenOpenR1:
-            ai_pipeline = qwen_open_r1_pipeline
-        case SupportedModelPipes.SmolLLM2:
-            ai_pipeline = smollm2_pipeline
-        case SupportedModelPipes.SmolLLM2Reasoning:
-            ai_pipeline = smollm2_reasoning_pipeline
+        case SupportedModels.Gemma3:
+            llm_response = gemma_3_pipeline.create_chat_completion(
+                messages=[
+                    {"role": "system", "content": "You are a helpful assistant."},
+                    {"role": "user", "content": f"{payload.message}"},
+                ],
+                max_tokens=512,
+                seed=8,
+            )["choices"][0]["message"]["content"]
+        case SupportedModels.Gemma2:
+            llm_response = process_groq_chat_request(
+                groq_client=request.app.groq_client,
+                message=payload.message,
+                model="gemma2-9b-it",
+            )
+        case SupportedModels.Llama3_3:
+            llm_response = process_groq_chat_request(
+                groq_client=request.app.groq_client,
+                message=payload.message,
+                model="llama-3.3-70b-versatile",
+            )
+        case SupportedModels.Llama3_1:
+            llm_response = process_groq_chat_request(
+                groq_client=request.app.groq_client,
+                message=payload.message,
+                model="llama-3.1-8b-instant",
+            )
+        case SupportedModels.Qwen2_5:
+            llm_response = process_groq_chat_request(
+                groq_client=request.app.groq_client,
+                message=payload.message,
+                model="qwen-2.5-32b",
+            )
+        case SupportedModels.Deepseek_R1:
+            llm_response = process_groq_chat_request(
+                groq_client=request.app.groq_client,
+                message=payload.message,
+                model="deepseek-r1-distill-qwen-32b",
+            )
 
-    inference_start_time = time.perf_counter()
-    ai_response = ai_pipeline.create_chat_completion(
-        messages=[{"role": "user", "content": f"{payload.message}"}],
-        max_tokens=512,
-        seed=8,
-    )["choices"][0]["message"]["content"].strip()
+    ai_response = "" if llm_response is None else llm_response.strip()
     inference_end_time = time.perf_counter()
 
-    elapsed_time = inference_end_time - inference_start_time
-    logger.info(f"Inference took: {elapsed_time:.4f} seconds")
+    inference_elapsed_time = inference_end_time - inference_start_time
+    logger.info(f"Inference took: {inference_elapsed_time:.4f} seconds")
 
     integrate_payload = AdUnitsIntegrateModel(
         ad_unit_ids=[
@@ -210,10 +244,15 @@ def chat(
         base_content=ai_response,
     )
 
+    integration_start_time = time.perf_counter()
     integration_result = client.integrate_ad_units(integrate_payload)
     integrated_content = integration_result.get("response_data", {}).get(
         "integrated_content"
     )
+    integration_end_time = time.perf_counter()
+
+    integration_elapsed_time = integration_end_time - integration_start_time
+    logger.info(f"Integration took: {integration_elapsed_time:.4f} seconds")
 
     request.app.mongo_db["logs"].insert_one(
         {

app/static/index.html

@@ -62,18 +62,25 @@
                     <div id="model-dropdown"
                         class="absolute z-10 w-full mt-1 bg-white border border-gray-300 rounded-md shadow-lg hidden">
                         <ul class="py-1">
-                            <li class="model-option px-4 py-2 hover:bg-gray-100 cursor-pointer" data-value="smollm2">
-                                SmolLM2</li>
+                            <li class="model-option px-4 py-2 hover:bg-gray-100 cursor-pointer" data-value="gemma2">
+                                Gemma2</li>
                         </ul>
                         <ul class="py-1">
-                            <li class="model-option px-4 py-2 hover:bg-gray-100 cursor-pointer"
-                                data-value="smollm2-reasoning">
-                                SmolLLM2Reasoning</li>
+                            <li class="model-option px-4 py-2 hover:bg-gray-100 cursor-pointer" data-value="llama3_1">
+                                Llama3.1</li>
+                        </ul>
+                        <ul class="py-1">
+                            <li class="model-option px-4 py-2 hover:bg-gray-100 cursor-pointer" data-value="llama3_3">
+                                Llama3.3</li>
+                        </ul>
+                        <ul class="py-1">
+                            <li class="model-option px-4 py-2 hover:bg-gray-100 cursor-pointer" data-value="qwen2_5">
+                                Qwen2.5</li>
                         </ul>
                         <ul class="py-1">
                             <li class="model-option px-4 py-2 hover:bg-gray-100 cursor-pointer"
-                                data-value="qwen-open-r1">
-                                QwenOpenR1</li>
+                                data-value="deepseek_r1">
+                                Deepseek_R1</li>
                         </ul>
                         <ul class="py-1">
                             <li class="model-option px-4 py-2 hover:bg-gray-100 cursor-pointer" data-value="gemma3">
@@ -325,7 +332,7 @@
             const selectedModelText = document.getElementById('selected-model-text');
 
             // Default model
-            let selectedModel = 'smollm2';
+            let selectedModel = 'gemma2';
 
             // Toggle dropdown
             modelDropdownButton.addEventListener('click', () => {

compose.yaml

@@ -12,6 +12,7 @@ services:
       - ./app:/app
     environment:
       ARCANA_API_KEY: "${ARCANA_API_KEY}"
+      GROQ_API_KEY: "${GROQ_API_KEY}"
       MONGO_URI: "${MONGO_URI}"
       OPENBLAS_NUM_THREADS: "${OPENBLAS_NUM_THREADS:-4}"
     healthcheck:

pyproject.toml

@@ -17,6 +17,7 @@ dependencies = [
     "huggingface-hub>=0.29",
     "pymongo>=4.11",
     "email-validator>=2.2",
+    "groq>=0.20",
 ]
 
 

uv.lock

@@ -45,6 +45,7 @@ dependencies = [
     { name = "arcana-codex" },
     { name = "email-validator" },
     { name = "fastapi-slim" },
+    { name = "groq" },
     { name = "huggingface-hub" },
     { name = "llama-cpp-python" },
     { name = "pillow" },
@@ -68,8 +69,9 @@ test = [
 [package.metadata]
 requires-dist = [
     { name = "arcana-codex", specifier = ">=0.2" },
-    { name = "email-validator", specifier = ">=2.2.0" },
+    { name = "email-validator", specifier = ">=2.2" },
     { name = "fastapi-slim", specifier = ">=0.115" },
+    { name = "groq", specifier = ">=0.20.0" },
     { name = "huggingface-hub", specifier = ">=0.29" },
     { name = "llama-cpp-python", specifier = ">=0.3" },
     { name = "pillow", specifier = ">=11.1" },
@@ -250,6 +252,15 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/91/a1/cf2472db20f7ce4a6be1253a81cfdf85ad9c7885ffbed7047fb72c24cf87/distlib-0.3.9-py2.py3-none-any.whl", hash = "sha256:47f8c22fd27c27e25a65601af709b38e4f0a45ea4fc2e710f65755fa8caaaf87", size = 468973 },
 ]
 
+[[package]]
+name = "distro"
+version = "1.9.0"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/fc/f8/98eea607f65de6527f8a2e8885fc8015d3e6f5775df186e443e0964a11c3/distro-1.9.0.tar.gz", hash = "sha256:2fa77c6fd8940f116ee1d6b94a2f90b13b5ea8d019b98bc8bafdcabcdd9bdbed", size = 60722 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/12/b3/231ffd4ab1fc9d679809f356cebee130ac7daa00d6d6f3206dd4fd137e9e/distro-1.9.0-py3-none-any.whl", hash = "sha256:7bffd925d65168f85027d8da9af6bddab658135b840670a223589bc0c8ef02b2", size = 20277 },
+]
+
 [[package]]
 name = "dnspython"
 version = "2.7.0"
@@ -304,6 +315,23 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/56/53/eb690efa8513166adef3e0669afd31e95ffde69fb3c52ec2ac7223ed6018/fsspec-2025.3.0-py3-none-any.whl", hash = "sha256:efb87af3efa9103f94ca91a7f8cb7a4df91af9f74fc106c9c7ea0efd7277c1b3", size = 193615 },
 ]
 
+[[package]]
+name = "groq"
+version = "0.20.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "anyio" },
+    { name = "distro" },
+    { name = "httpx" },
+    { name = "pydantic" },
+    { name = "sniffio" },
+    { name = "typing-extensions" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/8f/fc/e5a03586ffad7ae6c7996f388ca321a3bf8b9fa544a36a934ce4b6b44211/groq-0.20.0.tar.gz", hash = "sha256:2a201d41cae768c53d411dabcfea2333e2e138df22d909ed555ece426f1e016f", size = 121936 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/95/37/9b415df5dd1e6a685d3e8fd4e564a5e80f4f87c19d82829ad027fa2bb150/groq-0.20.0-py3-none-any.whl", hash = "sha256:c27b89903eb2b77f94ed95837ff3cadfc8c9e670953b1c5e5e2e855fea54b6c5", size = 124919 },
+]
+
 [[package]]
 name = "h11"
 version = "0.14.0"