apunkt/trustgraph
1
0
Fork 0
mirror of https://github.com/trustgraph-ai/trustgraph.git synced 2026-07-23 20:21:03 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions 3

Some streaming fixes

Browse source
This commit is contained in:
Cyber MacGeddon 2025-11-25 18:02:25 +00:00
parent cab6e95d80
commit 606f181c55
3 changed files with 95 additions and 28 deletions
Download patch file Download diff file Expand all files Collapse all files

2
trustgraph-base/trustgraph/base/__init__.py
View file

@ -12,7 +12,7 @@ from . parameter_spec import ParameterSpec
from . producer_spec import ProducerSpec
from . subscriber_spec import SubscriberSpec
from . request_response_spec import RequestResponseSpec
from . llm_service import LlmService, LlmResult
from . llm_service import LlmService, LlmResult, LlmChunk
from . chunking_service import ChunkingService
from . embeddings_service import EmbeddingsService
from . embeddings_client import EmbeddingsClientSpec

41
trustgraph-base/trustgraph/base/text_completion_client.py
View file

@ -3,18 +3,45 @@ from . request_response_spec import RequestResponse, RequestResponseSpec
from .. schema import TextCompletionRequest, TextCompletionResponse
class TextCompletionClient(RequestResponse):
async def text_completion(self, system, prompt, timeout=600):
resp = await self.request(
async def text_completion(self, system, prompt, streaming=False, timeout=600):
# If not streaming, use original behavior
if not streaming:
resp = await self.request(
TextCompletionRequest(
system = system, prompt = prompt, streaming = False
),
timeout=timeout
)
if resp.error:
raise RuntimeError(resp.error.message)
return resp.response
# For streaming: collect all chunks and return complete response
full_response = ""
async def collect_chunks(resp):
nonlocal full_response
if resp.error:
raise RuntimeError(resp.error.message)
if resp.response:
full_response += resp.response
# Return True when end_of_stream is reached
return getattr(resp, 'end_of_stream', False)
await self.request(
TextCompletionRequest(
system = system, prompt = prompt
system = system, prompt = prompt, streaming = True
),
recipient=collect_chunks,
timeout=timeout
)
if resp.error:
raise RuntimeError(resp.error.message)
return resp.response
return full_response
class TextCompletionClientSpec(RequestResponseSpec):
def __init__(

80
trustgraph-flow/trustgraph/prompt/template/service.py
View file

@ -101,6 +101,9 @@ class Processor(FlowProcessor):
kind = v.id
# Check if streaming is requested
streaming = getattr(v, 'streaming', False)
try:
logger.debug(f"Prompt terms: {v.terms}")
@ -109,16 +112,65 @@ class Processor(FlowProcessor):
k: json.loads(v)
for k, v in v.terms.items()
}
logger.debug(f"Handling prompt kind {kind}...")
logger.debug(f"Handling prompt kind {kind}... (streaming={streaming})")
# If streaming, we need to handle it differently
if streaming:
# For streaming, we need to intercept LLM responses
# and forward them as they arrive
async def llm_streaming(system, prompt):
logger.debug(f"System prompt: {system}")
logger.debug(f"User prompt: {prompt}")
# Use the text completion client with recipient handler
client = flow("text-completion-request")
async def forward_chunks(resp):
if resp.error:
raise RuntimeError(resp.error.message)
if resp.response:
# Forward each chunk immediately
r = PromptResponse(
text=resp.response,
object=None,
error=None,
end_of_stream=getattr(resp, 'end_of_stream', False),
)
await flow("response").send(r, properties={"id": id})
# Return True when end_of_stream
return getattr(resp, 'end_of_stream', False)
await client.request(
TextCompletionRequest(
system=system, prompt=prompt, streaming=True
),
recipient=forward_chunks,
timeout=600
)
# Return empty string since we already sent all chunks
return ""
try:
await self.manager.invoke(kind, input, llm_streaming)
except Exception as e:
logger.error(f"Prompt streaming exception: {e}", exc_info=True)
raise e
return
# Non-streaming path (original behavior)
async def llm(system, prompt):
logger.debug(f"System prompt: {system}")
logger.debug(f"User prompt: {prompt}")
resp = await flow("text-completion-request").text_completion(
system = system, prompt = prompt,
system = system, prompt = prompt, streaming = False,
)
try:
@ -143,6 +195,7 @@ class Processor(FlowProcessor):
text=resp,
object=None,
error=None,
end_of_stream=True,
)
await flow("response").send(r, properties={"id": id})
@ -158,6 +211,7 @@ class Processor(FlowProcessor):
text=None,
object=json.dumps(resp),
error=None,
end_of_stream=True,
)
await flow("response").send(r, properties={"id": id})
@ -175,27 +229,13 @@ class Processor(FlowProcessor):
type = "llm-error",
message = str(e),
),
response=None,
text=None,
object=None,
end_of_stream=True,
)
await flow("response").send(r, properties={"id": id})
except Exception as e:
logger.error(f"Prompt service exception: {e}", exc_info=True)
logger.debug("Sending error response...")
r = PromptResponse(
error=Error(
type = "llm-error",
message = str(e),
),
response=None,
)
await self.send(r, properties={"id": id})
@staticmethod
def add_args(parser):