Skip to content

Commit

Permalink
Fixed streamCompletionWithoutRetry sometimes dropping tokens
Browse files Browse the repository at this point in the history
When a json response was previously split across two "data: <JSON>" frames, the parser skipped the frames because the json parsing failed

E.g. the following response failed with the old code:
```
data: {"hello":
true}
```
  • Loading branch information
mpscholten committed Sep 13, 2023
1 parent 14c8c40 commit 8bdeb19
Showing 1 changed file with 9 additions and 15 deletions.
24 changes: 9 additions & 15 deletions ihp-openai/IHP/OpenAI.hs
Original file line number Diff line number Diff line change
Expand Up @@ -133,30 +133,24 @@ streamCompletionWithoutRetry secretKey completionRequest onStart callback = do
let status = getStatusCode p
if status == 200
then do
i2 <- Streams.mapM handleResponse i
x <- Streams.fold mappend mempty i2
return (Right x)
x <- Streams.foldM (parseResponseChunk callback) ("", "") i
return (Right (snd x))
else do
x <- Streams.fold mappend mempty i
return (Left $ "an error happend: " <> Text.pack (show x))

onStart
receiveResponse connection handler
where
handleResponse :: ByteString -> IO Text
handleResponse response = do
texts :: [Text] <- mapM handleLine (ByteString.lines response)
pure (Text.concat texts)
handleLine :: ByteString -> IO Text
handleLine line = do
case ByteString.stripPrefix "data: " line of
parseResponseChunk :: (Text -> IO ()) -> (ByteString, Text) -> ByteString -> IO (ByteString, Text)
parseResponseChunk callback (curBuffer, chunk) input = do
case ByteString.stripPrefix "data: " (ByteString.strip (curBuffer <> input)) of
Just json -> do
case decodeStrict json of
Just CompletionResult { choices } -> do
let tokens = mconcat $ map (.text) choices
let tokens :: Text = mconcat $ map (.text) choices
callback tokens
pure tokens
pure ("", chunk <> tokens)
otherwise -> do
pure ""
Nothing -> pure ""

pure (curBuffer <> json, chunk)
Nothing -> pure (curBuffer <> input, chunk)

0 comments on commit 8bdeb19

Please sign in to comment.