Typing improvements

Author: pamelafox

requirements-dev.txt

@@ -6,3 +6,4 @@ pip-compile-cross-platform
 pytest
 pytest-cov
 pytest-asyncio
+mypy

src/fastapi_app/openai_clients.py

@@ -8,24 +8,31 @@
 
 
 async def create_openai_chat_client(azure_credential):
+    openai_chat_client: openai.AsyncAzureOpenAI | openai.AsyncOpenAI
     OPENAI_CHAT_HOST = os.getenv("OPENAI_CHAT_HOST")
     if OPENAI_CHAT_HOST == "azure":
-        client_args = {}
+        api_version = os.environ["AZURE_OPENAI_VERSION"]
+        azure_endpoint = os.environ["AZURE_OPENAI_ENDPOINT"]
+        azure_deployment = os.environ["AZURE_OPENAI_EMBED_DEPLOYMENT"]
         if api_key := os.getenv("AZURE_OPENAI_KEY"):
             logger.info("Authenticating to Azure OpenAI using API key...")
-            client_args["api_key"] = api_key
+            openai_chat_client = openai.AsyncAzureOpenAI(
+                api_version=api_version,
+                azure_endpoint=azure_endpoint,
+                azure_deployment=azure_deployment,
+                api_key=api_key,
+            )
         else:
             logger.info("Authenticating to Azure OpenAI using Azure Identity...")
             token_provider = azure.identity.get_bearer_token_provider(
                 azure_credential, "https://cognitiveservices.azure.com/.default"
             )
-            client_args["azure_ad_token_provider"] = token_provider
-        openai_chat_client = openai.AsyncAzureOpenAI(
-            api_version=os.getenv("AZURE_OPENAI_VERSION"),
-            azure_endpoint=os.getenv("AZURE_OPENAI_ENDPOINT"),
-            azure_deployment=os.getenv("AZURE_OPENAI_CHAT_DEPLOYMENT"),
-            **client_args,
-        )
+            openai_chat_client = openai.AsyncAzureOpenAI(
+                api_version=api_version,
+                azure_endpoint=azure_endpoint,
+                azure_deployment=azure_deployment,
+                azure_ad_token_provider=token_provider,
+            )
         openai_chat_model = os.getenv("AZURE_OPENAI_CHAT_MODEL")
     elif OPENAI_CHAT_HOST == "ollama":
         logger.info("Authenticating to OpenAI using Ollama...")
@@ -43,24 +50,32 @@ async def create_openai_chat_client(azure_credential):
 
 
 async def create_openai_embed_client(azure_credential):
+    openai_embed_client: openai.AsyncAzureOpenAI | openai.AsyncOpenAI
     OPENAI_EMBED_HOST = os.getenv("OPENAI_EMBED_HOST")
     if OPENAI_EMBED_HOST == "azure":
-        client_args = {}
+        api_version = os.environ["AZURE_OPENAI_VERSION"]
+        azure_endpoint = os.environ["AZURE_OPENAI_ENDPOINT"]
+        azure_deployment = os.environ["AZURE_OPENAI_EMBED_DEPLOYMENT"]
         if api_key := os.getenv("AZURE_OPENAI_KEY"):
             logger.info("Authenticating to Azure OpenAI using API key...")
-            client_args["api_key"] = api_key
+            openai_embed_client = openai.AsyncAzureOpenAI(
+                api_version=api_version,
+                azure_endpoint=azure_endpoint,
+                azure_deployment=azure_deployment,
+                api_key=api_key,
+            )
         else:
             logger.info("Authenticating to Azure OpenAI using Azure Identity...")
             token_provider = azure.identity.get_bearer_token_provider(
                 azure_credential, "https://cognitiveservices.azure.com/.default"
             )
-            client_args["azure_ad_token_provider"] = token_provider
-        openai_embed_client = openai.AsyncAzureOpenAI(
-            api_version=os.getenv("AZURE_OPENAI_VERSION"),
-            azure_endpoint=os.getenv("AZURE_OPENAI_ENDPOINT"),
-            azure_deployment=os.getenv("AZURE_OPENAI_EMBED_DEPLOYMENT"),
-            **client_args,
-        )
+            openai_embed_client = openai.AsyncAzureOpenAI(
+                api_version=api_version,
+                azure_endpoint=azure_endpoint,
+                azure_deployment=azure_deployment,
+                azure_ad_token_provider=token_provider,
+            )
+
         openai_embed_model = os.getenv("AZURE_OPENAI_EMBED_MODEL")
         openai_embed_dimensions = os.getenv("AZURE_OPENAI_EMBED_DIMENSIONS")
     else:

src/fastapi_app/rag_advanced.py

@@ -32,13 +32,15 @@ def __init__(
         self.answer_prompt_template = open(current_dir / "prompts/answer.txt").read()
 
     async def run(
-        self, messages: list[dict], overrides: dict[str, Any] = {}
+        self, messages: list[ChatCompletionMessageParam], overrides: dict[str, Any] = {}
     ) -> RetrievalResponse | AsyncGenerator[dict[str, Any], None]:
         text_search = overrides.get("retrieval_mode") in ["text", "hybrid", None]
         vector_search = overrides.get("retrieval_mode") in ["vectors", "hybrid", None]
         top = overrides.get("top", 3)
 
         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]
 
         # Generate an optimized keyword search query based on the chat history and the last question

src/fastapi_app/rag_simple.py

@@ -1,8 +1,6 @@
 import pathlib
 from collections.abc import AsyncGenerator
-from typing import (
-    Any,
-)
+from typing import Any
 
 from openai import AsyncOpenAI
 from openai.types.chat import ChatCompletion, ChatCompletionMessageParam
@@ -30,13 +28,15 @@ def __init__(
         self.answer_prompt_template = open(current_dir / "prompts/answer.txt").read()
 
     async def run(
-        self, messages: list[dict], overrides: dict[str, Any] = {}
+        self, messages: list[ChatCompletionMessageParam], overrides: dict[str, Any] = {}
     ) -> RetrievalResponse | AsyncGenerator[dict[str, Any], None]:
         text_search = overrides.get("retrieval_mode") in ["text", "hybrid", None]
         vector_search = overrides.get("retrieval_mode") in ["vectors", "hybrid", None]
         top = overrides.get("top", 3)
 
         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]
 
         # Retrieve relevant items from the database