Add some None return types (#636)

Author: cbornet

examples/evaluation/tru_shared.py

@@ -192,7 +192,7 @@ def get_astra_vector_store(framework: Framework, collection_name: str):
     )
 
 
-def execute_query(framework: Framework, pipeline, query):
+def execute_query(framework: Framework, pipeline, query) -> None:
     if framework == Framework.LANG_CHAIN:
         pipeline.invoke(query)
     elif framework == Framework.LLAMA_INDEX:
@@ -204,7 +204,7 @@ def execute_query(framework: Framework, pipeline, query):
 
 
 # runs the pipeline across all queries in all known datasets
-def execute_experiment(framework: Framework, pipeline, experiment_name: str):
+def execute_experiment(framework: Framework, pipeline, experiment_name: str) -> None:
     init_tru()
 
     # use a short uuid to ensure that multiple experiments with the same name don't

examples/notebooks/advancedRAG.ipynb

@@ -101,7 +101,7 @@
     "import textwrap\n",
     "\n",
     "\n",
-    "def pprint_docs(docs):\n",
+    "def pprint_docs(docs) -> None:\n",
     "    print(\n",
     "        f\"\\n{'-' * 70}\\n\".join(\n",
     "            [\n",
@@ -112,7 +112,7 @@
     "    )\n",
     "\n",
     "\n",
-    "def pprint_result(result):\n",
+    "def pprint_result(result) -> None:\n",
     "    print(\"Answer: \" + \"\\n\".join(textwrap.wrap(result)))"
    ]
   },
@@ -244,7 +244,7 @@
     "from langchain.callbacks import get_openai_callback\n",
     "\n",
     "\n",
-    "def do_retrieval(chain):\n",
+    "def do_retrieval(chain) -> None:\n",
     "    for i in range(len(questions)):\n",
     "        print(\"-\" * 40)\n",
     "        print(f\"Question: {questions[i]}\\n\")\n",

examples/notebooks/conftest.py

@@ -29,7 +29,7 @@ def get_required_env(name) -> str:
 )
 
 
-def try_delete_with_backoff(collection: str, sleep=1, max_tries=2):
+def try_delete_with_backoff(collection: str, sleep=1, max_tries=2) -> None:
     try:
         logging.info("deleting collection %s", collection)
         response = client.delete_collection(collection)
@@ -46,7 +46,7 @@ def try_delete_with_backoff(collection: str, sleep=1, max_tries=2):
         try_delete_with_backoff(collection, sleep * 2, max_tries)
 
 
-def before_notebook():
+def before_notebook() -> None:
     collections = client.get_collections().get("status").get("collections")
     logging.info("Existing collections: %s", collections)
     for collection in collections:

examples/notebooks/nemo_guardrails.ipynb

@@ -192,7 +192,7 @@
     "\"\"\"\n",
     "\n",
     "\n",
-    "def yaml_config(engine, model):\n",
+    "def yaml_config(engine, model) -> str:\n",
     "    return f\"\"\"\n",
     "    models:\n",
     "      - type: main\n",
@@ -244,7 +244,7 @@
     "    return ActionResult(return_value=answer, context_updates=context_updates)\n",
     "\n",
     "\n",
-    "def init(app: LLMRails):\n",
+    "def init(app: LLMRails) -> None:\n",
     "    app.register_action(rag, \"rag\")"
    ]
   },

libs/e2e-tests/e2e_tests/conftest.py

@@ -156,13 +156,13 @@ def pytest_runtest_makereport(item, call):
         os.environ["RAGSTACK_E2E_TESTS_TEST_START"] = str(time.perf_counter_ns())
 
 
-def set_current_test_info(test_name: str, test_info: str):
+def set_current_test_info(test_name: str, test_info: str) -> None:
     test_info = test_info.replace("_", "-")
     os.environ["RAGSTACK_E2E_TESTS_TEST_INFO"] = f"{test_name}::{test_info}"
 
 
 @pytest.hookimpl()
-def pytest_sessionfinish():
+def pytest_sessionfinish() -> None:
     logging.info("All tests report:")
     logging.info("\n".join(all_report_lines))
     logging.info("Failed tests report:")
@@ -184,7 +184,7 @@ def pytest_sessionfinish():
     _report_to_file("", "llamaindex-tests-report.txt", llamaindex_report_lines)
 
 
-def _report_to_file(stats_str: str, filename: str, report_lines: list):
+def _report_to_file(stats_str: str, filename: str, report_lines: list) -> None:
     report_lines.sort()
     with open(filename, "w") as f:
         if stats_str:

libs/e2e-tests/e2e_tests/langchain/nemo_guardrails.py

@@ -53,7 +53,7 @@ def _colang() -> str:
 
 
 class NeMoRag:
-    def __init__(self, retriever):
+    def __init__(self, retriever) -> None:
         self.retriever = retriever
 
     async def rag_using_lc(self, context: dict, llm: BaseLLM) -> ActionResult:
@@ -77,7 +77,7 @@ async def rag_using_lc(self, context: dict, llm: BaseLLM) -> ActionResult:
 
         return ActionResult(return_value=answer, context_updates=context_updates)
 
-    def init(self, app: LLMRails):
+    def init(self, app: LLMRails) -> None:
         app.register_action(self.rag_using_lc, "rag")
 
 

libs/e2e-tests/e2e_tests/langchain/test_astra.py

@@ -25,14 +25,14 @@
 MINIMUM_ACCEPTABLE_SCORE = 0.1
 
 
-def test_basic_vector_search(vectorstore: AstraDBVectorStore):
+def test_basic_vector_search(vectorstore: AstraDBVectorStore) -> None:
     print("Running test_basic_vector_search")
     vectorstore.add_texts(["RAGStack is a framework to run LangChain in production"])
     retriever = vectorstore.as_retriever()
     assert len(retriever.get_relevant_documents("RAGStack")) > 0
 
 
-def test_ingest_errors(vectorstore: AstraDBVectorStore):
+def test_ingest_errors(vectorstore: AstraDBVectorStore) -> None:
     print("Running test_ingestion")
 
     empty_text = ""
@@ -75,7 +75,7 @@ def test_ingest_errors(vectorstore: AstraDBVectorStore):
             )
 
 
-def test_wrong_connection_parameters(vectorstore: AstraDBVectorStore):
+def test_wrong_connection_parameters(vectorstore: AstraDBVectorStore) -> None:
     try:
         AstraDBVectorStore(
             collection_name="something",
@@ -109,7 +109,7 @@ def test_wrong_connection_parameters(vectorstore: AstraDBVectorStore):
             )
 
 
-def test_basic_metadata_filtering_no_vector(vectorstore: AstraDBVectorStore):
+def test_basic_metadata_filtering_no_vector(vectorstore: AstraDBVectorStore) -> None:
     collection = vectorstore.collection
     vectorstore.add_texts(
         texts=["RAGStack is a framework to run LangChain in production"],
@@ -202,7 +202,7 @@ def test_basic_metadata_filtering_no_vector(vectorstore: AstraDBVectorStore):
             )
 
 
-def verify_document(document, expected_content, expected_metadata):
+def verify_document(document, expected_content, expected_metadata) -> None:
     if isinstance(document, Document):
         assert document.page_content == expected_content
         assert document.metadata == expected_metadata
@@ -211,7 +211,7 @@ def verify_document(document, expected_content, expected_metadata):
         assert document.get("metadata") == expected_metadata
 
 
-def test_vector_search_with_metadata(vectorstore: VectorStore):
+def test_vector_search_with_metadata(vectorstore: VectorStore) -> None:
     print("Running test_vector_search_with_metadata")
 
     document_ids = vectorstore.add_texts(
@@ -409,7 +409,7 @@ def test_vector_search_with_metadata(vectorstore: VectorStore):
 
 
 @pytest.mark.skip()
-def test_stress_astra():
+def test_stress_astra() -> None:
     handler = AstraDBVectorStoreHandler(VectorStoreImplementation.ASTRADB)
     while True:
         context = handler.before_test()
@@ -423,7 +423,7 @@ def test_stress_astra():
 
 
 class MockEmbeddings(Embeddings):
-    def __init__(self):
+    def __init__(self) -> None:
         self.embedded_documents = None
         self.embedded_query = None
 

libs/e2e-tests/e2e_tests/langchain/test_cassandra_tool.py

@@ -14,7 +14,7 @@
 
 
 @pytest.mark.usefixtures("cassandra")
-def test_tool_with_openai_tool():
+def test_tool_with_openai_tool() -> None:
     session = cassio.config.resolve_session()
     session.execute("DROP TABLE IF EXISTS default_keyspace.tool_table_users;")
 

libs/e2e-tests/e2e_tests/langchain/test_compatibility_rag.py

@@ -304,7 +304,7 @@ def llm():
         ),
     ],
 )
-def test_rag(test_case, vector_store, embedding, llm, request, record_property):
+def test_rag(test_case, vector_store, embedding, llm, request, record_property) -> None:
     set_current_test_info(
         "langchain::" + test_case,
         f"{llm},{embedding},{vector_store}",
@@ -327,7 +327,7 @@ def _run_test(
     embedding_fn,
     resolved_llm,
     record_property,
-):
+) -> None:
     # NeMo guardrails running only with certain LLMs
     if test_case == "nemo_guardrails" and not resolved_llm["nemo_config"]:
         skip_test_due_to_implementation_not_supported("nemo_guardrails")
@@ -408,7 +408,7 @@ def gemini_pro_llm():
         ("vertex_gemini_multimodal_embedding", "gemini_flash_llm"),
     ],
 )
-def test_multimodal(vector_store, embedding, llm, request, record_property):
+def test_multimodal(vector_store, embedding, llm, request, record_property) -> None:
     set_current_test_info(
         "langchain::multimodal_rag",
         f"{llm},{embedding},{vector_store}",
@@ -487,7 +487,7 @@ def embed_query(self, text: str) -> list[float]:
 
 
 @pytest.mark.parametrize("chat", ["vertex_gemini_pro_llm", "gemini_pro_llm"])
-def test_chat(chat, request, record_property):
+def test_chat(chat, request, record_property) -> None:
     set_current_test_info(
         "langchain::chat",
         chat,

libs/e2e-tests/e2e_tests/langchain/test_document_loaders.py

@@ -17,11 +17,11 @@
 from e2e_tests.test_utils.vector_store_handler import VectorStoreImplementation
 
 
-def set_current_test_info_document_loader(doc_loader: str):
+def set_current_test_info_document_loader(doc_loader: str) -> None:
     set_current_test_info("langchain::document_loader", doc_loader)
 
 
-def test_csv_loader():
+def test_csv_loader() -> None:
     set_current_test_info_document_loader("csv")
     with tempfile.NamedTemporaryFile(mode="w+", suffix=".csv") as temp_csv_file:
         with open(temp_csv_file.name, "w") as write:
@@ -41,7 +41,7 @@ def test_csv_loader():
         assert doc1.metadata == {"row": 0, "source": temp_csv_file.name}
 
 
-def test_web_based_loader():
+def test_web_based_loader() -> None:
     set_current_test_info_document_loader("web")
     loader = WebBaseLoader(
         ["https://langstream.ai/changelog/", "https://langstream.ai/faq/"]
@@ -69,7 +69,7 @@ def test_web_based_loader():
     }
 
 
-def test_s3_loader():
+def test_s3_loader() -> None:
     set_current_test_info_document_loader("s3")
     aws_region = "us-east-1"
     bucket_name = f"ragstack-ci-{uuid.uuid4()}"
@@ -95,7 +95,7 @@ def test_s3_loader():
         bucket.delete()
 
 
-def test_azure_blob_doc_loader():
+def test_azure_blob_doc_loader() -> None:
     set_current_test_info_document_loader("azure")
     from azure.storage.blob import BlobClient