Improve and tests

.github/workflows/app-tests.yaml

@@ -5,9 +5,13 @@ on:
     branches: [ main ]
   pull_request:
     branches: [ main ]
+  workflow_dispatch:
+
+permissions:
+  contents: read
 
 jobs:
-  test_package:
+  test-package:
     name: Test ${{ matrix.os }} Python ${{ matrix.python_version }}
     runs-on: ${{ matrix.os }}
     strategy:
@@ -65,4 +69,6 @@ jobs:
           run: |
             cd ./src/frontend
             npm install
-            npm run build
+            npm run build
+        - name: Run Pytest
+          run: python3 -m pytest

.github/workflows/python-code-quality.yaml

@@ -3,23 +3,34 @@ name: Python code quality
 on:
   push:
     branches: [ main ]
+    paths:
+      - '**.py'
+
   pull_request:
     branches: [ main ]
+    paths:
+      - '**.py'
+
+  workflow_dispatch:
+
+permissions:
+  contents: read
 
 jobs:
-  build:
+  checks-format-and-lint:
     runs-on: ubuntu-latest
     steps:
         - uses: actions/checkout@v4
         - name: Set up Python 3
           uses: actions/setup-python@v5
           with:
             python-version: "3.12"
+            cache: 'pip'
         - name: Install dependencies
           run: |
-            python -m pip install --upgrade pip
-            pip install -r requirements-dev.txt
+            python3 -m pip install --upgrade pip
+            python3 -m pip install ruff
         - name: Lint with ruff
           run: ruff check .
         - name: Check formatting with ruff
-          run: ruff format --check .
+          run: ruff format . --check

pyproject.toml

@@ -1,10 +1,20 @@
 [tool.ruff]
 line-length = 120
-target-version = "py311"
+target-version = "py312"
+lint.select = ["E", "F", "I", "UP"]
+lint.ignore = ["D203"]
+lint.isort.known-first-party = ["fastapi_app"]
 
-[tool.ruff.lint]
-select = ["E", "F", "I", "UP"]
-ignore = ["D203"]
+[tool.mypy]
+check_untyped_defs = true
+python_version = 3.12
+exclude = [".venv/*"]
 
-[tool.ruff.lint.isort]
-known-first-party = ["fastapi_app"]
+[tool.pytest.ini_options]
+addopts = "-ra --cov"
+testpaths = ["tests"]
+pythonpath = ['src']
+filterwarnings = ["ignore::DeprecationWarning"]
+
+[tool.coverage.report]
+show_missing = true

requirements-dev.txt

@@ -2,4 +2,8 @@
 ruff
 pre-commit
 pip-tools
-pip-compile-cross-platform
+pip-compile-cross-platform
+pytest
+pytest-cov
+pytest-asyncio
+psycopg2-binary

src/fastapi_app/api_models.py

@@ -17,3 +17,15 @@ class ThoughtStep(BaseModel):
     title: str
     description: Any
     props: dict = {}
+
+
+class RAGContext(BaseModel):
+    data_points: dict[int, dict[str, Any]]
+    thoughts: list[ThoughtStep]
+    followup_questions: list[str] | None = None
+
+
+class RetrievalResponse(BaseModel):
+    message: Message
+    context: RAGContext
+    session_state: Any | None = None

src/fastapi_app/api_routes.py

@@ -9,6 +9,8 @@
 from fastapi_app.rag_advanced import AdvancedRAGChat
 from fastapi_app.rag_simple import SimpleRAGChat
 
+from .api_models import RetrievalResponse
+
 router = fastapi.APIRouter()
 
 
@@ -52,7 +54,7 @@ async def search_handler(query: str, top: int = 5, enable_vector_search: bool =
     return [item.to_dict() for item in results]
 
 
-@router.post("/chat")
+@router.post("/chat", response_model=RetrievalResponse)
 async def chat_handler(chat_request: ChatRequest):
     messages = [message.model_dump() for message in chat_request.messages]
     overrides = chat_request.context.get("overrides", {})
@@ -79,5 +81,5 @@ async def chat_handler(chat_request: ChatRequest):
             chat_deployment=global_storage.openai_chat_deployment,
         )
 
-    response = await ragchat.run(messages, overrides=overrides)
+    response: RetrievalResponse = await ragchat.run(messages, overrides=overrides)
     return response

src/fastapi_app/embeddings.py

@@ -4,7 +4,7 @@
 
 
 async def compute_text_embedding(
-    q: str, openai_client, embed_model: str, embed_deployment: str = None, embedding_dimensions: int = 1536
+    q: str, openai_client, embed_model: str, embed_deployment: str | None = None, embedding_dimensions: int = 1536
 ):
     SUPPORTED_DIMENSIONS_MODEL = {
         "text-embedding-ada-002": False,

src/fastapi_app/postgres_searcher.py

@@ -103,7 +103,7 @@ async def search(
 
     async def search_and_embed(
         self,
-        query_text: str,
+        query_text: str | None = None,
         top: int = 5,
         enable_vector_search: bool = False,
         enable_text_search: bool = False,

src/fastapi_app/rag_advanced.py

@@ -5,12 +5,10 @@
 )
 
 from openai import AsyncOpenAI
-from openai.types.chat import (
-    ChatCompletion,
-)
+from openai.types.chat import ChatCompletion, ChatCompletionMessageParam
 from openai_messages_token_helper import build_messages, get_token_limit
 
-from .api_models import ThoughtStep
+from .api_models import Message, RAGContext, RetrievalResponse, ThoughtStep
 from .postgres_searcher import PostgresSearcher
 from .query_rewriter import build_search_function, extract_search_arguments
 
@@ -35,7 +33,7 @@ def __init__(
 
     async def run(
         self, messages: list[dict], overrides: dict[str, Any] = {}
-    ) -> dict[str, Any] | AsyncGenerator[dict[str, Any], None]:
+    ) -> 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)
@@ -45,7 +43,7 @@ async def run(
 
         # Generate an optimized keyword search query based on the chat history and the last question
         query_response_token_limit = 500
-        query_messages = build_messages(
+        query_messages: list[ChatCompletionMessageParam] = build_messages(
             model=self.chat_model,
             system_prompt=self.query_prompt_template,
             new_user_content=original_user_query,
@@ -55,7 +53,7 @@ async def run(
         )
 
         chat_completion: ChatCompletion = await self.openai_chat_client.chat.completions.create(
-            messages=query_messages,  # type: ignore
+            messages=query_messages,
             # Azure OpenAI takes the deployment name as the model name
             model=self.chat_deployment if self.chat_deployment else self.chat_model,
             temperature=0.0,  # Minimize creativity for search query generation
@@ -81,7 +79,7 @@ async def run(
 
         # Generate a contextual and content specific answer using the search results and chat history
         response_token_limit = 1024
-        messages = build_messages(
+        contextual_messages: list[ChatCompletionMessageParam] = build_messages(
             model=self.chat_model,
             system_prompt=overrides.get("prompt_template") or self.answer_prompt_template,
             new_user_content=original_user_query + "\n\nSources:\n" + content,
@@ -90,21 +88,21 @@ async def run(
             fallback_to_default=True,
         )
 
-        chat_completion_response = await self.openai_chat_client.chat.completions.create(
+        chat_completion_response: ChatCompletion = await self.openai_chat_client.chat.completions.create(
             # Azure OpenAI takes the deployment name as the model name
             model=self.chat_deployment if self.chat_deployment else self.chat_model,
-            messages=messages,
+            messages=contextual_messages,
             temperature=overrides.get("temperature", 0.3),
             max_tokens=response_token_limit,
             n=1,
             stream=False,
         )
-        first_choice = chat_completion_response.model_dump()["choices"][0]
-        return {
-            "message": first_choice["message"],
-            "context": {
-                "data_points": {item.id: item.to_dict() for item in results},
-                "thoughts": [
+        first_choice = chat_completion_response.choices[0]
+        return RetrievalResponse(
+            message=Message(content=first_choice.message.content, role=first_choice.message.role),
+            context=RAGContext(
+                data_points={item.id: item.to_dict() for item in results},
+                thoughts=[
                     ThoughtStep(
                         title="Prompt to generate search arguments",
                         description=[str(message) for message in query_messages],
@@ -130,13 +128,13 @@ async def run(
                     ),
                     ThoughtStep(
                         title="Prompt to generate answer",
-                        description=[str(message) for message in messages],
+                        description=[str(message) for message in contextual_messages],
                         props=(
                             {"model": self.chat_model, "deployment": self.chat_deployment}
                             if self.chat_deployment
                             else {"model": self.chat_model}
                         ),
                     ),
                 ],
-            },
-        }
+            ),
+        )

src/fastapi_app/rag_simple.py

@@ -5,9 +5,10 @@
 )
 
 from openai import AsyncOpenAI
+from openai.types.chat import ChatCompletion, ChatCompletionMessageParam
 from openai_messages_token_helper import build_messages, get_token_limit
 
-from .api_models import ThoughtStep
+from .api_models import Message, RAGContext, RetrievalResponse, ThoughtStep
 from .postgres_searcher import PostgresSearcher
 
 
@@ -30,7 +31,7 @@ def __init__(
 
     async def run(
         self, messages: list[dict], overrides: dict[str, Any] = {}
-    ) -> dict[str, Any] | AsyncGenerator[dict[str, Any], None]:
+    ) -> 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)
@@ -48,7 +49,7 @@ async def run(
 
         # Generate a contextual and content specific answer using the search results and chat history
         response_token_limit = 1024
-        messages = build_messages(
+        contextual_messages: list[ChatCompletionMessageParam] = build_messages(
             model=self.chat_model,
             system_prompt=overrides.get("prompt_template") or self.answer_prompt_template,
             new_user_content=original_user_query + "\n\nSources:\n" + content,
@@ -57,21 +58,21 @@ async def run(
             fallback_to_default=True,
         )
 
-        chat_completion_response = await self.openai_chat_client.chat.completions.create(
+        chat_completion_response: ChatCompletion = await self.openai_chat_client.chat.completions.create(
             # Azure OpenAI takes the deployment name as the model name
             model=self.chat_deployment if self.chat_deployment else self.chat_model,
-            messages=messages,
+            messages=contextual_messages,
             temperature=overrides.get("temperature", 0.3),
             max_tokens=response_token_limit,
             n=1,
             stream=False,
         )
-        first_choice = chat_completion_response.model_dump()["choices"][0]
-        return {
-            "message": first_choice["message"],
-            "context": {
-                "data_points": {item.id: item.to_dict() for item in results},
-                "thoughts": [
+        first_choice = chat_completion_response.choices[0]
+        return RetrievalResponse(
+            message=Message(content=first_choice.message.content, role=first_choice.message.role),
+            context=RAGContext(
+                data_points={item.id: item.to_dict() for item in results},
+                thoughts=[
                     ThoughtStep(
                         title="Search query for database",
                         description=original_user_query if text_search else None,
@@ -87,13 +88,13 @@ async def run(
                     ),
                     ThoughtStep(
                         title="Prompt to generate answer",
-                        description=[str(message) for message in messages],
+                        description=[str(message) for message in contextual_messages],
                         props=(
                             {"model": self.chat_model, "deployment": self.chat_deployment}
                             if self.chat_deployment
                             else {"model": self.chat_model}
                         ),
                     ),
                 ],
-            },
-        }
+            ),
+        )