mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-07-24 04:31:02 +02:00
Fixing some streamingg
This commit is contained in:
parent
215c1b9f75
commit
66170834f8
1 changed files with 52 additions and 15 deletions
|
|
@ -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(
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue