memmachine-server 0.3.2__tar.gz → 0.3.4__tar.gz

{memmachine_server-0.3.2/src/memmachine_server.egg-info → memmachine_server-0.3.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: memmachine-server
-Version: 0.3.2
+Version: 0.3.4
 Summary: MemMachine Server - The complete MemMachine memory system server with episodic and profile memory
 Author-email: MemMachine <noreply@memmachine.ai>
 License-Expression: Apache-2.0

{memmachine_server-0.3.2 → memmachine_server-0.3.4}/server_tests/memmachine_server/common/configuration/test_reranker_conf.py

@@ -7,6 +7,7 @@ from pydantic import SecretStr, ValidationError
 from memmachine_server.common.configuration.reranker_conf import (
     AmazonBedrockRerankerConf,
     BM25RerankerConf,
+    CohereRerankerConf,
     CrossEncoderRerankerConf,
     EmbedderRerankerConf,
     IdentityRerankerConf,
@@ -47,6 +48,18 @@ def amazon_bedrock_reranker_conf() -> dict:
     }
 
 
+@pytest.fixture
+def cohere_reranker_conf() -> dict:
+    return {
+        "provider": "cohere",
+        "config": {
+            "cohere_key": "test-cohere-key",
+            "model": "rerank-english-v3.0",
+            "base_url": "http://localhost:8000",
+        },
+    }
+
+
 @pytest.fixture
 def cross_encoder_reranker_conf() -> dict:
     return {
@@ -83,6 +96,7 @@ def full_reranker_input(
     bm25_reranker_conf,
     identity_reranker_conf,
     amazon_bedrock_reranker_conf,
+    cohere_reranker_conf,
     cross_encoder_reranker_conf,
     embedder_reranker_conf,
     rrf_hybrid_reranker_conf,
@@ -93,6 +107,7 @@ def full_reranker_input(
             "id_ranker_id": identity_reranker_conf,
             "bm_ranker_id": bm25_reranker_conf,
             "aws_reranker_id": amazon_bedrock_reranker_conf,
+            "cohere_reranker_id": cohere_reranker_conf,
             "cross_encoder_id": cross_encoder_reranker_conf,
             "embedder_id": embedder_reranker_conf,
         },
@@ -121,6 +136,13 @@ def test_valid_amazon_bedrock_reranker_conf(amazon_bedrock_reranker_conf):
     assert conf.model_id == "amazon.rerank-v1:0"
 
 
+def test_valid_cohere_reranker_conf(cohere_reranker_conf):
+    conf = CohereRerankerConf(**cohere_reranker_conf["config"])
+    assert conf.cohere_key == SecretStr("test-cohere-key")
+    assert conf.model == "rerank-english-v3.0"
+    assert conf.base_url == "http://localhost:8000"
+
+
 def test_valid_cross_encoder_reranker_conf(cross_encoder_reranker_conf):
     conf = CrossEncoderRerankerConf(**cross_encoder_reranker_conf["config"])
     assert conf.model_name == "cross-encoder/qnli-electra-base"
@@ -150,6 +172,8 @@ def test_full_reranker_conf(full_reranker_input):
 
     assert "aws_reranker_id" in conf.amazon_bedrock
     assert conf.amazon_bedrock["aws_reranker_id"].region == "us-west-2"
+    assert "cohere_reranker_id" in conf.cohere
+    assert conf.cohere["cohere_reranker_id"].base_url == "http://localhost:8000"
 
     assert "cross_encoder_id" in conf.cross_encoder
     assert (
@@ -182,3 +206,12 @@ def test_missing_required_field_in_bedrock_reranker():
         AmazonBedrockRerankerConf(**config)
     assert "missing" in str(exc_info.value)
     assert "aws_secret_access_key" in str(exc_info.value)
+
+
+def test_invalid_cohere_base_url():
+    with pytest.raises(ValidationError, match="Invalid base URL"):
+        CohereRerankerConf(
+            cohere_key=SecretStr("test-cohere-key"),
+            model="rerank-english-v3.0",
+            base_url="localhost:8000",
+        )

memmachine_server-0.3.4/server_tests/memmachine_server/common/embedder/test_openai_embedder.py

@@ -0,0 +1,126 @@
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import openai
+import pytest
+
+from memmachine_server.common.data_types import ExternalServiceAPIError
+from memmachine_server.common.embedder.openai_embedder import (
+    OpenAIEmbedder,
+    OpenAIEmbedderParams,
+)
+
+
+@pytest.fixture(
+    params=[
+        ["Are tomatoes fruits?", "Tomatoes are red."],
+        ["Are tomatoes fruits?", "Tomatoes are red.", ""],
+        ["."],
+        [" "],
+        [""],
+        [],
+    ],
+)
+def inputs(request):
+    return request.param
+
+
+@pytest.mark.integration
+@pytest.mark.asyncio
+async def test_ingest_embed(openai_embedder, inputs):
+    embeddings = await openai_embedder.ingest_embed(inputs)
+    assert isinstance(embeddings, list)
+    assert len(embeddings) == len(inputs)
+    assert all(len(embedding) == openai_embedder.dimensions for embedding in embeddings)
+
+
+@pytest.mark.integration
+@pytest.mark.asyncio
+async def test_search_embed(openai_embedder, inputs):
+    embeddings = await openai_embedder.search_embed(inputs)
+    assert isinstance(embeddings, list)
+    assert len(embeddings) == len(inputs)
+    assert all(len(embedding) == openai_embedder.dimensions for embedding in embeddings)
+
+
+@pytest.mark.integration
+@pytest.mark.asyncio
+async def test_large_input(openai_embedder):
+    input_text = "👩‍💻" * 10000
+
+    assert len(await openai_embedder.ingest_embed([input_text])) == 1
+    assert len(await openai_embedder.search_embed([input_text])) == 1
+
+
+@pytest.mark.integration
+@pytest.mark.asyncio
+async def test_many_inputs(openai_embedder):
+    input_texts = ["Hello, world!"] * 10000
+    assert len(await openai_embedder.ingest_embed(input_texts)) == 10000
+    assert len(await openai_embedder.search_embed(input_texts)) == 10000
+
+
+@pytest.mark.asyncio
+@patch("asyncio.sleep", new_callable=AsyncMock)
+async def test_embed_retry_on_internal_server_error(mock_sleep):
+    """Test retry logic on InternalServerError."""
+    mock_client = AsyncMock(spec=openai.AsyncOpenAI)
+
+    mock_embedding = MagicMock()
+    mock_embedding.embedding = [0.1] * 256
+
+    mock_response = MagicMock()
+    mock_response.data = [mock_embedding]
+    mock_response.usage.prompt_tokens = 10
+    mock_response.usage.total_tokens = 10
+
+    mock_client.embeddings.create = AsyncMock(
+        side_effect=[
+            openai.InternalServerError(
+                "internal server error", response=MagicMock(), body=None
+            ),
+            mock_response,
+        ]
+    )
+
+    embedder = OpenAIEmbedder(
+        OpenAIEmbedderParams(
+            client=mock_client,
+            model="test-model",
+            dimensions=256,
+        )
+    )
+
+    result = await embedder.ingest_embed(["test input"], max_attempts=2)
+
+    assert len(result) == 1
+    assert len(result[0]) == 256
+    assert mock_client.embeddings.create.call_count == 2
+    mock_sleep.assert_awaited_once_with(1)
+
+
+@pytest.mark.asyncio
+@patch("asyncio.sleep", new_callable=AsyncMock)
+async def test_embed_fail_after_max_retries_on_internal_server_error(mock_sleep):
+    """Test that ExternalServiceAPIError is raised after max_attempts on InternalServerError."""
+    mock_client = AsyncMock(spec=openai.AsyncOpenAI)
+    mock_client.embeddings.create = AsyncMock(
+        side_effect=openai.InternalServerError(
+            "internal server error", response=MagicMock(), body=None
+        ),
+    )
+
+    embedder = OpenAIEmbedder(
+        OpenAIEmbedderParams(
+            client=mock_client,
+            model="test-model",
+            dimensions=256,
+        )
+    )
+
+    with pytest.raises(ExternalServiceAPIError, match=r"max attempts"):
+        await embedder.ingest_embed(["test input"], max_attempts=3)
+
+    assert mock_client.embeddings.create.call_count == 3
+    assert mock_sleep.call_count == 2
+    mock_sleep.assert_any_await(1)
+    mock_sleep.assert_any_await(2)

{memmachine_server-0.3.2 → memmachine_server-0.3.4}/server_tests/memmachine_server/common/language_model/test_openai_chat_completions_language_model.py

@@ -1,5 +1,7 @@
 """Unit tests for OpenAIChatCompletionsLanguageModel."""
 
+from dataclasses import dataclass
+from typing import Any, cast
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import httpx
@@ -9,7 +11,7 @@ from openai.types import chat as openai_chat
 from openai.types.chat.chat_completion_message_function_tool_call import (
     Function as ToolCallFunction,
 )
-from pydantic import ValidationError
+from pydantic import BaseModel, ValidationError
 
 from memmachine_server.common.data_types import ExternalServiceAPIError
 from memmachine_server.common.language_model.openai_chat_completions_language_model import (
@@ -19,6 +21,38 @@ from memmachine_server.common.language_model.openai_chat_completions_language_mo
 from memmachine_server.common.metrics_factory import MetricsFactory
 
 
+class FakeAsyncStream:
+    def __init__(self, chunks):
+        self._chunks = iter(chunks)
+
+    def __aiter__(self):
+        return self
+
+    async def __anext__(self):
+        try:
+            return next(self._chunks)
+        except StopIteration as exc:
+            raise StopAsyncIteration from exc
+
+
+@dataclass
+class FakeReasoningChunk:
+    type: str = "response.reasoning.delta"
+
+
+class ParsedResponse(BaseModel):
+    answer: str
+    confidence: float
+
+
+def make_chat_completion(data: dict[str, Any]) -> openai_chat.ChatCompletion:
+    return openai_chat.ChatCompletion.model_validate(data)
+
+
+def make_chat_completion_chunk(data: dict[str, Any]) -> openai_chat.ChatCompletionChunk:
+    return openai_chat.ChatCompletionChunk.model_validate(data)
+
+
 @pytest.fixture
 def mock_metrics_factory():
     """Fixture for a mocked MetricsFactory."""
@@ -213,6 +247,339 @@ async def test_generate_response_success(mock_async_openai, minimal_config):
         {"role": "system", "content": "System prompt"},
         {"role": "user", "content": "User prompt"},
     ]
+    assert call_args.kwargs["store"] is False
+
+
+@pytest.mark.asyncio
+async def test_generate_parsed_response_success(mock_async_openai, minimal_config):
+    """Test a non-streamed structured response is parsed correctly."""
+    mock_response = make_chat_completion(
+        {
+            "id": "completion-1",
+            "object": "chat.completion",
+            "created": 1,
+            "model": "test-model",
+            "choices": [
+                {
+                    "index": 0,
+                    "finish_reason": "stop",
+                    "message": {
+                        "role": "assistant",
+                        "content": '{"answer":"ok","confidence":0.9}',
+                        "refusal": None,
+                    },
+                }
+            ],
+            "usage": {
+                "prompt_tokens": 7,
+                "completion_tokens": 5,
+                "total_tokens": 12,
+            },
+        }
+    )
+
+    mock_client = mock_async_openai.return_value
+    mock_client.chat.completions.create.return_value = mock_response
+
+    lm = OpenAIChatCompletionsLanguageModel(minimal_config)
+    parsed = await lm.generate_parsed_response(
+        ParsedResponse,
+        system_prompt="System prompt",
+        user_prompt="User prompt",
+    )
+
+    assert parsed == ParsedResponse(answer="ok", confidence=0.9)
+    call_args = mock_client.chat.completions.create.call_args
+    assert call_args.kwargs["model"] == "test-model"
+    assert call_args.kwargs["messages"] == [
+        {"role": "system", "content": "System prompt"},
+        {"role": "user", "content": "User prompt"},
+    ]
+    assert call_args.kwargs["response_format"]["type"] == "json_schema"
+
+
+@pytest.mark.asyncio
+async def test_generate_parsed_response_from_stream(
+    mock_async_openai,
+    minimal_config,
+):
+    """Test a streamed structured response is aggregated and validated."""
+    streamed_response = FakeAsyncStream(
+        [
+            FakeReasoningChunk(),
+            make_chat_completion_chunk(
+                {
+                    "id": "chunk-1",
+                    "object": "chat.completion.chunk",
+                    "created": 1,
+                    "model": "test-model",
+                    "choices": [
+                        {
+                            "index": 0,
+                            "delta": {
+                                "role": "assistant",
+                                "content": '{"answer":"ok",',
+                            },
+                            "finish_reason": None,
+                            "logprobs": None,
+                        }
+                    ],
+                    "usage": None,
+                }
+            ),
+            make_chat_completion_chunk(
+                {
+                    "id": "chunk-2",
+                    "object": "chat.completion.chunk",
+                    "created": 2,
+                    "model": "test-model",
+                    "choices": [
+                        {
+                            "index": 0,
+                            "delta": {"content": '"confidence":0.9}'},
+                            "finish_reason": "stop",
+                            "logprobs": None,
+                        }
+                    ],
+                    "usage": None,
+                }
+            ),
+        ]
+    )
+
+    mock_client = mock_async_openai.return_value
+    mock_client.chat.completions.create.return_value = streamed_response
+
+    lm = OpenAIChatCompletionsLanguageModel(minimal_config)
+    parsed = await lm.generate_parsed_response(ParsedResponse)
+
+    assert parsed == ParsedResponse(answer="ok", confidence=0.9)
+
+
+@pytest.mark.asyncio
+async def test_generate_parsed_response_from_stream_discards_reasoning_delta_chunks(
+    mock_async_openai,
+    minimal_config,
+):
+    """Test streamed structured responses ignore chunks with reasoning_content."""
+    reasoning_chunk = make_chat_completion_chunk(
+        {
+            "id": "chunk-0",
+            "object": "chat.completion.chunk",
+            "created": 0,
+            "model": "test-model",
+            "choices": [
+                {
+                    "index": 0,
+                    "delta": {"content": None},
+                    "finish_reason": None,
+                    "logprobs": None,
+                }
+            ],
+            "usage": None,
+        }
+    )
+    cast(Any, reasoning_chunk.choices[0].delta).reasoning_content = "chain-of-thought"
+
+    streamed_response = FakeAsyncStream(
+        [
+            reasoning_chunk,
+            make_chat_completion_chunk(
+                {
+                    "id": "chunk-1",
+                    "object": "chat.completion.chunk",
+                    "created": 1,
+                    "model": "test-model",
+                    "choices": [
+                        {
+                            "index": 0,
+                            "delta": {
+                                "role": "assistant",
+                                "content": '{"answer":"ok",',
+                            },
+                            "finish_reason": None,
+                            "logprobs": None,
+                        }
+                    ],
+                    "usage": None,
+                }
+            ),
+            make_chat_completion_chunk(
+                {
+                    "id": "chunk-2",
+                    "object": "chat.completion.chunk",
+                    "created": 2,
+                    "model": "test-model",
+                    "choices": [
+                        {
+                            "index": 0,
+                            "delta": {"content": '"confidence":0.9}'},
+                            "finish_reason": "stop",
+                            "logprobs": None,
+                        }
+                    ],
+                    "usage": None,
+                }
+            ),
+        ]
+    )
+
+    mock_client = mock_async_openai.return_value
+    mock_client.chat.completions.create.return_value = streamed_response
+
+    lm = OpenAIChatCompletionsLanguageModel(minimal_config)
+    parsed = await lm.generate_parsed_response(ParsedResponse)
+
+    assert parsed == ParsedResponse(answer="ok", confidence=0.9)
+
+
+@pytest.mark.asyncio
+async def test_generate_response_streamed_chat_completion_chunks(
+    mock_async_openai,
+    minimal_config,
+):
+    """Test a streamed chat completion is aggregated into a single response."""
+    streamed_response = FakeAsyncStream(
+        [
+            make_chat_completion_chunk(
+                {
+                    "id": "chunk-1",
+                    "object": "chat.completion.chunk",
+                    "created": 1,
+                    "model": "test-model",
+                    "choices": [
+                        {
+                            "index": 0,
+                            "delta": {"role": "assistant", "content": "Hello"},
+                            "finish_reason": None,
+                            "logprobs": None,
+                        }
+                    ],
+                    "usage": None,
+                }
+            ),
+            make_chat_completion_chunk(
+                {
+                    "id": "chunk-2",
+                    "object": "chat.completion.chunk",
+                    "created": 2,
+                    "model": "test-model",
+                    "choices": [
+                        {
+                            "index": 0,
+                            "delta": {"content": ", world!"},
+                            "finish_reason": "stop",
+                            "logprobs": None,
+                        }
+                    ],
+                    "usage": {
+                        "prompt_tokens": 10,
+                        "completion_tokens": 4,
+                        "total_tokens": 14,
+                    },
+                }
+            ),
+        ]
+    )
+
+    mock_client = mock_async_openai.return_value
+    mock_client.chat.completions.create.return_value = streamed_response
+
+    lm = OpenAIChatCompletionsLanguageModel(minimal_config)
+    (
+        content,
+        tool_calls,
+        input_tokens,
+        output_tokens,
+    ) = await lm.generate_response_with_token_usage()
+
+    assert content == "Hello, world!"
+    assert tool_calls == []
+    assert input_tokens == 10
+    assert output_tokens == 4
+
+
+@pytest.mark.asyncio
+async def test_generate_response_streamed_tool_calls_and_discards_reasoning_chunks(
+    mock_async_openai,
+    minimal_config,
+):
+    """Test streamed tool calls are reconstructed while reasoning chunks are ignored."""
+    streamed_response = FakeAsyncStream(
+        [
+            FakeReasoningChunk(),
+            make_chat_completion_chunk(
+                {
+                    "id": "chunk-1",
+                    "object": "chat.completion.chunk",
+                    "created": 1,
+                    "model": "test-model",
+                    "choices": [
+                        {
+                            "index": 0,
+                            "delta": {
+                                "tool_calls": [
+                                    {
+                                        "index": 0,
+                                        "id": "call_123",
+                                        "type": "function",
+                                        "function": {
+                                            "name": "get_weather",
+                                            "arguments": '{"location"',
+                                        },
+                                    }
+                                ]
+                            },
+                            "finish_reason": None,
+                            "logprobs": None,
+                        }
+                    ],
+                    "usage": None,
+                }
+            ),
+            make_chat_completion_chunk(
+                {
+                    "id": "chunk-2",
+                    "object": "chat.completion.chunk",
+                    "created": 2,
+                    "model": "test-model",
+                    "choices": [
+                        {
+                            "index": 0,
+                            "delta": {
+                                "tool_calls": [
+                                    {
+                                        "index": 0,
+                                        "function": {"arguments": ': "Boston"}'},
+                                    }
+                                ]
+                            },
+                            "finish_reason": "tool_calls",
+                            "logprobs": None,
+                        }
+                    ],
+                    "usage": None,
+                }
+            ),
+        ]
+    )
+
+    mock_client = mock_async_openai.return_value
+    mock_client.chat.completions.create.return_value = streamed_response
+
+    lm = OpenAIChatCompletionsLanguageModel(minimal_config)
+    content, tool_calls = await lm.generate_response()
+
+    assert content == ""
+    assert tool_calls == [
+        {
+            "call_id": "call_123",
+            "function": {
+                "name": "get_weather",
+                "arguments": {"location": "Boston"},
+            },
+        }
+    ]
 
 
 @pytest.mark.asyncio
@@ -397,6 +764,35 @@ async def test_generate_response_fail_after_max_retries(
     mock_sleep.assert_any_await(2)
 
 
+@pytest.mark.asyncio
+@patch("asyncio.sleep", new_callable=AsyncMock)
+async def test_generate_response_retry_on_internal_server_error(
+    mock_sleep,
+    mock_async_openai,
+    minimal_config,
+):
+    """Test retry logic on InternalServerError."""
+    mock_response = MagicMock()
+    mock_response.choices[0].message.content = "Success after retry"
+    mock_response.choices[0].message.tool_calls = None
+    mock_response.usage = None
+
+    mock_client = mock_async_openai.return_value
+    mock_client.chat.completions.create.side_effect = [
+        openai.InternalServerError(
+            "internal server error", response=MagicMock(), body=None
+        ),
+        mock_response,
+    ]
+
+    lm = OpenAIChatCompletionsLanguageModel(minimal_config)
+    content, _ = await lm.generate_response(max_attempts=2)
+
+    assert content == "Success after retry"
+    assert mock_client.chat.completions.create.call_count == 2
+    mock_sleep.assert_awaited_once_with(1)
+
+
 @pytest.mark.asyncio
 @pytest.mark.parametrize(
     "exception",

{memmachine_server-0.3.2 → memmachine_server-0.3.4}/server_tests/memmachine_server/common/language_model/test_openai_responses_language_model.py

@@ -184,6 +184,7 @@ async def test_generate_response_success(mock_async_openai, minimal_config):
         {"role": "system", "content": "System prompt"},
         {"role": "user", "content": "User prompt"},
     ]
+    assert call_args.kwargs["store"] is False
 
 
 @pytest.mark.asyncio
@@ -219,6 +220,8 @@ async def test_generate_response_with_tool_calls(mock_async_openai, minimal_conf
             },
         },
     ]
+    call_args = mock_client.responses.create.call_args
+    assert call_args.kwargs["store"] is False
 
 
 @pytest.mark.asyncio
@@ -353,6 +356,35 @@ async def test_generate_response_fail_after_max_retries(
     mock_sleep.assert_any_await(2)
 
 
+@pytest.mark.asyncio
+@patch("asyncio.sleep", new_callable=AsyncMock)
+async def test_generate_response_retry_on_internal_server_error(
+    mock_sleep,
+    mock_async_openai,
+    minimal_config,
+):
+    """Test retry logic on InternalServerError."""
+    mock_response = MagicMock()
+    mock_response.output_text = "Success after retry"
+    mock_response.output = None
+    mock_response.usage = None
+
+    mock_client = mock_async_openai.return_value
+    mock_client.responses.create.side_effect = [
+        openai.InternalServerError(
+            "internal server error", response=MagicMock(), body=None
+        ),
+        mock_response,
+    ]
+
+    lm = OpenAIResponsesLanguageModel(minimal_config)
+    content, _ = await lm.generate_response(max_attempts=2)
+
+    assert content == "Success after retry"
+    assert mock_client.responses.create.call_count == 2
+    mock_sleep.assert_awaited_once_with(1)
+
+
 @pytest.mark.asyncio
 @pytest.mark.asyncio
 @pytest.mark.parametrize(