[WIP] Notebook testing changes

Makefile

@@ -6,7 +6,7 @@
 
 SHELL := /bin/bash
 REPO_ROOT := $(shell dirname $(realpath $(firstword $(MAKEFILE_LIST))))
-PYTEST := poetry run pytest --rootdir=. -s
+PYTEST := poetry run pytest -n auto --rootdir=. -s
 POETRY_DIRS := $(shell find . \
 	-not -path "./dist/*" \
 	-maxdepth 4 \
@@ -29,6 +29,7 @@ env-%:
 env-tests:
 	poetry run pip install \
 		pytest \
+		pytest-xdist[psutil] \
 		nbconvert \
 		nbformat \
 		pytest-subtests \
@@ -63,8 +64,12 @@ env-tests-db: env-tests
 env-tests-notebook: env-tests env-tests-optional
 	poetry run pip install \
 		faiss-cpu \
+		langchainhub \
+		llama-hub \
+		rank_bm25 \
 		ipytree \
-		llama-index-readers-web
+		llama-index-readers-web \
+		llama-index-vector-stores-milvus
 
 # Lock the poetry dependencies for all the subprojects.
 lock: $(POETRY_DIRS)
@@ -127,7 +132,7 @@ codespell:
 
 # Generates a coverage report.
 coverage:
-	ALLOW_OPTIONALS=true poetry run pytest --rootdir=. tests/* --cov src --cov-report html
+	ALLOW_OPTIONALS=true poetry run pytest -n auto --rootdir=. tests/* --cov src --cov-report html
 
 # Run the static unit tests only, those in the static subfolder. They are run
 # for every tested python version while those outside of static are run only for
@@ -201,18 +206,21 @@ test-%-allow-optional: env
 test-%-optional: env-tests-optional
 	TEST_OPTIONAL=true make test-$*
 
+test-notebook: env-tests-notebook
+	make test-notebook-optional
+
 # Run the unit tests, those in the tests/unit. They are run in the CI pipeline
 # frequently.
 test-unit:
-	poetry run pytest --rootdir=. tests/unit/*
+	poetry run pytest -n auto --rootdir=. tests/unit/*
 # Tests in the e2e folder make use of possibly costly endpoints. They
 # are part of only the less frequently run release tests.
 test-e2e:
-	poetry run pytest --rootdir=. tests/e2e/*
+	poetry run pytest -n auto --rootdir=. tests/e2e/*
 
 # Runs the notebook test
 test-notebook:
-	poetry run pytest --rootdir=. tests/docs_notebooks/*
+	poetry run pytest -n auto --rootdir=. tests/docs_notebooks/*
 
 install-wheels:
 	pip install dist/*/*.whl

examples/expositional/frameworks/langchain/langchain_agents.ipynb

@@ -124,7 +124,7 @@
     "    description=\"useful for when you need to answer questions about current events\",\n",
     ")\n",
     "\n",
-    "llm = ChatOpenAI(model=\"gpt-3.5-turbo\", temperature=0)\n",
+    "llm = ChatOpenAI(model=\"gpt-4o-mini\", temperature=0)\n",
     "\n",
     "tools = [search_tool]\n",
     "\n",
@@ -152,7 +152,7 @@
     "        return (\n",
     "            float(\n",
     "                self.endpoint.client.chat.completions.create(\n",
-    "                    model=\"gpt-3.5-turbo\",\n",
+    "                    model=\"gpt-4o-mini\",\n",
     "                    messages=[\n",
     "                        {\n",
     "                            \"role\": \"system\",\n",

examples/expositional/frameworks/langchain/langchain_math_agent.ipynb

@@ -26,7 +26,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# !pip install trulens trulens-apps-langchain langchain==0.0.283"
+    "# !pip install trulens trulens-apps-langchain langchain langchain-openai"
    ]
   },
   {
@@ -35,11 +35,11 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from langchain import LLMMathChain\n",
+    "from langchain.chains.llm_math.base import LLMMathChain\n",
     "from langchain.agents import AgentType\n",
     "from langchain.agents import Tool\n",
     "from langchain.agents import initialize_agent\n",
-    "from langchain.chat_models import ChatOpenAI\n",
+    "from langchain_openai.chat_models import ChatOpenAI\n",
     "from trulens.core import TruSession\n",
     "from trulens.apps.langchain import TruChain\n",
     "\n",
@@ -80,7 +80,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "llm = ChatOpenAI(temperature=0, model=\"gpt-3.5-turbo-0613\")\n",
+    "llm = ChatOpenAI(temperature=0)\n",
     "\n",
     "llm_math_chain = LLMMathChain.from_llm(llm, verbose=True)\n",
     "\n",
@@ -139,7 +139,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3.11.4 ('agents')",
+   "display_name": "pkg_311",
    "language": "python",
    "name": "python3"
   },
@@ -153,11 +153,6 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3"
-  },
-  "vscode": {
-   "interpreter": {
-    "hash": "7d153714b979d5e6d08dd8ec90712dd93bff2c9b6c1f0c118169738af3430cd4"
-   }
   }
  },
  "nbformat": 4,

examples/expositional/frameworks/langchain/langchain_retrieval_agent.ipynb

@@ -130,7 +130,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "llm = ChatOpenAI(model_name=\"gpt-3.5-turbo-16k\", temperature=0.0)\n",
+    "llm = ChatOpenAI(model_name=\"gpt-4o-mini\", temperature=0.0)\n",
     "\n",
     "conversational_memory = ConversationSummaryBufferMemory(\n",
     "    k=4,\n",
@@ -245,7 +245,7 @@
     "        return (\n",
     "            float(\n",
     "                self.endpoint.client.chat.completions.create(\n",
-    "                    model=\"gpt-3.5-turbo\",\n",
+    "                    model=\"gpt-4o-mini\",\n",
     "                    messages=[\n",
     "                        {\n",
     "                            \"role\": \"system\",\n",
@@ -267,7 +267,7 @@
     "        return (\n",
     "            float(\n",
     "                self.endpoint.client.chat.completions.create(\n",
-    "                    model=\"gpt-3.5-turbo\",\n",
+    "                    model=\"gpt-4o-mini\",\n",
     "                    messages=[\n",
     "                        {\n",
     "                            \"role\": \"system\",\n",

examples/expositional/frameworks/llama_index/llama_index_agents.ipynb

@@ -145,7 +145,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "### Create a standalone GPT3.5 for comparison"
+    "### Create a standalone GPT for comparison"
    ]
   },
   {
@@ -154,9 +154,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "client = openai.OpenAI()\n",
-    "\n",
-    "chat_completion = client.chat.completions.create"
+    "client = openai.OpenAI()"
    ]
   },
   {
@@ -166,15 +164,15 @@
    "outputs": [],
    "source": [
     "from trulens.apps.custom import TruCustomApp\n",
-    "from trulens.core import instrument\n",
+    "from trulens.core.instruments import instrument\n",
     "\n",
     "\n",
     "class LLMStandaloneApp:\n",
     "    @instrument\n",
     "    def __call__(self, prompt):\n",
     "        return (\n",
-    "            chat_completion(\n",
-    "                model=\"gpt-3.5-turbo\",\n",
+    "            client.chat.completions.create(\n",
+    "                model=\"gpt-4o-mini\",\n",
     "                messages=[\n",
     "                    {\"role\": \"system\", \"content\": gordon_ramsay_prompt},\n",
     "                    {\"role\": \"user\", \"content\": prompt},\n",
@@ -486,7 +484,7 @@
    "provenance": []
   },
   "kernelspec": {
-   "display_name": "Python 3 (ipykernel)",
+   "display_name": "pkg_311",
    "language": "python",
    "name": "python3"
   },
@@ -500,11 +498,6 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3"
-  },
-  "vscode": {
-   "interpreter": {
-    "hash": "7d153714b979d5e6d08dd8ec90712dd93bff2c9b6c1f0c118169738af3430cd4"
-   }
   }
  },
  "nbformat": 4,

examples/expositional/frameworks/llama_index/llama_index_complex_evals.ipynb

@@ -95,7 +95,7 @@
    "outputs": [],
    "source": [
     "# Merge into a single large document rather than one document per-page\n",
-    "from llama_index import Document\n",
+    "from llama_index.core import Document\n",
     "\n",
     "document = Document(text=\"\\n\\n\".join([doc.text for doc in documents]))"
    ]
@@ -107,8 +107,8 @@
    "outputs": [],
    "source": [
     "from llama_index.core import ServiceContext\n",
-    "from llama_index.llms import OpenAI\n",
-    "from llama_index.node_parser import SentenceWindowNodeParser\n",
+    "from llama_index.llms.openai import OpenAI\n",
+    "from llama_index.core.node_parser import SentenceWindowNodeParser\n",
     "\n",
     "# create the sentence window node parser w/ default settings\n",
     "node_parser = SentenceWindowNodeParser.from_defaults(\n",
@@ -117,7 +117,7 @@
     "    original_text_metadata_key=\"original_text\",\n",
     ")\n",
     "\n",
-    "llm = OpenAI(model=\"gpt-3.5-turbo\", temperature=0.1)\n",
+    "llm = OpenAI(model=\"gpt-4o-mini\", temperature=0.1)\n",
     "sentence_context = ServiceContext.from_defaults(\n",
     "    llm=llm,\n",
     "    embed_model=\"local:BAAI/bge-small-en-v1.5\",\n",
@@ -340,7 +340,7 @@
    "provenance": []
   },
   "kernelspec": {
-   "display_name": "milvus",
+   "display_name": "pkg_311",
    "language": "python",
    "name": "python3"
   },

examples/expositional/frameworks/nemoguardrails/nemoguardrails_feedback_action_example.ipynb

@@ -109,7 +109,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from trulens.tru_rails import FeedbackActions\n",
+    "from trulens.apps.nemo.tru_rails import FeedbackActions\n",
     "\n",
     "FeedbackActions.register_feedback_functions(**fs_triad)\n",
     "FeedbackActions.register_feedback_functions(f_language_match)"
@@ -142,7 +142,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from trulens.dashboard.notebook_utils import writefileinterpolated"
+    "from trulens.dashboard.utils.notebook_utils import writefileinterpolated"
    ]
   },
   {
@@ -170,7 +170,7 @@
     "models:\n",
     "  - type: main\n",
     "    engine: openai\n",
-    "    model: gpt-3.5-turbo-instruct\n",
+    "    model: gpt-4o-mini\n",
     "\n",
     "rails:\n",
     "  output:\n",
@@ -479,7 +479,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "py311_trulens",
+   "display_name": "pkg_311",
    "language": "python",
    "name": "python3"
   },

examples/expositional/frameworks/nemoguardrails/nemoguardrails_trurails_example.ipynb

@@ -88,7 +88,7 @@
     "models:\n",
     "  - type: main\n",
     "    engine: openai\n",
-    "    model: gpt-3.5-turbo-instruct"
+    "    model: gpt-4o-mini"
    ]
   },
   {

examples/expositional/frameworks/openai_assistants/openai_assistants_api.ipynb

@@ -112,17 +112,26 @@
     "        self.client = client\n",
     "\n",
     "        # upload the file\\\n",
-    "        file = client.files.create(\n",
-    "            file=open(\"data/paul_graham_essay.txt\", \"rb\"), purpose=\"assistants\"\n",
-    "        )\n",
+    "        \n",
     "\n",
     "        # create the assistant with access to a retrieval tool\n",
     "        assistant = client.beta.assistants.create(\n",
     "            name=\"Paul Graham Essay Assistant\",\n",
     "            instructions=\"You are an assistant that answers questions about Paul Graham.\",\n",
-    "            tools=[{\"type\": \"retrieval\"}],\n",
     "            model=\"gpt-4-turbo-preview\",\n",
-    "            file_ids=[file.id],\n",
+    "            tools=[{\"type\": \"file_search\"}],\n",
+    "        )\n",
+    "\n",
+    "        vector_store = client.beta.vector_stores.create(name=\"Paul Graham\")\n",
+    "\n",
+    "        files = [open(\"data/paul_graham_essay.txt\", \"rb\")]\n",
+    "        client.beta.vector_stores.file_batches.upload_and_poll(\n",
+    "            vector_store_id=vector_store.id, files=files\n",
+    "        )\n",
+    "\n",
+    "        assistant = client.beta.assistants.update(\n",
+    "            assistant_id=assistant.id,\n",
+    "            tool_resources={\"file_search\": {\"vector_store_ids\": [vector_store.id]}},\n",
     "        )\n",
     "\n",
     "        self.assistant = assistant\n",

examples/expositional/models/anthropic/claude3_quickstart.ipynb

@@ -93,7 +93,7 @@
     "oai_client = OpenAI()\n",
     "\n",
     "oai_client.embeddings.create(\n",
-    "    model=\"text-embedding-ada-002\", input=university_info\n",
+    "    model=\"text-embedding-3-small\", input=university_info\n",
     ")"
    ]
   },
@@ -108,7 +108,7 @@
     "\n",
     "embedding_function = OpenAIEmbeddingFunction(\n",
     "    api_key=os.environ.get(\"OPENAI_API_KEY\"),\n",
-    "    model_name=\"text-embedding-ada-002\",\n",
+    "    model_name=\"text-embedding-3-small\",\n",
     ")\n",
     "\n",
     "\n",
@@ -178,7 +178,7 @@
     "        \"\"\"\n",
     "        completion = (\n",
     "            oai_client.chat.completions.create(\n",
-    "                model=\"gpt-3.5-turbo\",\n",
+    "                model=\"gpt-4o-mini\",\n",
     "                temperature=0,\n",
     "                messages=[\n",
     "                    {\n",

examples/expositional/models/azure/azure_openai_llama_index.ipynb

@@ -118,7 +118,7 @@
     "\n",
     "# You need to deploy your own embedding model as well as your own chat completion model\n",
     "embed_model = AzureOpenAIEmbedding(\n",
-    "    model=\"text-embedding-ada-002\",\n",
+    "    model=\"text-embedding-3-small\",\n",
     "    deployment_name=\"<your deployment>\",\n",
     "    api_key=os.environ[\"AZURE_OPENAI_API_KEY\"],\n",
     "    azure_endpoint=os.environ[\"AZURE_OPENAI_ENDPOINT\"],\n",