use continue and update tests

vocodedev · Aug 24, 2023 · bbfd34f · bbfd34f
1 parent c575af9
commit bbfd34f
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 7 deletions.
diff --git a/tests/streaming/agent/test_utils.py b/tests/streaming/agent/test_utils.py
@@ -28,11 +28,17 @@ async def _agen_from_list(l):
 
 
 def create_chatgpt_openai_object(
-    delta: Dict[str, str], finish_reason: Optional[Any] = None
+    delta: Optional[Dict[str, str]] = None,
+    finish_reason: Optional[Any] = None,
+    prompt_annotations=None,
 ):
-    return OpenAIObject.construct_from(
-        {"choices": [{"delta": delta, "finish_reason": finish_reason}]}
-    )
+    inner_obj = {}
+    if prompt_annotations:
+        inner_obj["prompt_annotations"] = prompt_annotations
+        inner_obj["choices"] = []
+    elif delta:
+        inner_obj["choices"] = [{"delta": delta, "finish_reason": finish_reason}]
+    return OpenAIObject.construct_from(inner_obj)
 
 
 class StreamOpenAIResponseTestCase(BaseModel):
@@ -43,6 +49,14 @@ class StreamOpenAIResponseTestCase(BaseModel):
 
 OPENAI_OBJECTS = [
     [
+        {
+            "prompt_annotations": [
+                {
+                    "prompt_index": 0,
+                    "content_filter_results": {},
+                }
+            ]
+        },
         {"delta": {"role": "assistant"}, "finish_reason": None},
         {"delta": {"content": "Hello"}, "finish_reason": None},
         {"delta": {"content": "!"}, "finish_reason": None},

diff --git a/vocode/streaming/agent/utils.py b/vocode/streaming/agent/utils.py
@@ -74,10 +74,8 @@ async def collate_response_async(
 async def openai_get_tokens(gen) -> AsyncGenerator[Union[str, FunctionFragment], None]:
     async for event in gen:
         choices = event.get("choices", [])
-        if event.get("prompt_annotations", []):
-            continue
         if len(choices) == 0:
-            break
+            continue
         choice = choices[0]
         if choice.finish_reason:
             break