Skip to content
This repository was archived by the owner on Jun 5, 2025. It is now read-only.

Commit 2efb453

Browse files
Record output from non-streaming responses in DB (#1078)
Closes: #938 We were not recording the output of the non-streamed responses in the DB. In our API we're filtering out all the conversations that don't contain an answer. Since we were not recording non-streamed responses we were filtering out the entire conversation if it was non-streamed. Note: This still doesn't add an output pipeline for non-streamed responses.
1 parent baeef8c commit 2efb453

File tree

1 file changed

+16
-8
lines changed

1 file changed

+16
-8
lines changed

src/codegate/providers/base.py

Lines changed: 16 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -133,12 +133,23 @@ async def _run_output_stream_pipeline(
133133
denormalized_stream = self._output_normalizer.denormalize_streaming(pipeline_output_stream)
134134
return denormalized_stream
135135

136-
def _run_output_pipeline(
136+
async def _run_output_pipeline(
137137
self,
138-
normalized_response: ModelResponse,
138+
input_context: PipelineContext,
139+
model_response: Any,
139140
) -> ModelResponse:
140-
# we don't have a pipeline for non-streamed output yet
141-
return normalized_response
141+
"""
142+
Run the output pipeline for a single response.
143+
144+
For the moment we don't have a pipeline for non-streamed output, so we
145+
just normalize the response and record the context. It is done here to match
146+
the behaviour of the streaming pipeline.
147+
"""
148+
normalized_response = self._output_normalizer.normalize(model_response)
149+
input_context.add_output(normalized_response)
150+
await self._db_recorder.record_context(input_context)
151+
output_result = self._output_normalizer.denormalize(normalized_response)
152+
return output_result
142153

143154
async def _run_input_pipeline(
144155
self,
@@ -263,10 +274,7 @@ async def complete(
263274
is_fim_request=is_fim_request,
264275
)
265276
if not streaming:
266-
normalized_response = self._output_normalizer.normalize(model_response)
267-
pipeline_output = self._run_output_pipeline(normalized_response)
268-
await self._db_recorder.record_context(input_pipeline_result.context)
269-
return self._output_normalizer.denormalize(pipeline_output)
277+
return await self._run_output_pipeline(input_pipeline_result.context, model_response)
270278

271279
pipeline_output_stream = await self._run_output_stream_pipeline(
272280
input_pipeline_result.context, model_response, is_fim_request=is_fim_request # type: ignore

0 commit comments

Comments
 (0)