more cleanup

ragna/assistants/_ai21labs.py

@@ -34,7 +34,7 @@ def _render_prompt(self, prompt: Union[str, list[Message]]) -> Union[str, list]:
         else:
             messages = prompt
         return [
-            {"role": message.role.value, "content": message.content}
+            {"text": message.content, "role": message.role.value}
             for message in messages
             if message.role is not MessageRole.SYSTEM
         ]

ragna/assistants/_google.py

@@ -50,6 +50,7 @@ async def generate(
         Returns:
             async streamed inference response string chunks
         """
+        # See https://ai.google.dev/api/generate-content#v1beta.models.streamGenerateContent
         async with self._call_api(
             "POST",
             f"https://generativelanguage.googleapis.com/v1beta/models/{self._MODEL}:streamGenerateContent",
@@ -76,7 +77,7 @@ async def generate(
                     "maxOutputTokens": max_new_tokens,
                 },
             },
-            parse_kwargs=dict(item="item"),  # .candidates.item.content.parts.item.text
+            parse_kwargs=dict(item="item"),
         ) as stream:
             async for data in stream:
                 yield data

ragna/assistants/_openai.py

@@ -1,12 +1,6 @@
 import abc
 from functools import cached_property
-from typing import (
-    Any,
-    AsyncIterator,
-    Optional,
-    Union,
-    cast,
-)
+from typing import Any, AsyncIterator, Optional, Union, cast
 
 from ragna.core import Message, MessageRole, Source
 

tests/assistants/test_api.py

@@ -94,15 +94,15 @@ async def generate(self, messages):
             content=messages[-1].content,
             parse_kwargs=parse_kwargs,
         ) as stream:
-            async for chunk in stream:
-                yield chunk
+            async for data in stream:
+                yield data
 
     async def answer(self, messages):
-        async for chunk in self.generate(messages):
-            if chunk.get("break"):
+        async for data in self.generate(messages):
+            if data.get("break"):
                 break
 
-            yield chunk
+            yield data
 
 
 @skip_on_windows