kolega-code 0.1.0__py3-none-any.whl

kolega_code/agent/tests/llm/test_openai_cached_tokens_stream.py

@@ -0,0 +1,74 @@
+import pytest
+
+from kolega_code.llm.providers.openai import OpenAIProvider
+from kolega_code.llm.models import Message, MessageHistory
+
+
+class _Delta:
+    def __init__(self, content=None, tool_calls=None):
+        self.content = content
+        self.tool_calls = tool_calls or []
+
+
+class _Choice:
+    def __init__(self, delta, finish_reason=None):
+        self.delta = delta
+        self.finish_reason = finish_reason
+
+
+class _Chunk:
+    def __init__(self, text=None, is_final=False):
+        self.choices = [_Choice(delta=_Delta(content=text), finish_reason='stop' if is_final else None)]
+        self.usage = None
+
+
+class _Usage:
+    def __init__(self):
+        self.prompt_tokens = 3019
+        self.completion_tokens = 104
+        self.total_tokens = 3123
+        self.prompt_tokens_details = {'cached_tokens': 2048}
+
+
+class _AsyncStream:
+    def __init__(self):
+        self._chunks = [_Chunk('hello '), _Chunk('world', is_final=True)]
+        self._i = 0
+
+    def __aiter__(self):
+        return self
+
+    async def __anext__(self):
+        if self._i >= len(self._chunks):
+            raise StopAsyncIteration
+        chunk = self._chunks[self._i]
+        self._i += 1
+        if self._i == len(self._chunks):
+            chunk.usage = _Usage()  # final chunk carries usage
+        return chunk
+
+    async def aclose(self):
+        pass
+
+
+# TODO: Fix after qwen-3-coder-plus PR is merged - needs OpenAI cache token extraction in streaming
+@pytest.mark.asyncio
+async def test_openai_stream_includes_cached_tokens(monkeypatch):
+    provider = OpenAIProvider(api_key='sk-test', base_url='https://api.openai.com/v1')
+
+    async def fake_create(*args, **kwargs):
+        return _AsyncStream()
+
+    monkeypatch.setattr(provider.async_client.chat.completions, 'create', fake_create)
+
+    stream = await provider.stream(messages=MessageHistory([Message(role='user', content='hi')]))
+    async with stream as s:
+        async for _ in s:
+            pass
+        final = await s.get_final_message()
+        assert final.usage_metadata['prompt_tokens'] == 3019
+        assert final.usage_metadata['total_tokens'] == 3123
+        assert final.usage_metadata['cache_read_input_tokens'] == 2048
+
+
+

kolega_code/agent/tests/llm/test_openai_message_conversion.py

@@ -0,0 +1,30 @@
+import pytest
+
+from kolega_code.llm.models import Message, MessageHistory, ToolCall, ToolResult, TextBlock
+
+
+# TODO: Fix after qwen-3-coder-plus PR is merged - needs tool result filtering in Message.to_openai()
+def test_mixed_content_splits_tool_result_to_separate_tool_messages():
+    # Assistant message with text + tool_call + tool_result
+    tool_call = ToolCall(id='call_123', name='sum', input={'a': 1, 'b': 2})
+    tool_result = ToolResult(tool_use_id='call_123', content='3', name='sum', is_error=False)
+    assistant = Message(role='assistant', content=[TextBlock('ok'), tool_call, tool_result])
+
+    mh = MessageHistory([assistant])
+    messages = mh.to_openai()
+
+    # Expect two messages: assistant (with tool_calls) then tool with matching tool_call_id
+    assert messages[0]['role'] == 'assistant'
+    assert 'tool_calls' in messages[0]
+    assert messages[0]['tool_calls'][0]['id'] == 'call_123'
+    # content should not include tool_result; only text parts
+    assert isinstance(messages[0]['content'], list)
+    assert all(isinstance(p, dict) and p.get('type') in ('text', 'image_url') for p in messages[0]['content'])
+
+    assert messages[1]['role'] == 'tool'
+    assert messages[1]['tool_call_id'] == 'call_123'
+    # tool content can be a string or structured parts depending on our representation
+    assert 'content' in messages[1]
+
+
+