mirror of
https://github.com/Monadical-SAS/reflector.git
synced 2025-12-21 20:59:05 +00:00
Mathieu Virbel
7d40305737
* server: implement retry that automatically detect httpx and backoff Closes #118 * server: fix formatting
43 lines
1.2 KiB
Python
43 lines
1.2 KiB
Python
from reflector.llm.base import LLM
|
|
from reflector.settings import settings
|
|
from reflector.utils.retry import retry
|
|
import httpx
|
|
|
|
|
|
class BananaLLM(LLM):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.timeout = settings.LLM_TIMEOUT
|
|
self.headers = {
|
|
"X-Banana-API-Key": settings.LLM_BANANA_API_KEY,
|
|
"X-Banana-Model-Key": settings.LLM_BANANA_MODEL_KEY,
|
|
}
|
|
|
|
async def _generate(self, prompt: str, **kwargs):
|
|
async with httpx.AsyncClient() as client:
|
|
response = await retry(client.post)(
|
|
settings.LLM_URL,
|
|
headers=self.headers,
|
|
json={"prompt": prompt},
|
|
timeout=self.timeout,
|
|
retry_timeout=300, # as per their sdk
|
|
)
|
|
response.raise_for_status()
|
|
text = response.json()["text"]
|
|
text = text[len(prompt) :] # remove prompt
|
|
return text
|
|
|
|
|
|
LLM.register("banana", BananaLLM)
|
|
|
|
if __name__ == "__main__":
|
|
|
|
async def main():
|
|
llm = BananaLLM()
|
|
result = await llm.generate("Hello, my name is")
|
|
print(result)
|
|
|
|
import asyncio
|
|
|
|
asyncio.run(main())
|