forked from Archives/langchain
You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
70 lines
2.1 KiB
Python
70 lines
2.1 KiB
Python
from __future__ import annotations
|
|
|
|
from typing import List
|
|
|
|
import aiohttp
|
|
import requests
|
|
from pydantic import BaseModel
|
|
|
|
from langchain.schema import BaseRetriever, Document
|
|
|
|
|
|
class ChatGPTPluginRetriever(BaseRetriever, BaseModel):
|
|
url: str
|
|
bearer_token: str
|
|
top_k: int = 3
|
|
filter: dict | None = None
|
|
aiosession: aiohttp.ClientSession | None = None
|
|
|
|
class Config:
|
|
"""Configuration for this pydantic object."""
|
|
|
|
arbitrary_types_allowed = True
|
|
|
|
def get_relevant_documents(self, query: str) -> List[Document]:
|
|
url, json, headers = self._create_request(query)
|
|
response = requests.post(url, json=json, headers=headers)
|
|
results = response.json()["results"][0]["results"]
|
|
docs = []
|
|
for d in results:
|
|
content = d.pop("text")
|
|
docs.append(Document(page_content=content, metadata=d))
|
|
return docs
|
|
|
|
async def aget_relevant_documents(self, query: str) -> List[Document]:
|
|
url, json, headers = self._create_request(query)
|
|
|
|
if not self.aiosession:
|
|
async with aiohttp.ClientSession() as session:
|
|
async with session.post(url, headers=headers, json=json) as response:
|
|
res = await response.json()
|
|
else:
|
|
async with self.aiosession.post(
|
|
url, headers=headers, json=json
|
|
) as response:
|
|
res = await response.json()
|
|
|
|
results = res["results"][0]["results"]
|
|
docs = []
|
|
for d in results:
|
|
content = d.pop("text")
|
|
docs.append(Document(page_content=content, metadata=d))
|
|
return docs
|
|
|
|
def _create_request(self, query: str) -> tuple[str, dict, dict]:
|
|
url = f"{self.url}/query"
|
|
json = {
|
|
"queries": [
|
|
{
|
|
"query": query,
|
|
"filter": self.filter,
|
|
"top_k": self.top_k,
|
|
}
|
|
]
|
|
}
|
|
headers = {
|
|
"Content-Type": "application/json",
|
|
"Authorization": f"Bearer {self.bearer_token}",
|
|
}
|
|
return url, json, headers
|