Add num_consistency_samples arg (#31)

Author: huiwengoh

services/chat-backend/pyproject.toml

@@ -57,4 +57,4 @@ show_missing = true
 default-groups = ["dev"]
 
 [tlm-core]
-version = "v0.0.47"
+version = "v0.0.48"

services/chat-backend/src/routers/chat.py

@@ -4,7 +4,7 @@
 from fastapi import APIRouter
 from litellm.types.utils import ModelResponse
 from tlm.chat import achat, acost
-from tlm.config.types import ModelProvider
+from tlm.config.types import ModelProvider, TLMOptions
 from tlm.types import Message, TokenCost, TLMScore
 
 from src.schemas.chat import (
@@ -111,6 +111,7 @@ async def _handle_request(
         model_provider=model_provider,
         custom_eval_config=input.options.custom_eval_criteria,
         completion=completion,
+        tlm_options=construct_tlm_options(input),
     )
 
     # TODO:properly handle multiple completions
@@ -133,6 +134,12 @@ async def _handle_request(
     return result
 
 
+def construct_tlm_options(input: ChatInput) -> TLMOptions:
+    return TLMOptions(
+        num_consistency_samples=input.options.num_consistency_samples,
+    )
+
+
 def maybe_add_log(tlm_score: TLMScore, log: list[ChatLogField]) -> dict[str, Any]:
     log_dict = {}
     if ChatLogField.PER_FIELD_SCORE in log:
@@ -150,8 +157,12 @@ async def chat_costs(
     input: ChatInput,
 ) -> TokenCost:
     reference_config = get_reference_config(input)
-    messages = [Message.model_validate(message.model_dump()) for message in input.messages]
-    model_provider = input.options.model_provider or model_registry.get(input.options.model)
+    messages = [
+        Message.model_validate(message.model_dump()) for message in input.messages
+    ]
+    model_provider = input.options.model_provider or model_registry.get(
+        input.options.model
+    )
 
     return await acost(
         messages,

services/chat-backend/src/schemas/chat/input.py

@@ -55,6 +55,7 @@ class ChatTlmInput(BaseModel):
 
     model_provider: ModelProviderInput | None = None
     custom_eval_criteria: list[CustomEvalConfig] = Field(default_factory=list)
+    num_consistency_samples: int | None = None
 
     log: list[ChatLogField] = Field(
         default_factory=lambda: [