add type for chat overrides

Author: john0isaac

src/backend/fastapi_app/api_models.py

@@ -16,9 +16,27 @@ class Message(BaseModel):
     role: AIChatRoles = AIChatRoles.USER
 
 
+class RetrievalMode(str, Enum):
+    TEXT = "text"
+    VECTORS = "vectors"
+    HYBRID = "hybrid"
+
+
+class ChatRequestOverrides(BaseModel):
+    top: int = 3
+    temperature: float = 0.3
+    retrieval_mode: RetrievalMode = RetrievalMode.HYBRID
+    use_advanced_flow: bool = True
+    prompt_template: str | None = None
+
+
+class ChatRequestContext(BaseModel):
+    overrides: ChatRequestOverrides
+
+
 class ChatRequest(BaseModel):
     messages: list[ChatCompletionMessageParam]
-    context: dict = {}
+    context: ChatRequestContext
 
 
 class ThoughtStep(BaseModel):

src/backend/fastapi_app/rag_advanced.py

@@ -7,6 +7,7 @@
 
 from fastapi_app.api_models import (
     AIChatRoles,
+    ChatRequestOverrides,
     Message,
     RAGContext,
     RetrievalResponse,
@@ -92,7 +93,7 @@ async def retrieve_and_build_context(
     async def run(
         self,
         messages: list[ChatCompletionMessageParam],
-        overrides: dict[str, Any] = {},
+        overrides: ChatRequestOverrides,
     ) -> RetrievalResponse:
         chat_params = self.get_params(messages, overrides)
 
@@ -165,7 +166,7 @@ async def run(
     async def run_stream(
         self,
         messages: list[ChatCompletionMessageParam],
-        overrides: dict[str, Any] = {},
+        overrides: ChatRequestOverrides,
     ) -> AsyncGenerator[RetrievalResponseDelta, None]:
         chat_params = self.get_params(messages, overrides)
 

src/backend/fastapi_app/rag_base.py

@@ -1,51 +1,44 @@
 import pathlib
 from abc import ABC, abstractmethod
 from collections.abc import AsyncGenerator
-from typing import Any
 
 from openai.types.chat import ChatCompletionMessageParam
-from pydantic import BaseModel
 
-from fastapi_app.api_models import (
-    RetrievalResponse,
-    RetrievalResponseDelta,
-)
+from fastapi_app.api_models import ChatRequestOverrides, RetrievalResponse, RetrievalResponseDelta
 from fastapi_app.postgres_models import Item
 
 
-class ChatParams(BaseModel):
-    top: int = 3
-    temperature: float = 0.3
+class ChatParams(ChatRequestOverrides):
+    prompt_template: str
     response_token_limit: int = 1024
     enable_text_search: bool
     enable_vector_search: bool
     original_user_query: str
     past_messages: list[ChatCompletionMessageParam]
-    prompt_template: str
 
 
 class RAGChatBase(ABC):
     current_dir = pathlib.Path(__file__).parent
     query_prompt_template = open(current_dir / "prompts/query.txt").read()
     answer_prompt_template = open(current_dir / "prompts/answer.txt").read()
 
-    def get_params(self, messages: list[ChatCompletionMessageParam], overrides: dict[str, Any]) -> ChatParams:
-        top: int = overrides.get("top", 3)
-        temperature: float = overrides.get("temperature", 0.3)
+    def get_params(self, messages: list[ChatCompletionMessageParam], overrides: ChatRequestOverrides) -> ChatParams:
         response_token_limit = 1024
-        prompt_template = overrides.get("prompt_template") or self.answer_prompt_template
+        prompt_template = overrides.prompt_template or self.answer_prompt_template
 
-        enable_text_search = overrides.get("retrieval_mode") in ["text", "hybrid", None]
-        enable_vector_search = overrides.get("retrieval_mode") in ["vectors", "hybrid", None]
+        enable_text_search = overrides.retrieval_mode in ["text", "hybrid", None]
+        enable_vector_search = overrides.retrieval_mode in ["vectors", "hybrid", None]
 
         original_user_query = messages[-1]["content"]
         if not isinstance(original_user_query, str):
             raise ValueError("The most recent message content must be a string.")
         past_messages = messages[:-1]
 
         return ChatParams(
-            top=top,
-            temperature=temperature,
+            top=overrides.top,
+            temperature=overrides.temperature,
+            retrieval_mode=overrides.retrieval_mode,
+            use_advanced_flow=overrides.use_advanced_flow,
             response_token_limit=response_token_limit,
             prompt_template=prompt_template,
             enable_text_search=enable_text_search,
@@ -67,15 +60,15 @@ async def retrieve_and_build_context(
     async def run(
         self,
         messages: list[ChatCompletionMessageParam],
-        overrides: dict[str, Any] = {},
+        overrides: ChatRequestOverrides,
     ) -> RetrievalResponse:
         raise NotImplementedError
 
     @abstractmethod
     async def run_stream(
         self,
         messages: list[ChatCompletionMessageParam],
-        overrides: dict[str, Any] = {},
+        overrides: ChatRequestOverrides,
     ) -> AsyncGenerator[RetrievalResponseDelta, None]:
         raise NotImplementedError
         if False:

src/backend/fastapi_app/rag_simple.py

@@ -1,12 +1,12 @@
 from collections.abc import AsyncGenerator
-from typing import Any
 
 from openai import AsyncAzureOpenAI, AsyncOpenAI, AsyncStream
 from openai.types.chat import ChatCompletion, ChatCompletionChunk, ChatCompletionMessageParam
 from openai_messages_token_helper import build_messages, get_token_limit
 
 from fastapi_app.api_models import (
     AIChatRoles,
+    ChatRequestOverrides,
     Message,
     RAGContext,
     RetrievalResponse,
@@ -63,7 +63,7 @@ async def retrieve_and_build_context(
     async def run(
         self,
         messages: list[ChatCompletionMessageParam],
-        overrides: dict[str, Any] = {},
+        overrides: ChatRequestOverrides,
     ) -> RetrievalResponse:
         chat_params = self.get_params(messages, overrides)
 
@@ -115,7 +115,7 @@ async def run(
     async def run_stream(
         self,
         messages: list[ChatCompletionMessageParam],
-        overrides: dict[str, Any] = {},
+        overrides: ChatRequestOverrides,
     ) -> AsyncGenerator[RetrievalResponseDelta, None]:
         chat_params = self.get_params(messages, overrides)
 

src/backend/fastapi_app/routes/api_routes.py

@@ -93,7 +93,7 @@ async def chat_handler(
     openai_chat: ChatClient,
     chat_request: ChatRequest,
 ):
-    overrides = chat_request.context.get("overrides", {})
+    overrides = chat_request.context.overrides
 
     searcher = PostgresSearcher(
         db_session=database_session,
@@ -102,7 +102,7 @@ async def chat_handler(
         embed_model=context.openai_embed_model,
         embed_dimensions=context.openai_embed_dimensions,
     )
-    if overrides.get("use_advanced_flow"):
+    if overrides.use_advanced_flow:
         run_ragchat = AdvancedRAGChat(
             searcher=searcher,
             openai_chat_client=openai_chat.client,
@@ -129,7 +129,7 @@ async def chat_stream_handler(
     openai_chat: ChatClient,
     chat_request: ChatRequest,
 ):
-    overrides = chat_request.context.get("overrides", {})
+    overrides = chat_request.context.overrides
 
     searcher = PostgresSearcher(
         db_session=database_session,
@@ -138,7 +138,7 @@ async def chat_stream_handler(
         embed_model=context.openai_embed_model,
         embed_dimensions=context.openai_embed_dimensions,
     )
-    if overrides.get("use_advanced_flow"):
+    if overrides.use_advanced_flow:
         run_ragchat = AdvancedRAGChat(
             searcher=searcher,
             openai_chat_client=openai_chat.client,

src/frontend/src/api/models.ts

@@ -1,4 +1,4 @@
-import { AIChatCompletion, AIChatCompletionDelta } from "@microsoft/ai-chat-protocol";
+import { AIChatCompletion, AIChatCompletionDelta, AIChatCompletionOperationOptions } from "@microsoft/ai-chat-protocol";
 
 export const enum RetrievalMode {
     Hybrid = "hybrid",
@@ -14,6 +14,14 @@ export type ChatAppRequestOverrides = {
     prompt_template?: string;
 };
 
+export type ChatAppRequestContext = {
+    overrides: ChatAppRequestOverrides;
+};
+
+export interface ChatAppRequestOptions extends AIChatCompletionOperationOptions {
+    context: ChatAppRequestContext
+}
+
 export type Thoughts = {
     title: string;
     description: any; // It can be any output from the api

src/frontend/src/pages/chat/Chat.tsx

@@ -4,7 +4,7 @@ import { SparkleFilled } from "@fluentui/react-icons";
 
 import styles from "./Chat.module.css";
 
-import { RetrievalMode, RAGChatCompletion, RAGChatCompletionDelta } from "../../api";
+import { RetrievalMode, RAGChatCompletion, RAGChatCompletionDelta, ChatAppRequestOptions } from "../../api";
 import { AIChatProtocolClient, AIChatMessage } from "@microsoft/ai-chat-protocol";
 import { Answer, AnswerError, AnswerLoading } from "../../components/Answer";
 import { QuestionInput } from "../../components/QuestionInput";
@@ -99,7 +99,7 @@ const Chat = () => {
                 { content: answer[1].message.content, role: "assistant" }
             ]);
             const allMessages: AIChatMessage[] = [...messages, { content: question, role: "user" }];
-            const options = {
+            const options: ChatAppRequestOptions = {
                 context: {
                     overrides: {
                         use_advanced_flow: useAdvancedFlow,