make requests more general (#2209)

This commit is contained in:
Harrison Chase 2023-03-30 20:41:56 -07:00 committed by GitHub
parent 1c03205cc2
commit 2d3918c152
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 22 additions and 21 deletions

View File

@ -18,25 +18,25 @@ class RequestsWrapper(BaseModel):
extra = Extra.forbid extra = Extra.forbid
arbitrary_types_allowed = True 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.""" """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.""" """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.""" """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.""" """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.""" """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: async def _arequest(self, method: str, url: str, **kwargs: Any) -> str:
"""Make an async request.""" """Make an async request."""
@ -52,22 +52,22 @@ class RequestsWrapper(BaseModel):
) as response: ) as response:
return await response.text() 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.""" """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.""" """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.""" """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.""" """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.""" """DELETE the URL and return the text asynchronously."""
return await self._arequest("DELETE", url) return await self._arequest("DELETE", url, **kwargs)

View File

@ -1,6 +1,7 @@
"""Test LLM Math functionality.""" """Test LLM Math functionality."""
import json import json
from typing import Any
import pytest import pytest
@ -16,7 +17,7 @@ class FakeRequestsChain(RequestsWrapper):
output: str output: str
def get(self, url: str) -> str: def get(self, url: str, **kwargs: Any) -> str:
"""Just return the specified output.""" """Just return the specified output."""
return self.output return self.output