Fixing some streamingg

This commit is contained in:
Cyber MacGeddon 2025-11-25 23:51:52 +00:00
parent 215c1b9f75
commit 66170834f8

View file

@ -8,24 +8,61 @@ class PromptClient(RequestResponse):
async def prompt(self, id, variables, timeout=600, streaming=False): async def prompt(self, id, variables, timeout=600, streaming=False):
resp = await self.request( if not streaming:
PromptRequest( # Non-streaming path
id = id, resp = await self.request(
terms = { PromptRequest(
k: json.dumps(v) id = id,
for k, v in variables.items() terms = {
}, k: json.dumps(v)
streaming = streaming for k, v in variables.items()
), },
timeout=timeout streaming = False
) ),
timeout=timeout
)
if resp.error: if resp.error:
raise RuntimeError(resp.error.message) raise RuntimeError(resp.error.message)
if resp.text: return resp.text if resp.text: return resp.text
return json.loads(resp.object) return json.loads(resp.object)
else:
# Streaming path - collect all chunks
full_text = ""
full_object = None
async def collect_chunks(resp):
nonlocal full_text, full_object
if resp.error:
raise RuntimeError(resp.error.message)
if resp.text:
full_text += resp.text
elif resp.object:
full_object = resp.object
return getattr(resp, 'end_of_stream', False)
await self.request(
PromptRequest(
id = id,
terms = {
k: json.dumps(v)
for k, v in variables.items()
},
streaming = True
),
recipient=collect_chunks,
timeout=timeout
)
if full_text: return full_text
return json.loads(full_object)
async def extract_definitions(self, text, timeout=600): async def extract_definitions(self, text, timeout=600):
return await self.prompt( return await self.prompt(