Files
reflector/server/reflector/llm/llm_oobabooga.py
Gokul Mohanarangan 2d686da15c pass schema as dict
2023-08-17 21:51:44 +05:30

22 lines
649 B
Python

import httpx
from reflector.llm.base import LLM
from reflector.settings import settings
class OobaboogaLLM(LLM):
async def _generate(self, prompt: str, schema: dict | None, **kwargs):
json_payload = {"prompt": prompt}
if schema:
json_payload["schema"] = schema
async with httpx.AsyncClient() as client:
response = await client.post(
settings.LLM_URL,
headers={"Content-Type": "application/json"},
json=json_payload,
)
response.raise_for_status()
return response.json()
LLM.register("oobabooga", OobaboogaLLM)