From 2d3918c15215144054661084095237ee1c5cec15 Mon Sep 17 00:00:00 2001 From: Harrison Chase Date: Thu, 30 Mar 2023 20:41:56 -0700 Subject: [PATCH] make requests more general (#2209) --- langchain/requests.py | 40 ++++++++++++++--------------- tests/unit_tests/chains/test_api.py | 3 ++- 2 files changed, 22 insertions(+), 21 deletions(-) diff --git a/langchain/requests.py b/langchain/requests.py index 60efc634..ac5091f0 100644 --- a/langchain/requests.py +++ b/langchain/requests.py @@ -18,25 +18,25 @@ class RequestsWrapper(BaseModel): extra = Extra.forbid arbitrary_types_allowed = True - def get(self, url: str) -> str: + def get(self, url: str, **kwargs: Any) -> str: """GET the URL and return the text.""" - return requests.get(url, headers=self.headers).text + return requests.get(url, headers=self.headers, **kwargs).text - def post(self, url: str, data: Dict[str, Any]) -> str: + def post(self, url: str, data: Dict[str, Any], **kwargs: Any) -> str: """POST to the URL and return the text.""" - return requests.post(url, json=data, headers=self.headers).text + return requests.post(url, json=data, headers=self.headers, **kwargs).text - def patch(self, url: str, data: Dict[str, Any]) -> str: + def patch(self, url: str, data: Dict[str, Any], **kwargs: Any) -> str: """PATCH the URL and return the text.""" - return requests.patch(url, json=data, headers=self.headers).text + return requests.patch(url, json=data, headers=self.headers, **kwargs).text - def put(self, url: str, data: Dict[str, Any]) -> str: + def put(self, url: str, data: Dict[str, Any], **kwargs: Any) -> str: """PUT the URL and return the text.""" - return requests.put(url, json=data, headers=self.headers).text + return requests.put(url, json=data, headers=self.headers, **kwargs).text - def delete(self, url: str) -> str: + def delete(self, url: str, **kwargs: Any) -> str: """DELETE the URL and return the text.""" - return requests.delete(url, headers=self.headers).text + return requests.delete(url, headers=self.headers, **kwargs).text async def _arequest(self, method: str, url: str, **kwargs: Any) -> str: """Make an async request.""" @@ -52,22 +52,22 @@ class RequestsWrapper(BaseModel): ) as response: return await response.text() - async def aget(self, url: str) -> str: + async def aget(self, url: str, **kwargs: Any) -> str: """GET the URL and return the text asynchronously.""" - return await self._arequest("GET", url) + return await self._arequest("GET", url, **kwargs) - async def apost(self, url: str, data: Dict[str, Any]) -> str: + async def apost(self, url: str, data: Dict[str, Any], **kwargs: Any) -> str: """POST to the URL and return the text asynchronously.""" - return await self._arequest("POST", url, json=data) + return await self._arequest("POST", url, json=data, **kwargs) - async def apatch(self, url: str, data: Dict[str, Any]) -> str: + async def apatch(self, url: str, data: Dict[str, Any], **kwargs: Any) -> str: """PATCH the URL and return the text asynchronously.""" - return await self._arequest("PATCH", url, json=data) + return await self._arequest("PATCH", url, json=data, **kwargs) - async def aput(self, url: str, data: Dict[str, Any]) -> str: + async def aput(self, url: str, data: Dict[str, Any], **kwargs: Any) -> str: """PUT the URL and return the text asynchronously.""" - return await self._arequest("PUT", url, json=data) + return await self._arequest("PUT", url, json=data, **kwargs) - async def adelete(self, url: str) -> str: + async def adelete(self, url: str, **kwargs: Any) -> str: """DELETE the URL and return the text asynchronously.""" - return await self._arequest("DELETE", url) + return await self._arequest("DELETE", url, **kwargs) diff --git a/tests/unit_tests/chains/test_api.py b/tests/unit_tests/chains/test_api.py index 8749d100..757960fb 100644 --- a/tests/unit_tests/chains/test_api.py +++ b/tests/unit_tests/chains/test_api.py @@ -1,6 +1,7 @@ """Test LLM Math functionality.""" import json +from typing import Any import pytest @@ -16,7 +17,7 @@ class FakeRequestsChain(RequestsWrapper): output: str - def get(self, url: str) -> str: + def get(self, url: str, **kwargs: Any) -> str: """Just return the specified output.""" return self.output