Feat/vllm endpoints (#36)

* added initial tests for vllm client class.

* added the ImprovedVllmClient class.

* added vLLM option to config, bootstrap, and api routes.

* added vLLM option to front-end provider dropdown.

* added api key to docker-compose file and removed it from the yaml config files.

* added error handling and design decision comments.

* removed example vllm client file.

* added model re-discovery on request failure.

* convert vLLM to a single client and format persona roles to be vLLM acceptable names (system/user/assistant).

* moved _clean_responses into LLMClient base class to eliminate duplication.

* changed vLLM API label color to red and also set collapsed button color from grey to the respective providers color.

docker-compose.yml

@@ -15,6 +15,7 @@ services:
       MONGODB_DATABASE: phd_advisor
       JWT_SECRET_KEY: ${JWT_SECRET_KEY:-CHANGEME-by-overriding-in-dot-env-file}
       GEMINI_API_KEY: ${GEMINI_API_KEY:-?}
+      VLLM_API_KEY: ${VLLM_API_KEY:-}
       CORS_ORIGINS: ${CORS_ORIGINS:-http://localhost:3000}
       GEMINI_MODEL: gemini-2.5-flash
       CONFIG_PATH: ${CONFIG_PATH:-/ccai/phd_config.yaml}

multi_llm_chatbot_backend/app/api/routes/provider.py

@@ -1,6 +1,8 @@
 from fastapi import APIRouter, Body, HTTPException
+from app.config import get_settings
 from app.llm.improved_gemini_client import ImprovedGeminiClient
 from app.llm.improved_ollama_client import ImprovedOllamaClient
+from app.llm.improved_vllm_client import ImprovedVllmClient
 from app.models.default_personas import get_default_personas
 from app.core.bootstrap import chat_orchestrator, llm, current_provider, available_providers
 from pydantic import BaseModel
@@ -24,6 +26,14 @@ def create_llm_client(provider: str = None):
             return ImprovedOllamaClient(model_name="llama3.2:1b")
     elif provider == "ollama":
         return ImprovedOllamaClient(model_name="llama3.2:1b")
+    elif provider == "vllm":
+        settings = get_settings()
+        if not settings.llm.vllm.api_url:
+            raise ValueError("No vLLM endpoint configured. Set llm.vllm.api_url in your config.")
+        return ImprovedVllmClient(
+            api_url=settings.llm.vllm.api_url,
+            api_key=settings.llm.vllm.api_key,
+        )
     else:
         raise ValueError(f"Unknown provider: {provider}")
 

multi_llm_chatbot_backend/app/config.py

@@ -218,9 +218,15 @@ class OllamaConfig(BaseModel):
     base_url: str = Field(default=os.getenv("OLLAMA_BASE_URL", "http://localhost:11434"))
 
 
+class VllmConfig(BaseModel):
+    api_url: str = ""
+    api_key: str = Field(default=os.getenv("VLLM_API_KEY", ""))
+
+
 class LLMConfig(BaseModel):
     gemini: GeminiConfig = GeminiConfig()
     ollama: OllamaConfig = OllamaConfig()
+    vllm: VllmConfig = VllmConfig()
 
 
 class RAGConfig(BaseModel):

multi_llm_chatbot_backend/app/core/bootstrap.py

@@ -2,19 +2,27 @@
 from app.config import get_settings
 from app.llm.improved_gemini_client import ImprovedGeminiClient
 from app.llm.improved_ollama_client import ImprovedOllamaClient
+from app.llm.improved_vllm_client import ImprovedVllmClient
 from app.core.improved_orchestrator import ImprovedChatOrchestrator
 from app.models.default_personas import get_default_personas
 
 settings = get_settings()
 
 current_provider = "gemini"
-available_providers = ["ollama", "gemini"]
+available_providers = ["ollama", "gemini", "vllm"]
 
 def create_llm_client(provider=None):
     if provider is None:
         provider = current_provider
     if provider == "gemini":
         return ImprovedGeminiClient(model_name=settings.llm.gemini.model)
+    elif provider == "vllm":
+        if not settings.llm.vllm.api_url:
+            raise ValueError("No vLLM endpoint configured. Set llm.vllm.api_url in your config.")
+        return ImprovedVllmClient(
+            api_url=settings.llm.vllm.api_url,
+            api_key=settings.llm.vllm.api_key,
+        )
     else:
         return ImprovedOllamaClient(
             model_name=settings.llm.ollama.model,

multi_llm_chatbot_backend/app/core/context_manager.py

@@ -134,10 +134,30 @@ class ContextManager:
             return self._format_for_gemini(messages, system_prompt)
         elif provider.lower() in ["ollama", "mistral"]:
             return self._format_for_ollama(messages, system_prompt)
+        elif provider.lower() == "vllm":
+            return self._format_for_vllm(messages, system_prompt)
         else:
             # Default format
             return [{"role": "system", "content": system_prompt}] + messages
-    
+
+    def _format_for_vllm(self, messages: List[dict], system_prompt: str) -> List[dict]:
+        """
+        Format messages for vLLM's OpenAI-compatible API.
+        Normalizes custom persona roles to 'assistant' since the API
+        only accepts system/user/assistant.
+        """
+        formatted = [{"role": "system", "content": system_prompt}]
+        for message in messages:
+            role = message["role"]
+            content = message["content"]
+            if role == "system":
+                continue
+            if role not in ("user", "assistant"):
+                content = f"[{role.title()} Advisor]: {content}"
+                role = "assistant"
+            formatted.append({"role": role, "content": content})
+        return formatted
+
     def _format_for_gemini(self, messages: List[dict], system_prompt: str) -> List[dict]:
         """
         Format messages for Gemini API (uses user/model roles with parts structure)

multi_llm_chatbot_backend/app/llm/improved_gemini_client.py

@@ -1,6 +1,5 @@
 import httpx
 import os
-import re
 from typing import List
 from app.llm.llm_client import LLMClient
 from app.core.context_manager import get_context_manager
@@ -121,11 +120,3 @@ class ImprovedGeminiClient(LLMClient):
         except Exception as e:
             logger.error(f"Unexpected error in Gemini client: {str(e)}")
             return "I encountered an unexpected error. Please try again."
-    
-    def _clean_response(self, response: str) -> str:
-        """Clean up response text, preserving Markdown formatting."""
-        response = response.replace("\r\n", "\n").replace("\r", "\n")
-        lines = [ln.rstrip() for ln in response.split("\n")]
-        response = re.sub(r"\n{3,}", "\n\n", "\n".join(lines)).strip()
-
-        return response

multi_llm_chatbot_backend/app/llm/improved_vllm_client.py

@@ -0,0 +1,67 @@
+from typing import List
+from openai import AsyncOpenAI, APIConnectionError, APIStatusError
+from app.llm.llm_client import LLMClient
+from app.core.context_manager import get_context_manager
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+class ImprovedVllmClient(LLMClient):
+    def __init__(self, api_url: str, api_key: str, model_name: str = None):
+        self.api_url = api_url
+        self.api_key = api_key
+        self.model_name = model_name
+        self.client = AsyncOpenAI(
+            base_url=f"{api_url}/v1",
+            api_key=api_key,
+            timeout=30.0,
+        )
+        self.context_manager = get_context_manager()
+
+    async def refresh_model(self):
+        """Query the vLLM endpoint to discover the currently loaded model."""
+        models = await self.client.models.list()
+        if not models.data:
+            raise ValueError("No models available at the vLLM endpoint")
+        self.model_name = models.data[0].id
+
+    async def generate(self, system_prompt: str, context: List[dict],
+                       temperature: float, max_tokens: int,
+                       response_mime_type: str = None) -> str:
+        try:
+            context_window = self.context_manager.prepare_context_for_llm(
+                messages=context,
+                system_prompt=system_prompt,
+                llm_provider="vllm"
+            )
+
+            logger.debug(f"Context prepared: {len(context_window.messages)} messages, "
+                        f"~{context_window.total_tokens} tokens, truncated={context_window.truncated}")
+
+            if not self.model_name:
+                await self.refresh_model()
+
+            response = await self.client.chat.completions.create(
+                model=self.model_name,
+                messages=context_window.messages,
+                temperature=temperature,
+                max_tokens=max_tokens,
+            )
+
+            text = response.choices[0].message.content.strip()
+            return self._clean_response(text)
+
+        except APIConnectionError:
+            logger.error(f"Unable to connect to vLLM at {self.api_url}")
+            return "I'm unable to connect to the AI service. Please ensure the vLLM endpoint is available."
+        except APIStatusError as e:
+            logger.error(f"vLLM API error: {e.status_code} - {e.message}")
+            if e.status_code == 404:
+                logger.info("Model not found, will re-discover on next request")
+                self.model_name = None
+            return "The AI service encountered an error. Please try again."
+        except Exception as e:
+            logger.error(f"Unexpected error in vLLM client: {str(e)}")
+            return "I encountered an unexpected error. Please try again."
+

multi_llm_chatbot_backend/app/llm/llm_client.py

@@ -1,5 +1,6 @@
 from abc import ABC, abstractmethod
 from typing import List
+import re
 
 class LLMClient(ABC):
     """Abstract base class for all LLM clients"""
@@ -19,4 +20,11 @@ class LLMClient(ABC):
         Returns:
             str: The generated response text
         """
-        pass
+        pass
+
+    def _clean_response(self, response: str) -> str:
+        """Clean up response text, preserving Markdown formatting."""
+        response = response.replace("\r\n", "\n").replace("\r", "\n")
+        lines = [ln.rstrip() for ln in response.split("\n")]
+        response = re.sub(r"\n{3,}", "\n\n", "\n".join(lines)).strip()
+        return response

multi_llm_chatbot_backend/app/tests/unit/test_vllm_client.py

@@ -0,0 +1,175 @@
+import asyncio
+import unittest
+from unittest.mock import AsyncMock, MagicMock, patch
+
+from openai import APIConnectionError, APIStatusError
+
+from app.llm.improved_vllm_client import ImprovedVllmClient
+
+
+FAKE_URL = "https://fake.example.com/vllm0"
+FAKE_KEY = "test-key"
+
+
+def _make_completion_mock(content="Response"):
+    """Build a mock that looks like an OpenAI ChatCompletion."""
+    mock_message = MagicMock()
+    mock_message.content = content
+    mock_choice = MagicMock()
+    mock_choice.message = mock_message
+    return MagicMock(choices=[mock_choice])
+
+
+@patch("app.llm.improved_vllm_client.get_context_manager")
+@patch("app.llm.improved_vllm_client.AsyncOpenAI")
+class TestImprovedVllmClient(unittest.TestCase):
+
+    # ------------------------------------------------------------------
+    # Construction
+    # ------------------------------------------------------------------
+
+    def test_constructor_stores_attributes(self, MockAsyncOpenAI, mock_get_ctx):
+        client = ImprovedVllmClient(
+            api_url=FAKE_URL, api_key=FAKE_KEY, model_name="test-model",
+        )
+        self.assertEqual(client.api_url, FAKE_URL)
+        self.assertEqual(client.api_key, FAKE_KEY)
+        self.assertEqual(client.model_name, "test-model")
+
+    def test_constructor_defaults_model_to_none(self, MockAsyncOpenAI, mock_get_ctx):
+        client = ImprovedVllmClient(api_url=FAKE_URL, api_key=FAKE_KEY)
+        self.assertIsNone(client.model_name)
+
+    # ------------------------------------------------------------------
+    # Model discovery
+    # ------------------------------------------------------------------
+
+    def test_refresh_model_discovers_model(self, MockAsyncOpenAI, mock_get_ctx):
+        client = ImprovedVllmClient(api_url=FAKE_URL, api_key=FAKE_KEY)
+
+        mock_model = MagicMock()
+        mock_model.id = "discovered-model"
+        client.client.models.list = AsyncMock(
+            return_value=MagicMock(data=[mock_model])
+        )
+
+        asyncio.run(client.refresh_model())
+        self.assertEqual(client.model_name, "discovered-model")
+
+    # ------------------------------------------------------------------
+    # generate – happy path
+    # ------------------------------------------------------------------
+
+    def test_generate_returns_cleaned_response(self, MockAsyncOpenAI, mock_get_ctx):
+        client = ImprovedVllmClient(
+            api_url=FAKE_URL, api_key=FAKE_KEY, model_name="test-model",
+        )
+        client.client.chat.completions.create = AsyncMock(
+            return_value=_make_completion_mock("  Here is my response.  ")
+        )
+
+        result = asyncio.run(client.generate(
+            system_prompt="You are helpful.",
+            context=[{"role": "user", "content": "Hello"}],
+            temperature=0.7,
+            max_tokens=100,
+        ))
+        self.assertEqual(result, "Here is my response.")
+
+    def test_generate_auto_discovers_model_when_none(self, MockAsyncOpenAI, mock_get_ctx):
+        client = ImprovedVllmClient(
+            api_url=FAKE_URL, api_key=FAKE_KEY, model_name=None,
+        )
+
+        mock_model = MagicMock()
+        mock_model.id = "auto-discovered"
+        client.client.models.list = AsyncMock(
+            return_value=MagicMock(data=[mock_model])
+        )
+        client.client.chat.completions.create = AsyncMock(
+            return_value=_make_completion_mock()
+        )
+
+        asyncio.run(client.generate(
+            system_prompt="Test",
+            context=[{"role": "user", "content": "Hi"}],
+            temperature=0.5,
+            max_tokens=50,
+        ))
+
+        client.client.models.list.assert_called_once()
+        self.assertEqual(client.model_name, "auto-discovered")
+
+    # ------------------------------------------------------------------
+    # generate – error handling
+    # ------------------------------------------------------------------
+
+    def test_generate_handles_connection_error(self, MockAsyncOpenAI, mock_get_ctx):
+        client = ImprovedVllmClient(
+            api_url=FAKE_URL, api_key=FAKE_KEY, model_name="test-model",
+        )
+        client.client.chat.completions.create = AsyncMock(
+            side_effect=APIConnectionError(request=MagicMock())
+        )
+
+        result = asyncio.run(client.generate(
+            system_prompt="Test",
+            context=[{"role": "user", "content": "Hi"}],
+            temperature=0.5,
+            max_tokens=50,
+        ))
+        self.assertIn("unable to connect", result.lower())
+
+    def test_generate_handles_status_error(self, MockAsyncOpenAI, mock_get_ctx):
+        client = ImprovedVllmClient(
+            api_url=FAKE_URL, api_key=FAKE_KEY, model_name="test-model",
+        )
+        mock_response = MagicMock()
+        mock_response.status_code = 500
+        client.client.chat.completions.create = AsyncMock(
+            side_effect=APIStatusError(
+                message="Server error", response=mock_response, body=None,
+            )
+        )
+
+        result = asyncio.run(client.generate(
+            system_prompt="Test",
+            context=[{"role": "user", "content": "Hi"}],
+            temperature=0.5,
+            max_tokens=50,
+        ))
+        self.assertIn("error", result.lower())
+
+    def test_generate_clears_model_on_404(self, MockAsyncOpenAI, mock_get_ctx):
+        client = ImprovedVllmClient(
+            api_url=FAKE_URL, api_key=FAKE_KEY, model_name="stale-model",
+        )
+        mock_response = MagicMock()
+        mock_response.status_code = 404
+        client.client.chat.completions.create = AsyncMock(
+            side_effect=APIStatusError(
+                message="Model not found", response=mock_response, body=None,
+            )
+        )
+
+        asyncio.run(client.generate(
+            system_prompt="Test",
+            context=[{"role": "user", "content": "Hi"}],
+            temperature=0.5,
+            max_tokens=50,
+        ))
+        self.assertIsNone(client.model_name)
+
+    # ------------------------------------------------------------------
+    # _clean_response
+    # ------------------------------------------------------------------
+
+    def test_clean_response_normalizes_whitespace(self, MockAsyncOpenAI, mock_get_ctx):
+        client = ImprovedVllmClient(
+            api_url=FAKE_URL, api_key=FAKE_KEY, model_name="test-model",
+        )
+        dirty = "Line one.\r\n\r\n\r\n\r\nLine two.  "
+        cleaned = client._clean_response(dirty)
+        self.assertNotIn("\r", cleaned)
+        self.assertNotIn("\n\n\n", cleaned)
+        self.assertEqual(cleaned, "Line one.\n\nLine two.")

multi_llm_chatbot_backend/requirements.txt

@@ -5,6 +5,7 @@ python-multipart
 
 # HTTP client for LLM APIs
 httpx
+openai~=2.30
 
 # Document processing
 PyPDF2

phd-advisor-frontend/src/components/ProviderDropdown.js

@@ -1,6 +1,6 @@
 // src/components/ProviderDropdown.js
 import React, { useState, useRef, useEffect } from 'react';
-import { ChevronDown, Cpu, Cloud, Loader2 } from 'lucide-react';
+import { ChevronDown, Cpu, Cloud, Server, Loader2 } from 'lucide-react';
 import { useTheme } from '../contexts/ThemeContext';
 
 const ProviderDropdown = ({ currentProvider, onProviderChange, isLoading = false }) => {
@@ -22,6 +22,13 @@ const ProviderDropdown = ({ currentProvider, onProviderChange, isLoading = false
       description: 'Local LLM via Ollama',
       icon: Cpu,
       badge: 'Local'
+    },
+    {
+      id: 'vllm',
+      name: 'vLLM',
+      description: 'vLLM inference endpoint',
+      icon: Server,
+      badge: 'API'
     }
   ];
 
@@ -69,7 +76,7 @@ const ProviderDropdown = ({ currentProvider, onProviderChange, isLoading = false
           )}
           <div className="provider-info">
             <span className="provider-name">{currentProviderInfo?.name || 'Unknown'}</span>
-            <span className="provider-badge">{currentProviderInfo?.badge}</span>
+            <span className={`provider-badge ${currentProvider}`}>{currentProviderInfo?.badge}</span>
           </div>
         </div>
         <ChevronDown 

phd-advisor-frontend/src/styles/ChatPage.css

@@ -882,6 +882,10 @@
   letter-spacing: 0.5px;
 }
 
+.provider-badge.gemini { color: #4285f4; }
+.provider-badge.ollama { color: #10b981; }
+.provider-badge.vllm   { color: #ef4444; }
+
 .clarification-message-container {
   display: flex;
   justify-content: center;
@@ -1112,6 +1116,11 @@
   color: #10b981;
 }
 
+.provider-option-badge.vllm {
+  background: rgba(239, 68, 68, 0.15);
+  color: #ef4444;
+}
+
 .provider-option-description {
   font-size: 11px;
   color: var(--text-tertiary);

phd_config.yaml

@@ -144,6 +144,8 @@ llm:
     model: "gemini-2.5-flash"
   ollama:
     model: "llama3.2:1b"
+  vllm:
+    api_url: https://rtx6000blackwell-1.neonaiservices2.com/vllm0
 
 rag:
   embedding_model: "all-MiniLM-L6-v2"

undergrad_config.yaml

@@ -416,6 +416,8 @@ llm:
     model: "gemini-2.5-flash"
   ollama:
     model: "llama3.2:1b"
+  vllm:
+    api_url: https://rtx6000blackwell-1.neonaiservices2.com/vllm0
 
 rag:
   embedding_model: "all-MiniLM-L6-v2"