diff --git a/memori/llm/_base.py b/memori/llm/_base.py index 7732e185..e1ed7a3a 100644 --- a/memori/llm/_base.py +++ b/memori/llm/_base.py @@ -211,9 +211,34 @@ def _format_response(self, raw_response): hasattr(formatted_response, "__dict__") and "_pb" in formatted_response.__dict__ ): + # Old google-generativeai format (protobuf) formatted_response = json.loads( json_format.MessageToJson(formatted_response.__dict__["_pb"]) ) + elif hasattr(formatted_response, "candidates"): + # New google-genai format (dict with candidates) + result = {} + if formatted_response.candidates: + candidates = [] + for candidate in formatted_response.candidates: + candidate_data = {} + if hasattr(candidate, "content") and candidate.content: + content_data = {} + if ( + hasattr(candidate.content, "parts") + and candidate.content.parts + ): + parts = [] + for part in candidate.content.parts: + if hasattr(part, "text"): + parts.append({"text": part.text}) + content_data["parts"] = parts + if hasattr(candidate.content, "role"): + content_data["role"] = candidate.content.role + candidate_data["content"] = content_data + candidates.append(candidate_data) + result["candidates"] = candidates + formatted_response = result else: formatted_response = {} @@ -720,11 +745,39 @@ def process_chunk(self, chunk): formatted_chunk = copy.deepcopy(chunk) if isinstance(self.raw_response, list): if "_pb" in formatted_chunk.__dict__: + # Old google-generativeai format (protobuf) self.raw_response.append( json.loads( json_format.MessageToJson(formatted_chunk.__dict__["_pb"]) ) ) + elif "candidates" in formatted_chunk.__dict__: + # New google-genai format (dict with candidates) + chunk_data = {} + if ( + hasattr(formatted_chunk, "candidates") + and formatted_chunk.candidates + ): + candidates = [] + for candidate in formatted_chunk.candidates: + candidate_data = {} + if hasattr(candidate, "content") and candidate.content: + content_data = {} + if ( + hasattr(candidate.content, "parts") + and candidate.content.parts + ): + parts = [] + for part in candidate.content.parts: + if hasattr(part, "text"): + parts.append({"text": part.text}) + content_data["parts"] = parts + if hasattr(candidate.content, "role"): + content_data["role"] = candidate.content.role + candidate_data["content"] = content_data + candidates.append(candidate_data) + chunk_data["candidates"] = candidates + self.raw_response.append(chunk_data) else: if isinstance(self.raw_response, dict): self.raw_response = merge_chunk(self.raw_response, chunk.__dict__) diff --git a/pyproject.toml b/pyproject.toml index ff961f60..b70a866e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "memori" -version = "3.1.2" +version = "3.1.3" description = "Memori Python SDK" authors = [{name = "Memori Labs Team", email = "noc@memorilabs.ai"}] license = {text = "Apache-2.0"} diff --git a/tests/llm/providers/__init__.py b/tests/llm/providers/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/llm/providers/azure_openai/__init__.py b/tests/llm/providers/azure_openai/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/llm/test_llm_azure_openai.py b/tests/llm/providers/azure_openai/test_azure_openai.py similarity index 100% rename from tests/llm/test_llm_azure_openai.py rename to tests/llm/providers/azure_openai/test_azure_openai.py diff --git a/tests/llm/providers/google_genai/__init__.py b/tests/llm/providers/google_genai/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/llm/providers/google_genai/test_google_genai.py b/tests/llm/providers/google_genai/test_google_genai.py new file mode 100644 index 00000000..61f89a8d --- /dev/null +++ b/tests/llm/providers/google_genai/test_google_genai.py @@ -0,0 +1,227 @@ +""" +Tests for google-genai SDK support. + +This ensures that the new google-genai SDK format (without _pb protobuf) +is properly handled for both streaming and non-streaming responses. + +""" + +import pytest + +from memori._config import Config +from memori.llm._base import BaseInvoke, BaseIterator + + +@pytest.fixture +def config(): + return Config() + + +class MockGoogleGenaiPart: + """Mock google-genai Part object.""" + + def __init__(self, text): + self.text = text + + +class MockGoogleGenaiContent: + """Mock google-genai Content object.""" + + def __init__(self, parts, role="model"): + self.parts = parts + self.role = role + + +class MockGoogleGenaiCandidate: + """Mock google-genai Candidate object.""" + + def __init__(self, content): + self.content = content + + +class MockGoogleGenaiResponse: + """Mock google-genai GenerateContentResponse (non-streaming).""" + + def __init__(self, text, role="model"): + part = MockGoogleGenaiPart(text) + content = MockGoogleGenaiContent([part], role) + self.candidates = [MockGoogleGenaiCandidate(content)] + # Note: No _pb attribute - this is the new format + + +class MockGoogleGenaiChunk: + """Mock google-genai streaming chunk.""" + + def __init__(self, text, role="model"): + part = MockGoogleGenaiPart(text) + content = MockGoogleGenaiContent([part], role) + self.candidates = [MockGoogleGenaiCandidate(content)] + # Note: No _pb attribute - this is the new format + + +class TestGoogleGenaiFormatDetection: + """Tests for detecting google-genai format (no _pb attribute).""" + + def test_response_has_no_pb_attribute(self): + """Verify mock response doesn't have _pb (like real google-genai).""" + response = MockGoogleGenaiResponse("Hello") + assert not hasattr(response, "_pb") + assert "_pb" not in response.__dict__ + + def test_response_has_candidates_attribute(self): + """Verify mock response has candidates (like real google-genai).""" + response = MockGoogleGenaiResponse("Hello") + assert hasattr(response, "candidates") + assert len(response.candidates) == 1 + + def test_chunk_has_no_pb_attribute(self): + """Verify mock chunk doesn't have _pb.""" + chunk = MockGoogleGenaiChunk("Hi") + assert not hasattr(chunk, "_pb") + assert "_pb" not in chunk.__dict__ + + def test_chunk_has_candidates_attribute(self): + """Verify mock chunk has candidates.""" + chunk = MockGoogleGenaiChunk("Hi") + assert hasattr(chunk, "candidates") + + +class TestGoogleGenaiNonStreamingFormat: + """Tests for non-streaming google-genai response formatting.""" + + def test_format_response_with_google_genai_format(self, config): + """Test _format_response handles google-genai format.""" + invoke = BaseInvoke(config, lambda **kwargs: None) + invoke._uses_protobuf = True + + response = MockGoogleGenaiResponse("Hello world", role="model") + formatted = invoke._format_response(response) + + assert "candidates" in formatted + assert len(formatted["candidates"]) == 1 + assert "content" in formatted["candidates"][0] + assert "parts" in formatted["candidates"][0]["content"] + assert ( + formatted["candidates"][0]["content"]["parts"][0]["text"] == "Hello world" + ) + assert formatted["candidates"][0]["content"]["role"] == "model" + + def test_format_response_with_empty_candidates(self, config): + """Test _format_response handles empty candidates.""" + invoke = BaseInvoke(config, lambda **kwargs: None) + invoke._uses_protobuf = True + + response = MockGoogleGenaiResponse("Test") + response.candidates = [] + formatted = invoke._format_response(response) + + # Empty candidates returns empty dict (no content to save) + assert formatted == {} + + def test_format_response_preserves_role(self, config): + """Test that role is preserved in formatted response.""" + invoke = BaseInvoke(config, lambda **kwargs: None) + invoke._uses_protobuf = True + + response = MockGoogleGenaiResponse("Hello", role="model") + formatted = invoke._format_response(response) + + assert formatted["candidates"][0]["content"]["role"] == "model" + + +class TestGoogleGenaiStreamingFormat: + """Tests for streaming google-genai chunk processing.""" + + def test_process_chunk_with_google_genai_format(self, config): + """Test process_chunk handles google-genai chunk format.""" + invoke = BaseInvoke(config, lambda **kwargs: None) + invoke._uses_protobuf = True + + iterator = BaseIterator(config, iter([])) + iterator.invoke = invoke + iterator.raw_response = [] + + chunk = MockGoogleGenaiChunk("Hello", role="model") + iterator.process_chunk(chunk) + + assert len(iterator.raw_response) == 1 + assert "candidates" in iterator.raw_response[0] + assert ( + iterator.raw_response[0]["candidates"][0]["content"]["parts"][0]["text"] + == "Hello" + ) + + def test_process_multiple_chunks(self, config): + """Test processing multiple streaming chunks.""" + invoke = BaseInvoke(config, lambda **kwargs: None) + invoke._uses_protobuf = True + + iterator = BaseIterator(config, iter([])) + iterator.invoke = invoke + iterator.raw_response = [] + + chunks = [ + MockGoogleGenaiChunk("Hello"), + MockGoogleGenaiChunk(" "), + MockGoogleGenaiChunk("World"), + ] + + for chunk in chunks: + iterator.process_chunk(chunk) + + assert len(iterator.raw_response) == 3 + assert ( + iterator.raw_response[0]["candidates"][0]["content"]["parts"][0]["text"] + == "Hello" + ) + assert ( + iterator.raw_response[1]["candidates"][0]["content"]["parts"][0]["text"] + == " " + ) + assert ( + iterator.raw_response[2]["candidates"][0]["content"]["parts"][0]["text"] + == "World" + ) + + def test_process_chunk_preserves_role(self, config): + """Test that role is preserved in processed chunk.""" + invoke = BaseInvoke(config, lambda **kwargs: None) + invoke._uses_protobuf = True + + iterator = BaseIterator(config, iter([])) + iterator.invoke = invoke + iterator.raw_response = [] + + chunk = MockGoogleGenaiChunk("Test", role="model") + iterator.process_chunk(chunk) + + assert iterator.raw_response[0]["candidates"][0]["content"]["role"] == "model" + + +class TestGoogleGenaiBackwardsCompatibility: + """Tests to ensure old google-generativeai format still works.""" + + def test_format_response_with_pb_format(self, config): + """Test _format_response still handles _pb format.""" + + class MockPbResponse: + def __init__(self): + self._pb = None # Would be protobuf in real usage + + invoke = BaseInvoke(config, lambda **kwargs: None) + invoke._uses_protobuf = True + + # When _pb exists but is None, it should try protobuf parsing + # This test just verifies the _pb path is still checked first + response = MockPbResponse() + assert "_pb" in response.__dict__ + + def test_non_protobuf_response_unchanged(self, config): + """Test non-protobuf responses pass through unchanged.""" + invoke = BaseInvoke(config, lambda **kwargs: None) + invoke._uses_protobuf = False + + response = {"choices": [{"message": {"content": "Hello"}}]} + formatted = invoke._format_response(response) + + assert formatted == response