Add run_manager and print the Google drive request.

langchain-ai · May 25, 2023 · 8096caf · 8096caf
1 parent 044d757
commit 8096caf
Show file tree

Hide file tree

Showing 8 changed files with 279 additions and 134 deletions.
diff --git a/docs/modules/agents/toolkits/examples/openapi.ipynb b/docs/modules/agents/toolkits/examples/openapi.ipynb
@@ -54,38 +54,38 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "--2023-03-31 15:45:56--  https://raw.githubusercontent.com/openai/openai-openapi/master/openapi.yaml\n",
-      "Resolving raw.githubusercontent.com (raw.githubusercontent.com)... 185.199.110.133, 185.199.109.133, 185.199.111.133, ...\n",
-      "Connecting to raw.githubusercontent.com (raw.githubusercontent.com)|185.199.110.133|:443... connected.\n",
-      "HTTP request sent, awaiting response... 200 OK\n",
-      "Length: 122995 (120K) [text/plain]\n",
-      "Saving to: ‘openapi.yaml’\n",
+      "--2023-05-24 15:59:36--  https://raw.githubusercontent.com/openai/openai-openapi/master/openapi.yaml\n",
+      "Résolution de raw.githubusercontent.com (raw.githubusercontent.com)… 2606:50c0:8001::154, 2606:50c0:8002::154, 2606:50c0:8003::154, ...\n",
+      "Connexion à raw.githubusercontent.com (raw.githubusercontent.com)|2606:50c0:8001::154|:443… connecté.\n",
+      "requête HTTP transmise, en attente de la réponse… 200 OK\n",
+      "Taille : 122995 (120K) [text/plain]\n",
+      "Enregistre : ‘openapi.yaml’\n",
       "\n",
-      "openapi.yaml        100%[===================>] 120.11K  --.-KB/s    in 0.01s   \n",
+      "openapi.yaml        100%[===================>] 120,11K  --.-KB/s    ds 0,03s   \n",
       "\n",
-      "2023-03-31 15:45:56 (10.4 MB/s) - ‘openapi.yaml’ saved [122995/122995]\n",
+      "2023-05-24 15:59:36 (4,53 MB/s) - ‘openapi.yaml’ enregistré [122995/122995]\n",
       "\n",
-      "--2023-03-31 15:45:57--  https://www.klarna.com/us/shopping/public/openai/v0/api-docs\n",
-      "Resolving www.klarna.com (www.klarna.com)... 52.84.150.34, 52.84.150.46, 52.84.150.61, ...\n",
-      "Connecting to www.klarna.com (www.klarna.com)|52.84.150.34|:443... connected.\n",
-      "HTTP request sent, awaiting response... 200 OK\n",
-      "Length: unspecified [application/json]\n",
-      "Saving to: ‘api-docs’\n",
+      "--2023-05-24 15:59:36--  https://www.klarna.com/us/shopping/public/openai/v0/api-docs\n",
+      "Résolution de www.klarna.com (www.klarna.com)… 2600:9000:2171:2200:c:51ac:89c0:93a1, 2600:9000:2171:c400:c:51ac:89c0:93a1, 2600:9000:2171:8200:c:51ac:89c0:93a1, ...\n",
+      "Connexion à www.klarna.com (www.klarna.com)|2600:9000:2171:2200:c:51ac:89c0:93a1|:443… connecté.\n",
+      "requête HTTP transmise, en attente de la réponse… 200 OK\n",
+      "Taille : non indiqué [application/json]\n",
+      "Enregistre : ‘api-docs’\n",
       "\n",
-      "api-docs                [ <=>                ]   1.87K  --.-KB/s    in 0s      \n",
+      "api-docs                [ <=>                ]   2,66K  --.-KB/s    ds 0s      \n",
       "\n",
-      "2023-03-31 15:45:57 (261 MB/s) - ‘api-docs’ saved [1916]\n",
+      "2023-05-24 15:59:37 (339 MB/s) - ‘api-docs’ enregistré [2723]\n",
       "\n",
-      "--2023-03-31 15:45:57--  https://raw.githubusercontent.com/APIs-guru/openapi-directory/main/APIs/spotify.com/1.0.0/openapi.yaml\n",
-      "Resolving raw.githubusercontent.com (raw.githubusercontent.com)... 185.199.110.133, 185.199.109.133, 185.199.111.133, ...\n",
-      "Connecting to raw.githubusercontent.com (raw.githubusercontent.com)|185.199.110.133|:443... connected.\n",
-      "HTTP request sent, awaiting response... 200 OK\n",
-      "Length: 286747 (280K) [text/plain]\n",
-      "Saving to: ‘openapi.yaml’\n",
+      "--2023-05-24 15:59:37--  https://raw.githubusercontent.com/APIs-guru/openapi-directory/main/APIs/spotify.com/1.0.0/openapi.yaml\n",
+      "Résolution de raw.githubusercontent.com (raw.githubusercontent.com)… 2606:50c0:8000::154, 2606:50c0:8003::154, 2606:50c0:8001::154, ...\n",
+      "Connexion à raw.githubusercontent.com (raw.githubusercontent.com)|2606:50c0:8000::154|:443… connecté.\n",
+      "requête HTTP transmise, en attente de la réponse… 200 OK\n",
+      "Taille : 289324 (283K) [text/plain]\n",
+      "Enregistre : ‘openapi.yaml’\n",
       "\n",
-      "openapi.yaml        100%[===================>] 280.03K  --.-KB/s    in 0.02s   \n",
+      "openapi.yaml        100%[===================>] 282,54K  --.-KB/s    ds 0,05s   \n",
       "\n",
-      "2023-03-31 15:45:58 (13.3 MB/s) - ‘openapi.yaml’ saved [286747/286747]\n",
+      "2023-05-24 15:59:37 (5,19 MB/s) - ‘openapi.yaml’ enregistré [289324/289324]\n",
       "\n"
      ]
     }
@@ -174,7 +174,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": 5,
    "id": "2a93271e",
    "metadata": {},
    "outputs": [
@@ -184,7 +184,7 @@
        "63"
       ]
      },
-     "execution_count": 6,
+     "execution_count": 5,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -202,18 +202,46 @@
   {
    "cell_type": "code",
    "execution_count": 7,
+   "id": "1d4cae4e",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Requirement already satisfied: tiktoken in /home/pprados/workspace.bda/propaleur/langchain/.venv/lib/python3.10/site-packages (0.3.3)\n",
+      "Requirement already satisfied: regex>=2022.1.18 in /home/pprados/workspace.bda/propaleur/langchain/.venv/lib/python3.10/site-packages (from tiktoken) (2023.5.5)\n",
+      "Requirement already satisfied: requests>=2.26.0 in /home/pprados/workspace.bda/propaleur/langchain/.venv/lib/python3.10/site-packages (from tiktoken) (2.28.2)\n",
+      "Requirement already satisfied: idna<4,>=2.5 in /home/pprados/workspace.bda/propaleur/langchain/.venv/lib/python3.10/site-packages (from requests>=2.26.0->tiktoken) (3.4)\n",
+      "Requirement already satisfied: charset-normalizer<4,>=2 in /home/pprados/workspace.bda/propaleur/langchain/.venv/lib/python3.10/site-packages (from requests>=2.26.0->tiktoken) (3.1.0)\n",
+      "Requirement already satisfied: certifi>=2017.4.17 in /home/pprados/workspace.bda/propaleur/langchain/.venv/lib/python3.10/site-packages (from requests>=2.26.0->tiktoken) (2023.5.7)\n",
+      "Requirement already satisfied: urllib3<1.27,>=1.21.1 in /home/pprados/workspace.bda/propaleur/langchain/.venv/lib/python3.10/site-packages (from requests>=2.26.0->tiktoken) (1.26.15)\n",
+      "\u001b[33mWARNING: You are using pip version 22.0.4; however, version 23.1.2 is available.\n",
+      "You should consider upgrading via the '/home/pprados/workspace.bda/propaleur/langchain/.venv/bin/python -m pip install --upgrade pip' command.\u001b[0m\u001b[33m\n",
+      "\u001b[0m"
+     ]
+    }
+   ],
+   "source": [
+    "!pip install tiktoken"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
    "id": "eb829190",
    "metadata": {},
    "outputs": [
     {
-     "data": {
-      "text/plain": [
-       "80326"
-      ]
-     },
-     "execution_count": 7,
-     "metadata": {},
-     "output_type": "execute_result"
+     "ename": "ModuleNotFoundError",
+     "evalue": "No module named 'tiktoken'",
+     "output_type": "error",
+     "traceback": [
+      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
+      "\u001b[0;31mModuleNotFoundError\u001b[0m                       Traceback (most recent call last)",
+      "Cell \u001b[0;32mIn[8], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mtiktoken\u001b[39;00m\n\u001b[1;32m      2\u001b[0m enc \u001b[38;5;241m=\u001b[39m tiktoken\u001b[38;5;241m.\u001b[39mencoding_for_model(\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mtext-davinci-003\u001b[39m\u001b[38;5;124m'\u001b[39m)\n\u001b[1;32m      3\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mcount_tokens\u001b[39m(s): \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mlen\u001b[39m(enc\u001b[38;5;241m.\u001b[39mencode(s))\n",
+      "\u001b[0;31mModuleNotFoundError\u001b[0m: No module named 'tiktoken'"
+     ]
     }
    ],
    "source": [
@@ -759,7 +787,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.9.1"
+   "version": "3.10.9"
   }
  },
  "nbformat": 4,

diff --git a/langchain/agents/load_tools.py b/langchain/agents/load_tools.py
@@ -273,8 +273,10 @@ def _get_openweathermap(**kwargs: Any) -> BaseTool:
 
 _EXTRA_OPTIONAL_TOOLS: Dict[str, Tuple[Callable[[KwArg(Any)], BaseTool], List[str]]] = {
     "wolfram-alpha": (_get_wolfram_alpha, ["wolfram_alpha_appid"]),
-    "google-drive-search": (_get_google_drive_search,
-                            ["gdrive_api_file", "mode", "gslide_mode", "gsheet_mode"]),
+    "google-drive-search": (
+        _get_google_drive_search,
+        ["gdrive_api_file", "mode", "gslide_mode", "gsheet_mode"],
+    ),
     "google-search": (_get_google_search, ["google_api_key", "google_cse_id"]),
     "google-search-results-json": (
         _get_google_search_results_json,
@@ -311,7 +313,7 @@ def _get_openweathermap(**kwargs: Any) -> BaseTool:
 
 
 def _handle_callbacks(
-        callback_manager: Optional[BaseCallbackManager], callbacks: Callbacks
+    callback_manager: Optional[BaseCallbackManager], callbacks: Callbacks
 ) -> Callbacks:
     if callback_manager is not None:
         warnings.warn(
@@ -327,11 +329,11 @@ def _handle_callbacks(
 
 
 def load_huggingface_tool(
-        task_or_repo_id: str,
-        model_repo_id: Optional[str] = None,
-        token: Optional[str] = None,
-        remote: bool = False,
-        **kwargs: Any,
+    task_or_repo_id: str,
+    model_repo_id: Optional[str] = None,
+    token: Optional[str] = None,
+    remote: bool = False,
+    **kwargs: Any,
 ) -> BaseTool:
     try:
         from transformers import load_tool
@@ -361,10 +363,10 @@ def load_huggingface_tool(
 
 
 def load_tools(
-        tool_names: List[str],
-        llm: Optional[BaseLanguageModel] = None,
-        callbacks: Callbacks = None,
-        **kwargs: Any,
+    tool_names: List[str],
+    llm: Optional[BaseLanguageModel] = None,
+    callbacks: Callbacks = None,
+    **kwargs: Any,
 ) -> List[BaseTool]:
     """Load tools based on their name.
 
@@ -430,8 +432,8 @@ def load_tools(
 def get_all_tool_names() -> List[str]:
     """Get a list of all possible tool names."""
     return (
-            list(_BASE_TOOLS)
-            + list(_EXTRA_OPTIONAL_TOOLS)
-            + list(_EXTRA_LLM_TOOLS)
-            + list(_LLM_TOOLS)
+        list(_BASE_TOOLS)
+        + list(_EXTRA_OPTIONAL_TOOLS)
+        + list(_EXTRA_LLM_TOOLS)
+        + list(_LLM_TOOLS)
     )
diff --git a/langchain/docstore/google_drive.py b/langchain/docstore/google_drive.py
@@ -1,7 +1,8 @@
-from typing import Any, Dict, Union
+from typing import Any, Dict, Optional, Union
 
 from pydantic.class_validators import root_validator
 
+from langchain.callbacks.manager import RunManager
 from langchain.docstore.base import Docstore
 from langchain.schema import Document
 from langchain.utilities.google_drive import GoogleDriveUtilities, get_template
@@ -27,15 +28,21 @@ def validate_template(cls, v: Dict[str, Any]) -> Dict[str, Any]:
             v["template"] = template
         return v
 
-    def search(self, search: str) -> Union[str, Document]:
+    def search(
+        self,
+        search: str,
+        run_manager: Optional[RunManager] = None,
+    ) -> Union[str, Document]:
         """Try to search for document.
 
         If document exists, return the document `description` and
         a PageWithLookups object.
         If page does not exist, return an error message.
         """
         try:
-            result = next(self.lazy_get_relevant_documents(query=search))
+            result = next(
+                self.lazy_get_relevant_documents(query=search, run_manager=run_manager)
+            )
             return Document(
                 page_content=result.page_content + "\n",
                 metadata={"page": result.metadata["source"]},

diff --git a/langchain/document_loaders/google_drive.py b/langchain/document_loaders/google_drive.py
@@ -15,6 +15,7 @@
 from pydantic.class_validators import root_validator
 
 from langchain.base_language import BaseLanguageModel
+from langchain.callbacks.manager import RunManager
 from langchain.document_loaders.base import BaseLoader
 from langchain.schema import Document
 from langchain.utilities.google_drive import (
@@ -118,13 +119,18 @@ def _lazy_load_files_ids(self) -> Iterator[Document]:
         )
 
     def lazy_get_relevant_documents(
-        self, query: Optional[str] = None, **kwargs: Any
+        self,
+        query: Optional[str] = None,
+        run_manager: Optional[RunManager] = None,
+        **kwargs: Any,
     ) -> Iterator[Document]:
         if self.document_ids:
             return self._lazy_load_documents_ids()
         if self.file_ids:
             return self._lazy_load_files_ids()
-        return super().lazy_get_relevant_documents(query=query, **kwargs)
+        return super().lazy_get_relevant_documents(
+            query=query, run_manager=run_manager, **kwargs
+        )
 
     def lazy_load(self) -> Iterator[Document]:
         return self.lazy_get_relevant_documents()

diff --git a/langchain/retrievers/google_drive.py b/langchain/retrievers/google_drive.py
@@ -1,8 +1,9 @@
-from typing import Any, Dict, List, Literal
+from typing import Any, Dict, List, Literal, Optional
 
 from pydantic.class_validators import root_validator
 from pydantic.config import Extra
 
+from langchain.callbacks.manager import RunManager
 from langchain.schema import BaseRetriever, Document
 from langchain.utilities.google_drive import (
     GoogleDriveUtilities,
@@ -41,7 +42,11 @@ def validate_template(cls, v: Dict[str, Any]) -> Dict[str, Any]:
             v["template"] = template
         return v
 
-    def get_relevant_documents(self, query: str) -> List[Document]:
+    def get_relevant_documents(
+        self,
+        query: str,
+        run_manager: Optional[RunManager] = None,
+    ) -> List[Document]:
         """Get documents relevant for a query.
 
         Args:
@@ -50,7 +55,9 @@ def get_relevant_documents(self, query: str) -> List[Document]:
         Returns:
             List of relevant documents
         """
-        return list(self.lazy_get_relevant_documents(query=query))
+        return list(
+            self.lazy_get_relevant_documents(query=query, run_manager=run_manager)
+        )
 
     async def aget_relevant_documents(self, query: str) -> List[Document]:
         """Get documents relevant for a query.

diff --git a/langchain/tools/google_drive/tool.py b/langchain/tools/google_drive/tool.py
@@ -25,7 +25,7 @@ def _run(
         run_manager: Optional[CallbackManagerForToolRun] = None,
     ) -> str:
         """Use the tool."""
-        return self.api_wrapper.run(query)
+        return self.api_wrapper.run(query, run_manager)
 
     async def _arun(
         self,