diff --git a/hindsight-integrations/litellm/hindsight_litellm/wrappers.py b/hindsight-integrations/litellm/hindsight_litellm/wrappers.py index ccfb5451..d672b6f4 100644 --- a/hindsight-integrations/litellm/hindsight_litellm/wrappers.py +++ b/hindsight-integrations/litellm/hindsight_litellm/wrappers.py @@ -12,7 +12,7 @@ import os import threading from dataclasses import dataclass, fields from datetime import datetime, timezone -from typing import Any, Dict, List, Optional +from typing import Any, Dict, Iterator, List, Optional from .config import ( DEFAULT_BANK_ID, @@ -678,6 +678,149 @@ async def aretain( ) +class _StreamWrapper: + """Wrapper for OpenAI stream that collects content and stores conversation when done.""" + + def __init__( + self, + stream: Any, + user_query: str, + model: str, + wrapper: "HindsightOpenAI", + settings: HindsightCallSettings, + ): + self._stream = stream + self._user_query = user_query + self._model = model + self._wrapper = wrapper + self._settings = settings + self._collected_content: List[str] = [] + self._finished = False + + def __iter__(self) -> Iterator[Any]: + return self + + def __next__(self) -> Any: + try: + chunk = next(self._stream) + # Collect content from the chunk + if hasattr(chunk, "choices") and chunk.choices: + delta = chunk.choices[0].delta + if hasattr(delta, "content") and delta.content: + self._collected_content.append(delta.content) + return chunk + except StopIteration: + # Stream exhausted - store conversation if we collected content + self._store_if_needed() + raise + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + self._store_if_needed() + if hasattr(self._stream, "__exit__"): + return self._stream.__exit__(exc_type, exc_val, exc_tb) + + def _store_if_needed(self): + """Store the collected conversation if not already done.""" + if self._finished or not self._settings.store_conversations: + return + + self._finished = True + if self._collected_content: + assistant_output = "".join(self._collected_content) + if assistant_output: + try: + self._wrapper._store_conversation( + self._user_query, assistant_output, self._model, self._settings + ) + except Exception as e: + if self._settings.verbose: + logger.warning(f"Failed to store streamed conversation: {e}") + + def close(self): + """Close the underlying stream if it has a close method.""" + self._store_if_needed() + if hasattr(self._stream, "close"): + self._stream.close() + + def __getattr__(self, name: str) -> Any: + """Proxy other attributes to the underlying stream.""" + return getattr(self._stream, name) + + +class _AnthropicStreamWrapper: + """Wrapper for Anthropic stream that collects content and stores conversation when done.""" + + def __init__( + self, + stream: Any, + user_query: str, + model: str, + wrapper: "HindsightAnthropic", + settings: HindsightCallSettings, + ): + self._stream = stream + self._user_query = user_query + self._model = model + self._wrapper = wrapper + self._settings = settings + self._collected_content: List[str] = [] + self._finished = False + + def __iter__(self) -> Iterator[Any]: + return self + + def __next__(self) -> Any: + try: + chunk = next(self._stream) + # Collect content from the chunk + if hasattr(chunk, "type") and chunk.type == "content_block_delta": + if hasattr(chunk, "delta") and hasattr(chunk.delta, "text"): + self._collected_content.append(chunk.delta.text) + return chunk + except StopIteration: + # Stream exhausted - store conversation if we collected content + self._store_if_needed() + raise + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + self._store_if_needed() + if hasattr(self._stream, "__exit__"): + return self._stream.__exit__(exc_type, exc_val, exc_tb) + + def _store_if_needed(self): + """Store the collected conversation if not already done.""" + if self._finished or not self._settings.store_conversations: + return + + self._finished = True + if self._collected_content: + assistant_output = "".join(self._collected_content) + if assistant_output: + try: + self._wrapper._store_conversation( + self._user_query, assistant_output, self._model, self._settings + ) + except Exception as e: + if self._settings.verbose: + logger.warning(f"Failed to store streamed conversation: {e}") + + def close(self): + """Close the underlying stream if it has a close method.""" + self._store_if_needed() + if hasattr(self._stream, "close"): + self._stream.close() + + def __getattr__(self, name: str) -> Any: + """Proxy other attributes to the underlying stream.""" + return getattr(self._stream, name) + + class HindsightOpenAI: """Wrapper for OpenAI client with Hindsight memory integration. @@ -1080,16 +1223,30 @@ class _WrappedCompletions: # Make the actual API call response = self._wrapper._client.chat.completions.create(**openai_kwargs) - # Store conversation - if user_query and settings.store_conversations: - if response.choices and response.choices[0].message: - assistant_output = response.choices[0].message.content or "" - if assistant_output: - self._wrapper._store_conversation( - user_query, assistant_output, model, settings - ) - - return response + # Handle streaming vs non-streaming responses + is_streaming = openai_kwargs.get("stream", False) + if is_streaming: + # Wrap the stream to collect content and store conversation when done + if user_query and settings.store_conversations: + return _StreamWrapper( + stream=response, + user_query=user_query, + model=model, + wrapper=self._wrapper, + settings=settings, + ) + else: + return response + else: + # Non-streaming: store conversation immediately + if user_query and settings.store_conversations: + if response.choices and response.choices[0].message: + assistant_output = response.choices[0].message.content or "" + if assistant_output: + self._wrapper._store_conversation( + user_query, assistant_output, model, settings + ) + return response class HindsightAnthropic: @@ -1477,19 +1634,33 @@ class _WrappedAnthropicMessages: # Make the actual API call response = self._wrapper._client.messages.create(**anthropic_kwargs) - # Store conversation - if user_query and settings.store_conversations: - if response.content: - assistant_output = "" - for block in response.content: - if hasattr(block, "text"): - assistant_output += block.text - if assistant_output: - self._wrapper._store_conversation( - user_query, assistant_output, model, settings - ) - - return response + # Handle streaming vs non-streaming responses + is_streaming = anthropic_kwargs.get("stream", False) + if is_streaming: + # Wrap the stream to collect content and store conversation when done + if user_query and settings.store_conversations: + return _AnthropicStreamWrapper( + stream=response, + user_query=user_query, + model=model, + wrapper=self._wrapper, + settings=settings, + ) + else: + return response + else: + # Non-streaming: store conversation immediately + if user_query and settings.store_conversations: + if response.content: + assistant_output = "" + for block in response.content: + if hasattr(block, "text"): + assistant_output += block.text + if assistant_output: + self._wrapper._store_conversation( + user_query, assistant_output, model, settings + ) + return response def wrap_openai( diff --git a/hindsight-integrations/litellm/tests/test_integration.py b/hindsight-integrations/litellm/tests/test_integration.py index 62902560..6d44e839 100644 --- a/hindsight-integrations/litellm/tests/test_integration.py +++ b/hindsight-integrations/litellm/tests/test_integration.py @@ -587,3 +587,182 @@ class TestSetDefaults: def test_get_defaults_returns_none_initially(self): """Test get_defaults returns None when not set.""" assert get_defaults() is None + + +class TestStreamingSupport: + """Tests for streaming support in wrappers.""" + + def test_wrap_openai_with_stream_no_error(self): + """Test that wrap_openai handles streaming without errors.""" + from unittest.mock import Mock, MagicMock + from hindsight_litellm.wrappers import wrap_openai + + # Create mock OpenAI client + mock_client = Mock() + mock_stream = MagicMock() + mock_client.chat.completions.create.return_value = mock_stream + + # Wrap the client with store_conversations=False + wrapped = wrap_openai( + mock_client, + hindsight_api_url="http://localhost:8888", + bank_id="test-agent", + store_conversations=False, # Disable storage for this test + ) + + # Call with stream=True + result = wrapped.chat.completions.create( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "Hello"}], + stream=True, + ) + + # Should return the stream without errors + assert result == mock_stream + # Verify the underlying client was called with stream=True + mock_client.chat.completions.create.assert_called_once() + call_kwargs = mock_client.chat.completions.create.call_args[1] + assert call_kwargs["stream"] is True + + def test_wrap_anthropic_with_stream_no_error(self): + """Test that wrap_anthropic handles streaming without errors.""" + from unittest.mock import Mock, MagicMock + from hindsight_litellm.wrappers import wrap_anthropic + + # Create mock Anthropic client + mock_client = Mock() + mock_stream = MagicMock() + mock_client.messages.create.return_value = mock_stream + + # Wrap the client with store_conversations=False + wrapped = wrap_anthropic( + mock_client, + hindsight_api_url="http://localhost:8888", + bank_id="test-agent", + store_conversations=False, # Disable storage for this test + ) + + # Call with stream=True + result = wrapped.messages.create( + model="claude-sonnet-4-20250514", + max_tokens=1024, + messages=[{"role": "user", "content": "Hello"}], + stream=True, + ) + + # Should return the stream without errors + assert result == mock_stream + # Verify the underlying client was called with stream=True + mock_client.messages.create.assert_called_once() + call_kwargs = mock_client.messages.create.call_args[1] + assert call_kwargs["stream"] is True + + def test_wrap_openai_stream_stores_conversation(self): + """Test that streaming stores conversation after all chunks are consumed.""" + from unittest.mock import Mock, MagicMock, patch + from hindsight_litellm.wrappers import wrap_openai + + # Create mock OpenAI client + mock_client = Mock() + + # Create mock stream chunks + class MockChunk: + def __init__(self, content): + self.choices = [MagicMock()] + self.choices[0].delta.content = content + + chunks = [ + MockChunk("Hello"), + MockChunk(" "), + MockChunk("world"), + MockChunk("!"), + ] + mock_client.chat.completions.create.return_value = iter(chunks) + + # Wrap the client + wrapped = wrap_openai( + mock_client, + hindsight_api_url="http://localhost:8888", + bank_id="test-agent", + store_conversations=True, # Enable storage + ) + + # Mock the hindsight client + mock_hindsight_client = MagicMock() + with patch.object(wrapped, "_get_hindsight_client", return_value=mock_hindsight_client): + # Call with stream=True + result = wrapped.chat.completions.create( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "Hello"}], + stream=True, + ) + + # Consume all chunks + collected = [] + for chunk in result: + collected.append(chunk) + + # Verify all chunks were yielded + assert len(collected) == 4 + + # Verify retain was called with the complete conversation + mock_hindsight_client.retain.assert_called_once() + call_kwargs = mock_hindsight_client.retain.call_args[1] + assert "USER: Hello" in call_kwargs["content"] + assert "ASSISTANT: Hello world!" in call_kwargs["content"] + + def test_wrap_anthropic_stream_stores_conversation(self): + """Test that streaming stores conversation after all chunks are consumed.""" + from unittest.mock import Mock, MagicMock, patch + from hindsight_litellm.wrappers import wrap_anthropic + + # Create mock Anthropic client + mock_client = Mock() + + # Create mock stream chunks + class MockChunk: + def __init__(self, content): + self.type = "content_block_delta" + self.delta = MagicMock() + self.delta.text = content + + chunks = [ + MockChunk("Hello"), + MockChunk(" "), + MockChunk("world"), + MockChunk("!"), + ] + mock_client.messages.create.return_value = iter(chunks) + + # Wrap the client + wrapped = wrap_anthropic( + mock_client, + hindsight_api_url="http://localhost:8888", + bank_id="test-agent", + store_conversations=True, # Enable storage + ) + + # Mock the hindsight client + mock_hindsight_client = MagicMock() + with patch.object(wrapped, "_get_hindsight_client", return_value=mock_hindsight_client): + # Call with stream=True + result = wrapped.messages.create( + model="claude-sonnet-4-20250514", + max_tokens=1024, + messages=[{"role": "user", "content": "Hello"}], + stream=True, + ) + + # Consume all chunks + collected = [] + for chunk in result: + collected.append(chunk) + + # Verify all chunks were yielded + assert len(collected) == 4 + + # Verify retain was called with the complete conversation + mock_hindsight_client.retain.assert_called_once() + call_kwargs = mock_hindsight_client.retain.call_args[1] + assert "USER: Hello" in call_kwargs["content"] + assert "ASSISTANT: Hello world!" in call_kwargs["content"] diff --git a/hindsight-integrations/litellm/uv.lock b/hindsight-integrations/litellm/uv.lock index 1f50c0e2..6935a138 100644 --- a/hindsight-integrations/litellm/uv.lock +++ b/hindsight-integrations/litellm/uv.lock @@ -691,7 +691,7 @@ wheels = [ [[package]] name = "hindsight-litellm" -version = "0.4.4" +version = "0.4.8" source = { editable = "." } dependencies = [ { name = "aiohttp" },