From 76b559bc1525610e4ec7802635b98ce8fe17808d Mon Sep 17 00:00:00 2001
From: Anton Dubovik <anton_dubovik@epam.com>
Date: Thu, 10 Oct 2024 12:20:41 +0100
Subject: [PATCH] feat: enabled type checking of examples

---
 Makefile                      |  3 ++-
 examples/langchain_rag/app.py | 12 +++++++-----
 noxfile.py                    | 34 +++++++++++++++++++++++++++++++++-
 pyproject.toml                |  1 -
 4 files changed, 42 insertions(+), 8 deletions(-)

diff --git a/Makefile b/Makefile
index e2d65ca..ef99034 100644
--- a/Makefile
+++ b/Makefile
@@ -17,7 +17,8 @@ publish: build
 	poetry publish -u __token__ -p ${PYPI_TOKEN} --skip-existing
 
 lint: install
-	poetry run nox -s lint
+	poetry run nox -s lint_sdk
+	poetry run nox -s lint_examples
 
 format: install
 	poetry run nox -s format
diff --git a/examples/langchain_rag/app.py b/examples/langchain_rag/app.py
index 3666eb1..1d45f81 100644
--- a/examples/langchain_rag/app.py
+++ b/examples/langchain_rag/app.py
@@ -14,14 +14,16 @@
 from langchain.storage import LocalFileStore
 from langchain_community.document_loaders import PyPDFLoader, WebBaseLoader
 from langchain_community.vectorstores import Chroma
+from langchain_core.pydantic_v1 import SecretStr
 from langchain_openai import AzureChatOpenAI, AzureOpenAIEmbeddings
 from langchain_text_splitters import RecursiveCharacterTextSplitter
-from utils import get_last_attachment_url, sanitize_namespace
 
 from aidial_sdk import DIALApp
 from aidial_sdk import HTTPException as DIALException
 from aidial_sdk.chat_completion import ChatCompletion, Choice, Request, Response
 
+from .utils import get_last_attachment_url, sanitize_namespace
+
 
 def get_env(name: str) -> str:
     value = os.getenv(name)
@@ -97,8 +99,8 @@ async def chat_completion(
                     azure_deployment=EMBEDDINGS_MODEL,
                     azure_endpoint=DIAL_URL,
                     # Header propagation automatically propagates the API key from the request headers.
-                    openai_api_key="-",
-                    openai_api_version=API_VERSION,
+                    api_key=SecretStr("-"),
+                    api_version=API_VERSION,
                     # The check leads to tokenization of the input strings.
                     # Tokenized input is only supported by OpenAI embedding models.
                     # For other models, the check should be disabled.
@@ -120,8 +122,8 @@ async def chat_completion(
                 azure_deployment=CHAT_MODEL,
                 azure_endpoint=DIAL_URL,
                 # Header propagation automatically propagates the API key from the request headers.
-                openai_api_key="-",
-                openai_api_version=API_VERSION,
+                api_key=SecretStr("-"),
+                api_version=API_VERSION,
                 temperature=0,
                 streaming=True,
                 callbacks=[CustomCallbackHandler(choice)],
diff --git a/noxfile.py b/noxfile.py
index fb9e6b6..fffe943 100644
--- a/noxfile.py
+++ b/noxfile.py
@@ -1,3 +1,6 @@
+import os
+from pathlib import Path
+
 import nox
 
 nox.options.reuse_existing_virtualenvs = True
@@ -12,7 +15,7 @@ def format_with_args(session: nox.Session, *args):
 
 
 @nox.session
-def lint(session: nox.Session):
+def lint_sdk(session: nox.Session):
     """Runs linters and fixers"""
     try:
         session.run("poetry", "install", "--all-extras", external=True)
@@ -26,6 +29,35 @@ def lint(session: nox.Session):
         )
 
 
+EXAMPLES = [
+    folder for folder in os.listdir("examples") if not folder.startswith("_")
+]
+
+
+@nox.session
+@nox.parametrize("example_dir", EXAMPLES)
+def lint_examples(session: nox.Session, example_dir: str):
+    try:
+        root = str(Path(__file__).parent)
+        project_dir = f"{root}/examples/{example_dir}"
+
+        session.run("poetry", "install", "--with=lint", external=True)
+        session.run("poetry", "check", "--lock", external=True)
+
+        session.run(
+            "pip", "install", "-r", f"{project_dir}/requirements.txt", "-q"
+        )
+
+        session.chdir(project_dir)
+        session.run("pyright", SRC)
+        session.run("flake8", SRC)
+        format_with_args(session, SRC, "--check")
+    except Exception:
+        session.error(
+            "linting has failed. Run 'make format' to fix formatting and fix other errors manually"
+        )
+
+
 @nox.session
 def format(session: nox.Session):
     """Runs linters and fixers"""
diff --git a/pyproject.toml b/pyproject.toml
index 2391342..8b6b760 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -81,7 +81,6 @@ exclude = [
     ".pytest_cache",
     "**/__pycache__",
     "build",
-    "examples/langchain_rag"
 ]
 
 [tool.black]