deployment

dynamodb_helper.py

@@ -6,6 +6,9 @@ from botocore.exceptions import ClientError
 from datetime import datetime, timezone
 from uuid import uuid4
 from decimal import Decimal
+from dotenv import load_dotenv
+
+load_dotenv()
 
 AWS_REGION = os.getenv("AWS_REGION", "us-east-1")
 AWS_ACCESS_KEY = os.getenv("AWS_ACCESS_KEY", None)
@@ -14,25 +17,29 @@ DYNAMODB_ENDPOINT = os.getenv("DYNAMODB_ENDPOINT", None)
 DDB_TABLE = os.getenv("DDB_TABLE", "chatbot-conversations")
 USE_LOCAL_DDB = os.getenv("USE_LOCAL_DDB", "false").lower() == "true"
 
+
 def get_dynamodb_client():
-    if USE_LOCAL_DDB: # only for local testing with DynamoDB Local
+    if USE_LOCAL_DDB:  # only for local testing with DynamoDB Local
         return boto3.resource(
             "dynamodb",
             endpoint_url=DYNAMODB_ENDPOINT,
             region_name=AWS_REGION,
             aws_access_key_id="fake",
-            aws_secret_access_key="fake"
+            aws_secret_access_key="fake",
+        )
+    else:  # production AWS DynamoDB
+        return boto3.resource(
+            "dynamodb",
+            region_name=AWS_REGION,
+            aws_access_key_id=AWS_ACCESS_KEY,
+            aws_secret_access_key=AWS_SECRET_ACCESS_KEY,
         )
-    else: # production AWS DynamoDB
-        return boto3.resource("dynamodb",
-                              region_name=AWS_REGION, 
-                              aws_access_key_id=AWS_ACCESS_KEY, 
-                              aws_secret_access_key=AWS_SECRET_ACCESS_KEY
-                            )
+
 
 dynamodb = get_dynamodb_client()
 table = None
 
+
 def create_table_if_not_exists(dynamodb):
     global table
     client = dynamodb.meta.client
@@ -55,29 +62,29 @@ def create_table_if_not_exists(dynamodb):
             TableName=DDB_TABLE,
             KeySchema=[
                 {"AttributeName": "PK", "KeyType": "HASH"},
-                {"AttributeName": "SK", "KeyType": "RANGE"}
+                {"AttributeName": "SK", "KeyType": "RANGE"},
             ],
             AttributeDefinitions=[
                 {"AttributeName": "PK", "AttributeType": "S"},
                 {"AttributeName": "SK", "AttributeType": "S"},
                 {"AttributeName": "GSI1_PK", "AttributeType": "S"},
-                {"AttributeName": "GSI1_SK", "AttributeType": "S"}
+                {"AttributeName": "GSI1_SK", "AttributeType": "S"},
             ],
             GlobalSecondaryIndexes=[
                 {
                     "IndexName": "GSI1",
                     "KeySchema": [
                         {"AttributeName": "GSI1_PK", "KeyType": "HASH"},
-                        {"AttributeName": "GSI1_SK", "KeyType": "RANGE"}
+                        {"AttributeName": "GSI1_SK", "KeyType": "RANGE"},
                     ],
                     "Projection": {"ProjectionType": "ALL"},
                     "ProvisionedThroughput": {
                         "ReadCapacityUnits": 5,
-                        "WriteCapacityUnits": 5
+                        "WriteCapacityUnits": 5,
                     },
                 }
             ],
-            BillingMode='PAY_PER_REQUEST'
+            BillingMode="PAY_PER_REQUEST",
             # ProvisionedThroughput={
             #     "ReadCapacityUnits": 5,
             #     "WriteCapacityUnits": 5
@@ -97,8 +104,10 @@ def iso_ts():
     # Return the current timestamp in ISO 8601 format
     return datetime.now(timezone.utc).isoformat()
 
+
 table = create_table_if_not_exists(dynamodb)
 
+
 def convert_floats(obj):
     if isinstance(obj, float):
         return Decimal(str(obj))
@@ -109,6 +118,7 @@ def convert_floats(obj):
     else:
         return obj
 
+
 def log_event(user_id, session_id, data):
     """
     Log conversation data to DynamoDB table.
@@ -125,12 +135,12 @@ def log_event(user_id, session_id, data):
     item = {
         "PK": f"SESSION#{session_id}",
         "SK": f"TS#{ts}#{uuid4().hex}",
-        'user_id': user_id,
+        "user_id": user_id,
         "GSI1_PK": f"USER#{user_id}",
         "GSI1_SK": f"TS#{ts}",
-        'session_id': session_id,
-        'timestamp': ts,
-        'data': convert_floats(data)
+        "session_id": session_id,
+        "timestamp": ts,
+        "data": convert_floats(data),
     }
     print(f"Logging conversation: {item}")
     try:

main.py

@@ -4,14 +4,12 @@ from contextlib import asynccontextmanager
 
 from pathlib import Path
 
-from typing import List, Literal, Optional, Tuple, Dict, Any
-from datetime import datetime, timezone
+from typing import AsyncGenerator, List, Literal, Optional, Tuple, Dict, Any, Generator
 
 from dotenv import load_dotenv
-load_dotenv()
 
 from fastapi import FastAPI, Request, BackgroundTasks
-from fastapi.responses import HTMLResponse, JSONResponse
+from fastapi.responses import HTMLResponse, JSONResponse, StreamingResponse
 from fastapi.staticfiles import StaticFiles
 from fastapi.templating import Jinja2Templates
 
@@ -19,7 +17,7 @@ from pydantic import BaseModel
 from dynamodb_helper import log_event
 
 from huggingface_hub import InferenceClient
-from openai import OpenAI
+from openai import AsyncOpenAI
 from google import genai
 
 
@@ -28,13 +26,15 @@ from langchain_core.messages import HumanMessage, AIMessage, SystemMessage
 from champ.prompts import DEFAULT_SYSTEM_PROMPT
 from champ.service import ChampService
 
+load_dotenv()
+
 # -------------------- Config --------------------
 BASE_DIR = Path(__file__).resolve().parent
 
 MODEL_MAP = {
     "champ": "champ-model/placeholder",
     "openai": "gpt-5-nano-2025-08-07",
-    "google": "gemini-2.5-flash-lite"
+    "google": "gemini-2.5-flash-lite",
 }
 
 HF_TOKEN = os.getenv("HF_TOKEN") or os.getenv("HF_API_TOKEN")
@@ -57,13 +57,13 @@ if GEMINI_API_KEY is None:
     )
 
 hf_client = InferenceClient(token=HF_TOKEN)
-openai_client = OpenAI(api_key=OPENAI_API_KEY) if OPENAI_API_KEY else None
+openai_client = AsyncOpenAI(api_key=OPENAI_API_KEY) if OPENAI_API_KEY else None
 gemini_client = genai.Client(api_key=GEMINI_API_KEY) if GEMINI_API_KEY else None
 
 
-
 # Max history messages to keep for context
-MAX_HISTORY = 20 
+MAX_HISTORY = 20
+
 
 class ChatMessage(BaseModel):
     role: Literal["user", "assistant", "system"]
@@ -105,15 +105,22 @@ def convert_messages_langchain(messages: List[ChatMessage]):
             list_chatmessages.append(SystemMessage(content=m.content))
     return list_chatmessages
 
+
 champ = ChampService(base_dir=BASE_DIR, hf_token=HF_TOKEN)
 
-def _call_openai(model_id: str, msgs: list[dict], temperature: float) -> str:
-    resp = openai_client.responses.create(
-        model=model_id,
-        input=msgs,
-        # no temperature for GPT-5 reasoning models  
+
+async def _call_openai(
+    model_id: str, msgs: list[dict], temperature: float
+) -> AsyncGenerator[str, None]:
+    # We are streaming the output because the model answers tend to be very long and slow to generate
+    stream = await openai_client.responses.create(
+        model=model_id, input=msgs, stream=True
     )
-    return (resp.output_text or "").strip()
+
+    async for chunk in stream:
+        if chunk.type == "response.output_text.delta":
+            yield chunk.delta
+
 
 def _call_gemini(model_id: str, msgs: list[dict], temperature: float) -> str:
     transcript = []
@@ -130,7 +137,12 @@ def _call_gemini(model_id: str, msgs: list[dict], temperature: float) -> str:
     )
     return (resp.text or "").strip()
 
-def _call_hf_client(model_id: str, msgs: list[dict], temperature: float,) -> str:
+
+def _call_hf_client(
+    model_id: str,
+    msgs: list[dict],
+    temperature: float,
+) -> str:
     resp = hf_client.chat.completions.create(
         model=model_id,
         messages=msgs,
@@ -142,7 +154,10 @@ def _call_hf_client(model_id: str, msgs: list[dict], temperature: float,) -> str
     except Exception:
         return str(resp)
 
-def call_llm(req: ChatRequest) -> Tuple[str, Dict[str, Any]]:
+
+def call_llm(
+    req: ChatRequest,
+) -> AsyncGenerator[str, None] | Tuple[str, Dict[str, Any]]:
     if req.model_type == "champ":
         msgs = convert_messages_langchain(req.messages)
         reply, triage_meta = champ.invoke(msgs)
@@ -155,7 +170,7 @@ def call_llm(req: ChatRequest) -> Tuple[str, Dict[str, Any]]:
     msgs = convert_messages(req.messages)
 
     if req.model_type == "openai":
-        return _call_openai(model_id, msgs, req.temperature), {}
+        return _call_openai(model_id, msgs, req.temperature)
 
     if req.model_type == "google":
         return _call_gemini(model_id, msgs, req.temperature), {}
@@ -173,6 +188,7 @@ def call_llm(req: ChatRequest) -> Tuple[str, Dict[str, Any]]:
 #     }
 #     conversations_collection.insert_one(record)
 
+
 # -------------------- FastAPI setup --------------------
 @asynccontextmanager
 async def lifespan(app: FastAPI):
@@ -180,6 +196,7 @@ async def lifespan(app: FastAPI):
     print("CHAMP RAG + agent initialized.")
     yield
 
+
 app = FastAPI(lifespan=lifespan)
 app.mount("/static", StaticFiles(directory="static"), name="static")
 templates = Jinja2Templates(directory="templates")
@@ -197,7 +214,34 @@ async def chat_endpoint(payload: ChatRequest, background_tasks: BackgroundTasks)
 
     try:
         loop = asyncio.get_running_loop()
-        reply, triage_meta = await loop.run_in_executor(None, call_llm, payload)
+        result = await loop.run_in_executor(None, call_llm, payload)
+
+        if isinstance(result, AsyncGenerator):
+
+            async def logging_wrapper():
+                reply = ""
+                async for token in result:
+                    reply += token
+                    yield token
+
+                background_tasks.add_task(
+                    log_event,
+                    user_id=payload.user_id,
+                    session_id=payload.session_id,
+                    data={
+                        "model_type": payload.model_type,
+                        "consent": payload.consent,
+                        "temperature": payload.temperature,
+                        "messages": payload.messages[-1].dict(),
+                        "reply": reply,
+                        "triage_meta": {},
+                    },
+                )
+
+            return StreamingResponse(logging_wrapper(), media_type="text/event-stream")
+
+        reply, triage_meta = result
+
     except Exception as e:
         background_tasks.add_task(
             log_event,
@@ -211,7 +255,6 @@ async def chat_endpoint(payload: ChatRequest, background_tasks: BackgroundTasks)
                 "messages": payload.messages[-1].dict(),
             },
         )
-        return JSONResponse({"error": str(e)}, status_code=500)
 
     background_tasks.add_task(
         log_event,
@@ -226,4 +269,4 @@ async def chat_endpoint(payload: ChatRequest, background_tasks: BackgroundTasks)
             **(triage_meta or {}),
         },
     )
-    return {"reply": reply}
+    return {"reply": reply}

requirements.txt

@@ -124,3 +124,4 @@ websockets==15.0.1
 xxhash==3.6.0
 yarl==1.22.0
 zstandard==0.25.0
+pytz==2025.2

static/app.js

@@ -15,14 +15,21 @@ const consentCheckbox = document.getElementById('consentCheckbox');
 const consentBtn = document.getElementById('consentBtn');
 
 // Local in-browser chat history
-let messages = [];
+// We store for each model its chat history.
+// We store the temperature of the google model as it can change.
+const modelChats = {};
+modelChats["champ"] = {"messages": []}
+modelChats["openai"] = {"messages": []}
+modelChats["google"] = {"messages": [], "temperature": 0.2}
+
 let consentGranted = false;
 let sessionId = 'session-' + crypto.randomUUID(); // Unique session ID, generated once per page load
 document.body.classList.add('no-scroll');
 
 function renderMessages() {
   chatWindow.innerHTML = '';
-  messages.forEach((m) => {
+  const modelType = systemPresetSelect.value;
+  modelChats[modelType]["messages"].forEach((m) => {
     const bubble = document.createElement('div');
     bubble.classList.add(
       'msg-bubble',
@@ -58,7 +65,8 @@ async function sendMessage() {
   if (!text) return;
 
   // Add user message locally
-  messages.push({ role: 'user', content: text });
+  const modelType = systemPresetSelect.value;
+  modelChats[modelType]["messages"].push({ role: 'user', content: text });
   renderMessages();
   userInput.value = '';
 
@@ -68,12 +76,11 @@ async function sendMessage() {
   const temperature = parseFloat(tempSlider.value);
   // const maxTokens = parseInt(maxTokensSlider.value, 10);
   // const systemPrompt = systemPresetSelect.value;
-  const modelType = systemPresetSelect.value;
 
   const payload = {
     user_id: getMachineId(),
     session_id: sessionId,
-    messages: messages.map((m) => ({ role: m.role, content: m.content })),
+    messages:  modelChats[modelType]["messages"].map((m) => ({ role: m.role, content: m.content })),
     temperature,
     // max_new_tokens: maxTokens,
     model_type: modelType,
@@ -87,17 +94,39 @@ async function sendMessage() {
       body: JSON.stringify(payload),
     });
 
-    const data = await res.json();
-
     if (!res.ok) {
       statusEl.textContent = data.error || 'Error from server.';
       statusEl.className = 'status status-error';
       return;
     }
 
-    const reply = data.reply || '(No reply)';
-    messages.push({ role: 'assistant', content: reply });
-    renderMessages();
+    const contentType = res.headers.get('content-type');
+
+    if (contentType && contentType.includes('application/json')) {
+      // Batch response
+      const data = await res.json();
+
+      const reply = data.reply || '(No reply)';
+      modelChats[modelType]["messages"].push({ role: 'assistant', content: reply });
+      renderMessages();
+    } else {
+      // Streaming response
+      const assistantMessage = { role: 'assistant', content: '' };
+      modelChats[modelType]["messages"].push(assistantMessage);
+
+      const reader = res.body.getReader();
+      const decoder = new TextDecoder();
+      let done = false;
+
+      while (!done) {
+        const { value, done: readerDone } = await reader.read();
+        done = readerDone;
+        const chunk = decoder.decode(value, { stream: true });
+        assistantMessage.content += chunk;
+        renderMessages();
+      }
+    }
+
 
     statusEl.textContent = 'Ready';
     statusEl.className = 'status status-ok';
@@ -113,8 +142,16 @@ function resetSession() {
 }
 
 function clearConversation() {
+  const modelType = systemPresetSelect.value;
+
   resetSession();
-  messages = [];
+  modelChats[modelType]["messages"] = [];
+  // If the model is google, we also have to clear the temperature
+  if (modelType === "google") {
+    modelChats["google"]["temperature"] = 0.2;
+    tempSlider.value = 0.2
+    updateSlidersUI();
+  }
   renderMessages();
   statusEl.textContent = 'Conversation cleared. Start a new chat!';
   statusEl.className = 'status status-ok';
@@ -147,15 +184,19 @@ userInput.addEventListener('keydown', (e) => {
 });
 
 tempSlider.addEventListener('input', () => {
-  if (!tempSlider.disabled) updateSlidersUI();
+  if (!tempSlider.disabled) {
+    updateSlidersUI();
+    const modelType = systemPresetSelect.value;
+    modelChats[modelType]["temperature"] = tempSlider.value;
+  }
 });
 // maxTokensSlider.addEventListener("input", updateSlidersUI);
 clearBtn.addEventListener('click', clearConversation);
 
 systemPresetSelect.addEventListener('change', () => {
   updateTempControlForModel();   // 👈 add this
-  clearConversation();
-  statusEl.textContent = 'Model changed. History cleared.';
+  renderMessages();
+  statusEl.textContent = 'Model changed.';
   statusEl.className = 'status status-ok';
 });
 
@@ -179,6 +220,7 @@ function updateTempControlForModel() {
     // Enable slider for other models
     tempSlider.disabled = false;
     tempSlider.classList.remove('disabled');
+    tempSlider.value = modelChats[model]["temperature"];
     updateSlidersUI(); // refresh displayed value
   }
 }